From 120a02a6e393ebeb298d4c805af657c5dfddaf60 Mon Sep 17 00:00:00 2001 From: hufenghong <19103376390@163.com> Date: Tue, 28 Jul 2026 21:53:35 +0800 Subject: [PATCH] fix(aionrs): scope tool call IDs by conversation --- .../src/capability/backend_output_sink.rs | 54 ++++++++++++++---- .../src/manager/aionrs/agent.rs | 3 +- .../src/repository/sqlite_conversation.rs | 56 +++++++++++++++++-- 3 files changed, 96 insertions(+), 17 deletions(-) diff --git a/crates/aionui-ai-agent/src/capability/backend_output_sink.rs b/crates/aionui-ai-agent/src/capability/backend_output_sink.rs index 087295aff..373855126 100644 --- a/crates/aionui-ai-agent/src/capability/backend_output_sink.rs +++ b/crates/aionui-ai-agent/src/capability/backend_output_sink.rs @@ -7,20 +7,24 @@ use crate::protocol::events::{ }; pub struct BackendOutputSink { + conversation_id: String, event_tx: broadcast::Sender, } impl BackendOutputSink { - pub fn new(event_tx: broadcast::Sender) -> Self { - Self { event_tx } + pub fn new(conversation_id: String, event_tx: broadcast::Sender) -> Self { + Self { + conversation_id, + event_tx, + } } - fn internal_call_id(tool_use_id: &str) -> Option { + fn internal_call_id(&self, tool_use_id: &str) -> Option { let id = tool_use_id.trim(); if id.is_empty() { None } else { - Some(format!("aionrs-{id}")) + Some(format!("aionrs-{}-{id}", self.conversation_id)) } } } @@ -42,7 +46,7 @@ impl OutputSink for BackendOutputSink { } fn emit_tool_call(&self, tool_use_id: &str, name: &str, input: &str) { - let Some(call_id) = Self::internal_call_id(tool_use_id) else { + let Some(call_id) = self.internal_call_id(tool_use_id) else { tracing::error!(tool = name, "Cannot emit tool_call with empty tool_use_id"); return; }; @@ -80,7 +84,7 @@ impl OutputSink for BackendOutputSink { } fn emit_tool_result(&self, tool_use_id: &str, name: &str, is_error: bool, content: &str) { - let Some(call_id) = Self::internal_call_id(tool_use_id) else { + let Some(call_id) = self.internal_call_id(tool_use_id) else { tracing::error!(tool = name, "Cannot emit tool_result with empty tool_use_id"); return; }; @@ -158,9 +162,11 @@ impl OutputSink for BackendOutputSink { mod tests { use super::*; + const CONVERSATION_ID: &str = "conv-test"; + fn make_sink() -> (BackendOutputSink, broadcast::Receiver) { let (tx, rx) = broadcast::channel(16); - (BackendOutputSink::new(tx), rx) + (BackendOutputSink::new(CONVERSATION_ID.to_string(), tx), rx) } #[test] @@ -260,14 +266,38 @@ mod tests { } } - assert_eq!(call_ids[0].0, "aionrs-call_a"); - assert_eq!(call_ids[1].0, "aionrs-call_b"); - assert_eq!(call_ids[2].0, "aionrs-call_a"); - assert_eq!(call_ids[3].0, "aionrs-call_b"); + assert_eq!(call_ids[0].0, "aionrs-conv-test-call_a"); + assert_eq!(call_ids[1].0, "aionrs-conv-test-call_b"); + assert_eq!(call_ids[2].0, "aionrs-conv-test-call_a"); + assert_eq!(call_ids[3].0, "aionrs-conv-test-call_b"); assert_eq!(call_ids[2].1, ToolCallStatus::Completed); assert_eq!(call_ids[3].1, ToolCallStatus::Completed); } + #[test] + fn identical_tool_use_ids_are_namespaced_by_conversation() { + let (first_tx, mut first_rx) = broadcast::channel(1); + let (second_tx, mut second_rx) = broadcast::channel(1); + let first = BackendOutputSink::new("conv-a".to_string(), first_tx); + let second = BackendOutputSink::new("conv-b".to_string(), second_tx); + + first.emit_tool_call("Glob_0", "Glob", r#"{"pattern":"*.rs"}"#); + second.emit_tool_call("Glob_0", "Glob", r#"{"pattern":"*.rs"}"#); + + let first_id = match first_rx.try_recv().unwrap() { + AgentStreamEvent::ToolCall(data) => data.call_id, + other => panic!("Expected ToolCall, got {other:?}"), + }; + let second_id = match second_rx.try_recv().unwrap() { + AgentStreamEvent::ToolCall(data) => data.call_id, + other => panic!("Expected ToolCall, got {other:?}"), + }; + + assert_eq!(first_id, "aionrs-conv-a-Glob_0"); + assert_eq!(second_id, "aionrs-conv-b-Glob_0"); + assert_ne!(first_id, second_id); + } + #[test] fn emit_stream_start_sends_start_event() { let (sink, mut rx) = make_sink(); @@ -370,7 +400,7 @@ mod tests { #[test] fn no_panic_when_no_receivers() { let (tx, _) = broadcast::channel(16); - let sink = BackendOutputSink::new(tx); + let sink = BackendOutputSink::new(CONVERSATION_ID.to_string(), tx); sink.emit_text_delta("hello", "msg-1"); sink.emit_thinking("thought", "msg-1"); sink.emit_tool_call("call_read_1", "Read", "{}"); diff --git a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs index 8e5c9677a..5d46c34d1 100644 --- a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs +++ b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs @@ -142,7 +142,8 @@ impl AionrsAgentManager { resume_session: Option, ) -> Result { let runtime = AgentRuntime::new(conversation_id.clone(), workspace.clone(), 128); - let sink: Arc = Arc::new(BackendOutputSink::new(runtime.event_sender())); + let sink: Arc = + Arc::new(BackendOutputSink::new(conversation_id.clone(), runtime.event_sender())); let runtime_env = config_extra.runtime_env.clone(); let image_input_override = config_extra.compat_overrides.image_input; let image_input_capability = image_input_override.unwrap_or_else(|| { diff --git a/crates/aionui-db/src/repository/sqlite_conversation.rs b/crates/aionui-db/src/repository/sqlite_conversation.rs index e01c8c803..4ba5aafb2 100644 --- a/crates/aionui-db/src/repository/sqlite_conversation.rs +++ b/crates/aionui-db/src/repository/sqlite_conversation.rs @@ -70,9 +70,9 @@ impl SqliteConversationRepository { Ok(()) } - async fn upsert_message_once(&self, message: &MessageRow) -> Result<(), sqlx::Error> { + async fn upsert_message_once(&self, message: &MessageRow) -> Result<(), DbError> { let mut tx = self.pool.begin().await?; - sqlx::query( + let result = sqlx::query( "INSERT INTO messages \ (id, conversation_id, msg_id, type, content, position, \ status, hidden, created_at) \ @@ -100,7 +100,8 @@ impl SqliteConversationRepository { END, \ position = COALESCE(messages.position, excluded.position), \ hidden = excluded.hidden, \ - created_at = MIN(messages.created_at, excluded.created_at)", + created_at = MIN(messages.created_at, excluded.created_at) \ + WHERE messages.conversation_id = excluded.conversation_id", ) .bind(&message.id) .bind(&message.conversation_id) @@ -114,6 +115,14 @@ impl SqliteConversationRepository { .execute(&mut *tx) .await?; + if result.rows_affected() == 0 { + tx.rollback().await?; + return Err(DbError::Conflict(format!( + "message id '{}' belongs to another conversation", + message.id + ))); + } + bump_conversation_updated_at(&mut tx, &message.conversation_id, message.created_at).await?; tx.commit().await?; @@ -692,7 +701,7 @@ impl IConversationRepository for SqliteConversationRepository { } async fn upsert_message(&self, message: &MessageRow) -> Result<(), DbError> { - self.upsert_message_once(message).await.map_err(DbError::from) + self.upsert_message_once(message).await } async fn update_message(&self, id: &str, updates: &MessageRowUpdate) -> Result<(), DbError> { @@ -1178,6 +1187,45 @@ mod tests { ); } + #[tokio::test] + async fn upsert_message_rejects_cross_conversation_id_collision() { + let (repo, _db) = setup().await; + let first_conversation = sample_conversation(SYSTEM_USER_ID); + let mut second_conversation = sample_conversation(SYSTEM_USER_ID); + second_conversation.updated_at = 1; + repo.create(&first_conversation).await.unwrap(); + repo.create(&second_conversation).await.unwrap(); + + let mut first_message = sample_message(&first_conversation.id); + first_message.id = "shared-tool-call-id".to_string(); + first_message.content = r#"{"name":"Glob","input":{"pattern":"*.rs"}}"#.to_string(); + repo.upsert_message(&first_message).await.unwrap(); + + let mut colliding_message = sample_message(&second_conversation.id); + colliding_message.id = first_message.id.clone(); + colliding_message.content = r#"{"name":"Glob","input":{"pattern":"*.md"}}"#.to_string(); + let error = repo.upsert_message(&colliding_message).await.unwrap_err(); + + assert!(matches!(error, DbError::Conflict(message) if message.contains("shared-tool-call-id"))); + let persisted = repo + .get_message(&first_conversation.id, &first_message.id) + .await + .unwrap() + .unwrap(); + assert_eq!(persisted.content, first_message.content); + assert!( + repo.get_message(&second_conversation.id, &colliding_message.id) + .await + .unwrap() + .is_none() + ); + assert_eq!( + repo.get(&second_conversation.id).await.unwrap().unwrap().updated_at, + second_conversation.updated_at, + "a rejected cross-conversation upsert must not bump conversation recency" + ); + } + #[tokio::test] async fn update_conversation_name() { let (repo, _db) = setup().await;