diff --git a/src/session/messenger.rs b/src/session/messenger.rs index 26fad85..2c6212a 100644 --- a/src/session/messenger.rs +++ b/src/session/messenger.rs @@ -2,10 +2,13 @@ use std::collections::HashMap; use crate::bus::{ChatMessage, MediaItem, MessageSource, OutboundMessage, SourceKind}; use crate::session::UnifiedSessionId; -use crate::tools::OutboundMessenger; +use crate::tools::{OutboundDelivery, OutboundMessenger}; -use super::persistence::append_persisted_messages; -use super::session::{CURRENT_SOURCE_SESSION, SessionManager}; +use super::persistence::{append_active_turn_message, append_persisted_messages}; +use super::session::{ + CURRENT_SOURCE_SESSION, CURRENT_TURN_DELIVERIES, CURRENT_TURN_ID, PendingTurnDelivery, + SessionManager, +}; #[async_trait::async_trait] impl OutboundMessenger for SessionManager { @@ -17,7 +20,7 @@ impl OutboundMessenger for SessionManager { content: &str, mut source: MessageSource, media: Vec, - ) -> Result<(), String> { + ) -> Result { if source.from_session.is_none() { source.from_session = CURRENT_SOURCE_SESSION .try_with(|value| value.clone()) @@ -47,6 +50,32 @@ impl OutboundMessenger for SessionManager { let origin = source.from_session.as_deref().unwrap_or("unknown"); let origin_id = source.from_session.clone(); let same_session = source.from_session.as_deref() == Some(target_sid.to_string().as_str()); + let current_turn_id = CURRENT_TURN_ID + .try_with(|value| value.clone()) + .ok() + .flatten(); + let current_turn_deliveries = CURRENT_TURN_DELIVERIES + .try_with(|value| value.clone()) + .ok() + .flatten(); + if same_session + && channel == "cli_chat" + && !media.is_empty() + && let (Some(turn_id), Some(deliveries)) = + (current_turn_id.clone(), current_turn_deliveries) + { + let guard = session.lock().await; + if guard.owns_active_turn(&turn_id) { + deliveries + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(PendingTurnDelivery { + content: content.to_string(), + media, + }); + return Ok(OutboundDelivery::AttachedToCurrentTurn); + } + } let marked_content = if content.trim().is_empty() && !media.is_empty() && same_session { String::new() } else { @@ -55,9 +84,14 @@ impl OutboundMessenger for SessionManager { let message = outbound_history_message(marked_content.clone(), source, &media); let message_id = message.id.clone(); - append_persisted_messages(&session, vec![message]) - .await - .map_err(|error| error.to_string())?; + if same_session && let Some(turn_id) = current_turn_id { + // Ownership is revalidated atomically with the in-memory append; + // a concurrent /stop therefore falls back to a versioned write. + append_active_turn_message(&session, message, turn_id).await + } else { + append_persisted_messages(&session, vec![message]).await + } + .map_err(|error| error.to_string())?; if let Some(origin_id) = origin_id { self.restore_origin_dialog(&origin_id, &target_sid).await; @@ -78,7 +112,8 @@ impl OutboundMessenger for SessionManager { delivery: None, }) .await - .map_err(|error| error.to_string()) + .map_err(|error| error.to_string())?; + Ok(OutboundDelivery::Delivered) } } @@ -108,6 +143,7 @@ impl SessionManager { Vec::new(), ) .await + .map(|_| ()) } } diff --git a/src/session/persistence.rs b/src/session/persistence.rs index f8723a1..5de8725 100644 --- a/src/session/persistence.rs +++ b/src/session/persistence.rs @@ -49,9 +49,42 @@ pub(super) async fn append_persisted_messages( .map(|_| ()) } +enum VersionPolicy { + Advance, + PreserveForOwnedTurn(String), +} + +/// Persist a message emitted as a side effect of the currently active Turn +/// without invalidating that same Turn's captured session version. +/// +/// The task-local Turn ID is revalidated while holding the session lock. If it +/// no longer owns the target session, the write advances `state_version` like +/// any unrelated mutation so stale model work is still rejected. +pub(super) async fn append_active_turn_message( + session: &Arc>, + message: ChatMessage, + turn_id: String, +) -> Result<(), StorageError> { + append_persisted_messages_inner( + session, + vec![message], + VersionPolicy::PreserveForOwnedTurn(turn_id), + ) + .await + .map(|_| ()) +} + pub(super) async fn append_persisted_messages_with_meta( session: &Arc>, messages: Vec, +) -> Result, StorageError> { + append_persisted_messages_inner(session, messages, VersionPolicy::Advance).await +} + +async fn append_persisted_messages_inner( + session: &Arc>, + messages: Vec, + version_policy: VersionPolicy, ) -> Result, StorageError> { if messages.is_empty() { return Ok(Vec::new()); @@ -60,12 +93,19 @@ pub(super) async fn append_persisted_messages_with_meta( let persistence_lock = { session.lock().await.persistence_lock.clone() }; let _persistence_guard = persistence_lock.lock().await; let message_ids: Vec<_> = messages.iter().map(|message| message.id.clone()).collect(); - let snapshots: Vec> = { + let (snapshots, advance_state_version): (Vec>, bool) = { let mut guard = session.lock().await; - messages + let advance_state_version = match &version_policy { + VersionPolicy::Advance => true, + VersionPolicy::PreserveForOwnedTurn(turn_id) => !guard.owns_active_turn(turn_id), + }; + let snapshots = messages .into_iter() - .map(|message| guard.add_message_in_memory(message, true)) - .collect() + .map(|message| { + guard.add_message_in_memory_with_version(message, true, advance_state_version) + }) + .collect(); + (snapshots, advance_state_version) }; let committed = snapshots .iter() @@ -74,7 +114,10 @@ pub(super) async fn append_persisted_messages_with_meta( .collect(); if let Err(error) = persist_added_messages(snapshots).await { - session.lock().await.rollback_message_suffix(&message_ids); + session + .lock() + .await + .rollback_message_suffix_with_version(&message_ids, advance_state_version); return Err(error); } Ok(committed) @@ -109,7 +152,12 @@ where #[cfg(test)] mod tests { use super::*; + use crate::config::LLMProviderConfig; + use crate::memory::MemoryManager; use crate::session::{TurnController, TurnStatus}; + use crate::tools::ToolRegistry; + use std::collections::HashMap; + use std::path::PathBuf; #[tokio::test] async fn completed_is_published_only_after_persistence_succeeds() { @@ -149,4 +197,75 @@ mod tests { Some("failed to persist turn: database down") ); } + + #[tokio::test] + async fn active_turn_side_effect_does_not_invalidate_its_session_version() { + let dir = tempfile::tempdir().unwrap(); + let storage = Arc::new( + crate::storage::Storage::new(&dir.path().join("memory.db")) + .await + .unwrap(), + ); + let memory_manager = Arc::new(MemoryManager::new( + storage, + "test".to_string(), + "test".to_string(), + )); + let config = LLMProviderConfig { + provider_type: "openai".to_string(), + name: "test".to_string(), + base_url: "http://127.0.0.1".to_string(), + api_key: "test".to_string(), + extra_headers: HashMap::new(), + model_id: "test".to_string(), + temperature: None, + max_tokens: None, + model_extra: HashMap::new(), + max_tool_iterations: 1, + token_limit: 8_192, + workspace_dir: PathBuf::from("."), + input_types: vec!["text".to_string(), "image".to_string()], + }; + let session = Arc::new(Mutex::new( + Session::new( + crate::session::UnifiedSessionId::new("cli_chat", "chat", "dialog"), + config, + Arc::new(ToolRegistry::new()), + None, + String::new(), + "test".to_string(), + memory_manager, + ) + .await + .unwrap(), + )); + let base_version = session.lock().await.state_version_for_test(); + + session.lock().await.set_active_turn_for_test("turn-1"); + append_active_turn_message( + &session, + ChatMessage::assistant("sent screenshot"), + "turn-1".to_string(), + ) + .await + .unwrap(); + + { + let guard = session.lock().await; + assert_eq!(guard.state_version_for_test(), base_version); + assert_eq!(guard.get_history().len(), 1); + } + + append_active_turn_message( + &session, + ChatMessage::assistant("late screenshot"), + "stale-turn".to_string(), + ) + .await + .unwrap(); + + let guard = session.lock().await; + assert_eq!(guard.state_version_for_test(), base_version + 1); + assert_eq!(guard.get_history().len(), 2); + } } diff --git a/src/session/session.rs b/src/session/session.rs index c861b8a..31a6d81 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -73,6 +73,117 @@ fn committed_turn_delta( tokio::task_local! { pub(super) static CURRENT_SOURCE_SESSION: Option; + pub(super) static CURRENT_TURN_ID: Option; + pub(super) static CURRENT_TURN_DELIVERIES: Option>>>; +} + +pub(super) struct PendingTurnDelivery { + pub content: String, + pub media: Vec, +} + +fn take_pending_turn_deliveries( + deliveries: &Arc>>, +) -> Vec { + std::mem::take( + &mut *deliveries + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()), + ) +} + +fn take_current_turn_deliveries() -> Vec { + CURRENT_TURN_DELIVERIES + .try_with(|deliveries| { + deliveries + .as_ref() + .map(take_pending_turn_deliveries) + .unwrap_or_default() + }) + .unwrap_or_default() +} + +fn collapse_pending_turn_deliveries(pending: Vec) -> (String, Vec) { + let fallback_content = pending + .iter() + .map(|delivery| delivery.content.trim()) + .filter(|content| !content.is_empty()) + .collect::>() + .join("\n\n"); + let mut media_refs = Vec::new(); + for media_ref in pending + .into_iter() + .flat_map(|delivery| delivery.media.into_iter().map(|media| media.to_media_ref())) + { + if !media_refs.iter().any(|existing: &MediaRef| { + existing.path == media_ref.path && existing.media_type == media_ref.media_type + }) { + media_refs.push(media_ref); + } + } + (fallback_content, media_refs) +} + +fn attach_pending_to_message( + message: &mut ChatMessage, + fallback_content: &str, + media_refs: &[MediaRef], +) { + if message.content.trim().is_empty() && !fallback_content.is_empty() { + message.content = fallback_content.to_string(); + } + for media_ref in media_refs { + if !message.media_refs.iter().any(|existing| { + existing.path == media_ref.path && existing.media_type == media_ref.media_type + }) { + message.media_refs.push(media_ref.clone()); + } + } +} + +fn attach_pending_turn_deliveries( + result: &mut crate::agent::AgentProcessResult, + pending: Vec, +) { + let (fallback_content, media_refs) = collapse_pending_turn_deliveries(pending); + if media_refs.is_empty() && fallback_content.is_empty() { + return; + } + + attach_pending_to_message(&mut result.final_response, &fallback_content, &media_refs); + + if let Some(final_message) = result + .emitted_messages + .iter_mut() + .rev() + .find(|message| message.id == result.final_response.id) + { + final_message + .content + .clone_from(&result.final_response.content); + final_message + .media_refs + .clone_from(&result.final_response.media_refs); + } +} + +fn partial_assistant_with_pending_deliveries( + snapshot: &TurnSnapshot, + completion_status: CompletionStatus, + pending: Vec, +) -> Option { + let (fallback_content, media_refs) = collapse_pending_turn_deliveries(pending); + let mut message = partial_assistant_message(snapshot, completion_status).or_else(|| { + (!media_refs.is_empty()).then(|| { + let mut message = ChatMessage::assistant(""); + message.id = snapshot.message_id.clone(); + message.turn_id = Some(snapshot.id.0.clone()); + message.completion_status = completion_status; + message + }) + })?; + attach_pending_to_message(&mut message, &fallback_content, &media_refs); + Some(message) } /// Result of handling a message - either an AI response or a command output @@ -182,7 +293,11 @@ async fn fail_turn_with_partial( error: String, ) { let snapshot = controller.snapshot(); - let partial = partial_assistant_message(&snapshot, CompletionStatus::Interrupted); + let partial = partial_assistant_with_pending_deliveries( + &snapshot, + CompletionStatus::Interrupted, + take_current_turn_deliveries(), + ); if let Some(partial) = partial { controller.begin_finalizing(); if let Err(persistence_error) = append_persisted_messages(session, vec![partial]).await { @@ -369,6 +484,60 @@ mod cancelled_partial_tests { .is_none() ); } + + #[test] + fn pending_same_turn_media_is_attached_to_the_final_response() { + let final_message = ChatMessage::assistant("截图已经准备好了"); + let final_id = final_message.id.clone(); + let mut result = crate::agent::AgentProcessResult { + final_response: final_message.clone(), + emitted_messages: vec![ + ChatMessage::tool("call-1", "send_message", "附件已加入当前回复"), + final_message, + ], + total_tokens: None, + usage: None, + }; + + attach_pending_turn_deliveries( + &mut result, + vec![PendingTurnDelivery { + content: "这是百度首页截图".to_string(), + media: vec![MediaItem::new("/tmp/baidu.png", "image")], + }], + ); + + assert_eq!(result.final_response.content, "截图已经准备好了"); + assert_eq!(result.final_response.media_refs.len(), 1); + assert_eq!(result.final_response.media_refs[0].path, "/tmp/baidu.png"); + let committed_final = result + .emitted_messages + .iter() + .find(|message| message.id == final_id) + .unwrap(); + assert_eq!(committed_final.media_refs.len(), 1); + assert_eq!(committed_final.media_refs[0].path, "/tmp/baidu.png"); + } + + #[test] + fn pending_media_survives_a_turn_without_partial_text() { + let (controller, _emitter, _) = TurnController::start("session", "message-id"); + + let message = partial_assistant_with_pending_deliveries( + &controller.snapshot(), + CompletionStatus::Interrupted, + vec![PendingTurnDelivery { + content: "这是已生成的截图".to_string(), + media: vec![MediaItem::new("/tmp/baidu.png", "image")], + }], + ) + .unwrap(); + + assert_eq!(message.id, "message-id"); + assert_eq!(message.content, "这是已生成的截图"); + assert_eq!(message.completion_status, CompletionStatus::Interrupted); + assert_eq!(message.media_refs.len(), 1); + } } use crate::bus::MessageBus; use crate::providers::{LLMProvider, create_provider}; @@ -710,10 +879,11 @@ impl Session { self.id.to_string() } - pub(super) fn add_message_in_memory( + pub(super) fn add_message_in_memory_with_version( &mut self, message: ChatMessage, persist: bool, + advance_state_version: bool, ) -> Option { let is_user = message.role == "user"; let now = chrono::Utc::now().timestamp_millis(); @@ -768,7 +938,9 @@ impl Session { self.message_count += 1; } self.last_active_at = now; - self.state_version = self.state_version.wrapping_add(1); + if advance_state_version { + self.state_version = self.state_version.wrapping_add(1); + } persist_snapshot.map(|(storage, session_id, msg_meta)| { let session_meta = crate::storage::session::SessionMeta { @@ -797,7 +969,11 @@ impl Session { /// Roll back messages that were appended in memory but whose atomic /// persistence failed. This is only called while holding the session lock, /// so the suffix check also protects against removing unrelated messages. - pub(super) fn rollback_message_suffix(&mut self, message_ids: &[String]) { + pub(super) fn rollback_message_suffix_with_version( + &mut self, + message_ids: &[String], + advance_state_version: bool, + ) { if message_ids.is_empty() || self.messages.len() < message_ids.len() { return; } @@ -819,7 +995,29 @@ impl Session { self.seq_counter -= message_ids.len() as i64; self.total_message_count -= message_ids.len() as i64; self.message_count -= removed_user_messages; - self.state_version = self.state_version.wrapping_add(1); + if advance_state_version { + self.state_version = self.state_version.wrapping_add(1); + } + } + + pub(super) fn owns_active_turn(&self, turn_id: &str) -> bool { + self.active_turn_emitter + .as_ref() + .is_some_and(|active| active.turn_id == turn_id) + } + + #[cfg(test)] + pub(super) fn state_version_for_test(&self) -> u64 { + self.state_version + } + + #[cfg(test)] + pub(super) fn set_active_turn_for_test(&mut self, turn_id: &str) { + let (_controller, emitter, _) = TurnController::start(self.id.to_string(), "test-message"); + self.active_turn_emitter = Some(ActiveTurnEmitter { + turn_id: turn_id.to_string(), + emitter, + }); } /// 获取消息历史 @@ -2767,6 +2965,8 @@ fn spawn_agent_worker( let commit_delivery = turn_delivery.clone(); let commit_target = turn_target.clone(); let turn_lifecycle = &turn_controller; + let pending_turn_deliveries = Arc::new(std::sync::Mutex::new(Vec::new())); + let scoped_turn_deliveries = pending_turn_deliveries.clone(); let process_future = async move { let response_session_id = unified_str2.clone(); let process_result = crate::agent::sub_agent::DELEGATE_CONTEXT.scope( @@ -2777,7 +2977,7 @@ fn spawn_agent_worker( }, agent.process_streaming(history_out.clone(), agent_turn.clone()), ).await; - let result = match process_result { + let mut result = match process_result { Ok(r) => r, Err(AgentError::LlmError(ref msg)) if is_context_overflow_error(msg) => @@ -2921,6 +3121,8 @@ fn spawn_agent_worker( } }; + let pending = take_current_turn_deliveries(); + attach_pending_turn_deliveries(&mut result, pending); let response_content = result.final_response.content; let total_tokens = result.total_tokens; let usage = result.usage; @@ -3005,15 +3207,20 @@ fn spawn_agent_worker( let _ = bus2.publish_outbound(outbound).await; } }; + let process_future = + CURRENT_TURN_ID.scope(Some(active_turn_id.clone()), process_future); + let process_future = CURRENT_TURN_DELIVERIES + .scope(Some(scoped_turn_deliveries), process_future); tokio::select! { () = process_future => {} _ = cancel_rx => { // cancelled — current_cancel already taken by /stop let snapshot = turn_controller.snapshot(); - if let Some(partial) = partial_assistant_message( + if let Some(partial) = partial_assistant_with_pending_deliveries( &snapshot, CompletionStatus::Cancelled, + take_pending_turn_deliveries(&pending_turn_deliveries), ) { turn_controller.begin_finalizing(); match append_persisted_messages(&session, vec![partial]).await { diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 9ddb076..0e789c4 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -41,7 +41,7 @@ pub use pty::{PtyManager, PtyTool}; pub use registry::ToolRegistry; pub use send_message::SendMessageTool; pub use todo::TodoTool; -pub use traits::{OutboundMessenger, Tool, ToolResult, ToolResultWithMedia}; +pub use traits::{OutboundDelivery, OutboundMessenger, Tool, ToolResult, ToolResultWithMedia}; pub use web_fetch::WebFetchTool; use crate::agent::SubAgentManager; diff --git a/src/tools/send_message.rs b/src/tools/send_message.rs index 9a507ee..f3e1624 100644 --- a/src/tools/send_message.rs +++ b/src/tools/send_message.rs @@ -6,7 +6,7 @@ use mime_guess::mime; use crate::bus::{MediaItem, MessageSource, SourceKind}; -use super::traits::{OutboundMessenger, Tool, ToolResult}; +use super::traits::{OutboundDelivery, OutboundMessenger, Tool, ToolResult}; pub struct SendMessageTool { messenger: Arc, @@ -144,11 +144,16 @@ target_chat_id 支持两种格式::(发送到该聊天下 .send_message(channel, chat_id, dialog_id, content, source, media) .await { - Ok(()) => Ok(ToolResult { + Ok(OutboundDelivery::Delivered) => Ok(ToolResult { success: true, output: "消息已发送".to_string(), error: None, }), + Ok(OutboundDelivery::AttachedToCurrentTurn) => Ok(ToolResult { + success: true, + output: "附件已加入当前回复".to_string(), + error: None, + }), Err(e) => Ok(ToolResult { success: false, output: String::new(), diff --git a/src/tools/traits.rs b/src/tools/traits.rs index 560dd08..a364f94 100644 --- a/src/tools/traits.rs +++ b/src/tools/traits.rs @@ -16,6 +16,12 @@ pub struct ToolResultWithMedia { pub media_refs: Vec, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutboundDelivery { + Delivered, + AttachedToCurrentTurn, +} + impl From for ToolResultWithMedia { fn from(result: ToolResult) -> Self { Self { @@ -67,5 +73,5 @@ pub trait OutboundMessenger: Send + Sync { content: &str, source: MessageSource, media: Vec, - ) -> Result<(), String>; + ) -> Result; }