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

use std::sync::Arc;

use pretty_assertions::assert_eq;
use serde_json::json;
use tempfile::TempDir;

use crate::model::CodeCellRuntimeStatus;
use crate::model::ConversationItemKind;
use crate::model::ExecutionStatus;
use crate::model::ProducerRef;
use crate::model::RolloutStatus;
use crate::model::ToolCallKind;
use crate::model::ToolCallSummary;
use crate::payload::RawPayloadKind;
use crate::raw_event::RawToolCallRequester;
use crate::raw_event::RawTraceEventPayload;
use crate::reducer::test_support::append_inference_request;
use crate::reducer::test_support::append_inference_start;
use crate::reducer::test_support::append_inference_start_for_thread;
use crate::reducer::test_support::create_started_writer;
use crate::reducer::test_support::generic_summary;
use crate::reducer::test_support::message;
use crate::reducer::test_support::start_turn;
use crate::reducer::test_support::start_turn_for_thread;
use crate::reducer::test_support::trace_context;
use crate::reducer::test_support::trace_context_for_thread;
use crate::replay_bundle;

#[test]
fn code_cell_lifecycle_links_nested_tools_and_outputs_without_inferring_legacy_wait()
-> anyhow::Result<()> {
    let temp = TempDir::new()?;
    let writer = create_started_writer(&temp)?;
    start_turn(&writer, "turn-1")?;

    let request = writer.write_json_payload(
        RawPayloadKind::InferenceRequest,
        &json!({
            "input": [message("user", "count files")]
        }),
    )?;
    append_inference_start(&writer, "inference-1", "turn-1", request)?;
    let response = writer.write_json_payload(
        RawPayloadKind::InferenceResponse,
        &json!({
            "response_id": "resp-1",
            "output_items": [{
                "type": "custom_tool_call",
                "name": "exec",
                "call_id": "call-code",
                "input": "text('hi')"
            }]
        }),
    )?;
    // Runtime tool dispatch starts before the stream-completion hook has
    // reduced the model response that requested `exec`.
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CodeCellStarted {
            runtime_cell_id: "1".to_string(),
            model_visible_call_id: "call-code".to_string(),
            source_js: "text('hi')".to_string(),
        },
    )?;
    writer.append(RawTraceEventPayload::InferenceCompleted {
        inference_call_id: "inference-1".to_string(),
        response_id: Some("resp-1".to_string()),
        upstream_request_id: None,
        response_payload: response,
    })?;
    let initial_response = writer.write_json_payload(
        RawPayloadKind::ToolResult,
        &json!({"response": "initial caller result"}),
    )?;
    let final_response = writer.write_json_payload(
        RawPayloadKind::ToolResult,
        &json!({"response": "later runtime result"}),
    )?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CodeCellInitialResponse {
            runtime_cell_id: "1".to_string(),
            status: CodeCellRuntimeStatus::Yielded,
            response_payload: Some(initial_response.clone()),
        },
    )?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::ToolCallStarted {
            tool_call_id: "nested-tool-1".to_string(),
            model_visible_call_id: None,
            code_mode_runtime_tool_id: Some("tool-1".to_string()),
            requester: RawToolCallRequester::CodeCell {
                runtime_cell_id: "1".to_string(),
            },
            kind: ToolCallKind::ExecCommand,
            summary: ToolCallSummary::Generic {
                label: "exec_command".to_string(),
                input_preview: Some("pwd".to_string()),
                output_preview: None,
            },
            invocation_payload: None,
        },
    )?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::ToolCallEnded {
            tool_call_id: "nested-tool-1".to_string(),
            status: ExecutionStatus::Completed,
            result_payload: None,
        },
    )?;

    start_turn(&writer, "turn-2")?;
    let followup = writer.write_json_payload(
        RawPayloadKind::InferenceRequest,
        &json!({
            "previous_response_id": "resp-1",
            "input": [{
                "type": "custom_tool_call_output",
                "call_id": "call-code",
                "output": "Script running with cell ID 1"
            }]
        }),
    )?;
    append_inference_start(&writer, "inference-2", "turn-2", followup)?;
    let wait_request = writer.write_json_payload(
        RawPayloadKind::ToolInvocation,
        &json!({
            "tool_name": "wait",
            "tool_namespace": null,
            "payload": {
                "type": "function",
                "arguments": "{\"cell_id\":\"1\"}"
            }
        }),
    )?;
    writer.append_with_context(
        trace_context("turn-2"),
        RawTraceEventPayload::ToolCallStarted {
            tool_call_id: "wait-tool-1".to_string(),
            model_visible_call_id: Some("wait-call".to_string()),
            code_mode_runtime_tool_id: None,
            requester: RawToolCallRequester::Model,
            kind: ToolCallKind::Other {
                name: "wait".to_string(),
            },
            summary: ToolCallSummary::Generic {
                label: "wait".to_string(),
                input_preview: Some("{\"cell_id\":\"1\"}".to_string()),
                output_preview: None,
            },
            invocation_payload: Some(wait_request),
        },
    )?;
    writer.append_with_context(
        trace_context("turn-2"),
        RawTraceEventPayload::CodeCellEnded {
            runtime_cell_id: "1".to_string(),
            status: CodeCellRuntimeStatus::Completed,
            response_payload: Some(final_response.clone()),
        },
    )?;

    let rollout = replay_bundle(temp.path())?;
    let code_cell_id = test_reduced_code_cell_id("call-code");
    let cell = &rollout.code_cells[&code_cell_id];
    let output_item_id = rollout.inference_calls["inference-2"]
        .request_item_ids
        .last()
        .expect("exec output item");

    assert_eq!(
        (
            &cell.initial_response_status,
            &cell.raw_initial_response_payload_id,
            &cell.raw_result_payload_id,
        ),
        (
            &Some(CodeCellRuntimeStatus::Yielded),
            &Some(initial_response.raw_payload_id),
            &Some(final_response.raw_payload_id),
        )
    );
    assert_eq!(cell.thread_id, "thread-root");
    assert_eq!(cell.runtime_status, CodeCellRuntimeStatus::Completed);
    assert_eq!(cell.execution.status, ExecutionStatus::Completed);
    assert_eq!(cell.runtime_cell_id, Some("1".to_string()));
    assert_eq!(cell.nested_tool_call_ids, vec!["nested-tool-1"]);
    // The legacy request alone does not establish that the host resolved a target.
    assert!(cell.wait_tool_call_ids.is_empty());
    assert_eq!(cell.output_item_ids, vec![output_item_id.clone()]);
    assert_eq!(
        rollout.conversation_items[output_item_id].produced_by,
        vec![ProducerRef::CodeCell {
            code_cell_id: code_cell_id.clone(),
        }]
    );
    assert_eq!(
        rollout.conversation_items[&cell.source_item_id].kind,
        ConversationItemKind::CustomToolCall,
    );

    Ok(())
}

#[test]
fn fast_code_cell_lifecycle_waits_for_source_item() -> anyhow::Result<()> {
    let temp = TempDir::new()?;
    let writer = create_started_writer(&temp)?;
    start_turn(&writer, "turn-1")?;

    // Put outputs on either side of ordinal 10, where map-key order differs from history order.
    let mut input: Vec<_> = (0..7)
        .map(|index| message("user", &format!("context {index}")))
        .collect();
    for (dispatch_id, call_id) in [("outer", "call-code"), ("no-cell", "other-exec")] {
        writer.append_with_context(
            trace_context("turn-1"),
            RawTraceEventPayload::ToolCallStarted {
                tool_call_id: dispatch_id.into(),
                model_visible_call_id: Some(call_id.into()),
                code_mode_runtime_tool_id: None,
                requester: RawToolCallRequester::Model,
                kind: ToolCallKind::Other {
                    name: "exec".into(),
                },
                summary: generic_summary("exec"),
                invocation_payload: None,
            },
        )?;
    }
    append_inference_request(
        &writer,
        "thread-root",
        "turn-1",
        "inference-1",
        input.clone(),
    )?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CodeCellStarted {
            runtime_cell_id: "1".to_string(),
            model_visible_call_id: "call-code".to_string(),
            source_js: "not valid js".to_string(),
        },
    )?;
    let initial_response = writer.write_json_payload(
        RawPayloadKind::ToolResult,
        &json!({"response": "initial caller result"}),
    )?;
    let final_response = writer.write_json_payload(
        RawPayloadKind::ToolResult,
        &json!({"response": "later runtime result"}),
    )?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CodeCellInitialResponse {
            runtime_cell_id: "1".to_string(),
            status: CodeCellRuntimeStatus::Failed,
            response_payload: Some(initial_response.clone()),
        },
    )?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CodeCellEnded {
            runtime_cell_id: "1".to_string(),
            status: CodeCellRuntimeStatus::Failed,
            response_payload: Some(final_response.clone()),
        },
    )?;
    let call = json!({
        "type": "custom_tool_call",
        "name": "exec",
        "call_id": "call-code",
        "input": "not valid js"
    });
    // Compaction exposes the source and original output before the response is recorded.
    // Its replacement output must be appended before the queued cell is materialized.
    input.push(call.clone());
    input.push(json!({
        "type": "custom_tool_call_output", "call_id": "call-code", "output": "initial caller result"
    }));
    let checkpoint = writer.write_json_payload(
        RawPayloadKind::CompactionCheckpoint,
        &json!({
            "input_history": input,
            "replacement_history": [call, {
                "type": "custom_tool_call_output", "call_id": "call-code", "output": "later runtime result"
            }]
        }),
    )?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CompactionInstalled {
            compaction_id: "compaction-1".to_string(),
            checkpoint_payload: checkpoint,
        },
    )?;
    let response = writer.write_json_payload(
        RawPayloadKind::InferenceResponse,
        &json!({
            "response_id": "resp-1",
            "output_items": [call]
        }),
    )?;
    writer.append(RawTraceEventPayload::InferenceCompleted {
        inference_call_id: "inference-1".to_string(),
        response_id: Some("resp-1".to_string()),
        upstream_request_id: None,
        response_payload: response,
    })?;

    let rollout = replay_bundle(temp.path())?;
    let code_cell_id = test_reduced_code_cell_id("call-code");
    let cell = &rollout.code_cells[&code_cell_id];

    assert_eq!(
        rollout.tool_calls["outer"].started_code_cell_id.as_ref(),
        Some(&cell.code_cell_id)
    );
    assert_eq!(rollout.tool_calls["no-cell"].started_code_cell_id, None);
    assert_eq!(
        cell.output_item_ids,
        vec!["conversation_item:9", "conversation_item:12"]
    );
    assert_eq!(
        (
            &cell.initial_response_status,
            &cell.raw_initial_response_payload_id,
            &cell.raw_result_payload_id,
        ),
        (
            &Some(CodeCellRuntimeStatus::Failed),
            &Some(initial_response.raw_payload_id),
            &Some(final_response.raw_payload_id),
        )
    );
    assert_eq!(cell.thread_id, "thread-root");
    assert_eq!(cell.runtime_status, CodeCellRuntimeStatus::Failed);
    assert_eq!(cell.execution.status, ExecutionStatus::Failed);
    assert_eq!(cell.runtime_cell_id, Some("1".to_string()));
    assert_eq!(
        rollout.conversation_items[&cell.source_item_id].kind,
        ConversationItemKind::CustomToolCall,
    );

    Ok(())
}

#[test]
fn cancelled_turn_terminates_unfinished_code_cell() -> anyhow::Result<()> {
    let temp = TempDir::new()?;
    let writer = create_started_writer(&temp)?;
    start_turn(&writer, "turn-1")?;

    let request = writer.write_json_payload(
        RawPayloadKind::InferenceRequest,
        &json!({
            "input": [message("user", "count files")]
        }),
    )?;
    append_inference_start(&writer, "inference-1", "turn-1", request)?;
    let response = writer.write_json_payload(
        RawPayloadKind::InferenceResponse,
        &json!({
            "response_id": "resp-1",
            "output_items": [{
                "type": "custom_tool_call",
                "name": "exec",
                "call_id": "call-code",
                "input": "await tools.exec_command({cmd: 'slow'});"
            }]
        }),
    )?;
    writer.append(RawTraceEventPayload::InferenceCompleted {
        inference_call_id: "inference-1".to_string(),
        response_id: Some("resp-1".to_string()),
        upstream_request_id: None,
        response_payload: response,
    })?;
    writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CodeCellStarted {
            runtime_cell_id: "1".to_string(),
            model_visible_call_id: "call-code".to_string(),
            source_js: "await tools.exec_command({cmd: 'slow'});".to_string(),
        },
    )?;
    let turn_end = writer.append_with_context(
        trace_context("turn-1"),
        RawTraceEventPayload::CodexTurnEnded {
            codex_turn_id: "turn-1".to_string(),
            status: ExecutionStatus::Cancelled,
        },
    )?;

    let rollout = replay_bundle(temp.path())?;
    let code_cell_id = test_reduced_code_cell_id("call-code");
    let cell = &rollout.code_cells[&code_cell_id];

    assert_eq!(cell.runtime_status, CodeCellRuntimeStatus::Terminated);
    assert_eq!(cell.execution.status, ExecutionStatus::Cancelled);
    assert_eq!(cell.execution.ended_seq, Some(turn_end.seq));

    Ok(())
}

#[test]
fn runtime_code_cell_ids_can_repeat_across_threads_and_after_resume() -> anyhow::Result<()> {
    let temp = TempDir::new()?;
    let writer = Arc::new(create_started_writer(&temp)?);
    let mut pending_dispatches = Vec::new();
    writer.append(RawTraceEventPayload::ThreadStarted {
        thread_id: "thread-child".to_string(),
        agent_path: "/root/child".to_string(),
        metadata_payload: None,
    })?;
    for (thread_id, turn_id, inference_call_id, call_id) in [
        ("thread-root", "turn-root", "inference-root", "call-root"),
        (
            "thread-child",
            "turn-child",
            "inference-child",
            "call-child",
        ),
        (
            "thread-child",
            "turn-resumed",
            "inference-resumed",
            "call-resumed",
        ),
    ] {
        if turn_id == "turn-resumed" {
            writer.append(RawTraceEventPayload::ThreadEnded {
                thread_id: thread_id.to_string(),
                status: RolloutStatus::Completed,
            })?;
            writer.append(RawTraceEventPayload::ThreadResumed {
                thread_id: thread_id.to_string(),
            })?;
        }
        start_turn_for_thread(&writer, thread_id, turn_id)?;
        let request = writer.write_json_payload(
            RawPayloadKind::InferenceRequest,
            &json!({
                "input": [message("user", "run code")]
            }),
        )?;
        append_inference_start_for_thread(&writer, thread_id, turn_id, inference_call_id, request)?;
        writer.append_with_context(
            trace_context_for_thread(thread_id, turn_id),
            RawTraceEventPayload::CodeCellStarted {
                runtime_cell_id: "1".to_string(),
                model_visible_call_id: call_id.to_string(),
                source_js: "text('hi')".to_string(),
            },
        )?;
        let waits: &[(&str, Option<&str>, &str)] = match turn_id {
            "turn-root" => &[("rejected", None, "{")],
            "turn-child" => &[
                ("confirmed", None, r#"{"cell_id":"1"}"#),
                ("late", None, r#"{"cell_id":"1"}"#),
            ],
            _ => &[("foreign", Some("fixture"), r#"{"cell_id":"1"}"#)],
        };
        for &(wait_id, namespace, arguments) in waits {
            let dispatch = crate::ToolDispatchTraceContext::start(
                Arc::clone(&writer),
                crate::ToolDispatchInvocation {
                    thread_id: thread_id.into(),
                    codex_turn_id: turn_id.into(),
                    tool_call_id: wait_id.into(),
                    tool_name: "wait".into(),
                    tool_namespace: namespace.map(str::to_owned),
                    requester: crate::ToolDispatchRequester::Model {
                        model_visible_call_id: wait_id.into(),
                    },
                    payload: crate::ToolDispatchPayload::Function {
                        arguments: arguments.into(),
                    },
                },
            );
            if wait_id == "rejected" {
                dispatch.record_failed("invalid wait arguments");
                continue;
            }
            // Replay these dispatches while they are still pending.
            pending_dispatches.push(dispatch);
            if wait_id == "confirmed" {
                crate::CodeCellTraceContext::enabled(Arc::clone(&writer), thread_id, turn_id, "1")
                    .record_wait_target(wait_id);
            }
        }
        let response = writer.write_json_payload(
            RawPayloadKind::InferenceResponse,
            &json!({
                "response_id": format!("resp-{turn_id}"),
                "output_items": [{
                    "type": "custom_tool_call",
                    "name": "exec",
                    "call_id": call_id,
                    "input": "text('hi')"
                }]
            }),
        )?;
        writer.append(RawTraceEventPayload::InferenceCompleted {
            inference_call_id: inference_call_id.to_string(),
            response_id: Some(format!("resp-{turn_id}")),
            upstream_request_id: None,
            response_payload: response,
        })?;
        if turn_id == "turn-resumed" {
            // The pre-resume wait is confirmed after the new cell has reused its handle.
            crate::CodeCellTraceContext::enabled(Arc::clone(&writer), thread_id, "turn-child", "1")
                .record_wait_target("late");
        }
        writer.append_with_context(
            trace_context_for_thread(thread_id, turn_id),
            RawTraceEventPayload::CodeCellEnded {
                runtime_cell_id: "1".to_string(),
                status: CodeCellRuntimeStatus::Completed,
                response_payload: None,
            },
        )?;
    }

    let rollout = replay_bundle(temp.path())?;
    let root_cell_id = test_reduced_code_cell_id("call-root");
    let child_cell_id = test_reduced_code_cell_id("call-child");
    let resumed_cell_id = test_reduced_code_cell_id("call-resumed");

    assert_eq!(rollout.code_cells[&root_cell_id].thread_id, "thread-root");
    assert_eq!(rollout.code_cells[&child_cell_id].thread_id, "thread-child");
    assert_eq!(
        rollout.code_cells[&resumed_cell_id].thread_id,
        "thread-child"
    );
    assert_eq!(
        rollout.code_cells[&root_cell_id].runtime_cell_id,
        Some("1".to_string())
    );
    assert_eq!(
        rollout.code_cells[&child_cell_id].runtime_cell_id,
        Some("1".to_string())
    );
    assert_eq!(
        rollout.code_cells[&resumed_cell_id].runtime_cell_id,
        Some("1".to_string())
    );
    assert_eq!(
        rollout.code_cells[&child_cell_id].wait_tool_call_ids,
        ["confirmed"]
    );
    assert!(
        rollout.code_cells[&root_cell_id]
            .wait_tool_call_ids
            .is_empty()
    );
    assert!(
        rollout.code_cells[&resumed_cell_id]
            .wait_tool_call_ids
            .is_empty()
    );
    for id in ["confirmed", "late"] {
        assert_eq!(
            rollout.tool_calls[id]
                .resolved_wait_runtime_cell_id
                .as_deref(),
            Some("1")
        );
    }
    assert!(
        rollout.tool_calls["confirmed"]
            .execution
            .ended_seq
            .is_none()
    );
    for id in ["rejected", "foreign"] {
        assert_eq!(rollout.tool_calls[id].resolved_wait_runtime_cell_id, None);
    }

    Ok(())
}

fn test_reduced_code_cell_id(model_visible_call_id: &str) -> String {
    format!("code_cell:{model_visible_call_id}")
}