openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-14 21:40:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Record host-confirmed Code Mode wait targets in rollout traces (#61338)

Browse Source 
## Why

Inferring wait targets from request arguments can abort trace replay on malformed arguments and link waits to cells the host never resolved.

## What changed

Record an explicit `CodeCellWaitTargetResolved` event when the host confirms a live cell, and expose its runtime handle as `resolved_wait_runtime_cell_id` on the tool call. Build cell wait links from these observations, preserving links while cell source is pending and preventing late confirmations from linking to recycled handles after thread resume.

Older traces no longer infer wait links from request arguments, so an empty `wait_tool_call_ids` list does not establish that nobody waited on a cell.

## Testing

Extend dispatch and replay tests to cover live, missing, and invalid wait targets, pending cell source, legacy traces, and runtime handle reuse across threads and resume.

GitOrigin-RevId: 0f13b291c197c5487258b6ce36276cf1f19e2b6d
This commit is contained in:
tongzhou wang
2026-11-13 18:43:21 +00:00
committed by copyberry
parent c79dfb1d9c
commit 9d12d007b2
10 changed files with 238 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/code_mode/wait_handler.rs
+13 -8
View File
@@ -154,11 +154,23 @@ impl CodeModeWaitHandler {
FunctionCallError::RespondToModel(error)
})?;
if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response {
// Dispatch already records failed waits. Re-parsing {"cell_id": 5} during
// replay used to abort the whole trace, and even a valid string can name
// a closed cell. Only the host can confirm which cell this wait reached.
let runtime_cell_id = match response {
RuntimeResponse::Yielded { cell_id, .. }
| RuntimeResponse::Terminated { cell_id, .. }
| RuntimeResponse::Result { cell_id, .. } => cell_id,
};
let cell_trace = exec
.session
.services
.rollout_thread_trace
.code_cell_trace_context(
exec.turn.sub_id.as_str(),
runtime_cell_id.as_str(),
);
cell_trace.record_wait_target(&call_id);
tracing::Span::current().record("cell.id", trace_id(runtime_cell_id.as_str()));
telemetry.cell_id = Some(runtime_cell_id.to_string());
exec.session
@@ -174,14 +186,7 @@ impl CodeModeWaitHandler {
.unwrap_or_default(),
},
);
exec.session
.services
.rollout_thread_trace
.code_cell_trace_context(
exec.turn.sub_id.as_str(),
runtime_cell_id.as_str(),
)
.record_ended(response);
cell_trace.record_ended(response);
exec.session
.services
.code_mode_service
codex-rs/core/src/tools/tool_dispatch_trace_tests.rs
+53 -30
View File
@@ -63,19 +63,19 @@ impl ToolExecutor<ToolInvocation> for TestHandler {
impl CoreToolRuntime for TestHandler {}
struct MissingCellCodeModeSessionProvider;
struct WaitTestSessionProvider;
impl codex_code_mode::CodeModeSessionProvider for MissingCellCodeModeSessionProvider {
impl codex_code_mode::CodeModeSessionProvider for WaitTestSessionProvider {
fn create_session(&self) -> codex_code_mode::CodeModeSessionProviderFuture<'_> {
Box::pin(async {
Ok(Arc::new(MissingCellCodeModeSession) as Arc<dyn codex_code_mode::CodeModeSession>)
Ok(Arc::new(WaitTestSession) as Arc<dyn codex_code_mode::CodeModeSession>)
})
}
}
struct MissingCellCodeModeSession;
struct WaitTestSession;
impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession {
impl codex_code_mode::CodeModeSession for WaitTestSession {
fn execute<'a>(
&'a self,
_request: codex_code_mode::ExecuteRequest,
@@ -90,6 +90,17 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession {
request: codex_code_mode::WaitRequest,
_preempt: Option<CancellationToken>,
) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::WaitOutcome> {
if request.cell_id.as_str() == "live" {
return Box::pin(async move {
Ok(codex_code_mode::WaitOutcome::LiveCell(
codex_code_mode::RuntimeResponse::Yielded {
cell_id: request.cell_id,
content_items: Vec::new(),
code_mode_host_duration: None,
},
))
});
}
self.terminate(request.cell_id.into())
}
@@ -287,12 +298,12 @@ async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> any
}
#[tokio::test]
async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Result<()> {
async fn wait_handler_records_only_host_confirmed_targets() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let (mut session, turn) = make_session_and_context().await;
session.services.code_mode_service = CodeModeService::new(
session.thread_id,
Arc::new(MissingCellCodeModeSessionProvider),
Arc::new(WaitTestSessionProvider),
&turn.config.code_mode,
session.services.executed_tool_calls.clone(),
);
@@ -305,33 +316,45 @@ async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Resu
let session = Arc::new(session);
let turn = Arc::new(turn);
let mut invocation = test_invocation(
session,
turn,
"wait-call",
WAIT_TOOL_NAME,
ToolCallSource::Direct,
r#"{"cell_id":"noop","terminate":true}"#,
);
invocation.tool_name = invocation.tool_name.with_default_namespace();
assert!(
super::tool_dispatch_invocation(&invocation)
.expect("wait calls should produce a trace invocation")
.tool_namespace
.is_none()
);
registry
.dispatch_any_with_state(invocation, /*call_state*/ None)
.await?;
let dispatch = |call_id, arguments| {
let mut invocation = test_invocation(
Arc::clone(&session),
Arc::clone(&turn),
call_id,
WAIT_TOOL_NAME,
ToolCallSource::Direct,
arguments,
);
invocation.tool_name = invocation.tool_name.with_default_namespace();
assert!(
super::tool_dispatch_invocation(&invocation)
.expect("wait calls should produce a trace invocation")
.tool_namespace
.is_none()
);
registry.dispatch_any_with_state(invocation, /*call_state*/ None)
};
dispatch("live-wait", r#"{"cell_id":"live"}"#).await?;
dispatch("missing-wait", r#"{"cell_id":"noop","terminate":true}"#).await?;
assert!(matches!(
dispatch("invalid-wait", r#"{"cell_id":5}"#).await,
Err(FunctionCallError::RespondToModel(_))
));
trace.flush()?;
let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?;
assert_eq!(replayed.code_cells.len(), 0);
assert!(replayed.code_cells.is_empty());
assert_eq!(
["live-wait", "missing-wait", "invalid-wait"].map(|id| replayed.tool_calls[id]
.resolved_wait_runtime_cell_id
.as_deref()),
[Some("live"), None, None]
);
assert!(
replayed.tool_calls["wait-call"]
.raw_result_payload_id
.is_some()
replayed
.tool_calls
.values()
.all(|call| call.raw_result_payload_id.is_some())
);
Ok(())
codex-rs/rollout-trace/src/code_cell.rs
+17
View File
@@ -121,6 +121,23 @@ impl CodeCellTraceContext {
);
}
/// Records a host-confirmed wait target before the handler returns. A missing cell
/// or rejected argument has no such observation. This does not imply caller delivery.
pub fn record_wait_target(&self, tool_call_id: &str) {
let CodeCellTraceContextState::Enabled(context) = &self.state else {
return;
};
if let Err(err) = context.writer.enqueue_event(
RawTraceEventContext::turn(&context.thread_id, &context.codex_turn_id),
RawTraceEventPayload::CodeCellWaitTargetResolved {
tool_call_id: tool_call_id.to_owned(),
runtime_cell_id: context.runtime_cell_id.clone(),
},
) {
warn!("failed to append rollout trace event: {err:#}");
}
}
/// Records the terminal lifecycle point for a code-mode runtime cell.
pub fn record_ended(&self, response: &RuntimeResponse) {
let CodeCellTraceContextState::Enabled(context) = &self.state else {
codex-rs/rollout-trace/src/model/runtime.rs
+5
View File
@@ -50,6 +50,8 @@ pub struct CodeCell {
pub yielded_seq: Option<RawEventSeq>,
pub source_js: String,
pub nested_tool_call_ids: Vec<ToolCallId>,
/// Wait dispatches confirmed by the host. Older traces did not record that confirmation,
/// so an empty list does not prove nobody waited on the cell.
pub wait_tool_call_ids: Vec<ToolCallId>,
}
@@ -126,6 +128,9 @@ pub struct ToolCall {
/// Cell started by this outer exec dispatch, if a matching cell was recorded.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub started_code_cell_id: Option<CodeCellId>,
/// Runtime handle confirmed by the Code Mode host, even if the cell source is absent.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub resolved_wait_runtime_cell_id: Option<String>,
pub thread_id: AgentThreadId,
/// Runtime activation that started the tool. Background work may outlive this turn.
pub started_by_codex_turn_id: Option<CodexTurnId>,
codex-rs/rollout-trace/src/raw_event.rs
+6
View File
@@ -179,6 +179,11 @@ pub enum RawTraceEventPayload {
tool_call_id: ToolCallId,
exec_tool_call_id: ToolCallId,
},
/// Live Code Mode cell reached by a wait or terminate call.
CodeCellWaitTargetResolved {
tool_call_id: ToolCallId,
runtime_cell_id: String,
},
ToolCallRuntimeStarted {
tool_call_id: ToolCallId,
/// Executed handler, which may differ from the requested callable.
@@ -283,6 +288,7 @@ impl RawTraceEventPayload {
| RawTraceEventPayload::CodeCellStarted { .. }
| RawTraceEventPayload::McpToolCallCorrelationAssigned { .. }
| RawTraceEventPayload::StdinTargetResolved { .. }
| RawTraceEventPayload::CodeCellWaitTargetResolved { .. }
| RawTraceEventPayload::AgentResultObserved {
carried_payload: None,
..
codex-rs/rollout-trace/src/reducer/code_cell.rs
+42 -45
View File
@@ -16,7 +16,7 @@
use anyhow::Context;
use anyhow::Result;
use anyhow::bail;
use serde_json::Value;
use anyhow::ensure;
use super::TraceReducer;
use crate::model::CodeCell;
@@ -54,6 +54,7 @@ pub(super) struct PendingCodeCellStart {
pub(super) wall_time_unix_ms: i64,
pub(super) thread_id: String,
pub(super) codex_turn_id: Option<String>,
pub(super) wait_tool_call_ids: Vec<ToolCallId>,
pub(super) started: StartedCodeCell,
}
@@ -136,6 +137,7 @@ impl TraceReducer {
wall_time_unix_ms,
thread_id,
codex_turn_id,
wait_tool_call_ids,
started,
} = pending;
if self.rollout.code_cells.contains_key(&started.code_cell_id) {
@@ -201,7 +203,7 @@ impl TraceReducer {
yielded_seq: None,
source_js: started.source_js,
nested_tool_call_ids,
wait_tool_call_ids: Vec::new(),
wait_tool_call_ids,
},
);
@@ -479,57 +481,52 @@ impl TraceReducer {
Ok(())
}
/// Records that a model-visible wait call is waiting on a runtime code cell.
///
/// Wait calls are not nested JavaScript tools, so the relationship is inferred
/// from the runtime cell id inside the function arguments.
pub(super) fn link_wait_tool_call_from_request_payload(
/// Links only a host-confirmed wait, using the cell binding from this activation.
pub(super) fn link_confirmed_wait(
&mut self,
thread_id: &str,
tool_call_id: &ToolCallId,
request_payload: Option<&RawPayloadRef>,
turn_id: &str,
tool_call_id: ToolCallId,
runtime_cell_id: String,
) -> Result<()> {
let Some(request_payload) = request_payload else {
return Ok(());
};
let payload = self.payload_reader.read_json(request_payload)?;
if payload.get("tool_name").and_then(Value::as_str) != Some("wait") {
let wait = self
.rollout
.tool_calls
.get_mut(&tool_call_id)
.with_context(|| {
format!("confirmed wait referenced unknown dispatch {tool_call_id}")
})?;
ensure!(
wait.thread_id == thread_id
&& wait.started_by_codex_turn_id.as_deref() == Some(turn_id),
"confirmed wait target has wrong dispatch context"
);
ensure!(
wait.resolved_wait_runtime_cell_id.is_none(),
"duplicate confirmed wait target for {tool_call_id}"
);
wait.resolved_wait_runtime_cell_id = Some(runtime_cell_id.clone());
// Retain the host observation but do not bind an older dispatch to a
// potentially recycled handle if it was recorded after thread reload.
if self
.last_thread_resume_seq
.get(thread_id)
.is_some_and(|seq| *seq > wait.execution.started_seq)
{
return Ok(());
}
// `wait` is a normal model-visible function call, not a nested JS tool
// request. The only stable edge back to the code cell is the runtime
// `cell_id` inside the function arguments.
let Some(arguments) = payload
.get("payload")
.and_then(|payload| payload.get("arguments"))
.and_then(Value::as_str)
else {
bail!(
"wait tool request payload {} did not contain function arguments",
request_payload.raw_payload_id
);
};
let arguments: Value = serde_json::from_str(arguments).with_context(|| {
format!(
"wait tool request payload {} had invalid JSON arguments",
request_payload.raw_payload_id
)
})?;
let Some(runtime_cell_id) = arguments.get("cell_id").and_then(Value::as_str) else {
bail!(
"wait tool request payload {} did not contain cell_id",
request_payload.raw_payload_id
);
};
let Some(code_cell_id) =
self.code_cell_id_for_runtime_cell_id_if_known(thread_id, runtime_cell_id)
let Some(cell_id) =
self.code_cell_id_for_runtime_cell_id_if_known(thread_id, &runtime_cell_id)
else {
return Ok(());
};
let Some(cell) = self.rollout.code_cells.get_mut(&code_cell_id) else {
return Ok(());
};
push_unique(&mut cell.wait_tool_call_ids, tool_call_id);
if let Some(cell) = self.rollout.code_cells.get_mut(&cell_id) {
cell.wait_tool_call_ids.push(tool_call_id);
} else if let Some(start) = self.pending_code_cell_starts.get_mut(&cell_id) {
// The runtime is known; its model-visible exec source may arrive later.
start.wait_tool_call_ids.push(tool_call_id);
}
Ok(())
}
codex-rs/rollout-trace/src/reducer/code_cell_tests.rs
+76 -3
View File
@@ -1,3 +1,5 @@
use std::sync::Arc;
use pretty_assertions::assert_eq;
use serde_json::json;
use tempfile::TempDir;
@@ -25,7 +27,8 @@ use crate::reducer::test_support::trace_context_for_thread;
use crate::replay_bundle;
#[test]
fn code_cell_lifecycle_links_nested_tools_waits_and_outputs() -> anyhow::Result<()> {
fn code_cell_lifecycle_links_nested_tools_and_outputs_without_inferring_legacy_wait()
-> anyhow::Result<()> {
let temp = TempDir::new()?;
let writer = create_started_writer(&temp)?;
start_turn(&writer, "turn-1")?;
@@ -184,7 +187,8 @@ fn code_cell_lifecycle_links_nested_tools_waits_and_outputs() -> anyhow::Result<
assert_eq!(cell.execution.status, ExecutionStatus::Completed);
assert_eq!(cell.runtime_cell_id, Some("1".to_string()));
assert_eq!(cell.nested_tool_call_ids, vec!["nested-tool-1"]);
assert_eq!(cell.wait_tool_call_ids, vec!["wait-tool-1"]);
// The legacy request alone does not establish that the host resolved a target.
assert!(cell.wait_tool_call_ids.is_empty());
assert_eq!(cell.output_item_ids, vec![output_item_id.clone()]);
assert_eq!(
rollout.conversation_items[output_item_id].produced_by,
@@ -405,7 +409,7 @@ fn cancelled_turn_terminates_unfinished_code_cell() -> anyhow::Result<()> {
#[test]
fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let writer = create_started_writer(&temp)?;
let writer = Arc::new(create_started_writer(&temp)?);
writer.append(RawTraceEventPayload::ThreadStarted {
thread_id: "thread-child".to_string(),
agent_path: "/root/child".to_string(),
@@ -451,6 +455,39 @@ fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow:
source_js: "text('hi')".to_string(),
},
)?;
let waits: &[(&str, Option<&str>, &str)] = match turn_id {
"turn-root" => &[("rejected", None, "{")],
"turn-child" => &[
("confirmed", None, r#"{"cell_id":"1"}"#),
("late", None, r#"{"cell_id":"1"}"#),
],
_ => &[("foreign", Some("fixture"), r#"{"cell_id":"1"}"#)],
};
for &(wait_id, namespace, arguments) in waits {
let dispatch = crate::ToolDispatchTraceContext::start(
Arc::clone(&writer),
crate::ToolDispatchInvocation {
thread_id: thread_id.into(),
codex_turn_id: turn_id.into(),
tool_call_id: wait_id.into(),
tool_name: "wait".into(),
tool_namespace: namespace.map(str::to_owned),
requester: crate::ToolDispatchRequester::Model {
model_visible_call_id: wait_id.into(),
},
payload: crate::ToolDispatchPayload::Function {
arguments: arguments.into(),
},
},
);
if wait_id == "rejected" {
dispatch.record_failed("invalid wait arguments");
} else if wait_id == "confirmed" {
// The source item is still pending. No registry return is required for this edge.
crate::CodeCellTraceContext::enabled(Arc::clone(&writer), thread_id, turn_id, "1")
.record_wait_target(wait_id);
}
}
let response = writer.write_json_payload(
RawPayloadKind::InferenceResponse,
&json!({
@@ -469,6 +506,11 @@ fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow:
upstream_request_id: None,
response_payload: response,
})?;
if turn_id == "turn-resumed" {
// The pre-resume wait is confirmed after the new cell has reused its handle.
crate::CodeCellTraceContext::enabled(Arc::clone(&writer), thread_id, "turn-child", "1")
.record_wait_target("late");
}
writer.append_with_context(
trace_context_for_thread(thread_id, turn_id),
RawTraceEventPayload::CodeCellEnded {
@@ -502,6 +544,37 @@ fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow:
rollout.code_cells[&resumed_cell_id].runtime_cell_id,
Some("1".to_string())
);
assert_eq!(
rollout.code_cells[&child_cell_id].wait_tool_call_ids,
["confirmed"]
);
assert!(
rollout.code_cells[&root_cell_id]
.wait_tool_call_ids
.is_empty()
);
assert!(
rollout.code_cells[&resumed_cell_id]
.wait_tool_call_ids
.is_empty()
);
for id in ["confirmed", "late"] {
assert_eq!(
rollout.tool_calls[id]
.resolved_wait_runtime_cell_id
.as_deref(),
Some("1")
);
}
assert!(
rollout.tool_calls["confirmed"]
.execution
.ended_seq
.is_none()
);
for id in ["rejected", "foreign"] {
assert_eq!(rollout.tool_calls[id].resolved_wait_runtime_cell_id, None);
}
Ok(())
}
codex-rs/rollout-trace/src/reducer/mod.rs
+22 -1
View File
@@ -75,6 +75,7 @@ pub fn replay_bundle(bundle_dir: impl AsRef<Path>) -> Result<RolloutTrace> {
thread_conversation_snapshots: BTreeMap::new(),
pending_compaction_replacement_item_ids: BTreeMap::new(),
code_cell_ids_by_runtime: BTreeMap::new(),
last_thread_resume_seq: BTreeMap::new(),
pending_code_cell_starts: BTreeMap::new(),
pending_code_cell_lifecycle_events: BTreeMap::new(),
pending_agent_interaction_edges: Vec::new(),
@@ -229,6 +230,8 @@ struct TraceReducer {
/// tools, and `wait` calls arrive with the runtime-local `cell_id`, so this
/// index is the one intentional bridge between those namespaces.
code_cell_ids_by_runtime: BTreeMap<(String, String), String>,
/// A late observation from an older dispatch cannot identify a cell in the new runtime.
last_thread_resume_seq: BTreeMap<String, crate::RawEventSeq>,
/// Code-cell starts whose model-visible `custom_tool_call` item has not
/// been reduced yet.
///
@@ -316,7 +319,7 @@ impl TraceReducer {
self.end_thread(event.seq, event.wall_time_unix_ms, thread_id, status)?;
}
RawTraceEventPayload::ThreadResumed { thread_id } => {
self.resume_thread(&thread_id)?;
self.resume_thread(&thread_id, event.seq)?;
}
RawTraceEventPayload::CodexTurnStarted {
codex_turn_id,
@@ -436,6 +439,23 @@ impl TraceReducer {
} => {
self.resolve_stdin_target(&tool_call_id, &exec_tool_call_id)?;
}
RawTraceEventPayload::CodeCellWaitTargetResolved {
tool_call_id,
runtime_cell_id,
} => {
self.link_confirmed_wait(
event
.thread_id
.as_deref()
.context("confirmed wait missing thread")?,
event
.codex_turn_id
.as_deref()
.context("confirmed wait missing turn")?,
tool_call_id,
runtime_cell_id,
)?;
}
RawTraceEventPayload::ToolCallRuntimeStarted {
tool_call_id,
kind,
@@ -495,6 +515,7 @@ impl TraceReducer {
wall_time_unix_ms: event.wall_time_unix_ms,
thread_id,
codex_turn_id: event.codex_turn_id,
wait_tool_call_ids: Vec::new(),
started: StartedCodeCell {
code_cell_id: reduced_code_cell_id,
runtime_cell_id,
codex-rs/rollout-trace/src/reducer/thread.rs
+3 -1
View File
@@ -105,7 +105,7 @@ impl TraceReducer {
}
/// Reopens a known thread after eviction, preserving its identity and history.
pub(super) fn resume_thread(&mut self, thread_id: &str) -> Result<()> {
pub(super) fn resume_thread(&mut self, thread_id: &str, seq: RawEventSeq) -> Result<()> {
let thread = self.thread_mut(thread_id).with_context(|| {
format!(
"cannot resume thread {thread_id}: it was not started in this trace; resuming threads from an earlier trace is not supported"
@@ -120,6 +120,8 @@ impl TraceReducer {
thread.execution.status = ExecutionStatus::Running;
// Reload creates fresh cell and process namespaces. Historical objects
// remain, but their runtime handles must not resolve in the new session.
self.last_thread_resume_seq
.insert(thread_id.to_owned(), seq);
self.code_cell_ids_by_runtime
.retain(|(owner, _), _| owner != thread_id);
self.terminal_ids_by_process
codex-rs/rollout-trace/src/reducer/tool.rs
+1 -5
View File
@@ -122,11 +122,6 @@ impl TraceReducer {
.invocation_payload
.as_ref()
.map(|payload| payload.raw_payload_id.clone());
self.link_wait_tool_call_from_request_payload(
&thread_id,
&tool_call_id,
started.invocation_payload.as_ref(),
)?;
self.rollout.tool_calls.insert(
tool_call_id.clone(),
@@ -136,6 +131,7 @@ impl TraceReducer {
model_visible_call_id,
code_mode_runtime_tool_id: started.code_mode_runtime_tool_id,
started_code_cell_id: None,
resolved_wait_runtime_cell_id: None,
thread_id,
started_by_codex_turn_id: codex_turn_id,
execution: ExecutionWindow {