core: generalize interrupted tool-result handling

This commit is contained in:
Charles Cunningham
2026-03-02 18:01:22 -08:00
parent a5bcf4a7b8
commit 2c1effe2c0
9 changed files with 120 additions and 65 deletions
+17 -32
View File
@@ -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<String>,
}
@@ -6180,22 +6162,24 @@ async fn handle_assistant_item_done_in_plan_mode(
}
async fn drain_in_flight(
in_flight: &mut FuturesOrdered<BoxFuture<'static, CodexResult<ResponseInputItem>>>,
in_flight: &mut FuturesOrdered<BoxFuture<'static, CodexResult<ToolDispatchOutput>>>,
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
) -> CodexResult<()> {
) -> CodexResult<bool> {
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<BoxFuture<'static, CodexResult<ResponseInputItem>>> =
let mut in_flight: FuturesOrdered<BoxFuture<'static, CodexResult<ToolDispatchOutput>>> =
FuturesOrdered::new();
let mut needs_follow_up = false;
let mut last_agent_message: Option<String> = 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 {
-10
View File
@@ -73,7 +73,6 @@ pub(crate) struct TurnState {
pending_user_input: HashMap<String, oneshot::Sender<RequestUserInputResponse>>,
pending_dynamic_tools: HashMap<String, oneshot::Sender<DynamicToolResponse>>,
pending_input: Vec<ResponseInputItem>,
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 {
+2 -1
View File
@@ -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<Box<dyn Future<Output = Result<ResponseInputItem>> + Send + 'f>>;
Pin<Box<dyn Future<Output = Result<ToolDispatchOutput>> + Send + 'f>>;
#[derive(Default)]
pub(crate) struct OutputItemResult {
+6
View File
@@ -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 {
@@ -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::<RequestUserInputResponse>(&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));
}
}
+23 -4
View File
@@ -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) => {
+9 -4
View File
@@ -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<Output = Result<ResponseInputItem, CodexErr>> {
) -> impl std::future::Future<Output = Result<ToolDispatchOutput, CodexErr>> {
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<Result<ResponseInputItem, FunctionCallError>> =
let handle: AbortOnDropHandle<Result<ToolDispatchOutput, FunctionCallError>> =
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,
}
}
+16 -6
View File
@@ -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<ToolOutput, FunctionCallError>;
/// 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<ResponseInputItem, FunctionCallError> {
) -> Result<ToolDispatchOutput, FunctionCallError> {
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),
}
+13 -8
View File
@@ -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<ResponseInputItem, FunctionCallError> {
) -> Result<ToolDispatchOutput, FunctionCallError> {
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!(