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::{