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 1d3bc5643f..3d43e0c653 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -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 diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index 8500ee309b..2006b8681d 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -63,19 +63,19 @@ impl ToolExecutor 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) + Ok(Arc::new(WaitTestSession) as Arc) }) } } -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, ) -> 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(()) diff --git a/codex-rs/rollout-trace/src/code_cell.rs b/codex-rs/rollout-trace/src/code_cell.rs index 63241e8a61..de9ae5daaa 100644 --- a/codex-rs/rollout-trace/src/code_cell.rs +++ b/codex-rs/rollout-trace/src/code_cell.rs @@ -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 { diff --git a/codex-rs/rollout-trace/src/model/runtime.rs b/codex-rs/rollout-trace/src/model/runtime.rs index 3907804686..46ecc7e814 100644 --- a/codex-rs/rollout-trace/src/model/runtime.rs +++ b/codex-rs/rollout-trace/src/model/runtime.rs @@ -50,6 +50,8 @@ pub struct CodeCell { pub yielded_seq: Option, pub source_js: String, pub nested_tool_call_ids: Vec, + /// 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, } @@ -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, + /// 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, pub thread_id: AgentThreadId, /// Runtime activation that started the tool. Background work may outlive this turn. pub started_by_codex_turn_id: Option, diff --git a/codex-rs/rollout-trace/src/raw_event.rs b/codex-rs/rollout-trace/src/raw_event.rs index 3fa06eb26f..7593bc6c1d 100644 --- a/codex-rs/rollout-trace/src/raw_event.rs +++ b/codex-rs/rollout-trace/src/raw_event.rs @@ -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, .. diff --git a/codex-rs/rollout-trace/src/reducer/code_cell.rs b/codex-rs/rollout-trace/src/reducer/code_cell.rs index f61279a5d2..3cdf5140eb 100644 --- a/codex-rs/rollout-trace/src/reducer/code_cell.rs +++ b/codex-rs/rollout-trace/src/reducer/code_cell.rs @@ -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, + pub(super) wait_tool_call_ids: Vec, 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(()) } diff --git a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs index 317c647658..c80f0511a8 100644 --- a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs @@ -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(()) } diff --git a/codex-rs/rollout-trace/src/reducer/mod.rs b/codex-rs/rollout-trace/src/reducer/mod.rs index 4710b3452a..7479578b1c 100644 --- a/codex-rs/rollout-trace/src/reducer/mod.rs +++ b/codex-rs/rollout-trace/src/reducer/mod.rs @@ -75,6 +75,7 @@ pub fn replay_bundle(bundle_dir: impl AsRef) -> Result { 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, /// 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, diff --git a/codex-rs/rollout-trace/src/reducer/thread.rs b/codex-rs/rollout-trace/src/reducer/thread.rs index c9b8f2b647..93a8121822 100644 --- a/codex-rs/rollout-trace/src/reducer/thread.rs +++ b/codex-rs/rollout-trace/src/reducer/thread.rs @@ -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 diff --git a/codex-rs/rollout-trace/src/reducer/tool.rs b/codex-rs/rollout-trace/src/reducer/tool.rs index bf7061e19e..ace015653c 100644 --- a/codex-rs/rollout-trace/src/reducer/tool.rs +++ b/codex-rs/rollout-trace/src/reducer/tool.rs @@ -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 {