codex/codex-rs/app-server/tests/suite/v2/cyber_access_program.rs

use anyhow::Result;
use app_test_support::ChatGptAuthFixture;
use app_test_support::TestAppServer;
use app_test_support::write_chatgpt_auth;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::CodexErrorInfo;
use codex_app_server_protocol::CyberAccessProgram;
use codex_app_server_protocol::ExperimentalFeatureEnablementSetParams;
use codex_app_server_protocol::ExperimentalFeatureEnablementSetResponse;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadForkParams;
use codex_app_server_protocol::ThreadForkResponse;
use codex_app_server_protocol::ThreadMetadataUpdateParams;
use codex_app_server_protocol::ThreadMetadataUpdateResponse;
use codex_app_server_protocol::ThreadReadParams;
use codex_app_server_protocol::ThreadReadResponse;
use codex_app_server_protocol::ThreadResumeParams;
use codex_app_server_protocol::ThreadResumeResponse;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartedNotification;
use codex_app_server_protocol::TurnError;
use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::TurnStatus;
use codex_app_server_protocol::UserInput;
use codex_login::AuthCredentialsStoreMode;
use codex_login::AuthKeyringBackendKind;
use codex_login::login_with_api_key;
use core_test_support::responses;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::path::Path;
use tempfile::TempDir;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;

#[tokio::test]
async fn thread_start_stages_daybreak_until_persistence() -> Result<()> {
    core_test_support::skip_if_no_network!(Ok(()));
    let server = responses::start_mock_server().await;
    responses::mount_sse_sequence(
        &server,
        (0..3)
            .map(|index| responses::sse(vec![responses::ev_completed(&format!("resp-{index}"))]))
            .collect(),
    )
    .await;
    let home = TempDir::new()?;
    let mut app = start_chatgpt_app(home.path(), &server).await?;
    let mut threads = Vec::new();
    for daybreak_enabled in [Some(true), Some(false), None] {
        let started = app
            .start_thread(ThreadStartParams {
                daybreak_enabled,
                ..Default::default()
            })
            .await?;
        assert_eq!(started.thread.daybreak_enabled, daybreak_enabled);
        let notification: ThreadStartedNotification =
            app.read_notification("thread/started").await?;
        assert_eq!(notification.thread, started.thread);
        let read: ThreadReadResponse = app
            .request(|request_id| ClientRequest::ThreadRead {
                request_id,
                params: ThreadReadParams {
                    thread_id: started.thread.id.clone(),
                    include_turns: false,
                },
            })
            .await?;
        assert_eq!(read.thread.daybreak_enabled, daybreak_enabled);
        let completed = app
            .start_turn_and_wait_for_completion(TurnStartParams {
                thread_id: started.thread.id.clone(),
                input: vec![UserInput::Text {
                    text: "start this task".to_owned(),
                    text_elements: Vec::new(),
                }],
                ..Default::default()
            })
            .await?;
        assert_eq!(completed.turn.status, TurnStatus::Completed);
        threads.push((started.thread.id, daybreak_enabled));
    }
    app.shutdown_gracefully().await?;

    let mut restarted = start_chatgpt_app(home.path(), &server).await?;
    for (thread_id, daybreak_enabled) in threads {
        let read: ThreadReadResponse = restarted
            .request(|request_id| ClientRequest::ThreadRead {
                request_id,
                params: ThreadReadParams {
                    thread_id: thread_id.clone(),
                    include_turns: false,
                },
            })
            .await?;
        assert_eq!(
            (read.thread.id, read.thread.daybreak_enabled),
            (thread_id, daybreak_enabled)
        );
    }
    Ok(())
}

#[tokio::test]
async fn thread_start_rejects_daybreak_for_ephemeral_threads() -> Result<()> {
    core_test_support::skip_if_no_network!(Ok(()));
    let server = responses::start_mock_server().await;
    let home = TempDir::new()?;
    let mut app = start_chatgpt_app(home.path(), &server).await?;
    let request_id = app
        .send_thread_start_request_with_auto_env(ThreadStartParams {
            daybreak_enabled: Some(true),
            ephemeral: Some(true),
            ..Default::default()
        })
        .await?;
    let error = app
        .read_stream_until_error_message(RequestId::Integer(request_id))
        .await?;
    assert_eq!(
        error.error,
        JSONRPCErrorError {
            code: -32600,
            message: "daybreakEnabled is not supported for ephemeral threads".to_owned(),
            data: None,
        }
    );
    Ok(())
}

#[tokio::test]
async fn turn_start_forwards_explicit_cyber_access_program() -> Result<()> {
    core_test_support::skip_if_no_network!(Ok(()));
    let server = responses::start_mock_server().await;
    let programs = [
        None,
        Some(CyberAccessProgram::DaybreakBlue),
        Some(CyberAccessProgram::DaybreakRed),
        None,
        Some(CyberAccessProgram::Standard),
        None,
    ];
    let requests = responses::mount_sse_sequence(
        &server,
        programs
            .iter()
            .enumerate()
            .map(|(index, _)| {
                responses::sse(vec![responses::ev_completed(&format!("resp-{index}"))])
            })
            .collect(),
    )
    .await;
    let home = TempDir::new()?;
    let mut app = start_chatgpt_app(home.path(), &server).await?;
    let thread = app
        .start_thread(ThreadStartParams {
            daybreak_enabled: Some(true),
            ..Default::default()
        })
        .await?
        .thread;
    for program in programs {
        let completed = app
            .start_turn_and_wait_for_completion(TurnStartParams {
                thread_id: thread.id.clone(),
                cyber_access_program: program,
                input: vec![UserInput::Text {
                    text: "hello".to_owned(),
                    text_elements: Vec::new(),
                }],
                ..Default::default()
            })
            .await?;
        assert_eq!(completed.turn.status, TurnStatus::Completed);
    }
    let requests = requests.requests();
    assert_eq!(
        requests
            .iter()
            .map(|request| request.body_json().get("access_programs").cloned())
            .collect::<Vec<_>>(),
        [
            None,
            Some(json!({"cyber": "daybreak_blue"})),
            Some(json!({"cyber": "daybreak_red"})),
            None,
            Some(json!({"cyber": "standard"})),
            None,
        ]
    );
    Ok(())
}

#[tokio::test]
async fn api_key_cyber_access_program_requires_only_forwarding_feature() -> Result<()> {
    core_test_support::skip_if_no_network!(Ok(()));
    let server = responses::start_mock_server().await;
    let expected_programs = [
        None,
        None,
        None,
        Some(json!({"cyber": "daybreak_blue"})),
        Some(json!({"cyber": "daybreak_red"})),
        Some(json!({"cyber": "standard"})),
        None,
        Some(json!({"cyber": "daybreak_blue"})),
        Some(json!({"cyber": "daybreak_red"})),
        Some(json!({"cyber": "standard"})),
    ];
    let requests = responses::mount_sse_sequence(
        &server,
        expected_programs
            .iter()
            .enumerate()
            .map(|(index, _)| responses::sse_completed(&format!("resp-{index}")))
            .collect(),
    )
    .await;
    let home = TempDir::new()?;
    let mut app = start_api_key_app(home.path(), &server).await?;
    set_api_key_cyber_access_programs(&mut app, /*enabled*/ true).await?;
    for (forwarding_enabled, discovery_enabled) in
        [(false, false), (false, true), (true, false), (true, true)]
    {
        let thread = app
            .start_thread(ThreadStartParams {
                config: Some(HashMap::from([
                    (
                        "features.api_key_cyber_access_programs".to_owned(),
                        json!(forwarding_enabled),
                    ),
                    (
                        "features.api_key_model_discovery".to_owned(),
                        json!(discovery_enabled),
                    ),
                ])),
                ..Default::default()
            })
            .await?
            .thread;
        for program in [
            None,
            Some(CyberAccessProgram::DaybreakBlue),
            Some(CyberAccessProgram::DaybreakRed),
            Some(CyberAccessProgram::Standard),
        ] {
            let completed = app
                .start_turn_and_wait_for_completion(TurnStartParams {
                    thread_id: thread.id.clone(),
                    cyber_access_program: program,
                    input: vec![UserInput::Text {
                        text: "hello".to_owned(),
                        text_elements: Vec::new(),
                    }],
                    ..Default::default()
                })
                .await?;
            if program.is_some() && !forwarding_enabled {
                assert_eq!(completed.turn.status, TurnStatus::Failed);
                assert_eq!(
                    completed.turn.error,
                    Some(TurnError {
                        message: "Cyber access programs are disabled for this API-key session."
                            .to_owned(),
                        codex_error_info: Some(CodexErrorInfo::Other),
                        additional_details: None,
                        misalignment: None,
                    })
                );
            } else {
                assert_eq!(completed.turn.status, TurnStatus::Completed);
            }
        }
    }
    assert_eq!(
        requests
            .requests()
            .iter()
            .map(|request| (
                request.header("authorization"),
                request.body_json()["model"].clone(),
                request.body_json().get("access_programs").cloned(),
            ))
            .collect::<Vec<_>>(),
        expected_programs
            .into_iter()
            .map(|program| (
                Some("Bearer test-key".to_owned()),
                json!("gpt-6-sol"),
                program
            ))
            .collect::<Vec<_>>()
    );
    app.shutdown_gracefully().await?;
    Ok(())
}

#[tokio::test]
async fn api_key_cyber_access_program_server_rejection_does_not_fall_back() -> Result<()> {
    core_test_support::skip_if_no_network!(Ok(()));
    let server = responses::start_mock_server().await;
    Mock::given(method("POST"))
        .and(path("/v1/responses"))
        .respond_with(ResponseTemplate::new(/*s*/ 403).set_body_json(json!({
            "error": {
                "type": "invalid_request_error",
                "code": "access_program_not_enabled",
                "message": "Daybreak Blue is not enabled for this project."
            }
        })))
        .expect(/*r*/ 1..)
        .mount(&server)
        .await;
    let home = TempDir::new()?;
    let mut app = start_api_key_app(home.path(), &server).await?;
    set_api_key_cyber_access_programs(&mut app, /*enabled*/ true).await?;
    let thread = app.start_thread(ThreadStartParams::default()).await?.thread;
    let completed = app
        .start_turn_and_wait_for_completion(TurnStartParams {
            thread_id: thread.id,
            cyber_access_program: Some(CyberAccessProgram::DaybreakBlue),
            input: vec![UserInput::Text {
                text: "hello".to_owned(),
                text_elements: Vec::new(),
            }],
            ..Default::default()
        })
        .await?;
    assert_eq!(completed.turn.status, TurnStatus::Failed);
    let error = completed.turn.error.expect("server rejection");
    assert_eq!(
        error.codex_error_info,
        Some(CodexErrorInfo::HttpConnectionFailed {
            http_status_code: Some(403),
        })
    );
    assert!(
        error
            .message
            .contains("Daybreak Blue is not enabled for this project.")
    );
    let requests = responses::received_responses_requests(&server).await;
    assert!(!requests.is_empty());
    // Existing transport retries may repeat the request; none may drop the program.
    assert_eq!(
        requests
            .iter()
            .map(|request| (
                request.header("authorization"),
                request.body_json()["access_programs"].clone()
            ))
            .collect::<Vec<_>>(),
        vec![
            (
                Some("Bearer test-key".to_owned()),
                json!({"cyber": "daybreak_blue"})
            );
            requests.len()
        ]
    );
    app.shutdown_gracefully().await?;
    Ok(())
}

#[tokio::test]
async fn daybreak_thread_metadata_persists_independently_across_restart_and_fork() -> Result<()> {
    core_test_support::skip_if_no_network!(Ok(()));
    let server = responses::start_mock_server().await;
    let requests = responses::mount_sse_sequence(
        &server,
        (0..4)
            .map(|index| responses::sse(vec![responses::ev_completed(&format!("resp-{index}"))]))
            .collect(),
    )
    .await;
    let home = TempDir::new()?;
    let mut app = start_chatgpt_app(home.path(), &server).await?;
    let first = app
        .start_thread(ThreadStartParams {
            daybreak_enabled: Some(false),
            ..Default::default()
        })
        .await?;
    let second = app
        .start_thread(ThreadStartParams {
            daybreak_enabled: Some(true),
            ..Default::default()
        })
        .await?;
    assert_eq!(
        (
            first.thread.daybreak_enabled,
            second.thread.daybreak_enabled
        ),
        (Some(false), Some(true))
    );

    for (thread_id, enabled) in [(&first.thread.id, true), (&second.thread.id, false)] {
        let updated: ThreadMetadataUpdateResponse = app
            .request(|request_id| ClientRequest::ThreadMetadataUpdate {
                request_id,
                params: ThreadMetadataUpdateParams {
                    thread_id: thread_id.clone(),
                    project_id: None,
                    git_info: None,
                    daybreak_enabled: Some(enabled),
                },
            })
            .await?;
        assert_eq!(
            (updated.thread.id, updated.thread.daybreak_enabled),
            (thread_id.clone(), Some(enabled))
        );
    }
    // Leave the second thread empty to exercise metadata durability without a rollout.
    let completed = app
        .start_turn_and_wait_for_completion(TurnStartParams {
            thread_id: first.thread.id.clone(),
            input: vec![UserInput::Text {
                text: "start this task".to_owned(),
                text_elements: Vec::new(),
            }],
            ..Default::default()
        })
        .await?;
    assert_eq!(completed.turn.status, TurnStatus::Completed);
    assert_eq!(requests.requests().len(), 1);
    app.shutdown_gracefully().await?;

    let mut restarted = start_chatgpt_app(home.path(), &server).await?;
    for (thread_id, enabled) in [(&first.thread.id, true), (&second.thread.id, false)] {
        let read: ThreadReadResponse = restarted
            .request(|request_id| ClientRequest::ThreadRead {
                request_id,
                params: ThreadReadParams {
                    thread_id: thread_id.clone(),
                    include_turns: false,
                },
            })
            .await?;
        assert_eq!(
            (read.thread.id, read.thread.daybreak_enabled),
            (thread_id.clone(), Some(enabled))
        );
    }
    let resumed: ThreadResumeResponse = restarted
        .request(|request_id| ClientRequest::ThreadResume {
            request_id,
            params: ThreadResumeParams {
                thread_id: first.thread.id.clone(),
                ..Default::default()
            },
        })
        .await?;
    assert_eq!(resumed.thread.daybreak_enabled, Some(true));
    let forked: ThreadForkResponse = restarted
        .request(|request_id| ClientRequest::ThreadFork {
            request_id,
            params: ThreadForkParams {
                thread_id: first.thread.id.clone(),
                ..Default::default()
            },
        })
        .await?;
    assert_eq!(forked.thread.daybreak_enabled, Some(true));
    let updated: ThreadMetadataUpdateResponse = restarted
        .request(|request_id| ClientRequest::ThreadMetadataUpdate {
            request_id,
            params: ThreadMetadataUpdateParams {
                thread_id: first.thread.id.clone(),
                project_id: None,
                git_info: None,
                daybreak_enabled: Some(false),
            },
        })
        .await?;
    assert_eq!(
        (updated.thread.id, updated.thread.daybreak_enabled),
        (first.thread.id.clone(), Some(false))
    );
    assert_eq!(requests.requests().len(), 1);
    restarted.shutdown_gracefully().await?;

    let mut restarted = start_chatgpt_app(home.path(), &server).await?;
    for (thread_id, program, enabled) in [
        (&forked.thread.id, Some(CyberAccessProgram::Standard), true),
        (&forked.thread.id, None, true),
        (&first.thread.id, None, false),
    ] {
        let resumed: ThreadResumeResponse = restarted
            .request(|request_id| ClientRequest::ThreadResume {
                request_id,
                params: ThreadResumeParams {
                    thread_id: thread_id.clone(),
                    ..Default::default()
                },
            })
            .await?;
        assert_eq!(resumed.thread.daybreak_enabled, Some(enabled));
        let completed = restarted
            .start_turn_and_wait_for_completion(TurnStartParams {
                thread_id: thread_id.clone(),
                cyber_access_program: program,
                input: vec![UserInput::Text {
                    text: "hello".to_owned(),
                    text_elements: Vec::new(),
                }],
                ..Default::default()
            })
            .await?;
        assert_eq!(completed.turn.status, TurnStatus::Completed);
        let resumed: ThreadResumeResponse = restarted
            .request(|request_id| ClientRequest::ThreadResume {
                request_id,
                params: ThreadResumeParams {
                    thread_id: thread_id.clone(),
                    ..Default::default()
                },
            })
            .await?;
        assert_eq!(resumed.thread.daybreak_enabled, Some(enabled));
    }
    assert_eq!(
        requests
            .requests()
            .iter()
            .map(|request| request.body_json()["access_programs"].clone())
            .collect::<Vec<_>>(),
        [
            json!(null),
            json!({"cyber": "standard"}),
            json!(null),
            json!(null),
        ]
    );
    Ok(())
}

#[tokio::test]
async fn turn_start_forwards_cyber_access_program_with_personal_access_token() -> Result<()> {
    core_test_support::skip_if_no_network!(Ok(()));
    let server = responses::start_mock_server().await;
    Mock::given(method("GET"))
        .and(path("/v1/responses"))
        .respond_with(ResponseTemplate::new(426))
        .mount(&server)
        .await;
    Mock::given(method("GET"))
        .and(path("/v1/user-auth-credential/whoami"))
        .and(header("Authorization", "Bearer at-test-token"))
        .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
            "email": null,
            "chatgpt_user_id": "user-123",
            "chatgpt_account_id": "account-123",
            "chatgpt_plan_type": "enterprise_cbp_automation",
            "chatgpt_account_is_fedramp": false,
        })))
        .mount(&server)
        .await;
    Mock::given(method("GET"))
        .and(path("/backend-api/wham/config/bundle"))
        .and(header("Authorization", "Bearer at-test-token"))
        .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({})))
        .mount(&server)
        .await;
    let request = responses::mount_sse_once(
        &server,
        responses::sse(vec![responses::ev_completed("resp-1")]),
    )
    .await;
    let home = TempDir::new()?;
    std::fs::write(
        home.path().join("config.toml"),
        format!(
            "model = \"gpt-5.5\"\napproval_policy = \"never\"\nopenai_base_url = \"{0}/v1\"\nchatgpt_base_url = \"{0}/backend-api\"\n",
            server.uri(),
        ),
    )?;
    let authapi_base_url = server.uri();
    let mut app = TestAppServer::builder()
        .with_codex_home(home.path())
        .without_managed_config()
        .with_env_overrides(&[
            ("OPENAI_API_KEY", None),
            ("CODEX_ACCESS_TOKEN", Some("at-test-token")),
            ("CODEX_AUTHAPI_BASE_URL", Some(authapi_base_url.as_str())),
        ])
        .build_initialized()
        .await?;
    let thread = app.start_thread(ThreadStartParams::default()).await?.thread;

    app.start_turn_and_wait_for_completion(TurnStartParams {
        thread_id: thread.id,
        cyber_access_program: Some(CyberAccessProgram::DaybreakBlue),
        input: vec![UserInput::Text {
            text: "hello".to_owned(),
            text_elements: Vec::new(),
        }],
        ..Default::default()
    })
    .await?;

    assert_eq!(
        request.single_request().body_json()["access_programs"],
        json!({"cyber": "daybreak_blue"})
    );
    Ok(())
}

async fn start_chatgpt_app(home: &Path, server: &MockServer) -> Result<TestAppServer> {
    Mock::given(method("GET"))
        .and(path("/v1/responses"))
        .respond_with(ResponseTemplate::new(/*s*/ 426))
        .mount(server)
        .await;
    std::fs::write(
        home.join("config.toml"),
        format!(
            "model = \"gpt-5.5\"\napproval_policy = \"never\"\nopenai_base_url = \"{}/v1\"\ncli_auth_credentials_store = \"file\"\n",
            server.uri()
        ),
    )?;
    write_chatgpt_auth(
        home,
        ChatGptAuthFixture::new("chatgpt-test-token").plan_type("pro"),
        AuthCredentialsStoreMode::File,
    )?;
    TestAppServer::builder()
        .with_codex_home(home)
        .without_managed_config()
        .build_initialized()
        .await
}

async fn start_api_key_app(home: &Path, server: &MockServer) -> Result<TestAppServer> {
    Mock::given(method("GET"))
        .and(path("/v1/responses"))
        .respond_with(ResponseTemplate::new(/*s*/ 426))
        .mount(server)
        .await;
    // The existing inference override disables remote catalog discovery. Forwarding
    // must therefore work without granting programs through a test catalog.
    std::fs::write(
        home.join("config.toml"),
        format!(
            "model = \"gpt-6-sol\"\nopenai_base_url = \"{}/v1\"\n[features]\napi_key_model_discovery = false\n",
            server.uri()
        ),
    )?;
    login_with_api_key(
        home,
        "test-key",
        AuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
    )?;
    TestAppServer::builder()
        .with_codex_home(home)
        .without_managed_config()
        .with_env_overrides(&[
            ("OPENAI_API_KEY", None),
            ("CODEX_API_KEY", None),
            ("OPENAI_BASE_URL", None),
        ])
        .build_initialized()
        .await
}

async fn set_api_key_cyber_access_programs(app: &mut TestAppServer, enabled: bool) -> Result<()> {
    let enablement = BTreeMap::from([("api_key_cyber_access_programs".to_owned(), enabled)]);
    let response: ExperimentalFeatureEnablementSetResponse = app
        .request(
            |request_id| ClientRequest::ExperimentalFeatureEnablementSet {
                request_id,
                params: ExperimentalFeatureEnablementSetParams {
                    enablement: enablement.clone(),
                },
            },
        )
        .await?;
    assert_eq!(
        response,
        ExperimentalFeatureEnablementSetResponse { enablement }
    );
    Ok(())
}