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"# ), )}