diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 880e84ef4a..e46405dfe2 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -24,6 +24,7 @@ pub use runtime::ExecuteRequest; pub use runtime::RuntimeResponse; pub use runtime::WaitRequest; pub use service::CodeModeService; +pub use service::CodeModeToolInvocation; pub use service::CodeModeTurnHost; pub use service::CodeModeTurnWorker; diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 0f50edd329..1897a1f004 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -25,6 +25,7 @@ const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; #[derive(Clone, Debug)] pub struct ExecuteRequest { + pub cell_id: Option, pub tool_call_id: String, pub enabled_tools: Vec, pub source: String, @@ -331,6 +332,7 @@ mod tests { fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { + cell_id: None, tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 23ca7a7460..4c1ef8d760 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -23,12 +23,26 @@ use crate::runtime::TurnMessage; use crate::runtime::WaitRequest; use crate::runtime::spawn_runtime; +/// A nested tool request made by one running code-mode cell. +/// +/// The model only sees the outer `exec` call. Trace reduction needs the runtime +/// cell id and runtime tool id so it can attach this nested tool to the right +/// code cell without guessing from event ordering. +#[derive(Debug)] +pub struct CodeModeToolInvocation { + /// Code-mode runtime cell that issued the request. + pub cell_id: String, + /// Per-cell tool invocation id generated by the JavaScript runtime. + pub runtime_tool_call_id: String, + pub tool_name: ToolName, + pub input: Option, +} + #[async_trait] pub trait CodeModeTurnHost: Send + Sync { async fn invoke_tool( &self, - tool_name: ToolName, - input: Option, + invocation: CodeModeToolInvocation, cancellation_token: CancellationToken, ) -> Result; @@ -76,24 +90,40 @@ impl CodeModeService { *self.inner.stored_values.lock().await = values; } - pub async fn execute(&self, request: ExecuteRequest) -> Result { - let cell_id = self - .inner + pub fn allocate_cell_id(&self) -> String { + self.inner .next_cell_id .fetch_add(1, Ordering::Relaxed) - .to_string(); + .to_string() + } + + pub async fn execute(&self, request: ExecuteRequest) -> Result { + let cell_id = request + .cell_id + .clone() + .unwrap_or_else(|| self.allocate_cell_id()); + if self.inner.sessions.lock().await.contains_key(&cell_id) { + return Err(format!("exec cell {cell_id} already exists")); + } + let (event_tx, event_rx) = mpsc::unbounded_channel(); let (runtime_tx, runtime_terminate_handle) = spawn_runtime(request.clone(), event_tx)?; let (control_tx, control_rx) = mpsc::unbounded_channel(); let (response_tx, response_rx) = oneshot::channel(); - self.inner.sessions.lock().await.insert( + let mut sessions = self.inner.sessions.lock().await; + if sessions.contains_key(&cell_id) { + let _ = runtime_terminate_handle.terminate_execution(); + return Err(format!("exec cell {cell_id} already exists")); + } + sessions.insert( cell_id.clone(), SessionHandle { control_tx: control_tx.clone(), runtime_tx: runtime_tx.clone(), }, ); + drop(sessions); tokio::spawn(run_session_control( Arc::clone(&self.inner), @@ -181,9 +211,14 @@ impl CodeModeService { let host = Arc::clone(&host); let inner = Arc::clone(&inner); tokio::spawn(async move { - let response = host - .invoke_tool(name, input, CancellationToken::new()) - .await; + let invocation = CodeModeToolInvocation { + cell_id: cell_id.clone(), + runtime_tool_call_id: id.clone(), + tool_name: name, + input, + }; + let response = + host.invoke_tool(invocation, CancellationToken::new()).await; let runtime_tx = inner .sessions .lock() @@ -490,6 +525,7 @@ mod tests { fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { + cell_id: None, tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 4805b03d01..1ef1937f3f 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -6027,6 +6027,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() { session: Arc::clone(&session), turn: Arc::clone(&turn_context), tracker: Arc::clone(&turn_diff_tracker), + source: crate::tools::context::ToolCallSource::Direct, call_id, tool_name: codex_tools::ToolName::plain(tool_name), payload: ToolPayload::Function { @@ -6105,6 +6106,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() session: Arc::clone(&session), turn: Arc::clone(&turn_context), tracker: Arc::clone(&tracker), + source: crate::tools::context::ToolCallSource::Direct, call_id: "exec-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), payload: ToolPayload::Function { diff --git a/codex-rs/core/src/codex_tests_guardian.rs b/codex-rs/core/src/codex_tests_guardian.rs index cad67dcc8f..d727a22164 100644 --- a/codex-rs/core/src/codex_tests_guardian.rs +++ b/codex-rs/core/src/codex_tests_guardian.rs @@ -143,6 +143,7 @@ async fn guardian_allows_shell_additional_permissions_requests_past_policy_valid session: Arc::clone(&session), turn: Arc::clone(&turn_context), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "test-call".to_string(), tool_name: codex_tools::ToolName::plain("shell"), payload: ToolPayload::Function { @@ -209,6 +210,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic session: Arc::clone(&session), turn: Arc::clone(&turn_context), tracker: Arc::clone(&tracker), + source: crate::tools::context::ToolCallSource::Direct, call_id: "exec-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), payload: ToolPayload::Function { @@ -322,6 +324,7 @@ async fn shell_handler_allows_sticky_turn_permissions_without_inline_request_per session: Arc::clone(&session), turn: Arc::clone(&turn_context), tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "sticky-turn-grant".to_string(), tool_name: codex_tools::ToolName::plain("shell"), payload: ToolPayload::Function { diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 8eaa85b564..aafe41a128 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -4,6 +4,7 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; +use serde_json::json; use super::ExecContext; use super::PUBLIC_TOOL_NAME; @@ -30,12 +31,34 @@ impl CodeModeExecuteHandler { .code_mode_service .stored_values() .await; + // Allocate the runtime cell id before starting V8. Nested tools can be + // requested before the model-visible `exec` item has been reduced, so + // the trace needs this explicit bridge up front. + let runtime_cell_id = exec.session.services.code_mode_service.allocate_cell_id(); + let invocation_payload = codex_trace::write_payload( + "code_cell_invocation", + &json!({ + "cell_id": runtime_cell_id, + "model_visible_call_id": call_id, + "source_js": args.code, + "yield_time_ms": args.yield_time_ms, + "max_output_tokens": args.max_output_tokens, + }), + ); + emit_code_cell_started( + &exec, + &runtime_cell_id, + &call_id, + &args.code, + invocation_payload.as_ref(), + ); let started_at = std::time::Instant::now(); - let response = exec + let response = match exec .session .services .code_mode_service .execute(codex_code_mode::ExecuteRequest { + cell_id: Some(runtime_cell_id.clone()), tool_call_id: call_id, enabled_tools, source: args.code, @@ -44,13 +67,146 @@ impl CodeModeExecuteHandler { max_output_tokens: args.max_output_tokens, }) .await - .map_err(FunctionCallError::RespondToModel)?; + { + Ok(response) => response, + Err(err) => { + let result_payload = codex_trace::write_payload( + "code_cell_result", + &json!({ "cell_id": runtime_cell_id, "error_text": err }), + ); + emit_code_cell_failed(&exec, &runtime_cell_id, result_payload.as_ref()); + return Err(FunctionCallError::RespondToModel(err)); + } + }; + if !matches!(&response, codex_code_mode::RuntimeResponse::Yielded { .. }) { + let result_payload = codex_trace::write_payload( + "code_cell_result", + &code_cell_response_payload(&response), + ); + emit_code_cell_ended( + &exec, + &runtime_cell_id, + &response, + /*model_visible_wait_call_id*/ None, + result_payload.as_ref(), + ); + } handle_runtime_response(&exec, response, args.max_output_tokens, started_at) .await .map_err(FunctionCallError::RespondToModel) } } +fn emit_code_cell_started( + exec: &ExecContext, + runtime_cell_id: &str, + model_visible_call_id: &str, + source_js: &str, + invocation_payload: Option<&codex_trace::RawPayloadRef>, +) { + tracing::event!( + target: codex_otel::OTEL_TRACE_SAFE_TARGET, + tracing::Level::INFO, + event.name = %"codex.code_cell.started", + thread.id = %exec.session.conversation_id, + turn.id = %exec.turn.sub_id, + code_cell.runtime_id = %runtime_cell_id, + model_visible_call.id = %model_visible_call_id, + code_cell.source_js = %source_js, + raw_payload.invocation.id = %invocation_payload.map(|payload| payload.raw_payload_id.as_str()).unwrap_or(""), + raw_payload.invocation.path = %invocation_payload.map(|payload| payload.path.as_str()).unwrap_or(""), + raw_payload.invocation.kind = %invocation_payload.map(|payload| payload.kind.as_str()).unwrap_or(""), + ); +} + +pub(super) fn emit_code_cell_ended( + exec: &ExecContext, + runtime_cell_id: &str, + response: &codex_code_mode::RuntimeResponse, + model_visible_wait_call_id: Option<&str>, + result_payload: Option<&codex_trace::RawPayloadRef>, +) { + tracing::event!( + target: codex_otel::OTEL_TRACE_SAFE_TARGET, + tracing::Level::INFO, + event.name = %"codex.code_cell.ended", + thread.id = %exec.session.conversation_id, + turn.id = %exec.turn.sub_id, + code_cell.runtime_id = %runtime_cell_id, + status = %code_cell_status(response), + model_visible_wait_call.id = %model_visible_wait_call_id.unwrap_or(""), + raw_payload.result.id = %result_payload.map(|payload| payload.raw_payload_id.as_str()).unwrap_or(""), + raw_payload.result.path = %result_payload.map(|payload| payload.path.as_str()).unwrap_or(""), + raw_payload.result.kind = %result_payload.map(|payload| payload.kind.as_str()).unwrap_or(""), + ); +} + +fn emit_code_cell_failed( + exec: &ExecContext, + runtime_cell_id: &str, + result_payload: Option<&codex_trace::RawPayloadRef>, +) { + tracing::event!( + target: codex_otel::OTEL_TRACE_SAFE_TARGET, + tracing::Level::INFO, + event.name = %"codex.code_cell.ended", + thread.id = %exec.session.conversation_id, + turn.id = %exec.turn.sub_id, + code_cell.runtime_id = %runtime_cell_id, + status = %"failed", + raw_payload.result.id = %result_payload.map(|payload| payload.raw_payload_id.as_str()).unwrap_or(""), + raw_payload.result.path = %result_payload.map(|payload| payload.path.as_str()).unwrap_or(""), + raw_payload.result.kind = %result_payload.map(|payload| payload.kind.as_str()).unwrap_or(""), + ); +} + +fn code_cell_status(response: &codex_code_mode::RuntimeResponse) -> &'static str { + match response { + codex_code_mode::RuntimeResponse::Yielded { .. } => "yielded", + codex_code_mode::RuntimeResponse::Terminated { .. } => "terminated", + codex_code_mode::RuntimeResponse::Result { error_text, .. } => { + if error_text.is_some() { + "failed" + } else { + "completed" + } + } + } +} + +fn code_cell_response_payload(response: &codex_code_mode::RuntimeResponse) -> serde_json::Value { + match response { + codex_code_mode::RuntimeResponse::Yielded { + cell_id, + content_items, + } => json!({ + "cell_id": cell_id, + "status": "yielded", + "content_items": content_items, + }), + codex_code_mode::RuntimeResponse::Terminated { + cell_id, + content_items, + } => json!({ + "cell_id": cell_id, + "status": "terminated", + "content_items": content_items, + }), + codex_code_mode::RuntimeResponse::Result { + cell_id, + content_items, + stored_values, + error_text, + } => json!({ + "cell_id": cell_id, + "status": if error_text.is_some() { "failed" } else { "completed" }, + "content_items": content_items, + "stored_values": stored_values, + "error_text": error_text, + }), + } +} + impl ToolHandler for CodeModeExecuteHandler { type Output = FunctionToolOutput; @@ -62,6 +218,23 @@ impl ToolHandler for CodeModeExecuteHandler { matches!(payload, ToolPayload::Custom { .. }) } + fn uses_first_class_trace_object(&self, invocation: &ToolInvocation) -> bool { + // Code-mode `exec` is represented by the first-class CodeCell runtime + // lifecycle. Emitting the generic dispatch ToolCall as well would give + // viewers two top-level objects for the same model-visible JS cell. + // + // Only suppress the generic ToolCall once parsing is known to succeed. + // Parse errors return before a CodeCell id exists, so those failures + // still need the ordinary tool trace object. + invocation.tool_name.namespace.is_none() + && invocation.tool_name.name.as_str() == PUBLIC_TOOL_NAME + && matches!( + &invocation.payload, + ToolPayload::Custom { input } + if codex_code_mode::parse_exec_source(input).is_ok() + ) + } + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 7d66d56c39..dcda0fa492 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -7,6 +7,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +use codex_code_mode::CodeModeToolInvocation; use codex_code_mode::CodeModeTurnHost; use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; @@ -73,6 +74,10 @@ impl CodeModeService { self.inner.replace_stored_values(values).await; } + pub(crate) fn allocate_cell_id(&self) -> String { + self.inner.allocate_cell_id() + } + pub(crate) async fn execute( &self, request: codex_code_mode::ExecuteRequest, @@ -118,15 +123,13 @@ struct CoreTurnHost { impl CodeModeTurnHost for CoreTurnHost { async fn invoke_tool( &self, - tool_name: ToolName, - input: Option, + invocation: CodeModeToolInvocation, cancellation_token: CancellationToken, ) -> Result { call_nested_tool( self.exec.clone(), self.tool_runtime.clone(), - tool_name, - input, + invocation, cancellation_token, ) .await @@ -298,10 +301,15 @@ async fn build_nested_router(exec: &ExecContext) -> ToolRouter { async fn call_nested_tool( exec: ExecContext, tool_runtime: ToolCallRuntime, - tool_name: ToolName, - input: Option, + invocation: CodeModeToolInvocation, cancellation_token: CancellationToken, ) -> Result { + let CodeModeToolInvocation { + cell_id, + runtime_tool_call_id, + tool_name, + input, + } = invocation; if tool_name.namespace.is_none() && tool_name.name == PUBLIC_TOOL_NAME { return Err(FunctionCallError::RespondToModel(format!( "{PUBLIC_TOOL_NAME} cannot invoke itself" @@ -335,7 +343,14 @@ async fn call_nested_tool( payload, }; let result = tool_runtime - .handle_tool_call_with_source(call, ToolCallSource::CodeMode, cancellation_token) + .handle_tool_call_with_source( + call, + ToolCallSource::CodeMode { + cell_id, + runtime_tool_call_id, + }, + cancellation_token, + ) .await?; Ok(result.code_mode_result()) } 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 dcc416b832..fb6fda1e49 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -10,6 +10,7 @@ use crate::tools::registry::ToolKind; use super::DEFAULT_WAIT_YIELD_TIME_MS; use super::ExecContext; use super::WAIT_TOOL_NAME; +use super::execute_handler::emit_code_cell_ended; use super::handle_runtime_response; pub struct CodeModeWaitHandler; @@ -49,6 +50,7 @@ impl ToolHandler for CodeModeWaitHandler { let ToolInvocation { session, turn, + call_id, tool_name, payload, .. @@ -66,12 +68,28 @@ impl ToolHandler for CodeModeWaitHandler { .services .code_mode_service .wait(codex_code_mode::WaitRequest { - cell_id: args.cell_id, + cell_id: args.cell_id.clone(), yield_time_ms: args.yield_time_ms, terminate: args.terminate, }) .await .map_err(FunctionCallError::RespondToModel)?; + if !matches!(&response, codex_code_mode::RuntimeResponse::Yielded { .. }) + && !is_missing_cell_response(&response, &args.cell_id) + { + // The initial `exec` call only records a terminal end when + // it finishes immediately. Yielded cells close later via + // `wait`, so the runtime object stays live until this point. + // Duplicate waits for a removed cell return a model-visible + // error but must not rewrite an already-finished CodeCell. + emit_code_cell_ended( + &exec, + &args.cell_id, + &response, + Some(&call_id), + /*result_payload*/ None, + ); + } handle_runtime_response(&exec, response, args.max_tokens, started_at) .await .map_err(FunctionCallError::RespondToModel) @@ -82,3 +100,14 @@ impl ToolHandler for CodeModeWaitHandler { } } } + +fn is_missing_cell_response(response: &codex_code_mode::RuntimeResponse, cell_id: &str) -> bool { + let codex_code_mode::RuntimeResponse::Result { + error_text: Some(error_text), + .. + } = response + else { + return false; + }; + error_text == &format!("exec cell {cell_id} not found") +} diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 60bcbcda23..15175f39a4 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -28,11 +28,16 @@ use tokio::sync::Mutex; pub type SharedTurnDiffTracker = Arc>; -#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[derive(Clone, Debug, Eq, PartialEq)] pub enum ToolCallSource { Direct, JsRepl, - CodeMode, + CodeMode { + /// Runtime cell that issued the nested tool request. + cell_id: String, + /// Code-mode's per-cell tool invocation id. + runtime_tool_call_id: String, + }, } #[derive(Clone)] @@ -40,6 +45,7 @@ pub struct ToolInvocation { pub session: Arc, pub turn: Arc, pub tracker: SharedTurnDiffTracker, + pub source: ToolCallSource, pub call_id: String, pub tool_name: ToolName, pub payload: ToolPayload, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 10f41d97eb..a5f628cc5c 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -68,6 +68,7 @@ fn invocation( session, turn, tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(tool_name), payload, diff --git a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs index 1dac095f1b..cf8b59b0e5 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs @@ -28,6 +28,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { session: Arc::new(session), turn: Arc::new(turn), tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME), payload: ToolPayload::Function { diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index e66133f961..a1ed33051b 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -226,6 +226,7 @@ async fn shell_pre_tool_use_payload_uses_joined_command() { session: session.into(), turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "call-41".to_string(), tool_name: codex_tools::ToolName::plain("shell"), payload, @@ -251,6 +252,7 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() { session: session.into(), turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "call-42".to_string(), tool_name: codex_tools::ToolName::plain("shell_command"), payload, diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index b641814516..51a82996af 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -211,6 +211,7 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() { session: session.into(), turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "call-43".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), payload, @@ -234,6 +235,7 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() { session: session.into(), turn: turn.into(), tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), + source: crate::tools::context::ToolCallSource::Direct, call_id: "call-44".to_string(), tool_name: codex_tools::ToolName::plain("write_stdin"), payload, diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 6403102733..8e731e965a 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -75,6 +75,10 @@ pub trait ToolHandler: Send + Sync { None } + fn uses_first_class_trace_object(&self, _invocation: &ToolInvocation) -> bool { + false + } + /// Perform the actual [ToolInvocation] and returns a [ToolOutput] containing /// the final output to return to the model. fn handle( @@ -133,6 +137,8 @@ trait AnyToolHandler: Send + Sync { result: &dyn ToolOutput, ) -> Option; + fn uses_first_class_trace_object(&self, invocation: &ToolInvocation) -> bool; + fn handle_any<'a>( &'a self, invocation: ToolInvocation, @@ -164,6 +170,10 @@ where ToolHandler::post_tool_use_payload(self, call_id, payload, result) } + fn uses_first_class_trace_object(&self, invocation: &ToolInvocation) -> bool { + ToolHandler::uses_first_class_trace_object(self, invocation) + } + fn handle_any<'a>( &'a self, invocation: ToolInvocation, @@ -224,16 +234,23 @@ impl ToolRegistry { ) -> Result { let call_id_owned = invocation.call_id.clone(); let trace_tool_call_id = format!("tool:{call_id_owned}"); - let invocation_payload = - codex_trace::write_payload("tool_invocation", &trace_tool_invocation(&invocation)); - emit_tool_started( - &invocation, - &trace_tool_call_id, - invocation_payload.as_ref(), - ); + let uses_first_class_trace_object = self + .handler(&invocation.tool_name) + .is_some_and(|handler| handler.uses_first_class_trace_object(&invocation)); + if !uses_first_class_trace_object { + let invocation_payload = + codex_trace::write_payload("tool_invocation", &trace_tool_invocation(&invocation)); + emit_tool_started( + &invocation, + &trace_tool_call_id, + invocation_payload.as_ref(), + ); + } let result = self.dispatch_any_inner(invocation).await; - emit_tool_ended(&trace_tool_call_id, &result); + if !uses_first_class_trace_object { + emit_tool_ended(&trace_tool_call_id, &result); + } result } @@ -466,6 +483,8 @@ fn emit_tool_started( tool_call_id: &str, invocation_payload: Option<&codex_trace::RawPayloadRef>, ) { + let (requester_type, model_visible_call_id, runtime_cell_id, runtime_tool_call_id) = + trace_tool_source_fields(invocation); tracing::event!( target: codex_otel::OTEL_TRACE_SAFE_TARGET, tracing::Level::INFO, @@ -474,13 +493,34 @@ fn emit_tool_started( turn.id = %invocation.turn.sub_id, tool.call_id = %tool_call_id, tool.name = %invocation.tool_name.display(), - model_visible_call.id = %invocation.call_id, + tool.requester.type = %requester_type, + model_visible_call.id = %model_visible_call_id, + code_cell.runtime_id = %runtime_cell_id, + code_mode_runtime_tool.id = %runtime_tool_call_id, raw_payload.invocation.id = %invocation_payload.map(|payload| payload.raw_payload_id.as_str()).unwrap_or(""), raw_payload.invocation.path = %invocation_payload.map(|payload| payload.path.as_str()).unwrap_or(""), raw_payload.invocation.kind = %invocation_payload.map(|payload| payload.kind.as_str()).unwrap_or(""), ); } +fn trace_tool_source_fields(invocation: &ToolInvocation) -> (&str, &str, &str, &str) { + match &invocation.source { + crate::tools::context::ToolCallSource::Direct => { + ("model", invocation.call_id.as_str(), "", "") + } + crate::tools::context::ToolCallSource::CodeMode { + cell_id, + runtime_tool_call_id, + } => ( + "code_cell", + "", + cell_id.as_str(), + runtime_tool_call_id.as_str(), + ), + crate::tools::context::ToolCallSource::JsRepl => ("runtime", "", "", ""), + } +} + fn emit_tool_ended(tool_call_id: &str, result: &Result) { let (status, output_preview) = match result { Ok(result) => { @@ -541,10 +581,26 @@ fn trace_tool_invocation(invocation: &ToolInvocation) -> Value { "call_id": &invocation.call_id, "tool_name": invocation.tool_name.display(), "turn_id": &invocation.turn.sub_id, + "source": trace_tool_invocation_source(invocation), "payload": payload, }) } +fn trace_tool_invocation_source(invocation: &ToolInvocation) -> Value { + match &invocation.source { + crate::tools::context::ToolCallSource::Direct => json!({ "type": "model" }), + crate::tools::context::ToolCallSource::CodeMode { + cell_id, + runtime_tool_call_id, + } => json!({ + "type": "code_cell", + "runtime_cell_id": cell_id, + "runtime_tool_call_id": runtime_tool_call_id, + }), + crate::tools::context::ToolCallSource::JsRepl => json!({ "type": "runtime" }), + } +} + fn trace_tool_result( result: &Result, status: &str, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index ef7ae32d4b..abe010a60b 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -285,6 +285,7 @@ impl ToolRouter { session, turn, tracker, + source, call_id, tool_name, payload, diff --git a/codex-rs/trace/src/reduce.rs b/codex-rs/trace/src/reduce.rs index 961a32c098..90eea54505 100644 --- a/codex-rs/trace/src/reduce.rs +++ b/codex-rs/trace/src/reduce.rs @@ -51,7 +51,9 @@ struct Reducer<'a> { // around only while reducing. It lets full request snapshots point back to // the canonical item without writing the raw item into state.json. raw_conversation_items: BTreeMap, + code_cell_ids_by_runtime: BTreeMap<(String, String), String>, next_conversation_item_ordinal: u64, + next_terminal_operation_ordinal: u64, } pub fn reduce_bundle_to_path(bundle_dir: impl AsRef, output: impl AsRef) -> Result<()> { @@ -97,7 +99,9 @@ fn reduce_bundle(bundle_dir: &Path) -> Result { raw_payloads: BTreeMap::new(), }, raw_conversation_items: BTreeMap::new(), + code_cell_ids_by_runtime: BTreeMap::new(), next_conversation_item_ordinal: 1, + next_terminal_operation_ordinal: 1, }; let event_log_path = bundle_dir.join("events.jsonl"); @@ -112,6 +116,10 @@ fn reduce_bundle(bundle_dir: &Path) -> Result { reducer.apply_event(event)?; } reducer.link_tools_to_conversation_items(); + reducer.link_code_cells_to_conversation_items(); + reducer.drop_redundant_code_cell_tool_calls(); + reducer.attach_wait_tool_payloads_to_code_cells(); + reducer.sync_terminal_model_observations(); reducer.attach_tool_payloads_to_interaction_edges(); if reducer.trace.started_at_unix_ms == 0 { reducer.trace.started_at_unix_ms = reducer @@ -156,8 +164,10 @@ impl Reducer<'_> { self.inference_completed(&event)?; } "codex.inference.failed" | "inference_failed" => self.inference_failed(&event), - "codex.tool.started" | "tool_started" => self.tool_started(&event), + "codex.tool.started" | "tool_started" => self.tool_started(&event)?, "codex.tool.ended" | "tool_ended" => self.tool_ended(&event), + "codex.code_cell.started" => self.code_cell_started(&event), + "codex.code_cell.ended" => self.code_cell_ended(&event), "codex.collab.spawn.started" => self.tool_to_thread_edge_started(&event, "spawn_agent"), "codex.collab.spawn.ended" => self.tool_to_thread_edge_ended(&event, "spawn_agent"), "codex.collab.message.started" => { @@ -358,18 +368,30 @@ impl Reducer<'_> { } } - fn tool_started(&mut self, event: &CapturedEvent) { + fn tool_started(&mut self, event: &CapturedEvent) -> Result<()> { let (Some(tool_call_id), Some(thread_id), Some(turn_id)) = ( trace_field_str(event, "tool", "call_id"), trace_field_str(event, "thread", "id"), trace_field_str(event, "turn", "id"), ) else { - return; + return Ok(()); }; self.ensure_thread(thread_id, event.wall_time_unix_ms); let tool_name = trace_field_str(event, "tool", "name").unwrap_or("tool"); let raw_invocation_payload_id = raw_payload_field_str(event, "invocation", "id").unwrap_or(""); + let model_visible_call_id = trace_field_str(event, "model_visible_call", "id") + .filter(|call_id| !call_id.is_empty()); + let code_mode_runtime_tool_id = + field_str(event, "code_mode_runtime_tool.id").filter(|tool_id| !tool_id.is_empty()); + let requester = self.tool_requester(event, thread_id); + let terminal_operation_id = self.start_terminal_operation_for_tool( + event, + thread_id, + tool_call_id, + tool_name, + raw_invocation_payload_id, + )?; self.trace.tool_calls.insert( tool_call_id.to_string(), json!({ @@ -377,11 +399,13 @@ impl Reducer<'_> { "thread_id": thread_id, "codex_turn_id": turn_id, "started_by_codex_turn_id": turn_id, - "model_visible_call_id": trace_field_str(event, "model_visible_call", "id").unwrap_or(""), + "model_visible_call_id": model_visible_call_id, + "code_mode_runtime_tool_id": code_mode_runtime_tool_id, "model_visible_call_item_ids": [], "model_visible_output_item_ids": [], - "requester": { "type": "model" }, - "kind": { "type": "other", "name": tool_name }, + "requester": requester, + "kind": tool_call_kind(tool_name), + "terminal_operation_id": terminal_operation_id, "summary": { "input_preview": "", "output_preview": "", @@ -392,12 +416,17 @@ impl Reducer<'_> { "raw_runtime_payload_ids": [], }), ); + self.link_tool_call_to_code_cell(tool_call_id); + Ok(()) } fn tool_ended(&mut self, event: &CapturedEvent) { let Some(tool_call_id) = trace_field_str(event, "tool", "call_id") else { return; }; + let raw_result_payload_id = raw_payload_field_str(event, "result", "id").unwrap_or(""); + let mut terminal_operation_id = None; + let mut thread_id = None; if let Some(tool) = self.trace.tool_calls.get_mut(tool_call_id) { let status = field_str(event, "status").unwrap_or("completed"); set_execution_end(tool, event.wall_time_unix_ms, event.seq, status); @@ -407,8 +436,96 @@ impl Reducer<'_> { Value::String(field_str(event, "output_preview").unwrap_or("").to_string()), ); } - tool["raw_result_payload_id"] = - empty_to_null(raw_payload_field_str(event, "result", "id").unwrap_or("")); + tool["raw_result_payload_id"] = empty_to_null(raw_result_payload_id); + terminal_operation_id = tool + .get("terminal_operation_id") + .and_then(Value::as_str) + .map(str::to_string); + thread_id = tool + .get("thread_id") + .and_then(Value::as_str) + .map(str::to_string); + } + if let Some(operation_id) = terminal_operation_id { + self.end_terminal_operation( + &operation_id, + thread_id.as_deref().unwrap_or(""), + event, + raw_result_payload_id, + ); + } + } + + fn code_cell_started(&mut self, event: &CapturedEvent) { + let (Some(thread_id), Some(turn_id), Some(runtime_cell_id), Some(model_visible_call_id)) = ( + trace_field_str(event, "thread", "id"), + trace_field_str(event, "turn", "id"), + field_str(event, "code_cell.runtime_id"), + trace_field_str(event, "model_visible_call", "id"), + ) else { + return; + }; + self.ensure_thread(thread_id, event.wall_time_unix_ms); + let code_cell_id = reduced_code_cell_id(model_visible_call_id); + self.code_cell_ids_by_runtime.insert( + (thread_id.to_string(), runtime_cell_id.to_string()), + code_cell_id.clone(), + ); + self.trace.code_cells.entry(code_cell_id.clone()).or_insert_with(|| { + let raw_invocation_payload_id = + raw_payload_field_str(event, "invocation", "id").unwrap_or(""); + json!({ + "code_cell_id": code_cell_id, + "model_visible_call_id": model_visible_call_id, + "thread_id": thread_id, + "codex_turn_id": turn_id, + "source_item_id": Value::Null, + "output_item_ids": [], + "runtime_cell_id": runtime_cell_id, + "execution": execution_json(event.wall_time_unix_ms, None, "running", event.seq, None), + "runtime_status": "running", + "source_js": field_str(event, "code_cell.source_js").unwrap_or(""), + "nested_tool_call_ids": [], + "wait_tool_call_ids": [], + "raw_payload_ids": non_empty_array(raw_invocation_payload_id), + }) + }); + } + + fn code_cell_ended(&mut self, event: &CapturedEvent) { + let (Some(thread_id), Some(runtime_cell_id)) = ( + trace_field_str(event, "thread", "id"), + field_str(event, "code_cell.runtime_id"), + ) else { + return; + }; + let Some(code_cell_id) = self + .code_cell_ids_by_runtime + .get(&(thread_id.to_string(), runtime_cell_id.to_string())) + .cloned() + else { + return; + }; + if let Some(cell) = self.trace.code_cells.get_mut(&code_cell_id) { + let runtime_status = field_str(event, "status").unwrap_or("completed"); + set_execution_end( + cell, + event.wall_time_unix_ms, + event.seq, + code_cell_execution_status(runtime_status), + ); + cell["runtime_status"] = Value::String(runtime_status.to_string()); + if let Some(wait_call_id) = trace_field_str(event, "model_visible_wait_call", "id") + .filter(|call_id| !call_id.is_empty()) + { + let tool_call_id = normalize_tool_call_id(wait_call_id); + push_unique_json_string(cell, "wait_tool_call_ids", &tool_call_id); + } + if let Some(raw_result_payload_id) = raw_payload_field_str(event, "result", "id") + .filter(|payload_id| !payload_id.is_empty()) + { + push_unique_json_string(cell, "raw_payload_ids", raw_result_payload_id); + } } } @@ -675,6 +792,172 @@ impl Reducer<'_> { item_id } + fn tool_requester(&self, event: &CapturedEvent, thread_id: &str) -> Value { + if field_str(event, "tool.requester.type") != Some("code_cell") { + return json!({ "type": "model" }); + } + + // Code-mode nested tools are not visible to the model as ordinary + // request/response items. The runtime cell id is the stable bridge + // from the tool event back to the model-visible `exec` code cell. + let Some(runtime_cell_id) = field_str(event, "code_cell.runtime_id") else { + return json!({ "type": "code_cell", "code_cell_id": Value::Null }); + }; + let code_cell_id = self + .code_cell_ids_by_runtime + .get(&(thread_id.to_string(), runtime_cell_id.to_string())) + .cloned(); + json!({ + "type": "code_cell", + "code_cell_id": code_cell_id, + "runtime_cell_id": runtime_cell_id, + }) + } + + fn link_tool_call_to_code_cell(&mut self, tool_call_id: &str) { + let Some(code_cell_id) = self + .trace + .tool_calls + .get(tool_call_id) + .and_then(|tool| tool.get("requester")) + .and_then(|requester| requester.get("code_cell_id")) + .and_then(Value::as_str) + .map(str::to_string) + else { + return; + }; + if let Some(cell) = self.trace.code_cells.get_mut(&code_cell_id) { + push_unique_json_string(cell, "nested_tool_call_ids", tool_call_id); + } + } + + fn start_terminal_operation_for_tool( + &mut self, + event: &CapturedEvent, + thread_id: &str, + tool_call_id: &str, + tool_name: &str, + raw_invocation_payload_id: &str, + ) -> Result { + if !is_terminal_tool(tool_name) { + return Ok(Value::Null); + } + + // Terminal operations are a runtime view over shell-like tools. They + // deliberately hang off the reduced tool call so direct model tools and + // nested code-cell tools share the same terminal model. + let operation_id = format!( + "terminal_operation:{}", + self.next_terminal_operation_ordinal + ); + self.next_terminal_operation_ordinal += 1; + let request = self + .payload_by_id(raw_invocation_payload_id)? + .map(|payload| terminal_request_from_invocation(&payload)) + .unwrap_or_else(|| json!({ "type": "exec_command" })); + let terminal_id = terminal_id_from_terminal_request(&request); + self.trace.terminal_operations.insert( + operation_id.clone(), + json!({ + "operation_id": operation_id, + "terminal_id": terminal_id, + "tool_call_id": tool_call_id, + "kind": terminal_operation_kind(tool_name), + "execution": execution_json(event.wall_time_unix_ms, None, "running", event.seq, None), + "request": request, + "result": Value::Null, + "model_observations": [], + "raw_payload_ids": non_empty_array(raw_invocation_payload_id), + }), + ); + // Continuation tools already name their terminal session in the request. + // Create the session immediately so interrupted traces still group the + // in-flight operation under the right terminal rail. + if let Some(terminal_id) = terminal_id { + self.ensure_terminal_session_for_operation(thread_id, &terminal_id, &operation_id); + } + Ok(json!(operation_id)) + } + + fn end_terminal_operation( + &mut self, + operation_id: &str, + thread_id: &str, + event: &CapturedEvent, + raw_result_payload_id: &str, + ) { + let result_payload = self + .payload_by_id(raw_result_payload_id) + .ok() + .flatten() + .unwrap_or_else(|| json!({})); + let request_terminal_id = self + .trace + .terminal_operations + .get(operation_id) + .and_then(terminal_id_from_operation_request); + // Prefer real runtime ids over synthetic ids. Some continuation + // results omit `session_id`, but their request still names the session. + let terminal_id = terminal_id_from_result_payload(&result_payload, "") + .or(request_terminal_id) + .or_else(|| terminal_id_from_result_payload(&result_payload, operation_id)); + if let Some(operation) = self.trace.terminal_operations.get_mut(operation_id) { + let status = field_str(event, "status").unwrap_or("completed"); + set_execution_end(operation, event.wall_time_unix_ms, event.seq, status); + if let Some(terminal_id) = terminal_id.as_deref() { + operation["terminal_id"] = Value::String(terminal_id.to_string()); + } + operation["result"] = + terminal_result_from_payload(&result_payload, field_str(event, "output_preview")); + push_unique_json_string(operation, "raw_payload_ids", raw_result_payload_id); + } + if let Some(terminal_id) = terminal_id { + self.ensure_terminal_session_for_operation(thread_id, &terminal_id, operation_id); + } + } + + fn ensure_terminal_session_for_operation( + &mut self, + thread_id: &str, + terminal_id: &str, + operation_id: &str, + ) { + if terminal_id.is_empty() || operation_id.is_empty() { + return; + } + let operation_execution = self + .trace + .terminal_operations + .get(operation_id) + .and_then(|operation| operation.get("execution")) + .cloned() + .unwrap_or_else(|| execution_json(0, None, "running", 0, None)); + let session_execution = if terminal_id.starts_with("terminal:terminal_operation:") { + operation_execution + } else { + let mut execution = operation_execution; + execution["ended_at_unix_ms"] = Value::Null; + execution["ended_seq"] = Value::Null; + execution["status"] = Value::String("running".to_string()); + execution + }; + self.trace + .terminal_sessions + .entry(terminal_id.to_string()) + .or_insert_with(|| { + json!({ + "terminal_id": terminal_id, + "thread_id": thread_id, + "created_by_operation_id": operation_id, + "operation_ids": [], + "execution": session_execution, + }) + }); + if let Some(session) = self.trace.terminal_sessions.get_mut(terminal_id) { + push_unique_json_string(session, "operation_ids", operation_id); + } + } + fn payload_by_id(&self, raw_payload_id: &str) -> Result> { if raw_payload_id.is_empty() { return Ok(None); @@ -721,26 +1004,236 @@ impl Reducer<'_> { tool["model_visible_call_item_ids"] = json!(call_items); tool["model_visible_output_item_ids"] = json!(output_items); } + + let links = self + .trace + .tool_calls + .iter() + .flat_map(|(tool_call_id, tool)| { + string_array_field(tool, "model_visible_output_item_ids") + .into_iter() + .map(|item_id| (tool_call_id.clone(), item_id)) + .collect::>() + }) + .collect::>(); + for (tool_call_id, item_id) in links { + self.add_conversation_item_producer( + &item_id, + json!({ "type": "tool", "tool_call_id": tool_call_id }), + ); + } + } + + fn link_code_cells_to_conversation_items(&mut self) { + let code_cell_ids = self.trace.code_cells.keys().cloned().collect::>(); + for code_cell_id in code_cell_ids { + let Some(model_visible_call_id) = self + .trace + .code_cells + .get(&code_cell_id) + .and_then(|cell| cell.get("model_visible_call_id")) + .and_then(Value::as_str) + .map(str::to_string) + else { + continue; + }; + + let mut source_item_id = None; + let mut output_item_ids = Vec::new(); + let mut output_call_ids = vec![model_visible_call_id.clone()]; + if let Some(cell) = self.trace.code_cells.get(&code_cell_id) { + for wait_tool_call_id in string_array_field(cell, "wait_tool_call_ids") { + if let Some(call_id) = wait_tool_call_id.strip_prefix("tool:") { + output_call_ids.push(call_id.to_string()); + } + } + } + // The code-cell start event gives us runtime identity, but the + // provider payload gives us the canonical conversation item ids. + // Link them after replay so we can see both sides of the turn. A + // yielded cell may emit its final model-visible output through a + // later `wait` call, so outputs are collected from both ids. + for item in self.trace.conversation_items.values() { + let Some(item_call_id) = item.get("call_id").and_then(Value::as_str) else { + continue; + }; + if !output_call_ids + .iter() + .any(|call_id| call_id.as_str() == item_call_id) + { + continue; + } + let Some(item_id) = item.get("item_id").and_then(Value::as_str) else { + continue; + }; + match item.get("kind").and_then(Value::as_str) { + Some("custom_tool_call") if item_call_id == model_visible_call_id => { + source_item_id = Some(item_id.to_string()); + } + Some("custom_tool_call_output" | "function_call_output") => { + output_item_ids.push(item_id.to_string()); + } + _ => {} + } + } + + if let Some(cell) = self.trace.code_cells.get_mut(&code_cell_id) { + if let Some(source_item_id) = source_item_id { + cell["source_item_id"] = Value::String(source_item_id); + } + cell["output_item_ids"] = json!(output_item_ids); + } + for output_item_id in output_item_ids { + self.add_conversation_item_producer( + &output_item_id, + json!({ "type": "code_cell", "code_cell_id": code_cell_id }), + ); + } + } + } + + fn drop_redundant_code_cell_tool_calls(&mut self) { + let mut redundant_tool_call_ids = Vec::new(); + for (tool_call_id, tool) in &self.trace.tool_calls { + if tool + .get("kind") + .and_then(|kind| kind.get("type")) + .and_then(Value::as_str) + != Some("code_cell") + { + continue; + } + let Some(model_visible_call_id) = + tool.get("model_visible_call_id").and_then(Value::as_str) + else { + continue; + }; + let code_cell_id = reduced_code_cell_id(model_visible_call_id); + if self.trace.code_cells.contains_key(&code_cell_id) { + redundant_tool_call_ids.push(tool_call_id.clone()); + } + } + + for tool_call_id in redundant_tool_call_ids { + // `exec` itself is represented by `code_cells`; keeping the generic + // dispatch ToolCall makes viewers render two roots for the same + // model-visible JavaScript cell. Copy its payload refs first so the + // code-cell details view still has a direct path to raw bytes. + if let Some(tool) = self.trace.tool_calls.remove(&tool_call_id) + && let Some(model_visible_call_id) = + tool.get("model_visible_call_id").and_then(Value::as_str) + { + let code_cell_id = reduced_code_cell_id(model_visible_call_id); + if let Some(cell) = self.trace.code_cells.get_mut(&code_cell_id) { + for raw_payload_id in tool_raw_payload_ids(&tool) { + push_unique_json_string(cell, "raw_payload_ids", &raw_payload_id); + } + } + } + self.remove_conversation_item_producer( + json!({ "type": "tool", "tool_call_id": tool_call_id }), + ); + } + } + + fn attach_wait_tool_payloads_to_code_cells(&mut self) { + for cell in self.trace.code_cells.values_mut() { + for wait_tool_call_id in string_array_field(cell, "wait_tool_call_ids") { + let Some(wait_tool) = self.trace.tool_calls.get(&wait_tool_call_id) else { + continue; + }; + for raw_payload_id in tool_raw_payload_ids(wait_tool) { + push_unique_json_string(cell, "raw_payload_ids", &raw_payload_id); + } + } + } + } + + fn sync_terminal_model_observations(&mut self) { + let operation_ids = self + .trace + .terminal_operations + .keys() + .cloned() + .collect::>(); + for operation_id in operation_ids { + let Some(tool_call_id) = self + .trace + .terminal_operations + .get(&operation_id) + .and_then(|operation| operation.get("tool_call_id")) + .and_then(Value::as_str) + .map(str::to_string) + else { + continue; + }; + let Some(tool) = self.trace.tool_calls.get(&tool_call_id) else { + continue; + }; + // Direct shell tools are observed by the model through their own + // call/output items. Nested shell tools are only observed through + // the enclosing code cell's output, so terminal playback needs to + // point at the code cell in that case. + let observation = if tool + .get("requester") + .and_then(|requester| requester.get("type")) + .and_then(Value::as_str) + == Some("code_cell") + { + let Some(code_cell_id) = tool + .get("requester") + .and_then(|requester| requester.get("code_cell_id")) + .and_then(Value::as_str) + else { + continue; + }; + let Some(cell) = self.trace.code_cells.get(code_cell_id) else { + continue; + }; + json!({ + "call_item_ids": non_empty_array(cell.get("source_item_id").and_then(Value::as_str).unwrap_or("")), + "output_item_ids": string_array_field(cell, "output_item_ids"), + "source": "code_cell_output", + }) + } else { + json!({ + "call_item_ids": string_array_field(tool, "model_visible_call_item_ids"), + "output_item_ids": string_array_field(tool, "model_visible_output_item_ids"), + "source": "direct_tool_call", + }) + }; + if let Some(operation) = self.trace.terminal_operations.get_mut(&operation_id) { + operation["model_observations"] = json!([observation]); + } + } + } + + fn add_conversation_item_producer(&mut self, item_id: &str, producer: Value) { + let Some(item) = self.trace.conversation_items.get_mut(item_id) else { + return; + }; + let Some(produced_by) = item.get_mut("produced_by").and_then(Value::as_array_mut) else { + return; + }; + if !produced_by.iter().any(|existing| existing == &producer) { + produced_by.push(producer); + } + } + + fn remove_conversation_item_producer(&mut self, producer: Value) { + for item in self.trace.conversation_items.values_mut() { + let Some(produced_by) = item.get_mut("produced_by").and_then(Value::as_array_mut) + else { + continue; + }; + produced_by.retain(|existing| existing != &producer); + } } fn attach_tool_payloads_to_interaction_edges(&mut self) { let mut tool_payloads = BTreeMap::new(); for (tool_call_id, tool) in &self.trace.tool_calls { - let mut raw_payload_ids = Vec::new(); - for field in ["raw_invocation_payload_id", "raw_result_payload_id"] { - if let Some(raw_payload_id) = tool.get(field).and_then(Value::as_str) { - push_unique(&mut raw_payload_ids, raw_payload_id); - } - } - if let Some(runtime_payload_ids) = tool - .get("raw_runtime_payload_ids") - .and_then(Value::as_array) - { - for raw_payload_id in runtime_payload_ids.iter().filter_map(Value::as_str) { - push_unique(&mut raw_payload_ids, raw_payload_id); - } - } - + let raw_payload_ids = tool_raw_payload_ids(tool); let mut item_ids = Vec::new(); for field in [ "model_visible_call_item_ids", @@ -877,6 +1370,234 @@ fn normalize_body(kind: &str, item: &Value, raw_payload_id: &str) -> Value { } } +fn reduced_code_cell_id(model_visible_call_id: &str) -> String { + format!("code_cell:{model_visible_call_id}") +} + +fn tool_call_kind(tool_name: &str) -> Value { + if is_terminal_tool(tool_name) { + return json!({ "type": terminal_operation_kind(tool_name) }); + } + + match tool_name { + "exec" => json!({ "type": "code_cell" }), + "spawn_agent" | "send_message" | "followup_task" | "wait_agent" | "close_agent" + | "list_agents" => json!({ "type": "agent_interaction", "name": tool_name }), + _ => json!({ "type": "other", "name": tool_name }), + } +} + +fn is_terminal_tool(tool_name: &str) -> bool { + matches!( + tool_name, + "exec_command" | "local_shell" | "shell" | "shell_command" | "write_stdin" + ) +} + +fn terminal_operation_kind(tool_name: &str) -> &'static str { + match tool_name { + "write_stdin" => "write_stdin", + "exec_command" | "local_shell" | "shell" | "shell_command" => "exec_command", + _ => "terminal", + } +} + +fn terminal_request_from_invocation(invocation: &Value) -> Value { + let tool_name = invocation + .get("tool_name") + .and_then(Value::as_str) + .unwrap_or(""); + let payload = invocation.get("payload").unwrap_or(&Value::Null); + match payload.get("type").and_then(Value::as_str) { + Some("function") => { + let arguments = payload + .get("arguments") + .and_then(Value::as_str) + .and_then(parse_json_object) + .unwrap_or_default(); + terminal_request_from_arguments(tool_name, &arguments) + } + Some("local_shell") => { + let command = payload.get("command").cloned().unwrap_or_else(|| json!([])); + json!({ + "type": "exec_command", + "command": command, + "display_command": display_command(&command), + "cwd": payload.get("workdir").and_then(Value::as_str).unwrap_or(""), + "timeout_ms": payload.get("timeout_ms").cloned().unwrap_or(Value::Null), + }) + } + _ => json!({ + "type": terminal_operation_kind(tool_name), + "raw_payload": payload, + }), + } +} + +fn terminal_request_from_arguments(tool_name: &str, arguments: &Map) -> Value { + if tool_name == "write_stdin" { + return json!({ + "type": "write_stdin", + "session_id": arguments.get("session_id").cloned().unwrap_or(Value::Null), + "chars": arguments.get("chars").and_then(Value::as_str).unwrap_or(""), + }); + } + + let command = arguments + .get("command") + .cloned() + .or_else(|| arguments.get("cmd").map(|cmd| json!([cmd]))) + .unwrap_or_else(|| json!([])); + json!({ + "type": "exec_command", + "command": command, + "display_command": display_command(&command), + "cwd": arguments + .get("workdir") + .or_else(|| arguments.get("cwd")) + .and_then(Value::as_str) + .unwrap_or(""), + "yield_time_ms": arguments.get("yield_time_ms").cloned().unwrap_or(Value::Null), + "max_output_tokens": arguments.get("max_output_tokens").cloned().unwrap_or(Value::Null), + }) +} + +fn parse_json_object(raw: &str) -> Option> { + serde_json::from_str::(raw) + .ok() + .and_then(|value| value.as_object().cloned()) +} + +fn display_command(command: &Value) -> String { + if let Some(text) = command.as_str() { + return text.to_string(); + } + + command + .as_array() + .into_iter() + .flatten() + .filter_map(Value::as_str) + .collect::>() + .join(" ") +} + +fn terminal_id_from_result_payload(payload: &Value, operation_id: &str) -> Option { + let runtime_terminal_id = payload + .get("code_mode_result") + .and_then(|result| result.get("session_id")) + .and_then(|session_id| { + session_id + .as_u64() + .map(|id| id.to_string()) + .or_else(|| session_id.as_str().map(str::to_string)) + }); + if runtime_terminal_id.is_some() { + return runtime_terminal_id; + } + + // One-shot commands finish before the runtime has a persistent session id. + // Still create a terminal session so every terminal operation has a parent + // session object for viewers to hang lifecycle UI from. + payload + .get("code_mode_result") + .filter(|result| result.is_object()) + .and_then(|_| (!operation_id.is_empty()).then(|| format!("terminal:{operation_id}"))) +} + +fn terminal_id_from_operation_request(operation: &Value) -> Option { + // Continuation calls such as `write_stdin` already name the runtime + // session in their request. Completed chunks may omit that id from the + // result payload, so falling back to the request keeps all operations in + // the same terminal session instead of inventing a synthetic one. + operation + .get("request") + .and_then(terminal_id_from_terminal_request) +} + +fn terminal_id_from_terminal_request(request: &Value) -> Option { + request.get("session_id").and_then(|session_id| { + session_id + .as_u64() + .map(|id| id.to_string()) + .or_else(|| session_id.as_str().map(str::to_string)) + }) +} + +fn tool_raw_payload_ids(tool: &Value) -> Vec { + let mut raw_payload_ids = Vec::new(); + for field in ["raw_invocation_payload_id", "raw_result_payload_id"] { + if let Some(raw_payload_id) = tool.get(field).and_then(Value::as_str) { + push_unique(&mut raw_payload_ids, raw_payload_id); + } + } + for raw_payload_id in tool + .get("raw_runtime_payload_ids") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_str) + { + push_unique(&mut raw_payload_ids, raw_payload_id); + } + raw_payload_ids +} + +fn terminal_result_from_payload(payload: &Value, output_preview: Option<&str>) -> Value { + let code_mode_result = payload.get("code_mode_result").unwrap_or(&Value::Null); + json!({ + "output_preview": output_preview.unwrap_or(""), + "output": code_mode_result + .get("output") + .and_then(Value::as_str) + .unwrap_or_else(|| payload.get("output_preview").and_then(Value::as_str).unwrap_or("")), + "chunk_id": code_mode_result.get("chunk_id").cloned().unwrap_or(Value::Null), + "exit_code": code_mode_result.get("exit_code").cloned().unwrap_or(Value::Null), + "wall_time_seconds": code_mode_result + .get("wall_time_seconds") + .cloned() + .unwrap_or(Value::Null), + "original_token_count": code_mode_result + .get("original_token_count") + .cloned() + .unwrap_or(Value::Null), + "success": payload.get("success").cloned().unwrap_or(Value::Null), + }) +} + +fn code_cell_execution_status(runtime_status: &str) -> &'static str { + match runtime_status { + "completed" => "completed", + "failed" => "failed", + "terminated" => "cancelled", + "yielded" | "running" => "running", + _ => "completed", + } +} + +fn non_empty_array(value: &str) -> Value { + if value.is_empty() { + json!([]) + } else { + json!([value]) + } +} + +fn push_unique_json_string(object: &mut Value, key: &str, value: &str) { + if value.is_empty() { + return; + } + let Some(array) = object.get_mut(key).and_then(Value::as_array_mut) else { + return; + }; + if !array + .iter() + .any(|existing| existing.as_str() == Some(value)) + { + array.push(Value::String(value.to_string())); + } +} + fn tool_output_text(output: Option<&Value>) -> String { let Some(output) = output else { return String::new(); @@ -884,6 +1605,13 @@ fn tool_output_text(output: Option<&Value>) -> String { if let Some(text) = output.as_str() { return text.to_string(); } + if let Some(items) = output.as_array() { + return items + .iter() + .filter_map(|item| item.get("text").and_then(Value::as_str)) + .collect::>() + .join(""); + } if let Some(text) = output.get("content").and_then(Value::as_str) { return text.to_string(); } @@ -953,11 +1681,15 @@ fn trace_field_str<'a>(event: &'a CapturedEvent, namespace: &str, name: &str) -> fn normalized_tool_call_id(event: &CapturedEvent) -> Option { let call_id = trace_field_str(event, "tool", "call_id").filter(|call_id| !call_id.is_empty())?; - Some(if call_id.starts_with("tool:") { + Some(normalize_tool_call_id(call_id)) +} + +fn normalize_tool_call_id(call_id: &str) -> String { + if call_id.starts_with("tool:") { call_id.to_string() } else { format!("tool:{call_id}") - }) + } } fn interaction_edge_id(kind: &str, tool_call_id: &str) -> String { @@ -1123,14 +1855,14 @@ mod tests { std::fs::write( temp.path().join("payloads/1.json"), - serde_json::to_vec(&json!({ "input": [user_message.clone()] }))?, + serde_json::to_vec(&json!({ "input": [user_message] }))?, )?; std::fs::write( temp.path().join("payloads/2.json"), serde_json::to_vec(&json!({ "response_id": "resp-1", "token_usage": { "input_tokens": 1, "output_tokens": 2 }, - "output_items": [tool_call.clone()] + "output_items": [tool_call] }))?, )?; // The next request includes the whole model-visible input again. The @@ -1412,6 +2144,647 @@ mod tests { Ok(()) } + #[test] + fn reduces_code_cell_nested_terminal_tool_relationships() -> anyhow::Result<()> { + let temp = TempDir::new()?; + std::fs::write( + temp.path().join("manifest.json"), + serde_json::to_vec(&json!({ + "trace_id": "trace-test", + "started_at_unix_ms": 10, + }))?, + )?; + std::fs::create_dir(temp.path().join("payloads"))?; + + let user_message = json!({ + "type": "message", + "role": "user", + "content": [{ "type": "input_text", "text": "count files" }] + }); + let code_call = json!({ + "type": "custom_tool_call", + "status": "completed", + "call_id": "call-code", + "name": "exec", + "input": "await tools.exec_command({ cmd: \"find . -type f | wc -l\" });" + }); + let code_output = json!({ + "type": "custom_tool_call_output", + "call_id": "call-code", + "output": [{ "type": "input_text", "text": "Script completed\nOutput:\n42" }] + }); + + std::fs::write( + temp.path().join("payloads/1.json"), + serde_json::to_vec(&json!({ "input": [user_message] }))?, + )?; + std::fs::write( + temp.path().join("payloads/2.json"), + serde_json::to_vec(&json!({ "output_items": [code_call] }))?, + )?; + std::fs::write( + temp.path().join("payloads/3.json"), + serde_json::to_vec(&json!({ + "call_id": "call-code", + "tool_name": "exec", + "payload": { "type": "custom", "input": code_call["input"] } + }))?, + )?; + std::fs::write( + temp.path().join("payloads/4.json"), + serde_json::to_vec(&json!({ + "call_id": "runtime-tool-1", + "tool_name": "exec_command", + "payload": { + "type": "function", + "arguments": "{\"cmd\":\"find . -type f | wc -l\",\"workdir\":\"/repo\",\"max_output_tokens\":200,\"yield_time_ms\":1000}" + }, + "source": { + "type": "code_cell", + "runtime_cell_id": "runtime-cell-1", + "runtime_tool_call_id": "runtime-tool-1" + } + }))?, + )?; + std::fs::write( + temp.path().join("payloads/5.json"), + serde_json::to_vec(&json!({ + "status": "completed", + "success": true, + "output_preview": "Chunk ID: chunk-1\nWall time: 1.0 seconds\nOutput:\n", + "code_mode_result": { + "chunk_id": "chunk-1", + "wall_time_seconds": 1.0, + "session_id": 55007, + "original_token_count": 0, + "output": "" + } + }))?, + )?; + std::fs::write( + temp.path().join("payloads/6.json"), + serde_json::to_vec(&json!({ + "status": "completed", + "success": true, + "output_preview": "Script completed\nOutput:\n42" + }))?, + )?; + // Code-cell outputs are model-visible only when the follow-up request + // sends the complete conversation snapshot back to the provider. + std::fs::write( + temp.path().join("payloads/7.json"), + serde_json::to_vec(&json!({ + "input": [user_message, code_call, code_output] + }))?, + )?; + + let mut events = std::fs::File::create(temp.path().join("events.jsonl"))?; + write_event( + &mut events, + 1, + "codex.thread.started", + json!({"thread.id": "thread-1"}), + )?; + write_event( + &mut events, + 2, + "codex.turn.started", + json!({"thread.id": "thread-1", "turn.id": "turn-1"}), + )?; + write_event( + &mut events, + 3, + "codex.inference.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "inference.id": "inf-1", + "raw_payload.request.id": "raw_payload:1", + "raw_payload.request.path": "payloads/1.json", + "raw_payload.request.kind": "inference_request" + }), + )?; + write_event( + &mut events, + 4, + "codex.inference.completed", + json!({ + "inference.id": "inf-1", + "response.id": "resp-1", + "raw_payload.response.id": "raw_payload:2", + "raw_payload.response.path": "payloads/2.json", + "raw_payload.response.kind": "inference_response_summary" + }), + )?; + write_event( + &mut events, + 5, + "codex.tool.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "tool.call_id": "tool:call-code", + "tool.name": "exec", + "model_visible_call.id": "call-code", + "raw_payload.invocation.id": "raw_payload:3", + "raw_payload.invocation.path": "payloads/3.json", + "raw_payload.invocation.kind": "tool_invocation" + }), + )?; + write_event( + &mut events, + 6, + "codex.code_cell.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "code_cell.runtime_id": "runtime-cell-1", + "model_visible_call.id": "call-code", + "code_cell.source_js": "await tools.exec_command({ cmd: \"find . -type f | wc -l\" });" + }), + )?; + write_event( + &mut events, + 7, + "codex.tool.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "tool.call_id": "tool:nested-exec", + "tool.name": "exec_command", + "tool.requester.type": "code_cell", + "code_cell.runtime_id": "runtime-cell-1", + "code_mode_runtime_tool.id": "runtime-tool-1", + "raw_payload.invocation.id": "raw_payload:4", + "raw_payload.invocation.path": "payloads/4.json", + "raw_payload.invocation.kind": "tool_invocation" + }), + )?; + write_event( + &mut events, + 8, + "codex.tool.ended", + json!({ + "tool.call_id": "tool:nested-exec", + "status": "completed", + "output_preview": "Chunk ID: chunk-1\nWall time: 1.0 seconds\nOutput:\n", + "raw_payload.result.id": "raw_payload:5", + "raw_payload.result.path": "payloads/5.json", + "raw_payload.result.kind": "tool_result" + }), + )?; + write_event( + &mut events, + 9, + "codex.code_cell.ended", + json!({ + "thread.id": "thread-1", + "code_cell.runtime_id": "runtime-cell-1", + "status": "completed" + }), + )?; + write_event( + &mut events, + 10, + "codex.tool.ended", + json!({ + "tool.call_id": "tool:call-code", + "status": "completed", + "raw_payload.result.id": "raw_payload:6", + "raw_payload.result.path": "payloads/6.json", + "raw_payload.result.kind": "tool_result" + }), + )?; + write_event( + &mut events, + 11, + "codex.inference.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "inference.id": "inf-2", + "raw_payload.request.id": "raw_payload:7", + "raw_payload.request.path": "payloads/7.json", + "raw_payload.request.kind": "inference_request" + }), + )?; + + let trace = reduce_bundle(temp.path())?; + + assert_eq!( + trace.code_cells["code_cell:call-code"], + json!({ + "code_cell_id": "code_cell:call-code", + "model_visible_call_id": "call-code", + "thread_id": "thread-1", + "codex_turn_id": "turn-1", + "source_item_id": "item:2", + "output_item_ids": ["item:3"], + "runtime_cell_id": "runtime-cell-1", + "execution": { + "started_at_unix_ms": 1776420000000i64, + "ended_at_unix_ms": 1776420000000i64, + "status": "completed", + "started_seq": 6, + "ended_seq": 9, + }, + "runtime_status": "completed", + "source_js": "await tools.exec_command({ cmd: \"find . -type f | wc -l\" });", + "nested_tool_call_ids": ["tool:nested-exec"], + "wait_tool_call_ids": [], + "raw_payload_ids": ["raw_payload:3", "raw_payload:6"], + }) + ); + assert_eq!( + json!({ + "has_outer_exec_tool": trace.tool_calls.contains_key("tool:call-code"), + "nested_tool": trace.tool_calls["tool:nested-exec"].clone(), + "terminal_operation": trace.terminal_operations["terminal_operation:1"].clone(), + "terminal_session": trace.terminal_sessions["55007"].clone(), + "output_producers": trace.conversation_items["item:3"]["produced_by"].clone(), + }), + json!({ + "has_outer_exec_tool": false, + "nested_tool": { + "tool_call_id": "tool:nested-exec", + "thread_id": "thread-1", + "codex_turn_id": "turn-1", + "started_by_codex_turn_id": "turn-1", + "model_visible_call_id": null, + "code_mode_runtime_tool_id": "runtime-tool-1", + "model_visible_call_item_ids": [], + "model_visible_output_item_ids": [], + "requester": { + "type": "code_cell", + "code_cell_id": "code_cell:call-code", + "runtime_cell_id": "runtime-cell-1", + }, + "kind": { "type": "exec_command" }, + "terminal_operation_id": "terminal_operation:1", + "summary": { + "input_preview": "", + "output_preview": "Chunk ID: chunk-1\nWall time: 1.0 seconds\nOutput:\n", + }, + "execution": { + "started_at_unix_ms": 1776420000000i64, + "ended_at_unix_ms": 1776420000000i64, + "status": "completed", + "started_seq": 7, + "ended_seq": 8, + }, + "raw_invocation_payload_id": "raw_payload:4", + "raw_result_payload_id": "raw_payload:5", + "raw_runtime_payload_ids": [], + }, + "terminal_operation": { + "operation_id": "terminal_operation:1", + "terminal_id": "55007", + "tool_call_id": "tool:nested-exec", + "kind": "exec_command", + "execution": { + "started_at_unix_ms": 1776420000000i64, + "ended_at_unix_ms": 1776420000000i64, + "status": "completed", + "started_seq": 7, + "ended_seq": 8, + }, + "request": { + "type": "exec_command", + "command": ["find . -type f | wc -l"], + "display_command": "find . -type f | wc -l", + "cwd": "/repo", + "yield_time_ms": 1000, + "max_output_tokens": 200, + }, + "result": { + "output_preview": "Chunk ID: chunk-1\nWall time: 1.0 seconds\nOutput:\n", + "output": "", + "chunk_id": "chunk-1", + "exit_code": null, + "wall_time_seconds": 1.0, + "original_token_count": 0, + "success": true, + }, + "model_observations": [{ + "call_item_ids": ["item:2"], + "output_item_ids": ["item:3"], + "source": "code_cell_output", + }], + "raw_payload_ids": ["raw_payload:4", "raw_payload:5"], + }, + "terminal_session": { + "terminal_id": "55007", + "thread_id": "thread-1", + "created_by_operation_id": "terminal_operation:1", + "operation_ids": ["terminal_operation:1"], + "execution": { + "started_at_unix_ms": 1776420000000i64, + "ended_at_unix_ms": null, + "status": "running", + "started_seq": 7, + "ended_seq": null, + }, + }, + "output_producers": [ + { "type": "code_cell", "code_cell_id": "code_cell:call-code" }, + ], + }) + ); + Ok(()) + } + + #[test] + fn links_wait_outputs_and_payloads_to_code_cell() -> anyhow::Result<()> { + let temp = TempDir::new()?; + std::fs::write( + temp.path().join("manifest.json"), + serde_json::to_vec(&json!({ + "trace_id": "trace-test", + "started_at_unix_ms": 10, + }))?, + )?; + std::fs::create_dir(temp.path().join("payloads"))?; + + let code_call = json!({ + "type": "custom_tool_call", + "status": "completed", + "call_id": "call-code", + "name": "exec", + "input": "await sleep(1); text('done');" + }); + let yielded_output = json!({ + "type": "custom_tool_call_output", + "call_id": "call-code", + "output": [{ "type": "input_text", "text": "Execution yielded. Call wait." }] + }); + let wait_call = json!({ + "type": "function_call", + "status": "completed", + "call_id": "call-wait", + "name": "wait", + "arguments": "{\"cell_id\":\"runtime-cell-1\"}" + }); + let wait_output = json!({ + "type": "function_call_output", + "call_id": "call-wait", + "output": [{ "type": "input_text", "text": "done" }] + }); + + std::fs::write( + temp.path().join("payloads/code-cell-invocation.json"), + serde_json::to_vec(&json!({ "source_js": code_call["input"] }))?, + )?; + std::fs::write( + temp.path().join("payloads/wait-invocation.json"), + serde_json::to_vec(&json!({ + "call_id": "call-wait", + "tool_name": "wait", + "payload": { + "type": "function", + "arguments": "{\"cell_id\":\"runtime-cell-1\"}" + } + }))?, + )?; + std::fs::write( + temp.path().join("payloads/wait-result.json"), + serde_json::to_vec(&json!({ + "status": "completed", + "success": true, + "output_preview": "done" + }))?, + )?; + // Wait outputs are only canonical conversation items after the next + // provider request carries the full snapshot back to the model. + std::fs::write( + temp.path().join("payloads/request.json"), + serde_json::to_vec(&json!({ + "input": [code_call, yielded_output, wait_call, wait_output] + }))?, + )?; + + let mut events = std::fs::File::create(temp.path().join("events.jsonl"))?; + write_event( + &mut events, + 1, + "codex.thread.started", + json!({"thread.id": "thread-1"}), + )?; + write_event( + &mut events, + 2, + "codex.turn.started", + json!({"thread.id": "thread-1", "turn.id": "turn-1"}), + )?; + write_event( + &mut events, + 3, + "codex.code_cell.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "code_cell.runtime_id": "runtime-cell-1", + "model_visible_call.id": "call-code", + "code_cell.source_js": "await sleep(1); text('done');", + "raw_payload.invocation.id": "raw_payload:code-cell-invocation", + "raw_payload.invocation.path": "payloads/code-cell-invocation.json", + "raw_payload.invocation.kind": "code_cell_invocation" + }), + )?; + write_event( + &mut events, + 4, + "codex.tool.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "tool.call_id": "tool:call-wait", + "tool.name": "wait", + "tool.requester.type": "model", + "model_visible_call.id": "call-wait", + "raw_payload.invocation.id": "raw_payload:wait-invocation", + "raw_payload.invocation.path": "payloads/wait-invocation.json", + "raw_payload.invocation.kind": "tool_invocation" + }), + )?; + write_event( + &mut events, + 5, + "codex.tool.ended", + json!({ + "tool.call_id": "tool:call-wait", + "status": "completed", + "output_preview": "done", + "raw_payload.result.id": "raw_payload:wait-result", + "raw_payload.result.path": "payloads/wait-result.json", + "raw_payload.result.kind": "tool_result" + }), + )?; + write_event( + &mut events, + 6, + "codex.code_cell.ended", + json!({ + "thread.id": "thread-1", + "code_cell.runtime_id": "runtime-cell-1", + "status": "completed", + "model_visible_wait_call.id": "call-wait" + }), + )?; + write_event( + &mut events, + 7, + "codex.inference.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "inference.id": "inf-1", + "raw_payload.request.id": "raw_payload:request", + "raw_payload.request.path": "payloads/request.json", + "raw_payload.request.kind": "inference_request" + }), + )?; + + let trace = reduce_bundle(temp.path())?; + + assert_eq!( + json!({ + "output_item_ids": trace.code_cells["code_cell:call-code"]["output_item_ids"].clone(), + "wait_tool_call_ids": trace.code_cells["code_cell:call-code"]["wait_tool_call_ids"].clone(), + "raw_payload_ids": trace.code_cells["code_cell:call-code"]["raw_payload_ids"].clone(), + "wait_output_producers": trace.conversation_items["item:4"]["produced_by"].clone(), + }), + json!({ + "output_item_ids": ["item:2", "item:4"], + "wait_tool_call_ids": ["tool:call-wait"], + "raw_payload_ids": [ + "raw_payload:code-cell-invocation", + "raw_payload:wait-invocation", + "raw_payload:wait-result", + ], + "wait_output_producers": [ + { "type": "tool", "tool_call_id": "tool:call-wait" }, + { "type": "code_cell", "code_cell_id": "code_cell:call-code" }, + ], + }) + ); + Ok(()) + } + + #[test] + fn write_stdin_terminal_operation_uses_request_session_id() -> anyhow::Result<()> { + let temp = TempDir::new()?; + std::fs::write( + temp.path().join("manifest.json"), + serde_json::to_vec(&json!({ + "trace_id": "trace-test", + "started_at_unix_ms": 10, + }))?, + )?; + std::fs::create_dir(temp.path().join("payloads"))?; + + std::fs::write( + temp.path().join("payloads/1.json"), + serde_json::to_vec(&json!({ + "call_id": "call-write", + "tool_name": "write_stdin", + "payload": { + "type": "function", + "arguments": "{\"session_id\":12121,\"chars\":\"\"}" + } + }))?, + )?; + std::fs::write( + temp.path().join("payloads/2.json"), + serde_json::to_vec(&json!({ + "status": "completed", + "success": true, + "output_preview": "Chunk ID: chunk-1\nWall time: 2.0 seconds\nProcess exited with code 0\nOutput:\n42\n", + "code_mode_result": { + "chunk_id": "chunk-1", + "wall_time_seconds": 2.0, + "exit_code": 0, + "original_token_count": 1, + "output": "42\n" + } + }))?, + )?; + + let mut events = std::fs::File::create(temp.path().join("events.jsonl"))?; + write_event( + &mut events, + 1, + "codex.thread.started", + json!({"thread.id": "thread-1"}), + )?; + write_event( + &mut events, + 2, + "codex.turn.started", + json!({"thread.id": "thread-1", "turn.id": "turn-1"}), + )?; + write_event( + &mut events, + 3, + "codex.tool.started", + json!({ + "thread.id": "thread-1", + "turn.id": "turn-1", + "tool.call_id": "tool:call-write", + "tool.name": "write_stdin", + "tool.requester.type": "model", + "model_visible_call.id": "call-write", + "raw_payload.invocation.id": "raw_payload:1", + "raw_payload.invocation.path": "payloads/1.json", + "raw_payload.invocation.kind": "tool_invocation" + }), + )?; + write_event( + &mut events, + 4, + "codex.tool.ended", + json!({ + "tool.call_id": "tool:call-write", + "status": "completed", + "output_preview": "Chunk ID: chunk-1\nWall time: 2.0 seconds\nProcess exited with code 0\nOutput:\n42\n", + "raw_payload.result.id": "raw_payload:2", + "raw_payload.result.path": "payloads/2.json", + "raw_payload.result.kind": "tool_result" + }), + )?; + + let trace = reduce_bundle(temp.path())?; + + // `write_stdin` is a continuation of an existing terminal session. The + // final result can omit `session_id`, so the reducer must preserve the + // request-side id instead of creating a one-shot synthetic session. + assert_eq!( + trace.terminal_operations["terminal_operation:1"]["terminal_id"], + json!("12121") + ); + assert_eq!( + trace.terminal_sessions["12121"], + json!({ + "terminal_id": "12121", + "thread_id": "thread-1", + "created_by_operation_id": "terminal_operation:1", + "operation_ids": ["terminal_operation:1"], + "execution": { + "started_at_unix_ms": 1776420000000i64, + "ended_at_unix_ms": null, + "status": "running", + "started_seq": 3, + "ended_seq": null, + }, + }) + ); + assert!( + !trace + .terminal_sessions + .contains_key("terminal:terminal_operation:1") + ); + Ok(()) + } + fn write_event( file: &mut std::fs::File, _seq: u64,