From 510ed56cd4f3a32e40d53d93080dabf12575d50f Mon Sep 17 00:00:00 2001 From: wsp Date: Tue, 10 Nov 2026 20:37:41 +0800 Subject: [PATCH] fix(session): repair incomplete tool history on restore Persisted sessions can contain an assistant tool call without its result when the application exits during tool execution. Repair the final active round during restore and before ordinary execution by appending an explicit unknown-result marker, without rerunning the interrupted tool or waiting during shutdown. - Rewrite repaired context snapshots at the newest represented turn index - Keep normalization idempotent and shared by restore, execution, and fork paths - Add regression coverage for in-memory and persisted snapshot repair --- .../{fork_agent => }/context_normalization.rs | 37 +++++- .../src/agentic/coordination/coordinator.rs | 4 +- .../src/agentic/execution/execution_engine.rs | 21 +++- .../core/src/agentic/fork_agent/mod.rs | 6 - src/crates/assembly/core/src/agentic/mod.rs | 1 + .../src/agentic/session/session_manager.rs | 118 +++++++++++++++++- 6 files changed, 173 insertions(+), 14 deletions(-) rename src/crates/assembly/core/src/agentic/{fork_agent => }/context_normalization.rs (81%) diff --git a/src/crates/assembly/core/src/agentic/fork_agent/context_normalization.rs b/src/crates/assembly/core/src/agentic/context_normalization.rs similarity index 81% rename from src/crates/assembly/core/src/agentic/fork_agent/context_normalization.rs rename to src/crates/assembly/core/src/agentic/context_normalization.rs index 334188fbec..0ceac4ba52 100644 --- a/src/crates/assembly/core/src/agentic/fork_agent/context_normalization.rs +++ b/src/crates/assembly/core/src/agentic/context_normalization.rs @@ -1,3 +1,9 @@ +//! Model-context normalization at provider-history boundaries. +//! +//! Session restore, ordinary execution, interruption recovery, and forked +//! execution all share the same invariant: every assistant tool call visible +//! to a provider must have a corresponding tool result. + use crate::agentic::core::{Message, MessageContent, ToolResult}; /// Ensure a fork never carries an assistant tool call without a matching @@ -23,15 +29,17 @@ pub(crate) fn normalize_fork_context_messages(mut messages: Vec) -> Vec /// malformed older history with an incomplete *earlier* round is not repaired /// by this fast path. Snapshots without round metadata are treated as legacy /// one-round contexts and use the same logic over the whole snapshot. -pub(crate) fn normalize_incomplete_tool_calls(messages: &mut Vec, result_text: &str) { +pub(crate) fn normalize_incomplete_tool_calls( + messages: &mut Vec, + result_text: &str, +) -> bool { let Some(last_round_id) = messages .iter() .rev() .find_map(|message| message.metadata.round_id.as_deref()) .map(str::to_string) else { - normalize_round_messages(messages, result_text, None); - return; + return normalize_round_messages(messages, result_text, None); }; let last_index = messages @@ -45,15 +53,16 @@ pub(crate) fn normalize_incomplete_tool_calls(messages: &mut Vec, resul start -= 1; } let mut round_messages = messages.split_off(start); - normalize_round_messages(&mut round_messages, result_text, Some(&last_round_id)); + let changed = normalize_round_messages(&mut round_messages, result_text, Some(&last_round_id)); messages.extend(round_messages); + changed } fn normalize_round_messages( messages: &mut Vec, result_text: &str, round_id: Option<&str>, -) { +) -> bool { let result_ids = messages .iter() .filter_map(|message| match &message.content { @@ -62,6 +71,7 @@ fn normalize_round_messages( }) .collect::>(); let mut normalized = Vec::with_capacity(messages.len()); + let mut changed = false; for message in messages.drain(..) { let missing = match &message.content { MessageContent::Mixed { tool_calls, .. } => tool_calls @@ -80,6 +90,7 @@ fn normalize_round_messages( .collect::>(), _ => Vec::new(), }; + changed |= !missing.is_empty(); normalized.push(message); normalized.extend(missing.into_iter().map(|result| { let message = Message::tool_result(result); @@ -90,6 +101,7 @@ fn normalize_round_messages( })); } *messages = normalized; + changed } #[cfg(test)] @@ -148,4 +160,19 @@ mod tests { assert_eq!(messages.len(), 2); } + + #[test] + fn normalization_reports_changes_and_is_idempotent() { + let call = ToolCall { + tool_id: "call-1".to_string(), + tool_name: "Read".to_string(), + arguments: serde_json::json!({}), + ..Default::default() + }; + let mut messages = vec![Message::assistant_with_tools(String::new(), vec![call])]; + + assert!(normalize_incomplete_tool_calls(&mut messages, "unknown")); + assert!(!normalize_incomplete_tool_calls(&mut messages, "unknown")); + assert_eq!(messages.len(), 2); + } } diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 9bfd4cff03..1ea5f5d2a9 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -11934,7 +11934,9 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .get_context_messages(parent_session_id) .await?; let context_messages = - crate::agentic::fork_agent::normalize_fork_context_messages(context_messages); + crate::agentic::context_normalization::normalize_fork_context_messages( + context_messages, + ); ForkAgentContextSnapshot::from_parent_session(&parent_session, context_messages) } diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index cb5e303588..f8ea8bf67d 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -17,6 +17,7 @@ use crate::agentic::agents::{ render_direct_tool_listing_body, PrependedPromptReminders, PromptBuilder, PromptBuilderContext, RuntimeContextNeeds, ToolListingSections, UserContextPolicy, UserContextSection, }; +use crate::agentic::context_normalization::normalize_incomplete_tool_calls; use crate::agentic::coordination::scheduler::agent_dialog_turn_image_contexts; use crate::agentic::core::{ render_system_reminder, InternalReminderKind, Message, MessageContent, MessageHelper, @@ -3270,7 +3271,7 @@ impl ExecutionEngine { start_time: std::time::Instant, ) -> OpenBitFunResult { let dialog_turn_id = context.dialog_turn_id.clone(); - let initial_count = initial_messages.len(); + let mut initial_messages = initial_messages; debug!( "Executing dialog turn implementation: dialog_turn_id={}", @@ -3311,6 +3312,24 @@ impl ExecutionEngine { OpenBitFunError::Session(format!("Session not found: {}", context.session_id)) })?; + // A restored process may have persisted an assistant tool call after + // model output but before the tool result. Repair that boundary before + // constructing any provider request and persist the canonical context so + // every subsequent request sees a valid call/result pair. + if normalize_incomplete_tool_calls( + &mut initial_messages, + "Tool execution was still in progress when the application closed; the result is unknown. Verify any external side effects before continuing.", + ) { + self.session_manager + .replace_context_messages(&context.session_id, initial_messages.clone()) + .await; + info!( + "Normalized incomplete tool calls before dialog execution: session_id={}, turn_id={}", + context.session_id, context.dialog_turn_id + ); + } + let initial_count = initial_messages.len(); + if context.subagent_parent_info.is_none() && super::round_model_routing::configuration::session_eligible(&session) && super::round_model_routing::configuration::release_notice_required( diff --git a/src/crates/assembly/core/src/agentic/fork_agent/mod.rs b/src/crates/assembly/core/src/agentic/fork_agent/mod.rs index f69835c267..c9b5fd11be 100644 --- a/src/crates/assembly/core/src/agentic/fork_agent/mod.rs +++ b/src/crates/assembly/core/src/agentic/fork_agent/mod.rs @@ -7,12 +7,6 @@ use crate::agentic::core::{Message, Session, SessionConfig}; use crate::util::errors::{OpenBitFunError, OpenBitFunResult}; -mod context_normalization; - -pub(crate) use context_normalization::{ - normalize_fork_context_messages, normalize_incomplete_tool_calls, -}; - /// Immutable snapshot of a parent session's runtime context at fork time. #[derive(Debug, Clone)] pub struct ForkAgentContextSnapshot { diff --git a/src/crates/assembly/core/src/agentic/mod.rs b/src/crates/assembly/core/src/agentic/mod.rs index 7364d40e24..c66659237f 100644 --- a/src/crates/assembly/core/src/agentic/mod.rs +++ b/src/crates/assembly/core/src/agentic/mod.rs @@ -4,6 +4,7 @@ //! here until it can be split without changing tool, session, or review flows. // Core module +pub(crate) mod context_normalization; pub mod core; pub mod events; pub mod persistence; diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index f8a30f59ba..1441831db0 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -3,13 +3,13 @@ //! Responsible for session CRUD, lifecycle management, and resource association use crate::agentic::agents::get_agent_registry; +use crate::agentic::context_normalization::normalize_incomplete_tool_calls; use crate::agentic::core::{ new_turn_id, CompressionContract, CompressionState, InternalReminderKind, Message, MessageContent, MessageRole, MessageSemanticKind, ProcessingPhase, Session, SessionAgentRouteOwner, SessionConfig, SessionKind, SessionModelBindingPolicy, SessionState, SessionSummary, ToolPresentationMode, TurnStats, }; -use crate::agentic::fork_agent::normalize_incomplete_tool_calls; use crate::agentic::image_analysis::ImageContextData; use crate::agentic::keyed_lock::{KeyedAsyncLock, KeyedAsyncLockGuard}; use crate::agentic::memories::db::{MemoryDatabase, MEMORY_PHASE2_GLOBAL_JOB_KEY}; @@ -6595,6 +6595,36 @@ impl SessionManager { } }; + // A process can exit after the assistant tool-call message is persisted + // and before the tool result is appended. Providers reject that history + // shape, so close the final active round during ordinary restore. The + // repaired snapshot becomes the durable recovery boundary for the next + // process as well; no interrupted tool is re-executed here. + let repaired_incomplete_tool_calls = normalize_incomplete_tool_calls( + &mut messages, + "Tool execution was still in progress when the application closed; the result is unknown. Verify any external side effects before continuing.", + ); + if repaired_incomplete_tool_calls { + let repaired_snapshot_turn_index = latest_turn_index + .into_iter() + .chain(persisted_turns.last().map(|turn| turn.turn_index)) + .max(); + if let Some(turn_index) = repaired_snapshot_turn_index { + self.persistence_manager + .save_turn_context_snapshot( + session_storage_path, + session_id, + turn_index, + &messages, + ) + .await?; + info!( + "Normalized incomplete tool calls during session restore: session_id={}, turn_index={}", + session_id, turn_index + ); + } + } + if messages.is_empty() { debug!( "Session {} has empty persisted messages (may be new session)", @@ -18379,6 +18409,92 @@ mod tests { .is_empty()); } + #[tokio::test] + async fn restore_session_repairs_and_persists_incomplete_tool_calls() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = test_manager(persistence_manager.clone()); + let session_id = Uuid::new_v4().to_string(); + let mut session = Session::new_with_id( + session_id.clone(), + "Interrupted tool call".to_string(), + "Standard".to_string(), + SessionConfig { + tool_presentation_mode: Default::default(), + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + session.dialog_turn_ids = vec!["turn-1".to_string()]; + persistence_manager + .save_session(workspace.path(), &session) + .await + .expect("session should save"); + + let mut turn = DialogTurnData::new( + "turn-1".to_string(), + 0, + session_id.clone(), + UserMessageData { + id: "turn-1-user".to_string(), + content: "read the file".to_string(), + timestamp: 1, + metadata: None, + }, + ); + turn.mark_completed(); + persistence_manager + .save_dialog_turn(workspace.path(), &turn) + .await + .expect("turn should save"); + + let assistant = Message::assistant_with_tools( + String::new(), + vec![ToolCall { + tool_id: "call-1".to_string(), + tool_name: "Read".to_string(), + arguments: json!({"path": "a.txt"}), + ..Default::default() + }], + ) + .with_turn_id("turn-1".to_string()) + .with_round_id("round-1".to_string()); + persistence_manager + .save_turn_context_snapshot(workspace.path(), &session_id, 0, &[assistant]) + .await + .expect("incomplete snapshot should save"); + + manager + .restore_session(workspace.path(), &session_id) + .await + .expect("session should restore"); + + let restored = manager + .get_context_messages(&session_id) + .await + .expect("context should load"); + assert_eq!(restored.len(), 2); + assert!(matches!( + &restored[1].content, + MessageContent::ToolResult { tool_id, result, .. } + if tool_id == "call-1" + && result == &json!("Tool execution was still in progress when the application closed; the result is unknown. Verify any external side effects before continuing.") + )); + + let persisted = persistence_manager + .load_turn_context_snapshot(workspace.path(), &session_id, 0) + .await + .expect("repaired snapshot should load") + .expect("repaired snapshot should exist"); + assert_eq!(persisted.len(), 2); + assert!(matches!( + &persisted[1].content, + MessageContent::ToolResult { tool_id, .. } if tool_id == "call-1" + )); + } + #[tokio::test] async fn restore_session_view_tail_with_staged_revert_reads_only_visible_tail() { use crate::agentic::session::revert::{