codex/codex-rs/rollout-trace/src/reducer/tool.rs

use anyhow::Context;
use anyhow::Result;
use anyhow::bail;

use super::TraceReducer;
use crate::model::CodeModeRuntimeToolId;
use crate::model::ConversationItemKind;
use crate::model::ExecutionStatus;
use crate::model::ExecutionWindow;
use crate::model::McpCallId;
use crate::model::ModelVisibleCallId;
use crate::model::ProducerRef;
use crate::model::ToolCall;
use crate::model::ToolCallId;
use crate::model::ToolCallKind;
use crate::model::ToolCallSummary;
use crate::payload::RawPayloadRef;
use crate::raw_event::RawEventSeq;
use crate::raw_event::RawToolCallRequester;

mod agents;
mod terminal;

pub(super) use agents::ObservedAgentResultEdge;
pub(super) use agents::PendingAgentInteractionEdge;
pub(super) use agents::spawn_edge_id;

/// Raw tool-start fields after dispatch has stripped the common event envelope.
///
/// Tool starts carry several optional identity namespaces: model-visible calls,
/// code-mode runtime tools, and canonical invocation payloads. Grouping them keeps
/// the reducer callsite readable and avoids positional argument mistakes.
pub(super) struct ToolCallStarted {
    pub(super) tool_call_id: ToolCallId,
    pub(super) model_visible_call_id: Option<ModelVisibleCallId>,
    pub(super) code_mode_runtime_tool_id: Option<CodeModeRuntimeToolId>,
    pub(super) requester: RawToolCallRequester,
    pub(super) kind: ToolCallKind,
    pub(super) summary: ToolCallSummary,
    pub(super) invocation_payload: Option<RawPayloadRef>,
}

impl TraceReducer {
    /// Starts a tool call and links it to model-visible items or runtime parents when available.
    ///
    /// Some tools also create richer domain objects, such as terminal operations, from
    /// the same invocation payload. The generic ToolCall remains the common index.
    pub(super) fn start_tool_call(
        &mut self,
        seq: RawEventSeq,
        wall_time_unix_ms: i64,
        thread_id: Option<String>,
        codex_turn_id: Option<String>,
        started: ToolCallStarted,
    ) -> Result<()> {
        let tool_call_id = started.tool_call_id.clone();
        if self.rollout.tool_calls.contains_key(&tool_call_id) {
            bail!("duplicate tool call start for {tool_call_id}");
        }
        self.ensure_unique_model_visible_tool_call(
            started.model_visible_call_id.as_deref(),
            &tool_call_id,
        )?;

        let thread_id = self.tool_thread_id(thread_id, codex_turn_id.as_deref())?;
        self.validate_tool_turn(&thread_id, codex_turn_id.as_deref())?;

        let model_visible_call_id = started.model_visible_call_id.clone();
        let requester = self.reduce_tool_call_requester(&thread_id, started.requester.clone())?;
        let model_visible_call_item_ids = model_visible_call_id
            .as_deref()
            .map(|call_id| {
                self.model_visible_tool_item_ids(
                    &thread_id,
                    call_id,
                    &[
                        ConversationItemKind::FunctionCall,
                        ConversationItemKind::CustomToolCall,
                    ],
                )
            })
            .unwrap_or_default();
        let model_visible_output_item_ids = model_visible_call_id
            .as_deref()
            .map(|call_id| {
                self.model_visible_tool_item_ids(
                    &thread_id,
                    call_id,
                    &[
                        ConversationItemKind::FunctionCallOutput,
                        ConversationItemKind::CustomToolCallOutput,
                    ],
                )
            })
            .unwrap_or_default();

        self.thread_mut(&thread_id)?;

        // Some terminal-like tools, notably write_stdin, do not emit a richer
        // runtime begin event. For those tools the canonical invocation is the
        // only place to recover the terminal/session join key.
        let terminal_operation_id = self.start_terminal_operation_from_invocation(
            seq,
            wall_time_unix_ms,
            &thread_id,
            &tool_call_id,
            &started.kind,
            started.invocation_payload.as_ref(),
        )?;
        // Terminal-backed tools should render through the richer terminal
        // operation instead of the generic tool summary captured by producers.
        let summary = terminal_operation_id
            .as_ref()
            .map(|operation_id| ToolCallSummary::Terminal {
                operation_id: operation_id.clone(),
            })
            .unwrap_or(started.summary);
        let raw_invocation_payload_id = started
            .invocation_payload
            .as_ref()
            .map(|payload| payload.raw_payload_id.clone());
        self.link_wait_tool_call_from_request_payload(
            &thread_id,
            &tool_call_id,
            started.invocation_payload.as_ref(),
        )?;

        self.rollout.tool_calls.insert(
            tool_call_id.clone(),
            ToolCall {
                tool_call_id: tool_call_id.clone(),
                mcp_call_id: None,
                model_visible_call_id,
                code_mode_runtime_tool_id: started.code_mode_runtime_tool_id,
                thread_id,
                started_by_codex_turn_id: codex_turn_id,
                execution: ExecutionWindow {
                    started_at_unix_ms: wall_time_unix_ms,
                    started_seq: seq,
                    ended_at_unix_ms: None,
                    ended_seq: None,
                    status: ExecutionStatus::Running,
                },
                requester: requester.clone(),
                kind: started.kind,
                model_visible_call_item_ids,
                model_visible_output_item_ids: Vec::new(),
                terminal_operation_id,
                summary,
                raw_invocation_payload_id,
                raw_result_payload_id: None,
                raw_runtime_payload_ids: Vec::new(),
            },
        );

        self.link_tool_call_to_code_cell(&tool_call_id, &requester)?;
        self.link_tool_to_inference_response(&tool_call_id);
        // Output items need the reverse ProducerRef edge as well, so attach
        // them after insertion through the same helper used by the transcript
        // reducer when the output is observed after the tool start.
        for item_id in model_visible_output_item_ids {
            self.add_tool_output_item(&tool_call_id, &item_id)?;
        }
        // The call/output items may have been observed before this tool start.
        // Re-sync after insertion so terminal observations get both directions
        // of the model-visible link.
        self.sync_terminal_model_observation(&tool_call_id)?;
        Ok(())
    }

    /// Attaches the bridge-visible MCP UUID after the generic tool call exists.
    pub(super) fn assign_mcp_tool_call_correlation(
        &mut self,
        tool_call_id: ToolCallId,
        mcp_call_id: McpCallId,
    ) -> Result<()> {
        let Some(tool_call) = self.rollout.tool_calls.get_mut(&tool_call_id) else {
            bail!("MCP correlation referenced unknown tool call {tool_call_id}");
        };
        if tool_call.mcp_call_id.replace(mcp_call_id).is_some() {
            bail!("duplicate MCP correlation for tool call {tool_call_id}");
        }
        Ok(())
    }

    /// Completes the canonical tool call and any terminal operation driven by dispatch output.
    ///
    /// Protocol-backed terminal tools end from runtime events; direct tools
    /// may only have the canonical result payload, so this method handles both paths.
    pub(super) fn end_tool_call(
        &mut self,
        seq: RawEventSeq,
        wall_time_unix_ms: i64,
        tool_call_id: ToolCallId,
        status: ExecutionStatus,
        result_payload: Option<RawPayloadRef>,
    ) -> Result<()> {
        let (terminal_operation_id, thread_id, end_terminal_from_result) = {
            let Some(tool_call) = self.rollout.tool_calls.get_mut(&tool_call_id) else {
                bail!("tool call end referenced unknown call {tool_call_id}");
            };
            tool_call.execution.ended_at_unix_ms = Some(wall_time_unix_ms);
            tool_call.execution.ended_seq = Some(seq);
            tool_call.execution.status = status.clone();
            tool_call.raw_result_payload_id = result_payload
                .as_ref()
                .map(|payload| payload.raw_payload_id.clone());
            (
                tool_call.terminal_operation_id.clone(),
                tool_call.thread_id.clone(),
                // Protocol-backed tools end terminal operations from
                // runtime observations. Dispatch result payloads are still kept
                // on ToolCall, but they are caller-facing and may be transformed
                // relative to the raw terminal output.
                tool_call.raw_runtime_payload_ids.is_empty(),
            )
        };
        if end_terminal_from_result && let Some(operation_id) = terminal_operation_id {
            self.end_terminal_operation(
                seq,
                wall_time_unix_ms,
                &thread_id,
                &operation_id,
                status,
                result_payload.as_ref(),
            )?;
        }
        self.attach_agent_interaction_tool_result(&tool_call_id, result_payload.as_ref())?;
        Ok(())
    }

    /// Records a runtime-begin observation for an already started tool call.
    ///
    /// Runtime observations enrich the generic tool with protocol facts and may
    /// create domain-specific children such as terminal operations or agent edges.
    pub(super) fn start_tool_runtime_observation(
        &mut self,
        seq: RawEventSeq,
        wall_time_unix_ms: i64,
        tool_call_id: ToolCallId,
        runtime_payload: RawPayloadRef,
    ) -> Result<()> {
        let (thread_id, _requester, kind, existing_terminal_operation_id) = {
            let Some(tool_call) = self.rollout.tool_calls.get_mut(&tool_call_id) else {
                bail!("tool runtime start referenced unknown call {tool_call_id}");
            };
            push_unique(
                &mut tool_call.raw_runtime_payload_ids,
                &runtime_payload.raw_payload_id,
            );
            (
                tool_call.thread_id.clone(),
                tool_call.requester.clone(),
                tool_call.kind.clone(),
                tool_call.terminal_operation_id.clone(),
            )
        };
        if existing_terminal_operation_id.is_some()
            && matches!(kind, ToolCallKind::ExecCommand | ToolCallKind::WriteStdin)
        {
            bail!("tool runtime start would create a second terminal operation for {tool_call_id}");
        }

        // Protocol begin events carry runtime facts such as process ids and
        // cwd. These facts should create terminal rows, but they must not
        // replace the canonical invocation payload captured at dispatch.
        let terminal_operation_id = self.start_terminal_operation_from_runtime(
            seq,
            wall_time_unix_ms,
            &thread_id,
            &tool_call_id,
            &kind,
            &runtime_payload,
        )?;

        if let Some(operation_id) = &terminal_operation_id {
            let Some(tool_call) = self.rollout.tool_calls.get_mut(&tool_call_id) else {
                bail!("tool call {tool_call_id} disappeared during runtime start reduction");
            };
            if tool_call.terminal_operation_id.is_none() {
                tool_call.terminal_operation_id = Some(operation_id.clone());
                tool_call.summary = ToolCallSummary::Terminal {
                    operation_id: operation_id.clone(),
                };
            }
        }

        if terminal_operation_id.is_some() {
            self.sync_terminal_model_observation(&tool_call_id)?;
        }
        self.start_agent_interaction_from_runtime(&tool_call_id, &runtime_payload)?;
        Ok(())
    }

    /// Records a runtime-end observation for an already started tool call.
    pub(super) fn end_tool_runtime_observation(
        &mut self,
        seq: RawEventSeq,
        wall_time_unix_ms: i64,
        tool_call_id: ToolCallId,
        status: ExecutionStatus,
        runtime_payload: RawPayloadRef,
    ) -> Result<()> {
        let (thread_id, terminal_operation_id) = {
            let Some(tool_call) = self.rollout.tool_calls.get_mut(&tool_call_id) else {
                bail!("tool runtime end referenced unknown call {tool_call_id}");
            };
            push_unique(
                &mut tool_call.raw_runtime_payload_ids,
                &runtime_payload.raw_payload_id,
            );
            (
                tool_call.thread_id.clone(),
                tool_call.terminal_operation_id.clone(),
            )
        };

        if let Some(operation_id) = terminal_operation_id {
            self.end_terminal_operation(
                seq,
                wall_time_unix_ms,
                &thread_id,
                &operation_id,
                status,
                Some(&runtime_payload),
            )?;
        }
        self.end_agent_interaction_from_runtime(
            wall_time_unix_ms,
            &tool_call_id,
            &runtime_payload,
        )?;
        Ok(())
    }

    /// Attaches a conversation item observed after the tool call was reduced.
    ///
    /// Inference request/response ordering can expose call/output items after the
    /// runtime tool object exists, so transcript reduction calls back here to add
    /// reverse links without duplicating matching logic.
    pub(super) fn attach_model_visible_tool_item(
        &mut self,
        item_id: &str,
        call_id: Option<&str>,
        kind: &ConversationItemKind,
    ) -> Result<()> {
        let Some(call_id) = call_id else {
            return Ok(());
        };
        match kind {
            ConversationItemKind::FunctionCall | ConversationItemKind::CustomToolCall => {
                if let Some(tool_call_id) = self.single_tool_for_model_visible_call(call_id)? {
                    self.add_tool_call_item(&tool_call_id, item_id)?;
                    self.link_tool_to_inference_response(&tool_call_id);
                    self.sync_terminal_model_observation(&tool_call_id)?;
                }
            }
            ConversationItemKind::FunctionCallOutput
            | ConversationItemKind::CustomToolCallOutput => {
                if let Some(tool_call_id) = self.single_tool_for_model_visible_call(call_id)? {
                    self.add_tool_output_item(&tool_call_id, item_id)?;
                    self.sync_terminal_model_observation(&tool_call_id)?;
                }
            }
            ConversationItemKind::Message
            | ConversationItemKind::Reasoning
            | ConversationItemKind::CompactionMarker => {}
        }
        Ok(())
    }

    fn tool_thread_id(
        &self,
        thread_id: Option<String>,
        codex_turn_id: Option<&str>,
    ) -> Result<String> {
        if let Some(thread_id) = thread_id {
            return Ok(thread_id);
        }
        let Some(codex_turn_id) = codex_turn_id else {
            bail!("tool call start did not include thread or Codex turn context");
        };
        self.rollout
            .codex_turns
            .get(codex_turn_id)
            .map(|turn| turn.thread_id.clone())
            .with_context(|| {
                format!("tool call start referenced unknown Codex turn {codex_turn_id}")
            })
    }

    fn validate_tool_turn(&self, thread_id: &str, codex_turn_id: Option<&str>) -> Result<()> {
        if !self.rollout.threads.contains_key(thread_id) {
            bail!("tool call start referenced unknown thread {thread_id}");
        }
        if let Some(codex_turn_id) = codex_turn_id {
            let Some(turn) = self.rollout.codex_turns.get(codex_turn_id) else {
                bail!("tool call start referenced unknown Codex turn {codex_turn_id}");
            };
            if turn.thread_id != thread_id {
                bail!(
                    "tool call start used thread {thread_id}, but Codex turn {codex_turn_id} \
                     belongs to {}",
                    turn.thread_id
                );
            }
        }
        Ok(())
    }

    fn ensure_unique_model_visible_tool_call(
        &self,
        model_visible_call_id: Option<&str>,
        tool_call_id: &str,
    ) -> Result<()> {
        let Some(model_visible_call_id) = model_visible_call_id else {
            return Ok(());
        };
        if let Some(existing) = self.single_tool_for_model_visible_call(model_visible_call_id)?
            && existing != tool_call_id
        {
            bail!("duplicate tool call for model-visible call id {model_visible_call_id}");
        }
        Ok(())
    }

    fn single_tool_for_model_visible_call(
        &self,
        model_visible_call_id: &str,
    ) -> Result<Option<ToolCallId>> {
        let mut matching = self
            .rollout
            .tool_calls
            .values()
            .filter(|tool| tool.model_visible_call_id.as_deref() == Some(model_visible_call_id))
            .map(|tool| tool.tool_call_id.clone());
        let first = matching.next();
        if matching.next().is_some() {
            bail!("multiple tool calls matched model-visible call id {model_visible_call_id}");
        }
        Ok(first)
    }

    fn model_visible_tool_item_ids(
        &self,
        thread_id: &str,
        call_id: &str,
        kinds: &[ConversationItemKind],
    ) -> Vec<String> {
        self.rollout
            .conversation_items
            .values()
            .filter(|item| {
                item.thread_id == thread_id
                    && item.call_id.as_deref() == Some(call_id)
                    && kinds.contains(&item.kind)
            })
            .map(|item| item.item_id.clone())
            .collect::<Vec<_>>()
    }

    fn add_tool_call_item(&mut self, tool_call_id: &str, item_id: &str) -> Result<()> {
        let Some(tool_call) = self.rollout.tool_calls.get_mut(tool_call_id) else {
            bail!("tool call {tool_call_id} disappeared during conversation linking");
        };
        push_unique(&mut tool_call.model_visible_call_item_ids, item_id);
        Ok(())
    }

    fn add_tool_output_item(&mut self, tool_call_id: &str, item_id: &str) -> Result<()> {
        let Some(tool_call) = self.rollout.tool_calls.get_mut(tool_call_id) else {
            bail!("tool call {tool_call_id} disappeared during output linking");
        };
        push_unique(&mut tool_call.model_visible_output_item_ids, item_id);

        let Some(item) = self.rollout.conversation_items.get_mut(item_id) else {
            bail!("conversation item {item_id} disappeared during output linking");
        };
        let producer = ProducerRef::Tool {
            tool_call_id: tool_call_id.to_string(),
        };
        if !item.produced_by.contains(&producer) {
            item.produced_by.push(producer);
        }
        Ok(())
    }

    fn link_tool_to_inference_response(&mut self, tool_call_id: &str) {
        let Some(tool_call) = self.rollout.tool_calls.get(tool_call_id) else {
            return;
        };
        let call_item_ids = tool_call.model_visible_call_item_ids.clone();
        if call_item_ids.is_empty() {
            return;
        }
        for inference in self.rollout.inference_calls.values_mut() {
            if inference
                .response_item_ids
                .iter()
                .any(|item_id| call_item_ids.contains(item_id))
                && !inference
                    .tool_call_ids_started_by_response
                    .contains(&tool_call_id.to_string())
            {
                inference
                    .tool_call_ids_started_by_response
                    .push(tool_call_id.to_string());
            }
        }
    }
}

fn push_unique(items: &mut Vec<String>, item_id: &str) {
    if !items.iter().any(|existing| existing == item_id) {
        items.push(item_id.to_string());
    }
}