github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add opt-in code-mode yielding on new user input (#48135)

Browse Source 
## 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
This commit is contained in:
pakrym-oai
2026-09-25 16:28:36 +00:00
committed by copyberry
parent d9275a63ed
commit 1bf73324ca
13 changed files with 1011 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/tests/stdio.rs
+1
View File
@@ -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();
codex-rs/core/config.schema.json
+6
View File
@@ -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"
},
codex-rs/core/src/session/input_queue.rs
+31
View File
@@ -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<Option<ActiveTurn>>,
sub_id: &str,
interrupt: CancellationToken,
) -> Option<AbortOnDropHandle<()>> {
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()
}
codex-rs/core/src/session/mod.rs
+5
View File
@@ -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,
codex-rs/core/src/session/step_context.rs
+3
View File
@@ -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<TurnContext>,
/// Yields this request's code-mode observations when a user message arrives.
pub(crate) preempt: Option<CancellationToken>,
/// Realtime call activity and instructions captured for this sampling request.
pub(crate) realtime: RealtimeConversationSnapshot,
/// One immutable settings version captured before request preparation.
codex-rs/core/src/session/tests.rs
+5
View File
@@ -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,
codex-rs/core/src/session/turn.rs
+7
View File
@@ -1588,6 +1588,13 @@ async fn run_sampling_request(
cancellation_token: CancellationToken,
) -> CodexResult<(SamplingRequestResult, Vec<ResponseItem>)> {
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(
codex-rs/core/src/tools/code_mode/mod.rs
+4 -2
View File
@@ -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<CancellationToken>,
) -> Result<codex_code_mode::WaitOutcome, String> {
self.session().await?.wait(request, /*preempt*/ None).await
self.session().await?.wait(request, preempt).await
}
pub(crate) async fn terminate(
codex-rs/core/src/tools/code_mode/wait_handler.rs
+7 -4
View File
@@ -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| {
codex-rs/core/tests/suite/pending_input.rs
+681
View File
@@ -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::<Vec<_>>();
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::<String>));
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::<Vec<_>>();
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::<Vec<_>>();
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::<Vec<_>>();
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<Option<oneshot::Sender<()>>>,
}
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::<codex_core::config::Config>::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<String> {
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::<Vec<_>>()
.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();
codex-rs/core/tests/suite/scenarios.rs
+133
View File
@@ -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 <CELL_ID>")
.replace("\"cell_id\":\"1\"", "\"cell_id\":\"<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(()));
codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_input_yields_code_mode_cell.snap
+120
View File
@@ -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]
<OMITTED 155 LINES; ~4935 TOKENS; hash=5573869716192298>
- 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] <permissions instructions>
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: &&, ||
<OMITTED 44 LINES; ~790 TOKENS; hash=3a8900343e968d61>
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"]
</permissions instructions>
[02] <collaboration_mode># 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 `<collaboration_mode>...<...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.
</collaboration_mode>
03:message/developer:
<multi_agent_role>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.
<OMITTED 17 LINES; ~246 TOKENS; hash=780031e2df907944>
- 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.</multi_agent_role> [hash=c7f52e90688299ce]
04:message/developer:
<multi_agent_mode>Any earlier instruction enabling proactive multi-agent delegation no longer applies. D...elegation, or parallel agent work.</multi_agent_mode> [hash=2b6e037067b4390c]
05:message/user:
<environment_context>
<cwd><CWD></cwd>
<shell><HOST_SHELL></shell>
<current_date><CURRENT_DATE></current_date>
<timezone><HOST_TIMEZONE></timezone>
<filesystem><workspace_roots><root><CWD></root></workspace_roots><permission_profile type="managed"><f...ntry></file_system></permission_profile></filesystem> [hash=e8c1de60307d0a3d]
</environment_context>
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 <CELL_ID>
Wall time <DURATION> seconds
Output:
09:message/user:
Please also highlight the customer impact.
-- request 3 (turn) --
10:function_call/functions.wait:{"cell_id":"<CELL_ID>","yield_time_ms":60000}
11:function_call_output:Script completed
Wall time <DURATION> seconds
Output:
| Release ready
codex-rs/features/src/lib.rs
+8
View File
@@ -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",