From 1bf73324cadc72a53ed467edc7d3fd2b145a6166 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Fri, 25 Sep 2026 16:11:48 +0000 Subject: [PATCH] Add opt-in code-mode yielding on new user input (#48135) ## Why User input queued during a long-running code-mode `exec` or `wait` can remain pending until the call returns. Allow these calls to yield early so Codex can receive the new input while the cell continues running. ## What changed - Add the `instant_interrupt` feature flag, disabled by default. - Watch queued user input for each sampling request and pass a shared preemption signal to code-mode `exec` and `wait` calls, including calls received later in the same response. - Yield running cell IDs without stopping the cells, allowing subsequent `wait` calls to collect their results. ## Testing Add integration coverage for enabled and disabled behavior, repeated steering during `exec` and `wait`, input deferred during compaction, and later calls in the same response. Verify that direct tool results and queued user messages are preserved, and add a scenario snapshot showing continuation after a cell yields. GitOrigin-RevId: e33466a980d98e113bfb6c6a45c59de0cc7b3c1d --- codex-rs/code-mode-host/tests/stdio.rs | 1 + codex-rs/core/config.schema.json | 6 + codex-rs/core/src/session/input_queue.rs | 31 + codex-rs/core/src/session/mod.rs | 5 + codex-rs/core/src/session/step_context.rs | 3 + codex-rs/core/src/session/tests.rs | 5 + codex-rs/core/src/session/turn.rs | 7 + codex-rs/core/src/tools/code_mode/mod.rs | 6 +- .../core/src/tools/code_mode/wait_handler.rs | 11 +- codex-rs/core/tests/suite/pending_input.rs | 681 ++++++++++++++++++ codex-rs/core/tests/suite/scenarios.rs | 133 ++++ ...os__astra_input_yields_code_mode_cell.snap | 120 +++ codex-rs/features/src/lib.rs | 8 + 13 files changed, 1011 insertions(+), 6 deletions(-) create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_input_yields_code_mode_cell.snap diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs index e9a0fa805a..3636f6b419 100644 --- a/codex-rs/code-mode-host/tests/stdio.rs +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -246,6 +246,7 @@ async fn interrupt_yields_observations_without_stopping_the_cell() { description: String::new(), kind: CodeModeToolKind::Function, input_schema: None, + input_schema_max_bytes: None, output_schema: None, }]; let signal = CancellationToken::new(); diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index c8ad1a3194..58418764a6 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -894,6 +894,9 @@ "in_app_updates": { "type": "boolean" }, + "instant_interrupt": { + "type": "boolean" + }, "item_ids": { "type": "boolean" }, @@ -6926,6 +6929,9 @@ "in_app_updates": { "type": "boolean" }, + "instant_interrupt": { + "type": "boolean" + }, "item_ids": { "type": "boolean" }, diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 6ddef17167..c70ff75008 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -14,6 +14,8 @@ use std::collections::VecDeque; use std::sync::Arc; use tokio::sync::Mutex; use tokio::sync::watch; +use tokio_util::sync::CancellationToken; +use tokio_util::task::AbortOnDropHandle; static PENDING_MAILBOX_MESSAGES: Gauge = Gauge::new("core.mailbox.pending"); @@ -215,6 +217,29 @@ impl InputQueue { }) } + /// Signal once a user message is queued for this sampling request. + pub(crate) async fn watch_user_input( + &self, + active_turn: &Mutex>, + sub_id: &str, + interrupt: CancellationToken, + ) -> Option> { + let turn_state = self.turn_state_for_sub_id(active_turn, sub_id).await?; + // Subscribe before inspecting the queue so an arrival cannot be missed. + let mut activity = self.activity_tx.subscribe(); + Some(AbortOnDropHandle::new(tokio::spawn(async move { + loop { + if turn_state.lock().await.pending_input.has_user_input() { + interrupt.cancel(); + return; + } + if activity.changed().await.is_err() { + return; + } + } + }))) + } + /// Clear any pending waiters and input buffered for the current turn. pub(crate) async fn clear_pending(&self, active_turn: &ActiveTurn) { let mut turn_state = active_turn.turn_state.lock().await; @@ -363,6 +388,12 @@ impl InputQueue { } impl TurnInputQueue { + fn has_user_input(&self) -> bool { + self.items + .iter() + .any(|input| matches!(input, TurnInput::UserInput { .. })) + } + pub(crate) fn is_empty(&self) -> bool { self.items.is_empty() } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6d1580bcc4..a04ce2fdf2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3942,6 +3942,11 @@ impl Session { ) = prepared_tools??; turn_context.extension_data.insert(selected_plugins); Ok(Arc::new(StepContext { + preempt: turn_context + .config + .features + .enabled(Feature::InstantInterrupt) + .then(CancellationToken::new), realtime: self.conversation.snapshot().await, settings, token_budget, diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 62c387509c..781c18fd50 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -15,10 +15,13 @@ use codex_mcp::McpBinding; use codex_otel::SessionTelemetry; use codex_protocol::items::ModelInvocationContext; use codex_protocol::protocol::TurnContextItem; +use tokio_util::sync::CancellationToken; /// Request-scoped state that may change between model sampling requests. pub(crate) struct StepContext { pub(crate) turn: Arc, + /// Yields this request's code-mode observations when a user message arrives. + pub(crate) preempt: Option, /// Realtime call activity and instructions captured for this sampling request. pub(crate) realtime: RealtimeConversationSnapshot, /// One immutable settings version captured before request preparation. diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 1b43d3932c..ad0167f65b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -268,6 +268,11 @@ impl StepContext { }); settings.service_tier = turn.config.service_tier.clone(); Arc::new(Self { + preempt: turn + .config + .features + .enabled(Feature::InstantInterrupt) + .then(tokio_util::sync::CancellationToken::new), token_budget: token_budget::resolve_token_budget( turn.configured_token_budget.as_ref(), turn.use_model_token_budget_defaults, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 4e86837831..d4de229150 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1588,6 +1588,13 @@ async fn run_sampling_request( cancellation_token: CancellationToken, ) -> CodexResult<(SamplingRequestResult, Vec)> { let turn_context = Arc::clone(&step_context.turn); + let _input_watch = if let Some(preempt) = &step_context.preempt { + sess.input_queue + .watch_user_input(&sess.active_turn, &turn_context.sub_id, preempt.clone()) + .await + } else { + None + }; let base_instructions = sess.get_prompt_base_instructions().await; let tool_runtime = ToolCallRuntime::new( diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 49615443e8..79a95b2ac2 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -133,21 +133,23 @@ impl CodeModeService { request .yield_time_ms .get_or_insert(self.default_exec_yield_time_ms); + let preempt = step_context.preempt.clone(); let delegate = Arc::new(CodeModeCellDelegate { broker: Arc::clone(&self.dispatch_broker), step_context, }); self.session() .await? - .execute(request, delegate, /*preempt*/ None) + .execute(request, delegate, preempt) .await } pub(crate) async fn wait( &self, request: codex_code_mode::WaitRequest, + preempt: Option, ) -> Result { - self.session().await?.wait(request, /*preempt*/ None).await + self.session().await?.wait(request, preempt).await } pub(crate) async fn terminate( diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 0a7a633b7a..b0e4479134 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -134,10 +134,13 @@ impl CodeModeWaitHandler { exec.session .services .code_mode_service - .wait(codex_code_mode::WaitRequest { - cell_id, - yield_time_ms: args.yield_time_ms, - }) + .wait( + codex_code_mode::WaitRequest { + cell_id, + yield_time_ms: args.yield_time_ms, + }, + step_context.preempt.clone(), + ) .await } .map_err(|error| { diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 0453bf811f..4be888f715 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use codex_core::CodexThread; use codex_core::StartIfIdleSubmission; +use codex_core::StartThreadOptions; use codex_core::TurnInput; use codex_core::TurnInputRequest; use codex_core::TurnInputSubmission; @@ -17,6 +18,10 @@ use codex_protocol::AgentPath; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; +use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; +use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; +use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::items::TurnItem; use codex_protocol::models::AgentMessageInputContent; use codex_protocol::models::ContentItem; @@ -36,6 +41,7 @@ use core_test_support::context_snapshot::SnapshotEntry; use core_test_support::responses; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_completed_with_tokens; +use core_test_support::responses::ev_custom_tool_call; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_message_item_added; @@ -56,6 +62,11 @@ use serde_json::from_slice; use serde_json::json; use test_case::test_case; use tokio::sync::oneshot; +use wiremock::Mock; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path_regex; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn idle_user_input_reaches_the_first_model_request() -> anyhow::Result<()> { @@ -1113,6 +1124,676 @@ async fn injected_response_item_reopens_turn_after_final_answer() { server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn steers_during_tool_drain_preserve_tool_output_and_each_input() { + const TOOL_CALL_ID: &str = "held-tool"; + let (release_response, response_gate) = oneshot::channel(); + let first_chunks = vec![ + chunk(ev_response_created("resp-1")), + chunk(ev_function_call(TOOL_CALL_ID, "held_tool", "{}")), + gated_chunk(response_gate, vec![ev_completed("resp-1")]), + ]; + let (server, _) = + start_streaming_sse_server(vec![first_chunks, response_completed_chunks("resp-2")]).await; + let test = test_codex() + .with_model("gpt-5.4") + .with_config(|config| { + config + .features + .enable(Feature::InstantInterrupt) + .expect("enable InstantInterrupt feature"); + }) + .build_with_streaming_server(&server) + .await + .expect("build Codex test session"); + let codex = test + .thread_manager + .start_thread(StartThreadOptions { + dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "held_tool".to_string(), + description: "A tool held by the test until both steers are queued.".to_string(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })], + ..StartThreadOptions::new(test.config.clone()) + }) + .await + .expect("start thread with dynamic tool") + .thread; + + submit_user_input(&codex, "first prompt").await; + let EventMsg::DynamicToolCallRequest(tool_request) = wait_for_event(&codex, |event| { + matches!(event, EventMsg::DynamicToolCallRequest(request) if request.call_id == TOOL_CALL_ID) + }) + .await else { + unreachable!("predicate guarantees tool request"); + }; + + steer_user_input(&codex, "second prompt").await; + steer_user_input(&codex, "third prompt").await; + release_response + .send(()) + .expect("release the model response"); + assert!( + tokio::time::timeout( + std::time::Duration::from_millis(200), + server.wait_for_request_count(/*count*/ 2), + ) + .await + .is_err(), + "the replacement request must wait for the direct tool" + ); + codex + .submit(Op::DynamicToolResponse { + id: tool_request.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "held tool result".to_string(), + }], + success: true, + }, + }) + .await + .expect("complete dynamic tool"); + wait_for_event(&codex, |event| { + assert!(!matches!( + event, + EventMsg::TurnAborted(_) | EventMsg::StreamError(_) + )); + matches!(event, EventMsg::TurnComplete(completed) if completed.error.is_none()) + }) + .await; + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let second: Value = from_slice(&requests[1]).expect("parse replacement request"); + let prompts = message_input_texts(&second, "user") + .into_iter() + .filter(|text| { + matches!( + text.as_str(), + "first prompt" | "second prompt" | "third prompt" + ) + }) + .collect::>(); + assert_eq!( + prompts, + vec!["first prompt", "second prompt", "third prompt"] + ); + let result = function_call_output_text(&second, TOOL_CALL_ID).expect("real tool output"); + assert!(result.contains("held tool result"), "tool output: {result}"); + assert_eq!( + second["input"] + .as_array() + .expect("request input is an array") + .iter() + .filter(|item| item["type"] == "function_call_output" && item["call_id"] == TOOL_CALL_ID) + .count(), + 1, + "the direct tool result should be recorded exactly once" + ); + + server.shutdown().await; +} + +#[test_case(false; "disabled")] +#[test_case(true; "enabled")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn steers_yield_exec_and_wait_without_stopping_the_cell(instant_interrupt: bool) { + const EXEC_ID: &str = "exec-call"; + const WAIT_ONE_ID: &str = "wait-one"; + const WAIT_TWO_ID: &str = "wait-two"; + let server = responses::start_mock_server().await; + let test = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(move |config| { + if instant_interrupt { + config + .features + .enable(Feature::InstantInterrupt) + .expect("enable InstantInterrupt feature"); + } else { + config + .features + .disable(Feature::InstantInterrupt) + .expect("disable InstantInterrupt feature"); + } + config + .features + .enable(Feature::CodeMode) + .expect("enable CodeMode feature"); + config + .features + .enable(Feature::CodeModeHost) + .expect("enable CodeModeHost feature"); + config.code_mode.disable_in_process_fallback = true; + config + .features + .disable(Feature::EnableRequestCompression) + .expect("disable request compression"); + config.model_provider.supports_websockets = false; + }) + .build_with_auto_env(&server) + .await + .expect("build Codex test session"); + let codex = test + .thread_manager + .start_thread(StartThreadOptions { + dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "held_tool".into(), + description: "A tool held until the test responds.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })], + ..StartThreadOptions::new(test.config.clone()) + }) + .await + .expect("start thread with dynamic tool") + .thread; + + responses::mount_sse_once( + &server, + responses::sse(vec![ + ev_response_created("resp-exec"), + ev_custom_tool_call( + EXEC_ID, + "exec", + "// @exec: {\"yield_time_ms\": 60000}\nawait tools.held_tool({}); text('background result');", + ), + ev_completed("resp-exec"), + ]), + ) + .await; + + // The cell ID comes from the real yielded exec result, so do not predict a + // runtime generation or rely on a fixed ID when issuing subsequent waits. + let cell_id = Arc::new(std::sync::Mutex::new(None::)); + if instant_interrupt { + let cell_id_for_first_wait = Arc::clone(&cell_id); + Mock::given(method("POST")) + .and(path_regex(".*/responses$")) + .respond_with(move |request: &Request| { + let body: Value = from_slice(&request.body).expect("parse replacement request"); + let output = call_output_text(&body, "custom_tool_call_output", EXEC_ID) + .expect("yielded exec output"); + let call_items = body["input"] + .as_array() + .expect("request input is an array") + .iter() + .filter(|item| item.get("call_id").is_some()) + .cloned() + .collect::>(); + let id = output + .strip_prefix("Script running with cell ID ") + .and_then(|rest| rest.lines().next()) + .unwrap_or_else(|| { + panic!("missing running cell ID in exec output: {output:?}; call items: {call_items:?}") + }) + .to_string(); + *cell_id_for_first_wait.lock().expect("lock cell ID") = Some(id.clone()); + ResponseTemplate::new(200).set_body_raw( + responses::sse(vec![ + ev_response_created("resp-wait-one"), + ev_function_call( + WAIT_ONE_ID, + "wait", + &json!({"cell_id": id, "yield_time_ms": 60000}).to_string(), + ), + ev_completed("resp-wait-one"), + ]), + "text/event-stream", + ) + }) + .up_to_n_times(1) + .mount(&server) + .await; + let cell_id_for_second_wait = Arc::clone(&cell_id); + Mock::given(method("POST")) + .and(path_regex(".*/responses$")) + .respond_with(move |_: &Request| { + let id = cell_id_for_second_wait + .lock() + .expect("lock cell ID") + .clone() + .expect("yielded cell ID"); + ResponseTemplate::new(200).set_body_raw( + responses::sse(vec![ + ev_response_created("resp-wait-two"), + ev_function_call( + WAIT_TWO_ID, + "wait", + &json!({"cell_id": id, "yield_time_ms": 60000}).to_string(), + ), + ev_completed("resp-wait-two"), + ]), + "text/event-stream", + ) + }) + .up_to_n_times(1) + .mount(&server) + .await; + } + let final_response = + responses::mount_sse_once(&server, responses::sse_completed("resp-done")).await; + + submit_user_input(&codex, "start the cell").await; + let EventMsg::DynamicToolCallRequest(tool_request) = wait_for_event(&codex, |event| { + matches!(event, EventMsg::DynamicToolCallRequest(request) if request.tool == "held_tool") + }) + .await else { + unreachable!("predicate guarantees tool request"); + }; + steer_user_input(&codex, "interrupt exec").await; + if !instant_interrupt { + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + assert!( + final_response.requests().is_empty(), + "a disabled interrupt must leave exec waiting for the nested tool" + ); + codex + .submit(Op::DynamicToolResponse { + id: tool_request.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "released".into(), + }], + success: true, + }, + }) + .await + .expect("release nested tool"); + wait_for_turn_complete(&codex).await; + let request = final_response.single_request().body_json(); + assert!( + call_output_text(&request, "custom_tool_call_output", EXEC_ID) + .expect("exec output") + .contains("background result") + ); + assert!(message_input_texts(&request, "user").contains(&"interrupt exec".to_string())); + return; + } + wait_for_event(&codex, |event| { + matches!(event, EventMsg::RawResponseItem(raw) if matches!(&raw.item, ResponseItem::FunctionCall { call_id, .. } if call_id == WAIT_ONE_ID)) + }) + .await; + steer_user_input(&codex, "interrupt wait").await; + wait_for_event(&codex, |event| { + matches!(event, EventMsg::RawResponseItem(raw) if matches!(&raw.item, ResponseItem::FunctionCall { call_id, .. } if call_id == WAIT_TWO_ID)) + }) + .await; + codex + .submit(Op::DynamicToolResponse { + id: tool_request.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "released".into(), + }], + success: true, + }, + }) + .await + .expect("release nested tool"); + wait_for_event(&codex, |event| { + assert!(!matches!( + event, + EventMsg::TurnAborted(_) | EventMsg::StreamError(_) + )); + matches!(event, EventMsg::TurnComplete(completed) if completed.error.is_none()) + }) + .await; + + let request = final_response.single_request().body_json(); + let id = cell_id + .lock() + .expect("lock cell ID") + .clone() + .expect("yielded cell ID"); + let exec_output = + call_output_text(&request, "custom_tool_call_output", EXEC_ID).expect("exec output"); + let wait_output = + call_output_text(&request, "function_call_output", WAIT_ONE_ID).expect("first wait output"); + assert!(exec_output.contains(&format!("Script running with cell ID {id}"))); + assert!(wait_output.contains(&format!("Script running with cell ID {id}"))); + let completed_output = call_output_text(&request, "function_call_output", WAIT_TWO_ID) + .expect("second wait output"); + assert!( + completed_output.contains("background result"), + "{completed_output}" + ); + let prompts = message_input_texts(&request, "user") + .into_iter() + .filter(|text| { + matches!( + text.as_str(), + "start the cell" | "interrupt exec" | "interrupt wait" + ) + }) + .collect::>(); + assert_eq!( + prompts, + vec!["start the cell", "interrupt exec", "interrupt wait"] + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn steer_deferred_during_compaction_yields_code_mode() { + const EXEC_ID: &str = "post-compact-exec"; + let (release_compact, compact_gate) = oneshot::channel(); + let (release_continuation, continuation_gate) = oneshot::channel(); + let (release_follow_up, follow_up_gate) = oneshot::channel(); + let (server, _) = start_streaming_sse_server(vec![ + vec![ + chunk(ev_response_created("resp-first")), + chunk(ev_custom_tool_call("initial-exec", "exec", "text('ready');")), + chunk(ev_completed_with_tokens("resp-first", /*total_tokens*/ 500)), + ], + vec![ + chunk(ev_response_created("resp-compact")), + gated_chunk( + compact_gate, + vec![ + ev_message_item_done("msg-compact", "AUTO_COMPACT_SUMMARY"), + ev_completed_with_tokens("resp-compact", /*total_tokens*/ 50), + ], + ), + ], + vec![ + chunk(ev_response_created("resp-continuation")), + chunk(ev_custom_tool_call( + EXEC_ID, + "exec", + "// @exec: {\"yield_time_ms\": 60000}\nawait tools.held_tool({}); text('finished');", + )), + gated_chunk( + continuation_gate, + vec![ev_completed_with_tokens("resp-continuation", /*total_tokens*/ 60)], + ), + ], + vec![ + chunk(ev_response_created("resp-follow-up")), + gated_chunk(follow_up_gate, vec![ev_completed("resp-follow-up")]), + ], + ]) + .await; + let test = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(|config| { + config.model_provider.name = "OpenAI (test)".into(); + config.model_provider.supports_websockets = false; + config.model_auto_compact_token_limit = Some(200); + config + .features + .enable(Feature::InstantInterrupt) + .expect("enable InstantInterrupt feature"); + config + .features + .enable(Feature::CodeMode) + .expect("enable CodeMode feature"); + config + .features + .enable(Feature::CodeModeHost) + .expect("enable CodeModeHost feature"); + config.code_mode.disable_in_process_fallback = true; + config + .features + .disable(Feature::EnableRequestCompression) + .expect("disable request compression"); + }) + .build_with_streaming_server(&server) + .await + .expect("build Codex test session"); + let codex = test + .thread_manager + .start_thread(StartThreadOptions { + dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "held_tool".into(), + description: "A tool held until the test responds.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })], + ..StartThreadOptions::new(test.config.clone()) + }) + .await + .expect("start thread with dynamic tool") + .thread; + + submit_user_input(&codex, "first prompt").await; + tokio::time::timeout( + std::time::Duration::from_secs(10), + server.wait_for_request_count(/*count*/ 2), + ) + .await + .expect("compaction request should begin"); + steer_user_input(&codex, "old steer").await; + release_compact.send(()).expect("finish compaction"); + let EventMsg::DynamicToolCallRequest(held_tool) = wait_for_event(&codex, |event| { + matches!(event, EventMsg::DynamicToolCallRequest(request) if request.tool == "held_tool") + }) + .await else { + unreachable!("predicate guarantees a dynamic tool request"); + }; + release_continuation + .send(()) + .expect("finish continuation response"); + wait_for_event(&codex, |event| { + matches!(event, EventMsg::RawResponseCompleted(completed) if completed.response_id == "resp-continuation") + }) + .await; + tokio::time::timeout( + std::time::Duration::from_secs(10), + server.wait_for_request_count(/*count*/ 4), + ) + .await + .expect("the deferred steer should yield the exec"); + let requests = server.requests().await; + let continuation: Value = from_slice(&requests[2]).expect("parse continuation request"); + assert!(!message_input_texts(&continuation, "user").contains(&"old steer".to_string())); + let follow_up: Value = from_slice(&requests[3]).expect("parse follow-up request"); + assert!( + call_output_text(&follow_up, "custom_tool_call_output", EXEC_ID) + .expect("exec output") + .contains("Script running with cell ID ") + ); + let steers = message_input_texts(&follow_up, "user") + .into_iter() + .filter(|text| text == "old steer") + .collect::>(); + assert_eq!(steers, vec!["old steer"]); + + codex + .submit(Op::DynamicToolResponse { + id: held_tool.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "released".into(), + }], + success: true, + }, + }) + .await + .expect("release nested tool"); + release_follow_up.send(()).expect("finish follow-up"); + wait_for_turn_complete(&codex).await; + codex + .shutdown_and_wait() + .await + .expect("shut down test session"); + server.shutdown().await; +} + +struct FirstExecYielded { + sender: std::sync::Mutex>>, +} + +impl codex_extension_api::ToolLifecycleContributor for FirstExecYielded { + fn on_tool_finish<'a>( + &'a self, + input: codex_extension_api::ToolFinishInput<'a>, + ) -> codex_extension_api::ToolLifecycleFuture<'a> { + Box::pin(async move { + if input.tool_name.name == "exec" + && input.call_id == "first-exec" + && let Some(sender) = self.sender.lock().expect("lock sender").take() + { + let _ = sender.send(()); + } + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn input_yields_a_code_mode_call_received_later_in_the_same_response() { + let (first_yielded_tx, first_yielded_rx) = oneshot::channel(); + let (release_second_tx, release_second_rx) = oneshot::channel(); + let (release_final_tx, release_final_rx) = oneshot::channel(); + let (server, _) = start_streaming_sse_server(vec![ + vec![ + chunk(ev_response_created("resp-calls")), + chunk(ev_custom_tool_call( + "first-exec", + "exec", + "// @exec: {\"yield_time_ms\": 60000}\nawait tools.held_tool({}); text('first completed');", + )), + // This call is not even received until the first observation has + // yielded, proving it sees the already-latched request signal. + gated_chunk( + release_second_rx, + vec![ + ev_custom_tool_call( + "second-exec", + "exec", + "// @exec: {\"yield_time_ms\": 60000}\nawait new Promise(() => {});", + ), + ev_completed("resp-calls"), + ], + ), + ], + vec![ + chunk(ev_response_created("resp-follow-up")), + gated_chunk(release_final_rx, vec![ev_completed("resp-follow-up")]), + ], + ]) + .await; + let mut extensions = + codex_extension_api::ExtensionRegistryBuilder::::new(); + extensions.tool_lifecycle_contributor(Arc::new(FirstExecYielded { + sender: std::sync::Mutex::new(Some(first_yielded_tx)), + })); + let test = test_codex() + .with_model("test-gpt-5.1-codex") + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config + .features + .enable(Feature::InstantInterrupt) + .expect("enable InstantInterrupt feature"); + config + .features + .enable(Feature::CodeMode) + .expect("enable CodeMode feature"); + config + .features + .enable(Feature::CodeModeHost) + .expect("enable CodeModeHost feature"); + config.code_mode.disable_in_process_fallback = true; + config + .features + .disable(Feature::EnableRequestCompression) + .expect("disable request compression"); + config.model_provider.supports_websockets = false; + }) + .build_with_streaming_server(&server) + .await + .expect("build Codex test session"); + let codex = test + .thread_manager + .start_thread(StartThreadOptions { + dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "held_tool".into(), + description: "A tool held until the test responds.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })], + ..StartThreadOptions::new(test.config.clone()) + }) + .await + .expect("start thread with dynamic tool") + .thread; + + submit_user_input(&codex, "start cells").await; + let EventMsg::DynamicToolCallRequest(held_tool) = wait_for_event(&codex, |event| { + matches!(event, EventMsg::DynamicToolCallRequest(request) if request.tool == "held_tool") + }) + .await else { + unreachable!("predicate guarantees a dynamic tool request"); + }; + steer_user_input(&codex, "yield cells").await; + tokio::time::timeout(std::time::Duration::from_secs(10), first_yielded_rx) + .await + .expect("first exec should yield while the stream stays open") + .expect("first exec completion observer"); + release_second_tx.send(()).expect("release second exec"); + tokio::time::timeout( + std::time::Duration::from_secs(10), + server.wait_for_request_count(/*count*/ 2), + ) + .await + .expect("both exec observations should yield before the follow-up"); + let requests = server.requests().await; + let follow_up: Value = from_slice(&requests[1]).expect("parse follow-up request"); + let cell_id = |call_id| { + let output = call_output_text(&follow_up, "custom_tool_call_output", call_id) + .expect("yielded exec output"); + output + .strip_prefix("Script running with cell ID ") + .and_then(|rest| rest.lines().next()) + .expect("running cell ID") + .to_string() + }; + assert_ne!(cell_id("first-exec"), cell_id("second-exec")); + assert!(message_input_texts(&follow_up, "user").contains(&"yield cells".to_string())); + + codex + .submit(Op::DynamicToolResponse { + id: held_tool.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "released".into(), + }], + success: true, + }, + }) + .await + .expect("release nested tool"); + release_final_tx + .send(()) + .expect("finish follow-up response"); + wait_for_turn_complete(&codex).await; + codex + .shutdown_and_wait() + .await + .expect("shut down test session"); + server.shutdown().await; +} + +fn call_output_text(body: &Value, output_type: &str, call_id: &str) -> Option { + let output = &body["input"] + .as_array()? + .iter() + .find(|item| item["type"] == output_type && item["call_id"] == call_id)?["output"]; + match output { + Value::String(text) => Some(text.clone()), + Value::Array(items) => Some( + items + .iter() + .filter_map(|item| item["text"].as_str()) + .collect::>() + .join("\n"), + ), + _ => None, + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn user_input_does_not_preempt_after_reasoning_item() { let (gate_reasoning_done_tx, gate_reasoning_done_rx) = oneshot::channel(); diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index b29d534a6d..121fce211f 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -30,9 +30,14 @@ use codex_login::CodexAuth; use codex_models_manager::bundled_models_response; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; +use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; +use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::items::AgentMessageDelivery; use codex_protocol::items::TurnItem; use codex_protocol::models::ImageReference; +use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::CodeModeToolMessages; use codex_protocol::openai_models::ToolMessage; use codex_protocol::openai_models::ToolMode; @@ -993,6 +998,134 @@ async fn astra_reads_code_mode_call_timing() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn astra_continues_after_input_yields_a_code_mode_cell() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let test = test_codex() + .with_model("gpt-6-astra") + .with_config(|config| { + configure_scenario_catalog(config); + config.workspace_roots = vec![config.cwd.clone()]; + config + .features + .enable(Feature::InstantInterrupt) + .expect("enable instant interrupt"); + config + .features + .enable(Feature::CodeMode) + .expect("enable code mode"); + config + .features + .enable(Feature::CodeModeOnly) + .expect("enable code-mode-only tools"); + config + .features + .enable(Feature::CodeModeHost) + .expect("enable the code-mode host"); + config + .features + .disable(Feature::EnableRequestCompression) + .expect("disable request compression"); + config.code_mode.disable_in_process_fallback = true; + config.model_provider.supports_websockets = false; + }) + .build_with_auto_env(&server) + .await?; + let codex = test + .thread_manager + .start_thread(StartThreadOptions { + dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "check_release".into(), + description: "Check whether the release is ready.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })], + ..StartThreadOptions::new(test.config.clone()) + }) + .await? + .thread; + // A fresh code-mode session allocates its first cell as "1", as in the + // neighboring code-mode scenario. Assert that assumption below. + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("exec-response"), + ev_custom_tool_call( + "exec-call", + "exec", + "// @exec: {\"yield_time_ms\": 60000}\ntext(await tools.check_release({}));", + ), + ev_completed("exec-response"), + ]), + sse(vec![ + ev_response_created("wait-response"), + ev_function_call_with_namespace( + "wait-call", + "functions", + "wait", + r#"{"cell_id":"1","yield_time_ms":60000}"#, + ), + ev_completed("wait-response"), + ]), + sse(vec![ + ev_assistant_message( + "final", + "The release is ready; I'll highlight the customer impact.", + ), + ev_completed("final-response"), + ]), + ], + ) + .await; + + codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![text( + "Check the release status.", + )])) + .await?; + let EventMsg::DynamicToolCallRequest(tool_request) = wait_for_event(&codex, |event| { + matches!(event, EventMsg::DynamicToolCallRequest(request) if request.tool == "check_release") + }) + .await else { + unreachable!("predicate guarantees a dynamic tool request"); + }; + codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![text( + "Please also highlight the customer impact.", + )])) + .await?; + wait_for_event(&codex, |event| { + matches!(event, EventMsg::RawResponseItem(raw) if matches!(&raw.item, ResponseItem::FunctionCall { call_id, .. } if call_id == "wait-call")) + }) + .await; + codex + .submit(Op::DynamicToolResponse { + id: tool_request.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "Release ready".into(), + }], + success: true, + }, + }) + .await?; + wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + + let requests = mock.requests(); + assert!(requests[1].body_contains_text("Script running with cell ID 1")); + let snapshot = context_snapshot::format_request_history_snapshot( + "Astra continues after new user input yields a running code cell, then waits for its result.", + &requests, + &ContextSnapshotOptions::default(), + ) + .replace("cell ID 1", "cell ID ") + .replace("\"cell_id\":\"1\"", "\"cell_id\":\"\""); + insta::assert_snapshot!("astra_input_yields_code_mode_cell", snapshot); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_catalog_messages() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_input_yields_code_mode_cell.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_input_yields_code_mode_cell.snap new file mode 100644 index 0000000000..4334c834d1 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_input_yields_code_mode_cell.snap @@ -0,0 +1,120 @@ +--- +source: core/tests/suite/scenarios.rs +expression: snapshot +--- +Scenario: Astra continues after new user input yields a running code cell, then waits for its result. + +## Window 1 +-- request 1 (turn) -- +00:additional_tools/developer (3; hash=07af97f334a32933): + - namespace/functions; hash=a0b6a7ae1ecf9f28 + - custom/exec + - function/wait + - function/request_user_input + - function/request_user_input_async + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep + - namespace/collaboration: Tools for spawning and managing sub-agents.; hash=fd994b11525da73c + - function/followup_task + - function/interrupt_agent + - function/list_agents + - function/send_message + - function/spawn_agent + - function/wait_agent +01:message/developer: + You are Codex, an agent based on GPT-6. You and the user share one workspace, and your job is to collabo...them until their intended goal is completely handled. [hash=efa50f074efa5803] + + # When to ask the user for permission + + Use your best judgement given task context for when you really need user permission, like a competent co...ork without ending the turn to clarify with the user. [hash=ca7367e658b51388] + + User authorization and preferences persist across turns. Do not request permission again when the user h... any guidelines provided in skills or external files. [hash=911f91c92c72882c] + + + + - Skill naming: If a plugin contributes skills, those skill entries are prefixed with plugin_name: in the Skills list. + - MCP naming: Plugin-provided MCP tools keep standard MCP identifiers such as mcp__server__tool; use tool provenance to tell which plugin they come from. + - Trigger rules: If the user explicitly names a plugin, prefer capabilities associated with that plugin for that turn. + - Relationship to capabilities: Plugins are not invoked directly. Use their underlying skills, MCP tools, and app tools to help solve the task. + - Relevance: Determine what a plugin can help with from explicit user mention or from the plugin-associa..., MCP tools, and apps exposed elsewhere in this turn. [hash=620459a4e7b590f0] + - Missing/blocked: If the user requests a plugin that does not have relevant callable capabilities for t..., say so briefly and continue with the best fallback. [hash=746b0469a8f6cb3e] + +02:message/developer[2]: + [01] + Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `read-only`: The san... permits reading files. Network access is restricted. [hash=be7da4425f89053c] + # Escalation Requests + + Commands are run outside the sandbox if they are approved by the user, or match an existing rule that al...hell control operators, including but not limited to: [hash=d24ebf5fa4e4b9bb] + + - Pipes: | + - Logical operators: &&, || + + NEVER provide a prefix_rule if your command uses a heredoc or herestring. + + ### Examples + Good examples of prefixes: + - ["npm", "run", "dev"] + - ["gh", "pr", "check"] + - ["cargo", "test"] + + [02] # Collaboration Mode: Default + + You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active. + + Your active mode changes only when new developer instructions with a different `...<...by themselves. Known mode names are Default and Plan. [hash=e75f6b12fc52d812] + + ## request_user_input availability + + Use the `request_user_input` tool only when it is listed in the available tools for this turn. + + Use the `request_user_input` tool only for optional questions where the answer would materially improve the quality of the work. + + If `request_user_input` returns no answers, continue with best judgment instead of asking again or treating the turn as blocked. + + Never use the `request_user_input` tool for permission requests or permission-related escalations. + +03:message/developer: + You are `/root`, the primary agent in a team of agents collaborating to fulfill the user's goals. + + At the start of your turn, you are the active agent. + You can spawn sub-agents to handle subtasks, and those sub-agents can spawn their own sub-agents. + All agents in the team, including the agents that you can assign tasks to, are equally intelligent and capable, and have access to the same set of tools. + + You can use `spawn_agent` to create a new agent, `followup_task` to give an existing agent a new task an...message to a running agent without triggering a turn. [hash=7d40ff8b58977202] + `send_message` calls may be read by a human, so ensure they are legible. Always put proper spaces between words and/or numbers. + + - All agents use the same current working directory. + - As a result, edits made by one agent are immediately visible to all other agents. + + When calling `wait_agent`, prefer longer waits (minutes) to avoid busy polling. + + There are 4 available concurrency slots, meaning that up to 4 agents can be active at once, including you. + + Full-history forks (`fork_turns` omitted or `"all"`) inherit the parent model and reasoning effort and d...ne"` or a positive integer string. [hash=c7f52e90688299ce] +04:message/developer: + Any earlier instruction enabling proactive multi-agent delegation no longer applies. D...elegation, or parallel agent work. [hash=2b6e037067b4390c] +05:message/user: + + + + + + [hash=e8c1de60307d0a3d] + +06:message/user: + Check the release status. +-- request 2 (turn) -- +07:custom_tool_call/exec:// @exec: {"yield_time_ms": 60000} + text(await tools.check_release({})); +08:custom_tool_call_output:Script running with cell ID + Wall time seconds + Output: + +09:message/user: + Please also highlight the customer impact. +-- request 3 (turn) -- +10:function_call/functions.wait:{"cell_id":"","yield_time_ms":60000} +11:function_call_output:Script completed + Wall time seconds + Output: + | Release ready diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 171e6dc2ea..071b9260e5 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -208,6 +208,8 @@ pub enum Feature { /// Keep sampling through reasoning and commentary boundaries when agent mail arrives. /// Pending mail is delivered at the next normal input boundary instead. DeferMailboxPreemption, + /// Yield foreground code-mode observations when new user input arrives. + InstantInterrupt, /// Enable shared discussion tools for an agent tree. AgentMessageBoard, /// Removed compatibility flag retained as a no-op. @@ -1076,6 +1078,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::InstantInterrupt, + key: "instant_interrupt", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::CodeModeOnly, key: "code_mode_only",