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

//! Hot-path helpers for recording code-mode runtime cell lifecycles.
//!
//! The public `exec` tool is reduced as a first-class `CodeCell` instead of a
//! generic tool call. This module keeps the runtime response serialization and
//! lifecycle event policy inside the trace crate while core carries a compact,
//! no-op capable handle through execution and waits.

use std::sync::Arc;

use codex_code_mode::RuntimeResponse;
use serde::Serialize;
use tracing::warn;

use crate::model::AgentThreadId;
use crate::model::CodeCellRuntimeStatus;
use crate::model::CodexTurnId;
use crate::model::ModelVisibleCallId;
use crate::payload::RawPayloadKind;
use crate::payload::RawPayloadRef;
use crate::raw_event::RawTraceEventContext;
use crate::raw_event::RawTraceEventPayload;
use crate::writer::TraceWriter;

/// No-op capable trace handle for one code-mode runtime cell.
#[derive(Clone, Debug)]
pub struct CodeCellTraceContext {
    state: CodeCellTraceContextState,
}

#[derive(Clone, Debug)]
enum CodeCellTraceContextState {
    Disabled,
    Enabled(EnabledCodeCellTraceContext),
}

#[derive(Clone, Debug)]
struct EnabledCodeCellTraceContext {
    writer: Arc<TraceWriter>,
    thread_id: AgentThreadId,
    codex_turn_id: CodexTurnId,
    runtime_cell_id: String,
}

/// Raw code-mode response captured at the runtime boundary.
///
/// This is not the model-visible custom-tool output. The reducer links that
/// output through `CodeCell.output_item_ids` once the conversation item appears.
/// Keeping the raw runtime payload here preserves stored-value and lifecycle
/// evidence without duplicating the model-facing transcript.
#[derive(Serialize)]
struct CodeCellResponseTracePayload<'a> {
    response: &'a RuntimeResponse,
}

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

    /// Builds a context for an already-known code-mode runtime cell.
    pub(crate) fn enabled(
        writer: Arc<TraceWriter>,
        thread_id: impl Into<AgentThreadId>,
        codex_turn_id: impl Into<CodexTurnId>,
        runtime_cell_id: impl Into<String>,
    ) -> Self {
        Self {
            state: CodeCellTraceContextState::Enabled(EnabledCodeCellTraceContext {
                writer,
                thread_id: thread_id.into(),
                codex_turn_id: codex_turn_id.into(),
                runtime_cell_id: runtime_cell_id.into(),
            }),
        }
    }

    /// Records the parent runtime object before JavaScript can issue nested tool calls.
    pub fn record_started(
        &self,
        model_visible_call_id: impl Into<ModelVisibleCallId>,
        source_js: impl Into<String>,
    ) {
        let CodeCellTraceContextState::Enabled(context) = &self.state else {
            return;
        };
        append_with_context_best_effort(
            context,
            RawTraceEventPayload::CodeCellStarted {
                runtime_cell_id: context.runtime_cell_id.clone(),
                model_visible_call_id: model_visible_call_id.into(),
                source_js: source_js.into(),
            },
        );
    }

    /// Records the first response returned by the public code-mode `exec` tool.
    ///
    /// A yielded response returns control to the model while the cell keeps
    /// running. Terminal initial responses should be followed by `record_ended`
    /// by the caller so the reducer can distinguish model-visible output from
    /// runtime completion.
    pub fn record_initial_response(&self, response: &RuntimeResponse) {
        let CodeCellTraceContextState::Enabled(context) = &self.state else {
            return;
        };
        append_with_context_best_effort(
            context,
            RawTraceEventPayload::CodeCellInitialResponse {
                runtime_cell_id: context.runtime_cell_id.clone(),
                status: code_cell_status_for_runtime_response(response),
                response_payload: code_cell_response_payload(context, response),
            },
        );
    }

    /// Records the terminal lifecycle point for a code-mode runtime cell.
    pub fn record_ended(&self, response: &RuntimeResponse) {
        let CodeCellTraceContextState::Enabled(context) = &self.state else {
            return;
        };
        append_with_context_best_effort(
            context,
            RawTraceEventPayload::CodeCellEnded {
                runtime_cell_id: context.runtime_cell_id.clone(),
                status: code_cell_status_for_runtime_response(response),
                response_payload: code_cell_response_payload(context, response),
            },
        );
    }
}

fn code_cell_status_for_runtime_response(response: &RuntimeResponse) -> CodeCellRuntimeStatus {
    match response {
        RuntimeResponse::Yielded { .. } => CodeCellRuntimeStatus::Yielded,
        RuntimeResponse::Terminated { .. } => CodeCellRuntimeStatus::Terminated,
        RuntimeResponse::Result { error_text, .. } => {
            if error_text.is_some() {
                CodeCellRuntimeStatus::Failed
            } else {
                CodeCellRuntimeStatus::Completed
            }
        }
    }
}

fn code_cell_response_payload(
    context: &EnabledCodeCellTraceContext,
    response: &RuntimeResponse,
) -> Option<RawPayloadRef> {
    write_json_payload_best_effort(
        &context.writer,
        RawPayloadKind::ToolResult,
        &CodeCellResponseTracePayload { response },
    )
}

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: &EnabledCodeCellTraceContext,
    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:#}");
    }
}