openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Trace tool and code-mode runtime boundaries

Browse Source
This commit is contained in:
Albin Cassirer
2026-04-15 16:09:07 -07:00
parent 856ef510c2
commit 2abf9dcd4e
16 changed files with 590 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/lib.rs
+2
View File
@@ -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;
codex-rs/code-mode/src/runtime/mod.rs
+35
View File
@@ -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<String>,
pub tool_call_id: String,
pub enabled_tools: Vec<ToolDefinition>,
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(),
codex-rs/code-mode/src/service.rs
+81 -15
View File
@@ -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<JsonValue>,
}
#[async_trait]
pub trait CodeModeTurnHost: Send + Sync {
async fn invoke_tool(
&self,
tool_name: ToolName,
input: Option<JsonValue>,
invocation: CodeModeToolInvocation,
cancellation_token: CancellationToken,
) -> Result<JsonValue, String>;
@@ -76,24 +88,44 @@ impl CodeModeService {
*self.inner.stored_values.lock().await = values;
}
pub async fn execute(&self, request: ExecuteRequest) -> Result<RuntimeResponse, String> {
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<RuntimeResponse, String> {
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<RuntimeResponse, String> {
pub async fn wait(&self, request: WaitRequest) -> Result<WaitResponse, String> {
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();
codex-rs/core/src/codex_tests.rs
+5
View File
@@ -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",
codex-rs/core/src/codex_tests_guardian.rs
+4
View File
@@ -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,
codex-rs/core/src/tools/code_mode/execute_handler.rs
+39
View File
@@ -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<Self::Output, FunctionCallError> {
let ToolInvocation {
session,
codex-rs/core/src/tools/code_mode/mod.rs
+23 -8
View File
@@ -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<RuntimeResponse, String> {
) -> Result<codex_code_mode::WaitResponse, String> {
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<JsonValue>,
invocation: CodeModeToolInvocation,
cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
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<JsonValue>,
invocation: CodeModeToolInvocation,
cancellation_token: CancellationToken,
) -> Result<JsonValue, FunctionCallError> {
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())
}
codex-rs/core/src/tools/code_mode/wait_handler.rs
+23 -4
View File
@@ -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"
codex-rs/core/src/tools/context.rs
+10 -2
View File
@@ -28,11 +28,18 @@ use tokio::sync::Mutex;
pub type SharedTurnDiffTracker = Arc<Mutex<TurnDiffTracker>>;
#[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,
}
codex-rs/core/src/tools/handlers/multi_agents_tests.rs
+1
View File
@@ -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,
}
}
codex-rs/core/src/tools/handlers/request_user_input_tests.rs
+1
View File
@@ -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": [{
codex-rs/core/src/tools/handlers/shell_tests.rs
+2
View File
@@ -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 {
codex-rs/core/src/tools/handlers/unified_exec_tests.rs
+2
View File
@@ -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
codex-rs/core/src/tools/registry.rs
+86 -3
View File
@@ -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<PostToolUsePayload>;
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<ToolName, Arc<dyn AnyToolHandler>>,
}
/// 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<RolloutTraceRecorder>,
}
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<ToolName, Arc<dyn AnyToolHandler>>) -> 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)
}
}
}
}
codex-rs/core/src/tools/registry_tests.rs
+274 -4
View File
@@ -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<Self::Output, FunctionCallError> {
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<dyn AnyToolHandler>;
let namespaced_handler = Arc::new(TestHandler) as Arc<dyn AnyToolHandler>;
let plain_handler = Arc::new(TestHandler::default()) as Arc<dyn AnyToolHandler>;
let namespaced_handler = Arc::new(TestHandler::default()) as Arc<dyn AnyToolHandler>;
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<dyn AnyToolHandler>,
)]));
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<dyn AnyToolHandler>,
ToolCallSource::JsRepl,
)
.await?;
assert_dispatch_trace_skips(
Arc::new(TestHandler {
first_class_trace_object: true,
}) as Arc<dyn AnyToolHandler>,
ToolCallSource::Direct,
)
.await
}
async fn assert_dispatch_trace_skips(
handler: Arc<dyn AnyToolHandler>,
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<dyn AnyToolHandler>,
)]));
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<crate::codex::Session>,
turn: Arc<crate::codex::TurnContext>,
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<PathBuf> {
let mut entries = fs::read_dir(root)?
.map(|entry| entry.map(|entry| entry.path()))
.collect::<Result<Vec<_>, _>>()?;
entries.sort();
assert_eq!(entries.len(), 1);
Ok(entries.remove(0))
}
codex-rs/core/src/tools/router.rs
+2 -1
View File
@@ -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,
};