Codex Handbook
app-server/tests/suite/v2/thread_memory_mode_set.rs 138 lines
use anyhow::Result;use app_test_support::TestAppServer;use app_test_support::create_fake_rollout;use app_test_support::create_mock_responses_server_repeating_assistant;use app_test_support::to_response;use codex_app_server_protocol::JSONRPCResponse;use codex_app_server_protocol::RequestId;use codex_app_server_protocol::ThreadMemoryMode;use codex_app_server_protocol::ThreadMemoryModeSetParams;use codex_app_server_protocol::ThreadMemoryModeSetResponse;use codex_app_server_protocol::ThreadStartParams;use codex_app_server_protocol::ThreadStartResponse;use codex_protocol::ThreadId;use codex_state::StateRuntime;use pretty_assertions::assert_eq;use std::path::Path;use std::sync::Arc;use tempfile::TempDir;use tokio::time::timeout;const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);#[tokio::test]async fn thread_memory_mode_set_updates_loaded_thread_state() -> Result<()> {    let server = create_mock_responses_server_repeating_assistant("Done").await;    let codex_home = TempDir::new()?;    create_config_toml(codex_home.path(), &server.uri())?;    let state_db = init_state_db(codex_home.path()).await?;    let mut mcp = TestAppServer::new(codex_home.path()).await?;    timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;    let start_id = mcp        .send_thread_start_request(ThreadStartParams {            model: Some("mock-model".to_string()),            ..Default::default()        })        .await?;    let start_resp: JSONRPCResponse = timeout(        DEFAULT_READ_TIMEOUT,        mcp.read_stream_until_response_message(RequestId::Integer(start_id)),    )    .await??;    let ThreadStartResponse { thread, .. } = to_response::<ThreadStartResponse>(start_resp)?;    let thread_uuid = ThreadId::from_string(&thread.id)?;    let set_id = mcp        .send_thread_memory_mode_set_request(ThreadMemoryModeSetParams {            thread_id: thread.id,            mode: ThreadMemoryMode::Disabled,        })        .await?;    let set_resp: JSONRPCResponse = timeout(        DEFAULT_READ_TIMEOUT,        mcp.read_stream_until_response_message(RequestId::Integer(set_id)),    )    .await??;    let _: ThreadMemoryModeSetResponse = to_response::<ThreadMemoryModeSetResponse>(set_resp)?;    let memory_mode = state_db.get_thread_memory_mode(thread_uuid).await?;    assert_eq!(memory_mode.as_deref(), Some("disabled"));    Ok(())}#[tokio::test]async fn thread_memory_mode_set_updates_stored_thread_state() -> Result<()> {    let server = create_mock_responses_server_repeating_assistant("Done").await;    let codex_home = TempDir::new()?;    create_config_toml(codex_home.path(), &server.uri())?;    let state_db = init_state_db(codex_home.path()).await?;    let thread_id = create_fake_rollout(        codex_home.path(),        "2025-01-06T08-30-00",        "2025-01-06T08:30:00Z",        "Stored thread preview",        Some("mock_provider"),        /*git_info*/ None,    )?;    let thread_uuid = ThreadId::from_string(&thread_id)?;    let mut mcp = TestAppServer::new(codex_home.path()).await?;    timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;    for mode in [ThreadMemoryMode::Disabled, ThreadMemoryMode::Enabled] {        let set_id = mcp            .send_thread_memory_mode_set_request(ThreadMemoryModeSetParams {                thread_id: thread_id.clone(),                mode,            })            .await?;        let set_resp: JSONRPCResponse = timeout(            DEFAULT_READ_TIMEOUT,            mcp.read_stream_until_response_message(RequestId::Integer(set_id)),        )        .await??;        let _: ThreadMemoryModeSetResponse = to_response::<ThreadMemoryModeSetResponse>(set_resp)?;    }    let memory_mode = state_db.get_thread_memory_mode(thread_uuid).await?;    assert_eq!(memory_mode.as_deref(), Some("enabled"));    Ok(())}async fn init_state_db(codex_home: &Path) -> Result<Arc<StateRuntime>> {    let state_db = StateRuntime::init(codex_home.to_path_buf(), "mock_provider".into()).await?;    state_db        .mark_backfill_complete(/*last_watermark*/ None)        .await?;    Ok(state_db)}fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> {    let config_toml = codex_home.join("config.toml");    std::fs::write(        config_toml,        format!(            r#"model = "mock-model"approval_policy = "never"sandbox_mode = "read-only"model_provider = "mock_provider"suppress_unstable_features_warning = true[features]sqlite = true[model_providers.mock_provider]name = "Mock provider for test"base_url = "{server_uri}/v1"wire_api = "responses"request_max_retries = 0stream_max_retries = 0"#        ),    )}