codex/codex-rs/rollout-trace/src/tool_dispatch.rs

//! Hot-path helpers for recording canonical tool dispatch boundaries.
//!
//! Core owns tool routing and result conversion. The trace crate owns the raw
//! event schema, payload shape, and no-op behavior, so core only adapts its
//! domain objects into the small request/result structs defined here.

use std::fmt::Display;
use std::sync::Arc;

use codex_protocol::models::AdditionalPermissionProfile;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::SandboxPermissions;
use codex_protocol::models::SearchToolCallParams;
use serde::Serialize;
use serde_json::Value as JsonValue;
use serde_json::json;
use tracing::warn;

use crate::model::AgentThreadId;
use crate::model::CodeModeRuntimeToolId;
use crate::model::CodexTurnId;
use crate::model::ExecutionStatus;
use crate::model::ModelVisibleCallId;
use crate::model::ToolCallId;
use crate::model::ToolCallKind;
use crate::model::ToolCallSummary;
use crate::payload::RawPayloadKind;
use crate::payload::RawPayloadRef;
use crate::raw_event::RawToolCallRequester;
use crate::raw_event::RawTraceEventContext;
use crate::raw_event::RawTraceEventPayload;
use crate::writer::TraceWriter;

/// No-op capable trace handle for one resolved tool dispatch.
#[derive(Clone, Debug)]
pub struct ToolDispatchTraceContext {
    state: ToolDispatchTraceContextState,
}

#[derive(Clone, Debug)]
enum ToolDispatchTraceContextState {
    Disabled,
    Enabled(EnabledToolDispatchTraceContext),
}

#[derive(Clone, Debug)]
struct EnabledToolDispatchTraceContext {
    writer: Arc<TraceWriter>,
    thread_id: AgentThreadId,
    codex_turn_id: CodexTurnId,
    tool_call_id: ToolCallId,
}

/// Core-facing request data for the canonical Codex tool boundary.
pub struct ToolDispatchInvocation {
    pub thread_id: AgentThreadId,
    pub codex_turn_id: CodexTurnId,
    pub tool_call_id: ToolCallId,
    pub tool_name: String,
    pub tool_namespace: Option<String>,
    pub requester: ToolDispatchRequester,
    pub payload: ToolDispatchPayload,
}

/// Runtime source that caused a dispatch-level tool call.
pub enum ToolDispatchRequester {
    Model {
        model_visible_call_id: ModelVisibleCallId,
    },
    CodeCell {
        runtime_cell_id: String,
        runtime_tool_call_id: CodeModeRuntimeToolId,
    },
}

/// Tool input observed at the registry boundary.
pub enum ToolDispatchPayload {
    Function {
        arguments: String,
    },
    ToolSearch {
        arguments: SearchToolCallParams,
    },
    Custom {
        input: String,
    },
    LocalShell {
        command: Vec<String>,
        workdir: Option<String>,
        timeout_ms: Option<u64>,
        sandbox_permissions: Option<SandboxPermissions>,
        prefix_rule: Option<Vec<String>>,
        additional_permissions: Option<AdditionalPermissionProfile>,
        justification: Option<String>,
    },
}

/// Result data returned from a dispatch-level tool call.
#[derive(Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum ToolDispatchResult {
    DirectResponse { response_item: ResponseInputItem },
    CodeModeResponse { value: JsonValue },
}

/// Raw invocation payload for the canonical Codex tool boundary.
#[derive(Serialize)]
struct DispatchedToolTraceRequest<'a> {
    tool_name: &'a str,
    tool_namespace: Option<&'a str>,
    payload: &'a JsonValue,
}

/// Raw response payload for dispatch-level tool trace events.
#[derive(Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
enum DispatchedToolTraceResponse<'a> {
    DirectResponse {
        response_item: &'a ResponseInputItem,
    },
    CodeModeResponse {
        value: &'a JsonValue,
    },
    Error {
        error: String,
    },
}

impl ToolDispatchTraceContext {
    /// Builds a context that accepts trace calls and records nothing.
    pub(crate) fn disabled() -> Self {
        Self {
            state: ToolDispatchTraceContextState::Disabled,
        }
    }

    /// Returns whether caller-side result conversion would be recorded.
    ///
    /// Core uses this to avoid formatting or cloning tool outputs when the
    /// dispatch lifecycle is suppressed or tracing is disabled.
    pub fn is_enabled(&self) -> bool {
        matches!(self.state, ToolDispatchTraceContextState::Enabled(_))
    }

    /// Starts one dispatch-level lifecycle and returns the handle for its result.
    pub(crate) fn start(writer: Arc<TraceWriter>, invocation: ToolDispatchInvocation) -> Self {
        if suppresses_tool_dispatch_trace(&invocation) {
            return Self::disabled();
        }

        let context = EnabledToolDispatchTraceContext {
            writer,
            thread_id: invocation.thread_id.clone(),
            codex_turn_id: invocation.codex_turn_id.clone(),
            tool_call_id: invocation.tool_call_id.clone(),
        };
        record_started(&context, invocation);
        Self {
            state: ToolDispatchTraceContextState::Enabled(context),
        }
    }

    /// Records the caller-facing successful or failed tool result.
    pub fn record_completed(&self, status: ExecutionStatus, result: ToolDispatchResult) {
        let ToolDispatchTraceContextState::Enabled(context) = &self.state else {
            return;
        };
        let response = match &result {
            ToolDispatchResult::DirectResponse { response_item } => {
                DispatchedToolTraceResponse::DirectResponse { response_item }
            }
            ToolDispatchResult::CodeModeResponse { value } => {
                DispatchedToolTraceResponse::CodeModeResponse { value }
            }
        };
        append_tool_call_ended(context, status, &response);
    }

    /// Records a dispatch failure before the tool produced a normal result payload.
    pub fn record_failed(&self, error: impl Display) {
        let ToolDispatchTraceContextState::Enabled(context) = &self.state else {
            return;
        };
        append_tool_call_ended(
            context,
            ExecutionStatus::Failed,
            &DispatchedToolTraceResponse::Error {
                error: error.to_string(),
            },
        );
    }
}

fn suppresses_tool_dispatch_trace(invocation: &ToolDispatchInvocation) -> bool {
    matches!(invocation.payload, ToolDispatchPayload::Custom { .. })
        && invocation.tool_namespace.is_none()
        && invocation.tool_name == codex_code_mode::PUBLIC_TOOL_NAME
}

fn record_started(context: &EnabledToolDispatchTraceContext, invocation: ToolDispatchInvocation) {
    let tool_name = invocation.tool_name;
    let tool_namespace = invocation.tool_namespace;
    let kind = dispatched_tool_kind(&tool_name, &invocation.payload);
    let label = dispatched_tool_label(&tool_name, tool_namespace.as_deref(), &invocation.payload);
    let input_preview = Some(invocation.payload.log_payload_preview());
    let payload = invocation.payload.into_json_payload();
    let request = DispatchedToolTraceRequest {
        tool_name: tool_name.as_str(),
        tool_namespace: tool_namespace.as_deref(),
        payload: &payload,
    };
    let request_payload =
        write_json_payload_best_effort(&context.writer, RawPayloadKind::ToolInvocation, &request);
    let (model_visible_call_id, code_mode_runtime_tool_id, requester) =
        requester_fields(invocation.requester);

    append_with_context_best_effort(
        context,
        RawTraceEventPayload::ToolCallStarted {
            tool_call_id: context.tool_call_id.clone(),
            model_visible_call_id,
            code_mode_runtime_tool_id,
            requester,
            kind,
            summary: ToolCallSummary::Generic {
                label,
                input_preview,
                output_preview: None,
            },
            invocation_payload: request_payload,
        },
    );
}

fn requester_fields(
    requester: ToolDispatchRequester,
) -> (
    Option<ModelVisibleCallId>,
    Option<CodeModeRuntimeToolId>,
    RawToolCallRequester,
) {
    match requester {
        ToolDispatchRequester::Model {
            model_visible_call_id,
        } => (
            Some(model_visible_call_id),
            None,
            RawToolCallRequester::Model,
        ),
        ToolDispatchRequester::CodeCell {
            runtime_cell_id,
            runtime_tool_call_id,
        } => (
            None,
            Some(runtime_tool_call_id),
            RawToolCallRequester::CodeCell { runtime_cell_id },
        ),
    }
}

fn dispatched_tool_kind(tool_name: &str, _payload: &ToolDispatchPayload) -> ToolCallKind {
    match tool_name {
        "exec_command" | "local_shell" | "shell" | "shell_command" => ToolCallKind::ExecCommand,
        "write_stdin" => ToolCallKind::WriteStdin,
        "apply_patch" => ToolCallKind::ApplyPatch,
        "web_search" | "web_search_preview" => ToolCallKind::Web,
        "image_generation" | "image_query" => ToolCallKind::ImageGeneration,
        "spawn_agent" => ToolCallKind::SpawnAgent,
        "send_message" => ToolCallKind::SendMessage,
        "followup_task" | "assign_task" => ToolCallKind::AssignAgentTask,
        "wait_agent" => ToolCallKind::WaitAgent,
        "close_agent" | "interrupt_agent" => ToolCallKind::CloseAgent,
        other => ToolCallKind::Other {
            name: other.to_string(),
        },
    }
}

fn dispatched_tool_label(
    tool_name: &str,
    tool_namespace: Option<&str>,
    _payload: &ToolDispatchPayload,
) -> String {
    match tool_namespace {
        Some(namespace) => format!("{namespace}.{tool_name}"),
        None => tool_name.to_string(),
    }
}

impl ToolDispatchPayload {
    fn log_payload_preview(&self) -> String {
        match self {
            ToolDispatchPayload::Function { arguments } => truncate_preview(arguments),
            ToolDispatchPayload::ToolSearch { arguments } => truncate_preview(&arguments.query),
            ToolDispatchPayload::Custom { input } => truncate_preview(input),
            ToolDispatchPayload::LocalShell { command, .. } => truncate_preview(&command.join(" ")),
        }
    }

    fn into_json_payload(self) -> JsonValue {
        match self {
            ToolDispatchPayload::Function { arguments } => json!({
                "type": "function",
                "arguments": arguments,
            }),
            ToolDispatchPayload::ToolSearch { arguments } => json!({
                "type": "tool_search",
                "arguments": arguments,
            }),
            ToolDispatchPayload::Custom { input } => json!({
                "type": "custom",
                "input": input,
            }),
            ToolDispatchPayload::LocalShell {
                command,
                workdir,
                timeout_ms,
                sandbox_permissions,
                prefix_rule,
                additional_permissions,
                justification,
            } => json!({
                "type": "local_shell",
                "command": command,
                "workdir": workdir,
                "timeout_ms": timeout_ms,
                "sandbox_permissions": sandbox_permissions,
                "prefix_rule": prefix_rule,
                "additional_permissions": additional_permissions,
                "justification": justification,
            }),
        }
    }
}

fn truncate_preview(value: &str) -> String {
    const MAX_PREVIEW_CHARS: usize = 160;
    let mut chars = value.chars();
    let mut preview = chars.by_ref().take(MAX_PREVIEW_CHARS).collect::<String>();
    if chars.next().is_some() {
        preview.push_str("...");
    }
    preview
}

fn append_tool_call_ended(
    context: &EnabledToolDispatchTraceContext,
    status: ExecutionStatus,
    response: &DispatchedToolTraceResponse<'_>,
) {
    let response_payload =
        write_json_payload_best_effort(&context.writer, RawPayloadKind::ToolResult, response);
    append_with_context_best_effort(
        context,
        RawTraceEventPayload::ToolCallEnded {
            tool_call_id: context.tool_call_id.clone(),
            status,
            result_payload: response_payload,
        },
    );
}

fn write_json_payload_best_effort(
    writer: &TraceWriter,
    kind: RawPayloadKind,
    payload: &impl Serialize,
) -> Option<RawPayloadRef> {
    match writer.write_json_payload(kind, payload) {
        Ok(payload_ref) => Some(payload_ref),
        Err(err) => {
            warn!("failed to write rollout trace payload: {err:#}");
            None
        }
    }
}

fn append_with_context_best_effort(
    context: &EnabledToolDispatchTraceContext,
    payload: RawTraceEventPayload,
) {
    let event_context = RawTraceEventContext {
        thread_id: Some(context.thread_id.clone()),
        codex_turn_id: Some(context.codex_turn_id.clone()),
    };
    if let Err(err) = context.writer.append_with_context(event_context, payload) {
        warn!("failed to append rollout trace event: {err:#}");
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn suppresses_only_noncanonical_dispatch_boundaries() {
        assert!(suppresses_tool_dispatch_trace(&invocation(
            codex_code_mode::PUBLIC_TOOL_NAME,
            /*tool_namespace*/ None,
            ToolDispatchRequester::Model {
                model_visible_call_id: "call-exec".to_string(),
            },
            ToolDispatchPayload::Custom {
                input: "1 + 1".to_string(),
            },
        )));
        assert!(!suppresses_tool_dispatch_trace(&invocation(
            "custom_tool",
            /*tool_namespace*/ None,
            ToolDispatchRequester::Model {
                model_visible_call_id: "call-custom".to_string(),
            },
            ToolDispatchPayload::Custom {
                input: "payload".to_string(),
            },
        )));
        assert!(!suppresses_tool_dispatch_trace(&invocation(
            codex_code_mode::PUBLIC_TOOL_NAME,
            Some("mcp__server".to_string()),
            ToolDispatchRequester::Model {
                model_visible_call_id: "call-namespaced".to_string(),
            },
            ToolDispatchPayload::Custom {
                input: "payload".to_string(),
            },
        )));
    }

    #[test]
    fn classifies_interrupt_agent_as_close_agent() {
        assert_eq!(
            dispatched_tool_kind(
                "interrupt_agent",
                &ToolDispatchPayload::Function {
                    arguments: r#"{"target":"/root/child"}"#.to_string(),
                },
            ),
            ToolCallKind::CloseAgent
        );
    }

    fn invocation(
        tool_name: &str,
        tool_namespace: Option<String>,
        requester: ToolDispatchRequester,
        payload: ToolDispatchPayload,
    ) -> ToolDispatchInvocation {
        ToolDispatchInvocation {
            thread_id: "thread-1".to_string(),
            codex_turn_id: "turn-1".to_string(),
            tool_call_id: "tool-call-1".to_string(),
            tool_name: tool_name.to_string(),
            tool_namespace,
            requester,
            payload,
        }
    }
}