From 2c1effe2c026d7012b89efe809c751a8cdfa4fe3 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Thu, 19 Feb 2026 20:36:25 -0800 Subject: [PATCH] core: generalize interrupted tool-result handling --- codex-rs/core/src/codex.rs | 49 +++++++------------ codex-rs/core/src/state/turn.rs | 10 ---- codex-rs/core/src/stream_events_utils.rs | 3 +- codex-rs/core/src/tools/context.rs | 6 +++ .../src/tools/handlers/request_user_input.rs | 34 +++++++++++++ codex-rs/core/src/tools/js_repl/mod.rs | 27 ++++++++-- codex-rs/core/src/tools/parallel.rs | 13 +++-- codex-rs/core/src/tools/registry.rs | 22 ++++++--- codex-rs/core/src/tools/router.rs | 21 +++++--- 9 files changed, 120 insertions(+), 65 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 0368858f01..9cbee6c1a8 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -270,6 +270,7 @@ use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::context::ToolDispatchOutput; use crate::tools::handlers::SEARCH_TOOL_BM25_TOOL_NAME; use crate::tools::js_repl::JsReplHandle; use crate::tools::js_repl::resolve_compatible_node; @@ -2782,14 +2783,6 @@ impl Session { sub_id: &str, response: RequestUserInputResponse, ) { - if response.interrupted { - let mut active = self.active_turn.lock().await; - if let Some(at) = active.as_mut() { - let mut ts = at.turn_state.lock().await; - ts.mark_request_user_input_interrupted(); - } - } - let entry = { let mut active = self.active_turn.lock().await; match active.as_mut() { @@ -2810,17 +2803,6 @@ impl Session { } } - pub async fn take_request_user_input_interrupted(&self) -> bool { - let mut active = self.active_turn.lock().await; - match active.as_mut() { - Some(at) => { - let mut ts = at.turn_state.lock().await; - ts.take_request_user_input_interrupted() - } - None => false, - } - } - pub async fn notify_dynamic_tool_response(&self, call_id: &str, response: DynamicToolResponse) { let entry = { let mut active = self.active_turn.lock().await; @@ -5054,10 +5036,10 @@ pub(crate) async fn run_turn( Ok(sampling_request_output) => { let SamplingRequestResult { needs_follow_up, - request_user_input_interrupted, + interrupted_tool_result, last_agent_message: sampling_request_last_agent_message, } = sampling_request_output; - if request_user_input_interrupted { + if interrupted_tool_result { cancellation_token.cancel(); sess.finish_turn_without_completion_event(turn_context.as_ref()) .await; @@ -5664,7 +5646,7 @@ async fn built_tools( #[derive(Debug)] struct SamplingRequestResult { needs_follow_up: bool, - request_user_input_interrupted: bool, + interrupted_tool_result: bool, last_agent_message: Option, } @@ -6180,22 +6162,24 @@ async fn handle_assistant_item_done_in_plan_mode( } async fn drain_in_flight( - in_flight: &mut FuturesOrdered>>, + in_flight: &mut FuturesOrdered>>, sess: Arc, turn_context: Arc, -) -> CodexResult<()> { +) -> CodexResult { + let mut interrupted_tool_result = false; while let Some(res) = in_flight.next().await { match res { - Ok(response_input) => { - sess.record_conversation_items(&turn_context, &[response_input.into()]) + Ok(output) => { + sess.record_conversation_items(&turn_context, &[output.response_input.into()]) .await; + interrupted_tool_result |= output.interrupt_turn; } Err(err) => { error_or_panic(format!("in-flight tool future failed during drain: {err}")); } } } - Ok(()) + Ok(interrupted_tool_result) } #[allow(clippy::too_many_arguments)] @@ -6244,7 +6228,7 @@ async fn try_run_sampling_request( Arc::clone(&turn_context), Arc::clone(&turn_diff_tracker), ); - let mut in_flight: FuturesOrdered>> = + let mut in_flight: FuturesOrdered>> = FuturesOrdered::new(); let mut needs_follow_up = false; let mut last_agent_message: Option = None; @@ -6428,7 +6412,7 @@ async fn try_run_sampling_request( break Ok(SamplingRequestResult { needs_follow_up, - request_user_input_interrupted: false, + interrupted_tool_result: false, last_agent_message, }); } @@ -6520,12 +6504,13 @@ async fn try_run_sampling_request( ) .await; - drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?; + let interrupted_tool_result = + drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?; if let Ok(result) = outcome.as_mut() - && sess.take_request_user_input_interrupted().await + && interrupted_tool_result { result.needs_follow_up = false; - result.request_user_input_interrupted = true; + result.interrupted_tool_result = true; } if should_emit_turn_diff { diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index a0efd5cbcb..ccc50d066b 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -73,7 +73,6 @@ pub(crate) struct TurnState { pending_user_input: HashMap>, pending_dynamic_tools: HashMap>, pending_input: Vec, - request_user_input_interrupted: bool, } impl TurnState { @@ -97,7 +96,6 @@ impl TurnState { self.pending_user_input.clear(); self.pending_dynamic_tools.clear(); self.pending_input.clear(); - self.request_user_input_interrupted = false; } pub(crate) fn insert_pending_user_input( @@ -147,14 +145,6 @@ impl TurnState { pub(crate) fn has_pending_input(&self) -> bool { !self.pending_input.is_empty() } - - pub(crate) fn mark_request_user_input_interrupted(&mut self) { - self.request_user_input_interrupted = true; - } - - pub(crate) fn take_request_user_input_interrupted(&mut self) -> bool { - std::mem::take(&mut self.request_user_input_interrupted) - } } impl ActiveTurn { diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index abb4aa86cd..c827c51685 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -14,6 +14,7 @@ use crate::function_tool::FunctionCallError; use crate::memories::citations::get_thread_id_from_citations; use crate::parse_turn_item; use crate::state_db; +use crate::tools::context::ToolDispatchOutput; use crate::tools::parallel::ToolCallRuntime; use crate::tools::router::ToolRouter; use codex_protocol::models::FunctionCallOutputBody; @@ -106,7 +107,7 @@ async fn record_stage1_output_usage_for_completed_item( /// queuing any tool execution futures. This records items immediately so /// history and rollout stay in sync even if the turn is later cancelled. pub(crate) type InFlightFuture<'f> = - Pin> + Send + 'f>>; + Pin> + Send + 'f>>; #[derive(Default)] pub(crate) struct OutputItemResult { diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 0700b4d013..d1fc21e1df 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -74,6 +74,12 @@ pub enum ToolOutput { }, } +#[derive(Clone, Debug)] +pub struct ToolDispatchOutput { + pub response_input: ResponseInputItem, + pub interrupt_turn: bool, +} + impl ToolOutput { pub fn log_preview(&self) -> String { match self { diff --git a/codex-rs/core/src/tools/handlers/request_user_input.rs b/codex-rs/core/src/tools/handlers/request_user_input.rs index f9b444186e..5d01f2e581 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -11,6 +11,7 @@ use crate::tools::registry::ToolKind; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::TUI_VISIBLE_COLLABORATION_MODES; use codex_protocol::request_user_input::RequestUserInputArgs; +use codex_protocol::request_user_input::RequestUserInputResponse; fn request_user_input_is_available(mode: ModeKind, default_mode_request_user_input: bool) -> bool { mode.allows_request_user_input() @@ -121,6 +122,17 @@ impl ToolHandler for RequestUserInputHandler { success: Some(true), }) } + + fn should_interrupt_turn(&self, output: &ToolOutput) -> bool { + let ToolOutput::Function { body, .. } = output else { + return false; + }; + let Some(content) = body.to_text() else { + return false; + }; + serde_json::from_str::(&content) + .is_ok_and(|response| response.interrupted) + } } #[cfg(test)] @@ -171,4 +183,26 @@ mod tests { "Request user input for one to three short questions and wait for the response. This tool is only available in Default or Plan mode.".to_string() ); } + + #[test] + fn interrupted_response_interrupts_turn() { + let handler = RequestUserInputHandler; + let output = ToolOutput::Function { + body: FunctionCallOutputBody::Text(r#"{"answers":{},"interrupted":true}"#.to_string()), + success: Some(true), + }; + + assert!(handler.should_interrupt_turn(&output)); + } + + #[test] + fn non_interrupted_response_does_not_interrupt_turn() { + let handler = RequestUserInputHandler; + let output = ToolOutput::Function { + body: FunctionCallOutputBody::Text(r#"{"answers":{}}"#.to_string()), + success: Some(true), + }; + + assert!(!handler.should_interrupt_turn(&output)); + } } diff --git a/codex-rs/core/src/tools/js_repl/mod.rs b/codex-rs/core/src/tools/js_repl/mod.rs index a234fd0eb8..e852127a20 100644 --- a/codex-rs/core/src/tools/js_repl/mod.rs +++ b/codex-rs/core/src/tools/js_repl/mod.rs @@ -1350,8 +1350,8 @@ impl JsReplManager { ) .await { - Ok(response) => { - if let Some(items) = response_content_items(&response) { + Ok(output) => { + if let Some(items) = response_content_items(&output.response_input) { Self::record_exec_tool_call_content_items( &exec_tool_calls, &req.exec_id, @@ -1360,8 +1360,8 @@ impl JsReplManager { .await; } - let summary = Self::summarize_tool_call_response(&response); - match serde_json::to_value(response) { + let summary = Self::summarize_tool_call_response(&output.response_input); + match serde_json::to_value(output.response_input) { Ok(value) => { Self::log_tool_call_response(&req, true, &summary, Some(&value), None); RunToolResult { @@ -1369,6 +1369,25 @@ impl JsReplManager { ok: true, response: Some(value), error: None, +======= + Ok(output) => { + if let ResponseInputItem::FunctionCallOutput { output, .. } = &output.response_input + && let Some(items) = output.content_items() + { + let mut has_image = false; + let mut content = Vec::with_capacity(items.len()); + for item in items { + match item { + FunctionCallOutputContentItem::InputText { text } => { + content.push(ContentItem::InputText { text: text.clone() }); + } + FunctionCallOutputContentItem::InputImage { image_url } => { + has_image = true; + content.push(ContentItem::InputImage { + image_url: image_url.clone(), + }); + } +>>>>>>> 2fbcb2ec4 (core: generalize interrupted tool-result handling) } } Err(err) => { diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 5f2fc89e5f..34a20d933f 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -14,6 +14,7 @@ use crate::codex::TurnContext; use crate::error::CodexErr; use crate::function_tool::FunctionCallError; use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::context::ToolDispatchOutput; use crate::tools::context::ToolPayload; use crate::tools::router::ToolCall; use crate::tools::router::ToolRouter; @@ -51,7 +52,7 @@ impl ToolCallRuntime { self, call: ToolCall, cancellation_token: CancellationToken, - ) -> impl std::future::Future> { + ) -> impl std::future::Future> { let supports_parallel = self.router.tool_supports_parallel(&call.tool_name); let router = Arc::clone(&self.router); @@ -69,7 +70,7 @@ impl ToolCallRuntime { aborted = false, ); - let handle: AbortOnDropHandle> = + let handle: AbortOnDropHandle> = AbortOnDropHandle::new(tokio::spawn(async move { tokio::select! { _ = cancellation_token.cancelled() => { @@ -113,8 +114,8 @@ impl ToolCallRuntime { } impl ToolCallRuntime { - fn aborted_response(call: &ToolCall, secs: f32) -> ResponseInputItem { - match &call.payload { + fn aborted_response(call: &ToolCall, secs: f32) -> ToolDispatchOutput { + let response_input = match &call.payload { ToolPayload::Custom { .. } => ResponseInputItem::CustomToolCallOutput { call_id: call.call_id.clone(), output: FunctionCallOutputPayload { @@ -133,6 +134,10 @@ impl ToolCallRuntime { ..Default::default() }, }, + }; + ToolDispatchOutput { + response_input, + interrupt_turn: false, } } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 3c021cc582..3127889f51 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -9,6 +9,7 @@ use crate::function_tool::FunctionCallError; use crate::memories::usage::emit_metric_for_tool_read; use crate::protocol::SandboxPolicy; use crate::sandbox_tags::sandbox_tag; +use crate::tools::context::ToolDispatchOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; @@ -20,7 +21,6 @@ use codex_hooks::HookResult; use codex_hooks::HookToolInput; use codex_hooks::HookToolInputLocalShell; use codex_hooks::HookToolKind; -use codex_protocol::models::ResponseInputItem; use codex_utils_readiness::Readiness; use tracing::warn; @@ -53,6 +53,12 @@ pub trait ToolHandler: Send + Sync { /// Perform the actual [ToolInvocation] and returns a [ToolOutput] containing /// the final output to return to the model. async fn handle(&self, invocation: ToolInvocation) -> Result; + + /// Classify whether a successful tool output should interrupt the turn after + /// persisting the tool result to history/rollout. + fn should_interrupt_turn(&self, _output: &ToolOutput) -> bool { + false + } } pub struct ToolRegistry { @@ -79,7 +85,7 @@ impl ToolRegistry { pub async fn dispatch( &self, invocation: ToolInvocation, - ) -> Result { + ) -> Result { let tool_name = invocation.tool_name.clone(); let call_id_owned = invocation.call_id.clone(); let otel = invocation.turn.otel_manager.clone(); @@ -155,7 +161,7 @@ impl ToolRegistry { } let is_mutating = handler.is_mutating(&invocation).await; - let output_cell = tokio::sync::Mutex::new(None); + let output_cell = tokio::sync::Mutex::new(None::<(ToolOutput, bool)>); let invocation_for_tool = invocation.clone(); let started = Instant::now(); @@ -180,8 +186,9 @@ impl ToolRegistry { Ok(output) => { let preview = output.log_preview(); let success = output.success_for_logging(); + let interrupt_turn = handler.should_interrupt_turn(&output); let mut guard = output_cell.lock().await; - *guard = Some(output); + *guard = Some((output, interrupt_turn)); Ok((preview, success)) } Err(err) => Err(err), @@ -213,10 +220,13 @@ impl ToolRegistry { match result { Ok(_) => { let mut guard = output_cell.lock().await; - let output = guard.take().ok_or_else(|| { + let (output, interrupt_turn) = guard.take().ok_or_else(|| { FunctionCallError::Fatal("tool produced no output".to_string()) })?; - Ok(output.into_response(&call_id_owned, &payload_for_response)) + Ok(ToolDispatchOutput { + response_input: output.into_response(&call_id_owned, &payload_for_response), + interrupt_turn, + }) } Err(err) => Err(err), } diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index a55fb5fd5a..f36bdbbb9e 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -5,6 +5,7 @@ use crate::function_tool::FunctionCallError; use crate::mcp_connection_manager::ToolInfo; use crate::sandboxing::SandboxPermissions; use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::context::ToolDispatchOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::ConfiguredToolSpec; @@ -144,7 +145,7 @@ impl ToolRouter { tracker: SharedTurnDiffTracker, call: ToolCall, source: ToolCallSource, - ) -> Result { + ) -> Result { let ToolCall { tool_name, call_id, @@ -161,7 +162,7 @@ impl ToolRouter { "direct tool calls are disabled; use js_repl and codex.tool(...) instead" .to_string(), ); - return Ok(Self::failure_response( + return Ok(Self::failure_output( failure_call_id, payload_outputs_custom, err, @@ -180,7 +181,7 @@ impl ToolRouter { match self.registry.dispatch(invocation).await { Ok(response) => Ok(response), Err(FunctionCallError::Fatal(message)) => Err(FunctionCallError::Fatal(message)), - Err(err) => Ok(Self::failure_response( + Err(err) => Ok(Self::failure_output( failure_call_id, payload_outputs_custom, err, @@ -188,13 +189,13 @@ impl ToolRouter { } } - fn failure_response( + fn failure_output( call_id: String, payload_outputs_custom: bool, err: FunctionCallError, - ) -> ResponseInputItem { + ) -> ToolDispatchOutput { let message = err.to_string(); - if payload_outputs_custom { + let response_input = if payload_outputs_custom { ResponseInputItem::CustomToolCallOutput { call_id, output: codex_protocol::models::FunctionCallOutputPayload { @@ -210,6 +211,10 @@ impl ToolRouter { success: Some(false), }, } + }; + ToolDispatchOutput { + response_input, + interrupt_turn: false, } } } @@ -265,7 +270,7 @@ mod tests { .dispatch_tool_call(session, turn, tracker, call, ToolCallSource::Direct) .await?; - match response { + match response.response_input { ResponseInputItem::FunctionCallOutput { output, .. } => { let content = output.text_content().unwrap_or_default(); assert!( @@ -318,7 +323,7 @@ mod tests { .dispatch_tool_call(session, turn, tracker, call, ToolCallSource::JsRepl) .await?; - match response { + match response.response_input { ResponseInputItem::FunctionCallOutput { output, .. } => { let content = output.text_content().unwrap_or_default(); assert!(