codex/codex-rs/core/tests/suite/guardian_retry.rs

//! Exercises recovery through the parent tool call, reviewer, and executor boundary.

use anyhow::Context;
use anyhow::Result;
use codex_core::config::Constrained;
use codex_core::context::UserGoalUpdate;
use codex_protocol::AgentPath;
use codex_protocol::approvals::GuardianAssessmentStatus;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::InterAgentCommunication;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::SandboxPolicy;
use core_test_support::responses::*;
use core_test_support::skip_if_no_network;
use core_test_support::skip_if_wine_exec;
use core_test_support::streaming_sse::StreamingSseChunk;
use core_test_support::streaming_sse::start_streaming_sse_server;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use std::time::Duration;
use test_case::test_case;

#[derive(Clone, Copy)]
enum ReviewInterruption {
    Cancellation,
    Timeout,
}

#[test_case(ReviewInterruption::Cancellation, None; "cancellation")]
#[test_case(ReviewInterruption::Cancellation, Some(0); "unlimited_cancellation")]
#[test_case(ReviewInterruption::Timeout, None; "timeout")]
#[test_case(ReviewInterruption::Timeout, Some(180_000); "configured_timeout")]
#[tokio::test]
async fn interrupted_guardian_review_preserves_evidence_once(
    interruption: ReviewInterruption,
    timeout_ms: Option<u32>,
) -> Result<()> {
    skip_if_no_network!(Ok(()));
    skip_if_wine_exec!(
        Ok(()),
        "Guardian approval actions require host-native paths"
    );
    let command = |id: &str| {
        sse(vec![
            ev_function_call(
                id,
                "exec_command",
                &json!({
                    "cmd": "echo unexpected > guardian-interrupted.txt",
                    "sandbox_permissions": "require_escalated",
                    "justification": "Exercise interrupted Guardian review recovery",
                })
                .to_string(),
            ),
            ev_completed(id),
        ])
    };
    let denial = |id: &str| {
        sse(vec![
            ev_assistant_message(
                id,
                r#"{"risk_level":"low","user_authorization":"low","outcome":"deny","rationale":"Fixture denial."}"#,
            ),
            ev_completed(id),
        ])
    };
    let (hold_review, review_gate) = tokio::sync::oneshot::channel();
    let mut responses = vec![
        (None, sse(vec![ev_completed("waiting-for-evidence")])),
        (None, command("committed-command")),
        (None, denial("msg_committed_review")),
        (None, sse(vec![ev_completed("first-parent-done")])),
        (None, command("interrupted-command")),
        (Some(review_gate), denial("msg_interrupted_review")),
    ];
    if matches!(interruption, ReviewInterruption::Timeout) {
        responses.push((None, sse(vec![ev_completed("timed-out-parent-done")])));
    }
    responses.extend([
        (None, command("followup-command")),
        (None, denial("msg_followup_review")),
        (None, sse(vec![ev_completed("last-parent-done")])),
    ]);
    let (model, _) = start_streaming_sse_server(
        responses
            .into_iter()
            .map(|(gate, body)| vec![StreamingSseChunk { gate, body }])
            .collect(),
    )
    .await;
    let server = start_mock_server().await;
    let base_url = format!("{}/v1", model.uri());
    let test = test_codex()
        .with_pre_build_hook(move |home| {
            if let Some(timeout_ms) = timeout_ms {
                std::fs::write(
                    home.join("config.toml"),
                    format!("[auto_review]\ntimeout_ms = {timeout_ms}\n"),
                )
                .expect("write timeout config");
            }
        })
        .with_config(move |config| {
            config.model_provider.base_url = Some(base_url);
            config.model_provider.stream_idle_timeout_ms = Some(600_000);
            config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
            config.approvals_reviewer = ApprovalsReviewer::AutoReview;
            config
                .set_legacy_sandbox_policy(SandboxPolicy::new_workspace_write_policy())
                .expect("set sandbox policy");
        })
        .build_with_auto_env(&server)
        .await?;
    test.submit_text_turn("Wait for the parent's instructions before requesting the command.")
        .await?;
    // Commit one review, then deliver new encrypted evidence for the interrupted attempt.
    for evidence in ["committed-evidence", "pending-evidence"] {
        test.codex
            .submit(Op::InterAgentCommunication {
                communication: InterAgentCommunication::new_encrypted(
                    AgentPath::root(),
                    AgentPath::root().join("worker").expect("worker path"),
                    Vec::new(),
                    evidence.into(),
                    /*trigger_turn*/ true,
                ),
                start_options: Default::default(),
            })
            .await?;
        if evidence == "committed-evidence" {
            wait_for_event(&test.codex, |event| {
                matches!(event, EventMsg::TurnComplete(_))
            })
            .await;
        }
    }
    // The request reaching the server proves Guardian has admitted the pending evidence.
    tokio::time::timeout(
        Duration::from_secs(20),
        model.wait_for_request_count(/*count*/ 6),
    )
    .await
    .context("wait for pending Guardian review")?;
    match interruption {
        ReviewInterruption::Cancellation => {
            if timeout_ms == Some(0) {
                tokio::time::pause();
                tokio::time::advance(Duration::from_secs(120)).await;
                tokio::time::resume();
            }
            test.codex.submit(Op::Interrupt).await?;
        }
        ReviewInterruption::Timeout => {
            tokio::time::pause();
            tokio::time::advance(
                Duration::from_millis(u64::from(timeout_ms.unwrap_or(90_000)))
                    + Duration::from_secs(1),
            )
            .await;
            tokio::time::resume();
        }
    }
    let expected_status = match interruption {
        ReviewInterruption::Cancellation => GuardianAssessmentStatus::Aborted,
        ReviewInterruption::Timeout => GuardianAssessmentStatus::TimedOut,
    };
    let mut review_finished = false;
    let mut parent_finished = false;
    while !review_finished || !parent_finished {
        match test.codex.next_event().await?.msg {
            EventMsg::GuardianAssessment(review) if review.status == expected_status => {
                review_finished = true;
            }
            EventMsg::TurnAborted(_)
                if matches!(interruption, ReviewInterruption::Cancellation) =>
            {
                parent_finished = true;
            }
            EventMsg::TurnComplete(_) if matches!(interruption, ReviewInterruption::Timeout) => {
                parent_finished = true;
            }
            _ => {}
        }
    }
    test.submit_text_turn("Request the command again.").await?;
    let requests = model
        .requests()
        .await
        .into_iter()
        .map(|body| serde_json::from_slice::<Value>(&body))
        .collect::<serde_json::Result<Vec<_>>>()?;
    let reviews = requests
        .iter()
        .filter(|request| request["client_metadata"]["x-openai-subagent"] == "guardian")
        .collect::<Vec<_>>();
    assert_eq!(reviews.len(), 3);
    let pending_messages = reviews[1]["input"]
        .as_array()
        .expect("pending input")
        .iter()
        .filter(|item| item["type"] == "agent_message")
        .collect::<Vec<_>>();
    assert_eq!(pending_messages.len(), 2);
    let followup = reviews[2]["input"].as_array().expect("followup input");
    assert_eq!(
        followup
            .iter()
            .filter(|item| item["type"] == "agent_message")
            .collect::<Vec<_>>(),
        pending_messages,
        "the next review must receive both encrypted messages exactly once",
    );
    assert_eq!(
        followup
            .iter()
            .filter(|item| item["id"] == "msg_committed_review")
            .count(),
        1,
        "the last completed review must survive rollback",
    );
    let marker_error = test
        .fs()
        .get_metadata(
            &test.workspace_path_uri("guardian-interrupted.txt")?,
            Default::default(),
            /*sandbox*/ None,
        )
        .await
        .expect_err("interrupted or denied commands must not execute");
    assert_eq!(marker_error.kind(), std::io::ErrorKind::NotFound);
    test.codex.shutdown_and_wait().await?;
    drop(hold_review);
    model.shutdown().await;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn guardian_retry_executes_only_after_a_completed_approval() -> Result<()> {
    skip_if_no_network!(Ok(()));
    skip_if_wine_exec!(
        Ok(()),
        "Guardian approval actions require host-native paths"
    );
    let server = start_mock_server().await;
    let mut builder = test_codex().with_config(|config| {
        config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
        config.approvals_reviewer = ApprovalsReviewer::AutoReview;
        config
            .set_legacy_sandbox_policy(SandboxPolicy::new_workspace_write_policy())
            .expect("set sandbox policy");
    });
    let test = builder.build_with_auto_env(&server).await?;
    let responses = vec![
        sse(vec![ev_completed("waiting-for-parent")]),
        sse(vec![
            ev_function_call(
                "write-marker",
                "exec_command",
                &json!({
                    "cmd": "echo executed >> guardian-retry.txt",
                    "sandbox_permissions": "require_escalated",
                    "justification": "Write the requested marker",
                })
                .to_string(),
            ),
            ev_completed("parent-call"),
        ]),
        sse_failed(
            "first-review",
            "rate_limit_exceeded",
            "temporary review error",
        ),
        sse_failed(
            "stream-retry",
            "rate_limit_exceeded",
            "temporary review error",
        ),
        sse(vec![
            ev_assistant_message(
                "approval",
                r#"{"risk_level":"low","user_authorization":"high","outcome":"allow","rationale":"requested write"}"#,
            ),
            ev_completed("review-approved"),
        ]),
        sse(vec![ev_completed("parent-done")]),
    ];
    let requests = mount_sse_sequence(&server, responses).await;
    test.submit_text_turn("Write the marker once after the parent confirms.")
        .await?;
    // Deliver encrypted evidence through the real mailbox before the action that
    // needs review. A failed reviewer turn must not admit this native item again.
    test.codex
        .submit(Op::InterAgentCommunication {
            communication: InterAgentCommunication::new_encrypted(
                AgentPath::root(),
                AgentPath::root().join("worker").expect("worker path"),
                Vec::new(),
                "opaque-parent-approval".into(),
                /*trigger_turn*/ true,
            ),
            start_options: Default::default(),
        })
        .await?;
    let mut reviews = Vec::new();
    let mut warnings = Vec::new();
    let mut command_started_at_ms = None;
    loop {
        match test.codex.next_event().await?.msg {
            EventMsg::GuardianAssessment(review) => reviews.push(review),
            EventMsg::ItemCompleted(event) if event.item.id() == "write-marker" => {
                command_started_at_ms = event.started_at_ms;
            }
            EventMsg::GuardianWarning(warning) => warnings.push(warning.message),
            EventMsg::TurnComplete(_) => break,
            _ => {}
        }
    }
    let requests = requests.requests();
    assert_eq!(requests.len(), 6);
    let guardian_requests = &requests[2..5];
    let expected_messages = guardian_requests[0].inputs_of_type("agent_message");
    assert_eq!(expected_messages.len(), 1);
    assert!(expected_messages[0]["id"].is_string());
    assert_eq!(
        expected_messages[0]["id"],
        requests[1].inputs_of_type("agent_message")[0]["id"],
    );
    let reviewer_thread_id =
        guardian_requests[0].body_json()["client_metadata"]["thread_id"].clone();
    assert!(reviewer_thread_id.is_string());
    for request in guardian_requests {
        assert_eq!(
            request.body_json()["client_metadata"]["thread_id"],
            reviewer_thread_id,
            "the failed review must reuse its Guardian session",
        );
        assert_eq!(request.inputs_of_type("agent_message"), expected_messages);
    }
    assert_eq!(
        reviews
            .iter()
            .map(|review| review.status)
            .collect::<Vec<_>>(),
        vec![
            GuardianAssessmentStatus::InProgress,
            GuardianAssessmentStatus::Approved,
        ]
    );
    assert_eq!(command_started_at_ms, Some(reviews[0].started_at_ms));
    assert_eq!(
        warnings.len(),
        1,
        "internal retries should not emit terminal warnings"
    );
    let contents = test
        .fs()
        .read_file_text(
            &test.workspace_path_uri("guardian-retry.txt")?,
            Default::default(),
            /*sandbox*/ None,
        )
        .await?;
    assert_eq!(contents.lines().collect::<Vec<_>>(), vec!["executed"]);
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn guardian_stale_approval_replaces_review_and_preserves_earlier_context() -> Result<()> {
    skip_if_no_network!(Ok(()));
    skip_if_wine_exec!(
        Ok(()),
        "Guardian approval actions require host-native paths"
    );
    let verdict = |id, rationale| {
        sse(vec![
            ev_assistant_message(
                id,
                &json!({
                    "risk_level": "low",
                    "user_authorization": "high",
                    "outcome": "allow",
                    "rationale": rationale,
                })
                .to_string(),
            ),
            ev_completed(id),
        ])
    };
    let command = |id, cmd| {
        sse(vec![
            ev_function_call(
                id,
                "exec_command",
                &json!({
                    "cmd": cmd,
                    "sandbox_permissions": "require_escalated",
                    "justification": "Write the requested marker once",
                })
                .to_string(),
            ),
            ev_completed(id),
        ])
    };
    let (release_review, review_gate) = tokio::sync::oneshot::channel();
    let (release_retry, retry_gate) = tokio::sync::oneshot::channel();
    let (model, _) = start_streaming_sse_server(
        vec![
            (
                None,
                command("earlier-command", "echo earlier >> earlier.txt"),
            ),
            (None, verdict("earlier-review", "Earlier valid approval")),
            (
                None,
                command("pending-command", "echo pending >> pending.txt"),
            ),
            (
                Some(review_gate),
                verdict("stale-review", "First invalidated approval"),
            ),
            (
                Some(retry_gate),
                verdict("stale-retry", "Second invalidated approval"),
            ),
            (
                None,
                verdict("replacement-review", "Decision based on updated guidance"),
            ),
            (None, sse(vec![ev_completed("parent-done")])),
        ]
        .into_iter()
        .map(|(gate, body)| vec![StreamingSseChunk { gate, body }])
        .collect(),
    )
    .await;
    let server = start_mock_server().await;
    let base_url = format!("{}/v1", model.uri());
    let test = test_codex()
        .with_config(move |config| {
            config.model_provider.base_url = Some(base_url);
            config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
            config.approvals_reviewer = ApprovalsReviewer::AutoReview;
            config
                .set_legacy_sandbox_policy(SandboxPolicy::new_workspace_write_policy())
                .expect("sandbox policy");
        })
        .build_with_auto_env(&server)
        .await?;
    test.codex
        .start_or_steer_turn(codex_core::TurnInputRequest::user_input(vec![
            codex_protocol::user_input::UserInput::Text {
                text: "Write the earlier marker, then the pending marker, once each.".to_owned(),
                text_elements: vec![],
            },
        ]))
        .await?;

    // Update authorization through the public goal API while each allow is held in flight.
    // The second invalidation must still preserve the checkpoint from the earlier action.
    let guidance = [
        "The pending marker should contain one line.",
        "Continue with the pending marker exactly once.",
    ];
    for (count, release, objective) in [
        (4, release_review, guidance[0]),
        (5, release_retry, guidance[1]),
    ] {
        tokio::time::timeout(Duration::from_secs(20), model.wait_for_request_count(count))
            .await
            .context("wait for pending review")?;
        test.codex
            .record_user_goal_update(UserGoalUpdate::Set {
                objective: Some(objective.to_owned()),
                status: None,
            })
            .await?;
        release.send(()).expect("release stale approval");
    }
    wait_for_event(&test.codex, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await;

    let requests = model
        .requests()
        .await
        .into_iter()
        .map(|body| serde_json::from_slice::<Value>(&body))
        .collect::<serde_json::Result<Vec<_>>>()?;
    let reviews = requests
        .iter()
        .filter(|request| request["client_metadata"]["x-openai-subagent"] == "guardian")
        .collect::<Vec<_>>();
    assert_eq!(reviews.len(), 4);
    let replacement = reviews[3]["input"].to_string();
    for expected in ["Earlier valid approval", guidance[0], guidance[1]] {
        assert!(
            replacement.contains(expected),
            "missing context: {expected}"
        );
    }
    for stale in ["First invalidated approval", "Second invalidated approval"] {
        assert!(
            !replacement.contains(stale),
            "stale review survived: {stale}"
        );
    }
    assert_eq!(
        replacement.matches(">>> APPROVAL REQUEST START").count(),
        2,
        "retain the earlier action and replace the pending action's stale review"
    );
    let contents = test
        .fs()
        .read_file_text(
            &test.workspace_path_uri("pending.txt")?,
            Default::default(),
            /*sandbox*/ None,
        )
        .await?;
    assert_eq!(contents.lines().collect::<Vec<_>>(), vec!["pending"]);
    test.codex.shutdown_and_wait().await?;
    model.shutdown().await;
    Ok(())
}