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::ToolCallKind;use crate::model::ToolCallSummary;use crate::payload::RawPayloadKind;use crate::raw_event::RawToolCallRequester;use crate::raw_event::RawTraceEventPayload;use crate::reducer::test_support::create_started_writer;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_waits_and_outputs() -> 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")]        }),    )?;    writer.append(RawTraceEventPayload::InferenceStarted {        inference_call_id: "inference-1".to_string(),        thread_id: "thread-root".to_string(),        codex_turn_id: "turn-1".to_string(),        model: "gpt-test".to_string(),        provider_name: "test-provider".to_string(),        request_payload: 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,    })?;    writer.append_with_context(        trace_context("turn-1"),        RawTraceEventPayload::CodeCellInitialResponse {            runtime_cell_id: "1".to_string(),            status: CodeCellRuntimeStatus::Yielded,            response_payload: None,        },    )?;    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"            }]        }),    )?;    writer.append(RawTraceEventPayload::InferenceStarted {        inference_call_id: "inference-2".to_string(),        thread_id: "thread-root".to_string(),        codex_turn_id: "turn-2".to_string(),        model: "gpt-test".to_string(),        provider_name: "test-provider".to_string(),        request_payload: 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: None,        },    )?;    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.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"]);    assert_eq!(cell.wait_tool_call_ids, vec!["wait-tool-1"]);    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")?;    let request = writer.write_json_payload(        RawPayloadKind::InferenceRequest,        &json!({            "input": [message("user", "count files")]        }),    )?;    writer.append(RawTraceEventPayload::InferenceStarted {        inference_call_id: "inference-1".to_string(),        thread_id: "thread-root".to_string(),        codex_turn_id: "turn-1".to_string(),        model: "gpt-test".to_string(),        provider_name: "test-provider".to_string(),        request_payload: request,    })?;    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(),        },    )?;    writer.append_with_context(        trace_context("turn-1"),        RawTraceEventPayload::CodeCellInitialResponse {            runtime_cell_id: "1".to_string(),            status: CodeCellRuntimeStatus::Failed,            response_payload: None,        },    )?;    writer.append_with_context(        trace_context("turn-1"),        RawTraceEventPayload::CodeCellEnded {            runtime_cell_id: "1".to_string(),            status: CodeCellRuntimeStatus::Failed,            response_payload: None,        },    )?;    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": "not valid js"            }]        }),    )?;    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!(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")]        }),    )?;    writer.append(RawTraceEventPayload::InferenceStarted {        inference_call_id: "inference-1".to_string(),        thread_id: "thread-root".to_string(),        codex_turn_id: "turn-1".to_string(),        model: "gpt-test".to_string(),        provider_name: "test-provider".to_string(),        request_payload: 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() -> anyhow::Result<()> {    let temp = TempDir::new()?;    let writer = create_started_writer(&temp)?;    writer.append(RawTraceEventPayload::ThreadStarted {        thread_id: "thread-child".to_string(),        agent_path: "/root/child".to_string(),        metadata_payload: None,    })?;    start_turn_for_thread(&writer, "thread-root", "turn-root")?;    start_turn_for_thread(&writer, "thread-child", "turn-child")?;    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",        ),    ] {        let request = writer.write_json_payload(            RawPayloadKind::InferenceRequest,            &json!({                "input": [message("user", "run code")]            }),        )?;        writer.append(RawTraceEventPayload::InferenceStarted {            inference_call_id: inference_call_id.to_string(),            thread_id: thread_id.to_string(),            codex_turn_id: turn_id.to_string(),            model: "gpt-test".to_string(),            provider_name: "test-provider".to_string(),            request_payload: 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 response = writer.write_json_payload(            RawPayloadKind::InferenceResponse,            &json!({                "response_id": format!("resp-{thread_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-{thread_id}")),            upstream_request_id: None,            response_payload: response,        })?;        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");    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[&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())    );    Ok(())}fn test_reduced_code_cell_id(model_visible_call_id: &str) -> String {    format!("code_cell:{model_visible_call_id}")}