From 42fe650785d3ef4ec919a89c22dda6454649c261 Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Sun, 19 Jul 2026 13:34:12 +0800 Subject: [PATCH] refactor: centralize turn input preparation --- docs/ARCHITECTURE.md | 2 +- src/session/mod.rs | 1 + src/session/session.rs | 153 +++++++++++--------------------------- src/session/turn_input.rs | 148 ++++++++++++++++++++++++++++++++++++ 4 files changed, 192 insertions(+), 112 deletions(-) create mode 100644 src/session/turn_input.rs diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 6b025fb..98956f6 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -183,7 +183,7 @@ Session ID 格式为： 5. 持久化写入由 `persistence_lock` 串行化；多条相关记录应使用 Storage 的原子接口。 6. 内存先变更但持久化失败时，必须回滚精确匹配的消息后缀，不能删除无关的新状态。 -SessionManager 负责组装会话上下文：系统提示、Skills、召回的 Knowledge、压缩后的 Timeline、可选的 active plan 摘要和当前消息历史。普通闲聊 session 没有 plan 摘要；计划状态由 `WorkManager` 从 SQLite 读取，在历史压缩之后追加，因此不以自然语言摘要作为权威来源。`AgentLoop` 接收完整输入执行一次模型/工具循环，本身不拥有会话状态。 +SessionManager 负责组装会话上下文：系统提示、Skills、召回的 Knowledge、压缩后的 Timeline、可选的 active plan 摘要和当前消息历史。`session::turn_input` 在 Session 锁外并行读取 Knowledge、active plan 并压缩历史，然后通过同一个 assembly 路径生成首次请求和 context-overflow 重试输入；重试不得复制系统提示或 runtime context 拼接逻辑。普通闲聊 session 没有 plan 摘要；计划状态由 `WorkManager` 从 SQLite 读取，因此不以自然语言摘要作为权威来源。`AgentLoop` 接收完整输入执行一次模型/工具循环，本身不拥有会话状态。 每个 session 最多有一个 active plan，但 plan 内多个 item 可以分别绑定不同子 Agent 并行执行。主 Agent 通过 `todo` 创建和维护计划，通过 `delegate.plan_item_id` 委托；子 Agent 不获得 `todo` 或 `delegate`。领取 item 使用条件更新防止重复执行，迟到结果只有在 plan 仍 active 且 execution ID 匹配时才允许提交。 diff --git a/src/session/mod.rs b/src/session/mod.rs index 7523334..08b9c1b 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -3,6 +3,7 @@ pub mod error; pub mod events; mod messenger; mod persistence; +mod turn_input; // The public `session::session` path is retained for API compatibility. #[allow(clippy::module_inception)] pub mod session; diff --git a/src/session/session.rs b/src/session/session.rs index c2b7bf2..ac70fc1 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -5,6 +5,7 @@ use tokio::sync::{Mutex, mpsc, oneshot}; use super::persistence::{append_persisted_messages, finalize_turn_after_persistence}; use super::turn::{TurnBlock, TurnController, TurnSnapshot}; +use super::turn_input::prepare_turn_input; use crate::bus::{ ChatMessage, CompletionStatus, MediaItem, MediaRef, MessageSource, OutboundMessage, SourceKind, }; @@ -48,7 +49,7 @@ pub enum HandleResult { AgentProcessing, } use crate::agent::context_compressor::ContextCompressionConfig; -use crate::agent::system_prompt::{build_runtime_context, build_system_prompt}; +use crate::agent::system_prompt::build_system_prompt; use crate::agent::{AgentError, AgentLoop, AgentTurnContext, ContextCompressor, TurnEmitter}; use crate::channels::slash_command::parse_slash_command; use crate::config::BrowserConfig; @@ -791,18 +792,6 @@ impl Session { } } - fn append_runtime_context_to_user_message(message: &mut ChatMessage, runtime_context: &str) { - if runtime_context.trim().is_empty() { - return; - } - - if message.content.trim().is_empty() { - message.content = runtime_context.to_string(); - } else { - message.content = format!("{}\n\n{}", message.content, runtime_context); - } - } - pub fn create_user_message_with_source( &self, content: &str, @@ -2584,7 +2573,14 @@ fn spawn_agent_worker( continue 'tasks; } - let (agent, history_raw, mut compressor, base_version, cancel_rx) = { + let ( + agent, + history_raw, + mut compressor, + system_prompt_out, + base_version, + cancel_rx, + ) = { let mut guard = session.lock().await; if guard.worker_generation != worker_gen { @@ -2623,100 +2619,49 @@ fn spawn_agent_worker( agent, history_raw, guard.fresh_context_compressor(), + guard.build_system_prompt(&skills_prompt), guard.state_version, cancel_rx, ) }; // lock released - let memory_context = match memory_manager - .recall( - &task.content, - 5, - Some(crate::memory::MemoryCategory::Knowledge), - None, - ) - .await - { - Ok(entries) if !entries.is_empty() => Some( - entries - .iter() - .map(|e| format!("- {}: {}", e.key, e.content)) - .collect::>() - .join("\n"), - ), - Err(e) => { - tracing::warn!(error = %e, "Failed to fetch memory context"); - None - } - _ => None, - }; - - let work_context = match work_manager.active_plan(&unified_str).await { - Ok(Some(plan)) => Some(plan.compact_context()), - Ok(None) => None, - Err(error) => { - tracing::warn!(error = %error, "Failed to load active task plan"); - None - } - }; - let runtime_context = build_runtime_context( - Some(unified_str.as_str()), - memory_context.as_deref(), - work_context.as_deref(), - ); - - let system_prompt_out = { - let guard = session.lock().await; + let prepared_input = prepare_turn_input( + memory_manager.clone(), + work_manager.clone(), + &unified_str, + &task.content, + system_prompt_out, + &mut compressor, + history_raw, + ) + .await; + let meta_snapshot = { + let mut guard = session.lock().await; if guard.worker_generation != worker_gen { return; } - guard.build_system_prompt(&skills_prompt) - }; - - let compression_result = compressor.compress_if_needed(history_raw).await; - let mut history_out = match compression_result { - Ok(result) => { - let meta_snapshot = { - let mut guard = session.lock().await; - if guard.worker_generation != worker_gen { - return; - } - if guard.state_version != base_version { - tracing::warn!( - session_id = %guard.id, - "Session changed while preparing agent history; dropping stale task" - ); - guard.current_cancel = None; - continue 'tasks; - } - if result.created_timelines { - guard.last_compressed_message_at = - Some(chrono::Utc::now().timestamp_millis()); - } - guard.last_consolidated_at = - Some(chrono::Utc::now().timestamp_millis()); - guard.session_meta_snapshot() - }; - if let Some((storage, meta)) = meta_snapshot - && let Err(e) = storage.upsert_session(&meta).await - { - tracing::warn!(error = %e, "Failed to persist session meta after compression"); - } - result.history + if guard.state_version != base_version { + tracing::warn!( + session_id = %guard.id, + "Session changed while preparing agent history; dropping stale task" + ); + guard.current_cancel = None; + continue 'tasks; } - Err(e) => { - tracing::warn!(error = %e, "Context compression failed in worker"); - let guard = session.lock().await; - if guard.worker_generation != worker_gen { - return; - } - guard.get_history().to_vec() + if prepared_input.created_timelines { + guard.last_compressed_message_at = + Some(chrono::Utc::now().timestamp_millis()); } + guard.last_consolidated_at = Some(chrono::Utc::now().timestamp_millis()); + guard.session_meta_snapshot() }; - history_out.insert(0, ChatMessage::system(system_prompt_out.clone())); - if let Some(last_msg) = history_out.iter_mut().rev().find(|m| m.role == "user") { - Session::append_runtime_context_to_user_message(last_msg, &runtime_context); + if let Some((storage, meta)) = meta_snapshot + && let Err(e) = storage.upsert_session(&meta).await + { + tracing::warn!(error = %e, "Failed to persist session meta after compression"); } + let history_out = prepared_input.messages; + let runtime_context = prepared_input.runtime; let (turn_controller, turn_emitter, turn_receiver) = TurnController::start( unified_str.clone(), @@ -2867,21 +2812,7 @@ fn spawn_agent_worker( tracing::warn!(error = %e, "Failed to persist session meta after retry compression"); } - let retry_history = { - let mut retry = retry_result.history; - retry.insert( - 0, - ChatMessage::system(system_prompt_out.clone()), - ); - if let Some(last_msg) = retry.iter_mut().rev().find(|m| m.role == "user") - { - Session::append_runtime_context_to_user_message( - last_msg, - &runtime_context, - ); - } - retry - }; + let retry_history = runtime_context.assemble(retry_result.history); match agent .process_streaming(retry_history, agent_turn.clone()) diff --git a/src/session/turn_input.rs b/src/session/turn_input.rs new file mode 100644 index 0000000..8c1ffdc --- /dev/null +++ b/src/session/turn_input.rs @@ -0,0 +1,148 @@ +use std::sync::Arc; + +use crate::agent::ContextCompressor; +use crate::agent::system_prompt::build_runtime_context; +use crate::bus::ChatMessage; +use crate::memory::{MemoryCategory, MemoryManager}; +use crate::work::WorkManager; + +/// Immutable context used to assemble provider input for both the initial call +/// and context-overflow recovery. +pub(super) struct TurnRuntimeContext { + system_prompt: String, + runtime_context: String, +} + +impl TurnRuntimeContext { + pub(super) fn assemble(&self, mut history: Vec) -> Vec { + history.insert(0, ChatMessage::system(self.system_prompt.clone())); + if let Some(last_user) = history + .iter_mut() + .rev() + .find(|message| message.role == "user") + { + append_runtime_context(last_user, &self.runtime_context); + } + history + } +} + +pub(super) struct PreparedTurnInput { + pub(super) messages: Vec, + pub(super) runtime: TurnRuntimeContext, + pub(super) created_timelines: bool, +} + +/// Builds the complete cross-turn provider input outside the Session lock. +/// Independent context sources and compression are fetched concurrently. +pub(super) async fn prepare_turn_input( + memory_manager: Arc, + work_manager: Arc, + session_id: &str, + query: &str, + system_prompt: String, + compressor: &mut ContextCompressor, + history: Vec, +) -> PreparedTurnInput { + let memory_future = memory_manager.recall(query, 5, Some(MemoryCategory::Knowledge), None); + let work_future = work_manager.active_plan(session_id); + let compression_future = compressor.compress_if_needed(history.clone()); + let (memory_result, work_result, compression_result) = + tokio::join!(memory_future, work_future, compression_future); + + let memory_context = match memory_result { + Ok(entries) if !entries.is_empty() => Some( + entries + .iter() + .map(|entry| format!("- {}: {}", entry.key, entry.content)) + .collect::>() + .join("\n"), + ), + Err(error) => { + tracing::warn!(error = %error, "Failed to fetch memory context"); + None + } + _ => None, + }; + let work_context = match work_result { + Ok(Some(plan)) => Some(plan.compact_context()), + Ok(None) => None, + Err(error) => { + tracing::warn!(error = %error, "Failed to load active task plan"); + None + } + }; + let compression = match compression_result { + Ok(result) => result, + Err(error) => { + tracing::warn!(error = %error, "Context compression failed while preparing turn input"); + crate::agent::context_compressor::CompressionResult { + history, + created_timelines: false, + } + } + }; + let runtime = TurnRuntimeContext { + system_prompt, + runtime_context: build_runtime_context( + Some(session_id), + memory_context.as_deref(), + work_context.as_deref(), + ), + }; + + PreparedTurnInput { + messages: runtime.assemble(compression.history), + runtime, + created_timelines: compression.created_timelines, + } +} + +fn append_runtime_context(message: &mut ChatMessage, runtime_context: &str) { + if runtime_context.trim().is_empty() { + return; + } + if message.content.trim().is_empty() { + message.content = runtime_context.to_string(); + } else { + message.content = format!("{}\n\n{}", message.content, runtime_context); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn runtime_context_is_added_only_to_latest_user_message() { + let runtime = TurnRuntimeContext { + system_prompt: "system".to_string(), + runtime_context: "runtime".to_string(), + }; + let messages = runtime.assemble(vec![ + ChatMessage::user("old"), + ChatMessage::assistant("answer"), + ChatMessage::user("new"), + ]); + + assert_eq!(messages[0].role, "system"); + assert_eq!(messages[1].content, "old"); + assert_eq!(messages[3].content, "new\n\nruntime"); + } + + #[test] + fn overflow_reassembly_does_not_duplicate_runtime_context() { + let runtime = TurnRuntimeContext { + system_prompt: "system".to_string(), + runtime_context: "runtime".to_string(), + }; + + let first = runtime.assemble(vec![ChatMessage::user("question")]); + let recovered = runtime.assemble(vec![ChatMessage::user("question")]); + + assert_eq!(first.len(), recovered.len()); + assert_eq!(first[0].content, recovered[0].content); + assert_eq!(first[1].content, recovered[1].content); + assert_eq!(recovered.len(), 2); + } +}