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

use anyhow::Result;
use codex_config::types::Personality;
use codex_core::TurnInputRequest;
use codex_history::InitialHistory;
use codex_history::ResumedHistory;
use codex_protocol::config_types::WindowsSandboxLevel;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ThreadHistoryMode;
use codex_protocol::protocol::ThreadSettingsOverrides;
use codex_protocol::user_input::UserInput;
use codex_thread_store::InMemoryThreadStore;
use codex_thread_store::LoadThreadHistoryParams;
use codex_thread_store::ThreadStoreResumePreparation;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::mount_sse_once;
use core_test_support::responses::mount_sse_sequence;
use core_test_support::responses::sse;
use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use std::sync::Arc;

/// A supplied preparation must select the prepared store path instead of silently falling back to
/// an ordinary resume.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_uses_prepared_store_path_when_preparation_is_present() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = start_mock_server().await;
    let initial = test_codex()
        .with_thread_store(Arc::new(InMemoryThreadStore::default()))
        .build_with_auto_env(&server)
        .await?;
    let thread_id = initial.session_configured.thread_id;
    initial.codex.shutdown_and_wait().await?;
    initial.thread_manager.remove_thread(&thread_id).await;
    let context = initial
        .thread_store
        .load_latest_model_context(LoadThreadHistoryParams {
            thread_id,
            include_archived: true,
        })
        .await?;

    let mut options = initial.start_thread_options().await;
    options.initial_history = InitialHistory::Resumed(ResumedHistory {
        history_revision: context.revision,
        conversation_id: context.thread_id,
        history: Arc::new(context.items),
        rollout_path: None,
    });
    options.resume_preparation = Some(ThreadStoreResumePreparation::new(()));

    let Err(err) = initial.thread_manager.start_thread(options).await else {
        panic!("prepared resume unexpectedly used the ordinary store path");
    };
    assert!(
        matches!(
            err.details(),
            CodexErrorDetails::UnsupportedOperation(message)
                if message == "resume_thread_with_preparation is not supported yet"
        ),
        "unexpected prepared resume error: {err}"
    );
    Ok(())
}

#[test_case::test_case(None, ThreadHistoryMode::Paginated; "store default")]
#[test_case::test_case(Some(ThreadHistoryMode::Legacy), ThreadHistoryMode::Legacy; "explicit legacy")]
#[test_case::test_case(Some(ThreadHistoryMode::Paginated), ThreadHistoryMode::Paginated; "explicit paginated")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn local_thread_history_mode_survives_restart(
    requested: Option<ThreadHistoryMode>,
    expected: ThreadHistoryMode,
) -> Result<()> {
    skip_if_no_network!(Ok(()));
    let server = start_mock_server().await;
    let mut builder = test_codex().with_history_mode(requested);
    let initial = builder.build_with_auto_env(&server).await?;
    assert_eq!(initial.codex.config_snapshot().await.history_mode, expected);
    mount_sse_once(
        &server,
        sse(vec![
            ev_assistant_message("seed", "Remembered answer"),
            ev_completed("seed"),
        ]),
    )
    .await;
    initial.submit_text_turn("Remember this task").await?;

    initial.codex.shutdown_and_wait().await?;
    let history = initial
        .thread_store
        .load_latest_model_context(LoadThreadHistoryParams {
            thread_id: initial.session_configured.thread_id,
            include_archived: true,
        })
        .await?;
    let resumed = initial
        .thread_manager
        .resume_thread_with_history(
            initial.config.clone(),
            InitialHistory::Resumed(ResumedHistory {
                history_revision: None,
                conversation_id: history.thread_id,
                history: Arc::new(history.items),
                rollout_path: initial.session_configured.rollout_path.clone(),
            }),
            initial.thread_manager.auth_manager(),
            /*parent_trace*/ None,
            ClientMcpExtensions::default(),
        )
        .await?;
    assert_eq!(
        resumed.thread.config_snapshot().await.history_mode,
        expected
    );
    let response = mount_sse_once(&server, sse(vec![ev_completed("followup")])).await;
    resumed
        .thread
        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
            text: "Continue".into(),
            text_elements: Vec::new(),
        }]))
        .await?;
    wait_for_event(&resumed.thread, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await;
    let request = response.single_request();
    assert!(request.body_contains_text("Remember this task"));
    assert!(request.body_contains_text("Remembered answer"));
    resumed.thread.shutdown_and_wait().await?;
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_restores_windows_sandbox_override() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = start_mock_server().await;
    let mut builder = test_codex();
    let initial = builder.build(&server).await?;
    core_test_support::submit_thread_settings(
        &initial.codex,
        ThreadSettingsOverrides {
            windows_sandbox_level: Some(WindowsSandboxLevel::Elevated),
            ..Default::default()
        },
    )
    .await?;
    initial.codex.ensure_rollout_materialized().await;
    let settings = initial.codex.restorable_thread_settings().await;

    let resumed = builder.restart(&server, &initial).await?;
    resumed.codex.restore_thread_settings(settings).await?;

    assert_eq!(
        resumed
            .codex
            .restorable_thread_settings()
            .await
            .windows_sandbox_level,
        Some(WindowsSandboxLevel::Elevated)
    );
    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_switches_models_preserves_base_instructions() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = start_mock_server().await;
    let legacy_instructions = "Legacy model instructions\n# Personality\nBe pragmatic.";
    let mut builder = test_codex()
        .with_model_info_override("gpt-5.2", move |model| {
            model
                .model_messages
                .get_or_insert_default()
                .instructions_template = Some(legacy_instructions.to_string());
        })
        .with_config(|config| {
            config.personality = Some(Personality::Pragmatic);
        });
    let initial = builder.build(&server).await?;
    let codex = Arc::clone(&initial.codex);

    let initial_sse = sse(vec![
        ev_response_created("resp-initial"),
        ev_assistant_message("msg-1", "Completed first turn"),
        ev_completed("resp-initial"),
    ]);
    let initial_mock = mount_sse_once(&server, initial_sse).await;

    codex
        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
            text: "Record initial instructions".into(),
            text_elements: Vec::new(),
        }]))
        .await?;
    wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await;
    let initial_instructions = initial_mock.single_request().instructions_text();
    assert_eq!(initial_instructions, legacy_instructions);

    let resumed_mock = mount_sse_sequence(
        &server,
        vec![
            sse(vec![
                ev_response_created("resp-resume-1"),
                ev_assistant_message("msg-2", "Resumed turn"),
                ev_completed("resp-resume-1"),
            ]),
            sse(vec![
                ev_response_created("resp-resume-2"),
                ev_assistant_message("msg-3", "Second resumed turn"),
                ev_completed("resp-resume-2"),
            ]),
        ],
    )
    .await;

    let mut resume_builder = test_codex().with_config(|config| {
        config.model = Some("gpt-5.4".to_string());
        config.personality = Some(Personality::None);
    });
    let resumed = resume_builder.restart(&server, &initial).await?;
    assert_eq!(
        resumed.session_configured.thread_id,
        initial.session_configured.thread_id
    );
    assert_eq!(
        resumed.codex.config_snapshot().await.personality,
        Some(Personality::None)
    );
    resumed
        .codex
        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
            text: "Resume with different model".into(),
            text_elements: Vec::new(),
        }]))
        .await?;
    wait_for_event(&resumed.codex, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await;

    resumed
        .codex
        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
            text: "Second turn after resume".into(),
            text_elements: Vec::new(),
        }]))
        .await?;
    wait_for_event(&resumed.codex, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await;

    let requests = resumed_mock.requests();
    assert_eq!(requests.len(), 2, "expected two resumed requests");

    let first_resumed = &requests[0];
    assert_eq!(first_resumed.instructions_text(), initial_instructions);
    assert!(
        first_resumed
            .message_input_texts("user")
            .iter()
            .any(|text| text == "Record initial instructions")
    );
    let first_developer_texts = first_resumed.message_input_texts("developer");
    let first_model_switch_count = first_developer_texts
        .iter()
        .filter(|text| text.contains("<model_switch>"))
        .count();
    assert!(
        first_model_switch_count >= 1,
        "expected model switch message on first post-resume turn"
    );

    let second_resumed = &requests[1];
    assert_eq!(second_resumed.instructions_text(), initial_instructions);
    let second_developer_texts = second_resumed.message_input_texts("developer");
    let second_model_switch_count = second_developer_texts
        .iter()
        .filter(|text| text.contains("<model_switch>"))
        .count();
    assert_eq!(
        second_model_switch_count, 1,
        "did not expect duplicate model switch message after first post-resume turn"
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_model_switch_is_not_duplicated_after_pre_turn_override() -> Result<()> {
    skip_if_no_network!(Ok(()));

    let server = start_mock_server().await;
    let mut builder = test_codex().with_config(|config| {
        config.model = Some("gpt-5.2".to_string());
    });
    let initial = builder.build(&server).await?;
    let codex = Arc::clone(&initial.codex);

    let initial_mock = mount_sse_once(
        &server,
        sse(vec![
            ev_response_created("resp-initial"),
            ev_assistant_message("msg-1", "Completed first turn"),
            ev_completed("resp-initial"),
        ]),
    )
    .await;
    codex
        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
            text: "Record initial instructions".into(),
            text_elements: Vec::new(),
        }]))
        .await?;
    wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await;
    let _ = initial_mock.single_request();

    let resumed_mock = mount_sse_once(
        &server,
        sse(vec![
            ev_response_created("resp-resume"),
            ev_assistant_message("msg-2", "Resumed turn"),
            ev_completed("resp-resume"),
        ]),
    )
    .await;

    let mut resume_builder = test_codex().with_config(|config| {
        config.model = Some("gpt-5.5".to_string());
    });
    let resumed = resume_builder.restart(&server, &initial).await?;
    core_test_support::submit_thread_settings(
        &resumed.codex,
        ThreadSettingsOverrides {
            model: Some("gpt-5.4".to_string()),
            ..Default::default()
        },
    )
    .await?;
    resumed
        .codex
        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
            text: "first turn after override".into(),
            text_elements: Vec::new(),
        }]))
        .await?;
    wait_for_event(&resumed.codex, |event| {
        matches!(event, EventMsg::TurnComplete(_))
    })
    .await;

    let request = resumed_mock.single_request();
    let developer_texts = request.message_input_texts("developer");
    let model_switch_count = developer_texts
        .iter()
        .filter(|text| text.contains("<model_switch>"))
        .count();
    assert_eq!(model_switch_count, 1);

    Ok(())
}