Merge 482f2b2353 into sapling-pr-archive-bolinfest

This commit is contained in:
Michael Bolin
2026-04-23 12:46:37 -07:00
committed by GitHub
39 changed files with 1638 additions and 114 deletions
+2
View File
@@ -6875,6 +6875,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(),
@@ -6954,6 +6955,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",
@@ -9,6 +9,7 @@ use crate::exec_policy::ExecPolicyManager;
use crate::guardian::GUARDIAN_REVIEWER_NAME;
use crate::sandboxing::SandboxPermissions;
use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolCallSource;
use crate::turn_diff_tracker::TurnDiffTracker;
use codex_app_server_protocol::ConfigLayerSource;
use codex_exec_server::EnvironmentManager;
@@ -341,6 +342,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(),
@@ -476,6 +478,7 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_shell_policy_skip() {
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
call_id: "strict-shell-call".to_string(),
tool_name: codex_tools::ToolName::plain("shell"),
source: ToolCallSource::Direct,
payload: ToolPayload::Function {
arguments: serde_json::json!({
"command": command,
@@ -521,6 +524,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",
@@ -639,6 +643,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": [
@@ -9,6 +9,7 @@ use super::ExecContext;
use super::PUBLIC_TOOL_NAME;
use super::build_enabled_tools;
use super::handle_runtime_response;
use super::is_exec_tool_name;
pub struct CodeModeExecuteHandler;
@@ -30,12 +31,23 @@ 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();
let code_cell_trace = exec.session.services.rollout_trace.start_code_cell_trace(
exec.session.conversation_id,
exec.turn.sub_id.as_str(),
runtime_cell_id.as_str(),
call_id.as_str(),
args.code.as_str(),
);
let started_at = std::time::Instant::now();
let response = exec
.session
.services
.code_mode_service
.execute(codex_code_mode::ExecuteRequest {
cell_id: runtime_cell_id,
tool_call_id: call_id,
enabled_tools,
source: args.code,
@@ -45,6 +57,15 @@ impl CodeModeExecuteHandler {
})
.await
.map_err(FunctionCallError::RespondToModel)?;
// Record the raw runtime boundary. The model-visible custom-tool output
// is produced by `handle_runtime_response` and later linked through
// `CodeCell.output_item_ids` in the reduced trace.
code_cell_trace.record_initial_response(&response);
// Yielded cells keep running, so terminal lifecycle is only emitted
// here when the first response also ended the runtime.
if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) {
code_cell_trace.record_ended(&response);
}
handle_runtime_response(&exec, response, args.max_output_tokens, started_at)
.await
.map_err(FunctionCallError::RespondToModel)
@@ -73,9 +94,7 @@ impl ToolHandler for CodeModeExecuteHandler {
} = invocation;
match payload {
ToolPayload::Custom { input }
if tool_name.namespace.is_none() && tool_name.name.as_str() == PUBLIC_TOOL_NAME =>
{
ToolPayload::Custom { input } if is_exec_tool_name(&tool_name) => {
self.execute(session, turn, call_id, input).await
}
_ => Err(FunctionCallError::RespondToModel(format!(
+29 -9
View File
@@ -7,6 +7,7 @@ use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use codex_code_mode::CodeModeNestedToolCall;
use codex_code_mode::CodeModeTurnHost;
use codex_code_mode::RuntimeResponse;
use codex_protocol::models::FunctionCallOutputContentItem;
@@ -45,6 +46,11 @@ pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode::PUBLIC_TOOL_NAME;
pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode::WAIT_TOOL_NAME;
pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = codex_code_mode::DEFAULT_WAIT_YIELD_TIME_MS;
/// Returns true for the un-namespaced code-mode `exec` tool.
pub(crate) fn is_exec_tool_name(tool_name: &ToolName) -> bool {
tool_name.namespace.is_none() && tool_name.name == PUBLIC_TOOL_NAME
}
#[derive(Clone)]
pub(crate) struct ExecContext {
pub(super) session: Arc<Session>,
@@ -73,6 +79,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 +93,7 @@ impl CodeModeService {
pub(crate) async fn wait(
&self,
request: codex_code_mode::WaitRequest,
) -> Result<RuntimeResponse, String> {
) -> Result<codex_code_mode::WaitOutcome, String> {
self.inner.wait(request).await
}
@@ -118,15 +128,13 @@ struct CoreTurnHost {
impl CodeModeTurnHost for CoreTurnHost {
async fn invoke_tool(
&self,
tool_name: ToolName,
input: Option<JsonValue>,
invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
call_nested_tool(
self.exec.clone(),
self.tool_runtime.clone(),
tool_name,
input,
invocation,
cancellation_token,
)
.await
@@ -302,11 +310,16 @@ 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: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> Result<JsonValue, FunctionCallError> {
if tool_name.namespace.is_none() && tool_name.name == PUBLIC_TOOL_NAME {
let CodeModeNestedToolCall {
cell_id,
runtime_tool_call_id,
tool_name,
input,
} = invocation;
if is_exec_tool_name(&tool_name) {
return Err(FunctionCallError::RespondToModel(format!(
"{PUBLIC_TOOL_NAME} cannot invoke itself"
)));
@@ -339,7 +352,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())
}
@@ -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,7 +72,28 @@ impl ToolHandler for CodeModeWaitHandler {
})
.await
.map_err(FunctionCallError::RespondToModel)?;
handle_runtime_response(&exec, response, args.max_tokens, started_at)
if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response
&& !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. })
{
// 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.
let runtime_cell_id = match response {
codex_code_mode::RuntimeResponse::Yielded { cell_id, .. }
| codex_code_mode::RuntimeResponse::Terminated { cell_id, .. }
| codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id,
};
exec.session
.services
.rollout_trace
.code_cell_trace_context(
exec.session.conversation_id,
exec.turn.sub_id.as_str(),
runtime_cell_id,
)
.record_ended(response);
}
handle_runtime_response(&exec, wait_response.into(), args.max_tokens, started_at)
.await
.map_err(FunctionCallError::RespondToModel)
}
+10 -2
View File
@@ -30,11 +30,18 @@ use tokio_util::sync::CancellationToken;
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)]
@@ -45,6 +52,7 @@ pub struct ToolInvocation {
pub tracker: SharedTurnDiffTracker,
pub call_id: String,
pub tool_name: ToolName,
pub source: ToolCallSource,
pub payload: ToolPayload,
}
@@ -37,6 +37,7 @@ async fn invocation_for_payload(payload: ToolPayload) -> ToolInvocation {
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-apply-patch".to_string(),
tool_name: codex_tools::ToolName::plain("apply_patch"),
source: crate::tools::context::ToolCallSource::Direct,
payload,
}
}
+3
View File
@@ -112,6 +112,7 @@ fn mcp_hook_tool_input(raw_arguments: &str) -> Value {
mod tests {
use super::*;
use crate::session::tests::make_session_and_context;
use crate::tools::context::ToolCallSource;
use crate::turn_diff_tracker::TurnDiffTracker;
use pretty_assertions::assert_eq;
use serde_json::json;
@@ -141,6 +142,7 @@ mod tests {
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-mcp-pre".to_string(),
tool_name: codex_tools::ToolName::namespaced("mcp__memory__", "create_entities"),
source: ToolCallSource::Direct,
payload,
}),
Some(PreToolUsePayload {
@@ -188,6 +190,7 @@ mod tests {
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-mcp-post".to_string(),
tool_name: codex_tools::ToolName::namespaced("mcp__filesystem__", "read_file"),
source: ToolCallSource::Direct,
payload,
};
assert_eq!(
@@ -72,6 +72,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,
}
}
@@ -31,6 +31,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": [{
@@ -13,6 +13,7 @@ use crate::shell::Shell;
use crate::shell::ShellType;
use crate::shell_snapshot::ShellSnapshot;
use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::ShellCommandHandler;
@@ -230,6 +231,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 {
@@ -257,6 +259,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 {
@@ -287,6 +290,7 @@ async fn build_post_tool_use_payload_uses_tool_output_wire_value() {
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-42".to_string(),
tool_name: codex_tools::ToolName::plain("shell_command"),
source: ToolCallSource::Direct,
payload,
};
assert_eq!(
@@ -15,6 +15,7 @@ use tempfile::tempdir;
use crate::session::tests::make_session_and_context;
use crate::tools::context::ExecCommandToolOutput;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::hook_names::HookToolName;
@@ -35,6 +36,7 @@ async fn invocation_for_payload(
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: call_id.to_string(),
tool_name: codex_tools::ToolName::plain(tool_name),
source: ToolCallSource::Direct,
payload,
}
}
@@ -232,6 +234,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 {
@@ -257,6 +260,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
+1
View File
@@ -12,6 +12,7 @@ pub(crate) mod router;
pub(crate) mod runtimes;
pub(crate) mod sandboxing;
pub(crate) mod spec;
pub(crate) mod tool_dispatch_trace;
pub(crate) mod tool_search_entry;
use codex_protocol::exec_output::ExecToolCallOutput;
+36 -4
View File
@@ -15,6 +15,7 @@ use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use crate::tools::hook_names::HookToolName;
use crate::tools::tool_dispatch_trace::ToolDispatchTrace;
use codex_hooks::HookEvent;
use codex_hooks::HookEventAfterToolUse;
use codex_hooks::HookPayload;
@@ -219,6 +220,19 @@ impl ToolRegistry {
Self { handlers }
}
#[cfg(test)]
pub(crate) fn empty_for_test() -> Self {
Self::new(HashMap::new())
}
#[cfg(test)]
pub(crate) fn with_handler_for_test<T>(name: ToolName, handler: Arc<T>) -> Self
where
T: ToolHandler + 'static,
{
Self::new(HashMap::from([(name, handler as Arc<dyn AnyToolHandler>)]))
}
fn handler(&self, name: &ToolName) -> Option<Arc<dyn AnyToolHandler>> {
self.handlers.get(name).map(Arc::clone)
}
@@ -294,6 +308,8 @@ impl ToolRegistry {
}
}
let dispatch_trace = ToolDispatchTrace::start(&invocation);
let handler = match self.handler(&tool_name) {
Some(handler) => handler,
None => {
@@ -309,7 +325,9 @@ impl ToolRegistry {
mcp_server_ref,
mcp_server_origin_ref,
);
return Err(FunctionCallError::RespondToModel(message));
let err = FunctionCallError::RespondToModel(message);
dispatch_trace.record_failed(&err);
return Err(err);
}
};
@@ -326,7 +344,9 @@ impl ToolRegistry {
mcp_server_ref,
mcp_server_origin_ref,
);
return Err(FunctionCallError::Fatal(message));
let err = FunctionCallError::Fatal(message);
dispatch_trace.record_failed(&err);
return Err(err);
}
if let Some(pre_tool_use_payload) = handler.pre_tool_use_payload(&invocation)
@@ -339,7 +359,9 @@ impl ToolRegistry {
)
.await
{
return Err(FunctionCallError::RespondToModel(message));
let err = FunctionCallError::RespondToModel(message);
dispatch_trace.record_failed(&err);
return Err(err);
}
let is_mutating = handler.is_mutating(&invocation).await;
@@ -420,6 +442,7 @@ impl ToolRegistry {
.await;
if let Some(err) = hook_abort_error {
dispatch_trace.record_failed(&err);
return Err(err);
}
@@ -459,9 +482,18 @@ impl ToolRegistry {
let result = guard.take().ok_or_else(|| {
FunctionCallError::Fatal("tool produced no output".to_string())
})?;
dispatch_trace.record_completed(
&invocation,
&result.call_id,
&result.payload,
result.result.as_ref(),
);
Ok(result)
}
Err(err) => Err(err),
Err(err) => {
dispatch_trace.record_failed(&err);
Err(err)
}
}
}
}
+5 -1
View File
@@ -1,6 +1,7 @@
use super::*;
use pretty_assertions::assert_eq;
#[derive(Default)]
struct TestHandler;
impl ToolHandler for TestHandler {
@@ -11,7 +12,10 @@ impl ToolHandler for TestHandler {
}
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),
))
}
}
+2 -1
View File
@@ -281,7 +281,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
{
@@ -298,6 +298,7 @@ impl ToolRouter {
tracker,
call_id,
tool_name,
source,
payload,
};
@@ -53,8 +53,8 @@ use codex_shell_escalation::EscalationPolicy;
use codex_shell_escalation::EscalationSession;
use codex_shell_escalation::ExecParams;
use codex_shell_escalation::ExecResult;
use codex_shell_escalation::Permissions as EscalatedPermissions;
use codex_shell_escalation::PreparedExec;
use codex_shell_escalation::ResolvedPermissionProfile;
use codex_shell_escalation::ShellCommandExecutor;
use codex_shell_escalation::Stopwatch;
use codex_utils_absolute_path::AbsolutePathBuf;
@@ -370,13 +370,17 @@ impl CoreShellActionProvider {
.map(|_| {
// Shell request additional permissions were already normalized and
// merged into the first-attempt sandbox policy.
EscalationExecution::Permissions(EscalationPermissions::Permissions(
EscalatedPermissions {
sandbox_policy: sandbox_policy.clone(),
file_system_sandbox_policy: file_system_sandbox_policy.clone(),
network_sandbox_policy,
},
))
EscalationExecution::Permissions(
EscalationPermissions::ResolvedPermissionProfile(
ResolvedPermissionProfile {
permission_profile: PermissionProfile::from_runtime_permissions(
file_system_sandbox_policy,
network_sandbox_policy,
),
sandbox_policy: sandbox_policy.clone(),
},
),
)
})
.unwrap_or(EscalationExecution::TurnDefault),
}
@@ -842,9 +846,9 @@ impl ShellCommandExecutor for CoreShellCommandExecutor {
additional_permissions: None,
})?
}
EscalationExecution::Permissions(EscalationPermissions::PermissionProfile(
permission_profile,
)) => {
EscalationExecution::Permissions(
EscalationPermissions::AdditionalPermissionProfile(permission_profile),
) => {
// Merge additive permissions into the existing turn/request sandbox policy.
self.prepare_sandboxed_exec(PrepareSandboxedExecParams {
command,
@@ -856,15 +860,19 @@ impl ShellCommandExecutor for CoreShellCommandExecutor {
additional_permissions: Some(permission_profile),
})?
}
EscalationExecution::Permissions(EscalationPermissions::Permissions(permissions)) => {
// Use a fully specified sandbox policy instead of merging into the turn policy.
EscalationExecution::Permissions(EscalationPermissions::ResolvedPermissionProfile(
permissions,
)) => {
// Use a fully specified permission profile instead of merging into the turn policy.
let (file_system_sandbox_policy, network_sandbox_policy) =
permissions.permission_profile.to_runtime_permissions();
self.prepare_sandboxed_exec(PrepareSandboxedExecParams {
command,
workdir,
env,
sandbox_policy: &permissions.sandbox_policy,
file_system_sandbox_policy: &permissions.file_system_sandbox_policy,
network_sandbox_policy: permissions.network_sandbox_policy,
file_system_sandbox_policy: &file_system_sandbox_policy,
network_sandbox_policy,
additional_permissions: None,
})?
}
@@ -33,7 +33,7 @@ use codex_sandboxing::SandboxType;
use codex_shell_escalation::EscalationExecution;
use codex_shell_escalation::EscalationPermissions;
use codex_shell_escalation::ExecResult;
use codex_shell_escalation::Permissions as EscalatedPermissions;
use codex_shell_escalation::ResolvedPermissionProfile;
use codex_utils_absolute_path::AbsolutePathBuf;
use pretty_assertions::assert_eq;
use serde_json::Value;
@@ -317,11 +317,13 @@ fn shell_request_escalation_execution_is_explicit() {
network_sandbox_policy,
Some(&requested_permissions),
),
EscalationExecution::Permissions(EscalationPermissions::Permissions(
EscalatedPermissions {
EscalationExecution::Permissions(EscalationPermissions::ResolvedPermissionProfile(
ResolvedPermissionProfile {
permission_profile: PermissionProfile::from_runtime_permissions(
&file_system_sandbox_policy,
network_sandbox_policy,
),
sandbox_policy,
file_system_sandbox_policy,
network_sandbox_policy,
},
)),
);
@@ -0,0 +1,139 @@
//! Adapter between core tool dispatch objects and rollout-trace events.
//!
//! `codex-rollout-trace` owns the event schema and writer behavior. This module
//! keeps the core-specific mapping from registry invocations/results out of the
//! registry control flow.
use crate::function_tool::FunctionCallError;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::ToolDispatchInvocation;
use codex_rollout_trace::ToolDispatchPayload;
use codex_rollout_trace::ToolDispatchRequester;
use codex_rollout_trace::ToolDispatchResult;
use codex_rollout_trace::ToolDispatchTraceContext;
/// Keeps registry early-return paths paired with trace end events.
pub(crate) struct ToolDispatchTrace {
context: ToolDispatchTraceContext,
}
impl ToolDispatchTrace {
pub(crate) fn start(invocation: &ToolInvocation) -> Self {
let context = invocation
.session
.services
.rollout_trace
.start_tool_dispatch_trace(|| tool_dispatch_invocation(invocation));
Self { context }
}
pub(crate) fn record_completed(
&self,
invocation: &ToolInvocation,
call_id: &str,
payload: &ToolPayload,
result: &dyn ToolOutput,
) {
if !self.context.is_enabled() {
return;
}
let Some(result_payload) = tool_dispatch_result(invocation, call_id, payload, result)
else {
return;
};
let status = if result.success_for_logging() {
ExecutionStatus::Completed
} else {
ExecutionStatus::Failed
};
self.context.record_completed(status, result_payload);
}
pub(crate) fn record_failed(&self, error: &FunctionCallError) {
self.context.record_failed(error);
}
}
fn tool_dispatch_invocation(invocation: &ToolInvocation) -> Option<ToolDispatchInvocation> {
let requester = match &invocation.source {
ToolCallSource::Direct => ToolDispatchRequester::Model {
model_visible_call_id: invocation.call_id.clone(),
},
ToolCallSource::CodeMode {
cell_id,
runtime_tool_call_id,
} => ToolDispatchRequester::CodeCell {
runtime_cell_id: cell_id.clone(),
runtime_tool_call_id: runtime_tool_call_id.clone(),
},
ToolCallSource::JsRepl => return None,
};
Some(ToolDispatchInvocation {
thread_id: invocation.session.conversation_id.to_string(),
codex_turn_id: invocation.turn.sub_id.clone(),
tool_call_id: invocation.call_id.clone(),
tool_name: invocation.tool_name.name.clone(),
tool_namespace: invocation.tool_name.namespace.clone(),
requester,
payload: tool_dispatch_payload(&invocation.payload),
})
}
fn tool_dispatch_result(
invocation: &ToolInvocation,
call_id: &str,
payload: &ToolPayload,
result: &dyn ToolOutput,
) -> Option<ToolDispatchResult> {
match invocation.source {
ToolCallSource::Direct => Some(ToolDispatchResult::DirectResponse {
response_item: result.to_response_item(call_id, payload),
}),
ToolCallSource::CodeMode { .. } => Some(ToolDispatchResult::CodeModeResponse {
value: result.code_mode_result(payload),
}),
ToolCallSource::JsRepl => None,
}
}
fn tool_dispatch_payload(payload: &ToolPayload) -> ToolDispatchPayload {
match payload {
ToolPayload::Function { arguments } => ToolDispatchPayload::Function {
arguments: arguments.clone(),
},
ToolPayload::ToolSearch { arguments } => ToolDispatchPayload::ToolSearch {
arguments: arguments.clone(),
},
ToolPayload::Custom { input } => ToolDispatchPayload::Custom {
input: input.clone(),
},
ToolPayload::LocalShell { params } => ToolDispatchPayload::LocalShell {
command: params.command.clone(),
workdir: params.workdir.clone(),
timeout_ms: params.timeout_ms,
sandbox_permissions: params.sandbox_permissions,
prefix_rule: params.prefix_rule.clone(),
additional_permissions: params.additional_permissions.clone(),
justification: params.justification.clone(),
},
ToolPayload::Mcp {
server,
tool,
raw_arguments,
} => ToolDispatchPayload::Mcp {
server: server.clone(),
tool: tool.clone(),
raw_arguments: raw_arguments.clone(),
},
}
}
#[cfg(test)]
#[path = "tool_dispatch_trace_tests.rs"]
mod tests;
@@ -0,0 +1,337 @@
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use codex_protocol::protocol::SessionSource;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::RolloutTraceRecorder;
use codex_rollout_trace::ThreadStartedTraceMetadata;
use codex_rollout_trace::ToolCallRequester;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
use tokio_util::sync::CancellationToken;
use crate::function_tool::FunctionCallError;
use crate::session::session::Session;
use crate::session::tests::make_session_and_context;
use crate::session::turn_context::TurnContext;
use crate::tools::code_mode::CodeModeWaitHandler;
use crate::tools::code_mode::WAIT_TOOL_NAME;
use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::registry::ToolHandler;
use crate::tools::registry::ToolKind;
use crate::tools::registry::ToolRegistry;
use crate::turn_diff_tracker::TurnDiffTracker;
#[derive(Default)]
struct TestHandler;
impl ToolHandler for TestHandler {
type Output = FunctionToolOutput;
fn kind(&self) -> ToolKind {
ToolKind::Function
}
async fn handle(&self, _invocation: ToolInvocation) -> Result<Self::Output, FunctionCallError> {
Ok(FunctionToolOutput::from_text("ok".to_string(), Some(true)))
}
}
#[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.start_code_cell_trace(
session.conversation_id,
turn.sub_id.as_str(),
"cell-1",
"call-code",
"await tools.test_tool({})",
);
let registry = ToolRegistry::with_handler_for_test(
codex_tools::ToolName::plain("test_tool"),
Arc::new(TestHandler),
);
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(ToolCallSource::JsRepl).await
}
#[tokio::test]
async fn dispatch_lifecycle_trace_records_unsupported_tool_failures() -> 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::empty_for_test();
let session = Arc::new(session);
let turn = Arc::new(turn);
let result = registry
.dispatch_any(test_invocation(
session,
turn,
"unsupported-call",
"missing_tool",
ToolCallSource::Direct,
"{}",
))
.await;
assert!(matches!(result, Err(FunctionCallError::RespondToModel(_))));
let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?;
let tool_call = &replayed.tool_calls["unsupported-call"];
assert_eq!(tool_call.execution.status, ExecutionStatus::Failed);
assert!(tool_call.raw_result_payload_id.is_some());
Ok(())
}
#[tokio::test]
async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> 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::with_handler_for_test(
codex_tools::ToolName::plain("test_tool"),
Arc::new(TestHandler),
);
let session = Arc::new(session);
let turn = Arc::new(turn);
let result = registry
.dispatch_any(test_invocation_with_payload(
session,
turn,
"incompatible-call",
codex_tools::ToolName::plain("test_tool"),
ToolCallSource::Direct,
ToolPayload::Custom {
input: "{}".to_string(),
},
))
.await;
assert!(matches!(result, Err(FunctionCallError::Fatal(_))));
let replayed = codex_rollout_trace::replay_bundle(single_bundle_dir(temp.path())?)?;
let tool_call = &replayed.tool_calls["incompatible-call"];
assert_eq!(tool_call.execution.status, ExecutionStatus::Failed);
assert!(tool_call.raw_result_payload_id.is_some());
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::with_handler_for_test(
codex_tools::ToolName::plain(WAIT_TOOL_NAME),
Arc::new(CodeModeWaitHandler),
);
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(())
}
async fn assert_dispatch_trace_skips(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::with_handler_for_test(
codex_tools::ToolName::plain("test_tool"),
Arc::new(TestHandler),
);
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(())
}
fn test_invocation(
session: Arc<Session>,
turn: Arc<TurnContext>,
call_id: &str,
tool_name: &str,
source: ToolCallSource,
arguments: &str,
) -> ToolInvocation {
test_invocation_with_payload(
session,
turn,
call_id,
codex_tools::ToolName::plain(tool_name),
source,
ToolPayload::Function {
arguments: arguments.to_string(),
},
)
}
fn test_invocation_with_payload(
session: Arc<Session>,
turn: Arc<TurnContext>,
call_id: &str,
tool_name: codex_tools::ToolName,
source: ToolCallSource,
payload: ToolPayload,
) -> ToolInvocation {
ToolInvocation {
session,
turn,
cancellation_token: CancellationToken::new(),
tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
call_id: call_id.to_string(),
tool_name,
source,
payload,
}
}
fn attach_test_trace(session: &mut Session, turn: &TurnContext, root: &Path) -> anyhow::Result<()> {
let thread_id = session.conversation_id;
let recorder = RolloutTraceRecorder::create_in_root_for_test(root, thread_id)?;
recorder.record_thread_started(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_started(thread_id, turn.sub_id.as_str());
session.services.rollout_trace = 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))
}