diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 880e84ef4a..148dd86d12 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -23,7 +23,9 @@ pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS; pub use runtime::ExecuteRequest; pub use runtime::RuntimeResponse; pub use runtime::WaitRequest; +pub use runtime::WaitResponse; 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..507a4082ed 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -10,6 +10,7 @@ use std::sync::mpsc as std_mpsc; use std::thread; use codex_protocol::ToolName; +use serde::Serialize; use serde_json::Value as JsonValue; use tokio::sync::mpsc; @@ -25,6 +26,12 @@ const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; #[derive(Clone, Debug)] pub struct ExecuteRequest { + /// Runtime cell id to use for this execution. + /// + /// Hosts that need to trace work before JavaScript starts can allocate an id + /// first and pass it here. `None` keeps the service-owned allocation path + /// for callers that only need the id once a runtime response is returned. + pub cell_id: Option, pub tool_call_id: String, pub enabled_tools: Vec, pub source: String, @@ -41,6 +48,33 @@ pub struct WaitRequest { } #[derive(Debug, PartialEq)] +pub enum WaitResponse { + /// The requested cell was live when the wait command was accepted. + /// + /// Non-yielding responses from this variant are terminal lifecycle points + /// for the matching code cell. + Cell(RuntimeResponse), + /// The requested cell was not live, so the response is only the result of + /// the `wait` tool call. It must not be treated as a code-cell lifecycle + /// event because there is no cell to complete. + MissingCell(RuntimeResponse), +} + +impl WaitResponse { + pub fn into_runtime_response(self) -> RuntimeResponse { + match self { + WaitResponse::Cell(response) | WaitResponse::MissingCell(response) => response, + } + } + + pub fn runtime_response(&self) -> &RuntimeResponse { + match self { + WaitResponse::Cell(response) | WaitResponse::MissingCell(response) => response, + } + } +} + +#[derive(Debug, PartialEq, Serialize)] pub enum RuntimeResponse { Yielded { cell_id: String, @@ -331,6 +365,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..1c8bc15559 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -21,14 +21,26 @@ use crate::runtime::RuntimeEvent; use crate::runtime::RuntimeResponse; use crate::runtime::TurnMessage; use crate::runtime::WaitRequest; +use crate::runtime::WaitResponse; use crate::runtime::spawn_runtime; +/// Nested tool request emitted by one code-mode cell. +/// +/// Code mode owns the per-cell runtime id. Hosts should preserve it for +/// provenance/debugging, but should still assign their own runtime tool call id +/// if their tool-call graph requires globally unique ids. +pub struct CodeModeToolInvocation { + pub cell_id: String, + 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 +88,44 @@ impl CodeModeService { *self.inner.stored_values.lock().await = values; } - pub async fn execute(&self, request: ExecuteRequest) -> Result { - let cell_id = self - .inner + /// Reserves the runtime cell id for a future `execute` request. + /// + /// The runtime can issue nested tool calls before the first `execute` + /// response is returned. Hosts that need a parent trace object for those + /// nested calls should allocate the cell id up front and pass it back on the + /// `ExecuteRequest`. + 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()); + let mut sessions = self.inner.sessions.lock().await; + if sessions.contains_key(&cell_id) { + return Err(format!("exec cell {cell_id} already exists")); + } + + // Keep the session registry locked through insertion so a caller-owned + // cell id cannot race with another execute and replace a live runtime. 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( + 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), @@ -113,7 +145,7 @@ impl CodeModeService { .map_err(|_| "exec runtime ended unexpectedly".to_string()) } - pub async fn wait(&self, request: WaitRequest) -> Result { + pub async fn wait(&self, request: WaitRequest) -> Result { let cell_id = request.cell_id.clone(); let handle = self .inner @@ -123,7 +155,7 @@ impl CodeModeService { .get(&request.cell_id) .cloned(); let Some(handle) = handle else { - return Ok(missing_cell_response(cell_id)); + return Ok(WaitResponse::MissingCell(missing_cell_response(cell_id))); }; let (response_tx, response_rx) = oneshot::channel(); let control_message = if request.terminate { @@ -135,11 +167,13 @@ impl CodeModeService { } }; if handle.control_tx.send(control_message).is_err() { - return Ok(missing_cell_response(cell_id)); + return Ok(WaitResponse::MissingCell(missing_cell_response(cell_id))); } match response_rx.await { - Ok(response) => Ok(response), - Err(_) => Ok(missing_cell_response(request.cell_id)), + Ok(response) => Ok(WaitResponse::Cell(response)), + Err(_) => Ok(WaitResponse::MissingCell(missing_cell_response( + request.cell_id, + ))), } } @@ -181,9 +215,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() @@ -482,6 +521,8 @@ mod tests { use super::RuntimeResponse; use super::SessionControlCommand; use super::SessionControlContext; + use super::WaitRequest; + use super::WaitResponse; use super::run_session_control; use crate::FunctionCallOutputContentItem; use crate::runtime::ExecuteRequest; @@ -490,6 +531,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(), @@ -832,6 +874,30 @@ image({ ); } + #[tokio::test] + async fn wait_reports_missing_cell_separately_from_runtime_results() { + let service = CodeModeService::new(); + + let response = service + .wait(WaitRequest { + cell_id: "missing".to_string(), + yield_time_ms: 1, + terminate: false, + }) + .await + .unwrap(); + + assert_eq!( + response, + WaitResponse::MissingCell(RuntimeResponse::Result { + cell_id: "missing".to_string(), + content_items: Vec::new(), + stored_values: HashMap::new(), + error_text: Some("exec cell missing not found".to_string()), + }) + ); + } + #[tokio::test] async fn terminate_waits_for_runtime_shutdown_before_responding() { let inner = test_inner(); diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 19c0917f67..74131e58a2 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -2776,6 +2776,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { .expect("create environment"), )), /*analytics_events_client*/ None, + /*inherited_rollout_trace*/ None, ) .await; @@ -2898,6 +2899,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ..HooksConfig::default() }), rollout: Mutex::new(None), + rollout_trace: None, user_shell: Arc::new(default_user_shell()), agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( config.as_ref(), @@ -3758,6 +3760,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( ..HooksConfig::default() }), rollout: Mutex::new(None), + rollout_trace: None, user_shell: Arc::new(default_user_shell()), agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( config.as_ref(), @@ -5720,6 +5723,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() { tracker: Arc::clone(&turn_diff_tracker), call_id, tool_name: codex_tools::ToolName::plain(tool_name), + source: crate::tools::context::ToolCallSource::Direct, payload: ToolPayload::Function { arguments: serde_json::json!({ "command": params.command.clone(), @@ -5798,6 +5802,7 @@ async fn unified_exec_rejects_escalated_permissions_when_policy_not_on_request() tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), + source: crate::tools::context::ToolCallSource::Direct, payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", diff --git a/codex-rs/core/src/codex_tests_guardian.rs b/codex-rs/core/src/codex_tests_guardian.rs index cad67dcc8f..a30a642254 100644 --- a/codex-rs/core/src/codex_tests_guardian.rs +++ b/codex-rs/core/src/codex_tests_guardian.rs @@ -145,6 +145,7 @@ async fn guardian_allows_shell_additional_permissions_requests_past_policy_valid tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "test-call".to_string(), tool_name: codex_tools::ToolName::plain("shell"), + source: crate::tools::context::ToolCallSource::Direct, payload: ToolPayload::Function { arguments: serde_json::json!({ "command": params.command.clone(), @@ -211,6 +212,7 @@ async fn guardian_allows_unified_exec_additional_permissions_requests_past_polic tracker: Arc::clone(&tracker), call_id: "exec-call".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), + source: crate::tools::context::ToolCallSource::Direct, payload: ToolPayload::Function { arguments: serde_json::json!({ "cmd": "echo hi", @@ -324,6 +326,7 @@ async fn shell_handler_allows_sticky_turn_permissions_without_inline_request_per tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), call_id: "sticky-turn-grant".to_string(), tool_name: codex_tools::ToolName::plain("shell"), + source: crate::tools::context::ToolCallSource::Direct, payload: ToolPayload::Function { arguments: serde_json::json!({ "command": [ @@ -449,6 +452,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { metrics_service_name: None, inherited_shell_snapshot: None, inherited_exec_policy: Some(Arc::new(parent_exec_policy)), + inherited_rollout_trace: None, user_shell_override: None, parent_trace: None, analytics_events_client: None, 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..b5397f63e7 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -30,12 +30,25 @@ impl CodeModeExecuteHandler { .code_mode_service .stored_values() .await; + // Allocate before starting V8 so the trace can create the parent + // CodeCell before model-authored JavaScript issues nested tool calls. + let runtime_cell_id = exec.session.services.code_mode_service.allocate_cell_id(); + if let Some(trace) = &exec.session.services.rollout_trace { + trace.record_code_cell_started( + exec.session.conversation_id.to_string(), + exec.turn.sub_id.clone(), + &runtime_cell_id, + &call_id, + &args.code, + ); + } let started_at = std::time::Instant::now(); let response = 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, @@ -45,6 +58,23 @@ impl CodeModeExecuteHandler { }) .await .map_err(FunctionCallError::RespondToModel)?; + if let Some(trace) = &exec.session.services.rollout_trace { + // The initial response is the model-visible custom-tool return. + // Yielded cells keep running, so terminal lifecycle is only emitted + // here when the first response also ended the runtime. + trace.record_code_cell_initial_response( + exec.session.conversation_id.to_string(), + exec.turn.sub_id.clone(), + &response, + ); + if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { + trace.record_code_cell_ended( + exec.session.conversation_id.to_string(), + exec.turn.sub_id.clone(), + &response, + ); + } + } handle_runtime_response(&exec, response, args.max_output_tokens, started_at) .await .map_err(FunctionCallError::RespondToModel) @@ -62,6 +92,15 @@ impl ToolHandler for CodeModeExecuteHandler { matches!(payload, ToolPayload::Custom { .. }) } + fn uses_first_class_trace_object(&self, invocation: &ToolInvocation) -> bool { + // `exec` is represented by the first-class CodeCell lifecycle. The + // dispatch-level ToolCall event would duplicate the same runtime + // boundary as a less precise object. + matches!(invocation.payload, ToolPayload::Custom { .. }) + && invocation.tool_name.namespace.is_none() + && invocation.tool_name.name == PUBLIC_TOOL_NAME + } + 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..938160c83f 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, @@ -83,7 +88,7 @@ impl CodeModeService { pub(crate) async fn wait( &self, request: codex_code_mode::WaitRequest, - ) -> Result { + ) -> Result { self.inner.wait(request).await } @@ -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..a2da500ade 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -61,7 +61,7 @@ impl ToolHandler for CodeModeWaitHandler { let args: ExecWaitArgs = parse_arguments(&arguments)?; let exec = ExecContext { session, turn }; let started_at = std::time::Instant::now(); - let response = exec + let wait_response = exec .session .services .code_mode_service @@ -72,9 +72,28 @@ impl ToolHandler for CodeModeWaitHandler { }) .await .map_err(FunctionCallError::RespondToModel)?; - handle_runtime_response(&exec, response, args.max_tokens, started_at) - .await - .map_err(FunctionCallError::RespondToModel) + let response = wait_response.runtime_response(); + if matches!(&wait_response, codex_code_mode::WaitResponse::Cell(_)) + && !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) + && let Some(trace) = &exec.session.services.rollout_trace + { + // Only a live-cell wait can close a CodeCell. A missing + // cell is still an ordinary `wait` tool result, but there + // is no runtime object for the reducer to complete. + trace.record_code_cell_ended( + exec.session.conversation_id.to_string(), + exec.turn.sub_id.clone(), + response, + ); + } + handle_runtime_response( + &exec, + wait_response.into_runtime_response(), + args.max_tokens, + started_at, + ) + .await + .map_err(FunctionCallError::RespondToModel) } _ => Err(FunctionCallError::RespondToModel(format!( "{WAIT_TOOL_NAME} expects JSON arguments" diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 60bcbcda23..e2a7883ed6 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -28,11 +28,18 @@ 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. This is useful for + /// debugging the JS/runtime bridge, but it is not the Codex tool call id + /// because the runtime id only needs to be unique within one cell. + runtime_tool_call_id: String, + }, } #[derive(Clone)] @@ -42,6 +49,7 @@ pub struct ToolInvocation { pub tracker: SharedTurnDiffTracker, pub call_id: String, pub tool_name: ToolName, + pub source: ToolCallSource, 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..a42c9d59a8 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -70,6 +70,7 @@ fn invocation( tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(tool_name), + source: crate::tools::context::ToolCallSource::Direct, 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..ffd1139ed6 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 @@ -30,6 +30,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { tracker: Arc::new(Mutex::new(TurnDiffTracker::default())), call_id: "call-1".to_string(), tool_name: codex_tools::ToolName::plain(REQUEST_USER_INPUT_TOOL_NAME), + source: crate::tools::context::ToolCallSource::Direct, payload: ToolPayload::Function { arguments: json!({ "questions": [{ diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index e66133f961..ac52b68b83 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -228,6 +228,7 @@ async fn shell_pre_tool_use_payload_uses_joined_command() { tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-41".to_string(), tool_name: codex_tools::ToolName::plain("shell"), + source: crate::tools::context::ToolCallSource::Direct, payload, }), Some(crate::tools::registry::PreToolUsePayload { @@ -253,6 +254,7 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() { tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-42".to_string(), tool_name: codex_tools::ToolName::plain("shell_command"), + source: crate::tools::context::ToolCallSource::Direct, payload, }), Some(crate::tools::registry::PreToolUsePayload { 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..730decc021 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -213,6 +213,7 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() { tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-43".to_string(), tool_name: codex_tools::ToolName::plain("exec_command"), + source: crate::tools::context::ToolCallSource::Direct, payload, }), Some(crate::tools::registry::PreToolUsePayload { @@ -236,6 +237,7 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() { tracker: Arc::new(Mutex::new(TurnDiffTracker::new())), call_id: "call-44".to_string(), tool_name: codex_tools::ToolName::plain("write_stdin"), + source: crate::tools::context::ToolCallSource::Direct, payload, }), None diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 7e3713dcb4..0a60789cec 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -8,8 +8,10 @@ use crate::hook_runtime::record_additional_contexts; use crate::hook_runtime::run_post_tool_use_hooks; use crate::hook_runtime::run_pre_tool_use_hooks; use crate::memories::usage::emit_metric_for_tool_read; +use crate::rollout_trace::RolloutTraceRecorder; use crate::sandbox_tags::sandbox_tag; use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ToolCallSource; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; @@ -22,6 +24,7 @@ use codex_hooks::HookToolInputLocalShell; use codex_hooks::HookToolKind; use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::SandboxPolicy; +use codex_rollout_trace::ExecutionStatus; use codex_tools::ConfiguredToolSpec; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -74,6 +77,16 @@ pub trait ToolHandler: Send + Sync { None } + /// Returns `true` when this handler is represented by a trace object other + /// than `ToolCall`. + /// + /// Protocol events are runtime observations on the `ToolCall`; they do not + /// suppress the canonical tool boundary. The public code-mode `exec` tool is + /// the exception because `CodeCell` owns that model-visible boundary. + 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( @@ -132,6 +145,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, @@ -163,6 +178,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, @@ -184,6 +203,60 @@ pub struct ToolRegistry { handlers: HashMap>, } +/// No-op capable trace handle for one resolved tool dispatch. +/// +/// The registry has several early-return paths after a handler is selected: +/// pre-use hooks, handler execution, after-use hooks, and result status all +/// affect the trace lifecycle. Keeping the trace eligibility and event writes +/// behind this helper makes those paths say what happened instead of repeating +/// the Direct/CodeMode/JsRepl/first-class-object policy at each branch. +struct DispatchTrace { + recorder: Option, +} + +impl DispatchTrace { + fn new(invocation: &ToolInvocation, handler: &dyn AnyToolHandler) -> Self { + let should_trace = matches!( + invocation.source, + ToolCallSource::Direct | ToolCallSource::CodeMode { .. } + ) && !handler.uses_first_class_trace_object(invocation); + + let recorder = should_trace + .then(|| invocation.session.services.rollout_trace.clone()) + .flatten(); + Self { recorder } + } + + fn record_started(&self, invocation: &ToolInvocation) { + if let Some(recorder) = &self.recorder { + recorder.record_dispatched_tool_call_started(invocation); + } + } + + fn record_completed(&self, invocation: &ToolInvocation, result: &AnyToolResult) { + if let Some(recorder) = &self.recorder { + let status = if result.result.success_for_logging() { + ExecutionStatus::Completed + } else { + ExecutionStatus::Failed + }; + recorder.record_dispatched_tool_call_ended( + invocation, + status, + result.result.as_ref(), + &result.call_id, + &result.payload, + ); + } + } + + fn record_failed(&self, invocation: &ToolInvocation, error: &FunctionCallError) { + if let Some(recorder) = &self.recorder { + recorder.record_dispatched_tool_call_failed(invocation, &error.to_string()); + } + } +} + impl ToolRegistry { fn new(handlers: HashMap>) -> Self { Self { handlers } @@ -288,6 +361,9 @@ impl ToolRegistry { return Err(FunctionCallError::Fatal(message)); } + let dispatch_trace = DispatchTrace::new(&invocation, handler.as_ref()); + dispatch_trace.record_started(&invocation); + if let Some(pre_tool_use_payload) = handler.pre_tool_use_payload(&invocation) && let Some(reason) = run_pre_tool_use_hooks( &invocation.session, @@ -297,10 +373,12 @@ impl ToolRegistry { ) .await { - return Err(FunctionCallError::RespondToModel(format!( + let err = FunctionCallError::RespondToModel(format!( "Command blocked by PreToolUse hook: {reason}. Command: {}", pre_tool_use_payload.command - ))); + )); + dispatch_trace.record_failed(&invocation, &err); + return Err(err); } let is_mutating = handler.is_mutating(&invocation).await; @@ -383,6 +461,7 @@ impl ToolRegistry { .await; if let Some(err) = hook_abort_error { + dispatch_trace.record_failed(&invocation, &err); return Err(err); } @@ -422,9 +501,13 @@ impl ToolRegistry { let result = guard.take().ok_or_else(|| { FunctionCallError::Fatal("tool produced no output".to_string()) })?; + dispatch_trace.record_completed(&invocation, &result); Ok(result) } - Err(err) => Err(err), + Err(err) => { + dispatch_trace.record_failed(&invocation, &err); + Err(err) + } } } } diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index f4fb72e791..8a9721d4d1 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -1,7 +1,25 @@ use super::*; +use crate::codex::make_session_and_context; +use crate::rollout_trace::RolloutTraceRecorder; +use crate::rollout_trace::ThreadStartedTraceMetadata; +use crate::tools::code_mode::CodeModeWaitHandler; +use crate::tools::code_mode::WAIT_TOOL_NAME; +use crate::turn_diff_tracker::TurnDiffTracker; +use codex_protocol::config_types::ModeKind; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TurnStartedEvent; +use codex_rollout_trace::ToolCallRequester; use pretty_assertions::assert_eq; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use tempfile::TempDir; -struct TestHandler; +#[derive(Default)] +struct TestHandler { + first_class_trace_object: bool, +} impl ToolHandler for TestHandler { type Output = crate::tools::context::FunctionToolOutput; @@ -10,15 +28,22 @@ impl ToolHandler for TestHandler { ToolKind::Function } + fn uses_first_class_trace_object(&self, _invocation: &ToolInvocation) -> bool { + self.first_class_trace_object + } + async fn handle(&self, _invocation: ToolInvocation) -> Result { - unreachable!("test handler should not be invoked") + Ok(crate::tools::context::FunctionToolOutput::from_text( + "ok".to_string(), + Some(true), + )) } } #[test] fn handler_looks_up_namespaced_aliases_explicitly() { - let plain_handler = Arc::new(TestHandler) as Arc; - let namespaced_handler = Arc::new(TestHandler) as Arc; + let plain_handler = Arc::new(TestHandler::default()) as Arc; + let namespaced_handler = Arc::new(TestHandler::default()) as Arc; let namespace = "mcp__codex_apps__gmail"; let tool_name = "gmail_get_recent_emails"; let plain_name = codex_tools::ToolName::plain(tool_name); @@ -49,3 +74,248 @@ fn handler_looks_up_namespaced_aliases_explicitly() { .is_some_and(|handler| Arc::ptr_eq(handler, &namespaced_handler)) ); } + +#[tokio::test] +async fn dispatch_lifecycle_trace_records_direct_and_code_mode_requesters() -> anyhow::Result<()> { + let temp = TempDir::new()?; + let (mut session, turn) = make_session_and_context().await; + attach_test_trace(&mut session, &turn, temp.path())?; + session + .services + .rollout_trace + .as_ref() + .expect("trace recorder") + .record_code_cell_started( + session.conversation_id.to_string(), + turn.sub_id.clone(), + "cell-1", + "call-code", + "await tools.test_tool({})", + ); + + let registry = ToolRegistry::new(HashMap::from([( + codex_tools::ToolName::plain("test_tool"), + Arc::new(TestHandler::default()) as Arc, + )])); + let session = Arc::new(session); + let turn = Arc::new(turn); + + registry + .dispatch_any(test_invocation( + Arc::clone(&session), + Arc::clone(&turn), + "direct-call", + "test_tool", + ToolCallSource::Direct, + "{}", + )) + .await?; + registry + .dispatch_any(test_invocation( + session, + turn, + "code-mode-call", + "test_tool", + ToolCallSource::CodeMode { + cell_id: "cell-1".to_string(), + runtime_tool_call_id: "tool-1".to_string(), + }, + "{}", + )) + .await?; + + let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?; + assert_eq!( + replayed.tool_calls["direct-call"].model_visible_call_id, + Some("direct-call".to_string()), + ); + assert_eq!( + replayed.tool_calls["direct-call"].requester, + ToolCallRequester::Model, + ); + assert!( + replayed.tool_calls["direct-call"] + .raw_invocation_payload_id + .is_some(), + "dispatch tracing should keep the tool invocation payload", + ); + assert!( + replayed.tool_calls["direct-call"] + .raw_result_payload_id + .is_some(), + "direct calls should keep the model-facing result payload", + ); + assert_eq!( + replayed.tool_calls["code-mode-call"].model_visible_call_id, + None, + ); + assert_eq!( + replayed.tool_calls["code-mode-call"].code_mode_runtime_tool_id, + Some("tool-1".to_string()), + ); + assert_eq!( + replayed.tool_calls["code-mode-call"].requester, + ToolCallRequester::CodeCell { + code_cell_id: "code_cell:call-code".to_string(), + }, + ); + assert!( + replayed.tool_calls["code-mode-call"] + .raw_result_payload_id + .is_some(), + "code-mode calls should keep the result returned to JavaScript", + ); + + Ok(()) +} + +#[tokio::test] +async fn dispatch_lifecycle_trace_skips_noncanonical_boundaries() -> anyhow::Result<()> { + assert_dispatch_trace_skips( + Arc::new(TestHandler::default()) as Arc, + ToolCallSource::JsRepl, + ) + .await?; + assert_dispatch_trace_skips( + Arc::new(TestHandler { + first_class_trace_object: true, + }) as Arc, + ToolCallSource::Direct, + ) + .await +} + +async fn assert_dispatch_trace_skips( + handler: Arc, + source: ToolCallSource, +) -> anyhow::Result<()> { + let temp = TempDir::new()?; + let (mut session, turn) = make_session_and_context().await; + attach_test_trace(&mut session, &turn, temp.path())?; + + let registry = ToolRegistry::new(HashMap::from([( + codex_tools::ToolName::plain("test_tool"), + handler, + )])); + let session = Arc::new(session); + let turn = Arc::new(turn); + + registry + .dispatch_any(test_invocation( + session, + turn, + "skipped-call", + "test_tool", + source, + "{}", + )) + .await?; + + let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?; + assert_eq!(replayed.tool_calls, Default::default()); + + Ok(()) +} + +#[tokio::test] +async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Result<()> { + let temp = TempDir::new()?; + let (mut session, turn) = make_session_and_context().await; + attach_test_trace(&mut session, &turn, temp.path())?; + + let registry = ToolRegistry::new(HashMap::from([( + codex_tools::ToolName::plain(WAIT_TOOL_NAME), + Arc::new(CodeModeWaitHandler) as Arc, + )])); + let session = Arc::new(session); + let turn = Arc::new(turn); + + registry + .dispatch_any(test_invocation( + session, + turn, + "wait-call", + WAIT_TOOL_NAME, + ToolCallSource::Direct, + r#"{"cell_id":"noop","terminate":true}"#, + )) + .await?; + + let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?; + assert_eq!(replayed.code_cells.len(), 0); + assert!( + replayed.tool_calls["wait-call"] + .raw_result_payload_id + .is_some() + ); + + Ok(()) +} + +fn test_invocation( + session: Arc, + turn: Arc, + call_id: &str, + tool_name: &str, + source: ToolCallSource, + arguments: &str, +) -> ToolInvocation { + ToolInvocation { + session, + turn, + tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), + call_id: call_id.to_string(), + tool_name: codex_tools::ToolName::plain(tool_name), + source, + payload: ToolPayload::Function { + arguments: arguments.to_string(), + }, + } +} + +fn attach_test_trace( + session: &mut crate::codex::Session, + turn: &crate::codex::TurnContext, + root: &Path, +) -> anyhow::Result<()> { + let thread_id = session.conversation_id; + let recorder = RolloutTraceRecorder::create_in_root_for_test( + root, + thread_id, + ThreadStartedTraceMetadata { + thread_id: thread_id.to_string(), + agent_path: "/root".to_string(), + task_name: None, + nickname: None, + agent_role: None, + session_source: SessionSource::Exec, + cwd: PathBuf::from("/workspace"), + rollout_path: None, + model: "gpt-test".to_string(), + provider_name: "test-provider".to_string(), + approval_policy: "never".to_string(), + sandbox_policy: "danger-full-access".to_string(), + }, + )?; + recorder.record_codex_turn_event( + thread_id.to_string(), + &turn.sub_id, + &EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn.sub_id.clone(), + started_at: None, + model_context_window: None, + collaboration_mode_kind: ModeKind::default(), + }), + ); + session.services.rollout_trace = Some(recorder); + Ok(()) +} + +fn single_bundle_dir(root: &Path) -> anyhow::Result { + let mut entries = fs::read_dir(root)? + .map(|entry| entry.map(|entry| entry.path())) + .collect::, _>>()?; + entries.sort(); + assert_eq!(entries.len(), 1); + Ok(entries.remove(0)) +} diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index ef7ae32d4b..6cdc8bfc03 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -271,7 +271,7 @@ impl ToolRouter { let direct_js_repl_call = tool_name.namespace.is_none() && matches!(tool_name.name.as_str(), "js_repl" | "js_repl_reset"); - if source == ToolCallSource::Direct + if matches!(source, ToolCallSource::Direct) && turn.tools_config.js_repl_tools_only && !direct_js_repl_call { @@ -287,6 +287,7 @@ impl ToolRouter { tracker, call_id, tool_name, + source, payload, };