diff --git a/codex-rs/app-server/tests/common/config.rs b/codex-rs/app-server/tests/common/config.rs index f3291d95e4c7..f359b40f89c9 100644 --- a/codex-rs/app-server/tests/common/config.rs +++ b/codex-rs/app-server/tests/common/config.rs @@ -3,6 +3,159 @@ use codex_features::Feature; use std::collections::BTreeMap; use std::path::Path; +/// Composes the standard mock Responses provider with test-specific configuration. +pub struct MockResponsesConfig { + provider_id: String, + provider_name: String, + provider_base_url: String, + model: String, + approval_policy: String, + sandbox_mode: String, + features: BTreeMap, + root_config: Vec, + provider_config: Vec, + extra_config: Vec, +} + +impl MockResponsesConfig { + pub fn new(server_uri: &str) -> Self { + Self { + provider_id: "mock_provider".to_string(), + provider_name: "Mock provider for test".to_string(), + provider_base_url: format!("{server_uri}/v1"), + model: "mock-model".to_string(), + approval_policy: "never".to_string(), + sandbox_mode: "read-only".to_string(), + features: BTreeMap::new(), + root_config: Vec::new(), + provider_config: Vec::new(), + extra_config: Vec::new(), + } + } + + pub fn with_model_provider(mut self, provider_id: &str) -> Self { + self.provider_id = provider_id.to_string(); + self + } + + pub fn with_provider_name(mut self, provider_name: &str) -> Self { + self.provider_name = provider_name.to_string(); + self + } + + pub fn with_provider_base_url(mut self, provider_base_url: &str) -> Self { + self.provider_base_url = provider_base_url.to_string(); + self + } + + pub fn with_model(mut self, model: &str) -> Self { + self.model = model.to_string(); + self + } + + pub fn with_approval_policy(mut self, approval_policy: &str) -> Self { + self.approval_policy = approval_policy.to_string(); + self + } + + pub fn with_sandbox_mode(mut self, sandbox_mode: &str) -> Self { + self.sandbox_mode = sandbox_mode.to_string(); + self + } + + pub fn enable_feature(mut self, feature: Feature) -> Self { + self.features.insert(feature, true); + self + } + + pub fn disable_feature(mut self, feature: Feature) -> Self { + self.features.insert(feature, false); + self + } + + pub fn with_features(mut self, features: &BTreeMap) -> Self { + self.features.extend( + features + .iter() + .map(|(&feature, &enabled)| (feature, enabled)), + ); + self + } + + pub fn with_root_config(mut self, config: &str) -> Self { + self.root_config.push(config.to_string()); + self + } + + pub fn with_provider_config(mut self, config: &str) -> Self { + self.provider_config.push(config.to_string()); + self + } + + pub fn with_extra_config(mut self, config: &str) -> Self { + self.extra_config.push(config.to_string()); + self + } + + pub fn write(self, codex_home: &Path) -> std::io::Result<()> { + let Self { + provider_id, + provider_name, + provider_base_url, + model, + approval_policy, + sandbox_mode, + features, + root_config, + provider_config, + extra_config, + } = self; + let root_config = root_config.join("\n"); + let provider_config = provider_config.join("\n"); + let extra_config = extra_config.join("\n"); + let feature_entries = features + .into_iter() + .map(|(feature, enabled)| { + let key = FEATURES + .iter() + .find(|spec| spec.id == feature) + .map(|spec| spec.key) + .expect("feature should have a config key"); + format!("{key} = {enabled}") + }) + .collect::>() + .join("\n"); + let feature_config = if feature_entries.is_empty() { + String::new() + } else { + format!("[features]\n{feature_entries}\n\n") + }; + + std::fs::write( + codex_home.join("config.toml"), + format!( + r#" +model = "{model}" +approval_policy = "{approval_policy}" +sandbox_mode = "{sandbox_mode}" +{root_config} +model_provider = "{provider_id}" + +{feature_config}[model_providers.{provider_id}] +name = "{provider_name}" +base_url = "{provider_base_url}" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +{provider_config} + +{extra_config} +"# + ), + ) + } +} + pub fn write_mock_responses_config_toml( codex_home: &Path, server_uri: &str, @@ -12,71 +165,24 @@ pub fn write_mock_responses_config_toml( model_provider_id: &str, compact_prompt: &str, ) -> std::io::Result<()> { - // Phase 1: build the features block for config.toml. - let mut features = BTreeMap::new(); - for (feature, enabled) in feature_flags { - features.insert(*feature, *enabled); - } - let feature_entries = features - .into_iter() - .map(|(feature, enabled)| { - let key = FEATURES - .iter() - .find(|spec| spec.id == feature) - .map(|spec| spec.key) - .expect("feature should have a config key"); - format!("{key} = {enabled}") - }) - .collect::>() - .join("\n"); - // Phase 2: build provider-specific config bits. - let requires_line = match requires_openai_auth { - Some(true) => "requires_openai_auth = true\n".to_string(), - Some(false) | None => String::new(), - }; - let provider_name = if matches!(requires_openai_auth, Some(true)) { - "OpenAI" - } else { - "Mock provider for test" - }; - let provider_block = format!( - r#" -[model_providers.{model_provider_id}] -name = "{provider_name}" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -supports_websockets = false -{requires_line} -"# - ); - let openai_base_url_line = if model_provider_id == "openai" { - format!("openai_base_url = \"{server_uri}/v1\"\n") - } else { - String::new() - }; - // Phase 3: write the final config file. - 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" -compact_prompt = "{compact_prompt}" -model_auto_compact_token_limit = {auto_compact_limit} - -model_provider = "{model_provider_id}" -{openai_base_url_line} - -[features] -{feature_entries} -{provider_block} -"# - ), - ) + let mut config = MockResponsesConfig::new(server_uri) + .with_model_provider(model_provider_id) + .with_features(feature_flags) + .with_root_config(&format!( + "compact_prompt = \"{compact_prompt}\"\nmodel_auto_compact_token_limit = {auto_compact_limit}" + )) + .with_provider_config("supports_websockets = false"); + + if model_provider_id == "openai" { + config = config.with_root_config(&format!("openai_base_url = \"{server_uri}/v1\"")); + } + if matches!(requires_openai_auth, Some(true)) { + config = config + .with_provider_name("OpenAI") + .with_provider_config("requires_openai_auth = true"); + } + + config.write(codex_home) } pub fn write_mock_responses_config_toml_with_chatgpt_base_url( @@ -84,25 +190,11 @@ pub fn write_mock_responses_config_toml_with_chatgpt_base_url( server_uri: &str, chatgpt_base_url: &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" -chatgpt_base_url = "{chatgpt_base_url}" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) + MockResponsesConfig::new(server_uri) + .with_root_config(&format!("chatgpt_base_url = \"{chatgpt_base_url}\"")) + .write(codex_home) } + +#[cfg(test)] +#[path = "config_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/tests/common/config_tests.rs b/codex-rs/app-server/tests/common/config_tests.rs new file mode 100644 index 000000000000..23cfec38e827 --- /dev/null +++ b/codex-rs/app-server/tests/common/config_tests.rs @@ -0,0 +1,72 @@ +use super::*; +use tempfile::TempDir; + +#[test] +fn mock_responses_config_composes_model_provider_features_and_extra_tables() { + let home = TempDir::new().expect("temporary CODEX_HOME"); + MockResponsesConfig::new("http://127.0.0.1:1234") + .with_model("custom-model") + .with_model_provider("openai-custom") + .with_provider_name("OpenAI") + .with_provider_base_url("http://127.0.0.1:1234/api/codex") + .with_approval_policy("on-request") + .with_sandbox_mode("workspace-write") + .enable_feature(Feature::Personality) + .disable_feature(Feature::ShellSnapshot) + .with_root_config("chatgpt_base_url = \"http://127.0.0.1:1234\"") + .with_provider_config("requires_openai_auth = true") + .with_extra_config("[extra]\nenabled = true") + .write(home.path()) + .expect("write composable mock Responses config"); + + let config = + std::fs::read_to_string(home.path().join("config.toml")).expect("read config.toml"); + for expected in [ + "model = \"custom-model\"", + "approval_policy = \"on-request\"", + "sandbox_mode = \"workspace-write\"", + "chatgpt_base_url = \"http://127.0.0.1:1234\"", + "model_provider = \"openai-custom\"", + "shell_snapshot = false", + "personality = true", + "[model_providers.openai-custom]\nname = \"OpenAI\"", + "base_url = \"http://127.0.0.1:1234/api/codex\"", + "requires_openai_auth = true", + "[extra]\nenabled = true", + ] { + assert!(config.contains(expected), "config is missing {expected}"); + } +} + +#[test] +fn legacy_mock_responses_writer_preserves_provider_auth_and_feature_overrides() { + let home = TempDir::new().expect("temporary CODEX_HOME"); + write_mock_responses_config_toml( + home.path(), + "http://127.0.0.1:1234", + &BTreeMap::from([ + (Feature::Personality, true), + (Feature::ShellSnapshot, false), + ]), + /*auto_compact_limit*/ 321, + Some(true), + "openai", + "compact this", + ) + .expect("write legacy-compatible mock Responses config"); + + let config = + std::fs::read_to_string(home.path().join("config.toml")).expect("read config.toml"); + for expected in [ + "compact_prompt = \"compact this\"", + "model_auto_compact_token_limit = 321", + "openai_base_url = \"http://127.0.0.1:1234/v1\"", + "shell_snapshot = false", + "personality = true", + "[model_providers.openai]\nname = \"OpenAI\"", + "supports_websockets = false", + "requires_openai_auth = true", + ] { + assert!(config.contains(expected), "config is missing {expected}"); + } +} diff --git a/codex-rs/app-server/tests/common/lib.rs b/codex-rs/app-server/tests/common/lib.rs index e065e77a9e12..1cb946128cbc 100644 --- a/codex-rs/app-server/tests/common/lib.rs +++ b/codex-rs/app-server/tests/common/lib.rs @@ -18,6 +18,7 @@ pub use auth_fixtures::ChatGptIdTokenClaims; pub use auth_fixtures::encode_id_token; pub use auth_fixtures::write_chatgpt_auth; use codex_app_server_protocol::JSONRPCResponse; +pub use config::MockResponsesConfig; pub use config::write_mock_responses_config_toml; pub use config::write_mock_responses_config_toml_with_chatgpt_base_url; pub use core_test_support::PathBufExt; diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index 26055f6aecb3..062d3a1e6ea6 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -21,6 +21,7 @@ use codex_app_server_protocol::AppsReadParams; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientNotification; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::CollaborationModeListParams; use codex_app_server_protocol::CommandExecParams; use codex_app_server_protocol::CommandExecResizeParams; @@ -106,6 +107,7 @@ use codex_app_server_protocol::ThreadSetNameParams; use codex_app_server_protocol::ThreadSettingsUpdateParams; use codex_app_server_protocol::ThreadShellCommandParams; use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::ThreadTurnsListParams; use codex_app_server_protocol::ThreadUnarchiveParams; use codex_app_server_protocol::ThreadUnsubscribeParams; @@ -125,6 +127,7 @@ use codex_login::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR; use core_test_support::is_remote_test_environment; use core_test_support::test_codex::TestEnv; use core_test_support::test_codex::test_env; +use serde::de::DeserializeOwned; use tempfile::TempDir; use tokio::process::Command; @@ -154,6 +157,10 @@ pub const DEFAULT_CLIENT_NAME: &str = "codex-app-server-tests"; pub const DISABLE_PLUGIN_STARTUP_TASKS_ARG: &str = "--disable-plugin-startup-tasks-for-tests"; const DISABLE_MANAGED_CONFIG_ENV_VAR: &str = "CODEX_APP_SERVER_DISABLE_MANAGED_CONFIG"; const CODE_MODE_HOST_PATH_ENV_VAR: &str = "CODEX_CODE_MODE_HOST_PATH"; +#[cfg(windows)] +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(25); +#[cfg(not(windows))] +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); impl TestAppServer { /// Starts building a server with a temporary CODEX_HOME and the standard @@ -461,6 +468,15 @@ impl TestAppServer { self.send_thread_start_request(params).await } + /// Starts a thread using the standard automatic test environment. + pub async fn start_thread( + &mut self, + params: ThreadStartParams, + ) -> anyhow::Result { + let request_id = self.send_thread_start_request_with_auto_env(params).await?; + tokio::time::timeout(DEFAULT_REQUEST_TIMEOUT, self.read_response(request_id)).await? + } + /// Send a `thread/resume` JSON-RPC request. pub async fn send_thread_resume_request( &mut self, @@ -1445,6 +1461,26 @@ impl TestAppServer { .await } + /// Sends a typed protocol request and waits for its deserialized response. + /// + /// The request builder receives a fresh ID so tests do not need to manage + /// the JSON-RPC request ID themselves. + pub async fn request( + &mut self, + make_request: impl FnOnce(RequestId) -> ClientRequest, + ) -> anyhow::Result { + let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed); + let request = make_request(RequestId::Integer(request_id)); + ensure!( + request.id() == &RequestId::Integer(request_id), + "typed request must use the supplied request ID" + ); + let request = serde_json::from_value::(serde_json::to_value(request)?)?; + self.send_jsonrpc_message(JSONRPCMessage::Request(request)) + .await?; + tokio::time::timeout(DEFAULT_REQUEST_TIMEOUT, self.read_response(request_id)).await? + } + async fn send_request( &mut self, method: &str, @@ -1549,6 +1585,21 @@ impl TestAppServer { Ok(response) } + /// Reads and deserializes the successful response for an integer request ID. + /// + /// This does not impose a timeout, so callers can retain suite-specific + /// timeout policies when requests need different latency budgets. + pub async fn read_response( + &mut self, + request_id: i64, + ) -> anyhow::Result { + let response = self + .read_stream_until_response_message(RequestId::Integer(request_id)) + .await?; + serde_json::from_value(response.result) + .with_context(|| format!("failed to deserialize response for request {request_id}")) + } + pub async fn read_stream_until_error_message( &mut self, request_id: RequestId, @@ -1586,6 +1637,22 @@ impl TestAppServer { Ok(notification) } + /// Reads and deserializes the parameters of the next matching notification. + /// + /// This does not impose a timeout, so callers can retain suite-specific + /// timeout policies when notifications need different latency budgets. + pub async fn read_notification( + &mut self, + method: &str, + ) -> anyhow::Result { + let notification = self.read_stream_until_notification_message(method).await?; + let params = notification + .params + .with_context(|| format!("notification `{method}` is missing parameters"))?; + serde_json::from_value(params) + .with_context(|| format!("failed to deserialize notification `{method}`")) + } + pub async fn read_stream_until_matching_notification( &mut self, description: &str, @@ -1775,6 +1842,22 @@ impl TestAppServerBuilder { self } + /// Builds a server and completes its standard initialization handshake. + pub async fn build_initialized(self) -> anyhow::Result { + self.build_initialized_with_timeout(DEFAULT_REQUEST_TIMEOUT) + .await + } + + /// Builds and initializes a server while preserving a suite-specific timeout. + pub async fn build_initialized_with_timeout( + self, + timeout: Duration, + ) -> anyhow::Result { + let mut server = self.build().await?; + tokio::time::timeout(timeout, server.initialize()).await??; + Ok(server) + } + /// Builds a server with a temporary CODEX_HOME and automatic environment /// by default. pub async fn build(self) -> anyhow::Result { diff --git a/codex-rs/app-server/tests/suite/auth.rs b/codex-rs/app-server/tests/suite/auth.rs index 4869748070e2..de40a45c5123 100644 --- a/codex-rs/app-server/tests/suite/auth.rs +++ b/codex-rs/app-server/tests/suite/auth.rs @@ -1,5 +1,6 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; @@ -16,6 +17,7 @@ use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; +use codex_features::Feature; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; use codex_protocol::account::PlanType as AccountPlanType; use pretty_assertions::assert_eq; @@ -37,33 +39,13 @@ fn create_config_toml_custom_provider( codex_home: &Path, requires_openai_auth: bool, ) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - let requires_line = if requires_openai_auth { - "requires_openai_auth = true\n" - } else { - "" - }; - let contents = format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "danger-full-access" - -model_provider = "mock_provider" - -[features] -shell_snapshot = false - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "http://127.0.0.1:0/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -{requires_line} -"# - ); - std::fs::write(config_toml, contents) + let mut config = MockResponsesConfig::new("http://127.0.0.1:0") + .with_sandbox_mode("danger-full-access") + .disable_feature(Feature::ShellSnapshot); + if requires_openai_auth { + config = config.with_provider_config("requires_openai_auth = true"); + } + config.write(codex_home) } fn create_config_toml(codex_home: &Path) -> std::io::Result<()> { @@ -100,12 +82,8 @@ shell_snapshot = false async fn login_with_api_key_via_request(mcp: &mut TestAppServer, api_key: &str) -> Result<()> { let request_id = mcp.send_login_account_api_key_request(api_key).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response, LoginAccountResponse::ApiKey {}); Ok(()) } @@ -119,9 +97,8 @@ async fn get_auth_status_no_auth() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_get_auth_status_request(GetAuthStatusParams { @@ -130,12 +107,8 @@ async fn get_auth_status_no_auth() -> Result<()> { }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(status.auth_method, None, "expected no auth method"); assert_eq!(status.auth_token, None, "expected no token"); Ok(()) @@ -149,9 +122,8 @@ async fn get_auth_status_with_api_key() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; login_with_api_key_via_request(&mut mcp, "sk-test-key").await?; @@ -162,12 +134,8 @@ async fn get_auth_status_with_api_key() -> Result<()> { }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(status.auth_method, Some(AuthMode::ApiKey)); assert_eq!(status.auth_token, Some("sk-test-key".to_string())); Ok(()) @@ -202,9 +170,8 @@ async fn personal_access_token_without_email_supports_auth_status_and_account_re ("CODEX_ACCESS_TOKEN", Some("at-test-token")), ("CODEX_AUTHAPI_BASE_URL", Some(authapi_base_url.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_get_auth_status_request(GetAuthStatusParams { @@ -213,12 +180,8 @@ async fn personal_access_token_without_email_supports_auth_status_and_account_re }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( status, GetAuthStatusResponse { @@ -268,9 +231,8 @@ async fn get_auth_status_with_api_key_when_auth_not_required() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; login_with_api_key_via_request(&mut mcp, "sk-test-key").await?; @@ -281,12 +243,8 @@ async fn get_auth_status_with_api_key_when_auth_not_required() -> Result<()> { }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(status.auth_method, None, "expected no auth method"); assert_eq!(status.auth_token, None, "expected no token"); assert_eq!( @@ -305,9 +263,8 @@ async fn get_auth_status_with_api_key_no_include_token() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; login_with_api_key_via_request(&mut mcp, "sk-test-key").await?; @@ -318,12 +275,8 @@ async fn get_auth_status_with_api_key_no_include_token() -> Result<()> { }; let request_id = mcp.send_get_auth_status_request(params).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(status.auth_method, Some(AuthMode::ApiKey)); assert!(status.auth_token.is_none(), "token must be omitted"); Ok(()) @@ -337,9 +290,8 @@ async fn get_auth_status_with_api_key_refresh_requested() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; login_with_api_key_via_request(&mut mcp, "sk-test-key").await?; @@ -350,12 +302,8 @@ async fn get_auth_status_with_api_key_refresh_requested() -> Result<()> { }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( status, GetAuthStatusResponse { @@ -404,9 +352,8 @@ async fn get_auth_status_omits_token_after_permanent_refresh_failure() -> Result Some(refresh_url.as_str()), ), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_get_auth_status_request(GetAuthStatusParams { @@ -415,12 +362,8 @@ async fn get_auth_status_omits_token_after_permanent_refresh_failure() -> Result }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( status, GetAuthStatusResponse { @@ -437,12 +380,8 @@ async fn get_auth_status_omits_token_after_permanent_refresh_failure() -> Result }) .await?; - let second_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_request_id)), - ) - .await??; - let second_status: GetAuthStatusResponse = to_response(second_resp)?; + let second_status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(second_request_id)).await??; assert_eq!(second_status, status); server.verify().await; @@ -487,9 +426,8 @@ async fn get_auth_status_omits_token_after_proactive_refresh_failure() -> Result Some(refresh_url.as_str()), ), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_get_auth_status_request(GetAuthStatusParams { @@ -498,12 +436,8 @@ async fn get_auth_status_omits_token_after_proactive_refresh_failure() -> Result }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let status: GetAuthStatusResponse = to_response(resp)?; + let status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( status, GetAuthStatusResponse { @@ -555,9 +489,8 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res Some(refresh_url.as_str()), ), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let failed_request_id = mcp .send_get_auth_status_request(GetAuthStatusParams { @@ -566,12 +499,8 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res }) .await?; - let failed_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(failed_request_id)), - ) - .await??; - let failed_status: GetAuthStatusResponse = to_response(failed_resp)?; + let failed_status: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(failed_request_id)).await??; assert_eq!( failed_status, GetAuthStatusResponse { @@ -599,12 +528,11 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res }) .await?; - let recovered_resp: JSONRPCResponse = timeout( + let recovered_status: GetAuthStatusResponse = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(recovered_request_id)), + mcp.read_response(recovered_request_id), ) .await??; - let recovered_status: GetAuthStatusResponse = to_response(recovered_resp)?; assert_eq!( recovered_status, GetAuthStatusResponse { @@ -626,9 +554,8 @@ async fn login_api_key_rejected_when_forced_chatgpt() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_api_key_request("sk-test-key") diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index 050a07ebe62c..228165836d64 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -1,8 +1,8 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_fake_rollout; use app_test_support::rollout_path; -use app_test_support::to_response; use codex_app_server::in_process; use codex_app_server::in_process::InProcessStartArgs; use codex_app_server_protocol::ClientInfo; @@ -12,7 +12,6 @@ use codex_app_server_protocol::GetConversationSummaryParams; use codex_app_server_protocol::GetConversationSummaryResponse; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeParams; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_arg0::Arg0DispatchPaths; use codex_config::CloudConfigBundleLoader; @@ -35,10 +34,8 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use tempfile::TempDir; -use tokio::time::timeout; use uuid::Uuid; -const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); const FILENAME_TS: &str = "2025-01-02T12-00-00"; const META_RFC3339: &str = "2025-01-02T12:00:00Z"; const CREATED_AT_RFC3339: &str = "2025-01-02T12:00:00.000Z"; @@ -96,21 +93,17 @@ async fn get_conversation_summary_by_thread_id_reads_rollout() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_get_conversation_summary_request(GetConversationSummaryParams::ThreadId { - conversation_id: thread_id, + let received: GetConversationSummaryResponse = mcp + .request(|request_id| ClientRequest::GetConversationSummary { + request_id, + params: GetConversationSummaryParams::ThreadId { + conversation_id: thread_id, + }, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetConversationSummaryResponse = to_response(response)?; assert_eq!(normalized_summary_path(received.summary)?, expected); Ok(()) @@ -226,21 +219,17 @@ async fn get_conversation_summary_by_relative_rollout_path_resolves_from_codex_h let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_get_conversation_summary_request(GetConversationSummaryParams::RolloutPath { - rollout_path: relative_path, + let received: GetConversationSummaryResponse = mcp + .request(|request_id| ClientRequest::GetConversationSummary { + request_id, + params: GetConversationSummaryParams::RolloutPath { + rollout_path: relative_path, + }, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetConversationSummaryResponse = to_response(response)?; assert_eq!(normalized_summary_path(received.summary)?, expected); Ok(()) @@ -260,24 +249,9 @@ fn create_config_toml_with_in_memory_thread_store( codex_home: &Path, store_id: &str, ) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }} - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "http://127.0.0.1:1/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) + MockResponsesConfig::new("http://127.0.0.1:1") + .with_root_config(&format!( + "experimental_thread_store = {{ type = \"in_memory\", id = \"{store_id}\" }}" + )) + .write(codex_home) } diff --git a/codex-rs/app-server/tests/suite/fuzzy_file_search.rs b/codex-rs/app-server/tests/suite/fuzzy_file_search.rs index 2a6c8808b7ca..34dd3cd9dc0c 100644 --- a/codex-rs/app-server/tests/suite/fuzzy_file_search.rs +++ b/codex-rs/app-server/tests/suite/fuzzy_file_search.rs @@ -46,13 +46,11 @@ shell_snapshot = false async fn initialized_mcp(codex_home: &TempDir) -> Result { create_config_toml(codex_home.path())?; - let mut mcp = TestAppServer::builder() + TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() - .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - Ok(mcp) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await } async fn wait_for_session_updated( @@ -243,9 +241,8 @@ async fn test_fuzzy_file_search_sorts_and_includes_indices() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let root_path = root.path().to_string_lossy().to_string(); // Send fuzzyFileSearch request. @@ -313,9 +310,8 @@ async fn test_fuzzy_file_search_accepts_cancellation_token() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let root_path = root.path().to_string_lossy().to_string(); let request_id = mcp diff --git a/codex-rs/app-server/tests/suite/logging.rs b/codex-rs/app-server/tests/suite/logging.rs index 5628b62b6b51..04486a70cef8 100644 --- a/codex-rs/app-server/tests/suite/logging.rs +++ b/codex-rs/app-server/tests/suite/logging.rs @@ -1,16 +1,13 @@ use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::app_server_json_shutdown_event; use app_test_support::create_exec_command_sse_response; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; -use app_test_support::to_response; -use app_test_support::write_mock_responses_config_toml; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ThreadStartParams; -use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput; @@ -19,7 +16,6 @@ use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; -use std::collections::BTreeMap; use tempfile::TempDir; use tokio::time::Duration; use tokio::time::timeout; @@ -58,52 +54,39 @@ async fn app_server_emits_structured_tool_call_timing_event() -> Result<()> { ]) .await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::from([(Feature::UnifiedExec, true)]), - /*auto_compact_limit*/ 100_000, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::UnifiedExec) + .with_root_config("compact_prompt = \"compact\"\nmodel_auto_compact_token_limit = 100000") + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_json_logging("warn,codex_core::tools::parallel=info") - .build() + .build_initialized() .await?; - timeout(READ_TIMEOUT, app_server.initialize()).await??; - let thread_start_id = app_server - .send_thread_start_request_with_auto_env(ThreadStartParams { + let thread = app_server + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) - .await?; - let thread_start_response: JSONRPCResponse = timeout( - READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + .await? + .thread; - let turn_start_id = app_server - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - input: vec![UserInput::Text { - text: "run a command".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let TurnStartResponse { turn } = app_server + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + input: vec![UserInput::Text { + text: "run a command".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_start_response: JSONRPCResponse = timeout( - READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response(turn_start_response)?; timeout( READ_TIMEOUT, diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 4c2bdf270d04..8b1d9585906e 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -185,26 +185,18 @@ async fn read_account(mcp: &mut TestAppServer) -> Result { refresh_token: false, }) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await? } async fn assert_account_updated( mcp: &mut TestAppServer, auth_mode: Option, ) -> Result<()> { - let notification = timeout( + let payload: AccountUpdatedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("account/updated"), + mcp.read_notification("account/updated"), ) .await??; - let ServerNotification::AccountUpdated(payload) = notification.try_into()? else { - bail!("unexpected notification") - }; assert_eq!( payload, AccountUpdatedNotification { @@ -284,17 +276,11 @@ async fn logout_account_removes_auth_and_notifies() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let id = mcp.send_logout_account_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(id)), - ) - .await??; - let _ok: LogoutAccountResponse = to_response(resp)?; + let _ok: LogoutAccountResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(id)).await??; let note = timeout( DEFAULT_READ_TIMEOUT, @@ -321,12 +307,8 @@ async fn logout_account_removes_auth_and_notifies() -> Result<()> { refresh_token: false, }) .await?; - let get_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(get_id)), - ) - .await??; - let account: GetAccountResponse = to_response(get_resp)?; + let account: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(get_id)).await??; assert_eq!(account.account, None); Ok(()) } @@ -346,9 +328,8 @@ async fn logout_account_succeeds_when_config_reload_fails() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; std::fs::write(codex_home.path().join("config.toml"), "invalid = [")?; @@ -393,9 +374,8 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( @@ -404,12 +384,8 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { Some("pro".to_string()), ) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(set_resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); let note = timeout( @@ -429,12 +405,8 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { refresh_token: false, }) .await?; - let get_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(get_id)), - ) - .await??; - let account: GetAccountResponse = to_response(get_resp)?; + let account: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(get_id)).await??; assert_eq!( account, GetAccountResponse { @@ -447,24 +419,16 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { ); let logout_id = mcp.send_logout_account_request().await?; - let logout_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(logout_id)), - ) - .await??; - let _: LogoutAccountResponse = to_response(logout_resp)?; + let _: LogoutAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(logout_id)).await??; let get_id = mcp .send_get_account_request(GetAccountParams { refresh_token: false, }) .await?; - let get_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(get_id)), - ) - .await??; - let account: GetAccountResponse = to_response(get_resp)?; + let account: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(get_id)).await??; assert_eq!(account.account, None); Ok(()) @@ -493,9 +457,8 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( @@ -504,12 +467,8 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { Some("pro".to_string()), ) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(set_resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); let _updated = timeout( DEFAULT_READ_TIMEOUT, @@ -522,12 +481,8 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { refresh_token: true, }) .await?; - let get_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(get_id)), - ) - .await??; - let account: GetAccountResponse = to_response(get_resp)?; + let account: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(get_id)).await??; assert_eq!( account, GetAccountResponse { @@ -622,9 +577,8 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( @@ -633,12 +587,8 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { Some("pro".to_string()), ) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(set_resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); let _updated = timeout( DEFAULT_READ_TIMEOUT, @@ -652,12 +602,8 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let thread = to_response::(thread_resp)?; + let thread: codex_app_server_protocol::ThreadStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let turn_req = mcp .send_turn_start_request(codex_app_server_protocol::TurnStartParams { @@ -677,11 +623,8 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { Some("pro"), ) .await?; - let _turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; + let _: codex_app_server_protocol::TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; let _turn_completed = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -733,9 +676,8 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( @@ -744,12 +686,8 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { Some("pro".to_string()), ) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(set_resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); let _updated = timeout( DEFAULT_READ_TIMEOUT, @@ -763,12 +701,8 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let thread = to_response::(thread_resp)?; + let thread: codex_app_server_protocol::ThreadStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let turn_req = mcp .send_turn_start_request(codex_app_server_protocol::TurnStartParams { @@ -801,11 +735,8 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { ) .await?; - let _turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; + let _: codex_app_server_protocol::TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; let completed_notif: JSONRPCNotification = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -860,9 +791,8 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( @@ -871,12 +801,8 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { Some("pro".to_string()), ) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(set_resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); let _updated = timeout( DEFAULT_READ_TIMEOUT, @@ -890,12 +816,8 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let thread = to_response::(thread_resp)?; + let thread: codex_app_server_protocol::ThreadStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let turn_req = mcp .send_turn_start_request(codex_app_server_protocol::TurnStartParams { @@ -928,11 +850,8 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { ) .await?; - let _turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; + let _: codex_app_server_protocol::TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; let completed_notif: JSONRPCNotification = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -980,9 +899,8 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( @@ -991,12 +909,8 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { Some("pro".to_string()), ) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(set_resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); let _updated = timeout( DEFAULT_READ_TIMEOUT, @@ -1010,12 +924,8 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let thread = to_response::(thread_resp)?; + let thread: codex_app_server_protocol::ThreadStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let turn_req = mcp .send_turn_start_request(codex_app_server_protocol::TurnStartParams { @@ -1048,11 +958,8 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { ) .await?; - let _turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; + let _: codex_app_server_protocol::TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; let completed_notif: JSONRPCNotification = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -1077,19 +984,14 @@ async fn login_account_api_key_succeeds_and_notifies() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_login_account_api_key_request("sk-test-key") .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(req_id)).await??; assert_eq!(login, LoginAccountResponse::ApiKey {}); let note = timeout( @@ -1135,9 +1037,8 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider() -> R .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let mut expected_config = read_config_toml(codex_home.path())?; expected_config .as_table_mut() @@ -1216,9 +1117,8 @@ async fn login_amazon_bedrock_rejects_non_bedrock_provider_override_without_chan .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) .with_args(&["-c", "model_provider=\"mock_provider\""]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") @@ -1275,9 +1175,8 @@ async fn login_amazon_bedrock_allows_bedrock_provider_override() -> Result<()> { .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) .with_args(&["-c", "model_provider=\"amazon-bedrock\""]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") @@ -1331,9 +1230,8 @@ async fn logout_managed_bedrock_restores_default_account() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") .await?; @@ -1402,9 +1300,8 @@ async fn logout_aws_managed_bedrock_errors_without_changing_auth_or_config() -> .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_logout_account_request().await?; let error = timeout( DEFAULT_READ_TIMEOUT, @@ -1529,9 +1426,8 @@ async fn login_managed_bedrock_updates_active_bedrock_account() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") .await?; @@ -1573,9 +1469,8 @@ async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_change .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let expected_config = read_config_toml(codex_home.path())?; let request_id = mcp @@ -1623,9 +1518,8 @@ async fn login_account_amazon_bedrock_rejected_when_forced_chatgpt() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") .await?; @@ -1657,9 +1551,8 @@ async fn login_account_amazon_bedrock_rejected_with_external_chatgpt_auth() -> R let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( access_token, @@ -1712,9 +1605,8 @@ async fn login_account_api_key_rejected_when_forced_chatgpt() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_api_key_request("sk-test-key") @@ -1746,9 +1638,8 @@ async fn login_account_chatgpt_rejected_when_forced_api() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_request().await?; let err: JSONRPCError = timeout( @@ -1787,9 +1678,8 @@ async fn login_account_chatgpt_device_code_returns_error_when_disabled() -> Resu ("OPENAI_API_KEY", None), (LOGIN_ISSUER_ENV_VAR, Some(issuer.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_device_code_request().await?; let err: JSONRPCError = timeout( @@ -1853,17 +1743,12 @@ async fn login_account_chatgpt_device_code_succeeds_and_notifies() -> Result<()> ("OPENAI_API_KEY", None), (LOGIN_ISSUER_ENV_VAR, Some(issuer.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_device_code_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::ChatgptDeviceCode { login_id, verification_url, @@ -1931,17 +1816,12 @@ async fn login_account_chatgpt_device_code_failure_notifies_without_account_upda ("OPENAI_API_KEY", None), (LOGIN_ISSUER_ENV_VAR, Some(issuer.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_device_code_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::ChatgptDeviceCode { login_id, .. } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2007,17 +1887,12 @@ async fn login_account_chatgpt_device_code_can_be_cancelled() -> Result<()> { ("OPENAI_API_KEY", None), (LOGIN_ISSUER_ENV_VAR, Some(issuer.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_device_code_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::ChatgptDeviceCode { login_id, .. } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2027,12 +1902,8 @@ async fn login_account_chatgpt_device_code_can_be_cancelled() -> Result<()> { login_id: login_id.clone(), }) .await?; - let cancel_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(cancel_id)), - ) - .await??; - let cancel: CancelLoginAccountResponse = to_response(cancel_resp)?; + let cancel: CancelLoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(cancel_id)).await??; assert_eq!(cancel.status, CancelLoginAccountStatus::Canceled); let note = timeout( @@ -2077,18 +1948,12 @@ async fn login_account_chatgpt_start_can_be_cancelled() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::Chatgpt { login_id, auth_url } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2102,12 +1967,8 @@ async fn login_account_chatgpt_start_can_be_cancelled() -> Result<()> { login_id: login_id.clone(), }) .await?; - let cancel_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(cancel_id)), - ) - .await??; - let _ok: CancelLoginAccountResponse = to_response(cancel_resp)?; + let _ok: CancelLoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(cancel_id)).await??; let note = timeout( DEFAULT_READ_TIMEOUT, @@ -2151,18 +2012,12 @@ async fn login_account_chatgpt_uses_debug_oauth_overrides() -> Result<()> { (CLIENT_ID_OVERRIDE_ENV_VAR, Some("staging-client")), (LOGIN_ISSUER_ENV_VAR, Some("https://auth.example.com")), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::Chatgpt { login_id, auth_url } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2181,12 +2036,8 @@ async fn login_account_chatgpt_uses_debug_oauth_overrides() -> Result<()> { let cancel_id = mcp .send_cancel_login_account_request(CancelLoginAccountParams { login_id }) .await?; - let cancel_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(cancel_id)), - ) - .await??; - let _: CancelLoginAccountResponse = to_response(cancel_resp)?; + let _: CancelLoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(cancel_id)).await??; Ok(()) } @@ -2216,9 +2067,8 @@ async fn login_account_chatgpt_redirects_to_hosted_success_page() -> Result<()> Some("http://localhost:3000/codex/open-app"), ), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_request(json!({ @@ -2227,12 +2077,8 @@ async fn login_account_chatgpt_redirects_to_hosted_success_page() -> Result<()> "useHostedLoginSuccessPage": true, })) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::Chatgpt { auth_url, .. } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2272,19 +2118,13 @@ async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // Initiate the ChatGPT login flow let request_id = mcp.send_login_account_chatgpt_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::Chatgpt { login_id, .. } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2304,12 +2144,8 @@ async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { Some("pro".to_string()), ) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let response: LoginAccountResponse = to_response(set_resp)?; + let response: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); let _updated = timeout( DEFAULT_READ_TIMEOUT, @@ -2324,12 +2160,8 @@ async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { login_id: login_id.clone(), }) .await?; - let cancel_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(cancel_id)), - ) - .await??; - let cancel: CancelLoginAccountResponse = to_response(cancel_resp)?; + let cancel: CancelLoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(cancel_id)).await??; assert_eq!(cancel.status, CancelLoginAccountStatus::NotFound); Ok(()) @@ -2351,18 +2183,12 @@ async fn login_account_chatgpt_includes_forced_workspace_query_param() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::Chatgpt { auth_url, .. } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2392,18 +2218,12 @@ async fn login_account_chatgpt_includes_forced_workspace_allowlist_query_param() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_login_account_chatgpt_request().await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let login: LoginAccountResponse = to_response(resp)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let LoginAccountResponse::Chatgpt { auth_url, .. } = login else { bail!("unexpected login response: {login:?}"); }; @@ -2436,21 +2256,16 @@ async fn get_account_no_auth() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let params = GetAccountParams { refresh_token: false, }; let request_id = mcp.send_get_account_request(params).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let account: GetAccountResponse = to_response(resp)?; + let account: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(account.account, None, "expected no account"); assert_eq!(account.requires_openai_auth, true); @@ -2471,31 +2286,22 @@ async fn get_account_with_api_key() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_login_account_api_key_request("sk-test-key") .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let _login_ok = to_response::(resp)?; + let _login_ok: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(req_id)).await??; let params = GetAccountParams { refresh_token: false, }; let request_id = mcp.send_get_account_request(params).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(resp)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { account: Some(Account::ApiKey {}), @@ -2519,21 +2325,16 @@ async fn get_account_when_auth_not_required() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let params = GetAccountParams { refresh_token: false, }; let request_id = mcp.send_get_account_request(params).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(resp)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { account: None, @@ -2564,21 +2365,16 @@ region = "us-west-2" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let params = GetAccountParams { refresh_token: false, }; let request_id = mcp.send_get_account_request(params).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(resp)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { account: Some(Account::AmazonBedrock { @@ -2613,9 +2409,8 @@ command = "print-token" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; assert_eq!( read_account(&mut mcp).await?, @@ -2650,9 +2445,8 @@ region = "us-west-2" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; std::fs::write(codex_home.path().join("config.toml"), "invalid = [")?; @@ -2710,21 +2504,16 @@ async fn get_account_with_managed_bedrock_provider() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_get_account_request(GetAccountParams { refresh_token: false, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(resp)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( received, @@ -2760,21 +2549,16 @@ async fn get_account_with_chatgpt() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let params = GetAccountParams { refresh_token: false, }; let request_id = mcp.send_get_account_request(params).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(resp)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { account: Some(Account::Chatgpt { @@ -2807,21 +2591,16 @@ async fn get_account_with_chatgpt_without_email() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_get_account_request(GetAccountParams { refresh_token: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(response)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( received, @@ -2880,9 +2659,8 @@ async fn get_account_omits_chatgpt_after_permanent_refresh_failure() -> Result<( Some(refresh_url.as_str()), ), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let auth_status_request_id = mcp .send_get_auth_status_request(GetAuthStatusParams { @@ -2890,12 +2668,11 @@ async fn get_account_omits_chatgpt_after_permanent_refresh_failure() -> Result<( refresh_token: Some(true), }) .await?; - let auth_status_resp: JSONRPCResponse = timeout( + let _: GetAuthStatusResponse = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(auth_status_request_id)), + mcp.read_response(auth_status_request_id), ) .await??; - let _: GetAuthStatusResponse = to_response(auth_status_resp)?; let request_id = mcp .send_get_account_request(GetAccountParams { @@ -2903,12 +2680,8 @@ async fn get_account_omits_chatgpt_after_permanent_refresh_failure() -> Result<( }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(resp)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( received, @@ -2941,21 +2714,16 @@ async fn get_account_with_chatgpt_missing_plan_claim_returns_unknown() -> Result .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let params = GetAccountParams { refresh_token: false, }; let request_id = mcp.send_get_account_request(params).await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountResponse = to_response(resp)?; + let received: GetAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { account: Some(Account::Chatgpt { diff --git a/codex-rs/app-server/tests/suite/v2/app_installed.rs b/codex-rs/app-server/tests/suite/v2/app_installed.rs index 5494916bebed..1327e21feb6d 100644 --- a/codex-rs/app-server/tests/suite/v2/app_installed.rs +++ b/codex-rs/app-server/tests/suite/v2/app_installed.rs @@ -10,7 +10,6 @@ use std::time::Duration; use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Json; use axum::Router; @@ -21,7 +20,6 @@ use codex_app_server_protocol::AppsInstalledParams; use codex_app_server_protocol::AppsInstalledResponse; use codex_app_server_protocol::InstalledApp; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -202,12 +200,8 @@ async fn installed_apps_thread_id_uses_effective_thread_config() -> Result<()> { ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(response)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_TIMEOUT, app_server.read_response(request_id)).await??; let request_id = app_server .send_apps_installed_request(AppsInstalledParams { @@ -215,12 +209,8 @@ async fn installed_apps_thread_id_uses_effective_thread_config() -> Result<()> { force_refresh: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: AppsInstalledResponse = to_response(response)?; + let response: AppsInstalledResponse = + timeout(DEFAULT_TIMEOUT, app_server.read_response(request_id)).await??; let alpha = expected .apps .iter_mut() @@ -261,13 +251,11 @@ async fn installed_apps_failed_force_refresh_retains_previous_snapshot() -> Resu } async fn start_app_server(codex_home: &Path) -> Result { - let mut app_server = TestAppServer::builder() + TestAppServer::builder() .with_codex_home(codex_home) .without_managed_config() - .build() - .await?; - timeout(DEFAULT_TIMEOUT, app_server.initialize()).await??; - Ok(app_server) + .build_initialized_with_timeout(DEFAULT_TIMEOUT) + .await } async fn send_installed_request( @@ -280,12 +268,7 @@ async fn send_installed_request( force_refresh, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + timeout(DEFAULT_TIMEOUT, app_server.read_response(request_id)).await? } fn configured_codex_home(base_url: &str) -> Result { diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index bb76561e22f6..8dab9a829bd5 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -6,12 +6,10 @@ use std::sync::Mutex as StdMutex; use std::time::Duration; use anyhow::Result; -use anyhow::bail; use app_test_support::ChatGptAuthFixture; use app_test_support::ChatGptIdTokenClaims; use app_test_support::TestAppServer; use app_test_support::encode_id_token; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Json; use axum::Router; @@ -30,10 +28,8 @@ use codex_app_server_protocol::AppScreenshot; use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; -use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_config::types::AuthCredentialsStoreMode; @@ -69,11 +65,9 @@ async fn list_apps_returns_empty_when_connectors_disabled() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp .send_apps_list_request(AppsListParams { limit: Some(50), @@ -83,13 +77,8 @@ async fn list_apps_returns_empty_when_connectors_disabled() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(data.is_empty()); assert!(next_cursor.is_none()); @@ -139,9 +128,8 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -152,13 +140,8 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(data.is_empty()); assert!(next_cursor.is_none()); @@ -209,9 +192,8 @@ async fn list_apps_uses_external_chatgpt_auth() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let login_id = mcp .send_chatgpt_auth_tokens_login_request( access_token, @@ -219,15 +201,9 @@ async fn list_apps_uses_external_chatgpt_auth() -> Result<()> { Some("pro".to_string()), ) .await?; - let login_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(login_id)), - ) - .await??; - assert_eq!( - to_response::(login_response)?, - LoginAccountResponse::ChatgptAuthTokens {} - ); + let login_response: LoginAccountResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(login_id)).await??; + assert_eq!(login_response, LoginAccountResponse::ChatgptAuthTokens {}); let request_id = mcp .send_apps_list_request(AppsListParams { @@ -237,12 +213,8 @@ async fn list_apps_uses_external_chatgpt_auth() -> Result<()> { force_refetch: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].id, "beta"); @@ -295,9 +267,8 @@ async fn list_apps_returns_empty_when_workspace_codex_plugins_disabled() -> Resu .with_codex_home(codex_home.path()) .without_auto_env() .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -308,13 +279,8 @@ async fn list_apps_returns_empty_when_workspace_codex_plugins_disabled() -> Resu }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(data.is_empty()); assert!(next_cursor.is_none()); @@ -344,9 +310,8 @@ async fn list_apps_includes_plugin_apps_for_chatgpt_auth() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -356,12 +321,8 @@ async fn list_apps_includes_plugin_apps_for_chatgpt_auth() -> Result<()> { force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(data.iter().any(|app| app.id == "connector_sample")); assert!(next_cursor.is_none()); @@ -407,19 +368,14 @@ async fn list_apps_uses_thread_feature_flag_when_thread_id_is_provided() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let start_request = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_request)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(start_response)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(start_request)).await??; std::fs::write( codex_home.path().join("config.toml"), @@ -442,15 +398,10 @@ connectors = false force_refetch: false, }) .await?; - let global_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(global_request)), - ) - .await??; let AppsListResponse { data: global_data, next_cursor: global_next_cursor, - } = to_response(global_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(global_request)).await??; assert!(global_data.is_empty()); assert!(global_next_cursor.is_none()); @@ -462,15 +413,10 @@ connectors = false force_refetch: false, }) .await?; - let thread_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_request)), - ) - .await??; let AppsListResponse { data: thread_data, next_cursor: thread_next_cursor, - } = to_response(thread_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_request)).await??; assert!(thread_data.iter().any(|app| app.id == "beta")); assert!(thread_next_cursor.is_none()); @@ -524,9 +470,8 @@ async fn list_apps_keeps_apps_with_app_only_tools_accessible() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -536,12 +481,8 @@ async fn list_apps_keeps_apps_with_app_only_tools_accessible() -> Result<()> { force_refetch: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].id, connector_id); @@ -603,9 +544,8 @@ enabled = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -616,15 +556,10 @@ enabled = false }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let AppsListResponse { data: response_data, next_cursor, - } = to_response(response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(next_cursor.is_none()); assert_eq!(response_data.len(), 1); assert_eq!(response_data[0].id, "beta"); @@ -730,9 +665,8 @@ async fn list_apps_emits_updates_and_returns_after_both_lists_load() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -804,16 +738,10 @@ async fn list_apps_emits_updates_and_returns_after_both_lists_load() -> Result<( let second_update = read_app_list_updated_notification(&mut mcp).await?; assert_eq!(second_update.data, expected_merged); - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let AppsListResponse { data: response_data, next_cursor, - } = to_response(response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response_data, expected_merged); assert!(next_cursor.is_none()); @@ -884,9 +812,8 @@ async fn list_apps_waits_for_accessible_data_before_emitting_directory_updates() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -946,12 +873,8 @@ async fn list_apps_waits_for_accessible_data_before_emitting_directory_updates() ); } - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data, expected); assert!(next_cursor.is_none()); @@ -1000,9 +923,8 @@ async fn list_apps_does_not_emit_empty_interim_updates() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_list_request(AppsListParams { @@ -1044,12 +966,8 @@ async fn list_apps_does_not_emit_empty_interim_updates() -> Result<()> { let update = read_app_list_updated_notification(&mut mcp).await?; assert_eq!(update.data, expected); - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let AppsListResponse { data, next_cursor } = to_response(response)?; + let AppsListResponse { data, next_cursor } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data, expected); assert!(next_cursor.is_none()); @@ -1119,9 +1037,8 @@ async fn list_apps_paginates_results() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let first_request = mcp .send_apps_list_request(AppsListParams { @@ -1131,15 +1048,10 @@ async fn list_apps_paginates_results() -> Result<()> { force_refetch: false, }) .await?; - let first_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_request)), - ) - .await??; let AppsListResponse { data: first_page, next_cursor: first_cursor, - } = to_response(first_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(first_request)).await??; let expected_first = vec![AppInfo { id: "beta".to_string(), @@ -1178,15 +1090,10 @@ async fn list_apps_paginates_results() -> Result<()> { force_refetch: false, }) .await?; - let second_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_request)), - ) - .await??; let AppsListResponse { data: second_page, next_cursor: second_cursor, - } = to_response(second_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(second_request)).await??; let expected_second = vec![AppInfo { id: "alpha".to_string(), @@ -1260,9 +1167,8 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let initial_request = mcp .send_apps_list_request(AppsListParams { @@ -1272,15 +1178,10 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result force_refetch: false, }) .await?; - let initial_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(initial_request)), - ) - .await??; let AppsListResponse { data: initial_data, next_cursor: initial_next_cursor, - } = to_response(initial_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(initial_request)).await??; assert!(initial_next_cursor.is_none()); assert_eq!(initial_data.len(), 1); assert!(initial_data.iter().all(|app| app.is_accessible)); @@ -1317,15 +1218,10 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result force_refetch: false, }) .await?; - let cached_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(cached_request)), - ) - .await??; let AppsListResponse { data: cached_data, next_cursor: cached_next_cursor, - } = to_response(cached_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(cached_request)).await??; assert_eq!(cached_data, initial_data); assert!(cached_next_cursor.is_none()); @@ -1394,9 +1290,8 @@ async fn list_apps_force_refetch_patches_updates_from_cached_snapshots() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let warm_request = mcp .send_apps_list_request(AppsListParams { @@ -1469,15 +1364,10 @@ async fn list_apps_force_refetch_patches_updates_from_cached_snapshots() -> Resu ] ); - let warm_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(warm_request)), - ) - .await??; let AppsListResponse { data: warm_data, next_cursor: warm_next_cursor, - } = to_response(warm_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(warm_request)).await??; assert_eq!(warm_data, warm_second_update.data); assert!(warm_next_cursor.is_none()); @@ -1580,15 +1470,10 @@ async fn list_apps_force_refetch_patches_updates_from_cached_snapshots() -> Resu let second_update = read_app_list_updated_notification(&mut mcp).await?; assert_eq!(second_update.data, expected_final); - let refetch_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(refetch_request)), - ) - .await??; let AppsListResponse { data: refetch_data, next_cursor: refetch_next_cursor, - } = to_response(refetch_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(refetch_request)).await??; assert_eq!(refetch_data, expected_final); assert!(refetch_next_cursor.is_none()); @@ -1601,15 +1486,10 @@ async fn list_apps_force_refetch_patches_updates_from_cached_snapshots() -> Resu force_refetch: false, }) .await?; - let cached_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(cached_request)), - ) - .await??; let AppsListResponse { data: cached_data, next_cursor: cached_next_cursor, - } = to_response(cached_response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(cached_request)).await??; assert_eq!(cached_data, expected_final); assert!(cached_next_cursor.is_none()); @@ -1633,16 +1513,7 @@ async fn list_apps_force_refetch_patches_updates_from_cached_snapshots() -> Resu async fn read_app_list_updated_notification( mcp: &mut TestAppServer, ) -> Result { - let notification = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("app/list/updated"), - ) - .await??; - let parsed: ServerNotification = notification.try_into()?; - let ServerNotification::AppListUpdated(payload) = parsed else { - bail!("unexpected notification variant"); - }; - Ok(payload) + timeout(DEFAULT_TIMEOUT, mcp.read_notification("app/list/updated")).await? } #[derive(Clone)] diff --git a/codex-rs/app-server/tests/suite/v2/app_read.rs b/codex-rs/app-server/tests/suite/v2/app_read.rs index 53e0b8a9f623..e81f01773b0b 100644 --- a/codex-rs/app-server/tests/suite/v2/app_read.rs +++ b/codex-rs/app-server/tests/suite/v2/app_read.rs @@ -8,7 +8,6 @@ use app_test_support::ChatGptAuthFixture; use app_test_support::ChatGptIdTokenClaims; use app_test_support::TestAppServer; use app_test_support::encode_id_token; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Json; use axum::Router; @@ -22,7 +21,6 @@ use codex_app_server_protocol::AppsReadParams; use codex_app_server_protocol::AppsReadResponse; use codex_app_server_protocol::ConnectorMetadata; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; @@ -72,9 +70,8 @@ async fn app_read_deduplicates_orders_partial_misses_and_reuses_cached_metadata( .with_codex_home(codex_home.path()) .without_auto_env() .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let login_id = mcp .send_chatgpt_auth_tokens_login_request( access_token, @@ -82,15 +79,9 @@ async fn app_read_deduplicates_orders_partial_misses_and_reuses_cached_metadata( Some("plus".to_string()), ) .await?; - let login_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(login_id)), - ) - .await??; - assert_eq!( - to_response::(login_response)?, - LoginAccountResponse::ChatgptAuthTokens {} - ); + let login_response: LoginAccountResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(login_id)).await??; + assert_eq!(login_response, LoginAccountResponse::ChatgptAuthTokens {}); let raw_response = read_apps_raw( &mut mcp, @@ -183,9 +174,8 @@ async fn app_read_refetches_metadata_only_cache_entries_when_tools_are_requested .with_codex_home(codex_home.path()) .without_auto_env() .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; assert_eq!( read_apps(&mut mcp, vec!["cached"], /*include_tools*/ false).await?, @@ -269,9 +259,8 @@ async fn app_read_backend_failure_preserves_fresh_cached_records() -> Result<()> .with_codex_home(codex_home.path()) .without_auto_env() .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; assert_eq!( read_apps(&mut mcp, vec!["cached"], /*include_tools*/ true).await?, @@ -359,9 +348,8 @@ enabled = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let response = read_apps( &mut mcp, @@ -401,9 +389,8 @@ async fn app_read_rejects_more_than_one_hundred_input_ids() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_apps_read_request(AppsReadParams { @@ -442,12 +429,7 @@ async fn read_apps_raw( include_tools, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - Ok(response.result) + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await? } fn metadata(id: &str, name: &str, icon_url: Option<&str>) -> ConnectorMetadata { diff --git a/codex-rs/app-server/tests/suite/v2/client_metadata.rs b/codex-rs/app-server/tests/suite/v2/client_metadata.rs index 78807a055a79..b6d7090e342c 100644 --- a/codex-rs/app-server/tests/suite/v2/client_metadata.rs +++ b/codex-rs/app-server/tests/suite/v2/client_metadata.rs @@ -1,10 +1,8 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_fake_parented_rollout_with_source; use app_test_support::create_fake_rollout; -use app_test_support::to_response; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ReviewDelivery; use codex_app_server_protocol::ReviewStartParams; use codex_app_server_protocol::ReviewStartResponse; @@ -29,7 +27,6 @@ use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use std::collections::HashMap; -use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; @@ -53,17 +50,14 @@ async fn turn_start_forwards_client_metadata_to_responses_request_v2() -> Result .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - /*supports_websockets*/ false, - )?; + MockResponsesConfig::new(&server.uri()) + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -71,12 +65,8 @@ async fn turn_start_forwards_client_metadata_to_responses_request_v2() -> Result ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let client_metadata = HashMap::from([ ("fiber_run_id".to_string(), "fiber-start-123".to_string()), @@ -94,12 +84,8 @@ async fn turn_start_forwards_client_metadata_to_responses_request_v2() -> Result ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, @@ -143,11 +129,9 @@ async fn turn_start_sends_fork_lineage_in_turn_metadata_for_thread_fork_v2() -> .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - /*supports_websockets*/ false, - )?; + MockResponsesConfig::new(&server.uri()) + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let source_thread_id = create_fake_rollout( codex_home.path(), @@ -161,9 +145,8 @@ async fn turn_start_sends_fork_lineage_in_turn_metadata_for_thread_fork_v2() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let ThreadForkResponse { thread, .. } = fork_fake_rollout_thread(&mut mcp, source_thread_id.clone()).await?; @@ -179,12 +162,8 @@ async fn turn_start_sends_fork_lineage_in_turn_metadata_for_thread_fork_v2() -> ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, @@ -231,11 +210,9 @@ async fn review_start_sends_parent_lineage_in_turn_metadata_for_thread_fork_v2() .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - /*supports_websockets*/ false, - )?; + MockResponsesConfig::new(&server.uri()) + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let source_thread_id = create_fake_rollout( codex_home.path(), @@ -249,9 +226,8 @@ async fn review_start_sends_parent_lineage_in_turn_metadata_for_thread_fork_v2() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let ThreadForkResponse { thread, .. } = fork_fake_rollout_thread(&mut mcp, source_thread_id.clone()).await?; @@ -265,14 +241,9 @@ async fn review_start_sends_parent_lineage_in_turn_metadata_for_thread_fork_v2() }, }) .await?; - let review_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(review_req)), - ) - .await??; let ReviewStartResponse { review_thread_id, .. - } = to_response::(review_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(review_req)).await??; assert_eq!(review_thread_id, thread.id); timeout( @@ -328,11 +299,9 @@ async fn turn_start_sends_nested_subagent_lineage_after_cold_thread_resume_v2() .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - /*supports_websockets*/ false, - )?; + MockResponsesConfig::new(&server.uri()) + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let root_thread_id = CoreThreadId::new(); let root_thread_id_str = root_thread_id.to_string(); @@ -353,9 +322,8 @@ async fn turn_start_sends_nested_subagent_lineage_after_cold_thread_resume_v2() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_req = mcp .send_thread_resume_request(ThreadResumeParams { @@ -363,12 +331,8 @@ async fn turn_start_sends_nested_subagent_lineage_after_cold_thread_resume_v2() ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_req)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_req)).await??; assert_eq!(thread.id, subagent_thread_id); assert_eq!(thread.session_id, root_thread_id_str); assert_eq!(thread.parent_thread_id, Some(parent_thread_id_str.clone())); @@ -387,12 +351,8 @@ async fn turn_start_sends_nested_subagent_lineage_after_cold_thread_resume_v2() ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, @@ -443,27 +403,20 @@ async fn turn_steer_updates_client_metadata_on_follow_up_responses_request_v2() let request_log = responses::mount_response_sequence(&server, vec![first_response, second_response]).await; - create_config_toml( - codex_home.path(), - &server.uri(), - /*supports_websockets*/ false, - )?; + MockResponsesConfig::new(&server.uri()) + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let start_metadata = HashMap::from([("fiber_run_id".to_string(), "fiber-start-123".to_string())]); @@ -479,12 +432,8 @@ async fn turn_steer_updates_client_metadata_on_follow_up_responses_request_v2() ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; let turn_id = turn.id.clone(); timeout( @@ -511,12 +460,8 @@ async fn turn_steer_updates_client_metadata_on_follow_up_responses_request_v2() expected_turn_id: turn_id.clone(), }) .await?; - let steer_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(steer_req)), - ) - .await??; - let _turn: TurnSteerResponse = to_response::(steer_resp)?; + let _turn: TurnSteerResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(steer_req)).await??; timeout( DEFAULT_READ_TIMEOUT, @@ -571,17 +516,14 @@ async fn turn_start_forwards_client_metadata_to_responses_websocket_request_body .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &websocket_server.uri().replacen("ws://", "http://", 1), - /*supports_websockets*/ true, - )?; + MockResponsesConfig::new(&websocket_server.uri().replacen("ws://", "http://", 1)) + .with_provider_config("supports_websockets = true") + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -589,12 +531,8 @@ async fn turn_start_forwards_client_metadata_to_responses_websocket_request_body ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let client_metadata = HashMap::from([ ("fiber_run_id".to_string(), "fiber-start-123".to_string()), @@ -612,12 +550,8 @@ async fn turn_start_forwards_client_metadata_to_responses_websocket_request_body ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, @@ -657,34 +591,6 @@ async fn turn_start_forwards_client_metadata_to_responses_websocket_request_body Ok(()) } -fn create_config_toml( - codex_home: &Path, - server_uri: &str, - supports_websockets: bool, -) -> 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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -supports_websockets = {supports_websockets} -"# - ), - ) -} - async fn fork_fake_rollout_thread( mcp: &mut TestAppServer, source_thread_id: String, @@ -696,12 +602,7 @@ async fn fork_fake_rollout_thread( ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_req)), - ) - .await??; - to_response::(fork_resp) + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_req)).await? } fn parse_json_header(value: &str) -> serde_json::Value { diff --git a/codex-rs/app-server/tests/suite/v2/command_exec.rs b/codex-rs/app-server/tests/suite/v2/command_exec.rs index edd704c9fb22..f3955e7ffc7a 100644 --- a/codex-rs/app-server/tests/suite/v2/command_exec.rs +++ b/codex-rs/app-server/tests/suite/v2/command_exec.rs @@ -2,7 +2,6 @@ use anyhow::Context; use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::to_response; use base64::Engine; use base64::engine::general_purpose::STANDARD; use codex_app_server_protocol::CommandExecOutputDeltaNotification; @@ -45,9 +44,8 @@ async fn command_exec_without_streams_can_be_terminated() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let process_id = "sleep-1".to_string(); let command_request_id = mcp @@ -77,10 +75,7 @@ async fn command_exec_without_streams_can_be_terminated() -> Result<()> { .await?; assert_eq!(terminate_response.result, serde_json::json!({})); - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_ne!( response.exit_code, 0, "terminated command should not succeed" @@ -99,9 +94,8 @@ async fn command_exec_without_process_id_keeps_buffered_compatibility() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -126,10 +120,7 @@ async fn command_exec_without_process_id_keeps_buffered_compatibility() -> Resul }) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response, CommandExecResponse { @@ -152,9 +143,8 @@ async fn command_exec_env_overrides_merge_with_server_environment_and_support_un .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("COMMAND_EXEC_BASELINE", Some("server"))]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -186,10 +176,7 @@ async fn command_exec_env_overrides_merge_with_server_environment_and_support_un }) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response, CommandExecResponse { @@ -210,9 +197,8 @@ async fn command_exec_accepts_permission_profile() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -237,10 +223,7 @@ async fn command_exec_accepts_permission_profile() -> Result<()> { }) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response, CommandExecResponse { @@ -265,9 +248,8 @@ async fn command_exec_permission_profile_starts_selected_network_proxy() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -292,10 +274,7 @@ async fn command_exec_permission_profile_starts_selected_network_proxy() -> Resu }) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response, CommandExecResponse { @@ -317,9 +296,8 @@ async fn command_exec_permission_profile_does_not_reuse_default_network_proxy() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -344,10 +322,7 @@ async fn command_exec_permission_profile_does_not_reuse_default_network_proxy() }) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response, CommandExecResponse { @@ -379,9 +354,8 @@ async fn command_exec_permission_profile_project_roots_use_command_cwd() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -406,10 +380,7 @@ async fn command_exec_permission_profile_project_roots_use_command_cwd() -> Resu }) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response.exit_code, 0, "parent cwd write should fail under command project-root profile: {response:?}" @@ -435,9 +406,8 @@ async fn command_exec_returns_error_when_local_environment_is_disabled() -> Resu .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(CODEX_EXEC_SERVER_URL_ENV_VAR, Some("none"))]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -474,9 +444,8 @@ async fn command_exec_rejects_sandbox_policy_with_permission_profile() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -516,9 +485,8 @@ async fn command_exec_rejects_disable_timeout_with_timeout_ms() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -558,9 +526,8 @@ async fn command_exec_rejects_disable_output_cap_with_output_bytes_cap() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -600,9 +567,8 @@ async fn command_exec_rejects_negative_timeout_ms() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -642,9 +608,8 @@ async fn command_exec_without_process_id_rejects_streaming() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -684,9 +649,8 @@ async fn command_exec_non_streaming_respects_output_cap() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let command_request_id = mcp .send_command_exec_request(CommandExecParams { @@ -711,10 +675,7 @@ async fn command_exec_non_streaming_respects_output_cap() -> Result<()> { }) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response, CommandExecResponse { @@ -735,9 +696,8 @@ async fn command_exec_streaming_does_not_buffer_output() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let process_id = "stream-cap-1".to_string(); let command_request_id = mcp @@ -781,10 +741,7 @@ async fn command_exec_streaming_does_not_buffer_output() -> Result<()> { .await?; assert_eq!(terminate_response.result, serde_json::json!({})); - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_ne!( response.exit_code, 0, "terminated command should not succeed" @@ -803,9 +760,8 @@ async fn command_exec_pipe_streams_output_and_accepts_write() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let process_id = "pipe-1".to_string(); let command_request_id = mcp @@ -859,10 +815,7 @@ async fn command_exec_pipe_streams_output_and_accepts_write() -> Result<()> { ) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!( response, CommandExecResponse { @@ -883,9 +836,8 @@ async fn command_exec_tty_implies_streaming_and_reports_pty_output() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let process_id = "tty-1".to_string(); let command_request_id = mcp @@ -939,10 +891,7 @@ async fn command_exec_tty_implies_streaming_and_reports_pty_output() -> Result<( ) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!(response.exit_code, 0); assert_eq!(response.stdout, ""); assert_eq!(response.stderr, ""); @@ -958,9 +907,8 @@ async fn command_exec_tty_supports_initial_size_and_resize() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let process_id = "tty-size-1".to_string(); let command_request_id = mcp @@ -1031,10 +979,7 @@ async fn command_exec_tty_supports_initial_size_and_resize() -> Result<()> { ) .await?; - let response = mcp - .read_stream_until_response_message(RequestId::Integer(command_request_id)) - .await?; - let response: CommandExecResponse = to_response(response)?; + let response: CommandExecResponse = mcp.read_response(command_request_id).await?; assert_eq!(response.exit_code, 0); assert_eq!(response.stdout, ""); assert_eq!(response.stderr, ""); @@ -1130,10 +1075,7 @@ async fn command_exec_process_ids_are_connection_scoped_and_disconnect_terminate async fn read_command_exec_delta( mcp: &mut TestAppServer, ) -> Result { - let notification = mcp - .read_stream_until_notification_message("command/exec/outputDelta") - .await?; - decode_delta_notification(notification) + mcp.read_notification("command/exec/outputDelta").await } async fn wait_for_command_exec_output_contains( diff --git a/codex-rs/app-server/tests/suite/v2/compaction.rs b/codex-rs/app-server/tests/suite/v2/compaction.rs index 494c8b89446a..f1637d13ba67 100644 --- a/codex-rs/app-server/tests/suite/v2/compaction.rs +++ b/codex-rs/app-server/tests/suite/v2/compaction.rs @@ -7,18 +7,14 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; -use app_test_support::write_mock_responses_config_toml; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RawResponseCompletedNotification; use codex_app_server_protocol::RequestId; -use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ThreadCompactStartParams; use codex_app_server_protocol::ThreadCompactStartResponse; use codex_app_server_protocol::ThreadItem; @@ -36,7 +32,6 @@ use codex_protocol::models::ResponseItem; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; -use std::collections::BTreeMap; use tempfile::TempDir; use tokio::time::timeout; @@ -74,21 +69,12 @@ async fn auto_compaction_local_emits_started_and_completed_items() -> Result<()> responses::mount_sse_sequence(&server, vec![sse1, sse2, sse3, sse4]).await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::default(), - AUTO_COMPACT_LIMIT, - /*requires_openai_auth*/ None, - "mock_provider", - COMPACT_PROMPT, - )?; + compaction_config(&server.uri(), AUTO_COMPACT_LIMIT).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_id = start_thread(&mut mcp).await?; for message in ["first", "second", "third"] { @@ -155,15 +141,11 @@ async fn auto_compaction_remote_emits_started_and_completed_items() -> Result<() .await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::from([(Feature::RemoteCompactionV2, false)]), - REMOTE_AUTO_COMPACT_LIMIT, - Some(true), - "mock_provider", - COMPACT_PROMPT, - )?; + compaction_config(&server.uri(), REMOTE_AUTO_COMPACT_LIMIT) + .disable_feature(Feature::RemoteCompactionV2) + .with_provider_name("OpenAI") + .with_provider_config("requires_openai_auth = true") + .write(codex_home.path())?; write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("access-chatgpt").plan_type("pro"), @@ -173,9 +155,8 @@ async fn auto_compaction_remote_emits_started_and_completed_items() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_id = start_thread(&mut mcp).await?; for message in ["first", "second", "third"] { @@ -270,21 +251,12 @@ async fn thread_compact_start_triggers_compaction_and_returns_empty_response() - responses::mount_sse_sequence(&server, vec![sse]).await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::default(), - AUTO_COMPACT_LIMIT, - /*requires_openai_auth*/ None, - "mock_provider", - COMPACT_PROMPT, - )?; + compaction_config(&server.uri(), AUTO_COMPACT_LIMIT).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -293,36 +265,23 @@ async fn thread_compact_start_triggers_compaction_and_returns_empty_response() - ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let thread_id = thread.id; let compact_id = mcp .send_thread_compact_start_request(ThreadCompactStartParams { thread_id: thread_id.clone(), }) .await?; - let compact_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(compact_id)), - ) - .await??; - let _compact: ThreadCompactStartResponse = - to_response::(compact_resp)?; + let _: ThreadCompactStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(compact_id)).await??; let started = wait_for_context_compaction_started(&mut mcp).await?; - let raw_completed = timeout( + let raw_completed: RawResponseCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("rawResponse/completed"), + mcp.read_notification("rawResponse/completed"), ) .await??; - let raw_completed: ServerNotification = raw_completed.try_into()?; - let ServerNotification::RawResponseCompleted(raw_completed) = raw_completed else { - anyhow::bail!("expected rawResponse/completed notification"); - }; let completed = wait_for_context_compaction_completed(&mut mcp).await?; let ThreadItem::ContextCompaction { id: started_id } = started.item else { @@ -361,21 +320,12 @@ async fn thread_compact_start_rejects_invalid_thread_id() -> Result<()> { let server = responses::start_mock_server().await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::default(), - AUTO_COMPACT_LIMIT, - /*requires_openai_auth*/ None, - "mock_provider", - COMPACT_PROMPT, - )?; + compaction_config(&server.uri(), AUTO_COMPACT_LIMIT).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_compact_start_request(ThreadCompactStartParams { @@ -400,21 +350,12 @@ async fn thread_compact_start_rejects_unknown_thread_id() -> Result<()> { let server = responses::start_mock_server().await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::default(), - AUTO_COMPACT_LIMIT, - /*requires_openai_auth*/ None, - "mock_provider", - COMPACT_PROMPT, - )?; + compaction_config(&server.uri(), AUTO_COMPACT_LIMIT).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_compact_start_request(ThreadCompactStartParams { @@ -440,12 +381,8 @@ async fn start_thread(mcp: &mut TestAppServer) -> Result { ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_id)).await??; Ok(thread.id) } @@ -465,25 +402,19 @@ async fn send_turn_and_wait( ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_id)).await??; wait_for_turn_completed(mcp, &turn.id).await?; Ok(turn.id) } async fn wait_for_turn_completed(mcp: &mut TestAppServer, turn_id: &str) -> Result<()> { loop { - let notification: JSONRPCNotification = timeout( + let completed: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed: TurnCompletedNotification = - serde_json::from_value(notification.params.clone().expect("turn/completed params"))?; if completed.turn.id == turn_id { return Ok(()); } @@ -494,13 +425,8 @@ async fn wait_for_context_compaction_started( mcp: &mut TestAppServer, ) -> Result { loop { - let notification: JSONRPCNotification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/started"), - ) - .await??; let started: ItemStartedNotification = - serde_json::from_value(notification.params.clone().expect("item/started params"))?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("item/started")).await??; if let ThreadItem::ContextCompaction { .. } = started.item { return Ok(started); } @@ -511,13 +437,11 @@ async fn wait_for_context_compaction_completed( mcp: &mut TestAppServer, ) -> Result { loop { - let notification: JSONRPCNotification = timeout( + let completed: ItemCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/completed"), + mcp.read_notification("item/completed"), ) .await??; - let completed: ItemCompletedNotification = - serde_json::from_value(notification.params.clone().expect("item/completed params"))?; if let ThreadItem::ContextCompaction { .. } = completed.item { return Ok(completed); } @@ -527,3 +451,11 @@ async fn wait_for_context_compaction_completed( fn parse_json_header(value: &str) -> serde_json::Value { serde_json::from_str(value).expect("turn metadata should be JSON") } + +fn compaction_config(server_uri: &str, auto_compact_limit: i64) -> MockResponsesConfig { + MockResponsesConfig::new(server_uri) + .with_root_config(&format!( + "compact_prompt = \"{COMPACT_PROMPT}\"\nmodel_auto_compact_token_limit = {auto_compact_limit}" + )) + .with_provider_config("supports_websockets = false") +} diff --git a/codex-rs/app-server/tests/suite/v2/config_rpc.rs b/codex-rs/app-server/tests/suite/v2/config_rpc.rs index 3c7704f2236b..6407005c19f9 100644 --- a/codex-rs/app-server/tests/suite/v2/config_rpc.rs +++ b/codex-rs/app-server/tests/suite/v2/config_rpc.rs @@ -2,7 +2,6 @@ use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::test_path_buf_with_windows; use app_test_support::test_tmp_path_buf; -use app_test_support::to_response; use codex_app_server_protocol::AppConfig; use codex_app_server_protocol::AppToolApproval; use codex_app_server_protocol::ApprovalsReviewer; @@ -20,7 +19,6 @@ use codex_app_server_protocol::ConfigWriteResponse; use codex_app_server_protocol::ConfiguredHookHandler; use codex_app_server_protocol::ForcedChatgptWorkspaceIds; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::MergeStrategy; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SandboxMode; @@ -69,17 +67,12 @@ additionalContextLimit = 4096 let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_config_requirements_read_request().await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ConfigRequirementsReadResponse = to_response(response)?; + let response: ConfigRequirementsReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let requirements = response .requirements .expect("managed requirements should be returned"); @@ -117,17 +110,12 @@ service_tier = "fast" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_config_requirements_read_request().await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ConfigRequirementsReadResponse = to_response(response)?; + let response: ConfigRequirementsReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let defaults = response .requirements @@ -159,9 +147,8 @@ sandbox_mode = "workspace-write" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -169,16 +156,11 @@ sandbox_mode = "workspace-write" cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ConfigReadResponse { config, origins, layers, - } = to_response(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(config.model.as_deref(), Some("gpt-user")); assert_eq!( @@ -213,9 +195,8 @@ allowed_domains = ["example.com"] let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -223,16 +204,11 @@ allowed_domains = ["example.com"] cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ConfigReadResponse { config, origins, layers, - } = to_response(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let tools = config.tools.expect("tools present"); assert_eq!( @@ -288,9 +264,8 @@ forced_chatgpt_workspace_id = "{WORKSPACE_ID}" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -298,12 +273,8 @@ forced_chatgpt_workspace_id = "{WORKSPACE_ID}" cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ConfigReadResponse { config, .. } = to_response(resp)?; + let ConfigReadResponse { config, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( config.forced_chatgpt_workspace_id, @@ -331,9 +302,8 @@ forced_chatgpt_workspace_id = ["{WORKSPACE_ID_A}", "{WORKSPACE_ID_B}"] let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -341,12 +311,8 @@ forced_chatgpt_workspace_id = ["{WORKSPACE_ID_A}", "{WORKSPACE_ID_B}"] cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ConfigReadResponse { config, .. } = to_response(resp)?; + let ConfigReadResponse { config, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( config.forced_chatgpt_workspace_id, @@ -377,9 +343,8 @@ location = { country = "US", city = "New York", timezone = "America/New_York" } let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -387,12 +352,8 @@ location = { country = "US", city = "New York", timezone = "America/New_York" } cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ConfigReadResponse { config, .. } = to_response(resp)?; + let ConfigReadResponse { config, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( config.tools.expect("tools present").web_search, @@ -425,9 +386,8 @@ web_search = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -435,12 +395,8 @@ web_search = true cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ConfigReadResponse { config, .. } = to_response(resp)?; + let ConfigReadResponse { config, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(config.tools.expect("tools present").web_search, None,); @@ -470,9 +426,8 @@ default_tools_approval_mode = "prompt" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -480,16 +435,11 @@ default_tools_approval_mode = "prompt" cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ConfigReadResponse { config, origins, layers, - } = to_response(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( config.apps, @@ -598,9 +548,8 @@ width = 320 let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -608,12 +557,8 @@ width = 320 cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ConfigReadResponse { config, .. } = to_response(resp)?; + let ConfigReadResponse { config, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let desktop = config.desktop.expect("desktop settings present"); assert_eq!(desktop.get("appearanceTheme"), Some(&json!("dark"))); @@ -649,9 +594,8 @@ model_reasoning_effort = "high" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -659,14 +603,9 @@ model_reasoning_effort = "high" cwd: Some(workspace.path().to_string_lossy().into_owned()), }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ConfigReadResponse { config, origins, .. - } = to_response(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(config.model_reasoning_effort, Some(ReasoningEffort::High)); assert_eq!( @@ -727,9 +666,8 @@ writable_roots = [{}] "CODEX_APP_SERVER_MANAGED_CONFIG_PATH", Some(&managed_path_str), )]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_config_read_request(ConfigReadParams { @@ -737,16 +675,11 @@ writable_roots = [{}] cwd: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ConfigReadResponse { config, origins, layers, - } = to_response(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(config.model.as_deref(), Some("gpt-system")); assert_eq!( @@ -820,9 +753,8 @@ model = "gpt-old" let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_config_read_request(ConfigReadParams { @@ -830,12 +762,8 @@ model = "gpt-old" cwd: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let read: ConfigReadResponse = to_response(read_resp)?; + let read: ConfigReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; let expected_version = read.origins.get("model").map(|m| m.version.clone()); let write_id = mcp @@ -847,12 +775,8 @@ model = "gpt-old" expected_version, }) .await?; - let write_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let write: ConfigWriteResponse = to_response(write_resp)?; + let write: ConfigWriteResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(write_id)).await??; let expected_file_path = AbsolutePathBuf::resolve_path_against_base("config.toml", codex_home); assert_eq!(write.status, WriteStatus::Ok); @@ -865,12 +789,8 @@ model = "gpt-old" cwd: None, }) .await?; - let verify_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(verify_id)), - ) - .await??; - let verify: ConfigReadResponse = to_response(verify_resp)?; + let verify: ConfigReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(verify_id)).await??; assert_eq!(verify.config.model.as_deref(), Some("gpt-new")); Ok(()) @@ -885,9 +805,8 @@ async fn config_value_write_updates_desktop_settings() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let write_id = mcp .send_config_value_write_request(ConfigValueWriteParams { @@ -898,12 +817,8 @@ async fn config_value_write_updates_desktop_settings() -> Result<()> { expected_version: None, }) .await?; - let write_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let write: ConfigWriteResponse = to_response(write_resp)?; + let write: ConfigWriteResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(write_id)).await??; assert_eq!(write.status, WriteStatus::Ok); let read_id = mcp @@ -912,12 +827,8 @@ async fn config_value_write_updates_desktop_settings() -> Result<()> { cwd: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let read: ConfigReadResponse = to_response(read_resp)?; + let read: ConfigReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; let desktop = read.config.desktop.expect("desktop settings present"); assert_eq!(desktop.get("appearanceTheme"), Some(&json!("dark"))); @@ -938,9 +849,8 @@ model = "gpt-old" let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let write_id = mcp .send_config_value_write_request(ConfigValueWriteParams { @@ -958,20 +868,12 @@ model = "gpt-old" }) .await?; - let write_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let write: ConfigWriteResponse = to_response(write_resp)?; + let write: ConfigWriteResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(write_id)).await??; assert_eq!(write.status, WriteStatus::Ok); - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let read: ConfigReadResponse = to_response(read_resp)?; + let read: ConfigReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(read.config.model.as_deref(), Some("gpt-new")); Ok(()) @@ -990,9 +892,8 @@ model = "gpt-old" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let write_id = mcp .send_config_value_write_request(ConfigValueWriteParams { @@ -1029,9 +930,8 @@ async fn config_batch_write_applies_multiple_edits() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let writable_root = test_tmp_path_buf(); let batch_id = mcp @@ -1056,12 +956,8 @@ async fn config_batch_write_applies_multiple_edits() -> Result<()> { reload_user_config: false, }) .await?; - let batch_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(batch_id)), - ) - .await??; - let batch_write: ConfigWriteResponse = to_response(batch_resp)?; + let batch_write: ConfigWriteResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(batch_id)).await??; assert_eq!(batch_write.status, WriteStatus::Ok); let expected_file_path = AbsolutePathBuf::resolve_path_against_base("config.toml", codex_home); assert_eq!(batch_write.file_path, expected_file_path); @@ -1072,12 +968,8 @@ async fn config_batch_write_applies_multiple_edits() -> Result<()> { cwd: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let read: ConfigReadResponse = to_response(read_resp)?; + let read: ConfigReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(read.config.sandbox_mode, Some(SandboxMode::WorkspaceWrite)); let sandbox = read .config @@ -1105,9 +997,8 @@ model = "gpt-5.3-spark" let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let batch_id = mcp .send_config_batch_write_request(ConfigBatchWriteParams { @@ -1165,9 +1056,8 @@ async fn config_batch_write_updates_multiple_desktop_settings() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let batch_id = mcp .send_config_batch_write_request(ConfigBatchWriteParams { @@ -1191,12 +1081,8 @@ async fn config_batch_write_updates_multiple_desktop_settings() -> Result<()> { reload_user_config: false, }) .await?; - let batch_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(batch_id)), - ) - .await??; - let batch_write: ConfigWriteResponse = to_response(batch_resp)?; + let batch_write: ConfigWriteResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(batch_id)).await??; assert_eq!(batch_write.status, WriteStatus::Ok); let read_id = mcp @@ -1205,12 +1091,8 @@ async fn config_batch_write_updates_multiple_desktop_settings() -> Result<()> { cwd: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let read: ConfigReadResponse = to_response(read_resp)?; + let read: ConfigReadResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; let desktop = read.config.desktop.expect("desktop settings present"); assert_eq!(desktop.get("selected-avatar-id"), Some(&json!("codex"))); assert_eq!( diff --git a/codex-rs/app-server/tests/suite/v2/current_time.rs b/codex-rs/app-server/tests/suite/v2/current_time.rs index a8a05f7fe1b5..d3a012f0ea80 100644 --- a/codex-rs/app-server/tests/suite/v2/current_time.rs +++ b/codex-rs/app-server/tests/suite/v2/current_time.rs @@ -1,12 +1,9 @@ -use std::path::Path; - use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::CurrentTimeReadResponse; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -38,40 +35,38 @@ async fn current_time_read_round_trip_adds_reminder_to_model_input() -> Result<( ) .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_extra_config( + r#"[features.current_time_reminder] +enabled = true +reminder_interval_seconds = 1 +clock_source = "external" +"#, + ) + .write(codex_home.path())?; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; - let thread_request_id = app_server - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) + let ThreadStartResponse { thread, .. } = app_server + .start_thread(ThreadStartParams::default()) .await?; - let thread_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(thread_request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_response)?; - let turn_request_id = app_server - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - input: vec![UserInput::Text { - text: "What time is it?".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = app_server + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + input: vec![UserInput::Text { + text: "What time is it?".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(turn_request_id)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_response)?; let server_request = timeout( DEFAULT_READ_TIMEOUT, @@ -105,29 +100,3 @@ async fn current_time_read_round_trip_adds_reminder_to_model_input() -> Result<( ); Ok(()) } - -fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -model_provider = "mock_provider" - -[features.current_time_reminder] -enabled = true -reminder_interval_seconds = 1 -clock_source = "external" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs index 48a983ca08a6..349f4573d183 100644 --- a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs +++ b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs @@ -1,5 +1,6 @@ use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; @@ -35,7 +36,6 @@ use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; -use std::path::Path; use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; @@ -60,7 +60,7 @@ async fn thread_start_normalizes_legacy_dynamic_tools_into_model_request() -> Re let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -178,7 +178,7 @@ async fn thread_start_rejects_hidden_dynamic_tools_without_namespace() -> Result let server = MockServer::start().await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -220,7 +220,7 @@ async fn thread_start_rejects_invalid_dynamic_tool_inputs() -> Result<()> { let server = MockServer::start().await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -366,7 +366,7 @@ async fn dynamic_tool_call_round_trip_sends_text_content_items_to_model() -> Res let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -596,7 +596,7 @@ async fn start_function_dynamic_tool_call(call_id: &str) -> Result 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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 4062ac9c0bb6..5067eea4b2fa 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -1,7 +1,6 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; -use app_test_support::write_mock_responses_config_toml; use axum::Json; use axum::Router; use axum::body::Bytes; @@ -19,6 +18,7 @@ use codex_app_server_protocol::SelectedCapabilityRoot; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput; use codex_utils_path_uri::PathUri; use core_test_support::responses; @@ -40,7 +40,6 @@ use rmcp::transport::StreamableHttpService; use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; use serde_json::json; use std::borrow::Cow; -use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; @@ -114,15 +113,10 @@ async fn selected_executor_plugin_exposes_its_mcps_only_to_that_thread() -> Resu let _ = axum::serve(http_listener, http_router).await; }); let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &responses_server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; + MockResponsesConfig::new(&responses_server.uri()) + .with_root_config("compact_prompt = \"compact\"\nmodel_auto_compact_token_limit = 1024") + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let codex_bin = toml::Value::String( codex_utils_cargo_bin::cargo_bin("codex")? .to_string_lossy() @@ -180,9 +174,8 @@ HTTP_PROXY = {http_proxy} .with_codex_home(codex_home.path()) // This suite owns environments.toml to exercise explicit executor selection. .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; let selected_thread = start_thread( &mut app_server, @@ -226,12 +219,8 @@ startup_timeout_sec = 10 })), ) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: McpServerOauthLoginResponse = to_response(response)?; + let response: McpServerOauthLoginResponse = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(request_id)).await??; assert!( response .authorization_url @@ -269,13 +258,11 @@ startup_timeout_sec = 10 assert!(token_request.contains("grant_type=authorization_code")); assert!(token_request.contains("code=executor-test-code")); assert!(token_request.contains("code_verifier=")); - let notification = timeout( + let completed: McpServerOauthLoginCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_notification_message("mcpServer/oauthLogin/completed"), + app_server.read_notification("mcpServer/oauthLogin/completed"), ) .await??; - let completed: McpServerOauthLoginCompletedNotification = - serde_json::from_value(notification.params.expect("notification params"))?; assert_eq!( completed, McpServerOauthLoginCompletedNotification { @@ -322,11 +309,8 @@ startup_timeout_sec = 10 ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(request_id)).await??; timeout( DEFAULT_READ_TIMEOUT, app_server.read_stream_until_notification_message("turn/completed"), @@ -353,12 +337,8 @@ startup_timeout_sec = 10 meta: None, }) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: McpServerToolCallResponse = to_response(response)?; + let response: McpServerToolCallResponse = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(request_id)).await??; assert_eq!( response.structured_content, Some(json!({"echo": "ECHOING: hello over executor HTTP"})) @@ -373,12 +353,8 @@ startup_timeout_sec = 10 meta: None, }) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: McpServerToolCallResponse = to_response(response)?; + let response: McpServerToolCallResponse = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(request_id)).await??; assert_eq!( response .structured_content @@ -479,12 +455,8 @@ async fn mcp_server_names( thread_id: Some(thread_id), }) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ListMcpServerStatusResponse = to_response(response)?; + let response: ListMcpServerStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(request_id)).await??; Ok(response .data .into_iter() @@ -503,11 +475,7 @@ async fn start_thread( ..Default::default() }) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(response)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(request_id)).await??; Ok(thread.id) } diff --git a/codex-rs/app-server/tests/suite/v2/experimental_api.rs b/codex-rs/app-server/tests/suite/v2/experimental_api.rs index ca120bfed2fc..3bc29d07ebda 100644 --- a/codex-rs/app-server/tests/suite/v2/experimental_api.rs +++ b/codex-rs/app-server/tests/suite/v2/experimental_api.rs @@ -1,5 +1,6 @@ use anyhow::Result; use app_test_support::DEFAULT_CLIENT_NAME; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; @@ -20,7 +21,6 @@ use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_protocol::protocol::RealtimeOutputModality; use pretty_assertions::assert_eq; -use std::path::Path; use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; @@ -251,7 +251,7 @@ async fn realtime_webrtc_start_requires_experimental_api_capability() -> Result< async fn thread_start_mock_field_requires_experimental_api_capability() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -293,7 +293,7 @@ async fn thread_start_without_dynamic_tools_allows_without_experimental_api_capa -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -334,7 +334,7 @@ async fn thread_start_granular_approval_policy_requires_experimental_api_capabil { let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -393,26 +393,3 @@ fn assert_experimental_capability_error(error: JSONRPCError, reason: &str) { ); assert_eq!(error.error.data, None); } - -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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs b/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs index 8d66855d523d..425603bf37e9 100644 --- a/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs +++ b/codex-rs/app-server/tests/suite/v2/experimental_feature_list.rs @@ -2,9 +2,9 @@ use std::time::Duration; use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_repeating_assistant; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::ConfigReadParams; use codex_app_server_protocol::ConfigReadResponse; @@ -15,7 +15,6 @@ use codex_app_server_protocol::ExperimentalFeatureListParams; use codex_app_server_protocol::ExperimentalFeatureListResponse; use codex_app_server_protocol::ExperimentalFeatureStage; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -53,11 +52,9 @@ async fn experimental_feature_list_returns_feature_metadata_with_stage() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp .send_experimental_feature_list_request(ExperimentalFeatureListParams::default()) .await?; @@ -142,9 +139,8 @@ async fn experimental_feature_list_marks_apps_and_plugins_disabled_by_workspace_ .with_codex_home(codex_home.path()) .without_auto_env() .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_experimental_feature_list_request(ExperimentalFeatureListParams::default()) @@ -173,28 +169,12 @@ async fn experimental_feature_list_resolves_thread_project_config() -> Result<() let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; let workspace = TempDir::new()?; - let server_uri = server.uri(); let workspace_key = workspace.path().to_string_lossy().replace('\\', "\\\\"); - std::fs::write( - codex_home.path().join("config.toml"), - format!( - r#"model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -model_provider = "mock_provider" - -[projects."{workspace_key}"] -trust_level = "trusted" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - )?; + MockResponsesConfig::new(&server.uri()) + .with_extra_config(&format!( + "[projects.\"{workspace_key}\"]\ntrust_level = \"trusted\"" + )) + .write(codex_home.path())?; let project_config_dir = workspace.path().join(".codex"); std::fs::create_dir_all(&project_config_dir)?; std::fs::write( @@ -207,9 +187,8 @@ memories = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let thread_start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -245,9 +224,8 @@ async fn experimental_feature_list_rejects_unknown_thread_id() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_experimental_feature_list_request(ExperimentalFeatureListParams { @@ -284,9 +262,8 @@ async fn experimental_feature_enablement_set_applies_to_global_and_thread_config let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let actual = set_experimental_feature_enablement( &mut mcp, @@ -325,9 +302,8 @@ async fn experimental_feature_enablement_set_does_not_override_user_config() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let actual = set_experimental_feature_enablement( &mut mcp, @@ -360,9 +336,8 @@ async fn experimental_feature_enablement_set_only_updates_named_features() -> Re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; set_experimental_feature_enablement( &mut mcp, @@ -439,9 +414,8 @@ async fn experimental_feature_enablement_set_allows_remote_control() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let remote_control_enabled = false; let enablement = BTreeMap::from([("remote_control".to_string(), remote_control_enabled)]); @@ -461,9 +435,8 @@ async fn experimental_feature_enablement_set_empty_map_is_no_op() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; set_experimental_feature_enablement( &mut mcp, @@ -498,9 +471,8 @@ async fn experimental_feature_enablement_set_ignores_invalid_features() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let actual = set_experimental_feature_enablement( &mut mcp, @@ -549,10 +521,5 @@ async fn read_config(mcp: &mut TestAppServer, cwd: Option) -> Result(mcp: &mut TestAppServer, request_id: i64) -> Result { - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await? } diff --git a/codex-rs/app-server/tests/suite/v2/external_agent_config.rs b/codex-rs/app-server/tests/suite/v2/external_agent_config.rs index fd03838e7c74..e34af5279ebe 100644 --- a/codex-rs/app-server/tests/suite/v2/external_agent_config.rs +++ b/codex-rs/app-server/tests/suite/v2/external_agent_config.rs @@ -2,12 +2,11 @@ use std::time::Duration; use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::start_analytics_events_server; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; -use app_test_support::write_mock_responses_config_toml; use codex_app_server_protocol::ExternalAgentConfigDetectResponse; use codex_app_server_protocol::ExternalAgentConfigImportCompletedNotification; use codex_app_server_protocol::ExternalAgentConfigImportHistoriesReadResponse; @@ -16,7 +15,6 @@ use codex_app_server_protocol::ExternalAgentConfigImportResponse; use codex_app_server_protocol::ExternalAgentConfigMigrationItemType; use codex_app_server_protocol::ExternalAgentImportedConnectorCandidate; use codex_app_server_protocol::ExternalAgentImportedConnectorSource; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PluginListParams; use codex_app_server_protocol::PluginListResponse; use codex_app_server_protocol::RequestId; @@ -28,11 +26,11 @@ use codex_app_server_protocol::ThreadReadResponse; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput; use codex_config::types::AuthCredentialsStoreMode; use core_test_support::responses; use pretty_assertions::assert_eq; -use std::collections::BTreeMap; use std::path::Path; use std::path::PathBuf; use tempfile::TempDir; @@ -85,9 +83,8 @@ async fn external_agent_config_detect_accepts_migration_source_and_defaults_unkn .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let mut responses = Vec::new(); for params in [ @@ -108,12 +105,13 @@ async fn external_agent_config_detect_accepts_migration_source_and_defaults_unkn let request_id = mcp .send_raw_request("externalAgentConfig/detect", Some(params)) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - responses.push(to_response::(response)?); + responses.push( + timeout( + DEFAULT_TIMEOUT, + mcp.read_response::(request_id), + ) + .await??, + ); } assert_eq!(responses[0].items.len(), 1); @@ -138,9 +136,8 @@ async fn external_agent_config_migration_source_drives_detect_and_import() -> Re .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -151,12 +148,8 @@ async fn external_agent_config_migration_source_drives_detect_and_import() -> Re })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items.len(), 1); assert_eq!( detected.items[0].item_type, @@ -172,20 +165,14 @@ async fn external_agent_config_migration_source_drives_detect_and_import() -> Re })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 1); assert_eq!(completed.item_type_results[0].successes.len(), 1); @@ -209,9 +196,8 @@ async fn external_agent_config_import_source_remains_attribution_only() -> Resul .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -219,12 +205,8 @@ async fn external_agent_config_import_source_remains_attribution_only() -> Resul Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items.len(), 1); assert_eq!( detected.items[0].item_type, @@ -240,20 +222,14 @@ async fn external_agent_config_import_source_remains_attribution_only() -> Resul })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 1); assert_eq!(completed.item_type_results[0].successes.len(), 1); @@ -373,9 +349,8 @@ source = {:?} .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -386,12 +361,8 @@ source = {:?} })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items.len(), 2); assert!( detected @@ -415,20 +386,14 @@ source = {:?} })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 2); assert!( @@ -454,12 +419,8 @@ source = {:?} ancestor_thread_id: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadListResponse = to_response(response)?; + let response: ThreadListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let thread = response.data.first().expect("imported session"); assert_eq!(thread.cwd.as_path(), project_root); assert_eq!(thread.preview, "first request"); @@ -471,12 +432,8 @@ source = {:?} include_turns: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadReadResponse = to_response(response)?; + let response: ThreadReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.thread.turns.len(), 1); let imported_items = &response.thread.turns[0].items; assert_eq!(imported_items.len(), 3); @@ -502,12 +459,8 @@ source = {:?} force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let marketplace = response .marketplaces .iter() @@ -546,9 +499,8 @@ async fn external_agent_config_import_sends_completion_notification_for_sync_onl ("HOME", Some(home_dir.as_str())), ("CODEX_SQLITE_HOME", Some(sqlite_home_dir.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -563,21 +515,14 @@ async fn external_agent_config_import_sends_completion_notification_for_sync_onl ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let progress = timeout( + let progress: ExternalAgentConfigImportProgressNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/progress"), + mcp.read_notification("externalAgentConfig/import/progress"), ) .await??; - assert_eq!(progress.method, "externalAgentConfig/import/progress"); - let progress: ExternalAgentConfigImportProgressNotification = - serde_json::from_value(progress.params.expect("progress params"))?; assert_eq!(progress.import_id, import_id); assert_eq!(progress.item_type_results.len(), 1); assert_eq!( @@ -585,14 +530,11 @@ async fn external_agent_config_import_sends_completion_notification_for_sync_onl ExternalAgentConfigMigrationItemType::Config ); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); let state_db = codex_state::StateRuntime::init(sqlite_home.path().to_path_buf(), "mock_provider".into()) @@ -626,12 +568,8 @@ async fn external_agent_config_import_sends_completion_notification_for_sync_onl /*params*/ None, ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportHistoriesReadResponse = to_response(response)?; + let response: ExternalAgentConfigImportHistoriesReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.connectors, Vec::new()); let entry = response .data @@ -664,9 +602,8 @@ async fn external_agent_memory_import_requires_feature_config() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -674,12 +611,8 @@ async fn external_agent_memory_import_requires_feature_config() -> Result<()> { Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items, Vec::new()); let request_id = mcp @@ -727,9 +660,8 @@ async fn external_agent_config_detects_non_memory_items_when_config_reload_fails .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; std::fs::write( codex_home.path().join("config.toml"), "this is not valid = [toml", @@ -741,12 +673,8 @@ async fn external_agent_config_detects_non_memory_items_when_config_reload_fails Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( detected .items @@ -792,9 +720,8 @@ async fn external_agent_config_detects_and_imports_project_memory_files() -> Res .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; for details in [serde_json::json!({}), serde_json::json!({ "memory": [] })] { let request_id = mcp @@ -827,12 +754,8 @@ async fn external_agent_config_detects_and_imports_project_memory_files() -> Res Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let mut detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let mut detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; detected .items .retain(|item| item.item_type == ExternalAgentConfigMigrationItemType::Memory); @@ -867,20 +790,14 @@ async fn external_agent_config_detects_and_imports_project_memory_files() -> Res Some(serde_json::json!({ "migrationItems": detected.items })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 1); let memory_result = &completed.item_type_results[0]; @@ -952,12 +869,8 @@ async fn external_agent_config_detects_and_imports_project_memory_files() -> Res Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let mut detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let mut detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; detected .items .retain(|item| item.item_type == ExternalAgentConfigMigrationItemType::Memory); @@ -977,20 +890,14 @@ async fn external_agent_config_detects_and_imports_project_memory_files() -> Res Some(serde_json::json!({ "migrationItems": detected.items })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 1); assert_eq!(completed.item_type_results[0].failures, Vec::new()); @@ -1060,9 +967,8 @@ async fn external_agent_config_import_reports_failed_sync_import_in_completion() Some(analytics_capture_file.as_str()), ), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -1086,21 +992,15 @@ async fn external_agent_config_import_reports_failed_sync_import_in_completion() ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); let config_result = completed .item_type_results @@ -1203,9 +1103,8 @@ async fn external_agent_config_import_completed_tracks_analytics_event() -> Resu .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -1229,21 +1128,15 @@ async fn external_agent_config_import_completed_tracks_analytics_event() -> Resu })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 1); assert_eq!(completed.item_type_results[0].successes.len(), 0); @@ -1365,9 +1258,8 @@ async fn external_agent_config_import_reinstalls_plugins_from_known_marketplaces .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -1375,12 +1267,8 @@ async fn external_agent_config_import_reinstalls_plugins_from_known_marketplaces Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items.len(), 1); assert_eq!( detected.items[0].item_type, @@ -1403,22 +1291,15 @@ async fn external_agent_config_import_reinstalls_plugins_from_known_marketplaces Some(serde_json::json!({ "migrationItems": detected.items })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 1); let plugin_result = &completed.item_type_results[0]; @@ -1477,12 +1358,8 @@ async fn external_agent_config_import_reinstalls_plugins_from_known_marketplaces force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let plugin = response .marketplaces .iter() @@ -1526,9 +1403,8 @@ async fn external_agent_config_import_sends_completion_notification_after_pendin .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -1549,21 +1425,14 @@ async fn external_agent_config_import_sends_completion_notification_after_pendin ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); Ok(()) @@ -1573,7 +1442,7 @@ async fn external_agent_config_import_sends_completion_notification_after_pendin async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("follow-up answer").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let project_root = codex_home.path().join("repo"); let recent_timestamp = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true); let session_dir = external_agent_home(codex_home.path()).join("projects/repo"); @@ -1630,9 +1499,8 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -1642,12 +1510,8 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items.len(), 1); assert_eq!( detected.items[0] @@ -1664,21 +1528,14 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { Some(serde_json::json!({ "migrationItems": detected.items })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); assert_eq!(completed.item_type_results.len(), 1); let session_result = &completed.item_type_results[0]; @@ -1711,12 +1568,8 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { /*params*/ None, ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportHistoriesReadResponse = to_response(response)?; + let response: ExternalAgentConfigImportHistoriesReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response.connectors, vec![ExternalAgentImportedConnectorCandidate { @@ -1742,12 +1595,8 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { ancestor_thread_id: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadListResponse = to_response(response)?; + let response: ThreadListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let thread = response .data .first() @@ -1763,12 +1612,8 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { include_turns: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadReadResponse = to_response(response)?; + let response: ThreadReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.thread.turns.len(), 2); let control_items = &response.thread.turns[0].items; assert_eq!(control_items.len(), 1); @@ -1814,12 +1659,7 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadResumeResponse = to_response(response)?; + let _: ThreadResumeResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let request_id = mcp .send_turn_start_request(TurnStartParams { @@ -1832,11 +1672,7 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { ..Default::default() }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -1849,12 +1685,8 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { include_turns: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadReadResponse = to_response(response)?; + let response: ThreadReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.thread.turns.len(), 3); match &response.thread.turns[2].items[1] { ThreadItem::AgentMessage { text, .. } => assert_eq!(text, "follow-up answer"), @@ -1868,7 +1700,7 @@ async fn external_agent_config_import_creates_session_rollouts() -> Result<()> { async fn external_agent_config_import_does_not_initialize_required_mcp() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("unused").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut config = std::fs::read_to_string(codex_home.path().join("config.toml"))?; config.push_str( r#" @@ -1900,9 +1732,8 @@ required = true .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -1923,11 +1754,8 @@ required = true })), ) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + let _: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), @@ -1950,12 +1778,8 @@ required = true ancestor_thread_id: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadListResponse = to_response(response)?; + let response: ThreadListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.data.len(), 1); Ok(()) @@ -1966,7 +1790,7 @@ async fn external_agent_config_import_accepts_detected_session_payload_after_res { let server = create_mock_responses_server_repeating_assistant("unused").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let project_root = codex_home.path().join("repo"); let recent_timestamp = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true); let session_dir = external_agent_home(codex_home.path()).join("projects/repo"); @@ -1989,9 +1813,8 @@ async fn external_agent_config_import_accepts_detected_session_payload_after_res .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -2012,21 +1835,14 @@ async fn external_agent_config_import_accepts_detected_session_payload_after_res })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); let request_id = mcp @@ -2045,12 +1861,8 @@ async fn external_agent_config_import_accepts_detected_session_payload_after_res ancestor_thread_id: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadListResponse = to_response(response)?; + let response: ThreadListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.data.len(), 1); Ok(()) @@ -2060,7 +1872,7 @@ async fn external_agent_config_import_accepts_detected_session_payload_after_res async fn external_agent_config_import_skips_already_imported_session_versions() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("unused").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let project_root = codex_home.path().join("repo"); let recent_timestamp = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true); let session_dir = external_agent_home(codex_home.path()).join("projects/repo"); @@ -2083,9 +1895,8 @@ async fn external_agent_config_import_skips_already_imported_session_versions() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -2093,12 +1904,8 @@ async fn external_agent_config_import_skips_already_imported_session_versions() Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; for _ in 0..2 { let request_id = mcp @@ -2107,21 +1914,14 @@ async fn external_agent_config_import_skips_already_imported_session_versions() Some(serde_json::json!({ "migrationItems": detected.items.clone() })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); } @@ -2141,12 +1941,8 @@ async fn external_agent_config_import_skips_already_imported_session_versions() ancestor_thread_id: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadListResponse = to_response(response)?; + let response: ThreadListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.data.len(), 1); Ok(()) @@ -2158,7 +1954,7 @@ async fn external_agent_config_import_returns_before_background_session_import_f -> Result<()> { let server = create_mock_responses_server_repeating_assistant("unused").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let project_root = codex_home.path().join("repo"); let recent_timestamp = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true); let session_dir = external_agent_home(codex_home.path()).join("projects/repo"); @@ -2179,9 +1975,8 @@ async fn external_agent_config_import_returns_before_background_session_import_f .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -2189,12 +1984,8 @@ async fn external_agent_config_import_returns_before_background_session_import_f Some(serde_json::json!({ "includeHome": true })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items.len(), 1); let detected_items = detected.items; @@ -2210,12 +2001,8 @@ async fn external_agent_config_import_returns_before_background_session_import_f Some(serde_json::json!({ "migrationItems": detected_items.clone() })), ) .await?; - let response: JSONRPCResponse = timeout( - Duration::from_secs(5), - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(Duration::from_secs(5), mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); assert!( @@ -2234,12 +2021,11 @@ async fn external_agent_config_import_returns_before_background_session_import_f Some(serde_json::json!({ "migrationItems": detected_items })), ) .await?; - let response: JSONRPCResponse = timeout( + let response: ExternalAgentConfigImportResponse = timeout( Duration::from_secs(5), - mcp.read_stream_until_response_message(RequestId::Integer(duplicate_request_id)), + mcp.read_response(duplicate_request_id), ) .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; let duplicate_import_id = assert_import_response(response); let mut completed_import_ids = Vec::new(); @@ -2253,14 +2039,11 @@ async fn external_agent_config_import_returns_before_background_session_import_f }) .await??; - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; completed_import_ids.push(completed.import_id); } completed_import_ids.sort(); @@ -2284,12 +2067,8 @@ async fn external_agent_config_import_returns_before_background_session_import_f ancestor_thread_id: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadListResponse = to_response(response)?; + let response: ThreadListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.data.len(), 1); Ok(()) @@ -2314,15 +2093,12 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ .await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::default(), - /*auto_compact_limit*/ 200, - /*requires_openai_auth*/ None, - "mock_provider", - "Summarize the conversation.", - )?; + MockResponsesConfig::new(&server.uri()) + .with_root_config( + "compact_prompt = \"Summarize the conversation.\"\nmodel_auto_compact_token_limit = 200", + ) + .with_provider_config("supports_websockets = false") + .write(codex_home.path())?; let project_root = codex_home.path().join("repo"); let recent_timestamp = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true); @@ -2358,9 +2134,8 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("HOME", Some(home_dir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -2370,12 +2145,8 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let detected: ExternalAgentConfigDetectResponse = to_response(response)?; + let detected: ExternalAgentConfigDetectResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(detected.items.len(), 1); let request_id = mcp @@ -2384,21 +2155,14 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ Some(serde_json::json!({ "migrationItems": detected.items })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ExternalAgentConfigImportResponse = to_response(response)?; + let response: ExternalAgentConfigImportResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let import_id = assert_import_response(response); - let notification = timeout( + let completed: ExternalAgentConfigImportCompletedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"), + mcp.read_notification("externalAgentConfig/import/completed"), ) .await??; - assert_eq!(notification.method, "externalAgentConfig/import/completed"); - let completed: ExternalAgentConfigImportCompletedNotification = - serde_json::from_value(notification.params.expect("completed params"))?; assert_eq!(completed.import_id, import_id); let request_id = mcp @@ -2417,12 +2181,8 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ ancestor_thread_id: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadListResponse = to_response(response)?; + let response: ThreadListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let thread = response .data .first() @@ -2435,12 +2195,7 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadResumeResponse = to_response(response)?; + let _: ThreadResumeResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let request_id = mcp .send_turn_start_request(TurnStartParams { @@ -2453,11 +2208,7 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ ..Default::default() }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -2475,28 +2226,6 @@ async fn external_agent_config_import_compacts_huge_session_before_first_follow_ Ok(()) } -fn create_config_toml(codex_home: &std::path::Path, server_uri: &str) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - fn write_analytics_config(codex_home: &std::path::Path, base_url: &str) -> std::io::Result<()> { std::fs::write( codex_home.join("config.toml"), diff --git a/codex-rs/app-server/tests/suite/v2/hooks_list.rs b/codex-rs/app-server/tests/suite/v2/hooks_list.rs index b02180c750b0..dd653e77d9ff 100644 --- a/codex-rs/app-server/tests/suite/v2/hooks_list.rs +++ b/codex-rs/app-server/tests/suite/v2/hooks_list.rs @@ -1,10 +1,10 @@ use std::time::Duration; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::to_response; use codex_app_server_protocol::ConfigBatchWriteParams; use codex_app_server_protocol::ConfigEdit; use codex_app_server_protocol::HookEventName; @@ -15,12 +15,11 @@ use codex_app_server_protocol::HookTrustStatus; use codex_app_server_protocol::HooksListEntry; use codex_app_server_protocol::HooksListParams; use codex_app_server_protocol::HooksListResponse; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::MergeStrategy; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; use codex_core::config::set_project_trust_level; use codex_protocol::config_types::TrustLevel; @@ -141,21 +140,16 @@ async fn hooks_list_shows_discovered_hook() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![cwd.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let config_path = AbsolutePathBuf::from_absolute_path(std::fs::canonicalize( codex_home.path().join("config.toml"), )?)?; @@ -223,21 +217,16 @@ async fn hooks_list_shows_discovered_plugin_hook() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![cwd.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let plugin_hooks_path = AbsolutePathBuf::from_absolute_path(std::fs::canonicalize( codex_home .path() @@ -316,33 +305,23 @@ async fn hooks_list_warms_plugin_capabilities_for_thread_start() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let hooks_list_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![cwd.path().to_path_buf()], }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(hooks_list_id)), - ) - .await??; + let _: HooksListResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(hooks_list_id)).await??; std::fs::remove_file(plugin_mcp_path)?; let thread_start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let _: ThreadStartResponse = to_response( - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??, - )?; + let _: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_matching_notification("plugin MCP server starting", |notification| { @@ -369,21 +348,16 @@ async fn hooks_list_shows_plugin_hook_load_warnings() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![cwd.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].hooks, Vec::new()); @@ -429,9 +403,8 @@ timeout = 5 let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_hooks_list_request(HooksListParams { @@ -441,12 +414,8 @@ timeout = 5 ], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let project_config_path = AbsolutePathBuf::try_from(workspace.path().join(".codex/config.toml"))?; assert_eq!( @@ -517,21 +486,15 @@ async fn hooks_list_uses_root_repo_hooks_for_linked_worktrees() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let list_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![repo_root.clone(), worktree_root.clone()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = timeout(DEFAULT_TIMEOUT, mcp.read_response(list_id)).await??; let repo_hook = data[0].hooks[0].clone(); let worktree_hook = data[1].hooks[0].clone(); let repo_config_path = @@ -559,24 +522,15 @@ async fn hooks_list_uses_root_repo_hooks_for_linked_worktrees() -> Result<()> { reload_user_config: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let _: codex_app_server_protocol::ConfigWriteResponse = to_response(response)?; + let _: codex_app_server_protocol::ConfigWriteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(write_id)).await??; let list_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![worktree_root], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = timeout(DEFAULT_TIMEOUT, mcp.read_response(list_id)).await??; assert_eq!(data[0].hooks[0].trust_status, HookTrustStatus::Trusted); Ok(()) @@ -591,21 +545,16 @@ async fn config_batch_write_toggles_user_hook() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![cwd.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let hook = &data[0].hooks[0]; assert_eq!(hook.enabled, true); @@ -625,24 +574,16 @@ async fn config_batch_write_toggles_user_hook() -> Result<()> { reload_user_config: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let _: codex_app_server_protocol::ConfigWriteResponse = to_response(response)?; + let _: codex_app_server_protocol::ConfigWriteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(write_id)).await??; let request_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![cwd.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data[0].hooks.len(), 1); assert_eq!(data[0].hooks[0].key, hook.key); assert_eq!(data[0].hooks[0].enabled, false); @@ -663,24 +604,16 @@ async fn config_batch_write_toggles_user_hook() -> Result<()> { reload_user_config: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let _: codex_app_server_protocol::ConfigWriteResponse = to_response(response)?; + let _: codex_app_server_protocol::ConfigWriteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(write_id)).await??; let request_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![cwd.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data[0].hooks[0].enabled, true); Ok(()) } @@ -715,53 +648,32 @@ with Path(r"{hook_log_path}").open("a", encoding="utf-8") as handle: hook_log_path = hook_log_path.display(), ), )?; - std::fs::write( - codex_home.path().join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - -[hooks] + MockResponsesConfig::new(&server.uri()) + .with_extra_config(&format!( + r#"[hooks] [[hooks.UserPromptSubmit]] [[hooks.UserPromptSubmit.hooks]] type = "command" -command = "python3 {hook_script_path}" +command = "python3 {}" "#, - server_uri = server.uri(), - hook_script_path = hook_script_path.display(), - ), - )?; + hook_script_path.display() + )) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let hook_list_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![codex_home.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(hook_list_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(hook_list_id)).await??; let hook = data[0].hooks[0].clone(); assert_eq!(hook.trust_status, HookTrustStatus::Untrusted); @@ -771,12 +683,8 @@ command = "python3 {hook_script_path}" ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(response)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_id)).await??; let first_turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -789,11 +697,7 @@ command = "python3 {hook_script_path}" ..Default::default() }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_turn_id)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(first_turn_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -817,24 +721,16 @@ command = "python3 {hook_script_path}" reload_user_config: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let _: codex_app_server_protocol::ConfigWriteResponse = to_response(response)?; + let _: codex_app_server_protocol::ConfigWriteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(write_id)).await??; let hook_list_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![codex_home.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(hook_list_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(hook_list_id)).await??; let trusted_hook = &data[0].hooks[0]; assert_eq!(trusted_hook.key, hook.key); assert_eq!(trusted_hook.current_hash, hook.current_hash); @@ -851,11 +747,8 @@ command = "python3 {hook_script_path}" ..Default::default() }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_turn_id)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(second_turn_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -887,24 +780,16 @@ command = "python3 {hook_script_path}" reload_user_config: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let _: codex_app_server_protocol::ConfigWriteResponse = to_response(response)?; + let _: codex_app_server_protocol::ConfigWriteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(write_id)).await??; let hook_list_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![codex_home.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(hook_list_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(hook_list_id)).await??; let modified_hook = &data[0].hooks[0]; assert_eq!(modified_hook.key, hook.key); assert_ne!(modified_hook.current_hash, hook.current_hash); @@ -921,11 +806,7 @@ command = "python3 {hook_script_path}" ..Default::default() }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(third_turn_id)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(third_turn_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -970,53 +851,32 @@ with Path(r"{hook_log_path}").open("a", encoding="utf-8") as handle: hook_log_path = hook_log_path.display(), ), )?; - std::fs::write( - codex_home.path().join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - -[hooks] + MockResponsesConfig::new(&server.uri()) + .with_extra_config(&format!( + r#"[hooks] [[hooks.UserPromptSubmit]] [[hooks.UserPromptSubmit.hooks]] type = "command" -command = "python3 {hook_script_path}" +command = "python3 {}" "#, - server_uri = server.uri(), - hook_script_path = hook_script_path.display(), - ), - )?; + hook_script_path.display() + )) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let hook_list_id = mcp .send_hooks_list_request(HooksListParams { cwds: vec![codex_home.path().to_path_buf()], }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(hook_list_id)), - ) - .await??; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(hook_list_id)).await??; let hook = &data[0].hooks[0]; assert_eq!(hook.enabled, true); @@ -1036,12 +896,8 @@ command = "python3 {hook_script_path}" reload_user_config: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let _: codex_app_server_protocol::ConfigWriteResponse = to_response(response)?; + let _: codex_app_server_protocol::ConfigWriteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(write_id)).await??; let thread_start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1049,12 +905,8 @@ command = "python3 {hook_script_path}" ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(response)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_id)).await??; let first_turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -1067,11 +919,7 @@ command = "python3 {hook_script_path}" ..Default::default() }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_turn_id)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(first_turn_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -1101,12 +949,8 @@ command = "python3 {hook_script_path}" reload_user_config: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(write_id)), - ) - .await??; - let _: codex_app_server_protocol::ConfigWriteResponse = to_response(response)?; + let _: codex_app_server_protocol::ConfigWriteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(write_id)).await??; let second_turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -1119,11 +963,8 @@ command = "python3 {hook_script_path}" ..Default::default() }) .await?; - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_turn_id)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(second_turn_id)).await??; timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), diff --git a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index 10bba06773fb..b56e78298b16 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -4,13 +4,11 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::ImageGenerationItem; use codex_app_server_protocol::ItemCompletedNotification; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -18,6 +16,7 @@ use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; use codex_config::types::AuthCredentialsStoreMode; +use codex_features::Feature; use core_test_support::responses; use core_test_support::skip_if_remote; use pretty_assertions::assert_eq; @@ -92,9 +91,8 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; start_image_generation_turn( &mut mcp, ThreadStartParams { @@ -221,9 +219,8 @@ async fn standalone_image_generation_failure_emits_terminal_item() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; start_image_generation_turn(&mut mcp, ThreadStartParams::default()).await?; let completed = timeout( @@ -350,9 +347,8 @@ async fn standalone_image_generation_is_exposed_in_code_mode_only() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; start_image_generation_turn(&mut mcp, ThreadStartParams::default()).await?; timeout( DEFAULT_READ_TIMEOUT, @@ -416,9 +412,8 @@ generatedImage(result); let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; start_image_generation_turn(&mut mcp, ThreadStartParams::default()).await?; timeout( DEFAULT_READ_TIMEOUT, @@ -503,9 +498,8 @@ async fn run_image_edit_test( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; start_turn( &mut mcp, ThreadStartParams { @@ -555,12 +549,8 @@ async fn start_turn( let thread_req = mcp .send_thread_start_request_with_auto_env(thread_start_params) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let turn_req = mcp .send_turn_start_request(TurnStartParams { @@ -570,12 +560,7 @@ async fn start_turn( ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _turn: TurnStartResponse = to_response::(turn_resp)?; + let _: TurnStartResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; Ok(()) } @@ -584,14 +569,7 @@ async fn wait_for_image_generation_completed( mcp: &mut TestAppServer, ) -> Result { loop { - let notification = mcp - .read_stream_until_notification_message("item/completed") - .await?; - let completed: ItemCompletedNotification = serde_json::from_value( - notification - .params - .context("item/completed notification should include params")?, - )?; + let completed: ItemCompletedNotification = mcp.read_notification("item/completed").await?; if matches!(&completed.item, ThreadItem::ImageGeneration(_)) { return Ok(completed); } @@ -627,32 +605,14 @@ fn create_config_toml( server_uri: &str, mode: ImagegenTestMode, ) -> std::io::Result<()> { - let code_mode_only = match mode { - ImagegenTestMode::Direct => "", - ImagegenTestMode::CodeModeOnly => "code_mode_only = true", - }; - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -model_provider = "openai-custom" -chatgpt_base_url = "{server_uri}" - -[features] -{code_mode_only} - -[model_providers.openai-custom] -name = "OpenAI" -base_url = "{server_uri}/api/codex" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -supports_websockets = false -requires_openai_auth = true -"# - ), - ) + let mut config = MockResponsesConfig::new(server_uri) + .with_model_provider("openai-custom") + .with_provider_name("OpenAI") + .with_provider_base_url(&format!("{server_uri}/api/codex")) + .with_root_config(&format!("chatgpt_base_url = \"{server_uri}\"")) + .with_provider_config("supports_websockets = false\nrequires_openai_auth = true"); + if matches!(mode, ImagegenTestMode::CodeModeOnly) { + config = config.enable_feature(Feature::CodeModeOnly); + } + config.write(codex_home) } diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index aba3f56078c6..047cd7ada204 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; @@ -7,19 +8,18 @@ use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeResponse; use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_features::Feature; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_cargo_bin::cargo_bin; use core_test_support::fs_wait; use pretty_assertions::assert_eq; use serde_json::Value; -use std::path::Path; use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; @@ -32,7 +32,9 @@ async fn initialize_uses_client_info_name_as_originator() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; let expected_codex_home = AbsolutePathBuf::try_from(codex_home.path().canonicalize()?)?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; + MockResponsesConfig::new(&server.uri()) + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() @@ -71,7 +73,9 @@ async fn initialize_probe_does_not_override_originator() -> Result<()> { let responses = Vec::new(); let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; + MockResponsesConfig::new(&server.uri()) + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() @@ -102,7 +106,9 @@ async fn initialize_codex_backend_does_not_override_originator() -> Result<()> { let responses = Vec::new(); let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; + MockResponsesConfig::new(&server.uri()) + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() @@ -134,7 +140,9 @@ async fn initialize_respects_originator_override_env_var() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; let expected_codex_home = AbsolutePathBuf::try_from(codex_home.path().canonicalize()?)?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; + MockResponsesConfig::new(&server.uri()) + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() @@ -177,7 +185,9 @@ async fn initialize_rejects_invalid_client_name() -> Result<()> { let responses = Vec::new(); let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; + MockResponsesConfig::new(&server.uri()) + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() @@ -213,7 +223,9 @@ async fn initialize_opt_out_notification_methods_filters_notifications() -> Resu let responses = Vec::new(); let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri(), "never")?; + MockResponsesConfig::new(&server.uri()) + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .build() @@ -289,16 +301,14 @@ async fn turn_start_notify_payload_includes_initialize_client_name() -> Result<( let notify_file_str = notify_file .to_str() .expect("notify file path should be valid UTF-8"); - create_config_toml_with_extra( - codex_home.path(), - &server.uri(), - "never", - &format!( + MockResponsesConfig::new(&server.uri()) + .with_root_config(&format!( "notify = [{}, {}]", toml_basic_string(notify_capture), toml_basic_string(notify_file_str) - ), - )?; + )) + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -317,12 +327,8 @@ async fn turn_start_notify_payload_includes_initialize_client_name() -> Result<( let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let turn_req = mcp .send_turn_start_request(TurnStartParams { @@ -335,12 +341,7 @@ async fn turn_start_notify_payload_includes_initialize_client_name() -> Result<( ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_resp)?; + let _: TurnStartResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, @@ -356,48 +357,6 @@ async fn turn_start_notify_payload_includes_initialize_client_name() -> Result<( Ok(()) } -// Helper to create a config.toml pointing at the mock model server. -fn create_config_toml( - codex_home: &Path, - server_uri: &str, - approval_policy: &str, -) -> std::io::Result<()> { - create_config_toml_with_extra(codex_home, server_uri, approval_policy, "") -} - -fn create_config_toml_with_extra( - codex_home: &Path, - server_uri: &str, - approval_policy: &str, - extra: &str, -) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "mock-model" -approval_policy = "{approval_policy}" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -{extra} - -[features] -shell_snapshot = false - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - fn toml_basic_string(value: &str) -> String { format!("\"{}\"", value.replace('\\', "\\\\").replace('"', "\\\"")) } diff --git a/codex-rs/app-server/tests/suite/v2/marketplace_add.rs b/codex-rs/app-server/tests/suite/v2/marketplace_add.rs index 89a930125585..0e82a802b2cb 100644 --- a/codex-rs/app-server/tests/suite/v2/marketplace_add.rs +++ b/codex-rs/app-server/tests/suite/v2/marketplace_add.rs @@ -1,10 +1,7 @@ use anyhow::Result; use app_test_support::TestAppServer; -use app_test_support::to_response; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::MarketplaceAddParams; use codex_app_server_protocol::MarketplaceAddResponse; -use codex_app_server_protocol::RequestId; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -31,9 +28,8 @@ async fn marketplace_add_local_directory_source() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_marketplace_add_request(MarketplaceAddParams { @@ -43,16 +39,11 @@ async fn marketplace_add_local_directory_source() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let MarketplaceAddResponse { marketplace_name, installed_root, already_added, - } = to_response(response)?; + } = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let expected_root = AbsolutePathBuf::from_absolute_path(source.canonicalize()?)?; assert_eq!(marketplace_name, "debug"); diff --git a/codex-rs/app-server/tests/suite/v2/marketplace_remove.rs b/codex-rs/app-server/tests/suite/v2/marketplace_remove.rs index 3f84180ae1a0..fdefc7570950 100644 --- a/codex-rs/app-server/tests/suite/v2/marketplace_remove.rs +++ b/codex-rs/app-server/tests/suite/v2/marketplace_remove.rs @@ -3,8 +3,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use app_test_support::TestAppServer; -use app_test_support::to_response; -use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::MarketplaceRemoveParams; use codex_app_server_protocol::MarketplaceRemoveResponse; use codex_app_server_protocol::RequestId; @@ -56,22 +55,16 @@ async fn marketplace_remove_deletes_config_and_installed_root() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - - let request_id = mcp - .send_marketplace_remove_request(MarketplaceRemoveParams { - marketplace_name: "debug".to_string(), + let response: MarketplaceRemoveResponse = mcp + .request(|request_id| ClientRequest::MarketplaceRemove { + request_id, + params: MarketplaceRemoveParams { + marketplace_name: "debug".to_string(), + }, }) .await?; - - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: MarketplaceRemoveResponse = to_response(response)?; assert_eq!(response.marketplace_name, "debug"); let removed_installed_root = response .installed_root @@ -98,9 +91,8 @@ async fn marketplace_remove_rejects_unknown_marketplace() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_marketplace_remove_request(MarketplaceRemoveParams { diff --git a/codex-rs/app-server/tests/suite/v2/marketplace_upgrade.rs b/codex-rs/app-server/tests/suite/v2/marketplace_upgrade.rs index 6e8f300d0d07..4ef4f9174e32 100644 --- a/codex-rs/app-server/tests/suite/v2/marketplace_upgrade.rs +++ b/codex-rs/app-server/tests/suite/v2/marketplace_upgrade.rs @@ -5,8 +5,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use app_test_support::TestAppServer; -use app_test_support::to_response; -use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::MarketplaceUpgradeParams; use codex_app_server_protocol::MarketplaceUpgradeResponse; use codex_app_server_protocol::RequestId; @@ -130,18 +129,13 @@ async fn send_marketplace_upgrade( mcp: &mut TestAppServer, marketplace_name: Option<&str>, ) -> Result { - let request_id = mcp - .send_marketplace_upgrade_request(MarketplaceUpgradeParams { + mcp.request(|request_id| ClientRequest::MarketplaceUpgrade { + request_id, + params: MarketplaceUpgradeParams { marketplace_name: marketplace_name.map(str::to_string), - }) - .await?; - - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + }, + }) + .await } #[tokio::test] @@ -172,9 +166,8 @@ async fn marketplace_upgrade_all_configured_git_marketplaces() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let debug_root = expected_installed_root(codex_home.path(), "debug")?; let tools_root = expected_installed_root(codex_home.path(), "tools")?; @@ -230,9 +223,8 @@ async fn marketplace_upgrade_named_marketplace_only() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let tools_root = expected_installed_root(codex_home.path(), "tools")?; let response = send_marketplace_upgrade(&mut mcp, Some("tools")).await?; @@ -275,9 +267,8 @@ async fn marketplace_upgrade_returns_empty_roots_when_already_up_to_date() -> Re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let first_response = send_marketplace_upgrade(&mut mcp, Some("debug")).await?; assert!(first_response.errors.is_empty()); @@ -307,9 +298,8 @@ async fn marketplace_upgrade_rejects_unknown_or_non_git_marketplace() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; for marketplace_name in ["missing", "local-only"] { let request_id = mcp diff --git a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs index f96808498276..e919a218d693 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs @@ -5,8 +5,8 @@ use std::time::Duration; use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Router; use codex_app_server::in_process; @@ -14,7 +14,6 @@ use codex_app_server::in_process::InProcessStartArgs; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::InitializeParams; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::McpResourceContent; use codex_app_server_protocol::McpResourceReadParams; use codex_app_server_protocol::McpResourceReadResponse; @@ -22,6 +21,7 @@ use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput; use codex_arg0::Arg0DispatchPaths; use codex_config::CloudConfigBundleLoader; @@ -29,6 +29,7 @@ use codex_config::LoaderOverrides; use codex_config::types::AuthCredentialsStoreMode; use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; +use codex_features::Feature; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; use core_test_support::responses; @@ -104,35 +105,23 @@ async fn mcp_resource_read_returns_resource_contents() -> Result<()> { ) .await?; - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; - - let read_request_id = mcp - .send_mcp_resource_read_request(McpResourceReadParams { - thread_id: Some(thread.id), - server: "codex_apps".to_string(), - uri: TEST_RESOURCE_URI.to_string(), + let read_response: McpResourceReadResponse = mcp + .request(|request_id| ClientRequest::McpResourceRead { + request_id, + params: McpResourceReadParams { + thread_id: Some(thread.id), + server: "codex_apps".to_string(), + uri: TEST_RESOURCE_URI.to_string(), + }, }) .await?; - let read_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_request_id)), - ) - .await??; - assert_eq!( - to_response::(read_response)?, - expected_resource_read_response() - ); + assert_eq!(read_response, expected_resource_read_response()); apps_server_handle.abort(); let _ = apps_server_handle.await; @@ -159,12 +148,8 @@ async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() - ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_start_id)).await??; let response_mock = responses::mount_sse_sequence( &responses_server, @@ -241,11 +226,8 @@ async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() - ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_start_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -353,9 +335,9 @@ async fn orchestrator_skill_can_read_referenced_resource_without_an_executor() - ..Default::default() }) .await?; - timeout( + let _: TurnStartResponse = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(refreshed_turn_start_id)), + mcp.read_response(refreshed_turn_start_id), ) .await??; timeout( @@ -399,12 +381,8 @@ async fn local_executor_does_not_expose_orchestrator_skills() -> Result<()> { ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_start_id)).await??; let response_mock = responses::mount_sse_once( &responses_server, @@ -425,11 +403,8 @@ async fn local_executor_does_not_expose_orchestrator_skills() -> Result<()> { ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_start_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -474,18 +449,12 @@ enabled = false ) .await?; - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; let response_mock = responses::mount_sse_once( &responses_server, @@ -506,11 +475,8 @@ enabled = false ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_start_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -578,42 +544,31 @@ apps = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let read_request_id = mcp - .send_mcp_resource_read_request(McpResourceReadParams { - thread_id: None, - server: "codex_apps".to_string(), - uri: TEST_RESOURCE_URI.to_string(), + let read_response: McpResourceReadResponse = mcp + .request(|request_id| ClientRequest::McpResourceRead { + request_id, + params: McpResourceReadParams { + thread_id: None, + server: "codex_apps".to_string(), + uri: TEST_RESOURCE_URI.to_string(), + }, }) .await?; - let read_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_request_id)), - ) - .await??; - - assert_eq!( - to_response::(read_response)?, - expected_resource_read_response() - ); - - let read_request_id = mcp - .send_mcp_resource_read_request(McpResourceReadParams { - thread_id: None, - server: "codex_apps".to_string(), - uri: TEST_ELICITATION_RESOURCE_URI.to_string(), + assert_eq!(read_response, expected_resource_read_response()); + let read_response: McpResourceReadResponse = mcp + .request(|request_id| ClientRequest::McpResourceRead { + request_id, + params: McpResourceReadParams { + thread_id: None, + server: "codex_apps".to_string(), + uri: TEST_ELICITATION_RESOURCE_URI.to_string(), + }, }) .await?; - let read_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_request_id)), - ) - .await??; assert_eq!( - to_response::(read_response)?, + read_response, McpResourceReadResponse { contents: vec![McpResourceContent::Text { uri: TEST_ELICITATION_RESOURCE_URI.to_string(), @@ -713,34 +668,16 @@ async fn start_resource_test_app_server_with_extra_config( environment: ResourceTestEnvironment, ) -> Result<(TempDir, TestAppServer)> { let codex_home = TempDir::new()?; - std::fs::write( - codex_home.path().join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "untrusted" -sandbox_mode = "read-only" - -model_provider = "mock_provider" -chatgpt_base_url = "{apps_server_url}" -mcp_oauth_credentials_store = "file" - -[features] -apps = true - -[skills] -include_instructions = true -{extra_config} - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{responses_server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - )?; + MockResponsesConfig::new(responses_server_uri) + .with_approval_policy("untrusted") + .with_root_config(&format!( + "chatgpt_base_url = \"{apps_server_url}\"\nmcp_oauth_credentials_store = \"file\"" + )) + .enable_feature(Feature::Apps) + .with_extra_config(&format!( + "[skills]\ninclude_instructions = true\n{extra_config}" + )) + .write(codex_home.path())?; write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("chatgpt-token") @@ -756,8 +693,7 @@ stream_max_retries = 0 // The Local caller explicitly exercises the implicit local executor. ResourceTestEnvironment::Local => builder.without_auto_env(), }; - let mut mcp = builder.build().await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let mcp = builder.build_initialized().await?; Ok((codex_home, mcp)) } diff --git a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs index 5d42cab0b674..e892b0491799 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs @@ -6,11 +6,11 @@ use std::sync::Arc; use std::time::Duration; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::to_response; -use app_test_support::write_mock_responses_config_toml; use axum::Router; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ListMcpServerStatusParams; use codex_app_server_protocol::ListMcpServerStatusResponse; use codex_app_server_protocol::McpServerStatusDetail; @@ -85,47 +85,28 @@ async fn mcp_server_status_list_returns_raw_server_and_tool_names() -> Result<() let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let (mcp_server_url, mcp_server_handle) = start_mcp_server("look-up.raw").await?; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.some-server] -url = "{mcp_server_url}/mcp" -"# - )); - std::fs::write(config_path, config_toml)?; + mock_responses_config(&server.uri()) + .with_extra_config(&format!( + "[mcp_servers.some-server]\nurl = \"{mcp_server_url}/mcp\"" + )) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let request_id = mcp - .send_list_mcp_server_status_request(ListMcpServerStatusParams { - cursor: None, - limit: None, - detail: None, - thread_id: None, + let response: ListMcpServerStatusResponse = mcp + .request(|request_id| ClientRequest::McpServerStatusList { + request_id, + params: ListMcpServerStatusParams { + cursor: None, + limit: None, + detail: None, + thread_id: None, + }, }) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ListMcpServerStatusResponse = to_response(response)?; assert_eq!(response.next_cursor, None); assert_eq!(response.data.len(), 1); @@ -161,24 +142,12 @@ async fn mcp_server_status_list_waits_for_live_stdio_metadata_before_using_cache -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - let barrier_file = codex_home.path().join("allow-initialize"); let pid_file = codex_home.path().join("mcp.pid"); std::fs::write(&barrier_file, "ready")?; - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.cached-stdio] + mock_responses_config(&server.uri()) + .with_extra_config(&format!( + r#"[mcp_servers.cached-stdio] command = {} enabled_tools = ["echo"] startup_timeout_sec = 10 @@ -188,33 +157,28 @@ MCP_TEST_DYNAMIC_SERVER_METADATA = "1" MCP_TEST_INITIALIZE_BARRIER_FILE = {} MCP_TEST_PID_FILE = {} "#, - toml::Value::String(stdio_server_bin()?), - toml::Value::String(barrier_file.to_string_lossy().into_owned()), - toml::Value::String(pid_file.to_string_lossy().into_owned()), - )); - std::fs::write(config_path, config_toml)?; + toml::Value::String(stdio_server_bin()?), + toml::Value::String(barrier_file.to_string_lossy().into_owned()), + toml::Value::String(pid_file.to_string_lossy().into_owned()), + )) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let first_request_id = mcp - .send_list_mcp_server_status_request(ListMcpServerStatusParams { - cursor: None, - limit: None, - detail: Some(McpServerStatusDetail::ToolsAndAuthOnly), - thread_id: None, + let first_response: ListMcpServerStatusResponse = mcp + .request(|request_id| ClientRequest::McpServerStatusList { + request_id, + params: ListMcpServerStatusParams { + cursor: None, + limit: None, + detail: Some(McpServerStatusDetail::ToolsAndAuthOnly), + thread_id: None, + }, }) .await?; - let first_response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_request_id)), - ) - .await??; - let first_response: ListMcpServerStatusResponse = to_response(first_response)?; let first_pid = wait_for_new_pid(&pid_file, /*previous_pid*/ None).await?; assert_dynamic_status(&first_response, &format!("rmcp-test-process-{first_pid}")); @@ -239,12 +203,8 @@ MCP_TEST_PID_FILE = {} ); std::fs::write(&barrier_file, "ready")?; - let second_response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_request_id)), - ) - .await??; - let second_response: ListMcpServerStatusResponse = to_response(second_response)?; + let second_response: ListMcpServerStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(second_request_id)).await??; assert_dynamic_status(&second_response, &format!("rmcp-test-process-{second_pid}")); Ok(()) @@ -256,36 +216,20 @@ async fn mcp_server_status_list_uses_thread_project_local_config() -> Result<()> let (mcp_server_url, mcp_server_handle) = start_mcp_server("project_lookup").await?; let codex_home = TempDir::new()?; let workspace = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; + mock_responses_config(&server.uri()).write(codex_home.path())?; std::fs::create_dir_all(workspace.path().join(".git"))?; set_project_trust_level(codex_home.path(), workspace.path(), TrustLevel::Trusted)?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { cwd: Some(workspace.path().to_string_lossy().into_owned()), ..Default::default() }) .await?; - let thread_start_response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; let project_config_dir = workspace.path().join(".codex"); std::fs::create_dir_all(&project_config_dir)?; @@ -299,36 +243,30 @@ url = "{mcp_server_url}/mcp" ), )?; - let threadless_request_id = mcp - .send_list_mcp_server_status_request(ListMcpServerStatusParams { - cursor: None, - limit: None, - detail: Some(McpServerStatusDetail::ToolsAndAuthOnly), - thread_id: None, + let threadless_response: ListMcpServerStatusResponse = mcp + .request(|request_id| ClientRequest::McpServerStatusList { + request_id, + params: ListMcpServerStatusParams { + cursor: None, + limit: None, + detail: Some(McpServerStatusDetail::ToolsAndAuthOnly), + thread_id: None, + }, }) .await?; - let threadless_response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(threadless_request_id)), - ) - .await??; - let threadless_response: ListMcpServerStatusResponse = to_response(threadless_response)?; assert_eq!(threadless_response.data, Vec::new()); - let thread_request_id = mcp - .send_list_mcp_server_status_request(ListMcpServerStatusParams { - cursor: None, - limit: None, - detail: Some(McpServerStatusDetail::ToolsAndAuthOnly), - thread_id: Some(thread.id), + let thread_response: ListMcpServerStatusResponse = mcp + .request(|request_id| ClientRequest::McpServerStatusList { + request_id, + params: ListMcpServerStatusParams { + cursor: None, + limit: None, + detail: Some(McpServerStatusDetail::ToolsAndAuthOnly), + thread_id: Some(thread.id), + }, }) .await?; - let thread_response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), - ) - .await??; - let thread_response: ListMcpServerStatusResponse = to_response(thread_response)?; assert_eq!(thread_response.next_cursor, None); assert_eq!(thread_response.data.len(), 1); @@ -455,32 +393,17 @@ async fn mcp_server_status_list_tools_and_auth_only_skips_slow_inventory_calls() let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await; let (mcp_server_url, mcp_server_handle) = start_slow_inventory_mcp_server("lookup").await?; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.some-server] -url = "{mcp_server_url}/mcp" -"# - )); - std::fs::write(config_path, config_toml)?; + mock_responses_config(&server.uri()) + .with_extra_config(&format!( + "[mcp_servers.some-server]\nurl = \"{mcp_server_url}/mcp\"" + )) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_list_mcp_server_status_request(ListMcpServerStatusParams { @@ -490,12 +413,8 @@ url = "{mcp_server_url}/mcp" thread_id: None, }) .await?; - let response = timeout( - Duration::from_millis(500), - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ListMcpServerStatusResponse = to_response(response)?; + let response: ListMcpServerStatusResponse = + timeout(Duration::from_millis(500), mcp.read_response(request_id)).await??; assert_eq!(response.next_cursor, None); assert_eq!(response.data.len(), 1); @@ -521,50 +440,33 @@ async fn mcp_server_status_list_keeps_tools_for_sanitized_name_collisions() -> R let (underscore_server_url, underscore_server_handle) = start_mcp_server("underscore_lookup").await?; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.some-server] + mock_responses_config(&server.uri()) + .with_extra_config(&format!( + r#"[mcp_servers.some-server] url = "{dash_server_url}/mcp" [mcp_servers.some_server] url = "{underscore_server_url}/mcp" "# - )); - std::fs::write(config_path, config_toml)?; + )) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let request_id = mcp - .send_list_mcp_server_status_request(ListMcpServerStatusParams { - cursor: None, - limit: None, - detail: None, - thread_id: None, + let response: ListMcpServerStatusResponse = mcp + .request(|request_id| ClientRequest::McpServerStatusList { + request_id, + params: ListMcpServerStatusParams { + cursor: None, + limit: None, + detail: None, + thread_id: None, + }, }) .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ListMcpServerStatusResponse = to_response(response)?; assert_eq!(response.next_cursor, None); assert_eq!(response.data.len(), 2); @@ -640,3 +542,9 @@ async fn start_slow_inventory_mcp_server(tool_name: &str) -> Result<(String, Joi Ok((format!("http://{addr}"), handle)) } + +fn mock_responses_config(server_uri: &str) -> MockResponsesConfig { + MockResponsesConfig::new(server_uri) + .with_root_config("compact_prompt = \"compact\"\nmodel_auto_compact_token_limit = 1024") + .with_provider_config("supports_websockets = false") +} diff --git a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs index 453b4c0abf48..814fccb827f1 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs @@ -1,20 +1,18 @@ use std::borrow::Cow; -use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; -use app_test_support::to_response; -use app_test_support::write_mock_responses_config_toml; use axum::Router; use codex_app_server_protocol::CapabilityRootLocation; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::EnvironmentAddResponse; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::McpElicitationSchema; use codex_app_server_protocol::McpServerElicitationAction; use codex_app_server_protocol::McpServerElicitationRequest; @@ -72,6 +70,8 @@ use super::exec_server_test_support::accept_exec_server_environment; use super::exec_server_test_support::read_exec_server_json; const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); +const AUTO_COMPACT_LIMIT: i64 = 1024; +const LARGE_OUTPUT_AUTO_COMPACT_LIMIT: i64 = 1_000_000; const TEST_SERVER_NAME: &str = "tool_server"; const TEST_TOOL_NAME: &str = "echo_tool"; const LARGE_RESPONSE_MESSAGE: &str = "large"; @@ -87,65 +87,36 @@ async fn mcp_server_tool_call_returns_tool_result() -> Result<()> { let responses_server = responses::start_mock_server().await; let (mcp_server_url, mcp_server_handle) = start_mcp_server().await?; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &responses_server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.{TEST_SERVER_NAME}] -url = "{mcp_server_url}/mcp" -"# - )); - std::fs::write(config_path, config_toml)?; + mcp_tool_config(&responses_server.uri(), &mcp_server_url, AUTO_COMPACT_LIMIT) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; let thread_id = thread.id.clone(); - - let tool_call_request_id = mcp - .send_mcp_server_tool_call_request(McpServerToolCallParams { - thread_id: thread_id.clone(), - server: TEST_SERVER_NAME.to_string(), - tool: TEST_TOOL_NAME.to_string(), - arguments: Some(json!({ - "message": "hello from app", - })), - meta: Some(json!({ - "source": "mcp-app", - })), + let response: McpServerToolCallResponse = mcp + .request(|request_id| ClientRequest::McpServerToolCall { + request_id, + params: McpServerToolCallParams { + thread_id: thread_id.clone(), + server: TEST_SERVER_NAME.to_string(), + tool: TEST_TOOL_NAME.to_string(), + arguments: Some(json!({ + "message": "hello from app", + })), + meta: Some(json!({ + "source": "mcp-app", + })), + }, }) .await?; - let tool_call_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(tool_call_request_id)), - ) - .await??; - let response: McpServerToolCallResponse = to_response(tool_call_response)?; assert_eq!(response.content.len(), 1); assert_eq!(response.content[0].get("type"), Some(&json!("text"))); @@ -180,9 +151,8 @@ async fn mcp_server_tool_call_returns_error_for_unknown_thread() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_mcp_server_tool_call_request(McpServerToolCallParams { @@ -212,45 +182,20 @@ async fn mcp_server_tool_call_round_trips_elicitation() -> Result<()> { let responses_server = responses::start_mock_server().await; let (mcp_server_url, mcp_server_handle) = start_mcp_server().await?; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &responses_server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.{TEST_SERVER_NAME}] -url = "{mcp_server_url}/mcp" -"# - )); - std::fs::write(config_path, config_toml)?; + mcp_tool_config(&responses_server.uri(), &mcp_server_url, AUTO_COMPACT_LIMIT) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), approval_policy: Some(codex_app_server_protocol::AskForApproval::UnlessTrusted), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; let tool_call_request_id = mcp .send_mcp_server_tool_call_request(McpServerToolCallParams { @@ -304,12 +249,11 @@ url = "{mcp_server_url}/mcp" ) .await?; - let tool_call_response: JSONRPCResponse = timeout( + let response: McpServerToolCallResponse = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(tool_call_request_id)), + mcp.read_response(tool_call_request_id), ) .await??; - let response: McpServerToolCallResponse = to_response(tool_call_response)?; assert_eq!(response.content.len(), 1); assert_eq!(response.content[0].get("type"), Some(&json!("text"))); assert_eq!(response.content[0].get("text"), Some(&json!("accepted"))); @@ -327,32 +271,16 @@ async fn mcp_server_elicitation_survives_environment_runtime_refresh() -> Result let exec_listener = TcpListener::bind("127.0.0.1:0").await?; let exec_server_url = format!("ws://{}", exec_listener.local_addr()?); let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &responses_server.uri(), - &BTreeMap::from([(Feature::DeferredExecutor, true)]), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.{TEST_SERVER_NAME}] -url = "{mcp_server_url}/mcp" -"# - )); - std::fs::write(config_path, config_toml)?; + mcp_tool_config(&responses_server.uri(), &mcp_server_url, AUTO_COMPACT_LIMIT) + .enable_feature(Feature::DeferredExecutor) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) // This test adds and refreshes an explicitly selected runtime environment. .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let add_environment_id = mcp .send_raw_request( "environment/add", @@ -363,12 +291,8 @@ url = "{mcp_server_url}/mcp" })), ) .await?; - let add_environment_response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(add_environment_id)), - ) - .await??; - let _: EnvironmentAddResponse = to_response(add_environment_response)?; + let _: EnvironmentAddResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(add_environment_id)).await??; let capability_root = TempDir::new()?; let thread_start_id = mcp @@ -393,12 +317,8 @@ url = "{mcp_server_url}/mcp" ..Default::default() }) .await?; - let thread_start_response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_start_id)).await??; let tool_call_request_id = mcp .send_mcp_server_tool_call_request(McpServerToolCallParams { @@ -450,12 +370,11 @@ url = "{mcp_server_url}/mcp" })?, ) .await?; - let tool_call_response = timeout( + let response: McpServerToolCallResponse = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(tool_call_request_id)), + mcp.read_response(tool_call_request_id), ) .await??; - let response: McpServerToolCallResponse = to_response(tool_call_response)?; assert_eq!(response.content[0].get("text"), Some(&json!("accepted"))); let _ = shutdown_tx.send(()); @@ -470,45 +389,20 @@ async fn mcp_server_tool_call_forwards_url_elicitation() -> Result<()> { let responses_server = responses::start_mock_server().await; let (mcp_server_url, mcp_server_handle) = start_mcp_server().await?; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &responses_server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1024, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.{TEST_SERVER_NAME}] -url = "{mcp_server_url}/mcp" -"# - )); - std::fs::write(config_path, config_toml)?; + mcp_tool_config(&responses_server.uri(), &mcp_server_url, AUTO_COMPACT_LIMIT) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), approval_policy: Some(codex_app_server_protocol::AskForApproval::UnlessTrusted), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; let tool_call_request_id = mcp .send_mcp_server_tool_call_request(McpServerToolCallParams { @@ -555,12 +449,11 @@ url = "{mcp_server_url}/mcp" ) .await?; - let tool_call_response: JSONRPCResponse = timeout( + let response: McpServerToolCallResponse = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(tool_call_request_id)), + mcp.read_response(tool_call_request_id), ) .await??; - let response: McpServerToolCallResponse = to_response(tool_call_response)?; assert_eq!(response.content.len(), 1); assert_eq!(response.content[0].get("type"), Some(&json!("text"))); assert_eq!(response.content[0].get("text"), Some(&json!("accepted"))); @@ -593,62 +486,37 @@ async fn mcp_tool_call_completion_notification_contains_truncated_large_result() let responses_server = create_mock_responses_server_sequence(responses).await; let (mcp_server_url, mcp_server_handle) = start_mcp_server().await?; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), + mcp_tool_config( &responses_server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 1_000_000, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; - - let config_path = codex_home.path().join("config.toml"); - let mut config_toml = std::fs::read_to_string(&config_path)?; - config_toml.push_str(&format!( - r#" -[mcp_servers.{TEST_SERVER_NAME}] -url = "{mcp_server_url}/mcp" -"# - )); - std::fs::write(config_path, config_toml)?; + &mcp_server_url, + LARGE_OUTPUT_AUTO_COMPACT_LIMIT, + ) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; - - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Call the large MCP tool".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let TurnStartResponse { turn, .. } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Call the large MCP tool".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let TurnStartResponse { turn, .. } = to_response(turn_start_resp)?; let completed = wait_for_mcp_tool_call_completed(&mut mcp, call_id).await?; assert_eq!(completed.turn_id, turn.id); @@ -901,17 +769,28 @@ async fn wait_for_mcp_tool_call_completed( call_id: &str, ) -> Result { loop { - let notification = timeout( + let completed: ItemCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/completed"), + mcp.read_notification("item/completed"), ) .await??; - let Some(params) = notification.params else { - continue; - }; - let completed: ItemCompletedNotification = serde_json::from_value(params)?; if matches!(&completed.item, ThreadItem::McpToolCall { id, .. } if id == call_id) { return Ok(completed); } } } + +fn mcp_tool_config( + server_uri: &str, + mcp_server_url: &str, + auto_compact_limit: i64, +) -> MockResponsesConfig { + MockResponsesConfig::new(server_uri) + .with_root_config(&format!( + "compact_prompt = \"compact\"\nmodel_auto_compact_token_limit = {auto_compact_limit}" + )) + .with_provider_config("supports_websockets = false") + .with_extra_config(&format!( + "[mcp_servers.{TEST_SERVER_NAME}]\nurl = \"{mcp_server_url}/mcp\"" + )) +} diff --git a/codex-rs/app-server/tests/suite/v2/memory_reset.rs b/codex-rs/app-server/tests/suite/v2/memory_reset.rs index ab8445378d7b..6408f318d4a3 100644 --- a/codex-rs/app-server/tests/suite/v2/memory_reset.rs +++ b/codex-rs/app-server/tests/suite/v2/memory_reset.rs @@ -1,10 +1,9 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; use chrono::Utc; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::MemoryResetResponse; -use codex_app_server_protocol::RequestId; +use codex_features::Feature; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use codex_state::Stage1JobClaimOutcome; @@ -22,7 +21,10 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs #[tokio::test] async fn memory_reset_clears_memory_files_and_rows_preserves_threads() -> Result<()> { let codex_home = TempDir::new()?; - create_config_toml(codex_home.path())?; + MockResponsesConfig::new("http://127.0.0.1:9") + .with_root_config("suppress_unstable_features_warning = true") + .enable_feature(Feature::Sqlite) + .write(codex_home.path())?; let state_db = init_state_db(codex_home.path()).await?; let memory_root = codex_home.path().join("memories"); @@ -39,19 +41,14 @@ async fn memory_reset_clears_memory_files_and_rows_preserves_threads() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request("memory/reset", /*params*/ None) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: MemoryResetResponse = to_response::(response)?; + let _: MemoryResetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let stage1_outputs = state_db .memories() @@ -129,27 +126,3 @@ async fn init_state_db(codex_home: &Path) -> Result> { .await?; Ok(state_db) } - -fn create_config_toml(codex_home: &Path) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - 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 = "http://127.0.0.1:9/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"#, - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/model_list.rs b/codex-rs/app-server/tests/suite/v2/model_list.rs index 65f563e15337..4c029b55c9e9 100644 --- a/codex-rs/app-server/tests/suite/v2/model_list.rs +++ b/codex-rs/app-server/tests/suite/v2/model_list.rs @@ -4,11 +4,10 @@ use anyhow::Error; use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use app_test_support::write_models_cache; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::Model; use codex_app_server_protocol::ModelListParams; use codex_app_server_protocol::ModelListResponse; @@ -99,29 +98,21 @@ async fn list_models_returns_all_models_with_large_limit() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - - let request_id = mcp - .send_list_models_request(ModelListParams { - limit: Some(100), - cursor: None, - include_hidden: None, - }) - .await?; - - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ModelListResponse { data: items, next_cursor, - } = to_response::(response)?; + } = mcp + .request(|request_id| ClientRequest::ModelList { + request_id, + params: ModelListParams { + limit: Some(100), + cursor: None, + include_hidden: None, + }, + }) + .await?; let expected_models = expected_visible_models(); @@ -137,29 +128,21 @@ async fn list_models_includes_hidden_models() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - - let request_id = mcp - .send_list_models_request(ModelListParams { - limit: Some(100), - cursor: None, - include_hidden: Some(true), - }) - .await?; - - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ModelListResponse { data: items, next_cursor, - } = to_response::(response)?; + } = mcp + .request(|request_id| ClientRequest::ModelList { + request_id, + params: ModelListParams { + limit: Some(100), + cursor: None, + include_hidden: Some(true), + }, + }) + .await?; assert!(items.iter().any(|item| item.hidden)); assert!(next_cursor.is_none()); @@ -227,28 +210,21 @@ openai_base_url = "{server_uri}/v1" .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - - let request_id = mcp - .send_list_models_request(ModelListParams { - limit: Some(100), - cursor: None, - include_hidden: None, - }) - .await?; - - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ModelListResponse { data: items, next_cursor, - } = to_response::(response)?; + } = mcp + .request(|request_id| ClientRequest::ModelList { + request_id, + params: ModelListParams { + limit: Some(100), + cursor: None, + include_hidden: None, + }, + }) + .await?; let mut expected_presets: Vec = vec![remote_model.into()]; ModelPreset::mark_default_by_picker_visibility(&mut expected_presets); let mut expected_items = expected_presets @@ -287,34 +263,27 @@ async fn list_models_pagination_works() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let expected_models = expected_visible_models(); let mut cursor = None; let mut items = Vec::new(); for _ in 0..expected_models.len() { - let request_id = mcp - .send_list_models_request(ModelListParams { - limit: Some(1), - cursor: cursor.clone(), - include_hidden: None, - }) - .await?; - - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ModelListResponse { data: page_items, next_cursor, - } = to_response::(response)?; + } = mcp + .request(|request_id| ClientRequest::ModelList { + request_id, + params: ModelListParams { + limit: Some(1), + cursor: cursor.clone(), + include_hidden: None, + }, + }) + .await?; assert_eq!(page_items.len(), 1); items.extend(page_items); @@ -340,11 +309,9 @@ async fn list_models_rejects_invalid_cursor() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp .send_list_models_request(ModelListParams { limit: None, diff --git a/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs b/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs index 58eb5ed6e79f..df842cf8566d 100644 --- a/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs +++ b/codex-rs/app-server/tests/suite/v2/model_provider_capabilities_read.rs @@ -2,11 +2,8 @@ use std::time::Duration; use anyhow::Result; use app_test_support::TestAppServer; -use app_test_support::to_response; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::ModelProviderCapabilitiesReadParams; use codex_app_server_protocol::ModelProviderCapabilitiesReadResponse; -use codex_app_server_protocol::RequestId; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::time::timeout; @@ -19,19 +16,14 @@ async fn read_default_provider_capabilities() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_model_provider_capabilities_read_request(ModelProviderCapabilitiesReadParams {}) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: ModelProviderCapabilitiesReadResponse = to_response(response)?; + let received: ModelProviderCapabilitiesReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let expected = ModelProviderCapabilitiesReadResponse { namespace_tools: true, @@ -53,19 +45,14 @@ async fn read_amazon_bedrock_provider_capabilities() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_model_provider_capabilities_read_request(ModelProviderCapabilitiesReadParams {}) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: ModelProviderCapabilitiesReadResponse = to_response(response)?; + let received: ModelProviderCapabilitiesReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let expected = ModelProviderCapabilitiesReadResponse { namespace_tools: true, diff --git a/codex-rs/app-server/tests/suite/v2/output_schema.rs b/codex-rs/app-server/tests/suite/v2/output_schema.rs index 668074bd1f08..00176f4d3a6c 100644 --- a/codex-rs/app-server/tests/suite/v2/output_schema.rs +++ b/codex-rs/app-server/tests/suite/v2/output_schema.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::to_response; use codex_app_server_protocol::JSONRPCResponse; @@ -11,7 +12,6 @@ use codex_app_server_protocol::UserInput as V2UserInput; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; -use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; @@ -30,7 +30,7 @@ async fn turn_start_accepts_output_schema_v2() -> Result<()> { let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -116,7 +116,7 @@ async fn turn_start_output_schema_is_per_turn_v2() -> Result<()> { let response_mock1 = responses::mount_sse_once(&server, body1).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -218,26 +218,3 @@ async fn turn_start_output_schema_is_per_turn_v2() -> Result<()> { Ok(()) } - -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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/permission_profile_list.rs b/codex-rs/app-server/tests/suite/v2/permission_profile_list.rs index d291dda172aa..791fc2a73f2b 100644 --- a/codex-rs/app-server/tests/suite/v2/permission_profile_list.rs +++ b/codex-rs/app-server/tests/suite/v2/permission_profile_list.rs @@ -2,12 +2,9 @@ use std::time::Duration; use anyhow::Result; use app_test_support::TestAppServer; -use app_test_support::to_response; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PermissionProfileListParams; use codex_app_server_protocol::PermissionProfileListResponse; use codex_app_server_protocol::PermissionProfileSummary; -use codex_app_server_protocol::RequestId; use codex_core::config::set_project_trust_level; use codex_protocol::config_types::TrustLevel; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS; @@ -44,9 +41,8 @@ description = "Inspect without writes." let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_permission_profile_list_request(PermissionProfileListParams { @@ -120,9 +116,8 @@ description = "Project-scoped profile." let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let first_request_id = mcp .send_permission_profile_list_request(PermissionProfileListParams { @@ -201,9 +196,8 @@ description = "Project-scoped profile." let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_permission_profile_list_request(PermissionProfileListParams { @@ -249,10 +243,5 @@ async fn read_response( mcp: &mut TestAppServer, request_id: i64, ) -> Result { - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await? } diff --git a/codex-rs/app-server/tests/suite/v2/plan_item.rs b/codex-rs/app-server/tests/suite/v2/plan_item.rs index 9717d20beee1..9709f2cb00b5 100644 --- a/codex-rs/app-server/tests/suite/v2/plan_item.rs +++ b/codex-rs/app-server/tests/suite/v2/plan_item.rs @@ -1,24 +1,21 @@ use anyhow::Result; use anyhow::anyhow; use anyhow::bail; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PlanDeltaNotification; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; -use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; -use codex_features::FEATURES; use codex_features::Feature; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; @@ -26,8 +23,6 @@ use codex_protocol::config_types::Settings; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; -use std::collections::BTreeMap; -use std::path::Path; use tempfile::TempDir; use tokio::time::sleep; use tokio::time::timeout; @@ -51,13 +46,14 @@ async fn plan_mode_uses_proposed_plan_block_for_plan_item() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::CollaborationModes) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let turn = start_plan_mode_turn(&mut mcp).await?; let (_, completed_items, plan_deltas, turn_completed) = @@ -112,13 +108,14 @@ async fn plan_mode_without_proposed_plan_does_not_emit_plan_item() -> Result<()> let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::CollaborationModes) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let _turn = start_plan_mode_turn(&mut mcp).await?; let (_, completed_items, plan_deltas, _) = collect_turn_notifications(&mut mcp).await?; @@ -134,18 +131,13 @@ async fn plan_mode_without_proposed_plan_does_not_emit_plan_item() -> Result<()> } async fn start_plan_mode_turn(mcp: &mut TestAppServer) -> Result { - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let thread = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) - .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let thread = to_response::(thread_resp)?.thread; + .await? + .thread; let collaboration_mode = CollaborationMode { mode: ModeKind::Plan, @@ -155,24 +147,22 @@ async fn start_plan_mode_turn(mcp: &mut TestAppServer) -> Result(turn_resp)?.turn) + Ok(response.turn) } async fn collect_turn_notifications( @@ -255,42 +245,3 @@ async fn wait_for_responses_request_count( .await??; Ok(()) } - -fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { - let features = BTreeMap::from([(Feature::CollaborationModes, true)]); - let feature_entries = features - .into_iter() - .map(|(feature, enabled)| { - let key = FEATURES - .iter() - .find(|spec| spec.id == feature) - .map(|spec| spec.key) - .expect("feature should have a config key"); - format!("{key} = {enabled}") - }) - .collect::>() - .join("\n"); - 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" - -[features] -{feature_entries} - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/plugin_install.rs b/codex-rs/app-server/tests/suite/v2/plugin_install.rs index 1d781cdc969d..3050ad9df766 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_install.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_install.rs @@ -11,7 +11,6 @@ use app_test_support::ChatGptAuthFixture; use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::TestAppServer; use app_test_support::start_analytics_events_server; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use axum::Json; use axum::Router; @@ -25,7 +24,6 @@ use codex_app_server_protocol::AppInfo; use codex_app_server_protocol::AppSummary; use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PluginAuthPolicy; use codex_app_server_protocol::PluginAvailability; use codex_app_server_protocol::PluginInstallParams; @@ -75,9 +73,8 @@ async fn plugin_install_rejects_relative_marketplace_paths() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -106,9 +103,8 @@ async fn plugin_install_rejects_missing_install_source() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -139,9 +135,8 @@ async fn plugin_install_rejects_multiple_install_sources() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -180,9 +175,8 @@ plugins = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -251,17 +245,12 @@ async fn plugin_install_writes_remote_plugin_to_cloud_and_cache() -> Result<()> .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -339,17 +328,12 @@ async fn plugin_install_uses_remote_apps_needing_auth_response() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -391,9 +375,8 @@ async fn plugin_install_rejects_missing_remote_bundle_url() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; let err = timeout( @@ -436,9 +419,8 @@ async fn plugin_install_rejects_plain_http_remote_bundle_url() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; let err = timeout( @@ -486,9 +468,8 @@ async fn plugin_install_rejects_invalid_remote_release_version() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; let err = timeout( @@ -522,9 +503,8 @@ async fn plugin_install_rejects_invalid_remote_plugin_name() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -556,9 +536,8 @@ async fn plugin_install_tracks_analytics_when_remote_detail_fetch_fails() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; let err = timeout( @@ -613,9 +592,8 @@ async fn plugin_install_rejects_remote_plugin_disabled_by_admin_before_download( .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; let err = timeout( @@ -666,9 +644,8 @@ async fn plugin_install_rejects_remote_plugin_not_available() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; let err = timeout( @@ -733,9 +710,8 @@ async fn plugin_install_rejects_when_workspace_codex_plugins_disabled() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -766,9 +742,8 @@ async fn plugin_install_returns_invalid_request_for_missing_marketplace_file() - let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -811,9 +786,8 @@ async fn plugin_install_returns_invalid_request_for_not_available_plugin() -> Re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -865,9 +839,8 @@ async fn plugin_install_returns_invalid_request_for_disallowed_product_plugin() .with_codex_home(codex_home.path()) .without_auto_env() .with_args(&["--session-source", "atlas"]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -918,9 +891,8 @@ async fn plugin_install_tracks_analytics_event() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -929,12 +901,8 @@ async fn plugin_install_tracks_analytics_event() -> Result<()> { plugin_name: "sample-plugin".to_string(), }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.apps_needing_auth, Vec::::new()); let payload = wait_for_plugin_analytics_payload(&analytics_server).await?; @@ -988,9 +956,8 @@ async fn plugin_install_failure_tracks_analytics_event() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -1045,17 +1012,12 @@ async fn plugin_install_tracks_remote_plugin_analytics_event() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.apps_needing_auth, Vec::::new()); let payload = wait_for_plugin_analytics_payload(&server).await?; @@ -1097,9 +1059,8 @@ async fn plugin_install_preserves_status_when_remote_bundle_error_body_is_too_la .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; let err = timeout( @@ -1226,9 +1187,8 @@ async fn plugin_install_returns_apps_needing_auth() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let directory_requests_before_install = server_control.directory_request_count(); let request_id = mcp @@ -1239,12 +1199,8 @@ async fn plugin_install_returns_apps_needing_auth() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1317,9 +1273,8 @@ async fn plugin_install_skips_mcp_oauth_for_chatgpt_dual_surface_plugin() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -1328,12 +1283,8 @@ async fn plugin_install_skips_mcp_oauth_for_chatgpt_dual_surface_plugin() -> Res plugin_name: "sample-plugin".to_string(), }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.auth_policy, PluginAuthPolicy::OnInstall); assert_eq!(oauth_discovery_request_count(&oauth_server).await, 0); @@ -1381,9 +1332,8 @@ async fn plugin_install_starts_mcp_oauth_with_formerly_disallowed_plugin_app() - let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -1392,12 +1342,8 @@ async fn plugin_install_starts_mcp_oauth_with_formerly_disallowed_plugin_app() - plugin_name: "sample-plugin".to_string(), }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1482,9 +1428,8 @@ async fn plugin_install_starts_mcp_oauth_through_protected_resource_metadata() - let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -1493,12 +1438,8 @@ async fn plugin_install_starts_mcp_oauth_through_protected_resource_metadata() - plugin_name: "sample-plugin".to_string(), }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: PluginInstallResponse = to_response(response)?; + let _: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; wait_for_remote_plugin_request_count( &authorization_server, "POST", @@ -1550,9 +1491,8 @@ connectors = true .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", Some("test-api-key"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -1561,12 +1501,8 @@ connectors = true plugin_name: "sample-plugin".to_string(), }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.auth_policy, PluginAuthPolicy::OnInstall); assert!(oauth_discovery_request_count(&oauth_server).await > 0); @@ -1593,17 +1529,12 @@ async fn plugin_install_starts_remote_mcp_oauth_for_install_response_only_app() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1646,17 +1577,12 @@ async fn plugin_install_skips_remote_mcp_oauth_for_bundled_same_name_app() -> Re .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = send_remote_plugin_install_request(&mut mcp, REMOTE_PLUGIN_ID).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1728,9 +1654,8 @@ async fn plugin_install_includes_formerly_disallowed_apps_needing_auth() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let directory_requests_before_install = warm_app_directory_cache(&mut mcp, &server_control, "Alpha").await?; @@ -1742,12 +1667,8 @@ async fn plugin_install_includes_formerly_disallowed_apps_needing_auth() -> Resu }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1815,9 +1736,8 @@ async fn plugin_install_makes_bundled_mcp_servers_available_to_followup_requests let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_install_request(PluginInstallParams { @@ -1826,12 +1746,8 @@ async fn plugin_install_makes_bundled_mcp_servers_available_to_followup_requests plugin_name: "sample-plugin".to_string(), }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstallResponse = to_response(response)?; + let response: PluginInstallResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.apps_needing_auth, Vec::::new()); let config = std::fs::read_to_string(codex_home.path().join("config.toml"))?; assert!(!config.contains("[mcp_servers.sample-mcp]")); @@ -1887,12 +1803,8 @@ async fn warm_app_directory_cache( ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(app_list_request_id)), - ) - .await??; - let response: AppsListResponse = to_response(response)?; + let response: AppsListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(app_list_request_id)).await??; assert!( response .data diff --git a/codex-rs/app-server/tests/suite/v2/plugin_list.rs b/codex-rs/app-server/tests/suite/v2/plugin_list.rs index e056fa60db3f..d5cf3d778683 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_list.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_list.rs @@ -4,13 +4,11 @@ use anyhow::Result; use anyhow::bail; use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::HookMetadata; use codex_app_server_protocol::HookTrustStatus; use codex_app_server_protocol::HooksListParams; use codex_app_server_protocol::HooksListResponse; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PluginAuthPolicy; use codex_app_server_protocol::PluginInstallPolicy; use codex_app_server_protocol::PluginInstallPolicySource; @@ -113,9 +111,8 @@ async fn plugin_list_skips_invalid_marketplace_file_and_reports_error() -> Resul ("HOME", Some(home.as_str())), ("USERPROFILE", Some(home.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -125,12 +122,8 @@ async fn plugin_list_skips_invalid_marketplace_file_and_reports_error() -> Resul }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!( response @@ -176,9 +169,8 @@ enabled = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_installed_request(PluginInstalledParams { @@ -187,12 +179,8 @@ enabled = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstalledResponse = to_response(response)?; + let response: PluginInstalledResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); assert_eq!(response.marketplaces[0].name, "openai-curated"); @@ -273,9 +261,8 @@ enabled = true let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, app_server.initialize()).await??; let request_id = app_server .send_plugin_installed_request(PluginInstalledParams { @@ -284,12 +271,8 @@ enabled = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstalledResponse = to_response(response)?; + let response: PluginInstalledResponse = + timeout(DEFAULT_TIMEOUT, app_server.read_response(request_id)).await??; let local_marketplace = response .marketplaces @@ -355,9 +338,8 @@ enabled = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_installed_request(PluginInstalledParams { @@ -366,12 +348,8 @@ enabled = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstalledResponse = to_response(response)?; + let response: PluginInstalledResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces, Vec::new()); assert_eq!(response.marketplace_load_errors, Vec::new()); @@ -384,9 +362,8 @@ async fn plugin_list_rejects_relative_cwds() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -469,9 +446,8 @@ async fn plugin_list_keeps_valid_marketplaces_when_another_marketplace_fails_to_ ("HOME", Some(home.as_str())), ("USERPROFILE", Some(home.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -484,12 +460,8 @@ async fn plugin_list_keeps_valid_marketplaces_when_another_marketplace_fails_to_ }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response.marketplaces, @@ -592,9 +564,8 @@ async fn plugin_list_returns_empty_when_workspace_codex_plugins_disabled() -> Re ("HOME", Some(home.as_str())), ("USERPROFILE", Some(home.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -604,12 +575,8 @@ async fn plugin_list_returns_empty_when_workspace_codex_plugins_disabled() -> Re }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -686,9 +653,8 @@ async fn plugin_list_reuses_cached_workspace_codex_plugins_setting() -> Result<( ("HOME", Some(home.as_str())), ("USERPROFILE", Some(home.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; for _ in 0..2 { let request_id = mcp @@ -699,12 +665,8 @@ async fn plugin_list_reuses_cached_workspace_codex_plugins_setting() -> Result<( }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); assert_eq!(response.marketplaces[0].name, "local-marketplace"); } @@ -773,9 +735,8 @@ async fn plugin_list_uses_alternate_discoverable_manifest_and_keeps_undiscoverab ("HOME", Some(home.as_str())), ("USERPROFILE", Some(home.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -785,12 +746,8 @@ async fn plugin_list_uses_alternate_discoverable_manifest_and_keeps_undiscoverab }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response.marketplaces, @@ -896,9 +853,8 @@ async fn plugin_list_accepts_omitted_cwds() -> Result<()> { ("HOME", Some(home.as_str())), ("USERPROFILE", Some(home.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -908,12 +864,7 @@ async fn plugin_list_accepts_omitted_cwds() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: PluginListResponse = to_response(response)?; + let _: PluginListResponse = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; Ok(()) } @@ -954,9 +905,8 @@ async fn plugin_list_returns_share_context_for_shared_local_plugin() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -966,12 +916,8 @@ async fn plugin_list_returns_share_context_for_shared_local_plugin() -> Result<( }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let plugin = response .marketplaces @@ -1051,9 +997,8 @@ enabled = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -1063,12 +1008,8 @@ enabled = false }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let marketplace = response .marketplaces @@ -1207,9 +1148,8 @@ enabled = false ("HOME", Some(home.as_str())), ("USERPROFILE", Some(home.as_str())), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -1222,12 +1162,8 @@ enabled = false }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let shared_plugin = response .marketplaces @@ -1299,9 +1235,8 @@ async fn plugin_list_returns_plugin_interface_with_absolute_asset_paths() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -1311,12 +1246,8 @@ async fn plugin_list_returns_plugin_interface_with_absolute_asset_paths() -> Res }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let plugin = response .marketplaces @@ -1417,9 +1348,8 @@ async fn plugin_list_accepts_legacy_string_default_prompt() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -1429,12 +1359,8 @@ async fn plugin_list_accepts_legacy_string_default_prompt() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let plugin = response .marketplaces @@ -1510,9 +1436,8 @@ enabled = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -1522,12 +1447,8 @@ enabled = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let plugin = response .marketplaces @@ -1619,14 +1540,13 @@ async fn app_server_startup_sync_downloads_remote_installed_plugin_bundles() -> let installed_path = codex_home .path() .join("plugins/cache/openai-curated-remote/linear/1.2.3"); - let mut mcp = TestAppServer::builder() + let _mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() .with_plugin_startup_tasks() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; wait_for_path_exists(&installed_path.join(".codex-plugin/plugin.json")).await?; let installed_plugin_manifest: serde_json::Value = serde_json::from_str( @@ -1704,9 +1624,8 @@ async fn plugin_list_sync_upgrades_and_removes_remote_installed_plugin_bundles() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -1715,12 +1634,8 @@ async fn plugin_list_sync_upgrades_and_removes_remote_installed_plugin_bundles() force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let remote_marketplace = response .marketplaces .into_iter() @@ -1896,9 +1811,8 @@ async fn plugin_list_includes_remote_marketplaces_when_remote_plugin_enabled() - let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -1908,12 +1822,8 @@ async fn plugin_list_includes_remote_marketplaces_when_remote_plugin_enabled() - }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let remote_marketplace = response .marketplaces @@ -2068,9 +1978,8 @@ async fn plugin_list_uses_cached_global_remote_catalog_and_refreshes_it() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2079,13 +1988,8 @@ async fn plugin_list_uses_cached_global_remote_catalog_and_refreshes_it() -> Res force_refetch: false, }) .await?; - let response: PluginListResponse = to_response( - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??, - )?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let remote_marketplace = response .marketplaces .iter() @@ -2117,13 +2021,8 @@ async fn plugin_list_uses_cached_global_remote_catalog_and_refreshes_it() -> Res force_refetch: false, }) .await?; - let response: PluginListResponse = to_response( - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??, - )?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let remote_marketplace = response .marketplaces .iter() @@ -2194,9 +2093,8 @@ async fn plugin_list_includes_openai_curated_remote_collection_when_remote_plugi let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2205,12 +2103,8 @@ async fn plugin_list_includes_openai_curated_remote_collection_when_remote_plugi force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let remote_marketplace = response .marketplaces @@ -2289,9 +2183,8 @@ async fn plugin_list_propagates_openai_curated_remote_collection_errors_when_rem let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2334,9 +2227,8 @@ async fn plugin_list_skips_openai_curated_remote_collection_for_api_auth_when_re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2345,12 +2237,8 @@ async fn plugin_list_skips_openai_curated_remote_collection_for_api_auth_when_re force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(response.marketplaces.is_empty()); assert!(response.marketplace_load_errors.is_empty()); @@ -2378,9 +2266,8 @@ async fn plugin_list_includes_api_curated_marketplace_for_api_auth_when_remote_p let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2389,12 +2276,8 @@ async fn plugin_list_includes_api_curated_marketplace_for_api_auth_when_remote_p force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let api_curated_marketplace = response .marketplaces @@ -2438,9 +2321,8 @@ async fn plugin_list_does_not_query_openai_curated_remote_collection_by_default( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2449,12 +2331,8 @@ async fn plugin_list_does_not_query_openai_curated_remote_collection_by_default( force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!( response @@ -2496,9 +2374,8 @@ async fn plugin_list_vertical_kind_noops_when_remote_plugin_enabled() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2507,12 +2384,8 @@ async fn plugin_list_vertical_kind_noops_when_remote_plugin_enabled() -> Result< force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!( response @@ -2555,9 +2428,8 @@ async fn plugin_list_does_not_append_global_remote_when_marketplace_kinds_are_ex let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -2567,12 +2439,8 @@ async fn plugin_list_does_not_append_global_remote_when_marketplace_kinds_are_ex }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!( response @@ -2639,9 +2507,8 @@ plugin_sharing = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_installed_request(PluginInstalledParams { @@ -2650,12 +2517,8 @@ plugin_sharing = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstalledResponse = to_response(response)?; + let response: PluginInstalledResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); let marketplace = &response.marketplaces[0]; @@ -2754,9 +2617,8 @@ plugin_sharing = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_installed_request(PluginInstalledParams { @@ -2765,12 +2627,8 @@ plugin_sharing = false }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstalledResponse = to_response(response)?; + let response: PluginInstalledResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); let marketplace = &response.marketplaces[0]; @@ -2846,9 +2704,8 @@ plugin_sharing = false .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_installed_request(PluginInstalledParams { @@ -2856,12 +2713,8 @@ plugin_sharing = false install_suggestion_plugin_names: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginInstalledResponse = to_response(response)?; + let response: PluginInstalledResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); assert_eq!(response.marketplaces[0].name, "created-by-me-remote"); @@ -2937,9 +2790,8 @@ trusted_hash = "sha256:unrelated" .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; trigger_plugin_installed_sync(&mut mcp).await?; let plugin_ids = [ @@ -3030,9 +2882,8 @@ async fn plugin_installed_hook_trust_write_failure_stays_untrusted() -> Result<( .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let original_permissions = std::fs::metadata(config_target_dir.path())?.permissions(); let _permission_guard = RestorePermissions( @@ -3104,9 +2955,8 @@ async fn plugin_list_fetches_workspace_directory_kind_when_remote_plugin_disable let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3116,12 +2966,8 @@ async fn plugin_list_fetches_workspace_directory_kind_when_remote_plugin_disable }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); let marketplace = &response.marketplaces[0]; @@ -3235,9 +3081,8 @@ plugin_sharing = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3247,12 +3092,8 @@ plugin_sharing = false }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); let marketplace = &response.marketplaces[0]; @@ -3367,9 +3208,8 @@ async fn plugin_list_fetches_shared_with_me_kind() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3379,12 +3219,8 @@ async fn plugin_list_fetches_shared_with_me_kind() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 2); let marketplace = response @@ -3529,9 +3365,8 @@ plugin_sharing = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3541,12 +3376,8 @@ plugin_sharing = false }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -3594,9 +3425,8 @@ plugin_sharing = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3606,12 +3436,8 @@ plugin_sharing = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -3729,9 +3555,8 @@ async fn plugin_list_marks_remote_plugin_disabled_by_admin() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3741,12 +3566,8 @@ async fn plugin_list_marks_remote_plugin_disabled_by_admin() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let remote_marketplace = response .marketplaces .into_iter() @@ -3794,9 +3615,8 @@ remote_plugin = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3806,12 +3626,8 @@ remote_plugin = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(response.marketplaces.is_empty()); wait_for_remote_plugin_request_count(&server, "/ps/plugins/list", /*expected_count*/ 0).await?; @@ -3835,9 +3651,8 @@ async fn plugin_list_fetches_featured_plugin_ids_without_chatgpt_auth() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_list_request(PluginListParams { @@ -3847,12 +3662,8 @@ async fn plugin_list_fetches_featured_plugin_ids_without_chatgpt_auth() -> Resul }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response.featured_plugin_ids, @@ -3880,9 +3691,8 @@ async fn plugin_list_uses_warmed_featured_plugin_ids_cache_on_first_request() -> .with_codex_home(codex_home.path()) .without_auto_env() .with_plugin_startup_tasks() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; wait_for_featured_plugin_request_count(&server, /*expected_count*/ 1).await?; let request_id = mcp @@ -3893,12 +3703,8 @@ async fn plugin_list_uses_warmed_featured_plugin_ids_cache_on_first_request() -> }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response.featured_plugin_ids, @@ -4042,12 +3848,8 @@ async fn trigger_plugin_installed_sync(mcp: &mut TestAppServer) -> Result<()> { install_suggestion_plugin_names: None, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: PluginInstalledResponse = to_response(response)?; + let _: PluginInstalledResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; Ok(()) } @@ -4064,10 +3866,7 @@ async fn wait_for_plugin_hooks( cwds: vec![cwd.to_path_buf()], }) .await?; - let response: JSONRPCResponse = mcp - .read_stream_until_response_message(RequestId::Integer(request_id)) - .await?; - let HooksListResponse { data } = to_response(response)?; + let HooksListResponse { data } = mcp.read_response(request_id).await?; let hooks = data .into_iter() .flat_map(|entry| entry.hooks) diff --git a/codex-rs/app-server/tests/suite/v2/plugin_read.rs b/codex-rs/app-server/tests/suite/v2/plugin_read.rs index aa0870ce7179..75735e19ea00 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_read.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_read.rs @@ -63,9 +63,8 @@ async fn plugin_read_rejects_missing_read_source() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -96,9 +95,8 @@ async fn plugin_read_rejects_multiple_read_sources() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -275,9 +273,8 @@ apps = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -287,12 +284,8 @@ apps = true }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin.marketplace_name, "openai-curated-remote"); assert_eq!( @@ -454,9 +447,8 @@ async fn plugin_read_returns_share_context_for_shared_remote_plugin() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; for remote_marketplace_name in [ "workspace-shared-with-me-private", @@ -470,12 +462,8 @@ async fn plugin_read_returns_share_context_for_shared_remote_plugin() -> Result< }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin.marketplace_name, "workspace-shared-with-me"); assert_eq!( @@ -666,9 +654,8 @@ async fn plugin_read_includes_share_url_for_admin_disabled_remote_plugin() -> Re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -678,12 +665,8 @@ async fn plugin_read_includes_share_url_for_admin_disabled_remote_plugin() -> Re }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin.marketplace_name, "openai-curated-remote"); assert_eq!(response.plugin.marketplace_path, None); @@ -803,9 +786,8 @@ async fn plugin_skill_read_reads_remote_skill_contents_when_remote_plugin_enable let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_skill_read_request(PluginSkillReadParams { @@ -815,12 +797,8 @@ async fn plugin_skill_read_reads_remote_skill_contents_when_remote_plugin_enable }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginSkillReadResponse = to_response(response)?; + let response: PluginSkillReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -859,9 +837,8 @@ async fn plugin_read_maps_missing_remote_plugin_to_invalid_request() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -915,9 +892,8 @@ remote_plugin = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -949,9 +925,8 @@ async fn plugin_read_rejects_invalid_remote_plugin_name() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1011,9 +986,8 @@ enabled = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let marketplace_path = AbsolutePathBuf::try_from(repo_root.path().join(".agents/plugins/marketplace.json"))?; @@ -1120,9 +1094,8 @@ async fn plugin_read_returns_share_context_for_shared_local_plugin() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1134,12 +1107,8 @@ async fn plugin_read_returns_share_context_for_shared_local_plugin() -> Result<( }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin.summary.remote_plugin_id, None); assert_eq!( @@ -1254,9 +1223,8 @@ async fn plugin_read_keeps_remote_version_when_share_principals_are_missing() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1268,12 +1236,8 @@ async fn plugin_read_keeps_remote_version_when_share_principals_are_missing() -> }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin.summary.remote_plugin_id, None); assert_eq!( @@ -1314,9 +1278,8 @@ async fn plugin_read_falls_back_to_local_share_context_without_remote_auth() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1328,12 +1291,8 @@ async fn plugin_read_falls_back_to_local_share_context_without_remote_auth() -> }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin.summary.remote_plugin_id, None); assert_eq!(response.plugin.summary.local_version, None); @@ -1376,9 +1335,8 @@ async fn plugin_read_fails_on_malformed_share_mapping() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1574,9 +1532,8 @@ enabled = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let marketplace_path = AbsolutePathBuf::try_from(repo_root.path().join(".agents/plugins/marketplace.json"))?; @@ -1588,12 +1545,8 @@ enabled = false }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin.marketplace_name, "codex-curated"); assert_eq!(response.plugin.marketplace_path, Some(marketplace_path)); @@ -1780,9 +1733,8 @@ async fn plugin_read_returns_app_metadata_category() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1792,12 +1744,8 @@ async fn plugin_read_returns_app_metadata_category() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response @@ -1878,9 +1826,8 @@ async fn plugin_read_hides_apps_for_api_key_auth() -> Result<()> { ("CODEX_API_KEY", None), ("OPENAI_API_KEY", None), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1890,12 +1837,8 @@ async fn plugin_read_hides_apps_for_api_key_auth() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert!(response.plugin.apps.is_empty()); assert_eq!(response.plugin.mcp_servers, vec!["alpha".to_string()]); @@ -1942,9 +1885,8 @@ async fn plugin_read_accepts_legacy_string_default_prompt() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -1956,12 +1898,8 @@ async fn plugin_read_accepts_legacy_string_default_prompt() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response @@ -2008,9 +1946,8 @@ async fn plugin_read_describes_uninstalled_git_source_without_cloning() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -2022,12 +1959,8 @@ async fn plugin_read_describes_uninstalled_git_source_without_cloning() -> Resul }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginReadResponse = to_response(response)?; + let response: PluginReadResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let expected_description = format!( "This is a cross-repo plugin. Install it to view more detailed information. The source of the plugin is {missing_remote_repo_url}, path `plugins/toolkit`." @@ -2075,9 +2008,8 @@ async fn plugin_read_returns_invalid_request_when_plugin_is_missing() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { @@ -2132,9 +2064,8 @@ async fn plugin_read_returns_invalid_request_when_plugin_manifest_is_missing() - let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_read_request(PluginReadParams { diff --git a/codex-rs/app-server/tests/suite/v2/plugin_share.rs b/codex-rs/app-server/tests/suite/v2/plugin_share.rs index 5e6196d12fc0..b2f34c8d2459 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_share.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_share.rs @@ -5,10 +5,8 @@ use std::time::Duration; use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PluginAuthPolicy; use codex_app_server_protocol::PluginInstallPolicy; use codex_app_server_protocol::PluginInterface; @@ -105,9 +103,8 @@ async fn plugin_share_save_uploads_local_plugin() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let expected_plugin_path = AbsolutePathBuf::try_from(plugin_path.clone())?; let request_id = mcp .send_raw_request( @@ -118,12 +115,8 @@ async fn plugin_share_save_uploads_local_plugin() -> Result<()> { ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareSaveResponse = to_response(response)?; + let response: PluginShareSaveResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -161,12 +154,8 @@ async fn plugin_share_save_uploads_local_plugin() -> Result<()> { let request_id = mcp .send_raw_request("plugin/share/list", Some(json!({}))) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareListResponse = to_response(response)?; + let response: PluginShareListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -261,9 +250,8 @@ async fn plugin_share_save_forwards_access_policy() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let expected_plugin_path = AbsolutePathBuf::try_from(plugin_path)?; let request_id = mcp .send_raw_request( @@ -282,12 +270,8 @@ async fn plugin_share_save_forwards_access_policy() -> Result<()> { ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareSaveResponse = to_response(response)?; + let response: PluginShareSaveResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -318,9 +302,8 @@ async fn plugin_share_save_rejects_listed_discoverability() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/save", @@ -376,9 +359,8 @@ plugin_sharing = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/save", @@ -425,9 +407,8 @@ async fn plugin_share_rejects_workspace_targets_from_client() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/save", @@ -507,9 +488,8 @@ async fn plugin_share_save_rejects_access_policy_for_existing_plugin() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/save", @@ -584,19 +564,14 @@ async fn plugin_share_list_returns_created_workspace_plugins() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request("plugin/share/list", Some(json!({}))) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareListResponse = to_response(response)?; + let response: PluginShareListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -667,9 +642,8 @@ async fn plugin_share_checkout_adds_personal_marketplace_entry() -> Result<()> { ("USERPROFILE", Some(home_env.as_str())), (TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1")), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -679,12 +653,8 @@ async fn plugin_share_checkout_adds_personal_marketplace_entry() -> Result<()> { })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareCheckoutResponse = to_response(response)?; + let response: PluginShareCheckoutResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; let plugin_path = AbsolutePathBuf::try_from(home.path().join("plugins/demo-plugin"))?; let marketplace_path = @@ -756,12 +726,8 @@ async fn plugin_share_checkout_adds_personal_marketplace_entry() -> Result<()> { force_refetch: false, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginListResponse = to_response(response)?; + let response: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.marketplaces.len(), 1); assert_eq!(response.marketplaces[0].name, "codex-curated"); assert_eq!(response.marketplaces[0].plugins[0].name, "demo-plugin"); @@ -782,12 +748,8 @@ async fn plugin_share_checkout_adds_personal_marketplace_entry() -> Result<()> { })), ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareCheckoutResponse = to_response(response)?; + let response: PluginShareCheckoutResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response.plugin_path, plugin_path); assert_eq!( std::fs::read_to_string(plugin_path.as_path().join("local-edit.txt"))?, @@ -832,9 +794,8 @@ async fn plugin_share_checkout_rejects_non_share_remote_plugin() -> Result<()> { ("USERPROFILE", Some(home_env.as_str())), (TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1")), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -924,9 +885,8 @@ async fn plugin_share_checkout_cleans_up_path_when_marketplace_update_fails() -> ("USERPROFILE", Some(home_env.as_str())), (TEST_ALLOW_HTTP_REMOTE_PLUGIN_BUNDLE_DOWNLOADS, Some("1")), ]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( @@ -1023,9 +983,8 @@ async fn plugin_share_update_targets_updates_share_targets() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/updateTargets", @@ -1043,12 +1002,8 @@ async fn plugin_share_update_targets_updates_share_targets() -> Result<()> { ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareUpdateTargetsResponse = to_response(response)?; + let response: PluginShareUpdateTargetsResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1119,9 +1074,8 @@ async fn plugin_share_update_targets_publishes_workspace_plugin() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/updateTargets", @@ -1133,12 +1087,8 @@ async fn plugin_share_update_targets_publishes_workspace_plugin() -> Result<()> ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareUpdateTargetsResponse = to_response(response)?; + let response: PluginShareUpdateTargetsResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1184,9 +1134,8 @@ plugin_sharing = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/updateTargets", @@ -1237,9 +1186,8 @@ async fn plugin_share_delete_removes_created_workspace_plugin() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_raw_request( "plugin/share/delete", @@ -1249,12 +1197,8 @@ async fn plugin_share_delete_removes_created_workspace_plugin() -> Result<()> { ) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareDeleteResponse = to_response(response)?; + let response: PluginShareDeleteResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(response, PluginShareDeleteResponse {}); @@ -1286,12 +1230,8 @@ async fn plugin_share_delete_removes_created_workspace_plugin() -> Result<()> { let request_id = mcp .send_raw_request("plugin/share/list", Some(json!({}))) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginShareListResponse = to_response(response)?; + let response: PluginShareListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, diff --git a/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs b/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs index 94efd617fc55..e21192356bce 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_uninstall.rs @@ -6,9 +6,7 @@ use app_test_support::ChatGptAuthFixture; use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::TestAppServer; use app_test_support::start_analytics_events_server; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PluginUninstallParams; use codex_app_server_protocol::PluginUninstallResponse; use codex_app_server_protocol::RequestId; @@ -45,21 +43,10 @@ enabled = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let params = PluginUninstallParams { - plugin_id: "sample-plugin@debug".to_string(), - }; - - let request_id = mcp.send_plugin_uninstall_request(params.clone()).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginUninstallResponse = to_response(response)?; + let response = uninstall_plugin(&mut mcp, "sample-plugin@debug").await?; assert_eq!(response, PluginUninstallResponse {}); assert!( @@ -71,13 +58,7 @@ enabled = true let config = std::fs::read_to_string(codex_home.path().join("config.toml"))?; assert!(!config.contains(r#"[plugins."sample-plugin@debug"]"#)); - let request_id = mcp.send_plugin_uninstall_request(params).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginUninstallResponse = to_response(response)?; + let response = uninstall_plugin(&mut mcp, "sample-plugin@debug").await?; assert_eq!(response, PluginUninstallResponse {}); Ok(()) @@ -107,21 +88,10 @@ async fn plugin_uninstall_tracks_analytics_event() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_plugin_uninstall_request(PluginUninstallParams { - plugin_id: "sample-plugin@debug".to_string(), - }) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginUninstallResponse = to_response(response)?; + let response = uninstall_plugin(&mut mcp, "sample-plugin@debug").await?; assert_eq!(response, PluginUninstallResponse {}); let payload = timeout(DEFAULT_TIMEOUT, async { @@ -173,9 +143,8 @@ plugins = false let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_uninstall_request(PluginUninstallParams { @@ -256,25 +225,14 @@ async fn plugin_uninstall_writes_remote_plugin_to_cloud_when_remote_plugin_enabl let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; // Simulate a background remote-cache refresh removing the local bundle // before the uninstall request captures its telemetry metadata. std::fs::remove_dir_all(remote_plugin_cache_root.join("1.0.0"))?; - let request_id = mcp - .send_plugin_uninstall_request(PluginUninstallParams { - plugin_id: REMOTE_PLUGIN_ID.to_string(), - }) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginUninstallResponse = to_response(response)?; + let response = uninstall_plugin(&mut mcp, REMOTE_PLUGIN_ID).await?; assert_eq!(response, PluginUninstallResponse {}); wait_for_remote_plugin_request_count( @@ -355,21 +313,10 @@ async fn plugin_uninstall_uses_detail_scope_for_cache_namespace() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_plugin_uninstall_request(PluginUninstallParams { - plugin_id: REMOTE_PLUGIN_ID.to_string(), - }) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginUninstallResponse = to_response(response)?; + let response = uninstall_plugin(&mut mcp, REMOTE_PLUGIN_ID).await?; assert_eq!(response, PluginUninstallResponse {}); wait_for_remote_plugin_request_count( @@ -433,21 +380,10 @@ async fn plugin_uninstall_accepts_workspace_remote_plugin_id_shape() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_plugin_uninstall_request(PluginUninstallParams { - plugin_id: WORKSPACE_REMOTE_PLUGIN_ID.to_string(), - }) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: PluginUninstallResponse = to_response(response)?; + let response = uninstall_plugin(&mut mcp, WORKSPACE_REMOTE_PLUGIN_ID).await?; assert_eq!(response, PluginUninstallResponse {}); wait_for_remote_plugin_request_count( @@ -486,9 +422,8 @@ async fn plugin_uninstall_rejects_before_post_when_remote_detail_fetch_fails() - let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_uninstall_request(PluginUninstallParams { @@ -532,9 +467,8 @@ async fn plugin_uninstall_rejects_remote_plugin_id_with_spaces_before_network_ca let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_uninstall_request(PluginUninstallParams { @@ -571,9 +505,8 @@ async fn plugin_uninstall_rejects_invalid_remote_plugin_id_before_network_call() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_uninstall_request(PluginUninstallParams { @@ -610,9 +543,8 @@ async fn plugin_uninstall_rejects_empty_remote_plugin_id() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_plugin_uninstall_request(PluginUninstallParams { @@ -631,6 +563,18 @@ async fn plugin_uninstall_rejects_empty_remote_plugin_id() -> Result<()> { Ok(()) } +async fn uninstall_plugin( + mcp: &mut TestAppServer, + plugin_id: &str, +) -> Result { + let request_id = mcp + .send_plugin_uninstall_request(PluginUninstallParams { + plugin_id: plugin_id.to_string(), + }) + .await?; + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await? +} + fn write_installed_plugin( codex_home: &TempDir, marketplace_name: &str, diff --git a/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs b/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs index 25da2cd85e03..288a6154f60d 100644 --- a/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs +++ b/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs @@ -3,19 +3,16 @@ use std::path::Path; use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::ConsumeAccountRateLimitResetCreditOutcome; use codex_app_server_protocol::ConsumeAccountRateLimitResetCreditParams; use codex_app_server_protocol::ConsumeAccountRateLimitResetCreditResponse; use codex_app_server_protocol::GetAccountParams; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; use pretty_assertions::assert_eq; -use serde::de::DeserializeOwned; use serde_json::json; use tempfile::TempDir; use tokio::time::timeout; @@ -151,7 +148,11 @@ async fn consume_account_rate_limit_reset_credit_forwards_selected_credit_id() - .await?; assert_eq!( - read_response::(&mut mcp, request_id).await?, + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_response::(request_id), + ) + .await??, ConsumeAccountRateLimitResetCreditResponse { outcome: ConsumeAccountRateLimitResetCreditOutcome::Reset, } @@ -244,9 +245,8 @@ async fn consume_timeout_releases_account_auth_queue() -> Result<()> { ("OPENAI_API_KEY", None), (RATE_LIMIT_RESET_REQUEST_TIMEOUT_ENV_VAR, Some("100")), ]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let consume_id = send_consume_reset_credit(&mut mcp, "request-timeout").await?; let account_id = mcp .send_get_account_request(GetAccountParams { @@ -288,14 +288,12 @@ async fn chatgpt_test_context() -> Result<(TempDir, MockServer)> { } async fn initialized_app_server(codex_home: &Path) -> Result { - let mut mcp = TestAppServer::builder() + TestAppServer::builder() .with_codex_home(codex_home) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() - .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - Ok(mcp) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await } async fn consume_reset_credit( @@ -303,7 +301,7 @@ async fn consume_reset_credit( idempotency_key: &str, ) -> Result { let request_id = send_consume_reset_credit(mcp, idempotency_key).await?; - read_response(mcp, request_id).await + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await? } async fn send_consume_reset_credit(mcp: &mut TestAppServer, idempotency_key: &str) -> Result { @@ -316,18 +314,6 @@ async fn send_consume_reset_credit(mcp: &mut TestAppServer, idempotency_key: &st .await } -async fn read_response(mcp: &mut TestAppServer, request_id: i64) -> Result -where - T: DeserializeOwned, -{ - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) -} - async fn read_error_response(mcp: &mut TestAppServer, request_id: i64) -> Result { let error = timeout( DEFAULT_READ_TIMEOUT, @@ -340,7 +326,11 @@ async fn read_error_response(mcp: &mut TestAppServer, request_id: i64) -> Result async fn login_with_api_key(mcp: &mut TestAppServer, api_key: &str) -> Result<()> { let request_id = mcp.send_login_account_api_key_request(api_key).await?; assert_eq!( - read_response::(mcp, request_id).await?, + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_response::(request_id), + ) + .await??, LoginAccountResponse::ApiKey {} ); Ok(()) diff --git a/codex-rs/app-server/tests/suite/v2/rate_limits.rs b/codex-rs/app-server/tests/suite/v2/rate_limits.rs index 5518ed132cb9..3194f2cf404a 100644 --- a/codex-rs/app-server/tests/suite/v2/rate_limits.rs +++ b/codex-rs/app-server/tests/suite/v2/rate_limits.rs @@ -1,13 +1,11 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::AddCreditsNudgeCreditType; use codex_app_server_protocol::AddCreditsNudgeEmailStatus; use codex_app_server_protocol::GetAccountRateLimitsResponse; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RateLimitReachedType; use codex_app_server_protocol::RateLimitResetCredit; @@ -46,9 +44,8 @@ async fn get_account_rate_limits_requires_auth() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_get_account_rate_limits_request().await?; @@ -75,9 +72,8 @@ async fn get_account_rate_limits_requires_chatgpt_auth() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; @@ -225,19 +221,13 @@ async fn get_account_rate_limits_returns_snapshot() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_get_account_rate_limits_request().await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - - let received: GetAccountRateLimitsResponse = to_response(response)?; + let received: GetAccountRateLimitsResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountRateLimitsResponse { rate_limits: RateLimitSnapshot { @@ -390,17 +380,12 @@ async fn get_account_rate_limits_preserves_count_when_reset_credit_details_fail( .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_get_account_rate_limits_request().await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: GetAccountRateLimitsResponse = to_response(response)?; + let received: GetAccountRateLimitsResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( received.rate_limit_reset_credits, @@ -421,9 +406,8 @@ async fn send_add_credits_nudge_email_requires_auth() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_add_credits_nudge_email_request(SendAddCreditsNudgeEmailParams { @@ -454,9 +438,8 @@ async fn send_add_credits_nudge_email_requires_chatgpt_auth() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; @@ -513,9 +496,8 @@ async fn send_add_credits_nudge_email_posts_expected_body() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_add_credits_nudge_email_request(SendAddCreditsNudgeEmailParams { @@ -523,12 +505,8 @@ async fn send_add_credits_nudge_email_posts_expected_body() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: SendAddCreditsNudgeEmailResponse = to_response(response)?; + let received: SendAddCreditsNudgeEmailResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(received.status, AddCreditsNudgeEmailStatus::Sent); @@ -561,9 +539,8 @@ async fn send_add_credits_nudge_email_maps_cooldown() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_add_credits_nudge_email_request(SendAddCreditsNudgeEmailParams { @@ -571,12 +548,8 @@ async fn send_add_credits_nudge_email_maps_cooldown() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: SendAddCreditsNudgeEmailResponse = to_response(response)?; + let received: SendAddCreditsNudgeEmailResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(received.status, AddCreditsNudgeEmailStatus::CooldownActive); @@ -609,9 +582,8 @@ async fn send_add_credits_nudge_email_surfaces_backend_failure() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_add_credits_nudge_email_request(SendAddCreditsNudgeEmailParams { @@ -642,12 +614,8 @@ async fn send_add_credits_nudge_email_surfaces_backend_failure() -> Result<()> { async fn login_with_api_key(mcp: &mut TestAppServer, api_key: &str) -> Result<()> { let request_id = mcp.send_login_account_api_key_request(api_key).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let login: LoginAccountResponse = to_response(response)?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(login, LoginAccountResponse::ApiKey {}); Ok(()) diff --git a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs index bb7a77f0548a..bb3957cf9bd3 100644 --- a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs +++ b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs @@ -1,15 +1,14 @@ use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::create_shell_command_sse_response; -use app_test_support::to_response; use codex_app_server_protocol::CommandExecutionStatus; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; @@ -43,7 +42,6 @@ use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStartedNotification; use codex_app_server_protocol::UserInput as V2UserInput; -use codex_features::FEATURES; use codex_features::Feature; use codex_protocol::protocol::CodexResponseHandoffMode; use codex_protocol::protocol::ConversationTextRole; @@ -304,20 +302,15 @@ impl RealtimeE2eHarness { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; Ok(Self { mcp, @@ -388,13 +381,8 @@ impl RealtimeE2eHarness { voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - self.mcp - .read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, self.mcp.read_response(start_request_id)).await??; let started = self .read_notification::("thread/realtime/started") @@ -448,13 +436,8 @@ impl RealtimeE2eHarness { voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - self.mcp - .read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, self.mcp.read_response(start_request_id)).await??; self.read_notification::("thread/realtime/started") .await @@ -485,13 +468,8 @@ impl RealtimeE2eHarness { voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - self.mcp - .read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, self.mcp.read_response(start_request_id)).await??; self.read_notification::("thread/realtime/started") .await @@ -528,13 +506,8 @@ impl RealtimeE2eHarness { }, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - self.mcp - .read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadRealtimeAppendAudioResponse = to_response(response)?; + let _: ThreadRealtimeAppendAudioResponse = + timeout(DEFAULT_TIMEOUT, self.mcp.read_response(request_id)).await??; Ok(()) } @@ -547,13 +520,8 @@ impl RealtimeE2eHarness { role: ConversationTextRole::User, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - self.mcp - .read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadRealtimeAppendTextResponse = to_response(response)?; + let _: ThreadRealtimeAppendTextResponse = + timeout(DEFAULT_TIMEOUT, self.mcp.read_response(request_id)).await??; Ok(()) } @@ -565,13 +533,8 @@ impl RealtimeE2eHarness { text: text.to_string(), }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - self.mcp - .read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadRealtimeAppendSpeechResponse = to_response(response)?; + let _: ThreadRealtimeAppendSpeechResponse = + timeout(DEFAULT_TIMEOUT, self.mcp.read_response(request_id)).await??; Ok(()) } @@ -724,20 +687,15 @@ async fn realtime_conversation_streams_v2_notifications() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { @@ -758,12 +716,8 @@ async fn realtime_conversation_streams_v2_notifications() -> Result<()> { voice: Some(RealtimeVoice::Cedar), }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(start_request_id)).await??; let started = read_notification::(&mut mcp, "thread/realtime/started") @@ -811,12 +765,8 @@ async fn realtime_conversation_streams_v2_notifications() -> Result<()> { }, }) .await?; - let audio_append_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(audio_append_request_id)), - ) - .await??; - let _: ThreadRealtimeAppendAudioResponse = to_response(audio_append_response)?; + let _: ThreadRealtimeAppendAudioResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(audio_append_request_id)).await??; let text_append_request_id = mcp .send_thread_realtime_append_text_request(ThreadRealtimeAppendTextParams { @@ -825,12 +775,8 @@ async fn realtime_conversation_streams_v2_notifications() -> Result<()> { role: ConversationTextRole::Developer, }) .await?; - let text_append_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(text_append_request_id)), - ) - .await??; - let _: ThreadRealtimeAppendTextResponse = to_response(text_append_response)?; + let _: ThreadRealtimeAppendTextResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(text_append_request_id)).await??; let assistant_append_request_id = mcp .send_thread_realtime_append_text_request(ThreadRealtimeAppendTextParams { @@ -839,12 +785,11 @@ async fn realtime_conversation_streams_v2_notifications() -> Result<()> { role: ConversationTextRole::Assistant, }) .await?; - let assistant_append_response: JSONRPCResponse = timeout( + let _: ThreadRealtimeAppendTextResponse = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(assistant_append_request_id)), + mcp.read_response(assistant_append_request_id), ) .await??; - let _: ThreadRealtimeAppendTextResponse = to_response(assistant_append_response)?; let output_audio = read_notification::( &mut mcp, @@ -1020,20 +965,15 @@ async fn realtime_start_can_skip_startup_context() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { @@ -1054,12 +994,8 @@ async fn realtime_start_can_skip_startup_context() -> Result<()> { voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(start_request_id)).await??; read_notification::(&mut mcp, "thread/realtime/started") .await?; @@ -1123,20 +1059,15 @@ async fn realtime_text_output_modality_requests_text_output_and_final_transcript let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { @@ -1157,12 +1088,8 @@ async fn realtime_text_output_modality_requests_text_output_and_final_transcript voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(start_request_id)).await??; let session_update = realtime_server .wait_for_request(/*connection_index*/ 0, /*request_index*/ 0) @@ -1237,19 +1164,14 @@ async fn realtime_list_voices_returns_supported_names() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_realtime_list_voices_request(ThreadRealtimeListVoicesParams {}) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadRealtimeListVoicesResponse = to_response(response)?; + let response: ThreadRealtimeListVoicesResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( response, @@ -1312,20 +1234,15 @@ async fn realtime_conversation_stop_emits_closed_notification() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { @@ -1346,12 +1263,8 @@ async fn realtime_conversation_stop_emits_closed_notification() -> Result<()> { voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(start_request_id)).await??; let started = read_notification::(&mut mcp, "thread/realtime/started") @@ -1362,12 +1275,8 @@ async fn realtime_conversation_stop_emits_closed_notification() -> Result<()> { thread_id: started.thread_id.clone(), }) .await?; - let stop_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(stop_request_id)), - ) - .await??; - let _: ThreadRealtimeStopResponse = to_response(stop_response)?; + let _: ThreadRealtimeStopResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(stop_request_id)).await??; let closed = read_notification::(&mut mcp, "thread/realtime/closed") @@ -1420,20 +1329,15 @@ async fn realtime_webrtc_start_emits_sdp_notification() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let thread_id = thread_start.thread.id; let start_request_id = mcp @@ -1457,12 +1361,8 @@ async fn realtime_webrtc_start_emits_sdp_notification() -> Result<()> { voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(start_request_id)).await??; let started = read_notification::(&mut mcp, "thread/realtime/started") @@ -1503,12 +1403,8 @@ async fn realtime_webrtc_start_emits_sdp_notification() -> Result<()> { thread_id: thread_id.clone(), }) .await?; - let stop_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(stop_request_id)), - ) - .await??; - let _: ThreadRealtimeStopResponse = to_response(stop_response)?; + let _: ThreadRealtimeStopResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(stop_request_id)).await??; let closed_notification = read_notification::(&mut mcp, "thread/realtime/closed") @@ -1722,14 +1618,8 @@ async fn webrtc_v1_default_automatic_output_uses_handoff_append() -> Result<()> ..Default::default() }) .await?; - let turn_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - harness - .mcp - .read_stream_until_response_message(RequestId::Integer(turn_request_id)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_response)?; + let _: TurnStartResponse = + timeout(DEFAULT_TIMEOUT, harness.mcp.read_response(turn_request_id)).await??; let _ = harness .read_notification::("turn/completed") .await?; @@ -1786,14 +1676,8 @@ async fn webrtc_v1_client_managed_handoffs_disable_automatic_output() -> Result< ..Default::default() }) .await?; - let turn_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - harness - .mcp - .read_stream_until_response_message(RequestId::Integer(turn_request_id)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_response)?; + let _: TurnStartResponse = + timeout(DEFAULT_TIMEOUT, harness.mcp.read_response(turn_request_id)).await??; let _ = harness .read_notification::("turn/completed") .await?; @@ -2037,14 +1921,8 @@ async fn realtime_automatic_standalone_output_is_item_and_append_speaks() -> Res ..Default::default() }) .await?; - let turn_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - harness - .mcp - .read_stream_until_response_message(RequestId::Integer(turn_request_id)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_response)?; + let _: TurnStartResponse = + timeout(DEFAULT_TIMEOUT, harness.mcp.read_response(turn_request_id)).await??; let _ = harness .read_notification::("turn/completed") .await?; @@ -2199,14 +2077,8 @@ async fn websocket_v2_assistant_output_without_handoff_reaches_realtime_context( ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - harness - .mcp - .read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: TurnStartResponse = to_response(response)?; + let _: TurnStartResponse = + timeout(DEFAULT_TIMEOUT, harness.mcp.read_response(request_id)).await??; let _ = harness .read_notification::("turn/completed") .await?; @@ -3047,21 +2919,16 @@ async fn realtime_webrtc_start_surfaces_backend_error() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; login_with_api_key(&mut mcp, "sk-test-key").await?; // Phase 2: start a normal app-server thread and request realtime over WebRTC. let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { @@ -3084,12 +2951,8 @@ async fn realtime_webrtc_start_surfaces_backend_error() -> Result<()> { voice: None, }) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_request_id)), - ) - .await??; - let _: ThreadRealtimeStartResponse = to_response(start_response)?; + let _: ThreadRealtimeStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(start_request_id)).await??; // Phase 3: the JSON-RPC start request returns, and the realtime failure is delivered as the // typed realtime error notification. @@ -3120,19 +2983,14 @@ async fn realtime_conversation_requires_feature_flag() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let thread_start_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; - let thread_start: ThreadStartResponse = to_response(thread_start_response)?; + let thread_start: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let start_request_id = mcp .send_thread_realtime_start_request(ThreadRealtimeStartParams { @@ -3182,25 +3040,13 @@ async fn read_notification_with_timeout( method: &str, timeout_duration: Duration, ) -> Result { - let notification = timeout( - timeout_duration, - mcp.read_stream_until_notification_message(method), - ) - .await??; - let params = notification - .params - .context("expected notification params to be present")?; - Ok(serde_json::from_value(params)?) + timeout(timeout_duration, mcp.read_notification(method)).await? } async fn login_with_api_key(mcp: &mut TestAppServer, api_key: &str) -> Result<()> { let request_id = mcp.send_login_account_api_key_request(api_key).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let login: LoginAccountResponse = to_response(response)?; + let login: LoginAccountResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(login, LoginAccountResponse::ApiKey {}); Ok(()) @@ -3466,48 +3312,28 @@ fn create_config_toml_with_realtime_version( realtime_version: RealtimeTestVersion, sandbox: RealtimeTestSandbox, ) -> std::io::Result<()> { - let realtime_feature_key = FEATURES - .iter() - .find(|spec| spec.id == Feature::RealtimeConversation) - .map(|spec| spec.key) - .unwrap_or("realtime_conversation"); - let realtime_version = realtime_version.config_value(); - let sandbox = sandbox.config_value(); - let startup_context = match startup_context { - StartupContextConfig::Generated => String::new(), - StartupContextConfig::Override(context) => { - format!("experimental_realtime_ws_startup_context = {context:?}\n") - } + let mut config = MockResponsesConfig::new(responses_server_uri) + .with_sandbox_mode(sandbox.config_value()) + .with_root_config(&format!( + "experimental_realtime_ws_base_url = \"{realtime_server_uri}\"\n\ + experimental_realtime_ws_backend_prompt = \"backend prompt\"" + )) + .with_extra_config(&format!( + "[realtime]\nversion = \"{}\"\ntype = \"conversational\"", + realtime_version.config_value() + )); + + if let StartupContextConfig::Override(context) = startup_context { + config = config.with_root_config(&format!( + "experimental_realtime_ws_startup_context = {context:?}" + )); + } + config = if realtime_enabled { + config.enable_feature(Feature::RealtimeConversation) + } else { + config.disable_feature(Feature::RealtimeConversation) }; - - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "{sandbox}" -model_provider = "mock_provider" -experimental_realtime_ws_base_url = "{realtime_server_uri}" -experimental_realtime_ws_backend_prompt = "backend prompt" -{startup_context} - -[realtime] -version = "{realtime_version}" -type = "conversational" - -[features] -{realtime_feature_key} = {realtime_enabled} - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{responses_server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) + config.write(codex_home) } fn assert_invalid_request(error: JSONRPCError, message: String) { diff --git a/codex-rs/app-server/tests/suite/v2/remote_control.rs b/codex-rs/app-server/tests/suite/v2/remote_control.rs index e6093ad0eea2..5f9ab44592ed 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_control.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_control.rs @@ -7,16 +7,17 @@ use anyhow::Context; use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::DEFAULT_CLIENT_NAME; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; -use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; use codex_app_server::AppServerRuntimeOptions; use codex_app_server::AppServerTransport; use codex_app_server::AppServerWebsocketAuthSettings; use codex_app_server::PluginStartupTasks; use codex_app_server::RemoteControlStartupMode; use codex_app_server::run_main_with_transport_options; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RemoteControlClient; @@ -97,12 +98,9 @@ async fn remote_control_preference( .remote_control_enabled) } -async fn wait_for_response(mcp: &mut TestAppServer, request_id: i64) -> Result { - timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await? +async fn wait_for_response(mcp: &mut TestAppServer, request_id: i64) -> Result<()> { + let _: serde_json::Value = timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; + Ok(()) } async fn assert_remote_control_disabled_by_requirements( @@ -132,20 +130,14 @@ async fn managed_requirements_reject_all_remote_control_rpcs() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let notification = timeout( + let status: RemoteControlStatusChangedNotification = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("remoteControl/status/changed"), + mcp.read_notification("remoteControl/status/changed"), ) .await??; - let status: RemoteControlStatusChangedNotification = serde_json::from_value( - notification - .params - .context("remote-control status notification should include params")?, - )?; assert_eq!(status.status, RemoteControlConnectionStatus::Disabled); assert_eq!(status.environment_id, None); @@ -194,17 +186,15 @@ async fn managed_requirements_allow_remote_control_true_does_not_enable_or_block let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp.send_remote_control_status_read_request().await?; - let response = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: RemoteControlStatusReadResponse = to_response(response)?; + let received: RemoteControlStatusReadResponse = mcp + .request(|request_id| ClientRequest::RemoteControlStatusRead { + request_id, + params: None, + }) + .await?; assert_eq!(received.status, RemoteControlConnectionStatus::Disabled); Ok(()) } @@ -393,17 +383,15 @@ async fn remote_control_disable_returns_disabled_status() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp.send_remote_control_disable_request().await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: RemoteControlDisableResponse = to_response(response)?; + let received: RemoteControlDisableResponse = mcp + .request(|request_id| ClientRequest::RemoteControlDisable { + request_id, + params: None, + }) + .await?; assert_eq!(received.status, RemoteControlConnectionStatus::Disabled); assert!(!received.server_name.is_empty()); @@ -418,17 +406,15 @@ async fn remote_control_status_read_returns_disabled_status() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp.send_remote_control_status_read_request().await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: RemoteControlStatusReadResponse = to_response(response)?; + let received: RemoteControlStatusReadResponse = mcp + .request(|request_id| ClientRequest::RemoteControlStatusRead { + request_id, + params: None, + }) + .await?; assert_eq!(received.status, RemoteControlConnectionStatus::Disabled); assert!(!received.server_name.is_empty()); @@ -444,9 +430,8 @@ async fn remote_control_enable_returns_connecting_status() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_remote_control_enable_request().await?; assert_eq!( @@ -460,12 +445,8 @@ async fn remote_control_enable_returns_connecting_status() -> Result<()> { .await .expect_err("enable response should wait for enrollment"); backend.complete_enrollment()?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: RemoteControlEnableResponse = to_response(response)?; + let received: RemoteControlEnableResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(received.status, RemoteControlConnectionStatus::Connecting); assert!(!received.server_name.is_empty()); @@ -484,9 +465,8 @@ async fn disable_waits_for_in_flight_durable_enable() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; mcp.send_remote_control_enable_request().await?; timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??; @@ -499,8 +479,8 @@ async fn disable_waits_for_in_flight_durable_enable() -> Result<()> { .expect_err("disable response should wait for the in-flight enable"); backend.complete_enrollment()?; - let response = wait_for_response(&mut mcp, disable_request_id).await?; - let received: RemoteControlDisableResponse = to_response(response)?; + let received: RemoteControlDisableResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(disable_request_id)).await??; assert_eq!(received.status, RemoteControlConnectionStatus::Disabled); assert_eq!( remote_control_preference(&state_db, &websocket_url).await?, @@ -520,9 +500,8 @@ async fn rpc_updates_durable_preference_but_ephemeral_does_not() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_remote_control_enable_request().await?; assert_eq!( @@ -581,9 +560,8 @@ async fn remote_control_status_read_returns_connecting_status_after_enable() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_remote_control_enable_request().await?; let enroll_request = timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??; @@ -592,19 +570,15 @@ async fn remote_control_status_read_returns_connecting_status_after_enable() -> "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" ); backend.complete_enrollment()?; - let _: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + let _: RemoteControlEnableResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; - let request_id = mcp.send_remote_control_status_read_request().await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: RemoteControlStatusReadResponse = to_response(response)?; + let received: RemoteControlStatusReadResponse = mcp + .request(|request_id| ClientRequest::RemoteControlStatusRead { + request_id, + params: None, + }) + .await?; assert_eq!(received.status, RemoteControlConnectionStatus::Connecting); assert!(!received.server_name.is_empty()); @@ -620,16 +594,12 @@ async fn remote_control_pairing_start_returns_pairing_artifacts() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_remote_control_enable_request().await?; - let _: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + let _: RemoteControlEnableResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??, "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" @@ -721,9 +691,8 @@ async fn pairing_start_works_after_ephemeral_enable() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_remote_control_ephemeral_enable_request().await?; wait_for_response(&mut mcp, request_id).await?; @@ -763,24 +732,20 @@ async fn remote_control_client_management_works_while_disabled() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_remote_control_clients_list_request(RemoteControlClientsListParams { - environment_id: "environment-id".to_string(), - cursor: Some("cursor-id".to_string()), - limit: Some(10), - order: Some(RemoteControlClientsListOrder::Desc), + let received: RemoteControlClientsListResponse = mcp + .request(|request_id| ClientRequest::RemoteControlClientsList { + request_id, + params: RemoteControlClientsListParams { + environment_id: "environment-id".to_string(), + cursor: Some("cursor-id".to_string()), + limit: Some(10), + order: Some(RemoteControlClientsListOrder::Desc), + }, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: RemoteControlClientsListResponse = to_response(response)?; assert_eq!( received, RemoteControlClientsListResponse { @@ -798,18 +763,15 @@ async fn remote_control_client_management_works_while_disabled() -> Result<()> { } ); - let request_id = mcp - .send_remote_control_clients_revoke_request(RemoteControlClientsRevokeParams { - environment_id: "environment-id".to_string(), - client_id: "client-id".to_string(), + let received: RemoteControlClientsRevokeResponse = mcp + .request(|request_id| ClientRequest::RemoteControlClientsRevoke { + request_id, + params: RemoteControlClientsRevokeParams { + environment_id: "environment-id".to_string(), + client_id: "client-id".to_string(), + }, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let received: RemoteControlClientsRevokeResponse = to_response(response)?; assert_eq!(received, RemoteControlClientsRevokeResponse {}); assert_eq!( timeout(DEFAULT_TIMEOUT, backend.wait_for_requests()).await??, @@ -1080,11 +1042,9 @@ struct HttpRequest { async fn configured_remote_control_listener(codex_home: &std::path::Path) -> Result { let listener = TcpListener::bind("127.0.0.1:0").await?; let remote_control_url = format!("http://{}/backend-api/", listener.local_addr()?); - write_mock_responses_config_toml_with_chatgpt_base_url( - codex_home, - &remote_control_url, - &remote_control_url, - )?; + MockResponsesConfig::new(&remote_control_url) + .with_root_config(&format!("chatgpt_base_url = \"{remote_control_url}\"")) + .write(codex_home)?; write_chatgpt_auth( codex_home, ChatGptAuthFixture::new("chatgpt-token") diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 1cb7e53a0b62..22da2b3d25a3 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -18,6 +18,7 @@ use std::path::Path; use std::sync::Arc; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_repeating_assistant; use codex_app_server::in_process; @@ -47,6 +48,7 @@ use codex_config::NoopThreadConfigLoader; use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; +use codex_features::Feature; use codex_feedback::CodexFeedback; use codex_protocol::ThreadId; use codex_protocol::models::BaseInstructions; @@ -467,27 +469,10 @@ fn create_config_toml_with_thread_store( server_uri: &str, store_id: &str, ) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }} - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - -[features] -plugins = false -"# - ), - ) + MockResponsesConfig::new(server_uri) + .with_root_config(&format!( + "experimental_thread_store = {{ type = \"in_memory\", id = \"{store_id}\" }}" + )) + .disable_feature(Feature::Plugins) + .write(codex_home) } diff --git a/codex-rs/app-server/tests/suite/v2/request_permissions.rs b/codex-rs/app-server/tests/suite/v2/request_permissions.rs index b30e63edb48c..787989b6d814 100644 --- a/codex-rs/app-server/tests/suite/v2/request_permissions.rs +++ b/codex-rs/app-server/tests/suite/v2/request_permissions.rs @@ -1,14 +1,13 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_request_permissions_sse_response; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PermissionGrantScope; use codex_app_server_protocol::PermissionsRequestApprovalResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestResolvedNotification; use codex_app_server_protocol::ThreadStartParams; @@ -16,6 +15,7 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_features::Feature; use core_test_support::skip_if_wine_exec; use tokio::time::timeout; @@ -35,45 +35,38 @@ async fn request_permissions_round_trip() -> Result<()> { create_final_assistant_message_sse_response("done")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .enable_feature(Feature::RequestPermissionsTool) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "pick a directory".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model".to_string()), - ..Default::default() + let TurnStartResponse { turn, .. } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "pick a directory".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + ..Default::default() + }, }) .await?; - let turn_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let TurnStartResponse { turn, .. } = to_response(turn_start_resp)?; let server_req = timeout( DEFAULT_READ_TIMEOUT, @@ -163,29 +156,3 @@ async fn request_permissions_round_trip() -> Result<()> { Ok(()) } - -fn create_config_toml(codex_home: &std::path::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 = "untrusted" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - -[features] -request_permissions_tool = true -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/request_user_input.rs b/codex-rs/app-server/tests/suite/v2/request_user_input.rs index d28b8a3dbcd7..5e3b99f44d23 100644 --- a/codex-rs/app-server/tests/suite/v2/request_user_input.rs +++ b/codex-rs/app-server/tests/suite/v2/request_user_input.rs @@ -1,11 +1,10 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestResolvedNotification; use codex_app_server_protocol::ThreadStartParams; @@ -60,54 +59,46 @@ async fn request_user_input_round_trip() -> Result<()> { create_final_assistant_message_sse_response("done")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "ask something".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model".to_string()), - effort: Some(ReasoningEffort::Medium), - collaboration_mode: Some(CollaborationMode { - mode: ModeKind::Plan, - settings: Settings { - model: "mock-model".to_string(), - reasoning_effort: Some(ReasoningEffort::Medium), - developer_instructions: None, - }, - }), - ..Default::default() + let TurnStartResponse { turn, .. } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "ask something".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + effort: Some(ReasoningEffort::Medium), + collaboration_mode: Some(CollaborationMode { + mode: ModeKind::Plan, + settings: Settings { + model: "mock-model".to_string(), + reasoning_effort: Some(ReasoningEffort::Medium), + developer_instructions: None, + }, + }), + ..Default::default() + }, }) .await?; - let turn_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let TurnStartResponse { turn, .. } = to_response(turn_start_resp)?; let server_req = timeout( DEFAULT_READ_TIMEOUT, @@ -162,25 +153,3 @@ async fn request_user_input_round_trip() -> Result<()> { Ok(()) } -fn create_config_toml(codex_home: &std::path::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 = "untrusted" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/request_validation.rs b/codex-rs/app-server/tests/suite/v2/request_validation.rs index 1ffdf8cab323..788c2c3c4b73 100644 --- a/codex-rs/app-server/tests/suite/v2/request_validation.rs +++ b/codex-rs/app-server/tests/suite/v2/request_validation.rs @@ -1,10 +1,8 @@ use anyhow::Result; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -32,19 +30,14 @@ async fn request_handlers_reject_remote_image_urls() -> Result<()> { )?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_request_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let thread_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_response)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_request_id)).await??; let thread_id = thread.id; let remote_tool_output = serde_json::to_value(ResponseItem::FunctionCallOutput { diff --git a/codex-rs/app-server/tests/suite/v2/review.rs b/codex-rs/app-server/tests/suite/v2/review.rs index 9c78cb8a63d2..48af99641abe 100644 --- a/codex-rs/app-server/tests/suite/v2/review.rs +++ b/codex-rs/app-server/tests/suite/v2/review.rs @@ -1,16 +1,15 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_shell_command_sse_response; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ReviewDelivery; use codex_app_server_protocol::ReviewStartParams; @@ -25,8 +24,10 @@ use codex_app_server_protocol::ThreadStartedNotification; use codex_app_server_protocol::ThreadStatusChangedNotification; use codex_app_server_protocol::TurnItemsView; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_features::Feature; use codex_skills::system_cache_root_dir; use core_test_support::responses; use pretty_assertions::assert_eq; @@ -47,21 +48,17 @@ async fn review_start_rejects_detached_delivery_for_paginated_parent() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_id = mcp - .send_thread_start_request(ThreadStartParams { - history_mode: Some(ThreadHistoryMode::Paginated), - ..Default::default() + let ThreadStartResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadStart { + request_id, + params: ThreadStartParams { + history_mode: Some(ThreadHistoryMode::Paginated), + ..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(start_resp)?; let review_id = mcp .send_review_start_request(ReviewStartParams { @@ -113,31 +110,25 @@ async fn review_start_runs_review_turn_and_emits_code_review_item() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_id = start_default_thread(&mut mcp).await?; - - let review_req = mcp - .send_review_start_request(ReviewStartParams { - thread_id: thread_id.clone(), - delivery: Some(ReviewDelivery::Inline), - target: ReviewTarget::Commit { - sha: "1234567deadbeef".to_string(), - title: Some("Tidy UI colors".to_string()), - }, - }) - .await?; - let review_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(review_req)), - ) - .await??; let ReviewStartResponse { turn, review_thread_id, - } = to_response::(review_resp)?; + } = mcp + .request(|request_id| ClientRequest::ReviewStart { + request_id, + params: ReviewStartParams { + thread_id: thread_id.clone(), + delivery: Some(ReviewDelivery::Inline), + target: ReviewTarget::Commit { + sha: "1234567deadbeef".to_string(), + title: Some("Tidy UI colors".to_string()), + }, + }, + }) + .await?; assert_eq!(review_thread_id, thread_id.clone()); let turn_id = turn.id.clone(); assert_eq!(turn.status, TurnStatus::InProgress); @@ -157,13 +148,8 @@ async fn review_start_runs_review_turn_and_emits_code_review_item() -> Result<() // Confirm we see the EnteredReviewMode marker on the main thread. let mut saw_entered_review_mode = false; for _ in 0..10 { - let item_started: JSONRPCNotification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/started"), - ) - .await??; let started: ItemStartedNotification = - serde_json::from_value(item_started.params.expect("params must be present"))?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("item/started")).await??; match started.item { ThreadItem::EnteredReviewMode { review, .. } => { assert_eq!(started.turn_id, turn_id); @@ -183,13 +169,11 @@ async fn review_start_runs_review_turn_and_emits_code_review_item() -> Result<() // on the same turn. Ignore any other items the stream surfaces. let mut review_body: Option = None; for _ in 0..10 { - let review_notif: JSONRPCNotification = timeout( + let completed: ItemCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/completed"), + mcp.read_notification("item/completed"), ) .await??; - let completed: ItemCompletedNotification = - serde_json::from_value(review_notif.params.expect("params must be present"))?; match completed.item { ThreadItem::ExitedReviewMode { review, .. } => { assert_eq!(completed.turn_id, turn_id); @@ -226,32 +210,30 @@ async fn review_start_exec_approval_item_id_matches_command_execution_item() -> let server = create_mock_responses_server_sequence(responses).await; let codex_home = TempDir::new()?; - create_config_toml_with_approval_policy(codex_home.path(), &server.uri(), "untrusted")?; + MockResponsesConfig::new(&server.uri()) + .with_provider_name("Mock provider") + .with_approval_policy("untrusted") + .disable_feature(Feature::ShellSnapshot) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_id = start_default_thread(&mut mcp).await?; - - let review_req = mcp - .send_review_start_request(ReviewStartParams { - thread_id, - delivery: Some(ReviewDelivery::Inline), - target: ReviewTarget::Commit { - sha: "1234567deadbeef".to_string(), - title: Some("Check review approvals".to_string()), + let ReviewStartResponse { turn, .. } = mcp + .request(|request_id| ClientRequest::ReviewStart { + request_id, + params: ReviewStartParams { + thread_id, + delivery: Some(ReviewDelivery::Inline), + target: ReviewTarget::Commit { + sha: "1234567deadbeef".to_string(), + title: Some("Check review approvals".to_string()), + }, }, }) .await?; - let review_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(review_req)), - ) - .await??; - let ReviewStartResponse { turn, .. } = to_response::(review_resp)?; let turn_id = turn.id.clone(); assert_eq!(turn.items_view, TurnItemsView::NotLoaded); assert_eq!( @@ -279,13 +261,8 @@ async fn review_start_exec_approval_item_id_matches_command_execution_item() -> let mut command_item_id = None; for _ in 0..10 { - let item_started: JSONRPCNotification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/started"), - ) - .await??; let started: ItemStartedNotification = - serde_json::from_value(item_started.params.expect("params must be present"))?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("item/started")).await??; if let ThreadItem::CommandExecution { id, .. } = started.item { command_item_id = Some(id); break; @@ -316,9 +293,8 @@ async fn review_start_rejects_empty_base_branch() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_id = start_default_thread(&mut mcp).await?; let request_id = mcp @@ -387,31 +363,25 @@ async fn review_start_with_detached_delivery_returns_new_thread_id() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_id = start_default_thread(&mut mcp).await?; materialize_thread_rollout(&mut mcp, &thread_id).await?; - - let review_req = mcp - .send_review_start_request(ReviewStartParams { - thread_id: thread_id.clone(), - delivery: Some(ReviewDelivery::Detached), - target: ReviewTarget::Custom { - instructions: "detached review".to_string(), - }, - }) - .await?; - let review_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(review_req)), - ) - .await??; let ReviewStartResponse { turn, review_thread_id, - } = to_response::(review_resp)?; + } = mcp + .request(|request_id| ClientRequest::ReviewStart { + request_id, + params: ReviewStartParams { + thread_id: thread_id.clone(), + delivery: Some(ReviewDelivery::Detached), + target: ReviewTarget::Custom { + instructions: "detached review".to_string(), + }, + }, + }) + .await?; assert_eq!(turn.status, TurnStatus::InProgress); assert_eq!(turn.items_view, TurnItemsView::NotLoaded); @@ -488,9 +458,8 @@ async fn review_start_rejects_empty_commit_sha() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_id = start_default_thread(&mut mcp).await?; let request_id = mcp @@ -526,9 +495,8 @@ async fn review_start_rejects_empty_custom_instructions() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_id = start_default_thread(&mut mcp).await?; let request_id = mcp @@ -559,18 +527,12 @@ async fn review_start_rejects_empty_custom_instructions() -> Result<()> { } async fn start_default_thread(mcp: &mut TestAppServer) -> Result { - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/started"), @@ -580,22 +542,20 @@ async fn start_default_thread(mcp: &mut TestAppServer) -> Result { } async fn materialize_thread_rollout(mcp: &mut TestAppServer, thread_id: &str) -> Result<()> { - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread_id.to_string(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "materialize rollout".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread_id.to_string(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "materialize rollout".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -605,35 +565,8 @@ async fn materialize_thread_rollout(mcp: &mut TestAppServer, thread_id: &str) -> } fn create_config_toml(codex_home: &std::path::Path, server_uri: &str) -> std::io::Result<()> { - create_config_toml_with_approval_policy(codex_home, server_uri, "never") -} - -fn create_config_toml_with_approval_policy( - codex_home: &std::path::Path, - server_uri: &str, - approval_policy: &str, -) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "mock-model" -approval_policy = "{approval_policy}" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -shell_snapshot = false - -[model_providers.mock_provider] -name = "Mock provider" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) + MockResponsesConfig::new(server_uri) + .with_provider_name("Mock provider") + .disable_feature(Feature::ShellSnapshot) + .write(codex_home) } diff --git a/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs b/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs index f7839fee9e31..8a0194e96e22 100644 --- a/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs +++ b/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs @@ -1,17 +1,16 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::CodexErrorInfo; use codex_app_server_protocol::ErrorNotification; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::ModelRerouteReason; use codex_app_server_protocol::ModelReroutedNotification; use codex_app_server_protocol::ModelVerification; use codex_app_server_protocol::ModelVerificationNotification; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -19,6 +18,7 @@ use codex_app_server_protocol::TurnModerationMetadataNotification; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput; +use codex_features::Feature; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; @@ -51,40 +51,28 @@ async fn openai_model_header_mismatch_emits_model_rerouted_notification_v2() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some(REQUESTED_MODEL.to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![UserInput::Text { - text: "trigger safeguard".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let turn_start: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "trigger safeguard".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let turn_start: TurnStartResponse = to_response(turn_resp)?; let rerouted = collect_turn_notifications_and_validate_no_warning_item(&mut mcp).await?; assert_eq!( @@ -121,40 +109,28 @@ async fn cyber_policy_response_emits_typed_error_notification_v2() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some(REQUESTED_MODEL.to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![UserInput::Text { - text: "trigger cyber policy error".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let turn_start: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "trigger cyber policy error".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let turn_start: TurnStartResponse = to_response(turn_resp)?; let error = collect_cyber_policy_error_and_validate_no_reroute(&mut mcp).await?; assert_eq!( @@ -201,40 +177,28 @@ async fn response_model_field_mismatch_emits_model_rerouted_notification_v2_when let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some(REQUESTED_MODEL.to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![UserInput::Text { - text: "trigger response model check".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let turn_start: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "trigger response model check".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let turn_start: TurnStartResponse = to_response(turn_resp)?; let rerouted = collect_turn_notifications_and_validate_no_warning_item(&mut mcp).await?; assert_eq!( @@ -273,40 +237,28 @@ async fn model_verification_emits_typed_notification_and_warning_v2() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some(REQUESTED_MODEL.to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![UserInput::Text { - text: "trigger model verification".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let turn_start: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "trigger model verification".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let turn_start: TurnStartResponse = to_response(turn_resp)?; let verification = collect_model_verification_notifications_and_validate_no_warning_item(&mut mcp).await?; @@ -350,50 +302,33 @@ async fn turn_moderation_metadata_emits_typed_notification_v2() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some(REQUESTED_MODEL.to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![UserInput::Text { - text: "trigger moderation metadata".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let turn_start: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "trigger moderation metadata".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( + let metadata: TurnModerationMetadataNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + mcp.read_notification("turn/moderationMetadata"), ) .await??; - let turn_start: TurnStartResponse = to_response(turn_resp)?; - - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/moderationMetadata"), - ) - .await??; - let metadata: TurnModerationMetadataNotification = - serde_json::from_value(notification.params.ok_or_else(|| { - anyhow::anyhow!("turn/moderationMetadata notifications must include params") - })?)?; assert_eq!( metadata, TurnModerationMetadataNotification { @@ -548,28 +483,9 @@ fn is_warning_user_message_item(item: &ThreadItem) -> bool { } fn create_config_toml(codex_home: &std::path::Path, server_uri: &str) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "{REQUESTED_MODEL}" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -remote_models = false -personality = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) + MockResponsesConfig::new(server_uri) + .with_model(REQUESTED_MODEL) + .disable_feature(Feature::RemoteModels) + .enable_feature(Feature::Personality) + .write(codex_home) } diff --git a/codex-rs/app-server/tests/suite/v2/selected_environment.rs b/codex-rs/app-server/tests/suite/v2/selected_environment.rs index 6dfeea20f195..f7fb2c82eb69 100644 --- a/codex-rs/app-server/tests/suite/v2/selected_environment.rs +++ b/codex-rs/app-server/tests/suite/v2/selected_environment.rs @@ -1,14 +1,11 @@ -use std::collections::BTreeMap; +use std::path::Path; use std::time::Duration; use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::PathBufExt; use app_test_support::TestAppServer; -use app_test_support::to_response; -use app_test_support::write_mock_responses_config_toml; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; @@ -21,6 +18,13 @@ use tokio::time::timeout; const AGENTS_INSTRUCTIONS: &str = "selected environment workspace instructions"; const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); +fn write_mock_config(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { + MockResponsesConfig::new(server_uri) + .with_root_config("compact_prompt = \"compact\"\nmodel_auto_compact_token_limit = 100000") + .with_provider_config("supports_websockets = false") + .write(codex_home) +} + fn text_turn_params(thread_id: String, prompt: &str) -> TurnStartParams { TurnStartParams { thread_id, @@ -36,20 +40,11 @@ fn text_turn_params(thread_id: String, prompt: &str) -> TurnStartParams { async fn thread_start_reports_selected_environment_metadata() -> Result<()> { let server = responses::start_mock_server().await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 100_000, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; + write_mock_config(codex_home.path(), &server.uri())?; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; let selected_workspace_roots = app_server .auto_env()? .selection() @@ -58,20 +53,14 @@ async fn thread_start_reports_selected_environment_metadata() -> Result<()> { .filter_map(|root| root.to_abs_path().ok()) .collect::>(); - let request_id = app_server - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ThreadStartResponse { cwd, runtime_workspace_roots, active_permission_profile, .. - } = to_response(response)?; + } = app_server + .start_thread(ThreadStartParams::default()) + .await?; let host_cwd = codex_home.path().to_path_buf().abs().canonicalize()?; let cwd = cwd.canonicalize()?; assert_eq!( @@ -101,20 +90,11 @@ async fn thread_start_reports_selected_environment_instruction_source() -> Resul ) .await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 100_000, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; + write_mock_config(codex_home.path(), &server.uri())?; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; let (agents_source, environment_cwd) = { let auto_env = app_server.auto_env()?; @@ -132,15 +112,9 @@ async fn thread_start_reports_selected_environment_instruction_source() -> Resul (agents_source, environment_cwd) }; - let request_id = app_server - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) + let response = app_server + .start_thread(ThreadStartParams::default()) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadStartResponse = to_response(response)?; assert_eq!(response.instruction_sources, vec![agents_source.into()]); timeout( @@ -179,20 +153,11 @@ async fn turn_model_context_uses_selected_environment() -> Result<()> { ) .await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml( - codex_home.path(), - &server.uri(), - &BTreeMap::new(), - /*auto_compact_limit*/ 100_000, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - )?; + write_mock_config(codex_home.path(), &server.uri())?; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; let (environment_cwd, environment_shell) = { let auto_env = app_server.auto_env()?; ( @@ -201,15 +166,10 @@ async fn turn_model_context_uses_selected_environment() -> Result<()> { ) }; - let request_id = app_server - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(response)?; + let thread = app_server + .start_thread(ThreadStartParams::default()) + .await? + .thread; timeout( DEFAULT_READ_TIMEOUT, app_server.start_turn_and_wait_for_completion(text_turn_params( diff --git a/codex-rs/app-server/tests/suite/v2/session_end.rs b/codex-rs/app-server/tests/suite/v2/session_end.rs index c7d59c04dd3c..93592f296745 100644 --- a/codex-rs/app-server/tests/suite/v2/session_end.rs +++ b/codex-rs/app-server/tests/suite/v2/session_end.rs @@ -4,11 +4,9 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; 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::ThreadArchiveParams; use codex_app_server_protocol::ThreadArchiveResponse; use codex_app_server_protocol::ThreadDeleteParams; @@ -18,6 +16,7 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput; +use codex_features::Feature; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -42,9 +41,8 @@ async fn run_removal_session_end_test(operation: &str) -> Result<()> { let log_path = write_config_and_hook(codex_home.path(), &server.uri())?; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(READ_TIMEOUT) .await?; - timeout(READ_TIMEOUT, app_server.initialize()).await??; let thread_id = start_thread(&mut app_server).await?; let turn_id = app_server @@ -57,12 +55,7 @@ async fn run_removal_session_end_test(operation: &str) -> Result<()> { ..Default::default() }) .await?; - let response = timeout( - READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let _: TurnStartResponse = to_response(response)?; + let _: TurnStartResponse = timeout(READ_TIMEOUT, app_server.read_response(turn_id)).await??; timeout( READ_TIMEOUT, app_server.read_stream_until_notification_message("turn/completed"), @@ -75,24 +68,16 @@ async fn run_removal_session_end_test(operation: &str) -> Result<()> { thread_id: thread_id.clone(), }) .await?; - let response: JSONRPCResponse = timeout( - READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadArchiveResponse = to_response(response)?; + let _: ThreadArchiveResponse = + timeout(READ_TIMEOUT, app_server.read_response(request_id)).await??; } else { let request_id = app_server .send_thread_delete_request(ThreadDeleteParams { thread_id: thread_id.clone(), }) .await?; - let response: JSONRPCResponse = timeout( - READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadDeleteResponse = to_response(response)?; + let _: ThreadDeleteResponse = + timeout(READ_TIMEOUT, app_server.read_response(request_id)).await??; } let payloads = read_hook_log(&log_path)?; @@ -116,9 +101,8 @@ async fn app_server_shutdown_runs_session_end_for_all_loaded_threads() -> Result let log_path = write_config_and_hook(codex_home.path(), &server.uri())?; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(READ_TIMEOUT) .await?; - timeout(READ_TIMEOUT, app_server.initialize()).await??; let first = start_thread(&mut app_server).await?; let second = start_thread(&mut app_server).await?; @@ -152,12 +136,9 @@ async fn start_thread(app_server: &mut TestAppServer) -> Result { ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - Ok(to_response::(response)?.thread.id) + let response: ThreadStartResponse = + timeout(READ_TIMEOUT, app_server.read_response(request_id)).await??; + Ok(response.thread.id) } fn write_config_and_hook(codex_home: &Path, server_uri: &str) -> Result { @@ -181,25 +162,11 @@ with Path(r"{}").open("a", encoding="utf-8") as handle: log_path.display() ), )?; - std::fs::write( - codex_home.join("config.toml"), - format!( - r#"model = "mock-model" -approval_policy = "never" -sandbox_mode = "danger-full-access" -model_provider = "mock_provider" - -[features] -hooks = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - -[[hooks.SessionEnd]] + MockResponsesConfig::new(server_uri) + .with_sandbox_mode("danger-full-access") + .enable_feature(Feature::CodexHooks) + .with_extra_config(&format!( + r#"[[hooks.SessionEnd]] matcher = "other" [[hooks.SessionEnd.hooks]] @@ -208,8 +175,8 @@ command = "python3 {script_path}" timeout = 3 "#, script_path = script_path.display(), - ), - )?; + )) + .write(codex_home)?; Ok(log_path) } diff --git a/codex-rs/app-server/tests/suite/v2/skills_list.rs b/codex-rs/app-server/tests/suite/v2/skills_list.rs index 285c8b459377..7937cc8792b4 100644 --- a/codex-rs/app-server/tests/suite/v2/skills_list.rs +++ b/codex-rs/app-server/tests/suite/v2/skills_list.rs @@ -1,26 +1,23 @@ use std::collections::BTreeMap; use std::time::Duration; -use anyhow::Context; use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_repeating_assistant; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; -use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; use codex_app_server_protocol::ExperimentalFeatureEnablementSetParams; use codex_app_server_protocol::ExperimentalFeatureEnablementSetResponse; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PluginListParams; use codex_app_server_protocol::PluginListResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SkillsChangedNotification; use codex_app_server_protocol::SkillsExtraRootsSetParams; use codex_app_server_protocol::SkillsExtraRootsSetResponse; use codex_app_server_protocol::SkillsListParams; use codex_app_server_protocol::SkillsListResponse; use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; use codex_config::types::AuthCredentialsStoreMode; use codex_exec_server::CODEX_EXEC_SERVER_URL_ENV_VAR; use codex_utils_absolute_path::AbsolutePathBuf; @@ -51,15 +48,8 @@ async fn expect_skills_changed_notification( mcp: &mut TestAppServer, timeout_duration: Duration, ) -> Result<()> { - let notification = timeout( - timeout_duration, - mcp.read_stream_until_notification_message("skills/changed"), - ) - .await??; - let params = notification - .params - .context("skills/changed params must be present")?; - let notification: SkillsChangedNotification = serde_json::from_value(params)?; + let notification: SkillsChangedNotification = + timeout(timeout_duration, mcp.read_notification("skills/changed")).await??; assert_eq!(notification, SkillsChangedNotification {}); Ok(()) } @@ -190,21 +180,16 @@ enabled = true let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let disablement_request_id = mcp .send_experimental_feature_enablement_set_request(ExperimentalFeatureEnablementSetParams { enablement: BTreeMap::from([("remote_plugin".to_string(), false)]), }) .await?; - let disablement_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(disablement_request_id)), - ) - .await??; - let _: ExperimentalFeatureEnablementSetResponse = to_response(disablement_response)?; + let _: ExperimentalFeatureEnablementSetResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(disablement_request_id)).await??; let initial_skills_list_request_id = mcp .send_skills_list_request(SkillsListParams { @@ -212,12 +197,11 @@ enabled = true force_reload: true, }) .await?; - let initial_skills_list_response: JSONRPCResponse = timeout( + let SkillsListResponse { data } = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(initial_skills_list_request_id)), + mcp.read_response(initial_skills_list_request_id), ) .await??; - let SkillsListResponse { data } = to_response(initial_skills_list_response)?; assert!(data.iter().any(|entry| { entry .skills @@ -230,12 +214,8 @@ enabled = true enablement: BTreeMap::from([("remote_plugin".to_string(), true)]), }) .await?; - let enablement_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(enablement_request_id)), - ) - .await??; - let _: ExperimentalFeatureEnablementSetResponse = to_response(enablement_response)?; + let _: ExperimentalFeatureEnablementSetResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(enablement_request_id)).await??; let skills_list_request_id = mcp .send_skills_list_request(SkillsListParams { @@ -243,12 +223,8 @@ enabled = true force_reload: true, }) .await?; - let skills_list_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(skills_list_request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(skills_list_response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(skills_list_request_id)).await??; assert!(data.iter().all(|entry| { entry @@ -350,9 +326,8 @@ async fn skills_list_loads_remote_installed_plugin_skills_from_cache() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let stale_skills_list_request_id = mcp .send_skills_list_request(SkillsListParams { @@ -360,12 +335,11 @@ async fn skills_list_loads_remote_installed_plugin_skills_from_cache() -> Result force_reload: true, }) .await?; - let stale_skills_list_response: JSONRPCResponse = timeout( + let SkillsListResponse { data } = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(stale_skills_list_request_id)), + mcp.read_response(stale_skills_list_request_id), ) .await??; - let SkillsListResponse { data } = to_response(stale_skills_list_response)?; assert_eq!(data.len(), 1); assert!( data[0] @@ -397,12 +371,8 @@ async fn skills_list_loads_remote_installed_plugin_skills_from_cache() -> Result force_refetch: false, }) .await?; - let plugin_list_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(plugin_list_request_id)), - ) - .await??; - let _: PluginListResponse = to_response(plugin_list_response)?; + let _: PluginListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(plugin_list_request_id)).await??; let SkillsListResponse { data } = timeout(DEFAULT_TIMEOUT, async { loop { @@ -412,12 +382,8 @@ async fn skills_list_loads_remote_installed_plugin_skills_from_cache() -> Result force_reload: false, }) .await?; - let skills_list_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(skills_list_request_id)), - ) - .await??; - let response: SkillsListResponse = to_response(skills_list_response)?; + let response: SkillsListResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(skills_list_request_id)).await??; if response.data.iter().any(|entry| { entry .skills @@ -481,9 +447,8 @@ async fn skills_list_excludes_plugin_skills_when_workspace_codex_plugins_disable .with_codex_home(codex_home.path()) .without_auto_env() .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_skills_list_request(SkillsListParams { @@ -492,12 +457,8 @@ async fn skills_list_excludes_plugin_skills_when_workspace_codex_plugins_disable }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data.len(), 1); assert!( data[0] @@ -532,9 +493,8 @@ async fn skills_list_skips_cwd_roots_when_environment_disabled() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[(CODEX_EXEC_SERVER_URL_ENV_VAR, Some("none"))]) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_skills_list_request(SkillsListParams { @@ -543,12 +503,8 @@ async fn skills_list_skips_cwd_roots_when_environment_disabled() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].cwd, cwd.path().to_path_buf()); assert_eq!(data[0].errors, Vec::new()); @@ -576,9 +532,8 @@ async fn skills_list_accepts_relative_cwds() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_skills_list_request(SkillsListParams { @@ -587,12 +542,8 @@ async fn skills_list_accepts_relative_cwds() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].cwd, relative_cwd); assert_eq!(data[0].errors, Vec::new()); @@ -608,9 +559,8 @@ async fn skills_list_preserves_requested_cwd_order() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_skills_list_request(SkillsListParams { @@ -622,12 +572,8 @@ async fn skills_list_preserves_requested_cwd_order() -> Result<()> { }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( data.iter() .map(|entry| entry.cwd.clone()) @@ -648,9 +594,8 @@ async fn skills_list_uses_cached_result_until_force_reload() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; // Seed the cwd cache before the cwd-local skill exists. let first_request_id = mcp @@ -659,12 +604,8 @@ async fn skills_list_uses_cached_result_until_force_reload() -> Result<()> { force_reload: false, }) .await?; - let first_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_request_id)), - ) - .await??; - let SkillsListResponse { data: first_data } = to_response(first_response)?; + let SkillsListResponse { data: first_data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(first_request_id)).await??; assert_eq!(first_data.len(), 1); assert!( first_data[0] @@ -686,12 +627,8 @@ async fn skills_list_uses_cached_result_until_force_reload() -> Result<()> { force_reload: false, }) .await?; - let second_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_request_id)), - ) - .await??; - let SkillsListResponse { data: second_data } = to_response(second_response)?; + let SkillsListResponse { data: second_data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(second_request_id)).await??; assert_eq!(second_data.len(), 1); assert!( second_data[0] @@ -706,12 +643,8 @@ async fn skills_list_uses_cached_result_until_force_reload() -> Result<()> { force_reload: true, }) .await?; - let third_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(third_request_id)), - ) - .await??; - let SkillsListResponse { data: third_data } = to_response(third_response)?; + let SkillsListResponse { data: third_data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(third_request_id)).await??; assert_eq!(third_data.len(), 1); assert!( third_data[0] @@ -738,21 +671,16 @@ async fn skills_extra_roots_set_updates_process_runtime_roots() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let set_request_id = mcp .send_skills_extra_roots_set_request(SkillsExtraRootsSetParams { extra_roots: vec![AbsolutePathBuf::from_absolute_path(&extra_skills_root)?], }) .await?; - let set_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_request_id)), - ) - .await??; - let _: SkillsExtraRootsSetResponse = to_response(set_response)?; + let _: SkillsExtraRootsSetResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(set_request_id)).await??; expect_skills_changed_notification(&mut mcp, DEFAULT_TIMEOUT).await?; let skills_request_id = mcp @@ -761,12 +689,8 @@ async fn skills_extra_roots_set_updates_process_runtime_roots() -> Result<()> { force_reload: false, }) .await?; - let skills_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(skills_request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(skills_response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(skills_request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].errors, Vec::new()); assert!( @@ -782,12 +706,8 @@ async fn skills_extra_roots_set_updates_process_runtime_roots() -> Result<()> { extra_roots: vec![AbsolutePathBuf::from_absolute_path(&missing_root)?], }) .await?; - let reset_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(reset_request_id)), - ) - .await??; - let _: SkillsExtraRootsSetResponse = to_response(reset_response)?; + let _: SkillsExtraRootsSetResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(reset_request_id)).await??; expect_skills_changed_notification(&mut mcp, DEFAULT_TIMEOUT).await?; let skills_request_id = mcp @@ -796,12 +716,8 @@ async fn skills_extra_roots_set_updates_process_runtime_roots() -> Result<()> { force_reload: false, }) .await?; - let skills_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(skills_request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(skills_response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(skills_request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].errors, Vec::new()); assert!( @@ -816,12 +732,8 @@ async fn skills_extra_roots_set_updates_process_runtime_roots() -> Result<()> { extra_roots: Vec::new(), }) .await?; - let clear_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(clear_request_id)), - ) - .await??; - let _: SkillsExtraRootsSetResponse = to_response(clear_response)?; + let _: SkillsExtraRootsSetResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(clear_request_id)).await??; expect_skills_changed_notification(&mut mcp, DEFAULT_TIMEOUT).await?; let skills_request_id = mcp .send_skills_list_request(SkillsListParams { @@ -829,12 +741,8 @@ async fn skills_extra_roots_set_updates_process_runtime_roots() -> Result<()> { force_reload: false, }) .await?; - let skills_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(skills_request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(skills_response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(skills_request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].errors, Vec::new()); assert!( @@ -848,21 +756,16 @@ async fn skills_extra_roots_set_updates_process_runtime_roots() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let skills_request_id = mcp .send_skills_list_request(SkillsListParams { cwds: vec![cwd.path().to_path_buf()], force_reload: false, }) .await?; - let skills_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(skills_request_id)), - ) - .await??; - let SkillsListResponse { data } = to_response(skills_response)?; + let SkillsListResponse { data } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(skills_request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].errors, Vec::new()); assert!( @@ -884,30 +787,26 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - write_mock_responses_config_toml_with_chatgpt_base_url( - codex_home.path(), - &server.uri(), - &server.uri(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_root_config(&format!("chatgpt_base_url = \"{}\"", server.uri())) + .write(codex_home.path())?; write_skill(&codex_home, "demo")?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let initial_skills_request_id = mcp .send_skills_list_request(SkillsListParams { cwds: vec![codex_home.path().to_path_buf()], force_reload: true, }) .await?; - let initial_skills_response: JSONRPCResponse = timeout( + let SkillsListResponse { data } = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(initial_skills_request_id)), + mcp.read_response(initial_skills_request_id), ) .await??; - let SkillsListResponse { data } = to_response(initial_skills_response)?; assert_eq!(data.len(), 1); assert!( data[0] @@ -945,11 +844,8 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( experimental_raw_events: false, }) .await?; - let _: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_request_id)), - ) - .await??; + let _: ThreadStartResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(thread_start_request_id)).await??; let skill_path = codex_home .path() @@ -961,29 +857,18 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( "---\nname: demo\ndescription: updated\n---\n\n# Updated\n", )?; - let notification = timeout( - WATCHER_TIMEOUT, - mcp.read_stream_until_notification_message("skills/changed"), - ) - .await??; - let params = notification - .params - .context("skills/changed params must be present")?; - let notification: SkillsChangedNotification = serde_json::from_value(params)?; - - assert_eq!(notification, SkillsChangedNotification {}); + expect_skills_changed_notification(&mut mcp, WATCHER_TIMEOUT).await?; let updated_skills_request_id = mcp .send_skills_list_request(SkillsListParams { cwds: vec![codex_home.path().to_path_buf()], force_reload: false, }) .await?; - let updated_skills_response: JSONRPCResponse = timeout( + let SkillsListResponse { data } = timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(updated_skills_request_id)), + mcp.read_response(updated_skills_request_id), ) .await??; - let SkillsListResponse { data } = to_response(updated_skills_response)?; assert_eq!(data.len(), 1); assert!( data[0] diff --git a/codex-rs/app-server/tests/suite/v2/sleep.rs b/codex-rs/app-server/tests/suite/v2/sleep.rs index bdbd29f63ac4..a88612564569 100644 --- a/codex-rs/app-server/tests/suite/v2/sleep.rs +++ b/codex-rs/app-server/tests/suite/v2/sleep.rs @@ -1,11 +1,10 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::CurrentTimeReadResponse; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::SleepItem; use codex_app_server_protocol::ThreadItem; @@ -16,7 +15,6 @@ use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; use core_test_support::responses; use pretty_assertions::assert_eq; -use std::path::Path; use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; @@ -55,44 +53,42 @@ async fn external_sleep_polls_current_time_and_emits_items() -> Result<()> { .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_extra_config( + r#"[features.current_time_reminder] +enabled = true +sleep_tool = true +clock_source = "external" +"#, + ) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Sleep briefly".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let TurnStartResponse { turn, .. } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Sleep briefly".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let TurnStartResponse { turn, .. } = to_response(turn_start_response)?; // Read once for the initial reminder, then once to establish the sleep deadline. respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?; @@ -145,13 +141,8 @@ async fn wait_for_sleep_started( call_id: &str, ) -> Result { loop { - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/started"), - ) - .await??; let started: ItemStartedNotification = - serde_json::from_value(notification.params.expect("item/started params"))?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("item/started")).await??; if matches!(&started.item, ThreadItem::Sleep(item) if item.id == call_id) { return Ok(started); } @@ -163,13 +154,11 @@ async fn wait_for_sleep_completed( call_id: &str, ) -> Result { loop { - let notification = timeout( + let completed: ItemCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/completed"), + mcp.read_notification("item/completed"), ) .await??; - let completed: ItemCompletedNotification = - serde_json::from_value(notification.params.expect("item/completed params"))?; if matches!(&completed.item, ThreadItem::Sleep(item) if item.id == call_id) { return Ok(completed); } @@ -197,29 +186,3 @@ async fn respond_to_current_time_read( .await?; Ok(()) } - -fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - -[features.current_time_reminder] -enabled = true -sleep_tool = true -clock_source = "external" -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/thread_archive.rs b/codex-rs/app-server/tests/suite/v2/thread_archive.rs index 31e2c9067ddc..51c1a993779c 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_archive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_archive.rs @@ -1,10 +1,10 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; 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::ClientRequest; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadArchiveParams; use codex_app_server_protocol::ThreadArchiveResponse; @@ -36,27 +36,20 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs async fn thread_archive_requires_materialized_rollout() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // Start a thread. - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(start_resp)?; assert!(!thread.id.is_empty()); let rollout_path = thread.path.clone().expect("thread path"); @@ -93,23 +86,20 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> { ); // Materialize rollout via a real user turn and confirm archive succeeds. - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![UserInput::Text { - text: "materialize".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "materialize".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_start_response)?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -128,27 +118,19 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> { .expect("expected rollout path for thread id to exist after materialization"); assert_paths_match_on_disk(&discovered_path, &rollout_path)?; - let archive_id = mcp - .send_thread_archive_request(ThreadArchiveParams { - thread_id: thread.id.clone(), + let _: ThreadArchiveResponse = mcp + .request(|request_id| ClientRequest::ThreadArchive { + request_id, + params: ThreadArchiveParams { + thread_id: thread.id.clone(), + }, }) .await?; - let archive_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(archive_id)), - ) - .await??; - let _: ThreadArchiveResponse = to_response::(archive_resp)?; - let archive_notification = timeout( + let archived_notification: ThreadArchivedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/archived"), + mcp.read_notification("thread/archived"), ) .await??; - let archived_notification: ThreadArchivedNotification = serde_json::from_value( - archive_notification - .params - .expect("thread/archived notification params"), - )?; assert_eq!(archived_notification.thread_id, thread.id); // Verify file moved. @@ -174,7 +156,7 @@ async fn thread_archive_requires_materialized_rollout() -> Result<()> { async fn thread_archive_archives_spawned_descendants() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let parent_id = create_fake_rollout( codex_home.path(), @@ -227,34 +209,25 @@ async fn thread_archive_archives_spawned_descendants() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let archive_id = mcp - .send_thread_archive_request(ThreadArchiveParams { - thread_id: parent_id.clone(), + let _: ThreadArchiveResponse = mcp + .request(|request_id| ClientRequest::ThreadArchive { + request_id, + params: ThreadArchiveParams { + thread_id: parent_id.clone(), + }, }) .await?; - let archive_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(archive_id)), - ) - .await??; - let _: ThreadArchiveResponse = to_response::(archive_resp)?; let mut archived_ids = Vec::new(); for _ in 0..3 { - let notification = timeout( + let archived_notification: ThreadArchivedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/archived"), + mcp.read_notification("thread/archived"), ) .await??; - let archived_notification: ThreadArchivedNotification = serde_json::from_value( - notification - .params - .expect("thread/archived notification params"), - )?; archived_ids.push(archived_notification.thread_id); } assert_eq!(archived_ids, vec![parent_id, grandchild_id, child_id]); @@ -289,7 +262,7 @@ async fn thread_archive_archives_spawned_descendants() -> Result<()> { async fn thread_archive_succeeds_when_descendant_archive_fails() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let parent_id = create_fake_rollout( codex_home.path(), @@ -352,34 +325,25 @@ async fn thread_archive_succeeds_when_descendant_archive_fails() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let archive_id = mcp - .send_thread_archive_request(ThreadArchiveParams { - thread_id: parent_id.clone(), + let _: ThreadArchiveResponse = mcp + .request(|request_id| ClientRequest::ThreadArchive { + request_id, + params: ThreadArchiveParams { + thread_id: parent_id.clone(), + }, }) .await?; - let archive_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(archive_id)), - ) - .await??; - let _: ThreadArchiveResponse = to_response::(archive_resp)?; let mut archived_ids = Vec::new(); for _ in 0..2 { - let notification = timeout( + let archived_notification: ThreadArchivedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/archived"), + mcp.read_notification("thread/archived"), ) .await??; - let archived_notification: ThreadArchivedNotification = serde_json::from_value( - notification - .params - .expect("thread/archived notification params"), - )?; archived_ids.push(archived_notification.thread_id); } assert_eq!(archived_ids, vec![parent_id, grandchild_id]); @@ -431,7 +395,7 @@ async fn thread_archive_succeeds_when_descendant_archive_fails() -> Result<()> { async fn thread_archive_succeeds_when_spawned_descendant_is_missing() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let parent_id = create_fake_rollout( codex_home.path(), @@ -460,32 +424,23 @@ async fn thread_archive_succeeds_when_spawned_descendant_is_missing() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let archive_id = mcp - .send_thread_archive_request(ThreadArchiveParams { - thread_id: parent_id.clone(), + let _: ThreadArchiveResponse = mcp + .request(|request_id| ClientRequest::ThreadArchive { + request_id, + params: ThreadArchiveParams { + thread_id: parent_id.clone(), + }, }) .await?; - let archive_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(archive_id)), - ) - .await??; - let _: ThreadArchiveResponse = to_response::(archive_resp)?; - let notification = timeout( + let archived_notification: ThreadArchivedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/archived"), + mcp.read_notification("thread/archived"), ) .await??; - let archived_notification: ThreadArchivedNotification = serde_json::from_value( - notification - .params - .expect("thread/archived notification params"), - )?; assert_eq!(archived_notification.thread_id, parent_id); assert!( @@ -512,44 +467,34 @@ async fn thread_archive_succeeds_when_spawned_descendant_is_missing() -> Result< async fn thread_archive_clears_stale_subscriptions_before_resume() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; - let start_id = primary - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = primary + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; - let turn_start_id = primary - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![UserInput::Text { - text: "materialize".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = primary + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "materialize".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_start_response)?; timeout( DEFAULT_READ_TIMEOUT, primary.read_stream_until_notification_message("turn/completed"), @@ -559,38 +504,31 @@ async fn thread_archive_clears_stale_subscriptions_before_resume() -> Result<()> let mut secondary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, secondary.initialize()).await??; - let archive_id = primary - .send_thread_archive_request(ThreadArchiveParams { - thread_id: thread.id.clone(), + let _: ThreadArchiveResponse = primary + .request(|request_id| ClientRequest::ThreadArchive { + request_id, + params: ThreadArchiveParams { + thread_id: thread.id.clone(), + }, }) .await?; - let archive_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(archive_id)), - ) - .await??; - let _: ThreadArchiveResponse = to_response::(archive_resp)?; timeout( DEFAULT_READ_TIMEOUT, primary.read_stream_until_notification_message("thread/archived"), ) .await??; - let unarchive_id = primary - .send_thread_unarchive_request(ThreadUnarchiveParams { - thread_id: thread.id.clone(), + let _: ThreadUnarchiveResponse = primary + .request(|request_id| ClientRequest::ThreadUnarchive { + request_id, + params: ThreadUnarchiveParams { + thread_id: thread.id.clone(), + }, }) .await?; - let unarchive_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(unarchive_id)), - ) - .await??; - let _: ThreadUnarchiveResponse = to_response::(unarchive_resp)?; timeout( DEFAULT_READ_TIMEOUT, primary.read_stream_until_notification_message("thread/unarchived"), @@ -598,39 +536,33 @@ async fn thread_archive_clears_stale_subscriptions_before_resume() -> Result<()> .await??; primary.clear_message_buffer(); - let resume_id = secondary - .send_thread_resume_request(ThreadResumeParams { - thread_id: thread.id.clone(), - ..Default::default() + let resume: ThreadResumeResponse = secondary + .request(|request_id| ClientRequest::ThreadResume { + request_id, + params: ThreadResumeParams { + thread_id: thread.id.clone(), + ..Default::default() + }, }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - secondary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let resume: ThreadResumeResponse = to_response::(resume_resp)?; assert_eq!(resume.thread.status, ThreadStatus::Idle); primary.clear_message_buffer(); secondary.clear_message_buffer(); - let resumed_turn_id = secondary - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![UserInput::Text { - text: "secondary turn".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = secondary + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![UserInput::Text { + text: "secondary turn".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let resumed_turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - secondary.read_stream_until_response_message(RequestId::Integer(resumed_turn_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(resumed_turn_resp)?; assert!( timeout( @@ -650,29 +582,6 @@ async fn thread_archive_clears_stale_subscriptions_before_resume() -> Result<()> Ok(()) } -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, config_contents(server_uri)) -} - -fn config_contents(server_uri: &str) -> String { - format!( - r#"model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ) -} - fn assert_paths_match_on_disk(actual: &Path, expected: &Path) -> std::io::Result<()> { let actual = actual.canonicalize()?; let expected = expected.canonicalize()?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_delete.rs b/codex-rs/app-server/tests/suite/v2/thread_delete.rs index 96cf2add636a..c3aa91cb2573 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_delete.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_delete.rs @@ -1,9 +1,8 @@ use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::create_fake_rollout; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadDeleteParams; use codex_app_server_protocol::ThreadDeleteResponse; @@ -59,34 +58,25 @@ async fn thread_delete_deletes_spawned_descendants() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let delete_id = mcp - .send_thread_delete_request(ThreadDeleteParams { - thread_id: parent_id.clone(), + let _: ThreadDeleteResponse = mcp + .request(|request_id| ClientRequest::ThreadDelete { + request_id, + params: ThreadDeleteParams { + thread_id: parent_id.clone(), + }, }) .await?; - let delete_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(delete_id)), - ) - .await??; - let _: ThreadDeleteResponse = to_response::(delete_resp)?; let mut deleted_ids = Vec::new(); for _ in 0..3 { - let notification = timeout( + let deleted_notification: ThreadDeletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/deleted"), + mcp.read_notification("thread/deleted"), ) .await??; - let deleted_notification: ThreadDeletedNotification = serde_json::from_value( - notification - .params - .expect("thread/deleted notification params"), - )?; deleted_ids.push(deleted_notification.thread_id); } assert_eq!(deleted_ids, vec![grandchild_id, child_id, parent_id]); @@ -162,9 +152,8 @@ async fn thread_delete_preflights_external_fork_references_for_spawned_subtrees( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let delete_id = mcp .send_thread_delete_request(ThreadDeleteParams { @@ -219,19 +208,10 @@ async fn thread_delete_handles_live_threads_before_rollout_exists() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) - .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let persisted_thread = to_response::(start_resp)?.thread; + let persisted_thread = mcp.start_thread(ThreadStartParams::default()).await?.thread; let rollout_path = find_thread_path_by_id_str( codex_home.path(), &persisted_thread.id, @@ -240,30 +220,21 @@ async fn thread_delete_handles_live_threads_before_rollout_exists() -> Result<() .await?; assert_eq!(rollout_path, None); - let delete_id = mcp - .send_thread_delete_request(ThreadDeleteParams { - thread_id: persisted_thread.id, + let _: ThreadDeleteResponse = mcp + .request(|request_id| ClientRequest::ThreadDelete { + request_id, + params: ThreadDeleteParams { + thread_id: persisted_thread.id, + }, }) .await?; - let delete_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(delete_id)), - ) - .await??; - let _: ThreadDeleteResponse = to_response::(delete_resp)?; - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { ephemeral: Some(true), ..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::(start_resp)?; let delete_id = mcp .send_thread_delete_request(ThreadDeleteParams { @@ -281,16 +252,12 @@ async fn thread_delete_handles_live_threads_before_rollout_exists() -> Result<() ); assert_eq!(delete_err.error.message, expected_message); - let list_id = mcp - .send_thread_loaded_list_request(ThreadLoadedListParams::default()) + let ThreadLoadedListResponse { mut data, .. } = mcp + .request(|request_id| ClientRequest::ThreadLoadedList { + request_id, + params: ThreadLoadedListParams::default(), + }) .await?; - let list_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; - let ThreadLoadedListResponse { mut data, .. } = - to_response::(list_resp)?; data.sort(); assert_eq!(data, vec![thread.id]); diff --git a/codex-rs/app-server/tests/suite/v2/thread_fork.rs b/codex-rs/app-server/tests/suite/v2/thread_fork.rs index 948937e793cd..3df18fa86587 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_fork.rs @@ -1,5 +1,6 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_fake_rollout; use app_test_support::create_fake_rollout_with_token_usage; @@ -49,7 +50,6 @@ use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; -use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; use wiremock::Mock; @@ -97,7 +97,7 @@ async fn list_threads(mcp: &mut TestAppServer) -> Result { async fn thread_fork_creates_new_thread_and_emits_started() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let preview = "Saved user message"; let conversation_id = create_fake_rollout( @@ -130,9 +130,8 @@ async fn thread_fork_creates_new_thread_and_emits_started() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -268,14 +267,13 @@ async fn thread_fork_creates_new_thread_and_emits_started() -> Result<()> { async fn thread_fork_preserves_persisted_approvals_reviewer() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let (source_thread_id, source_turn_id) = { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) @@ -336,9 +334,8 @@ async fn thread_fork_preserves_persisted_approvals_reviewer() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { thread_id: source_thread_id, @@ -364,26 +361,20 @@ async fn thread_fork_preserves_persisted_approvals_reviewer() -> Result<()> { async fn thread_fork_at_last_turn_id_keeps_only_terminal_prefix() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; let ThreadStartResponse { thread: source_thread, .. - } = to_response::(start_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let source_thread_id = source_thread.id.clone(); let source_path = source_thread.path.expect("source thread path"); @@ -400,12 +391,8 @@ async fn thread_fork_at_last_turn_id_keeps_only_terminal_prefix() -> Result<()> ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_request_id)).await??; turn_ids.push(turn.id); timeout( DEFAULT_READ_TIMEOUT, @@ -422,15 +409,10 @@ async fn thread_fork_at_last_turn_id_keeps_only_terminal_prefix() -> Result<()> ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; let ThreadForkResponse { thread: forked_thread, .. - } = to_response::(fork_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; assert_eq!( forked_thread @@ -498,7 +480,7 @@ async fn thread_fork_defers_inherited_active_goal_until_next_turn() -> Result<() ]) .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -509,22 +491,16 @@ async fn thread_fork_defers_inherited_active_goal_until_next_turn() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) .await?; - let start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; let ThreadStartResponse { thread: source_thread, .. - } = to_response::(start_response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let source_thread_id = ThreadId::from_string(&source_thread.id)?; let mut turn_ids = Vec::new(); @@ -578,15 +554,10 @@ async fn thread_fork_defers_inherited_active_goal_until_next_turn() -> Result<() ..Default::default() }) .await?; - let ordinary_fork_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(ordinary_fork_id)), - ) - .await??; let ThreadForkResponse { thread: ordinary_fork, .. - } = to_response::(ordinary_fork_response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(ordinary_fork_id)).await??; assert_eq!( state_db .thread_goals() @@ -610,15 +581,10 @@ async fn thread_fork_defers_inherited_active_goal_until_next_turn() -> Result<() ..Default::default() }) .await?; - let fork_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; let ThreadForkResponse { thread: forked_thread, .. - } = to_response::(fork_response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; let forked_thread_id = ThreadId::from_string(&forked_thread.id)?; assert_eq!(forked_thread.turns.len(), expected_turn_count); let mut expected_goal = source_goal.clone(); @@ -670,9 +636,8 @@ async fn thread_fork_defers_inherited_active_goal_until_next_turn() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id: forked_thread.id.clone(), @@ -750,7 +715,7 @@ async fn thread_fork_defers_inherited_active_goal_until_next_turn() -> Result<() async fn thread_fork_inherits_explicit_source_name_from_session_index() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout( codex_home.path(), @@ -767,9 +732,8 @@ async fn thread_fork_inherits_explicit_source_name_from_session_index() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -777,12 +741,8 @@ async fn thread_fork_inherits_explicit_source_name_from_session_index() -> Resul ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; let ThreadListResponse { data, .. } = list_threads(&mut mcp).await?; let listed = data @@ -798,7 +758,7 @@ async fn thread_fork_inherits_explicit_source_name_from_session_index() -> Resul async fn thread_fork_can_load_source_by_path() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let preview = "Saved user message"; let conversation_id = create_fake_rollout( @@ -822,9 +782,8 @@ async fn thread_fork_can_load_source_by_path() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -833,12 +792,8 @@ async fn thread_fork_can_load_source_by_path() -> Result<()> { ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; assert_ne!(thread.id, conversation_id); assert_eq!(thread.forked_from_id, Some(conversation_id)); @@ -853,7 +808,7 @@ async fn thread_fork_can_load_source_by_path() -> Result<()> { async fn thread_fork_can_cut_before_unfinished_stored_turn() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let conversation_id = create_fake_rollout( @@ -889,9 +844,8 @@ async fn thread_fork_can_cut_before_unfinished_stored_turn() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_read_request(ThreadReadParams { @@ -899,14 +853,9 @@ async fn thread_fork_can_cut_before_unfinished_stored_turn() -> Result<()> { include_turns: true, }) .await?; - let read_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; let ThreadReadResponse { thread: source_thread, - } = to_response::(read_response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(source_thread.turns.len(), 2); assert_eq!(source_thread.turns[1].id, unfinished_turn_id); assert_eq!(source_thread.turns[1].status, TurnStatus::Interrupted); @@ -918,15 +867,10 @@ async fn thread_fork_can_cut_before_unfinished_stored_turn() -> Result<()> { ..Default::default() }) .await?; - let fork_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; let ThreadForkResponse { thread: forked_thread, .. - } = to_response::(fork_response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; assert_eq!(forked_thread.turns.len(), 1); assert_eq!(forked_thread.preview, "Saved user message"); @@ -937,7 +881,7 @@ async fn thread_fork_can_cut_before_unfinished_stored_turn() -> Result<()> { async fn thread_fork_emits_restored_token_usage_before_next_turn() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout_with_token_usage( codex_home.path(), @@ -950,9 +894,8 @@ async fn thread_fork_emits_restored_token_usage_before_next_turn() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -961,12 +904,8 @@ async fn thread_fork_emits_restored_token_usage_before_next_turn() -> Result<()> ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; let note = timeout( DEFAULT_READ_TIMEOUT, @@ -995,7 +934,7 @@ async fn thread_fork_emits_restored_token_usage_before_next_turn() -> Result<()> async fn thread_fork_can_exclude_turns_and_skip_restored_token_usage() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout_with_token_usage( codex_home.path(), @@ -1008,9 +947,8 @@ async fn thread_fork_can_exclude_turns_and_skip_restored_token_usage() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1019,12 +957,8 @@ async fn thread_fork_can_exclude_turns_and_skip_restored_token_usage() -> Result ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; assert_eq!(thread.forked_from_id, Some(conversation_id)); assert_eq!(thread.preview, "Saved user message"); @@ -1048,7 +982,9 @@ async fn thread_fork_tracks_thread_initialized_analytics() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml_with_chatgpt_base_url(codex_home.path(), &server.uri(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_root_config(&format!(r#"chatgpt_base_url = "{}""#, server.uri())) + .write(codex_home.path())?; mount_analytics_capture(&server, codex_home.path()).await?; let conversation_id = create_fake_rollout( @@ -1064,9 +1000,8 @@ async fn thread_fork_tracks_thread_initialized_analytics() -> Result<()> { .with_codex_home(codex_home.path()) .without_auto_env() .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1075,12 +1010,8 @@ async fn thread_fork_tracks_thread_initialized_analytics() -> Result<()> { ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; let payload = wait_for_analytics_payload(&server, DEFAULT_READ_TIMEOUT).await?; let event = thread_initialized_event(&payload)?; @@ -1107,13 +1038,12 @@ async fn thread_fork_tracks_thread_initialized_analytics() -> Result<()> { async fn thread_fork_rejects_unmaterialized_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1121,12 +1051,8 @@ async fn thread_fork_rejects_unmaterialized_thread() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1155,7 +1081,7 @@ async fn thread_fork_rejects_unmaterialized_thread() -> Result<()> { async fn thread_fork_rejects_paginated_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout( codex_home.path(), @@ -1189,9 +1115,8 @@ async fn thread_fork_rejects_paginated_thread() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1217,7 +1142,7 @@ async fn thread_fork_rejects_paginated_thread() -> Result<()> { async fn thread_fork_with_empty_path_uses_thread_id() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout( codex_home.path(), @@ -1231,9 +1156,8 @@ async fn thread_fork_with_empty_path_uses_thread_id() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1243,12 +1167,8 @@ async fn thread_fork_with_empty_path_uses_thread_id() -> Result<()> { ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; assert_eq!( thread.forked_from_id.as_deref(), @@ -1280,11 +1200,9 @@ async fn thread_fork_surfaces_cloud_config_bundle_load_errors() -> Result<()> { let codex_home = TempDir::new()?; let model_server = create_mock_responses_server_repeating_assistant("Done").await; let chatgpt_base_url = format!("{}/backend-api", server.uri()); - create_config_toml_with_chatgpt_base_url( - codex_home.path(), - &model_server.uri(), - &chatgpt_base_url, - )?; + MockResponsesConfig::new(&model_server.uri()) + .with_root_config(&format!(r#"chatgpt_base_url = "{chatgpt_base_url}""#)) + .write(codex_home.path())?; write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("chatgpt-token") @@ -1316,9 +1234,8 @@ async fn thread_fork_surfaces_cloud_config_bundle_load_errors() -> Result<()> { Some(refresh_token_url.as_str()), ), ]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1358,7 +1275,7 @@ async fn thread_fork_surfaces_cloud_config_bundle_load_errors() -> Result<()> { async fn thread_fork_ephemeral_remains_pathless_and_omits_listing() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let preview = "Saved user message"; let conversation_id = create_fake_rollout( @@ -1373,9 +1290,8 @@ async fn thread_fork_ephemeral_remains_pathless_and_omits_listing() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1497,12 +1413,7 @@ async fn thread_fork_ephemeral_remains_pathless_and_omits_listing() -> Result<() ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_resp)?; + let _: TurnStartResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -1516,7 +1427,7 @@ async fn thread_fork_ephemeral_remains_pathless_and_omits_listing() -> Result<() async fn thread_fork_rejects_incompatible_boundaries_and_ephemeral_goal_deferral() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let thread_id = create_fake_rollout( codex_home.path(), "2025-01-05T12-00-00", @@ -1528,9 +1439,8 @@ async fn thread_fork_rejects_incompatible_boundaries_and_ephemeral_goal_deferral let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; for (params, expected_message) in [ ( @@ -1568,7 +1478,7 @@ async fn thread_fork_rejects_incompatible_boundaries_and_ephemeral_goal_deferral async fn pathless_ephemeral_thread_rejects_codex_home_path_after_reload() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let parent_thread_id = create_fake_rollout( codex_home.path(), @@ -1583,9 +1493,8 @@ async fn pathless_ephemeral_thread_rejects_codex_home_path_after_reload() -> Res let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; let fork_id = app_server .send_thread_fork_request(ThreadForkParams { @@ -1594,12 +1503,8 @@ async fn pathless_ephemeral_thread_rejects_codex_home_path_after_reload() -> Res ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(fork_id)).await??; assert!(thread.ephemeral); assert_eq!(thread.path, None); @@ -1614,12 +1519,8 @@ async fn pathless_ephemeral_thread_rejects_codex_home_path_after_reload() -> Res ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_resp)?; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, app_server.read_response(turn_id)).await??; timeout( DEFAULT_READ_TIMEOUT, app_server.read_stream_until_notification_message("turn/completed"), @@ -1632,9 +1533,8 @@ async fn pathless_ephemeral_thread_rejects_codex_home_path_after_reload() -> Res let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; let codex_home_path = codex_home.path().to_path_buf(); let resume_id = app_server @@ -1685,55 +1585,3 @@ async fn pathless_ephemeral_thread_rejects_codex_home_path_after_reload() -> Res Ok(()) } - -// Helper to create a config.toml pointing at the mock model server. -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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - -fn create_config_toml_with_chatgpt_base_url( - codex_home: &Path, - server_uri: &str, - chatgpt_base_url: &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" -chatgpt_base_url = "{chatgpt_base_url}" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/thread_inject_items.rs b/codex-rs/app-server/tests/suite/v2/thread_inject_items.rs index 9ec91072dce1..13ed15ed1932 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_inject_items.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_inject_items.rs @@ -1,14 +1,13 @@ use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadInjectItemsParams; use codex_app_server_protocol::ThreadInjectItemsResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; use codex_core::RolloutRecorder; use codex_protocol::models::ContentItem; @@ -19,7 +18,6 @@ use core_test_support::responses; use core_test_support::responses::strip_response_item_id; use core_test_support::responses::strip_response_item_ids_from_json; use serde_json::Value; -use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; @@ -36,13 +34,12 @@ async fn thread_inject_items_adds_raw_response_items_to_thread_history() -> Resu let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -50,12 +47,8 @@ async fn thread_inject_items_adds_raw_response_items_to_thread_history() -> Resu ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let injected_text = "Injected assistant context"; let injected_item = ResponseItem::Message { @@ -74,13 +67,8 @@ async fn thread_inject_items_adds_raw_response_items_to_thread_history() -> Resu items: vec![serde_json::to_value(&injected_item)?], }) .await?; - let inject_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(inject_req)), - ) - .await??; let _response: ThreadInjectItemsResponse = - to_response::(inject_resp)?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(inject_req)).await??; let rollout_path = thread.path.as_ref().context("thread path missing")?; let history = RolloutRecorder::get_rollout_history(rollout_path).await?; @@ -106,11 +94,7 @@ async fn thread_inject_items_adds_raw_response_items_to_thread_history() -> Resu ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -161,13 +145,12 @@ async fn thread_inject_items_adds_raw_response_items_after_a_turn() -> Result<() let response_mock = responses::mount_sse_sequence(&server, vec![first_body, second_body]).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -175,12 +158,8 @@ async fn thread_inject_items_adds_raw_response_items_after_a_turn() -> Result<() ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let first_turn_req = mcp .send_turn_start_request(TurnStartParams { @@ -193,11 +172,8 @@ async fn thread_inject_items_adds_raw_response_items_after_a_turn() -> Result<() ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_turn_req)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(first_turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -221,13 +197,8 @@ async fn thread_inject_items_adds_raw_response_items_after_a_turn() -> Result<() items: vec![injected_value.clone()], }) .await?; - let inject_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(inject_req)), - ) - .await??; let _response: ThreadInjectItemsResponse = - to_response::(inject_resp)?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(inject_req)).await??; let second_turn_req = mcp .send_turn_start_request(TurnStartParams { @@ -240,11 +211,8 @@ async fn thread_inject_items_adds_raw_response_items_after_a_turn() -> Result<() ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_turn_req)), - ) - .await??; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(second_turn_req)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -273,29 +241,6 @@ async fn thread_inject_items_adds_raw_response_items_after_a_turn() -> Result<() Ok(()) } -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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - fn response_item_text_position(items: &[Value], needle: &str) -> Option { items.iter().position(|item| { item.get("content") diff --git a/codex-rs/app-server/tests/suite/v2/thread_list.rs b/codex-rs/app-server/tests/suite/v2/thread_list.rs index bfc076fe526d..49e1f5558546 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_list.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_list.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_fake_parented_rollout_with_source; use app_test_support::create_fake_rollout; @@ -7,12 +8,11 @@ use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; use app_test_support::rollout_path; use app_test_support::test_absolute_path; -use app_test_support::to_response; use chrono::DateTime; use chrono::Utc; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::GitInfo as ApiGitInfo; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SessionSource; use codex_app_server_protocol::SortDirection; @@ -50,12 +50,10 @@ use uuid::Uuid; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); async fn init_mcp(codex_home: &Path) -> Result { - let mut mcp = TestAppServer::builder() + TestAppServer::builder() .with_codex_home(codex_home) - .build() - .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - Ok(mcp) + .build_initialized() + .await } async fn list_threads( @@ -87,8 +85,9 @@ async fn list_threads_with_sort( sort_key: Option, archived: Option, ) -> Result { - let request_id = mcp - .send_thread_list_request(codex_app_server_protocol::ThreadListParams { + mcp.request(|request_id| ClientRequest::ThreadList { + request_id, + params: codex_app_server_protocol::ThreadListParams { cursor, limit, sort_key, @@ -101,14 +100,9 @@ async fn list_threads_with_sort( search_term: None, parent_thread_id: None, ancestor_thread_id: None, - }) - .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response::(resp) + }, + }) + .await } enum ThreadListRelation { @@ -128,8 +122,9 @@ async fn list_threads_for_relation( ThreadListRelation::DirectChildrenOf(thread_id) => (Some(thread_id.to_string()), None), ThreadListRelation::DescendantsOf(thread_id) => (None, Some(thread_id.to_string())), }; - let request_id = mcp - .send_thread_list_request(codex_app_server_protocol::ThreadListParams { + mcp.request(|request_id| ClientRequest::ThreadList { + request_id, + params: codex_app_server_protocol::ThreadListParams { cursor, limit: Some(limit), sort_key: None, @@ -142,14 +137,9 @@ async fn list_threads_for_relation( search_term: None, parent_thread_id, ancestor_thread_id, - }) - .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response::(response) + }, + }) + .await } fn create_fake_rollouts( @@ -259,18 +249,12 @@ async fn thread_list_reports_system_error_idle_flag_after_failed_turn() -> Resul create_runtime_config(codex_home.path(), &server.uri())?; let mut mcp = init_mcp(codex_home.path()).await?; - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(start_resp)?; let seed_turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -283,12 +267,8 @@ async fn thread_list_reports_system_error_idle_flag_after_failed_turn() -> Resul ..Default::default() }) .await?; - let seed_turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(seed_turn_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(seed_turn_resp)?; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(seed_turn_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -306,12 +286,8 @@ async fn thread_list_reports_system_error_idle_flag_after_failed_turn() -> Resul ..Default::default() }) .await?; - let failed_turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(failed_turn_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(failed_turn_resp)?; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(failed_turn_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("error"), @@ -353,26 +329,7 @@ approval_policy = "never" } fn create_runtime_config(codex_home: &std::path::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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) + MockResponsesConfig::new(server_uri).write(codex_home) } #[tokio::test] @@ -588,14 +545,9 @@ async fn thread_list_respects_cwd_filters() -> Result<()> { ancestor_thread_id: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ThreadListResponse { data, next_cursor, .. - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(next_cursor, None); let filtered_ids: Vec<_> = data.iter().map(|thread| thread.id.as_str()).collect(); @@ -699,14 +651,9 @@ sqlite = true ancestor_thread_id: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ThreadListResponse { data, next_cursor, .. - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(next_cursor, None); let ids: Vec<_> = data.iter().map(|thread| thread.id.as_str()).collect(); @@ -757,14 +704,9 @@ async fn thread_search_returns_content_matches() -> Result<()> { search_term: "needle".to_string(), }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ThreadSearchResponse { data, next_cursor, .. - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(next_cursor, None); let ids: Vec<_> = data @@ -804,12 +746,8 @@ async fn thread_search_matches_json_escaped_content() -> Result<()> { search_term: search_term.to_string(), }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadSearchResponse { data, .. } = to_response::(resp)?; + let ThreadSearchResponse { data, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(data.len(), 1); assert_eq!(data[0].thread.id, thread_id); @@ -853,12 +791,8 @@ async fn thread_search_filters_by_source_kind() -> Result<()> { search_term: "needle".to_string(), }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadSearchResponse { data, .. } = to_response::(resp)?; + let ThreadSearchResponse { data, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let ids: Vec<_> = data .iter() @@ -917,12 +851,8 @@ sqlite = true ancestor_thread_id: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let repaired_response = to_response::(resp)?; + let repaired_response: ThreadListResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let ids: Vec<_> = repaired_response .data .iter() @@ -957,12 +887,8 @@ sqlite = true ancestor_thread_id: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let state_db_only_response = to_response::(resp)?; + let state_db_only_response: ThreadListResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let ids: Vec<_> = state_db_only_response .data .iter() @@ -988,12 +914,8 @@ sqlite = true ancestor_thread_id: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let scanned_response = to_response::(resp)?; + let scanned_response: ThreadListResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!(scanned_response.data.len(), 0); Ok(()) @@ -1986,12 +1908,7 @@ async fn thread_list_backwards_cursor_can_seed_forward_delta_sync() -> Result<() ancestor_thread_id: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response::(resp)? + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await?? }; let ids_page1: Vec<_> = page1.iter().map(|thread| thread.id.as_str()).collect(); assert_eq!(ids_page1, vec![id_watermark.as_str()]); @@ -2030,12 +1947,7 @@ async fn thread_list_backwards_cursor_can_seed_forward_delta_sync() -> Result<() ancestor_thread_id: None, }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response::(resp)? + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await?? }; let ids_delta: Vec<_> = delta_page.iter().map(|thread| thread.id.as_str()).collect(); assert_eq!(ids_delta, vec![id_watermark.as_str(), id_new.as_str()]); diff --git a/codex-rs/app-server/tests/suite/v2/thread_loaded_list.rs b/codex-rs/app-server/tests/suite/v2/thread_loaded_list.rs index d083d8279f3c..2c074cf67368 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_loaded_list.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_loaded_list.rs @@ -1,15 +1,12 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; 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::ThreadLoadedListParams; use codex_app_server_protocol::ThreadLoadedListResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use pretty_assertions::assert_eq; -use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; @@ -19,28 +16,22 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs async fn thread_loaded_list_returns_loaded_thread_ids() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_id = start_thread(&mut mcp).await?; let list_id = mcp .send_thread_loaded_list_request(ThreadLoadedListParams::default()) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; let ThreadLoadedListResponse { mut data, next_cursor, - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(list_id)).await??; data.sort(); assert_eq!(data, vec![thread_id]); assert_eq!(next_cursor, None); @@ -52,13 +43,12 @@ async fn thread_loaded_list_returns_loaded_thread_ids() -> Result<()> { async fn thread_loaded_list_paginates() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let first = start_thread(&mut mcp).await?; let second = start_thread(&mut mcp).await?; @@ -72,15 +62,10 @@ async fn thread_loaded_list_paginates() -> Result<()> { limit: Some(1), }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; let ThreadLoadedListResponse { data: first_page, next_cursor, - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(list_id)).await??; assert_eq!(first_page, vec![expected[0].clone()]); assert_eq!(next_cursor, Some(expected[0].clone())); @@ -90,44 +75,16 @@ async fn thread_loaded_list_paginates() -> Result<()> { limit: Some(1), }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; let ThreadLoadedListResponse { data: second_page, next_cursor, - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(list_id)).await??; assert_eq!(second_page, vec![expected[1].clone()]); assert_eq!(next_cursor, None); Ok(()) } -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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - async fn start_thread(mcp: &mut TestAppServer) -> Result { let req_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -135,11 +92,7 @@ async fn start_thread(mcp: &mut TestAppServer) -> Result { ..Default::default() }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(req_id)).await??; Ok(thread.id) } diff --git a/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs b/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs index 98395414c762..c43caaff120d 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_memory_mode_set.rs @@ -1,64 +1,54 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; 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::ClientRequest; 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_features::Feature; 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())?; + MockResponsesConfig::new(&server.uri()) + .with_root_config("suppress_unstable_features_warning = true") + .enable_feature(Feature::Sqlite) + .write(codex_home.path())?; let state_db = init_state_db(codex_home.path()).await?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(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, + let _: ThreadMemoryModeSetResponse = mcp + .request(|request_id| ClientRequest::ThreadMemoryModeSet { + request_id, + params: 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::(set_resp)?; let memory_mode = state_db.get_thread_memory_mode(thread_uuid).await?; assert_eq!(memory_mode.as_deref(), Some("disabled")); @@ -69,7 +59,10 @@ async fn thread_memory_mode_set_updates_loaded_thread_state() -> Result<()> { 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())?; + MockResponsesConfig::new(&server.uri()) + .with_root_config("suppress_unstable_features_warning = true") + .enable_feature(Feature::Sqlite) + .write(codex_home.path())?; let state_db = init_state_db(codex_home.path()).await?; let thread_id = create_fake_rollout( @@ -85,23 +78,19 @@ async fn thread_memory_mode_set_updates_stored_thread_state() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .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, + let _: ThreadMemoryModeSetResponse = mcp + .request(|request_id| ClientRequest::ThreadMemoryModeSet { + request_id, + params: 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::(set_resp)?; } let memory_mode = state_db.get_thread_memory_mode(thread_uuid).await?; @@ -116,30 +105,3 @@ async fn init_state_db(codex_home: &Path) -> Result> { .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 = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs b/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs index fd3f6ab16b5c..49fcb7ff55d5 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_metadata_update.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_fake_rollout; use app_test_support::create_mock_responses_server_repeating_assistant; @@ -19,6 +20,7 @@ use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::ThreadStatus; use codex_core::ARCHIVED_SESSIONS_SUBDIR; +use codex_features::Feature; use codex_git_utils::GitSha; use codex_protocol::ThreadId; use codex_protocol::protocol::GitInfo as RolloutGitInfo; @@ -39,7 +41,7 @@ const INVALID_REQUEST_ERROR_CODE: i64 = -32600; async fn thread_metadata_update_patches_git_branch_and_returns_updated_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -137,7 +139,7 @@ async fn thread_metadata_update_patches_git_branch_and_returns_updated_thread() async fn thread_metadata_update_rejects_empty_git_info_patch() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -186,7 +188,7 @@ async fn thread_metadata_update_rejects_empty_git_info_patch() -> Result<()> { async fn thread_metadata_update_rejects_ephemeral_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -240,7 +242,7 @@ async fn thread_metadata_update_rejects_ephemeral_thread() -> Result<()> { async fn thread_metadata_update_repairs_missing_sqlite_row_for_stored_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let _state_db = init_state_db(codex_home.path()).await?; let preview = "Stored thread preview"; @@ -297,7 +299,7 @@ async fn thread_metadata_update_repairs_missing_sqlite_row_for_stored_thread() - async fn thread_metadata_update_repairs_loaded_thread_without_resetting_summary() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let state_db = init_state_db(codex_home.path()).await?; let preview = "Loaded thread preview"; @@ -381,7 +383,7 @@ async fn thread_metadata_update_repairs_loaded_thread_without_resetting_summary( async fn thread_metadata_update_repairs_missing_sqlite_row_for_archived_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let _state_db = init_state_db(codex_home.path()).await?; let preview = "Archived thread preview"; @@ -448,7 +450,7 @@ async fn thread_metadata_update_repairs_missing_sqlite_row_for_archived_thread() async fn thread_metadata_update_can_clear_stored_git_fields() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let thread_id = create_fake_rollout( codex_home.path(), @@ -518,29 +520,8 @@ async fn init_state_db(codex_home: &Path) -> Result> { 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 = 0 -stream_max_retries = 0 -"# - ), - ) +fn mock_responses_config(server_uri: &str) -> MockResponsesConfig { + MockResponsesConfig::new(server_uri) + .with_root_config("suppress_unstable_features_warning = true") + .enable_feature(Feature::Sqlite) } diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 1a44dcecbbab..9bd4d1be4b4c 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_fake_paginated_rollout; use app_test_support::create_fake_rollout_with_text_elements; @@ -99,7 +100,7 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs async fn thread_read_returns_summary_without_turns() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let preview = "Saved user message"; let text_elements = [TextElement::new( @@ -122,9 +123,8 @@ async fn thread_read_returns_summary_without_turns() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_read_request(ThreadReadParams { @@ -132,12 +132,8 @@ async fn thread_read_returns_summary_without_turns() -> Result<()> { include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(thread.id, conversation_id); assert_eq!(thread.preview, preview); @@ -158,7 +154,7 @@ async fn thread_read_returns_summary_without_turns() -> Result<()> { async fn thread_read_can_include_turns() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let preview = "Saved user message"; let text_elements = vec![TextElement::new( @@ -181,9 +177,8 @@ async fn thread_read_can_include_turns() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_read_request(ThreadReadParams { @@ -191,12 +186,8 @@ async fn thread_read_can_include_turns() -> Result<()> { include_turns: true, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(thread.turns.len(), 1); let turn = &thread.turns[0]; @@ -225,7 +216,7 @@ async fn paginated_stored_thread_routes_projected_turns_and_rejects_legacy_histo -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_paginated_rollout( codex_home.path(), @@ -239,9 +230,8 @@ async fn paginated_stored_thread_routes_projected_turns_and_rejects_legacy_histo let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_read_request(ThreadReadParams { @@ -249,12 +239,8 @@ async fn paginated_stored_thread_routes_projected_turns_and_rejects_legacy_histo include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread } = to_response::(read_resp)?; + let ThreadReadResponse { thread } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); assert!(thread.turns.is_empty()); @@ -274,12 +260,8 @@ async fn paginated_stored_thread_routes_projected_turns_and_rejects_legacy_histo ancestor_thread_id: None, }) .await?; - let list_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; - let ThreadListResponse { data, .. } = to_response::(list_resp)?; + let ThreadListResponse { data, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(list_id)).await??; let listed = data .iter() .find(|thread| thread.id == conversation_id) @@ -333,7 +315,7 @@ async fn paginated_stored_thread_routes_projected_turns_and_rejects_legacy_histo async fn thread_turns_list_can_page_backward_and_forward() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let conversation_id = create_fake_rollout_with_text_elements( @@ -352,9 +334,8 @@ async fn thread_turns_list_can_page_backward_and_forward() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_turns_list_request(ThreadTurnsListParams { @@ -365,16 +346,11 @@ async fn thread_turns_list_can_page_backward_and_forward() -> Result<()> { items_view: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; let ThreadTurnsListResponse { data, next_cursor, backwards_cursor, - } = to_response::(read_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(turn_user_texts(&data), vec!["third", "second"]); assert!( data.iter() @@ -392,12 +368,8 @@ async fn thread_turns_list_can_page_backward_and_forward() -> Result<()> { items_view: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadTurnsListResponse { data, .. } = to_response::(read_resp)?; + let ThreadTurnsListResponse { data, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(turn_user_texts(&data), vec!["first"]); append_user_message(rollout_path.as_path(), "2025-01-05T12:03:00Z", "fourth")?; @@ -411,12 +383,8 @@ async fn thread_turns_list_can_page_backward_and_forward() -> Result<()> { items_view: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadTurnsListResponse { data, .. } = to_response::(read_resp)?; + let ThreadTurnsListResponse { data, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(turn_user_texts(&data), vec!["third", "fourth"]); Ok(()) @@ -426,7 +394,7 @@ async fn thread_turns_list_can_page_backward_and_forward() -> Result<()> { async fn thread_turns_list_supports_requested_items_view() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let conversation_id = create_fake_rollout_with_text_elements( @@ -445,9 +413,8 @@ async fn thread_turns_list_supports_requested_items_view() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let full = read_single_turn_items_view( &mut mcp, @@ -498,7 +465,7 @@ async fn thread_turns_list_supports_requested_items_view() -> Result<()> { async fn thread_search_occurrences_reads_paginated_projection() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let thread_id = codex_protocol::ThreadId::default(); let state_db = codex_state::StateRuntime::init( codex_home.path().to_path_buf(), @@ -606,9 +573,8 @@ async fn thread_search_occurrences_reads_paginated_projection() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_search_occurrences_request(ThreadSearchOccurrencesParams { thread_id: thread_id.to_string(), @@ -684,7 +650,11 @@ async fn thread_turns_list_reads_store_history_without_rollout_path() -> Result< let codex_home = TempDir::new()?; let thread_id = codex_protocol::ThreadId::from_string("00000000-0000-4000-8000-000000000123")?; let store_id = Uuid::new_v4().to_string(); - create_config_toml_with_thread_store(codex_home.path(), &store_id)?; + MockResponsesConfig::new("http://127.0.0.1:1") + .with_root_config(&format!( + r#"experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }}"# + )) + .write(codex_home.path())?; let store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; seed_pathless_store_thread(&store, thread_id).await?; @@ -751,7 +721,11 @@ async fn thread_turns_list_reads_store_history_without_rollout_path() -> Result< async fn thread_read_loaded_include_turns_reads_store_history_without_rollout_path() -> Result<()> { let codex_home = TempDir::new()?; let store_id = Uuid::new_v4().to_string(); - create_config_toml_with_thread_store(codex_home.path(), &store_id)?; + MockResponsesConfig::new("http://127.0.0.1:1") + .with_root_config(&format!( + r#"experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }}"# + )) + .write(codex_home.path())?; let store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; @@ -854,7 +828,11 @@ async fn thread_list_includes_store_thread_without_rollout_path() -> Result<()> let codex_home = TempDir::new()?; let thread_id = codex_protocol::ThreadId::from_string("00000000-0000-4000-8000-000000000124")?; let store_id = Uuid::new_v4().to_string(); - create_config_toml_with_thread_store(codex_home.path(), &store_id)?; + MockResponsesConfig::new("http://127.0.0.1:1") + .with_root_config(&format!( + r#"experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }}"# + )) + .write(codex_home.path())?; let store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; seed_pathless_store_thread(&store, thread_id).await?; @@ -933,7 +911,7 @@ async fn thread_list_includes_store_thread_without_rollout_path() -> Result<()> async fn thread_read_can_return_archived_threads_by_id() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let preview = "Archived saved user message"; @@ -956,9 +934,8 @@ async fn thread_read_can_return_archived_threads_by_id() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_read_request(ThreadReadParams { @@ -966,12 +943,8 @@ async fn thread_read_can_return_archived_threads_by_id() -> Result<()> { include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread } = to_response::(read_resp)?; + let ThreadReadResponse { thread } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(thread.id, conversation_id); assert_eq!(thread.preview, preview); @@ -985,7 +958,7 @@ async fn thread_read_can_return_archived_threads_by_id() -> Result<()> { async fn thread_resume_initial_turns_page_matches_requested_turns_list_page() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let conversation_id = create_fake_rollout_with_text_elements( @@ -1004,9 +977,8 @@ async fn thread_resume_initial_turns_page_matches_requested_turns_list_page() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let turns_list_id = mcp .send_thread_turns_list_request(ThreadTurnsListParams { @@ -1036,16 +1008,11 @@ async fn thread_resume_initial_turns_page_matches_requested_turns_list_page() -> ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread, initial_turns_page, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert!(thread.turns.is_empty()); assert_eq!( @@ -1060,7 +1027,7 @@ async fn thread_resume_initial_turns_page_matches_requested_turns_list_page() -> async fn thread_turns_list_rejects_cursor_when_anchor_turn_is_rolled_back() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let conversation_id = create_fake_rollout_with_text_elements( @@ -1079,9 +1046,8 @@ async fn thread_turns_list_rejects_cursor_when_anchor_turn_is_rolled_back() -> R let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_turns_list_request(ThreadTurnsListParams { @@ -1092,14 +1058,9 @@ async fn thread_turns_list_rejects_cursor_when_anchor_turn_is_rolled_back() -> R items_view: None, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; let ThreadTurnsListResponse { backwards_cursor, .. - } = to_response::(read_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; let backwards_cursor = backwards_cursor.expect("expected backwardsCursor for newest turn"); append_thread_rollback( @@ -1135,7 +1096,7 @@ async fn thread_turns_list_rejects_cursor_when_anchor_turn_is_rolled_back() -> R async fn thread_read_returns_forked_from_id_for_forked_threads() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout_with_text_elements( codex_home.path(), @@ -1150,9 +1111,8 @@ async fn thread_read_returns_forked_from_id_for_forked_threads() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let fork_id = mcp .send_thread_fork_request(ThreadForkParams { @@ -1160,12 +1120,8 @@ async fn thread_read_returns_forked_from_id_for_forked_threads() -> Result<()> { ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread: forked, .. } = to_response::(fork_resp)?; + let ThreadForkResponse { thread: forked, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; let read_id = mcp .send_thread_read_request(ThreadReadParams { @@ -1173,12 +1129,8 @@ async fn thread_read_returns_forked_from_id_for_forked_threads() -> Result<()> { include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(thread.forked_from_id, Some(conversation_id)); @@ -1189,13 +1141,12 @@ async fn thread_read_returns_forked_from_id_for_forked_threads() -> Result<()> { async fn thread_read_loaded_thread_returns_precomputed_path_before_materialization() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1203,12 +1154,8 @@ async fn thread_read_loaded_thread_returns_precomputed_path_before_materializati ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let thread_path = thread.path.clone().expect("thread path"); assert!( !thread_path.exists(), @@ -1221,12 +1168,8 @@ async fn thread_read_loaded_thread_returns_precomputed_path_before_materializati include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread: read, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread: read, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(read.id, thread.id); assert_eq!(read.path, Some(thread_path)); @@ -1241,7 +1184,7 @@ async fn thread_read_loaded_thread_returns_precomputed_path_before_materializati async fn paginated_thread_name_set_is_reflected_in_read_list_and_metadata_resume() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_paginated_rollout( codex_home.path(), @@ -1255,9 +1198,8 @@ async fn paginated_thread_name_set_is_reflected_in_read_list_and_metadata_resume let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // Set a user-facing thread title. let new_name = "Saved user message"; @@ -1267,12 +1209,8 @@ async fn paginated_thread_name_set_is_reflected_in_read_list_and_metadata_resume name: new_name.to_string(), }) .await?; - let set_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(set_id)), - ) - .await??; - let _: ThreadSetNameResponse = to_response::(set_resp)?; + let _: ThreadSetNameResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(set_id)).await??; let notification = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/name/updated"), @@ -1404,13 +1342,12 @@ async fn paginated_thread_name_set_is_reflected_in_read_list_and_metadata_resume async fn thread_read_include_turns_rejects_unmaterialized_loaded_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1418,12 +1355,8 @@ async fn thread_read_include_turns_rejects_unmaterialized_loaded_thread() -> Res ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let thread_path = thread.path.clone().expect("thread path"); assert!( !thread_path.exists(), @@ -1458,13 +1391,12 @@ async fn thread_read_include_turns_rejects_unmaterialized_loaded_thread() -> Res async fn thread_turns_list_rejects_unmaterialized_loaded_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1472,12 +1404,8 @@ async fn thread_turns_list_rejects_unmaterialized_loaded_thread() -> Result<()> ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let thread_path = thread.path.clone().expect("thread path"); assert!( !thread_path.exists(), @@ -1515,7 +1443,7 @@ async fn thread_turns_list_rejects_unmaterialized_loaded_thread() -> Result<()> async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let thread_id = codex_protocol::ThreadId::default(); let state_db = codex_state::StateRuntime::init( codex_home.path().to_path_buf(), @@ -1609,9 +1537,8 @@ async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let expected_turn_1_full = Turn { id: "turn-1".to_string(), @@ -1662,15 +1589,10 @@ async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { ..Default::default() }) .await?; - let legacy_resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(legacy_resume_id)), - ) - .await??; let ThreadResumeResponse { thread: legacy_thread, .. - } = to_response::(legacy_resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(legacy_resume_id)).await??; assert_eq!(legacy_thread.turns, expected_full_turns); let initial_page_resume_id = mcp @@ -1685,16 +1607,15 @@ async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { ..Default::default() }) .await?; - let initial_page_resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(initial_page_resume_id)), - ) - .await??; let ThreadResumeResponse { thread: initial_page_thread, initial_turns_page, .. - } = to_response::(initial_page_resume_resp)?; + } = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_response(initial_page_resume_id), + ) + .await??; assert!(initial_page_thread.turns.is_empty()); assert_eq!( initial_turns_page.expect("initial turns page").data, @@ -1708,17 +1629,12 @@ async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread, turns_backwards_cursor, items_backwards_cursor, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert!(thread.turns.is_empty()); let turns_backwards_cursor = turns_backwards_cursor.expect("resume should return a turn head cursor"); @@ -1732,16 +1648,11 @@ async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { ..Default::default() }) .await?; - let rejoin_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(rejoin_id)), - ) - .await??; let ThreadResumeResponse { turns_backwards_cursor: rejoin_turns_backwards_cursor, items_backwards_cursor: rejoin_items_backwards_cursor, .. - } = to_response::(rejoin_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(rejoin_id)).await??; assert_eq!( rejoin_turns_backwards_cursor.as_deref(), Some(turns_backwards_cursor.as_str()) @@ -1916,12 +1827,8 @@ async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { ..Default::default() }) .await?; - let turn_start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_start_response)?; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_start_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -1935,14 +1842,13 @@ async fn paginated_history_lists_use_projected_turns_and_items() -> Result<()> { async fn thread_items_list_returns_unsupported() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_items_list_request(ThreadItemsListParams { @@ -1977,13 +1883,12 @@ async fn thread_read_reports_system_error_idle_flag_after_failed_turn() -> Resul ) .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1991,12 +1896,8 @@ async fn thread_read_reports_system_error_idle_flag_after_failed_turn() -> Resul ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_start_id = mcp .send_turn_start_request(TurnStartParams { @@ -2009,12 +1910,8 @@ async fn thread_read_reports_system_error_idle_flag_after_failed_turn() -> Resul ..Default::default() }) .await?; - let turn_start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_start_response)?; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_start_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("error"), @@ -2027,12 +1924,8 @@ async fn thread_read_reports_system_error_idle_flag_after_failed_turn() -> Resul include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(thread.status, ThreadStatus::SystemError,); @@ -2301,51 +2194,3 @@ fn store_history_items() -> Vec { }, ))] } - -fn create_config_toml_with_thread_store(codex_home: &Path, store_id: &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" -experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }} - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "http://127.0.0.1:1/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - -// Helper to create a config.toml pointing at the mock model server. -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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 96182b375eae..21018e9b576f 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -1,5 +1,6 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_apply_patch_sse_response; use app_test_support::create_fake_paginated_rollout; @@ -62,6 +63,7 @@ use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput; use codex_config::types::AuthCredentialsStoreMode; use codex_core::ARCHIVED_SESSIONS_SUBDIR; +use codex_features::Feature; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationMode; @@ -133,7 +135,7 @@ const CODEX_5_2_INSTRUCTIONS_TEMPLATE_DEFAULT: &str = "You are Codex, a coding a async fn thread_resume_paginated_metadata_only_uses_model_context() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_paginated_rollout( codex_home.path(), "2025-01-05T12-00-00", @@ -146,9 +148,8 @@ async fn thread_resume_paginated_metadata_only_uses_model_context() -> Result<() let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let resume_id = primary .send_thread_resume_request(ThreadResumeParams { @@ -157,14 +158,9 @@ async fn thread_resume_paginated_metadata_only_uses_model_context() -> Result<() ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, primary.read_response(resume_id)).await??; assert_eq!(resumed.id, conversation_id); assert_eq!(resumed.history_mode, ThreadHistoryMode::Paginated); assert!(resumed.turns.is_empty()); @@ -209,13 +205,12 @@ async fn wait_for_responses_request_count( async fn thread_resume_rejects_unmaterialized_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // Start a thread. let start_id = mcp @@ -224,12 +219,8 @@ async fn thread_resume_rejects_unmaterialized_thread() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; // Resume should fail before the first user message materializes rollout storage. let resume_id = mcp @@ -259,13 +250,12 @@ async fn thread_resume_rejects_unmaterialized_thread() -> Result<()> { async fn thread_resume_with_empty_path_uses_running_thread_id() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -273,12 +263,8 @@ async fn thread_resume_with_empty_path_uses_running_thread_id() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -310,14 +296,9 @@ async fn thread_resume_with_empty_path_uses_running_thread_id() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(resumed.id, thread.id); Ok(()) @@ -332,7 +313,7 @@ async fn thread_resume_running_thread_uses_cached_instruction_sources() -> Resul let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let workspace = TempDir::new()?; let project_agents = workspace.path().join("AGENTS.md"); std::fs::write(&project_agents, "project instructions")?; @@ -341,9 +322,8 @@ async fn thread_resume_running_thread_uses_cached_instruction_sources() -> Resul .with_codex_home(codex_home.path()) // TODO(anp): Move the cached instruction-source fixture into the auto environment cwd. .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request(ThreadStartParams { @@ -351,16 +331,11 @@ async fn thread_resume_running_thread_uses_cached_instruction_sources() -> Resul ..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, instruction_sources, .. - } = to_response::(start_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let project_agents = AbsolutePathBuf::try_from(project_agents)?; let project_agents_source = LegacyAppPathString::from_abs_path(&project_agents); assert_eq!(instruction_sources, vec![project_agents_source.clone()]); @@ -395,15 +370,10 @@ async fn thread_resume_running_thread_uses_cached_instruction_sources() -> Resul ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { instruction_sources, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(instruction_sources, vec![project_agents_source]); @@ -414,7 +384,7 @@ async fn thread_resume_running_thread_uses_cached_instruction_sources() -> Resul async fn turn_start_updates_runtime_workspace_roots_for_loaded_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let extra_root_tmp = TempDir::new()?; let extra_root = extra_root_tmp.path().join("extra-root"); @@ -422,9 +392,8 @@ async fn turn_start_updates_runtime_workspace_roots_for_loaded_thread() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -432,12 +401,8 @@ async fn turn_start_updates_runtime_workspace_roots_for_loaded_thread() -> Resul ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -472,15 +437,10 @@ async fn turn_start_updates_runtime_workspace_roots_for_loaded_thread() -> Resul ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { runtime_workspace_roots, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!( runtime_workspace_roots, @@ -494,15 +454,14 @@ async fn turn_start_updates_runtime_workspace_roots_for_loaded_thread() -> Resul async fn thread_resume_preserves_persisted_approvals_reviewer() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let thread_id = { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request(ThreadStartParams { @@ -511,12 +470,8 @@ async fn thread_resume_preserves_persisted_approvals_reviewer() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -556,23 +511,17 @@ async fn thread_resume_preserves_persisted_approvals_reviewer() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id, ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { approvals_reviewer, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(approvals_reviewer, ApprovalsReviewer::AutoReview); @@ -583,7 +532,7 @@ async fn thread_resume_preserves_persisted_approvals_reviewer() -> Result<()> { async fn thread_resume_preserves_goal_first_and_fork_approvals_reviewer() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -595,9 +544,8 @@ async fn thread_resume_preserves_goal_first_and_fork_approvals_reviewer() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -606,12 +554,8 @@ async fn thread_resume_preserves_goal_first_and_fork_approvals_reviewer() -> Res ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let rollout_path = thread.path.clone().expect("thread path"); for objective in [ @@ -628,12 +572,8 @@ async fn thread_resume_preserves_goal_first_and_fork_approvals_reviewer() -> Res })), ) .await?; - let goal_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let _: ThreadGoalSetResponse = to_response(goal_resp)?; + let _: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(goal_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/goal/updated"), @@ -656,16 +596,11 @@ async fn thread_resume_preserves_goal_first_and_fork_approvals_reviewer() -> Res ..Default::default() }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; let ThreadForkResponse { thread: fork_thread, approvals_reviewer, .. - } = to_response::(fork_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(fork_id)).await??; assert_eq!(approvals_reviewer, ApprovalsReviewer::User); (thread.id, fork_thread.id) @@ -683,9 +618,8 @@ async fn thread_resume_preserves_goal_first_and_fork_approvals_reviewer() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; for (thread_id, expected_reviewer) in [ (thread_id, ApprovalsReviewer::AutoReview), (fork_thread_id, ApprovalsReviewer::User), @@ -696,14 +630,9 @@ async fn thread_resume_preserves_goal_first_and_fork_approvals_reviewer() -> Res ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { approvals_reviewer, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(approvals_reviewer, expected_reviewer); } @@ -716,7 +645,7 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config_toml = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -730,9 +659,8 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe let thread_id = { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -740,12 +668,8 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -778,12 +702,8 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe ..Default::default() }) .await?; - let update_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(update_id)), - ) - .await??; - let _: ThreadSettingsUpdateResponse = to_response(update_resp)?; + let _: ThreadSettingsUpdateResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(update_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/settings/updated"), @@ -795,26 +715,20 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id: thread_id.clone(), ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { model, reasoning_effort, approvals_reviewer, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(model, "gpt-5.2-codex"); assert_eq!(reasoning_effort, Some(ReasoningEffort::Ultra)); @@ -834,12 +748,8 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe ..Default::default() }) .await?; - let update_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(update_id)), - ) - .await??; - let _: ThreadSettingsUpdateResponse = to_response(update_resp)?; + let _: ThreadSettingsUpdateResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(update_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/settings/updated"), @@ -849,23 +759,17 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id, ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { reasoning_effort, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(reasoning_effort, None); @@ -876,7 +780,7 @@ async fn thread_resume_preserves_acknowledged_model_effort_and_approvals_reviewe async fn thread_goal_get_rejects_unmaterialized_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -887,9 +791,8 @@ async fn thread_goal_get_rejects_unmaterialized_thread() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -898,12 +801,8 @@ async fn thread_goal_get_rejects_unmaterialized_thread() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let goal_id = mcp .send_raw_request( @@ -935,7 +834,7 @@ async fn goal_first_live_thread_appears_in_state_db_thread_list() -> Result<()> let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; let codex_home_path = normalized_existing_path(codex_home.path())?; - create_config_toml(&codex_home_path, &server.uri())?; + mock_responses_config(&server.uri()).write(&codex_home_path)?; let config_path = codex_home_path.join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -951,9 +850,8 @@ async fn goal_first_live_thread_appears_in_state_db_thread_list() -> Result<()> .with_codex_home(&codex_home_path) .without_managed_config() .with_env_overrides(&[("CODEX_SQLITE_HOME", Some(sqlite_home))]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -961,12 +859,8 @@ async fn goal_first_live_thread_appears_in_state_db_thread_list() -> Result<()> ..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, cwd, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, cwd, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let goal_id = mcp .send_raw_request( @@ -978,12 +872,8 @@ async fn goal_first_live_thread_appears_in_state_db_thread_list() -> Result<()> })), ) .await?; - let goal_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let _goal: ThreadGoalSetResponse = to_response(goal_resp)?; + let _goal: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(goal_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/goal/updated"), @@ -1003,12 +893,8 @@ async fn goal_first_live_thread_appears_in_state_db_thread_list() -> Result<()> })), ) .await?; - let list_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; - let list: ThreadListResponse = to_response(list_resp)?; + let list: ThreadListResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(list_id)).await??; assert_eq!( list.data .iter() @@ -1025,7 +911,9 @@ async fn thread_resume_tracks_thread_initialized_analytics() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml_with_chatgpt_base_url(codex_home.path(), &server.uri(), &server.uri())?; + mock_responses_config(&server.uri()) + .with_root_config(&format!(r#"chatgpt_base_url = "{}""#, server.uri())) + .write(codex_home.path())?; mount_analytics_capture(&server, codex_home.path()).await?; let conversation_id = create_fake_rollout( @@ -1047,9 +935,8 @@ async fn thread_resume_tracks_thread_initialized_analytics() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -1057,12 +944,8 @@ async fn thread_resume_tracks_thread_initialized_analytics() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert!( !thread.session_id.is_empty(), "session id should not be empty" @@ -1089,15 +972,16 @@ async fn thread_resume_running_thread_tracks_thread_originator_in_analytics() -> let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml_with_chatgpt_base_url(codex_home.path(), &server.uri(), &server.uri())?; + mock_responses_config(&server.uri()) + .with_root_config(&format!(r#"chatgpt_base_url = "{}""#, server.uri())) + .write(codex_home.path())?; mount_analytics_capture(&server, codex_home.path()).await?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1107,12 +991,8 @@ async fn thread_resume_running_thread_tracks_thread_originator_in_analytics() -> ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -1143,14 +1023,9 @@ async fn thread_resume_running_thread_tracks_thread_originator_in_analytics() -> ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; let event = wait_for_matching_analytics_event(&server, DEFAULT_READ_TIMEOUT, |event| { event["event_type"] == "codex_thread_initialized" @@ -1195,7 +1070,7 @@ fn set_session_meta_on_fake_rollout( async fn thread_resume_returns_rollout_history() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let preview = "Saved user message"; let text_elements = vec![TextElement::new( @@ -1217,9 +1092,8 @@ async fn thread_resume_returns_rollout_history() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -1227,12 +1101,8 @@ async fn thread_resume_returns_rollout_history() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(thread.id, conversation_id); assert_eq!(thread.preview, preview); @@ -1392,7 +1262,7 @@ async fn thread_resume_redacts_payloads_for_chatgpt_remote_clients() -> Result<( async fn resume_redaction_fixture(client_name: Option<&str>) -> Result { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let meta_rfc3339 = "2025-01-05T12:00:00Z"; @@ -1511,7 +1381,7 @@ fn append_resume_redaction_history( async fn thread_resume_can_skip_turns_for_metadata_only_resume() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout_with_text_elements( codex_home.path(), @@ -1525,9 +1395,8 @@ async fn thread_resume_can_skip_turns_for_metadata_only_resume() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -1536,12 +1405,8 @@ async fn thread_resume_can_skip_turns_for_metadata_only_resume() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(thread.id, conversation_id); assert!(thread.turns.is_empty()); @@ -1553,7 +1418,7 @@ async fn thread_resume_can_skip_turns_for_metadata_only_resume() -> Result<()> { async fn thread_resume_rejects_archived_session_by_id() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let conversation_id = create_fake_rollout_with_text_elements( @@ -1575,9 +1440,8 @@ async fn thread_resume_rejects_archived_session_by_id() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -1607,7 +1471,7 @@ async fn thread_resume_rejects_archived_session_by_id() -> Result<()> { async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -1618,9 +1482,8 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1628,12 +1491,8 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -1667,12 +1526,8 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { })), ) .await?; - let goal_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let _goal: ThreadGoalSetResponse = to_response(goal_resp)?; + let _goal: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(goal_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/goal/updated"), @@ -1686,12 +1541,8 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let _resume: ThreadResumeResponse = to_response(resume_resp)?; + let _resume: ThreadResumeResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; let notification = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/goal/updated"), @@ -1716,7 +1567,7 @@ async fn thread_resume_keeps_paused_goal_paused() -> Result<()> { async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -1727,9 +1578,8 @@ async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1737,12 +1587,8 @@ async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -1777,12 +1623,8 @@ async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> })), ) .await?; - let goal_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let goal: ThreadGoalSetResponse = to_response(goal_resp)?; + let goal: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(goal_id)).await??; assert_eq!(goal.goal.status, ThreadGoalStatus::BudgetLimited); timeout( @@ -1800,12 +1642,8 @@ async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> })), ) .await?; - let replacement_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(replacement_id)), - ) - .await??; - let replacement: ThreadGoalSetResponse = to_response(replacement_resp)?; + let replacement: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(replacement_id)).await??; assert_eq!(replacement.goal.status, ThreadGoalStatus::BudgetLimited); assert_eq!(replacement.goal.token_budget, Some(10)); @@ -1819,7 +1657,7 @@ async fn thread_goal_set_preserves_budget_limited_same_objective() -> Result<()> async fn thread_goal_set_persists_resumable_stopped_statuses() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -1830,9 +1668,8 @@ async fn thread_goal_set_persists_resumable_stopped_statuses() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1840,12 +1677,8 @@ async fn thread_goal_set_persists_resumable_stopped_statuses() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -1883,12 +1716,8 @@ async fn thread_goal_set_persists_resumable_stopped_statuses() -> Result<()> { })), ) .await?; - let goal_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let goal: ThreadGoalSetResponse = to_response(goal_resp)?; + let goal: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(goal_id)).await??; assert_eq!(goal.goal.status, expected_status); let notification = timeout( @@ -1910,7 +1739,7 @@ async fn thread_goal_set_persists_resumable_stopped_statuses() -> Result<()> { async fn thread_goal_set_edits_objective_without_resetting_usage() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -1929,9 +1758,8 @@ async fn thread_goal_set_edits_objective_without_resetting_usage() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let goal_id = mcp .send_raw_request( @@ -1944,12 +1772,8 @@ async fn thread_goal_set_edits_objective_without_resetting_usage() -> Result<()> })), ) .await?; - let goal_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let goal: ThreadGoalSetResponse = to_response(goal_resp)?; + let goal: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(goal_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/goal/updated"), @@ -1991,12 +1815,8 @@ async fn thread_goal_set_edits_objective_without_resetting_usage() -> Result<()> })), ) .await?; - let edit_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(edit_id)), - ) - .await??; - let edit: ThreadGoalSetResponse = to_response(edit_resp)?; + let edit: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(edit_id)).await??; let updated_goal = state_db .thread_goals() .get_thread_goal(thread_id) @@ -2033,7 +1853,9 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul ]) .await; let codex_home = TempDir::new()?; - create_config_toml_with_chatgpt_base_url(codex_home.path(), &server.uri(), &server.uri())?; + mock_responses_config(&server.uri()) + .with_root_config(&format!(r#"chatgpt_base_url = "{}""#, server.uri())) + .write(codex_home.path())?; let config_path = codex_home.path().join("config.toml"); let config = std::fs::read_to_string(&config_path)?; std::fs::write( @@ -2045,9 +1867,8 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT.saturating_mul(2)) .await?; - timeout(DEFAULT_READ_TIMEOUT.saturating_mul(2), mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -2055,12 +1876,8 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -2094,12 +1911,8 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul })), ) .await?; - let goal_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(goal_id)), - ) - .await??; - let _goal: ThreadGoalSetResponse = to_response(goal_resp)?; + let _goal: ThreadGoalSetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(goal_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("thread/goal/updated"), @@ -2163,12 +1976,8 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul })), ) .await?; - let clear_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(clear_id)), - ) - .await??; - let clear: ThreadGoalClearResponse = to_response(clear_resp)?; + let clear: ThreadGoalClearResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(clear_id)).await??; assert!(clear.cleared); timeout( @@ -2190,12 +1999,8 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul })), ) .await?; - let get_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(get_id)), - ) - .await??; - let get: codex_app_server_protocol::ThreadGoalGetResponse = to_response(get_resp)?; + let get: codex_app_server_protocol::ThreadGoalGetResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(get_id)).await??; assert_eq!(None, get.goal); let clear_again_id = mcp @@ -2206,12 +2011,8 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul })), ) .await?; - let clear_again_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(clear_again_id)), - ) - .await??; - let clear_again: ThreadGoalClearResponse = to_response(clear_again_resp)?; + let clear_again: ThreadGoalClearResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(clear_again_id)).await??; assert!(!clear_again.cleared); Ok(()) @@ -2221,7 +2022,7 @@ async fn thread_goal_lifecycle_emits_analytics_and_clear_deletes_goal() -> Resul async fn thread_resume_emits_restored_token_usage_before_next_turn() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout_with_token_usage( codex_home.path(), @@ -2233,9 +2034,8 @@ async fn thread_resume_emits_restored_token_usage_before_next_turn() -> Result<( let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -2243,12 +2043,8 @@ async fn thread_resume_emits_restored_token_usage_before_next_turn() -> Result<( ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; let note = timeout( DEFAULT_READ_TIMEOUT, @@ -2277,7 +2073,7 @@ async fn thread_resume_emits_restored_token_usage_before_next_turn() -> Result<( async fn thread_resume_skips_restored_token_usage_when_turns_are_excluded() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let conversation_id = create_fake_rollout_with_token_usage( codex_home.path(), @@ -2289,9 +2085,8 @@ async fn thread_resume_skips_restored_token_usage_when_turns_are_excluded() -> R let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let first_resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -2326,15 +2121,10 @@ async fn thread_resume_skips_restored_token_usage_when_turns_are_excluded() -> R ..Default::default() }) .await?; - let second_resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed_again, .. - } = to_response::(second_resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(second_resume_id)).await??; assert!(resumed_again.turns.is_empty()); let second_note = timeout( @@ -2354,7 +2144,7 @@ async fn thread_resume_skips_restored_token_usage_when_turns_are_excluded() -> R async fn thread_resume_token_usage_replay_ignores_stale_interrupted_tail_turn() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let meta_rfc3339 = "2025-01-05T12:00:00Z"; @@ -2400,9 +2190,8 @@ async fn thread_resume_token_usage_replay_ignores_stale_interrupted_tail_turn() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -2410,12 +2199,8 @@ async fn thread_resume_token_usage_replay_ignores_stale_interrupted_tail_turn() ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(thread.turns.len(), 2); assert_eq!(thread.turns[0].status, TurnStatus::Completed); @@ -2445,7 +2230,7 @@ async fn thread_resume_token_usage_replay_ignores_stale_interrupted_tail_turn() async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let meta_rfc3339 = "2025-01-05T12:00:00Z"; @@ -2530,9 +2315,8 @@ async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -2540,12 +2324,8 @@ async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Re ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(thread.turns.len(), 2); assert_eq!(thread.turns[0].status, TurnStatus::Completed); @@ -2574,31 +2354,9 @@ async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Re async fn thread_resume_prefers_persisted_git_metadata_for_local_threads() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - let config_toml = codex_home.path().join("config.toml"); - std::fs::write( - &config_toml, - format!( - r#" -model = "gpt-5.4" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -personality = true -sqlite = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"#, - server.uri() - ), - )?; + mock_responses_config(&server.uri()) + .enable_feature(Feature::Sqlite) + .write(codex_home.path())?; let repo_path = codex_home.path().join("repo"); std::fs::create_dir_all(&repo_path)?; @@ -2729,9 +2487,8 @@ stream_max_retries = 0 let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let update_id = mcp .send_thread_metadata_update_request(ThreadMetadataUpdateParams { @@ -2755,12 +2512,8 @@ stream_max_retries = 0 ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!( thread @@ -2778,7 +2531,7 @@ async fn thread_resume_and_read_interrupt_incomplete_rollout_turn_when_thread_is { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let filename_ts = "2025-01-05T12-00-00"; let meta_rfc3339 = "2025-01-05T12:00:00Z"; @@ -2826,9 +2579,8 @@ async fn thread_resume_and_read_interrupt_incomplete_rollout_turn_when_thread_is let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -2836,12 +2588,8 @@ async fn thread_resume_and_read_interrupt_incomplete_rollout_turn_when_thread_is ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(thread.status, ThreadStatus::Idle); assert_eq!(thread.turns.len(), 2); @@ -2855,15 +2603,10 @@ async fn thread_resume_and_read_interrupt_incomplete_rollout_turn_when_thread_is ..Default::default() }) .await?; - let second_resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed_again, .. - } = to_response::(second_resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(second_resume_id)).await??; assert_eq!(resumed_again.status, ThreadStatus::Idle); assert_eq!(resumed_again.turns.len(), 2); @@ -2876,15 +2619,10 @@ async fn thread_resume_and_read_interrupt_incomplete_rollout_turn_when_thread_is include_turns: true, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; let ThreadReadResponse { thread: read_thread, .. - } = to_response::(read_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert_eq!(read_thread.status, ThreadStatus::Idle); assert_eq!(read_thread.turns.len(), 2); @@ -2903,9 +2641,8 @@ async fn thread_resume_defers_updated_at_until_turn_start() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let read_id = mcp .send_thread_read_request(ThreadReadParams { @@ -2913,15 +2650,10 @@ async fn thread_resume_defers_updated_at_until_turn_start() -> Result<()> { include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; let ThreadReadResponse { thread: before_resume, .. - } = to_response::(read_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -2929,12 +2661,8 @@ async fn thread_resume_defers_updated_at_until_turn_start() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(thread.updated_at, before_resume.updated_at); assert_eq!(thread.recency_at, before_resume.recency_at); @@ -2962,12 +2690,8 @@ async fn thread_resume_defers_updated_at_until_turn_start() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let ThreadResumeResponse { cwd, .. } = to_response::(resume_resp)?; + let ThreadResumeResponse { cwd, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(cwd, AbsolutePathBuf::from_absolute_path(codex_home.path())?); let turn_id = mcp @@ -2997,15 +2721,10 @@ async fn thread_resume_defers_updated_at_until_turn_start() -> Result<()> { include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; let ThreadReadResponse { thread: after_turn_start, .. - } = to_response::(read_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(read_id)).await??; assert!(after_turn_start.recency_at > before_resume.recency_at); timeout( @@ -3024,13 +2743,12 @@ async fn thread_resume_defers_updated_at_until_turn_start() -> Result<()> { async fn thread_resume_keeps_in_flight_turn_streaming() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -3038,12 +2756,8 @@ async fn thread_resume_keeps_in_flight_turn_streaming() -> Result<()> { ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let seed_turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3070,9 +2784,8 @@ async fn thread_resume_keeps_in_flight_turn_streaming() -> Result<()> { let mut secondary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, secondary.initialize()).await??; let turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3102,15 +2815,10 @@ async fn thread_resume_keeps_in_flight_turn_streaming() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - secondary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed_thread, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, secondary.read_response(resume_id)).await??; assert_ne!(resumed_thread.status, ThreadStatus::NotLoaded); timeout( @@ -3139,13 +2847,12 @@ async fn thread_resume_rejects_history_when_thread_is_running() -> Result<()> { let _first_response_mock = responses::mount_sse_once(&server, first_body).await; let _second_response_mock = responses::mount_response_once(&server, second_response).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -3153,12 +2860,8 @@ async fn thread_resume_rejects_history_when_thread_is_running() -> Result<()> { ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let seed_turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3261,13 +2964,12 @@ async fn thread_resume_rejects_mismatched_path_for_running_thread_id() -> Result let _first_response_mock = responses::mount_sse_once(&server, first_body).await; let _second_response_mock = responses::mount_response_once(&server, second_response).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -3275,12 +2977,8 @@ async fn thread_resume_rejects_mismatched_path_for_running_thread_id() -> Result ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let seed_turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3433,13 +3131,12 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R let _response_mock = responses::mount_response_sequence(&server, vec![first_response, second_response]).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -3448,12 +3145,8 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let seed_turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3466,12 +3159,8 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R ..Default::default() }) .await?; - let seed_turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(seed_turn_id)), - ) - .await??; - let TurnStartResponse { turn: seed_turn } = to_response::(seed_turn_resp)?; + let TurnStartResponse { turn: seed_turn } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(seed_turn_id)).await??; timeout( DEFAULT_READ_TIMEOUT, primary.read_stream_until_notification_message("turn/completed"), @@ -3516,17 +3205,12 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread, model, initial_turns_page, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, primary.read_response(resume_id)).await??; assert_eq!(model, "gpt-5.4"); let initial_turns_page = initial_turns_page.expect("resume should include initial turns page"); assert_eq!(initial_turns_page.data.len(), 1); @@ -3553,14 +3237,9 @@ async fn thread_resume_rejoins_running_paginated_thread_with_initial_page() -> R ..Default::default() }) .await?; - let asc_resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(asc_resume_id)), - ) - .await??; let ThreadResumeResponse { initial_turns_page, .. - } = to_response::(asc_resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, primary.read_response(asc_resume_id)).await??; let initial_turns_page = initial_turns_page.expect("resume should include initial turns page"); assert_eq!(initial_turns_page.data.len(), 1); assert_eq!(initial_turns_page.data[0].id, seed_turn.id); @@ -3595,13 +3274,12 @@ async fn thread_resume_can_skip_turns_when_thread_is_running() -> Result<()> { ) .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -3609,12 +3287,8 @@ async fn thread_resume_can_skip_turns_when_thread_is_running() -> Result<()> { ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3640,9 +3314,8 @@ async fn thread_resume_can_skip_turns_when_thread_is_running() -> Result<()> { let mut secondary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, secondary.initialize()).await??; let resume_id = secondary .send_thread_resume_request(ThreadResumeParams { @@ -3651,14 +3324,9 @@ async fn thread_resume_can_skip_turns_when_thread_is_running() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - secondary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, secondary.read_response(resume_id)).await??; assert_eq!(resumed.id, thread.id); assert_eq!(resumed.status, ThreadStatus::Idle); @@ -3691,13 +3359,12 @@ async fn thread_resume_replays_pending_command_execution_request_approval() -> R ]; let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -3705,12 +3372,8 @@ async fn thread_resume_replays_pending_command_execution_request_approval() -> R ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let seed_turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3768,15 +3431,10 @@ async fn thread_resume_replays_pending_command_execution_request_approval() -> R ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed_thread, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, primary.read_response(resume_id)).await??; assert_eq!(resumed_thread.id, thread.id); assert!( resumed_thread @@ -3839,13 +3497,12 @@ async fn thread_resume_replays_pending_file_change_request_approval() -> Result< create_final_assistant_message_sse_response("done")?, ]; let server = create_mock_responses_server_sequence_unchecked(responses).await; - create_config_toml(&codex_home, &server.uri())?; + mock_responses_config(&server.uri()).write(&codex_home)?; let mut primary = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -3854,12 +3511,8 @@ async fn thread_resume_replays_pending_file_change_request_approval() -> Result< ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let seed_turn_id = primary .send_turn_start_request(TurnStartParams { @@ -3945,15 +3598,10 @@ async fn thread_resume_replays_pending_file_change_request_approval() -> Result< ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed_thread, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, primary.read_response(resume_id)).await??; assert_eq!(resumed_thread.id, thread.id); assert!( resumed_thread @@ -3995,7 +3643,7 @@ async fn thread_resume_replays_pending_file_change_request_approval() -> Result< async fn thread_resume_with_overrides_defers_updated_at_until_turn_start() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let RestartedThreadFixture { mut mcp, @@ -4014,15 +3662,10 @@ async fn thread_resume_with_overrides_defers_updated_at_until_turn_start() -> Re ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed_thread, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(resumed_thread.updated_at, updated_at); assert_eq!(resumed_thread.status, ThreadStatus::Idle); @@ -4063,13 +3706,18 @@ async fn thread_resume_fails_when_required_mcp_server_fails_to_initialize() -> R let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; let rollout = setup_rollout_fixture(codex_home.path(), &server.uri()).await?; - create_config_toml_with_required_broken_mcp(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()) + .with_extra_config( + r#"[mcp_servers.required_broken] +command = "codex-definitely-not-a-real-binary" +required = true"#, + ) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -4122,11 +3770,9 @@ async fn thread_resume_surfaces_cloud_config_bundle_load_errors() -> Result<()> let codex_home = TempDir::new()?; let model_server = create_mock_responses_server_repeating_assistant("Done").await; let chatgpt_base_url = format!("{}/backend-api", server.uri()); - create_config_toml_with_chatgpt_base_url( - codex_home.path(), - &model_server.uri(), - &chatgpt_base_url, - )?; + mock_responses_config(&model_server.uri()) + .with_root_config(&format!(r#"chatgpt_base_url = "{chatgpt_base_url}""#)) + .write(codex_home.path())?; write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("chatgpt-token") @@ -4156,9 +3802,8 @@ async fn thread_resume_surfaces_cloud_config_bundle_load_errors() -> Result<()> Some(refresh_token_url.as_str()), ), ]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { @@ -4195,7 +3840,7 @@ async fn thread_resume_surfaces_cloud_config_bundle_load_errors() -> Result<()> async fn thread_resume_uses_path_over_non_running_thread_id() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let RestartedThreadFixture { mut mcp, @@ -4212,14 +3857,9 @@ async fn thread_resume_uses_path_over_non_running_thread_id() -> Result<()> { }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(resumed.id, thread_id); Ok(()) @@ -4230,7 +3870,7 @@ async fn thread_resume_can_load_source_by_external_path() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; let external_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let thread_id = create_fake_rollout( external_home.path(), "2025-01-05T12-00-00", @@ -4243,9 +3883,8 @@ async fn thread_resume_can_load_source_by_external_path() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let resume_id = mcp .send_thread_resume_request(ThreadResumeParams { thread_id: "not-a-valid-thread-id".to_string(), @@ -4254,14 +3893,9 @@ async fn thread_resume_can_load_source_by_external_path() -> Result<()> { }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert_eq!(resumed.id, thread_id); let resumed_path = resumed.path.as_ref().expect("resumed thread path"); assert_eq!( @@ -4278,7 +3912,7 @@ async fn thread_resume_can_load_source_by_external_path() -> Result<()> { async fn thread_resume_supports_history_and_overrides() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let RestartedThreadFixture { mut mcp, thread_id, .. @@ -4305,16 +3939,11 @@ async fn thread_resume_supports_history_and_overrides() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; let ThreadResumeResponse { thread: resumed, model_provider, .. - } = to_response::(resume_resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(resume_id)).await??; assert!(!resumed.id.is_empty()); assert_eq!(model_provider, "mock_provider"); assert_eq!(resumed.preview, history_text); @@ -4336,9 +3965,8 @@ async fn start_materialized_thread_and_restart( ) -> Result { let mut first_mcp = TestAppServer::builder() .with_codex_home(codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, first_mcp.initialize()).await??; let start_id = first_mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -4346,12 +3974,8 @@ async fn start_materialized_thread_and_restart( ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, first_mcp.read_response(start_id)).await??; let materialize_turn_id = first_mcp .send_turn_start_request(TurnStartParams { @@ -4381,12 +4005,8 @@ async fn start_materialized_thread_and_restart( include_turns: false, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - first_mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, first_mcp.read_response(read_id)).await??; let thread_id = thread.id; let rollout_file_path = thread @@ -4396,11 +4016,10 @@ async fn start_materialized_thread_and_restart( drop(first_mcp); - let mut second_mcp = TestAppServer::builder() + let second_mcp = TestAppServer::builder() .with_codex_home(codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, second_mcp.initialize()).await??; Ok(RestartedThreadFixture { mcp: second_mcp, @@ -4428,13 +4047,12 @@ async fn thread_resume_accepts_personality_override() -> Result<()> { let response_mock = responses::mount_sse_sequence(&server, vec![first_body, second_body]).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + mock_responses_config(&server.uri()).write(codex_home.path())?; let mut primary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, primary.initialize()).await??; let start_id = primary .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -4442,12 +4060,8 @@ async fn thread_resume_accepts_personality_override() -> Result<()> { ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - primary.read_stream_until_response_message(RequestId::Integer(start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, primary.read_response(start_id)).await??; let materialize_id = primary .send_turn_start_request(TurnStartParams { @@ -4473,9 +4087,8 @@ async fn thread_resume_accepts_personality_override() -> Result<()> { let mut secondary = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, secondary.initialize()).await??; let resume_id = secondary .send_thread_resume_request(ThreadResumeParams { @@ -4485,12 +4098,8 @@ async fn thread_resume_accepts_personality_override() -> Result<()> { ..Default::default() }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - secondary.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let resume: ThreadResumeResponse = to_response::(resume_resp)?; + let resume: ThreadResumeResponse = + timeout(DEFAULT_READ_TIMEOUT, secondary.read_response(resume_id)).await??; assert_eq!(resume.thread.status, ThreadStatus::Idle); let turn_id = secondary @@ -4536,95 +4145,10 @@ async fn thread_resume_accepts_personality_override() -> Result<()> { Ok(()) } -// Helper to create a config.toml pointing at the mock model server. -fn create_config_toml(codex_home: &std::path::Path, server_uri: &str) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "gpt-5.4" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -personality = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - -fn create_config_toml_with_chatgpt_base_url( - codex_home: &std::path::Path, - server_uri: &str, - chatgpt_base_url: &str, -) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "gpt-5.4" -approval_policy = "never" -sandbox_mode = "read-only" -chatgpt_base_url = "{chatgpt_base_url}" - -model_provider = "mock_provider" - -[features] -personality = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - -fn create_config_toml_with_required_broken_mcp( - codex_home: &std::path::Path, - server_uri: &str, -) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "gpt-5.4" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -personality = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - -[mcp_servers.required_broken] -command = "codex-definitely-not-a-real-binary" -required = true -"# - ), - ) +fn mock_responses_config(server_uri: &str) -> MockResponsesConfig { + MockResponsesConfig::new(server_uri) + .with_model("gpt-5.4") + .enable_feature(Feature::Personality) } #[allow(dead_code)] @@ -4645,7 +4169,7 @@ struct RolloutFixture { } async fn setup_rollout_fixture(codex_home: &Path, server_uri: &str) -> Result { - create_config_toml(codex_home, server_uri)?; + mock_responses_config(server_uri).write(codex_home)?; let preview = "Saved user message"; let filename_ts = "2025-01-05T12-00-00"; diff --git a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs index a9237f9a9a56..29196c1f3f55 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_rollback.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_rollback.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_repeating_assistant; @@ -32,7 +33,7 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs async fn thread_rollback_rejects_paginated_thread() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -129,7 +130,7 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -297,26 +298,3 @@ async fn thread_rollback_drops_last_turns_and_persists_to_rollout() -> Result<() Ok(()) } - -fn create_config_toml(codex_home: &std::path::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" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/thread_settings_update.rs b/codex-rs/app-server/tests/suite/v2/thread_settings_update.rs index ed6196d7aa27..bbab6e3d6ccd 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_settings_update.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_settings_update.rs @@ -1,14 +1,11 @@ use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::to_response; -use app_test_support::write_mock_responses_config_toml; use app_test_support::write_models_cache; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SandboxPolicy; use codex_app_server_protocol::ThreadReadParams; @@ -26,7 +23,6 @@ use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::Value; -use std::collections::BTreeMap; use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; @@ -46,9 +42,8 @@ async fn thread_settings_update_emits_notification_and_updates_future_turns() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let thread = start_thread(&mut mcp).await?.thread; send_thread_settings_update( @@ -113,9 +108,8 @@ async fn thread_settings_update_cwd_retargets_default_environment() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_start_request(ThreadStartParams { cwd: Some(initial_workspace.path().to_string_lossy().into_owned()), @@ -123,12 +117,8 @@ async fn thread_settings_update_cwd_retargets_default_environment() -> Result<() ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let thread = to_response::(response)?.thread; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; send_thread_settings_update( &mut mcp, @@ -185,9 +175,8 @@ async fn thread_settings_update_while_turn_is_active_emits_notification() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let thread = start_thread(&mut mcp).await?.thread; start_text_turn(&mut mcp, thread.id.clone()).await?; timeout( @@ -231,9 +220,8 @@ async fn thread_settings_update_null_service_tier_uses_default() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let thread = start_thread(&mut mcp).await?.thread; send_thread_settings_update( @@ -300,9 +288,8 @@ async fn thread_settings_update_rejects_sandbox_policy_with_permissions() -> Res let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let thread = start_thread(&mut mcp).await?.thread; let request_id = mcp @@ -337,9 +324,8 @@ async fn turn_start_settings_override_emits_thread_settings_updated() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_TIMEOUT) .await?; - timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let thread = start_thread(&mut mcp).await?.thread; timeout( DEFAULT_TIMEOUT, @@ -359,12 +345,8 @@ async fn turn_start_settings_override_emits_thread_settings_updated() -> Result< ..Default::default() }) .await?; - let turn_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response(turn_response)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(turn_request_id)).await??; assert!(!turn.id.is_empty()); let updated = read_thread_settings_updated(&mut mcp).await?; @@ -384,12 +366,8 @@ async fn send_thread_settings_update( params: ThreadSettingsUpdateParams, ) -> Result<()> { let request_id = mcp.send_thread_settings_update_request(params).await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadSettingsUpdateResponse = to_response(response)?; + let _: ThreadSettingsUpdateResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await??; Ok(()) } @@ -404,12 +382,8 @@ async fn start_text_turn(mcp: &mut TestAppServer, thread_id: String) -> Result<( ..Default::default() }) .await?; - let turn_response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response(turn_response)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_TIMEOUT, mcp.read_response(turn_request_id)).await??; assert!(!turn.id.is_empty()); Ok(()) } @@ -421,12 +395,7 @@ async fn start_thread(mcp: &mut TestAppServer) -> Result { ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await? } async fn read_thread_with_turns( @@ -439,26 +408,17 @@ async fn read_thread_with_turns( include_turns: true, }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + timeout(DEFAULT_TIMEOUT, mcp.read_response(request_id)).await? } async fn read_thread_settings_updated( mcp: &mut TestAppServer, ) -> Result { - let notification: JSONRPCNotification = timeout( + timeout( DEFAULT_TIMEOUT, - mcp.read_stream_until_notification_message("thread/settings/updated"), + mcp.read_notification("thread/settings/updated"), ) - .await??; - let params = notification - .params - .context("thread/settings/updated should include params")?; - Ok(serde_json::from_value(params)?) + .await? } async fn received_response_bodies(server: &wiremock::MockServer) -> Result> { @@ -484,13 +444,8 @@ fn service_tier_model_and_tier_id() -> Result<(String, String)> { } fn create_config_toml(codex_home: &std::path::Path, server_uri: &str) -> std::io::Result<()> { - write_mock_responses_config_toml( - codex_home, - server_uri, - &BTreeMap::default(), - /*auto_compact_limit*/ 200_000, - /*requires_openai_auth*/ None, - "mock_provider", - "compact", - ) + MockResponsesConfig::new(server_uri) + .with_root_config("compact_prompt = \"compact\"\nmodel_auto_compact_token_limit = 200000") + .with_provider_config("supports_websockets = false") + .write(codex_home) } diff --git a/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs b/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs index 88628f124454..5d9ef748a751 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_shell_command.rs @@ -1,10 +1,11 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_shell_command_sse_response; use app_test_support::format_with_current_shell_display; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::CommandExecutionApprovalDecision; use codex_app_server_protocol::CommandExecutionOutputDeltaNotification; use codex_app_server_protocol::CommandExecutionRequestApprovalResponse; @@ -12,7 +13,6 @@ use codex_app_server_protocol::CommandExecutionSource; use codex_app_server_protocol::CommandExecutionStatus; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::SortDirection; @@ -33,11 +33,7 @@ use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; use codex_core::shell::default_user_shell; use codex_exec_server::CODEX_EXEC_SERVER_URL_ENV_VAR; -use codex_features::FEATURES; -use codex_features::Feature; use pretty_assertions::assert_eq; -use std::collections::BTreeMap; -use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; @@ -53,44 +49,31 @@ async fn thread_shell_command_history_responses_exclude_persisted_command_execut std::fs::create_dir(&workspace)?; let server = create_mock_responses_server_sequence(vec![]).await; - create_config_toml( - codex_home.as_path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.as_path()) // thread/shellCommand intentionally executes on the app-server host. .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let start_id = mcp - .send_thread_start_request(ThreadStartParams::default()) + let ThreadStartResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadStart { + request_id, + params: ThreadStartParams::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::(start_resp)?; let (shell_command, expected_output) = current_shell_output_command("hello from bang")?; - let shell_id = mcp - .send_thread_shell_command_request(ThreadShellCommandParams { - thread_id: thread.id.clone(), - command: shell_command, + let _: ThreadShellCommandResponse = mcp + .request(|request_id| ClientRequest::ThreadShellCommand { + request_id, + params: ThreadShellCommandParams { + thread_id: thread.id.clone(), + command: shell_command, + }, }) .await?; - let shell_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(shell_id)), - ) - .await??; - let _: ThreadShellCommandResponse = to_response::(shell_resp)?; let started = wait_for_command_execution_started(&mut mcp, /*expected_id*/ None).await?; let ThreadItem::CommandExecution { @@ -133,52 +116,42 @@ async fn thread_shell_command_history_responses_exclude_persisted_command_execut ) .await??; - let read_id = mcp - .send_thread_read_request(ThreadReadParams { - thread_id: thread.id.clone(), - include_turns: true, + let ThreadReadResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadRead { + request_id, + params: ThreadReadParams { + thread_id: thread.id.clone(), + include_turns: true, + }, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; assert_eq!(thread.turns.len(), 1); assert_no_command_executions(&thread.turns[0].items, "thread/read"); - let turns_list_id = mcp - .send_thread_turns_list_request(ThreadTurnsListParams { - thread_id: thread.id.clone(), - cursor: None, - limit: None, - sort_direction: Some(SortDirection::Asc), - items_view: None, + let ThreadTurnsListResponse { data, .. } = mcp + .request(|request_id| ClientRequest::ThreadTurnsList { + request_id, + params: ThreadTurnsListParams { + thread_id: thread.id.clone(), + cursor: None, + limit: None, + sort_direction: Some(SortDirection::Asc), + items_view: None, + }, }) .await?; - let turns_list_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turns_list_id)), - ) - .await??; - let ThreadTurnsListResponse { data, .. } = - to_response::(turns_list_resp)?; assert_eq!(data.len(), 1); assert_no_command_executions(&data[0].items, "thread/turns/list"); - let fork_id = mcp - .send_thread_fork_request(ThreadForkParams { - thread_id: thread.id, - ..Default::default() + let ThreadForkResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadFork { + request_id, + params: ThreadForkParams { + thread_id: thread.id, + ..Default::default() + }, }) .await?; - let fork_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), - ) - .await??; - let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; assert_eq!(thread.turns.len(), 1); assert_no_command_executions(&thread.turns[0].items, "thread/fork"); @@ -191,31 +164,21 @@ async fn thread_shell_command_returns_error_when_local_environment_is_disabled() let codex_home = tmp.path().join("codex_home"); std::fs::create_dir(&codex_home)?; let server = create_mock_responses_server_sequence(vec![]).await; - create_config_toml( - codex_home.as_path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.as_path()) // This test intentionally exercises thread/shellCommand without a local host environment. .without_auto_env() .with_env_overrides(&[(CODEX_EXEC_SERVER_URL_ENV_VAR, Some("none"))]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let start_id = mcp - .send_thread_start_request(ThreadStartParams::default()) + let ThreadStartResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadStart { + request_id, + params: ThreadStartParams::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::(start_resp)?; let shell_id = mcp .send_thread_shell_command_request(ThreadShellCommandParams { thread_id: thread.id, @@ -252,50 +215,39 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { create_final_assistant_message_sse_response("done")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml( - codex_home.as_path(), - &server.uri(), - "untrusted", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.as_path()) // thread/shellCommand intentionally joins the app-server's host-local active turn. .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - - let start_id = mcp - .send_thread_start_request(ThreadStartParams::default()) + let ThreadStartResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadStart { + request_id, + params: ThreadStartParams::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::(start_resp)?; let (shell_command, expected_output) = current_shell_output_command("active turn bang")?; - let turn_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run python".to_string(), - text_elements: Vec::new(), - }], - cwd: Some(workspace.clone()), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run python".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let agent_started = wait_for_command_execution_started(&mut mcp, Some("call-approve")).await?; let ThreadItem::CommandExecution { @@ -319,18 +271,15 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { panic!("expected approval request"); }; - let shell_id = mcp - .send_thread_shell_command_request(ThreadShellCommandParams { - thread_id: thread.id.clone(), - command: shell_command, + let _: ThreadShellCommandResponse = mcp + .request(|request_id| ClientRequest::ThreadShellCommand { + request_id, + params: ThreadShellCommandParams { + thread_id: thread.id.clone(), + command: shell_command, + }, }) .await?; - let shell_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(shell_id)), - ) - .await??; - let _: ThreadShellCommandResponse = to_response::(shell_resp)?; let started = wait_for_command_execution_started_by_source(&mut mcp, CommandExecutionSource::UserShell) @@ -360,28 +309,20 @@ async fn thread_shell_command_uses_existing_active_turn() -> Result<()> { })?, ) .await?; - let _: TurnCompletedNotification = serde_json::from_value( - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), - ) - .await?? - .params - .expect("turn/completed params"), - )?; - - let read_id = mcp - .send_thread_read_request(ThreadReadParams { - thread_id: thread.id, - include_turns: true, - }) - .await?; - let read_resp: JSONRPCResponse = timeout( + let _: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), + mcp.read_notification("turn/completed"), ) .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadRead { + request_id, + params: ThreadReadParams { + thread_id: thread.id, + include_turns: true, + }, + }) + .await?; assert_eq!(thread.turns.len(), 1); assert_no_command_executions(&thread.turns[0].items, "thread/read"); @@ -420,14 +361,7 @@ async fn wait_for_command_execution_started( expected_id: Option<&str>, ) -> Result { loop { - let notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = serde_json::from_value( - notif - .params - .ok_or_else(|| anyhow::anyhow!("missing item/started params"))?, - )?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; let ThreadItem::CommandExecution { id, .. } = &started.item else { continue; }; @@ -457,14 +391,7 @@ async fn wait_for_command_execution_completed( expected_id: Option<&str>, ) -> Result { loop { - let notif = mcp - .read_stream_until_notification_message("item/completed") - .await?; - let completed: ItemCompletedNotification = serde_json::from_value( - notif - .params - .ok_or_else(|| anyhow::anyhow!("missing item/completed params"))?, - )?; + let completed: ItemCompletedNotification = mcp.read_notification("item/completed").await?; let ThreadItem::CommandExecution { id, .. } = &completed.item else { continue; }; @@ -479,58 +406,11 @@ async fn wait_for_command_execution_output_delta( item_id: &str, ) -> Result { loop { - let notif = mcp - .read_stream_until_notification_message("item/commandExecution/outputDelta") + let delta: CommandExecutionOutputDeltaNotification = mcp + .read_notification("item/commandExecution/outputDelta") .await?; - let delta: CommandExecutionOutputDeltaNotification = serde_json::from_value( - notif - .params - .ok_or_else(|| anyhow::anyhow!("missing output delta params"))?, - )?; if delta.item_id == item_id { return Ok(delta); } } } - -fn create_config_toml( - codex_home: &Path, - server_uri: &str, - approval_policy: &str, - feature_flags: &BTreeMap, -) -> std::io::Result<()> { - let feature_entries = feature_flags - .iter() - .map(|(feature, enabled)| { - let key = FEATURES - .iter() - .find(|spec| spec.id == *feature) - .map(|spec| spec.key) - .expect("feature should have a config key"); - format!("{key} = {enabled}") - }) - .collect::>() - .join("\n"); - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "{approval_policy}" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -{feature_entries} - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index cf5cc13332e3..da8ac7be215f 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -79,19 +79,12 @@ async fn start_thread_with_model( model: &str, allow_provider_model_fallback: bool, ) -> Result { - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { - model: Some(model.to_string()), - allow_provider_model_fallback, - ..Default::default() - }) - .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - to_response(response) + mcp.start_thread(ThreadStartParams { + model: Some(model.to_string()), + allow_provider_model_fallback, + ..Default::default() + }) + .await } #[tokio::test] @@ -105,22 +98,15 @@ model = "gpt-5.4-mini" )?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let response = mcp + .start_thread(ThreadStartParams { allow_provider_model_fallback: true, ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let response: ThreadStartResponse = to_response(response)?; assert_eq!(response.model, "openai.gpt-5.6-sol"); Ok(()) @@ -131,9 +117,8 @@ async fn thread_start_warns_for_exec_policy_parse_failure_after_initialize() -> let codex_home = TempDir::new()?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let rules_dir = codex_home.path().join("rules"); std::fs::create_dir_all(&rules_dir)?; @@ -141,15 +126,7 @@ async fn thread_start_warns_for_exec_policy_parse_failure_after_initialize() -> std::fs::write(&rules_path, "prefix_rule(")?; let rules_path = std::fs::canonicalize(rules_path)?; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) - .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadStartResponse = to_response(response)?; + mcp.start_thread(ThreadStartParams::default()).await?; let notification = timeout( DEFAULT_READ_TIMEOUT, @@ -212,9 +189,8 @@ async fn thread_start_does_not_repeat_initialize_exec_policy_warning() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_matching_notification( @@ -224,15 +200,7 @@ async fn thread_start_does_not_repeat_initialize_exec_policy_warning() -> Result ) .await??; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) - .await?; - let response = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let _: ThreadStartResponse = to_response(response)?; + mcp.start_thread(ThreadStartParams::default()).await?; let duplicate_warning = timeout( std::time::Duration::from_millis(250), @@ -260,9 +228,8 @@ async fn thread_start_provider_model_fallback_uses_bedrock_static_catalog() -> R )?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let unsupported_with_fallback = start_thread_with_model( &mut mcp, @@ -301,9 +268,8 @@ async fn thread_start_provider_model_fallback_ignores_dynamic_catalog() -> Resul create_config_toml_without_approval_policy(codex_home.path(), &server.uri())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let response = start_thread_with_model( &mut mcp, @@ -327,9 +293,8 @@ async fn thread_start_creates_thread_and_emits_started() -> Result<()> { // Start server and initialize. let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // Start a v2 thread with an explicit model override. let req_id = mcp @@ -477,37 +442,24 @@ async fn thread_start_history_mode_accepts_legacy_and_paginated() -> Result<()> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { history_mode: Some(ThreadHistoryMode::Legacy), ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(response)?; assert_eq!(thread.history_mode, ThreadHistoryMode::Legacy); - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { history_mode: Some(ThreadHistoryMode::Paginated), ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(response)?; assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); Ok(()) @@ -526,9 +478,8 @@ async fn thread_start_accepts_absolute_runtime_workspace_roots() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_thread_start_request(ThreadStartParams { @@ -539,17 +490,12 @@ async fn thread_start_accepts_absolute_runtime_workspace_roots() -> Result<()> { }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; let ThreadStartResponse { cwd: response_cwd, runtime_workspace_roots, sandbox, .. - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(req_id)).await??; assert_eq!(response_cwd, cwd.abs()); assert_eq!(runtime_workspace_roots, vec![extra_root.abs()]); @@ -578,16 +524,11 @@ async fn thread_start_accepts_absolute_runtime_workspace_roots() -> Result<()> { ..Default::default() }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; let ThreadStartResponse { runtime_workspace_roots, sandbox, .. - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(req_id)).await??; assert_eq!(runtime_workspace_roots, vec![environment_root.abs()]); #[cfg(windows)] let _ = sandbox; @@ -620,9 +561,8 @@ async fn thread_start_excludes_profile_workspace_roots_from_runtime_workspace_ro let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_thread_start_request(ThreadStartParams { @@ -631,15 +571,10 @@ async fn thread_start_excludes_profile_workspace_roots_from_runtime_workspace_ro }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; let ThreadStartResponse { runtime_workspace_roots, .. - } = to_response::(resp)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(req_id)).await??; assert_eq!( runtime_workspace_roots, @@ -662,9 +597,8 @@ async fn thread_start_rejects_unknown_environment_as_invalid_request() -> Result let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_start_request(ThreadStartParams { @@ -701,9 +635,8 @@ async fn thread_start_rejects_relative_environment_cwd_as_invalid_request() -> R let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let environment_id = mcp.auto_env_params()?.environment_id; let request_id = mcp @@ -749,9 +682,8 @@ async fn thread_start_response_includes_loaded_instruction_sources() -> Result<( .with_codex_home(codex_home.path()) // TODO(anp): Move the instruction-source fixture into the auto environment cwd. .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_start_request(ThreadStartParams { @@ -759,15 +691,10 @@ async fn thread_start_response_includes_loaded_instruction_sources() -> Result<( ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ThreadStartResponse { instruction_sources, .. - } = to_response::(response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let instruction_sources = instruction_sources .into_iter() @@ -801,9 +728,8 @@ async fn thread_start_response_excludes_empty_project_instruction_source() -> Re .with_codex_home(codex_home.path()) // TODO(anp): Move the instruction-source fixture into the auto environment cwd. .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_start_request(ThreadStartParams { @@ -811,15 +737,10 @@ async fn thread_start_response_excludes_empty_project_instruction_source() -> Re ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ThreadStartResponse { instruction_sources, .. - } = to_response::(response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let instruction_sources = instruction_sources .into_iter() @@ -847,9 +768,8 @@ async fn thread_start_without_selected_environment_includes_only_global_instruct let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_thread_start_request(ThreadStartParams { @@ -858,16 +778,11 @@ async fn thread_start_without_selected_environment_includes_only_global_instruct ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; let ThreadStartResponse { thread, instruction_sources, .. - } = to_response::(response)?; + } = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; assert_eq!( instruction_sources @@ -941,23 +856,16 @@ async fn thread_start_tracks_thread_initialized_analytics() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let req_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { thread_source: Some(ThreadSource::User), service_name: Some("codex_work_desktop".to_string()), ..Default::default() }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(resp)?; let payload = wait_for_analytics_payload(&server, DEFAULT_READ_TIMEOUT).await?; assert_eq!(payload["events"].as_array().expect("events array").len(), 1); @@ -994,26 +902,18 @@ model_reasoning_effort = "high" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let req_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { + reasoning_effort, .. + } = mcp + .start_thread(ThreadStartParams { cwd: Some(workspace.path().to_string_lossy().into_owned()), ..Default::default() }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let ThreadStartResponse { - reasoning_effort, .. - } = to_response::(resp)?; - assert_eq!(reasoning_effort, Some(ReasoningEffort::High)); Ok(()) } @@ -1027,25 +927,17 @@ async fn thread_start_drops_unsupported_service_tier_id() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let service_tier_id = "experimental-tier-id".to_string(); - let req_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { service_tier, .. } = mcp + .start_thread(ThreadStartParams { service_tier: Some(Some(service_tier_id.clone())), ..Default::default() }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let ThreadStartResponse { service_tier, .. } = to_response::(resp)?; - // Unsupported catalog ids are dropped at session config time instead of echoed back. assert_eq!(service_tier, None); Ok(()) @@ -1060,24 +952,16 @@ async fn thread_start_accepts_default_service_tier() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let req_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { service_tier, .. } = mcp + .start_thread(ThreadStartParams { service_tier: Some(Some(SERVICE_TIER_DEFAULT_REQUEST_VALUE.to_string())), ..Default::default() }) .await?; - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let ThreadStartResponse { service_tier, .. } = to_response::(resp)?; - assert_eq!( service_tier, Some(SERVICE_TIER_DEFAULT_REQUEST_VALUE.to_string()) @@ -1094,23 +978,15 @@ async fn thread_start_accepts_metrics_service_name() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let req_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { service_name: Some("my_app_server_client".to_string()), ..Default::default() }) .await?; - - let resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(resp)?; assert!(!thread.id.is_empty(), "thread id should not be empty"); Ok(()) @@ -1124,9 +1000,8 @@ async fn thread_start_ephemeral_remains_pathless() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -1173,9 +1048,8 @@ async fn thread_start_fails_when_required_mcp_server_fails_to_initialize() -> Re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) @@ -1212,21 +1086,10 @@ async fn thread_start_emits_mcp_server_status_updated_notifications() -> Result< let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let req_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) - .await?; - - let start_response: ThreadStartResponse = to_response( - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(req_id)), - ) - .await??, - )?; + let start_response = mcp.start_thread(ThreadStartParams::default()).await?; let starting = timeout( DEFAULT_READ_TIMEOUT, @@ -1331,9 +1194,8 @@ async fn thread_start_does_not_wait_for_optional_http_mcp_auth_discovery() -> Re let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) @@ -1405,9 +1267,8 @@ async fn thread_start_surfaces_cloud_config_bundle_load_errors() -> Result<()> { Some(refresh_token_url.as_str()), ), ]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let req_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams::default()) @@ -1458,39 +1319,26 @@ model_reasoning_effort = "high" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let first_request = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { - cwd: Some(workspace.path().display().to_string()), - sandbox: Some(SandboxMode::WorkspaceWrite), - ..Default::default() - }) - .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_request)), - ) - .await??; + mcp.start_thread(ThreadStartParams { + cwd: Some(workspace.path().display().to_string()), + sandbox: Some(SandboxMode::WorkspaceWrite), + ..Default::default() + }) + .await?; - let second_request = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { - cwd: Some(workspace.path().display().to_string()), - ..Default::default() - }) - .await?; - let second_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_request)), - ) - .await??; let ThreadStartResponse { approval_policy, reasoning_effort, .. - } = to_response::(second_response)?; + } = mcp + .start_thread(ThreadStartParams { + cwd: Some(workspace.path().display().to_string()), + ..Default::default() + }) + .await?; assert_eq!(approval_policy, AskForApproval::OnRequest); assert_eq!(reasoning_effort, Some(ReasoningEffort::High)); @@ -1521,22 +1369,15 @@ async fn thread_start_with_nested_git_cwd_trusts_repo_root() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { - cwd: Some(nested.display().to_string()), - sandbox: Some(SandboxMode::WorkspaceWrite), - ..Default::default() - }) - .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + mcp.start_thread(ThreadStartParams { + cwd: Some(nested.display().to_string()), + sandbox: Some(SandboxMode::WorkspaceWrite), + ..Default::default() + }) + .await?; let config_toml = std::fs::read_to_string(codex_home.path().join("config.toml"))?; let nested_abs = nested.abs(); @@ -1562,21 +1403,14 @@ async fn thread_start_with_read_only_sandbox_does_not_persist_project_trust() -> let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { - cwd: Some(workspace.path().display().to_string()), - ..Default::default() - }) - .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + mcp.start_thread(ThreadStartParams { + cwd: Some(workspace.path().display().to_string()), + ..Default::default() + }) + .await?; let config_toml = std::fs::read_to_string(codex_home.path().join("config.toml"))?; assert!(!config_toml.contains("trust_level = \"trusted\"")); @@ -1603,22 +1437,15 @@ async fn thread_start_preserves_untrusted_project_trust() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { - cwd: Some(workspace.path().display().to_string()), - sandbox: Some(SandboxMode::WorkspaceWrite), - ..Default::default() - }) - .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; + mcp.start_thread(ThreadStartParams { + cwd: Some(workspace.path().display().to_string()), + sandbox: Some(SandboxMode::WorkspaceWrite), + ..Default::default() + }) + .await?; let config_after = std::fs::read_to_string(&config_path)?; assert_eq!(config_after, config_before); @@ -1647,27 +1474,20 @@ model_reasoning_effort = "high" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { + approval_policy, + reasoning_effort, + .. + } = mcp + .start_thread(ThreadStartParams { cwd: Some(workspace.path().display().to_string()), sandbox: Some(SandboxMode::WorkspaceWrite), ..Default::default() }) .await?; - let response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let ThreadStartResponse { - approval_policy, - reasoning_effort, - .. - } = to_response::(response)?; assert_eq!(approval_policy, AskForApproval::OnRequest); assert_eq!(reasoning_effort, Some(ReasoningEffort::High)); @@ -1682,26 +1502,21 @@ fn create_config_toml_without_approval_policy( codex_home: &Path, server_uri: &str, ) -> std::io::Result<()> { - create_config_toml_with_optional_approval_policy( - codex_home, server_uri, /*approval_policy*/ None, - ) + create_config_toml(codex_home, server_uri, "sandbox_mode = \"read-only\"", "") } -fn create_config_toml_with_optional_approval_policy( +fn create_config_toml( codex_home: &Path, server_uri: &str, - approval_policy: Option<&str>, + top_level_config: &str, + additional_tables: &str, ) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - let approval_policy = approval_policy - .map(|policy| format!("approval_policy = \"{policy}\"\n")) - .unwrap_or_default(); std::fs::write( - config_toml, + codex_home.join("config.toml"), format!( r#" model = "mock-model" -{approval_policy}sandbox_mode = "read-only" +{top_level_config} model_provider = "mock_provider" @@ -1711,6 +1526,7 @@ base_url = "{server_uri}/v1" wire_api = "responses" request_max_retries = 0 stream_max_retries = 0 +{additional_tables} "# ), ) @@ -1721,27 +1537,17 @@ fn create_config_toml_with_profile_workspace_root( server_uri: &str, profile_root: &Path, ) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); let profile_root_key = profile_root .display() .to_string() .replace('\\', "\\\\") .replace('"', "\\\""); - std::fs::write( - config_toml, - format!( + create_config_toml( + codex_home, + server_uri, + "default_permissions = \"dev\"", + &format!( r#" -model = "mock-model" -default_permissions = "dev" -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - [permissions.dev.workspace_roots] "{profile_root_key}" = true @@ -1757,26 +1563,13 @@ fn create_config_toml_with_chatgpt_base_url( server_uri: &str, chatgpt_base_url: &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" -chatgpt_base_url = "{chatgpt_base_url}" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# + create_config_toml( + codex_home, + server_uri, + &format!( + "approval_policy = \"never\"\nsandbox_mode = \"read-only\"\nchatgpt_base_url = \"{chatgpt_base_url}\"" ), + "", ) } @@ -1784,24 +1577,12 @@ fn create_config_toml_with_required_broken_mcp( codex_home: &Path, server_uri: &str, ) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( + create_config_toml( + codex_home, + server_uri, + "approval_policy = \"never\"\nsandbox_mode = \"read-only\"", + &format!( r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - [mcp_servers.required_broken] {required_broken_transport} required = true @@ -1815,24 +1596,12 @@ fn create_config_toml_with_optional_broken_mcp( codex_home: &Path, server_uri: &str, ) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( + create_config_toml( + codex_home, + server_uri, + "approval_policy = \"never\"\nsandbox_mode = \"read-only\"", + &format!( r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - [mcp_servers.optional_broken] {optional_broken_transport} "#, @@ -1846,28 +1615,16 @@ fn create_config_toml_with_optional_http_mcp( server_uri: &str, mcp_uri: &str, ) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( + create_config_toml( + codex_home, + server_uri, + "approval_policy = \"never\"\nsandbox_mode = \"read-only\"", + &format!( r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 - [mcp_servers.optional_http] url = "{mcp_uri}" startup_timeout_sec = 60 -"# +"#, ), ) } diff --git a/codex-rs/app-server/tests/suite/v2/thread_status.rs b/codex-rs/app-server/tests/suite/v2/thread_status.rs index 6c54740443b0..999b285f97fb 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_status.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_status.rs @@ -1,14 +1,13 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; -use app_test_support::to_response; use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::ThreadStatus; @@ -16,6 +15,7 @@ use codex_app_server_protocol::ThreadStatusChangedNotification; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_features::Feature; use tempfile::TempDir; use tokio::time::timeout; @@ -26,46 +26,39 @@ async fn thread_status_changed_emits_runtime_updates() -> Result<()> { let codex_home = TempDir::new()?; let responses = vec![create_final_assistant_message_sse_response("done")?]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .enable_feature(Feature::CollaborationModes) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("RUST_LOG", Some("info"))]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; - - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "collect status updates".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model".to_string()), - ..Default::default() + + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "collect status updates".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + ..Default::default() + }, }) .await?; - let turn_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_start_resp)?; let mut saw_active_running = false; let mut saw_idle_after_turn = false; @@ -137,7 +130,10 @@ async fn thread_status_changed_can_be_opted_out() -> Result<()> { let codex_home = TempDir::new()?; let responses = vec![create_final_assistant_message_sse_response("done")?]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .enable_feature(Feature::CollaborationModes) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -164,37 +160,28 @@ async fn thread_status_changed_can_be_opted_out() -> Result<()> { anyhow::bail!("expected initialize response, got {message:?}"); }; - let thread_start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response(thread_start_resp)?; - - let turn_start_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run once".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model".to_string()), - ..Default::default() + + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run once".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model".to_string()), + ..Default::default() + }, }) .await?; - let turn_start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_start_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -223,29 +210,3 @@ async fn thread_status_changed_can_be_opted_out() -> Result<()> { Ok(()) } - -fn create_config_toml(codex_home: &std::path::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 = "untrusted" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -collaboration_modes = true - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index 8715d7154f2b..433fe88afd6c 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::to_response; @@ -57,13 +58,12 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -71,12 +71,8 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let rollout_path = thread.path.clone().expect("thread path"); @@ -91,12 +87,8 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result ..Default::default() }) .await?; - let turn_start_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_start_id)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_start_response)?; + let _: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_start_id)).await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -114,12 +106,8 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result thread_id: thread.id.clone(), }) .await?; - let archive_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(archive_id)), - ) - .await??; - let _: ThreadArchiveResponse = to_response::(archive_resp)?; + let _: ThreadArchiveResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(archive_id)).await??; let archived_path = find_archived_thread_path_by_id_str( codex_home.path(), @@ -158,16 +146,11 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result let ThreadUnarchiveResponse { thread: unarchived_thread, } = to_response::(unarchive_resp)?; - let unarchive_notification = timeout( + let unarchived_notification: ThreadUnarchivedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/unarchived"), + mcp.read_notification("thread/unarchived"), ) .await??; - let unarchived_notification: ThreadUnarchivedNotification = serde_json::from_value( - unarchive_notification - .params - .expect("thread/unarchived notification params"), - )?; assert_eq!(unarchived_notification.thread_id, thread.id); assert!( unarchived_thread.updated_at > old_timestamp, @@ -204,7 +187,11 @@ async fn thread_unarchive_moves_rollout_back_into_sessions_directory() -> Result async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { let codex_home = TempDir::new()?; let store_id = Uuid::new_v4().to_string(); - create_config_toml_with_in_memory_thread_store(codex_home.path(), &store_id)?; + MockResponsesConfig::new("http://127.0.0.1:1") + .with_root_config(&format!( + r#"experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }}"# + )) + .write(codex_home.path())?; let store = InMemoryThreadStore::for_id(store_id.clone()); let _in_memory_store = InMemoryThreadStoreId { store_id }; let thread_id = ThreadId::from_string("00000000-0000-4000-8000-000000000126")?; @@ -301,11 +288,6 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { Ok(()) } -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, config_contents(server_uri)) -} - struct InMemoryThreadStoreId { store_id: String, } @@ -316,50 +298,6 @@ impl Drop for InMemoryThreadStoreId { } } -fn create_config_toml_with_in_memory_thread_store( - codex_home: &Path, - store_id: &str, -) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -experimental_thread_store = {{ type = "in_memory", id = "{store_id}" }} - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "http://127.0.0.1:1/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - -fn config_contents(server_uri: &str) -> String { - format!( - r#"model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ) -} - fn assert_paths_match_on_disk(actual: &Path, expected: &Path) -> std::io::Result<()> { let actual = actual.canonicalize()?; let expected = expected.canonicalize()?; diff --git a/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs b/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs index 2dc12319eca4..da456ac820c8 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs @@ -1,15 +1,14 @@ use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_repeating_assistant; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::DynamicToolCallOutputContentItem; use codex_app_server_protocol::DynamicToolCallParams; use codex_app_server_protocol::DynamicToolCallResponse; use codex_app_server_protocol::DynamicToolFunctionSpec; use codex_app_server_protocol::DynamicToolSpec; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadLoadedListParams; @@ -40,39 +39,31 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs async fn thread_unsubscribe_keeps_thread_loaded_until_idle_timeout() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_sandbox_mode("danger-full-access") + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let thread_id = thread.id; - let unsubscribe_id = mcp - .send_thread_unsubscribe_request(ThreadUnsubscribeParams { - thread_id: thread_id.clone(), + let unsubscribe: ThreadUnsubscribeResponse = mcp + .request(|request_id| ClientRequest::ThreadUnsubscribe { + request_id, + params: ThreadUnsubscribeParams { + thread_id: thread_id.clone(), + }, }) .await?; - let unsubscribe_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(unsubscribe_id)), - ) - .await??; - let unsubscribe = to_response::(unsubscribe_resp)?; assert_eq!(unsubscribe.status, ThreadUnsubscribeStatus::Unsubscribed); assert!( @@ -84,16 +75,12 @@ async fn thread_unsubscribe_keeps_thread_loaded_until_idle_timeout() -> Result<( .is_err() ); - let list_id = mcp - .send_thread_loaded_list_request(ThreadLoadedListParams::default()) + let ThreadLoadedListResponse { data, next_cursor } = mcp + .request(|request_id| ClientRequest::ThreadLoadedList { + request_id, + params: ThreadLoadedListParams::default(), + }) .await?; - let list_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_id)), - ) - .await??; - let ThreadLoadedListResponse { data, next_cursor } = - to_response::(list_resp)?; assert_eq!(data, vec![thread_id]); assert_eq!(next_cursor, None); @@ -132,16 +119,17 @@ async fn thread_unsubscribe_during_turn_keeps_turn_running() -> Result<()> { .await; let first_response_completed = completions.remove(0); let final_response_completed = completions.remove(0); - create_config_toml(&codex_home, server.uri())?; + MockResponsesConfig::new(server.uri()) + .with_sandbox_mode("danger-full-access") + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), dynamic_tools: Some(vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { name: tool_name.to_string(), @@ -156,31 +144,22 @@ async fn thread_unsubscribe_during_turn_keeps_turn_running() -> Result<()> { ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let thread_id = thread.id; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread_id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run deterministic tool".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread_id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run deterministic tool".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -217,17 +196,14 @@ async fn thread_unsubscribe_during_turn_keeps_turn_running() -> Result<()> { } ); - let unsubscribe_id = mcp - .send_thread_unsubscribe_request(ThreadUnsubscribeParams { - thread_id: thread_id.clone(), + let unsubscribe: ThreadUnsubscribeResponse = mcp + .request(|request_id| ClientRequest::ThreadUnsubscribe { + request_id, + params: ThreadUnsubscribeParams { + thread_id: thread_id.clone(), + }, }) .await?; - let unsubscribe_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(unsubscribe_id)), - ) - .await??; - let unsubscribe = to_response::(unsubscribe_resp)?; assert_eq!(unsubscribe.status, ThreadUnsubscribeStatus::Unsubscribed); let closed_while_tool_call_blocked = timeout( @@ -266,76 +242,62 @@ async fn thread_unsubscribe_preserves_cached_status_before_idle_unload() -> Resu ) .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_sandbox_mode("danger-full-access") + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let thread_id = thread.id; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread_id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "fail this turn".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread_id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "fail this turn".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _: TurnStartResponse = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("error"), ) .await??; - let read_id = mcp - .send_thread_read_request(ThreadReadParams { - thread_id: thread_id.clone(), - include_turns: false, + let ThreadReadResponse { thread, .. } = mcp + .request(|request_id| ClientRequest::ThreadRead { + request_id, + params: ThreadReadParams { + thread_id: thread_id.clone(), + include_turns: false, + }, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(read_id)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; assert_eq!(thread.status, ThreadStatus::SystemError); - let unsubscribe_id = mcp - .send_thread_unsubscribe_request(ThreadUnsubscribeParams { - thread_id: thread_id.clone(), + let unsubscribe: ThreadUnsubscribeResponse = mcp + .request(|request_id| ClientRequest::ThreadUnsubscribe { + request_id, + params: ThreadUnsubscribeParams { + thread_id: thread_id.clone(), + }, }) .await?; - let unsubscribe_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(unsubscribe_id)), - ) - .await??; - let unsubscribe = to_response::(unsubscribe_resp)?; assert_eq!(unsubscribe.status, ThreadUnsubscribeStatus::Unsubscribed); assert!( timeout( @@ -346,19 +308,16 @@ async fn thread_unsubscribe_preserves_cached_status_before_idle_unload() -> Resu .is_err() ); - let resume_id = mcp - .send_thread_resume_request(ThreadResumeParams { - thread_id, - cwd: Some(codex_home.path().to_string_lossy().to_string()), - ..Default::default() + let resume: ThreadResumeResponse = mcp + .request(|request_id| ClientRequest::ThreadResume { + request_id, + params: ThreadResumeParams { + thread_id, + cwd: Some(codex_home.path().to_string_lossy().to_string()), + ..Default::default() + }, }) .await?; - let resume_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), - ) - .await??; - let resume: ThreadResumeResponse = to_response::(resume_resp)?; assert_eq!(resume.thread.status, ThreadStatus::SystemError); Ok(()) @@ -368,53 +327,42 @@ async fn thread_unsubscribe_preserves_cached_status_before_idle_unload() -> Resu async fn thread_unsubscribe_reports_not_subscribed_before_idle_unload() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_sandbox_mode("danger-full-access") + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let thread_id = thread.id; - let first_unsubscribe_id = mcp - .send_thread_unsubscribe_request(ThreadUnsubscribeParams { - thread_id: thread_id.clone(), + let first_unsubscribe: ThreadUnsubscribeResponse = mcp + .request(|request_id| ClientRequest::ThreadUnsubscribe { + request_id, + params: ThreadUnsubscribeParams { + thread_id: thread_id.clone(), + }, }) .await?; - let first_unsubscribe_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_unsubscribe_id)), - ) - .await??; - let first_unsubscribe = to_response::(first_unsubscribe_resp)?; assert_eq!( first_unsubscribe.status, ThreadUnsubscribeStatus::Unsubscribed ); - let second_unsubscribe_id = mcp - .send_thread_unsubscribe_request(ThreadUnsubscribeParams { thread_id }) + let second_unsubscribe: ThreadUnsubscribeResponse = mcp + .request(|request_id| ClientRequest::ThreadUnsubscribe { + request_id, + params: ThreadUnsubscribeParams { thread_id }, + }) .await?; - let second_unsubscribe_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_unsubscribe_id)), - ) - .await??; - let second_unsubscribe = to_response::(second_unsubscribe_resp)?; assert_eq!( second_unsubscribe.status, ThreadUnsubscribeStatus::NotSubscribed @@ -440,26 +388,3 @@ async fn wait_for_dynamic_tool_started( } } } - -fn create_config_toml(codex_home: &std::path::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 = "danger-full-access" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs b/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs index 2d47dd8ceb0f..6caf01d203d6 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs @@ -1,15 +1,14 @@ #![cfg(unix)] use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::create_shell_command_sse_response; -use app_test_support::to_response; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ServerRequestResolvedNotification; @@ -62,47 +61,40 @@ async fn turn_interrupt_aborts_running_turn() -> Result<()> { "call_sleep", )?]) .await; - create_config_toml(&codex_home, &server.uri(), "never", "workspace-write")?; + MockResponsesConfig::new(&server.uri()) + .with_sandbox_mode("workspace-write") + .with_root_config(r#"approvals_reviewer = "user""#) + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // Start a v2 thread and capture its id. - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; // Start a turn that triggers a long-running command. - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run sleep".to_string(), - text_elements: Vec::new(), - }], - cwd: Some(working_directory.clone()), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run sleep".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(working_directory.clone()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let turn_id = turn.id.clone(); // Give the command a brief moment to start. @@ -110,29 +102,21 @@ async fn turn_interrupt_aborts_running_turn() -> Result<()> { let thread_id = thread.id.clone(); // Interrupt the in-progress turn by id (v2 API). - let interrupt_id = mcp - .send_turn_interrupt_request(TurnInterruptParams { - thread_id: thread_id.clone(), - turn_id: turn_id.clone(), + let _: TurnInterruptResponse = mcp + .request(|request_id| ClientRequest::TurnInterrupt { + request_id, + params: TurnInterruptParams { + thread_id: thread_id.clone(), + turn_id: turn_id.clone(), + }, }) .await?; - let interrupt_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(interrupt_id)), - ) - .await??; - let _resp: TurnInterruptResponse = to_response::(interrupt_resp)?; - let completed_notif: JSONRPCNotification = timeout( + let completed: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed: TurnCompletedNotification = serde_json::from_value( - completed_notif - .params - .expect("turn/completed params must be present"), - )?; assert_eq!(completed.thread_id, thread_id); assert_eq!(completed.turn.status, TurnStatus::Interrupted); @@ -149,55 +133,43 @@ async fn turn_interrupt_rejects_completed_turn() -> Result<()> { create_final_assistant_message_sse_response("done")?, ]) .await; - create_config_toml(&codex_home, &server.uri(), "never", "workspace-write")?; + MockResponsesConfig::new(&server.uri()) + .with_sandbox_mode("workspace-write") + .with_root_config(r#"approvals_reviewer = "user""#) + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "say done".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "say done".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; - let completed_notif: JSONRPCNotification = timeout( + let completed: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed: TurnCompletedNotification = serde_json::from_value( - completed_notif - .params - .expect("turn/completed params must be present"), - )?; assert_eq!(completed.thread_id, thread.id); assert_eq!(completed.turn.id, turn.id); assert_eq!(completed.turn.status, TurnStatus::Completed); @@ -253,46 +225,39 @@ async fn turn_interrupt_resolves_pending_command_approval_request() -> Result<() "call_sleep_approval", )?]) .await; - create_config_toml(&codex_home, &server.uri(), "untrusted", "read-only")?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .with_root_config(r#"approvals_reviewer = "user""#) + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run python".to_string(), - text_elements: Vec::new(), - }], - cwd: Some(working_directory), - approval_policy: Some(codex_app_server_protocol::AskForApproval::UnlessTrusted), - ..Default::default() + + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run python".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(working_directory), + approval_policy: Some(codex_app_server_protocol::AskForApproval::UnlessTrusted), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let request = timeout( DEFAULT_READ_TIMEOUT, @@ -306,75 +271,31 @@ async fn turn_interrupt_resolves_pending_command_approval_request() -> Result<() assert_eq!(params.thread_id, thread.id); assert_eq!(params.turn_id, turn.id); - let interrupt_id = mcp - .send_turn_interrupt_request(TurnInterruptParams { - thread_id: thread.id.clone(), - turn_id: turn.id.clone(), + let _: TurnInterruptResponse = mcp + .request(|request_id| ClientRequest::TurnInterrupt { + request_id, + params: TurnInterruptParams { + thread_id: thread.id.clone(), + turn_id: turn.id.clone(), + }, }) .await?; - let interrupt_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(interrupt_id)), - ) - .await??; - let _resp: TurnInterruptResponse = to_response::(interrupt_resp)?; - let resolved_notification = timeout( + let resolved: ServerRequestResolvedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("serverRequest/resolved"), + mcp.read_notification("serverRequest/resolved"), ) .await??; - let resolved: ServerRequestResolvedNotification = serde_json::from_value( - resolved_notification - .params - .clone() - .expect("serverRequest/resolved params must be present"), - )?; assert_eq!(resolved.thread_id, thread.id); assert_eq!(resolved.request_id, request_id); - let completed_notif: JSONRPCNotification = timeout( + let completed: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed: TurnCompletedNotification = serde_json::from_value( - completed_notif - .params - .expect("turn/completed params must be present"), - )?; assert_eq!(completed.thread_id, thread.id); assert_eq!(completed.turn.status, TurnStatus::Interrupted); Ok(()) } - -// Helper to create a config.toml pointing at the mock model server. -fn create_config_toml( - codex_home: &std::path::Path, - server_uri: &str, - approval_policy: &str, - sandbox_mode: &str, -) -> std::io::Result<()> { - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "mock-model" -approval_policy = "{approval_policy}" -approvals_reviewer = "user" -sandbox_mode = "{sandbox_mode}" - -model_provider = "mock_provider" - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index ae8cca7f3493..6fbc1fc24384 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -1,5 +1,6 @@ use anyhow::Context; use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_apply_patch_sse_response; use app_test_support::create_exec_command_sse_response; @@ -10,7 +11,6 @@ use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::create_request_user_input_sse_response; use app_test_support::create_shell_command_sse_response; use app_test_support::format_with_current_shell_display; -use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; use app_test_support::write_models_cache; use codex_app_server::INPUT_TOO_LARGE_ERROR_CODE; @@ -19,6 +19,7 @@ use codex_app_server_protocol::AdditionalContextEntry; use codex_app_server_protocol::AdditionalContextKind; use codex_app_server_protocol::ByteRange; use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::CollabAgentStatus; use codex_app_server_protocol::CollabAgentTool; use codex_app_server_protocol::CollabAgentToolCallStatus; @@ -32,8 +33,6 @@ use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PatchApplyStatus; use codex_app_server_protocol::PatchChangeKind; use codex_app_server_protocol::RawResponseCompletedNotification; @@ -65,7 +64,6 @@ use codex_app_server_protocol::UserInput as V2UserInput; use codex_app_server_protocol::WarningNotification; use codex_core::test_support::all_model_presets; use codex_exec_server::LOCAL_ENVIRONMENT_ID; -use codex_features::FEATURES; use codex_features::Feature; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; @@ -86,7 +84,6 @@ use core_test_support::skip_if_wine_exec; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; -use std::collections::BTreeMap; use std::collections::HashMap; use std::path::Path; use tempfile::TempDir; @@ -127,52 +124,37 @@ async fn run_local_image_turn(detail: Option) -> Result> let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let image_path = codex_home.path().join("image.png"); std::fs::write(&image_path, TINY_PNG_BYTES)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::LocalImage { - path: image_path, - detail, - }], - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::LocalImage { + path: image_path, + detail, + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; assert!(!turn.id.is_empty()); timeout( @@ -227,70 +209,45 @@ async fn turn_start_with_empty_input_runs_model_request() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), thread_source: Some(ThreadSource::User), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: Vec::new(), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: Vec::new(), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; assert!(!turn.id.is_empty()); - let started_notif: JSONRPCNotification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/started"), - ) - .await??; let started: TurnStartedNotification = - serde_json::from_value(started_notif.params.expect("params must be present"))?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("turn/started")).await??; assert_eq!(started.thread_id, thread.id); assert_eq!(started.turn.id, turn.id); assert_eq!(started.turn.status, TurnStatus::InProgress); - let completed_notif: JSONRPCNotification = timeout( + let completed: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed: TurnCompletedNotification = serde_json::from_value( - completed_notif - .params - .expect("turn/completed params must be present"), - )?; assert_eq!(completed.thread_id, thread.id); assert_eq!(completed.turn.id, turn.id); assert_eq!(completed.turn.status, TurnStatus::Completed); @@ -337,55 +294,41 @@ async fn turn_start_additional_context_flows_to_model_input() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "inspect tab".to_string(), - text_elements: Vec::new(), - }], - additional_context: Some(HashMap::from([( - "custom_source".to_string(), - AdditionalContextEntry { - value: "source value".to_string(), - kind: AdditionalContextKind::Untrusted, - }, - )])), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "inspect tab".to_string(), + text_elements: Vec::new(), + }], + additional_context: Some(HashMap::from([( + "custom_source".to_string(), + AdditionalContextEntry { + value: "source value".to_string(), + kind: AdditionalContextKind::Untrusted, + }, + )])), + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -417,12 +360,9 @@ async fn turn_start_sends_originator_header() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) @@ -438,36 +378,28 @@ async fn turn_start_sends_originator_header() -> Result<()> { ) .await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), thread_source: Some(ThreadSource::User), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, @@ -497,62 +429,46 @@ async fn turn_start_emits_user_message_item_with_text_elements() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), thread_source: Some(ThreadSource::User), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let text_elements = vec![TextElement::new( ByteRange { start: 0, end: 5 }, Some("".to_string()), )]; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: Some("client-message-1".to_string()), - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: text_elements.clone(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: Some("client-message-1".to_string()), + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: text_elements.clone(), + }], + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; let user_message_item = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let notification = mcp - .read_stream_until_notification_message("item/started") - .await?; - let params = notification.params.expect("item/started params"); let item_started: ItemStartedNotification = - serde_json::from_value(params).expect("deserialize item/started notification"); + mcp.read_notification("item/started").await?; if let ThreadItem::UserMessage { .. } = item_started.item { return Ok::(item_started.item); } @@ -591,12 +507,9 @@ async fn turn_start_emits_thread_scoped_warning_notification_for_trimmed_skills( let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; write_models_cache(codex_home.path())?; let cache_path = codex_home.path().join("models_cache.json"); let mut cache: serde_json::Value = @@ -629,45 +542,28 @@ async fn turn_start_emits_thread_scoped_warning_notification_for_trimmed_skills( ("HOME", Some(isolated_home.as_ref())), ("USERPROFILE", Some(isolated_home.as_ref())), ]) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams::default()) - .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = mcp.start_thread(ThreadStartParams::default()).await?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("warning"), - ) - .await??; - let params = notification.params.expect("warning params"); let warning: WarningNotification = - serde_json::from_value(params).expect("deserialize warning notification"); + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("warning")).await??; assert_eq!(warning.thread_id.as_deref(), Some(thread.id.as_str())); assert_eq!( warning.message, @@ -710,12 +606,7 @@ async fn turn_start_sends_service_tier_id_to_model_request() -> Result<()> { let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; write_models_cache(codex_home.path())?; let service_tier_model = all_model_presets() .iter() @@ -725,39 +616,30 @@ async fn turn_start_sends_service_tier_id_to_model_request() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some(service_tier_model.id.clone()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - service_tier: Some(Some(service_tier_id.clone())), - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + service_tier: Some(Some(service_tier_id.clone())), + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -795,49 +677,34 @@ async fn turn_start_emits_raw_response_completed_with_upstream_usage() -> Result responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; write_models_cache(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { experimental_raw_events: true, ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let notification = timeout( DEFAULT_READ_TIMEOUT, @@ -882,21 +749,15 @@ async fn thread_start_omits_empty_instruction_overrides_from_model_request() -> let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { // TODO(aibrahim): Replace empty string instruction overrides with explicit tri-state // app-server semantics: omitted, explicitly none, or explicit value. config: Some(HashMap::from([( @@ -908,29 +769,21 @@ async fn thread_start_omits_empty_instruction_overrides_from_model_request() -> ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -997,46 +850,36 @@ async fn turn_start_tracks_thread_originator_in_analytics() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), thread_source: Some(ThreadSource::User), service_name: Some("codex_work_desktop".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Image { - url: TINY_PNG_DATA_URL.to_string(), - detail: None, - }], - responsesapi_client_metadata: Some(HashMap::from([( - "workspace_kind".to_string(), - "projectless".to_string(), - )])), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Image { + url: TINY_PNG_DATA_URL.to_string(), + detail: None, + }], + responsesapi_client_metadata: Some(HashMap::from([( + "workspace_kind".to_string(), + "projectless".to_string(), + )])), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -1126,47 +969,38 @@ async fn turn_profile_tracks_blocking_tool_and_follow_up_sampling() -> Result<() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "ask something".to_string(), - text_elements: Vec::new(), - }], - collaboration_mode: Some(CollaborationMode { - mode: ModeKind::Plan, - settings: Settings { - model: "mock-model".to_string(), - reasoning_effort: Some(ReasoningEffort::Medium), - developer_instructions: None, - }, - }), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "ask something".to_string(), + text_elements: Vec::new(), + }], + collaboration_mode: Some(CollaborationMode { + mode: ModeKind::Plan, + settings: Settings { + model: "mock-model".to_string(), + reasoning_effort: Some(ReasoningEffort::Medium), + developer_instructions: None, + }, + }), + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; let server_req = timeout( DEFAULT_READ_TIMEOUT, @@ -1221,55 +1055,42 @@ async fn turn_start_accepts_text_at_limit_with_mention_item() -> Result<()> { let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![ - V2UserInput::Text { - text: "x".repeat(MAX_USER_INPUT_TEXT_CHARS), - text_elements: Vec::new(), - }, - V2UserInput::Mention { - name: "Demo App".to_string(), - path: "app://demo-app".to_string(), - }, - ], - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![ + V2UserInput::Text { + text: "x".repeat(MAX_USER_INPUT_TEXT_CHARS), + text_elements: Vec::new(), + }, + V2UserInput::Mention { + name: "Demo App".to_string(), + path: "app://demo-app".to_string(), + }, + ], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; assert_eq!(turn.status, TurnStatus::InProgress); timeout( @@ -1284,31 +1105,21 @@ async fn turn_start_accepts_text_at_limit_with_mention_item() -> Result<()> { #[tokio::test] async fn turn_start_rejects_combined_oversized_text_input() -> Result<()> { let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - "http://localhost/unused", - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new("http://localhost/unused") + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let first = "x".repeat(MAX_USER_INPUT_TEXT_CHARS / 2); let second = "y".repeat(MAX_USER_INPUT_TEXT_CHARS / 2 + 1); @@ -1363,12 +1174,9 @@ async fn turn_start_rejects_combined_oversized_text_input() -> Result<()> { #[tokio::test] async fn turn_start_rejects_invalid_permission_selection_before_starting_turn() -> Result<()> { let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - "http://localhost/unused", - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new("http://localhost/unused") + .enable_feature(Feature::Personality) + .write(codex_home.path())?; std::fs::write( codex_home.path().join("managed_config.toml"), "sandbox_mode = \"read-only\"\n", @@ -1376,22 +1184,15 @@ async fn turn_start_rejects_invalid_permission_selection_before_starting_turn() let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let turn_req = mcp .send_turn_start_request(TurnStartParams { thread_id: thread.id, @@ -1444,12 +1245,9 @@ async fn turn_start_accepts_managed_network_profile_from_requirements() -> Resul let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::NetworkProxy, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::NetworkProxy) + .write(codex_home.path())?; std::fs::write( codex_home.path().join("requirements.toml"), r#" @@ -1473,48 +1271,38 @@ allow_local_binding = false let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; - let thread_req = app_server - .send_thread_start_request_with_auto_env(ThreadStartParams { - model: Some("mock-model".to_string()), - ..Default::default() - }) - .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; let ThreadStartResponse { thread, active_permission_profile, .. - } = to_response::(thread_resp)?; + } = app_server + .start_thread(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; let active_permission_profile = active_permission_profile.context("expected active permission profile")?; assert_eq!(active_permission_profile.id, "managed-network"); - let turn_req = app_server - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Use the managed network profile".to_string(), - text_elements: Vec::new(), - }], - permissions: Some("managed-network".to_string()), - ..Default::default() + let TurnStartResponse { turn } = app_server + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Use the managed network profile".to_string(), + text_elements: Vec::new(), + }], + permissions: Some("managed-network".to_string()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; assert!( !turn.id.is_empty(), "turn/start should resolve the managed profile's network configuration" @@ -1532,31 +1320,19 @@ allow_local_binding = false async fn turn_start_rejects_unknown_environment_before_starting_turn() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let turn_req = mcp .send_turn_start_request(TurnStartParams { @@ -1611,61 +1387,43 @@ async fn turn_start_emits_notifications_and_accepts_model_override() -> Result<( let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // Start a thread (v2) and capture its id. - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; // Start a turn with only input and thread_id set (no overrides). - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; assert!(!turn.id.is_empty()); // Expect a turn/started notification. - let notif: JSONRPCNotification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/started"), - ) - .await??; let started: TurnStartedNotification = - serde_json::from_value(notif.params.expect("params must be present"))?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("turn/started")).await??; assert_eq!(started.thread_id, thread.id); assert_eq!( started.turn.status, @@ -1675,66 +1433,48 @@ async fn turn_start_emits_notifications_and_accepts_model_override() -> Result<( assert_eq!(started.turn.items_view, TurnItemsView::NotLoaded); assert!(started.turn.items.is_empty()); - let completed_notif: JSONRPCNotification = timeout( + let completed: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed: TurnCompletedNotification = serde_json::from_value( - completed_notif - .params - .expect("turn/completed params must be present"), - )?; assert_eq!(completed.thread_id, thread.id); assert_eq!(completed.turn.id, turn.id); assert_eq!(completed.turn.status, TurnStatus::Completed); // Send a second turn that exercises the overrides path: change the model. - let turn_req2 = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Second".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model-override".to_string()), - ..Default::default() + let TurnStartResponse { turn: turn2 } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Second".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model-override".to_string()), + ..Default::default() + }, }) .await?; - let turn_resp2: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req2)), - ) - .await??; - let TurnStartResponse { turn: turn2 } = to_response::(turn_resp2)?; assert!(!turn2.id.is_empty()); // Ensure the second turn has a different id than the first. assert_ne!(turn.id, turn2.id); - let notif2: JSONRPCNotification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/started"), - ) - .await??; let started2: TurnStartedNotification = - serde_json::from_value(notif2.params.expect("params must be present"))?; + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("turn/started")).await??; assert_eq!(started2.thread_id, thread.id); assert_eq!(started2.turn.id, turn2.id); assert_eq!(started2.turn.status, TurnStatus::InProgress); assert_eq!(started2.turn.items_view, TurnItemsView::NotLoaded); assert!(started2.turn.items.is_empty()); - let completed_notif2: JSONRPCNotification = timeout( + let completed2: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed2: TurnCompletedNotification = serde_json::from_value( - completed_notif2 - .params - .expect("turn/completed params must be present"), - )?; assert_eq!(completed2.thread_id, thread.id); assert_eq!(completed2.turn.id, turn2.id); assert_eq!(completed2.turn.status, TurnStatus::Completed); @@ -1754,32 +1494,20 @@ async fn turn_start_accepts_collaboration_mode_override_v2() -> Result<()> { ]); let response_mock = responses::mount_sse_once(&server, body).await; - let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + let codex_home = TempDir::new()?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("gpt-5.4".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let collaboration_mode = CollaborationMode { mode: ModeKind::Default, @@ -1790,28 +1518,25 @@ async fn turn_start_accepts_collaboration_mode_override_v2() -> Result<()> { }, }; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model-override".to_string()), - effort: Some(ReasoningEffort::Low), - summary: Some(ReasoningSummary::Auto), - output_schema: None, - collaboration_mode: Some(collaboration_mode), - ..Default::default() + let _turn: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model-override".to_string()), + effort: Some(ReasoningEffort::Low), + summary: Some(ReasoningSummary::Auto), + output_schema: None, + collaboration_mode: Some(collaboration_mode), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _turn: TurnStartResponse = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -1844,21 +1569,15 @@ async fn turn_start_uses_thread_feature_overrides_for_request_user_input_tool_de let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("gpt-5.4".to_string()), config: Some(HashMap::from([( "features.default_mode_request_user_input".to_string(), @@ -1867,12 +1586,6 @@ async fn turn_start_uses_thread_feature_overrides_for_request_user_input_tool_de ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let collaboration_mode = CollaborationMode { mode: ModeKind::Default, @@ -1883,28 +1596,25 @@ async fn turn_start_uses_thread_feature_overrides_for_request_user_input_tool_de }, }; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - model: Some("mock-model-override".to_string()), - effort: Some(ReasoningEffort::Low), - summary: Some(ReasoningSummary::Auto), - output_schema: None, - collaboration_mode: Some(collaboration_mode), - ..Default::default() + let _turn: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + model: Some("mock-model-override".to_string()), + effort: Some(ReasoningEffort::Low), + summary: Some(ReasoningSummary::Auto), + output_schema: None, + collaboration_mode: Some(collaboration_mode), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _turn: TurnStartResponse = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -1932,50 +1642,37 @@ async fn turn_start_accepts_personality_override_v2() -> Result<()> { let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("exp-codex-personality".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - personality: Some(Personality::Friendly), - ..Default::default() + let _turn: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + personality: Some(Personality::Friendly), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _turn: TurnStartResponse = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -2012,49 +1709,36 @@ async fn turn_start_ignores_deprecated_multi_agent_mode() -> Result<()> { let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::MultiAgentV2, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::MultiAgentV2) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - multi_agent_mode: Some(MultiAgentMode::Proactive), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + multi_agent_mode: Some(MultiAgentMode::Proactive), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -2092,54 +1776,41 @@ async fn thread_start_ignores_deprecated_multi_agent_mode() -> Result<()> { let response_mock = responses::mount_sse_once(&server, body).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::MultiAgentV2, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::MultiAgentV2) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { + thread, + multi_agent_mode, + .. + } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), multi_agent_mode: Some(MultiAgentMode::Proactive), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { - thread, - multi_agent_mode, - .. - } = to_response::(thread_resp)?; assert_eq!(multi_agent_mode, MultiAgentMode::ExplicitRequestOnly); - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -2183,50 +1854,37 @@ async fn turn_start_change_personality_mid_thread_v2() -> Result<()> { let response_mock = responses::mount_sse_sequence(&server, vec![sse1, sse2]).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("exp-codex-personality".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - personality: None, - ..Default::default() + let _turn: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + personality: None, + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _turn: TurnStartResponse = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, @@ -2234,24 +1892,21 @@ async fn turn_start_change_personality_mid_thread_v2() -> Result<()> { ) .await??; - let turn_req2 = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "Hello again".to_string(), - text_elements: Vec::new(), - }], - personality: Some(Personality::Friendly), - ..Default::default() + let _turn2: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "Hello again".to_string(), + text_elements: Vec::new(), + }], + personality: Some(Personality::Friendly), + ..Default::default() + }, }) .await?; - let turn_resp2: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req2)), - ) - .await??; - let _turn2: TurnStartResponse = to_response::(turn_resp2)?; timeout( DEFAULT_READ_TIMEOUT, @@ -2347,33 +2002,23 @@ async fn turn_start_exec_approval_toggle_v2() -> Result<()> { ]; let server = create_mock_responses_server_sequence(responses).await; // Default approval is untrusted to force elicitation on first turn. - create_config_toml( - codex_home.as_path(), - &server.uri(), - "untrusted", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(codex_home.as_path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.as_path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let expected_environment_id = mcp.auto_env_params()?.environment_id; // thread/start - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(start_resp)?; // turn/start — expect CommandExecutionRequestApproval request from server let first_turn_id = mcp @@ -2445,27 +2090,25 @@ async fn turn_start_exec_approval_toggle_v2() -> Result<()> { } // Second turn with approval_policy=never should not elicit approval - let second_turn_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run python again".to_string(), - text_elements: Vec::new(), - }], - approval_policy: Some(codex_app_server_protocol::AskForApproval::Never), - sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess), - model: Some("mock-model".to_string()), - effort: Some(ReasoningEffort::Medium), - summary: Some(ReasoningSummary::Auto), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run python again".to_string(), + text_elements: Vec::new(), + }], + approval_policy: Some(codex_app_server_protocol::AskForApproval::Never), + sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess), + model: Some("mock-model".to_string()), + effort: Some(ReasoningEffort::Medium), + summary: Some(ReasoningSummary::Auto), + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_turn_id)), - ) - .await??; // Ensure we do NOT receive a CommandExecutionRequestApproval request before task completes timeout( @@ -2528,57 +2171,40 @@ async fn run_turn_start_exec_approval_rejection_v2( create_final_assistant_message_sse_response("done")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml( - codex_home.as_path(), - &server.uri(), - "untrusted", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(codex_home.as_path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.as_path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(start_resp)?; - let turn_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run python".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run python".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let started_command_execution = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let started_notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = - serde_json::from_value(started_notif.params.clone().expect("item/started params"))?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if let ThreadItem::CommandExecution { .. } = started.item { return Ok::(started.item); } @@ -2607,15 +2233,8 @@ async fn run_turn_start_exec_approval_rejection_v2( let completed_command_execution = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let completed_notif = mcp - .read_stream_until_notification_message("item/completed") - .await?; - let completed: ItemCompletedNotification = serde_json::from_value( - completed_notif - .params - .clone() - .expect("item/completed params"), - )?; + let completed: ItemCompletedNotification = + mcp.read_notification("item/completed").await?; if let ThreadItem::CommandExecution { .. } = completed.item { return Ok::(completed.item); } @@ -2690,71 +2309,61 @@ async fn turn_start_explicit_local_environment_updates_legacy_cwd_between_turns( create_final_assistant_message_sse_response("done second")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml( - &codex_home, - &server.uri(), - "untrusted", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; // thread/start - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(start_resp)?; // first turn with workspace-write sandbox and first_cwd - let first_turn = mcp - .send_turn_start_request(TurnStartParams { - environments: None, - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "first turn".to_string(), - text_elements: Vec::new(), - }], - responsesapi_client_metadata: None, - additional_context: None, - cwd: Some(first_cwd.clone()), - runtime_workspace_roots: None, - approval_policy: Some(codex_app_server_protocol::AskForApproval::Never), - approvals_reviewer: None, - sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::WorkspaceWrite { - writable_roots: vec![first_cwd.try_into()?], - network_access: false, - exclude_tmpdir_env_var: true, - exclude_slash_tmp: true, - }), - permissions: None, - model: Some("mock-model".to_string()), - effort: Some(ReasoningEffort::Medium), - summary: Some(ReasoningSummary::Auto), - service_tier: None, - personality: None, - output_schema: None, - collaboration_mode: None, - multi_agent_mode: None, + let first_writable_root = + codex_utils_absolute_path::AbsolutePathBuf::try_from(first_cwd.clone())?; + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + environments: None, + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "first turn".to_string(), + text_elements: Vec::new(), + }], + responsesapi_client_metadata: None, + additional_context: None, + cwd: Some(first_cwd.clone()), + runtime_workspace_roots: None, + approval_policy: Some(codex_app_server_protocol::AskForApproval::Never), + approvals_reviewer: None, + sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::WorkspaceWrite { + writable_roots: vec![first_writable_root], + network_access: false, + exclude_tmpdir_env_var: true, + exclude_slash_tmp: true, + }), + permissions: None, + model: Some("mock-model".to_string()), + effort: Some(ReasoningEffort::Medium), + summary: Some(ReasoningSummary::Auto), + service_tier: None, + personality: None, + output_schema: None, + collaboration_mode: None, + multi_agent_mode: None, + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_turn)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -2764,65 +2373,51 @@ async fn turn_start_explicit_local_environment_updates_legacy_cwd_between_turns( // Select a new local cwd without the top-level compatibility parameter. The inherited // workspace-write sandbox must follow the local environment cwd. - let second_turn = mcp - .send_turn_start_request(TurnStartParams { - environments: Some(vec![TurnEnvironmentParams { - environment_id: LOCAL_ENVIRONMENT_ID.to_string(), - cwd: second_cwd.abs().into(), + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + environments: Some(vec![TurnEnvironmentParams { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: second_cwd.abs().into(), + runtime_workspace_roots: None, + }]), + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "second turn".to_string(), + text_elements: Vec::new(), + }], + responsesapi_client_metadata: None, + additional_context: None, + cwd: None, runtime_workspace_roots: None, - }]), - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "second turn".to_string(), - text_elements: Vec::new(), - }], - responsesapi_client_metadata: None, - additional_context: None, - cwd: None, - runtime_workspace_roots: None, - approval_policy: Some(codex_app_server_protocol::AskForApproval::Never), - approvals_reviewer: None, - sandbox_policy: None, - permissions: None, - model: Some("mock-model".to_string()), - effort: Some(ReasoningEffort::Medium), - summary: Some(ReasoningSummary::Auto), - service_tier: None, - personality: None, - output_schema: None, - collaboration_mode: None, - multi_agent_mode: None, + approval_policy: Some(codex_app_server_protocol::AskForApproval::Never), + approvals_reviewer: None, + sandbox_policy: None, + permissions: None, + model: Some("mock-model".to_string()), + effort: Some(ReasoningEffort::Medium), + summary: Some(ReasoningSummary::Auto), + service_tier: None, + personality: None, + output_schema: None, + collaboration_mode: None, + multi_agent_mode: None, + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_turn)), - ) - .await??; - let settings_updated = timeout( + let settings_updated: ThreadSettingsUpdatedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/settings/updated"), + mcp.read_notification("thread/settings/updated"), ) .await??; - let settings_updated: ThreadSettingsUpdatedNotification = serde_json::from_value( - settings_updated - .params - .context("thread/settings/updated should include params")?, - )?; assert_eq!(settings_updated.thread_settings.cwd, second_cwd.abs()); let command_exec_item = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let item_started_notification = mcp - .read_stream_until_notification_message("item/started") - .await?; - let params = item_started_notification - .params - .clone() - .expect("item/started params"); let item_started: ItemStartedNotification = - serde_json::from_value(params).expect("deserialize item/started notification"); + mcp.read_notification("item/started").await?; if matches!(item_started.item, ThreadItem::CommandExecution { .. }) { return Ok::(item_started.item); } @@ -2912,64 +2507,53 @@ stream_max_retries = 0 let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(start_resp)?; - let first_turn_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "select dev profile".to_string(), - text_elements: Vec::new(), - }], - runtime_workspace_roots: Some(vec![old_root]), - permissions: Some("dev".to_string()), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "select dev profile".to_string(), + text_elements: Vec::new(), + }], + runtime_workspace_roots: Some(vec![old_root]), + permissions: Some("dev".to_string()), + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(first_turn_id)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), ) .await??; - let second_turn_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "write in new root".to_string(), - text_elements: Vec::new(), - }], - runtime_workspace_roots: Some(vec![new_root]), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "write in new root".to_string(), + text_elements: Vec::new(), + }], + runtime_workspace_roots: Some(vec![new_root]), + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(second_turn_id)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, @@ -3014,7 +2598,7 @@ async fn turn_start_resolves_sticky_thread_local_environment_and_turn_overrides( std::fs::create_dir(&workspace)?; let server = create_mock_responses_server_repeating_assistant("done").await; - create_config_toml(&codex_home, &server.uri(), "never", &BTreeMap::default())?; + MockResponsesConfig::new(&server.uri()).write(&codex_home)?; std::fs::write( codex_home.join("environments.toml"), r#" @@ -3029,9 +2613,8 @@ url = "ws://127.0.0.1:1" // This test owns environments.toml and explicitly compares local selections // with a configured remote environment, so auto env would change its subject. .without_auto_env() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; for case in [ EnvironmentSelectionCase { @@ -3085,55 +2668,36 @@ async fn run_environment_selection_case( ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: format!("run {}", case.name), - text_elements: Vec::new(), - }], - environments: environment_params(case.turn, workspace), - cwd: Some(workspace.to_path_buf()), - model: Some("mock-model".to_string()), - ..Default::default() + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; + + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: format!("run {}", case.name), + text_elements: Vec::new(), + }], + environments: environment_params(case.turn, workspace), + cwd: Some(workspace.to_path_buf()), + model: Some("mock-model".to_string()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; - let started_notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/started"), - ) - .await??; - let started: TurnStartedNotification = serde_json::from_value( - started_notification - .params - .ok_or_else(|| anyhow::anyhow!("turn/started notification should include params"))?, - )?; + let started: TurnStartedNotification = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("turn/started")).await??; assert_eq!(started.turn.id, turn.id, "{}", case.name); - let completed_notification = timeout( + let completed: TurnCompletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), + mcp.read_notification("turn/completed"), ) .await??; - let completed: TurnCompletedNotification = - serde_json::from_value(completed_notification.params.ok_or_else(|| { - anyhow::anyhow!("turn/completed notification should include params") - })?)?; assert_eq!(completed.turn.id, turn.id, "{}", case.name); assert_eq!( completed.turn.status, @@ -3184,59 +2748,42 @@ async fn turn_start_file_change_approval_v2() -> Result<()> { create_final_assistant_message_sse_response("patch applied")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml( - &codex_home, - &server.uri(), - "untrusted", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), cwd: Some(workspace.to_string_lossy().into_owned()), ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "apply patch".into(), - text_elements: Vec::new(), - }], - cwd: Some(workspace.clone()), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "apply patch".into(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let started_file_change = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let started_notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = - serde_json::from_value(started_notif.params.clone().expect("item/started params"))?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if let ThreadItem::FileChange { .. } = started.item { return Ok::(started.item); } @@ -3388,45 +2935,36 @@ async fn turn_start_does_not_stream_apply_patch_change_updates_without_feature_v create_final_assistant_message_sse_response("patch applied")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml(&codex_home, &server.uri(), "never", &BTreeMap::default())?; + MockResponsesConfig::new(&server.uri()).write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), cwd: Some(workspace.to_string_lossy().into_owned()), ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "apply patch".into(), - text_elements: Vec::new(), - }], - cwd: Some(workspace), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "apply patch".into(), + text_elements: Vec::new(), + }], + cwd: Some(workspace), + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -3510,17 +3048,12 @@ async fn turn_start_streams_apply_patch_change_updates_v2() -> Result<()> { create_final_assistant_message_sse_response("patch applied")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml( - &codex_home, - &server.uri(), - "never", - &BTreeMap::from([ - (Feature::ApplyPatchStreamingEvents, true), - (Feature::Plugins, false), - (Feature::RemoteModels, false), - (Feature::ShellSnapshot, false), - ]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::ApplyPatchStreamingEvents) + .disable_feature(Feature::Plugins) + .disable_feature(Feature::RemoteModels) + .disable_feature(Feature::ShellSnapshot) + .write(&codex_home)?; write_models_cache(&codex_home)?; let cache_path = codex_home.join("models_cache.json"); let mut cache: serde_json::Value = @@ -3538,56 +3071,40 @@ async fn turn_start_streams_apply_patch_change_updates_v2() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), cwd: Some(workspace.to_string_lossy().into_owned()), ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "apply patch".into(), - text_elements: Vec::new(), - }], - cwd: Some(workspace.clone()), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "apply patch".into(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let mut streamed_content = String::new(); while streamed_content != "live line\n" { - let delta_notif = timeout( + let delta: FileChangePatchUpdatedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/fileChange/patchUpdated"), + mcp.read_notification("item/fileChange/patchUpdated"), ) .await??; - let delta: FileChangePatchUpdatedNotification = serde_json::from_value( - delta_notif - .params - .clone() - .expect("item/fileChange/patchUpdated params"), - )?; assert_eq!(delta.thread_id, thread.id); assert_eq!(delta.turn_id, turn.id); assert_eq!(delta.item_id, call_id); @@ -3664,57 +3181,40 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Collab, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Collab) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("gpt-5.4".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: PARENT_PROMPT.to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let turn: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: PARENT_PROMPT.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let turn: TurnStartResponse = to_response::(turn_resp)?; let spawn_started = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let started_notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = - serde_json::from_value(started_notif.params.expect("item/started params"))?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if let ThreadItem::CollabAgentToolCall { id, .. } = &started.item && id == SPAWN_CALL_ID { @@ -3740,11 +3240,8 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() let spawn_completed = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let completed_notif = mcp - .read_stream_until_notification_message("item/completed") - .await?; let completed: ItemCompletedNotification = - serde_json::from_value(completed_notif.params.expect("item/completed params"))?; + mcp.read_notification("item/completed").await?; if let ThreadItem::CollabAgentToolCall { id, .. } = &completed.item && id == SPAWN_CALL_ID { @@ -3794,12 +3291,8 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() let turn_completed = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let turn_completed_notif = mcp - .read_stream_until_notification_message("turn/completed") - .await?; - let turn_completed: TurnCompletedNotification = serde_json::from_value( - turn_completed_notif.params.expect("turn/completed params"), - )?; + let turn_completed: TurnCompletedNotification = + mcp.read_notification("turn/completed").await?; if turn_completed.thread_id == thread.id && turn_completed.turn.id == turn.turn.id { return Ok::(turn_completed); } @@ -3810,30 +3303,22 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() assert_eq!(turn_completed.turn.id, turn.turn.id); // Reuse this live spawn setup to cover thread/delete's ThreadManager descendant path. - let delete_req = mcp - .send_thread_delete_request(ThreadDeleteParams { - thread_id: thread.id.clone(), + let _: ThreadDeleteResponse = mcp + .request(|request_id| ClientRequest::ThreadDelete { + request_id, + params: ThreadDeleteParams { + thread_id: thread.id.clone(), + }, }) .await?; - let delete_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(delete_req)), - ) - .await??; - let _: ThreadDeleteResponse = to_response::(delete_resp)?; let mut deleted_thread_ids = Vec::new(); for _ in 0..2 { - let deleted_notif = timeout( + let deleted: ThreadDeletedNotification = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("thread/deleted"), + mcp.read_notification("thread/deleted"), ) .await??; - let deleted: ThreadDeletedNotification = serde_json::from_value( - deleted_notif - .params - .expect("thread/deleted notification params"), - )?; deleted_thread_ids.push(deleted.thread_id); } assert_eq!( @@ -3841,15 +3326,12 @@ async fn turn_start_emits_spawn_agent_item_with_model_metadata_v2() -> Result<() vec![receiver_thread_id, thread.id.clone()] ); - let list_req = mcp - .send_thread_loaded_list_request(ThreadLoadedListParams::default()) + let ThreadLoadedListResponse { data, .. } = mcp + .request(|request_id| ClientRequest::ThreadLoadedList { + request_id, + params: ThreadLoadedListParams::default(), + }) .await?; - let list_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(list_req)), - ) - .await??; - let ThreadLoadedListResponse { data, .. } = to_response::(list_resp)?; assert_eq!(data, Vec::::new()); Ok(()) @@ -3884,57 +3366,41 @@ async fn direct_input_to_multi_agent_v2_subagent_is_rejected() -> Result<()> { ) .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::MultiAgentV2, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::MultiAgentV2) + .write(codex_home.path())?; write_models_cache(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("gpt-5.4".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - input: vec![V2UserInput::Text { - text: PARENT_PROMPT.to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: PARENT_PROMPT.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _: TurnStartResponse = to_response(turn_resp)?; let child_thread_id = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let completed_notif = mcp - .read_stream_until_notification_message("item/completed") - .await?; let completed: ItemCompletedNotification = - serde_json::from_value(completed_notif.params.expect("item/completed params"))?; + mcp.read_notification("item/completed").await?; if let ThreadItem::SubAgentActivity { id, kind: SubAgentActivityKind::Started, @@ -4049,12 +3515,9 @@ async fn turn_start_emits_spawn_agent_item_with_effective_role_model_metadata_v2 .await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Collab, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Collab) + .write(codex_home.path())?; std::fs::write( codex_home.path().join("custom-role.toml"), format!("model = \"{ROLE_MODEL}\"\nmodel_reasoning_effort = \"{ROLE_REASONING_EFFORT}\"\n",), @@ -4075,48 +3538,35 @@ config_file = "./custom-role.toml" let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("gpt-5.4".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: PARENT_PROMPT.to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let turn: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: PARENT_PROMPT.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let turn: TurnStartResponse = to_response::(turn_resp)?; let spawn_completed = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let completed_notif = mcp - .read_stream_until_notification_message("item/completed") - .await?; let completed: ItemCompletedNotification = - serde_json::from_value(completed_notif.params.expect("item/completed params"))?; + mcp.read_notification("item/completed").await?; if let ThreadItem::CollabAgentToolCall { id, .. } = &completed.item && id == SPAWN_CALL_ID { @@ -4166,12 +3616,8 @@ config_file = "./custom-role.toml" let turn_completed = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let turn_completed_notif = mcp - .read_stream_until_notification_message("turn/completed") - .await?; - let turn_completed: TurnCompletedNotification = serde_json::from_value( - turn_completed_notif.params.expect("turn/completed params"), - )?; + let turn_completed: TurnCompletedNotification = + mcp.read_notification("turn/completed").await?; if turn_completed.thread_id == thread.id && turn_completed.turn.id == turn.turn.id { return Ok::(turn_completed); } @@ -4218,60 +3664,43 @@ async fn turn_start_file_change_approval_accept_for_session_persists_v2() -> Res create_final_assistant_message_sse_response("patch 2 applied")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml( - &codex_home, - &server.uri(), - "untrusted", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), cwd: Some(workspace.to_string_lossy().into_owned()), ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; // First turn: expect FileChangeRequestApproval, respond with AcceptForSession, and verify the file exists. - let turn_1_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "apply patch 1".into(), - text_elements: Vec::new(), - }], - cwd: Some(workspace.clone()), - ..Default::default() + let TurnStartResponse { turn: turn_1 } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "apply patch 1".into(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + ..Default::default() + }, }) .await?; - let turn_1_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_1_req)), - ) - .await??; - let TurnStartResponse { turn: turn_1 } = to_response::(turn_1_resp)?; let started_file_change_1 = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let started_notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = - serde_json::from_value(started_notif.params.clone().expect("item/started params"))?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if let ThreadItem::FileChange { .. } = started.item { return Ok::(started.item); } @@ -4319,31 +3748,25 @@ async fn turn_start_file_change_approval_accept_for_session_persists_v2() -> Res assert_eq!(std::fs::read_to_string(&readme_path)?, "new line\n"); // Second turn: apply a patch to the same file. Approval should be skipped due to AcceptForSession. - let turn_2_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "apply patch 2".into(), - text_elements: Vec::new(), - }], - cwd: Some(workspace.clone()), - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "apply patch 2".into(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_2_req)), - ) - .await??; let started_file_change_2 = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let started_notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = - serde_json::from_value(started_notif.params.clone().expect("item/started params"))?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if let ThreadItem::FileChange { .. } = started.item { return Ok::(started.item); } @@ -4421,59 +3844,42 @@ async fn run_turn_start_file_change_approval_rejection_v2( create_final_assistant_message_sse_response("patch declined")?, ]; let server = create_mock_responses_server_sequence(responses).await; - create_config_toml( - &codex_home, - &server.uri(), - "untrusted", - &BTreeMap::default(), - )?; + MockResponsesConfig::new(&server.uri()) + .with_approval_policy("untrusted") + .write(&codex_home)?; let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), cwd: Some(workspace.to_string_lossy().into_owned()), ..Default::default() }) .await?; - let start_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(start_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "apply patch".into(), - text_elements: Vec::new(), - }], - cwd: Some(workspace.clone()), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "apply patch".into(), + text_elements: Vec::new(), + }], + cwd: Some(workspace.clone()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let started_file_change = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let started_notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = - serde_json::from_value(started_notif.params.clone().expect("item/started params"))?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if let ThreadItem::FileChange { .. } = started.item { return Ok::(started.item); } @@ -4518,15 +3924,8 @@ async fn run_turn_start_file_change_approval_rejection_v2( let completed_file_change = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let completed_notif = mcp - .read_stream_until_notification_message("item/completed") - .await?; - let completed: ItemCompletedNotification = serde_json::from_value( - completed_notif - .params - .clone() - .expect("item/completed params"), - )?; + let completed: ItemCompletedNotification = + mcp.read_notification("item/completed").await?; if let ThreadItem::FileChange { .. } = completed.item { return Ok::(completed.item); } @@ -4580,63 +3979,42 @@ async fn command_execution_notifications_include_process_id() -> Result<()> { ]; let server = create_mock_responses_server_sequence(responses).await; let codex_home = TempDir::new()?; - create_config_toml_with_sandbox( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::UnifiedExec, true)]), - "danger-full-access", - )?; + MockResponsesConfig::new(&server.uri()) + .with_sandbox_mode("danger-full-access") + .enable_feature(Feature::UnifiedExec) + .write(codex_home.path())?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let start_id = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(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::(start_resp)?; - let turn_id = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run a command".to_string(), - text_elements: Vec::new(), - }], - sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess), - ..Default::default() + let TurnStartResponse { turn: _turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run a command".to_string(), + text_elements: Vec::new(), + }], + sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let TurnStartResponse { turn: _turn } = to_response::(turn_resp)?; let started_command = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let notif = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = serde_json::from_value( - notif - .params - .clone() - .expect("item/started should include params"), - )?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if let ThreadItem::CommandExecution { .. } = started.item { return Ok::(started.item); } @@ -4658,15 +4036,8 @@ async fn command_execution_notifications_include_process_id() -> Result<()> { let completed_command = timeout(DEFAULT_READ_TIMEOUT, async { loop { - let notif = mcp - .read_stream_until_notification_message("item/completed") - .await?; - let completed: ItemCompletedNotification = serde_json::from_value( - notif - .params - .clone() - .expect("item/completed should include params"), - )?; + let completed: ItemCompletedNotification = + mcp.read_notification("item/completed").await?; if let ThreadItem::CommandExecution { .. } = completed.item { return Ok::(completed.item); } @@ -4716,51 +4087,39 @@ async fn turn_start_with_elevated_override_does_not_persist_project_trust() -> R let server = create_mock_responses_server_sequence_unchecked(responses).await; let codex_home = TempDir::new()?; - create_config_toml( - codex_home.path(), - &server.uri(), - "never", - &BTreeMap::from([(Feature::Personality, true)]), - )?; + MockResponsesConfig::new(&server.uri()) + .enable_feature(Feature::Personality) + .write(codex_home.path())?; let workspace = TempDir::new()?; let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_request = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { cwd: Some(workspace.path().display().to_string()), ..Default::default() }) .await?; - let thread_response: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_request)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_response)?; - let turn_request = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id, - cwd: Some(workspace.path().to_path_buf()), - sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess), - input: vec![V2UserInput::Text { - text: "Hello".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() + let _: TurnStartResponse = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id, + cwd: Some(workspace.path().to_path_buf()), + sandbox_policy: Some(codex_app_server_protocol::SandboxPolicy::DangerFullAccess), + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }, }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_request)), - ) - .await??; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -4774,70 +4133,6 @@ async fn turn_start_with_elevated_override_does_not_persist_project_trust() -> R Ok(()) } -// Helper to create a config.toml pointing at the mock model server. -fn create_config_toml( - codex_home: &Path, - server_uri: &str, - approval_policy: &str, - feature_flags: &BTreeMap, -) -> std::io::Result<()> { - create_config_toml_with_sandbox( - codex_home, - server_uri, - approval_policy, - feature_flags, - "read-only", - ) -} - -fn create_config_toml_with_sandbox( - codex_home: &Path, - server_uri: &str, - approval_policy: &str, - feature_flags: &BTreeMap, - sandbox_mode: &str, -) -> std::io::Result<()> { - let mut features = BTreeMap::new(); - for (feature, enabled) in feature_flags { - features.insert(*feature, *enabled); - } - let feature_entries = features - .into_iter() - .map(|(feature, enabled)| { - let key = FEATURES - .iter() - .find(|spec| spec.id == feature) - .map(|spec| spec.key) - .expect("feature should have a config key"); - format!("{key} = {enabled}") - }) - .collect::>() - .join("\n"); - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "mock-model" -approval_policy = "{approval_policy}" -sandbox_mode = "{sandbox_mode}" - -model_provider = "mock_provider" - -[features] -{feature_entries} - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) -} - fn write_test_skill(codex_home: &Path, name: &str) -> std::io::Result<()> { let skill_dir = codex_home.join("skills").join(name); std::fs::create_dir_all(&skill_dir)?; diff --git a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs index 0dbc516b8772..99c5951a238d 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs @@ -7,20 +7,18 @@ // network access are required the first time the artifact is fetched. use anyhow::Result; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::create_shell_command_sse_response; -use app_test_support::to_response; use codex_app_server_protocol::CommandAction; use codex_app_server_protocol::CommandExecutionApprovalDecision; use codex_app_server_protocol::CommandExecutionRequestApprovalResponse; use codex_app_server_protocol::CommandExecutionStatus; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; @@ -30,7 +28,6 @@ use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; -use codex_features::FEATURES; use codex_features::Feature; use core_test_support::responses; use core_test_support::skip_if_no_network; @@ -104,7 +101,6 @@ async fn turn_start_shell_zsh_fork_executes_command_v2() -> Result<()> { )?; let mut mcp = create_zsh_test_mcp_process(&codex_home, &workspace, &zsh_path).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -113,12 +109,8 @@ async fn turn_start_shell_zsh_fork_executes_command_v2() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -137,12 +129,8 @@ async fn turn_start_shell_zsh_fork_executes_command_v2() -> Result<()> { ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_id)).await??; let started_command_execution = timeout(DEFAULT_READ_TIMEOUT, async { loop { @@ -228,7 +216,6 @@ async fn turn_start_shell_zsh_fork_exec_approval_decline_v2() -> Result<()> { )?; let mut mcp = create_zsh_test_mcp_process(&codex_home, &workspace, &zsh_path).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -237,12 +224,8 @@ async fn turn_start_shell_zsh_fork_exec_approval_decline_v2() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -256,11 +239,7 @@ async fn turn_start_shell_zsh_fork_exec_approval_decline_v2() -> Result<()> { ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_id)).await??; let server_req = timeout( DEFAULT_READ_TIMEOUT, @@ -366,7 +345,6 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { )?; let mut mcp = create_zsh_test_mcp_process(&codex_home, &workspace, &zsh_path).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -375,12 +353,8 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -394,11 +368,7 @@ async fn turn_start_shell_zsh_fork_exec_approval_cancel_v2() -> Result<()> { ..Default::default() }) .await?; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; + let _: TurnStartResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_id)).await??; let server_req = timeout( DEFAULT_READ_TIMEOUT, @@ -530,7 +500,6 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() )?; let mut mcp = create_zsh_test_mcp_process(&codex_home, &workspace, &zsh_path).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let start_id = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -539,12 +508,8 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() ..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::(start_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(start_id)).await??; let turn_id = mcp .send_turn_start_request(TurnStartParams { @@ -567,12 +532,8 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; + let TurnStartResponse { turn } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_id)).await??; let mut approved_subcommand_strings = Vec::new(); let mut approved_subcommand_ids = Vec::new(); @@ -771,7 +732,7 @@ async fn create_zsh_test_mcp_process( .with_codex_home(codex_home) .with_program(&app_server) .with_env_overrides(&[("ZDOTDIR", Some(zdotdir.as_str()))]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await } @@ -820,45 +781,11 @@ fn create_config_toml( approval_policy: &str, feature_flags: &BTreeMap, ) -> std::io::Result<()> { - let mut features = BTreeMap::from([(Feature::RemoteModels, false)]); - for (feature, enabled) in feature_flags { - features.insert(*feature, *enabled); - } - let feature_entries = features - .into_iter() - .map(|(feature, enabled)| { - let key = FEATURES - .iter() - .find(|spec| spec.id == feature) - .map(|spec| spec.key) - .expect("feature should have a config key"); - format!("{key} = {enabled}") - }) - .collect::>() - .join("\n"); - let config_toml = codex_home.join("config.toml"); - std::fs::write( - config_toml, - format!( - r#" -model = "mock-model" -approval_policy = "{approval_policy}" -sandbox_mode = "read-only" - -model_provider = "mock_provider" - -[features] -{feature_entries} - -[model_providers.mock_provider] -name = "Mock provider for test" -base_url = "{server_uri}/v1" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -"# - ), - ) + MockResponsesConfig::new(server_uri) + .with_approval_policy(approval_policy) + .disable_feature(Feature::RemoteModels) + .with_features(feature_flags) + .write(codex_home) } fn find_test_zsh_path() -> Result> { diff --git a/codex-rs/app-server/tests/suite/v2/turn_steer.rs b/codex-rs/app-server/tests/suite/v2/turn_steer.rs index 5fb44fb92bec..240b66048bd0 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_steer.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_steer.rs @@ -6,16 +6,15 @@ use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::create_shell_command_sse_response; -use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; use codex_app_server::INPUT_TOO_LARGE_ERROR_CODE; use codex_app_server::INVALID_PARAMS_ERROR_CODE; use codex_app_server_protocol::AdditionalContextEntry; use codex_app_server_protocol::AdditionalContextKind; +use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; @@ -54,22 +53,15 @@ async fn turn_steer_requires_active_turn() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; let steer_req = mcp .send_turn_steer_request(TurnSteerParams { @@ -150,41 +142,31 @@ async fn turn_steer_rejects_oversized_text_input() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run sleep".to_string(), - text_elements: Vec::new(), - }], - cwd: Some(working_directory.clone()), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run sleep".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(working_directory.clone()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let _task_started: JSONRPCNotification = timeout( DEFAULT_READ_TIMEOUT, @@ -274,41 +256,31 @@ async fn turn_steer_returns_active_turn_id() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run sleep".to_string(), - text_elements: Vec::new(), - }], - cwd: Some(working_directory.clone()), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run sleep".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(working_directory.clone()), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; let _task_started: JSONRPCNotification = timeout( DEFAULT_READ_TIMEOUT, @@ -316,25 +288,22 @@ async fn turn_steer_returns_active_turn_id() -> Result<()> { ) .await??; - let steer_req = mcp - .send_turn_steer_request(TurnSteerParams { - thread_id: thread.id.clone(), - client_user_message_id: Some("client-steer-message-1".to_string()), - input: vec![V2UserInput::Text { - text: "steer".to_string(), - text_elements: Vec::new(), - }], - responsesapi_client_metadata: None, - additional_context: None, - expected_turn_id: turn.id.clone(), + let steer: TurnSteerResponse = mcp + .request(|request_id| ClientRequest::TurnSteer { + request_id, + params: TurnSteerParams { + thread_id: thread.id.clone(), + client_user_message_id: Some("client-steer-message-1".to_string()), + input: vec![V2UserInput::Text { + text: "steer".to_string(), + text_elements: Vec::new(), + }], + responsesapi_client_metadata: None, + additional_context: None, + expected_turn_id: turn.id.clone(), + }, }) .await?; - let steer_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(steer_req)), - ) - .await??; - let steer: TurnSteerResponse = to_response::(steer_resp)?; assert_eq!(steer.turn_id, turn.id); timeout(DEFAULT_READ_TIMEOUT, async { @@ -421,41 +390,31 @@ async fn turn_steer_rejects_context_only_input_without_merging_context() -> Resu let mut mcp = TestAppServer::builder() .with_codex_home(&codex_home) .without_managed_config() - .build() + .build_initialized() .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_req = mcp - .send_thread_start_request_with_auto_env(ThreadStartParams { + let ThreadStartResponse { thread, .. } = mcp + .start_thread(ThreadStartParams { model: Some("mock-model".to_string()), ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread.id.clone(), - client_user_message_id: None, - input: vec![V2UserInput::Text { - text: "run sleep".to_string(), - text_elements: Vec::new(), - }], - cwd: Some(working_directory), - ..Default::default() + let TurnStartResponse { turn } = mcp + .request(|request_id| ClientRequest::TurnStart { + request_id, + params: TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![V2UserInput::Text { + text: "run sleep".to_string(), + text_elements: Vec::new(), + }], + cwd: Some(working_directory), + ..Default::default() + }, }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/started"), diff --git a/codex-rs/app-server/tests/suite/v2/web_search.rs b/codex-rs/app-server/tests/suite/v2/web_search.rs index f4893d3a1c58..41450cec7f19 100644 --- a/codex-rs/app-server/tests/suite/v2/web_search.rs +++ b/codex-rs/app-server/tests/suite/v2/web_search.rs @@ -1,17 +1,14 @@ use std::collections::HashMap; -use std::path::Path; use std::time::Duration; use anyhow::Context; use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; -use app_test_support::to_response; use app_test_support::write_chatgpt_auth; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadReadParams; use codex_app_server_protocol::ThreadReadResponse; @@ -23,6 +20,7 @@ use codex_app_server_protocol::UserInput as V2UserInput; use codex_app_server_protocol::WebSearchAction; use codex_app_server_protocol::WebSearchItem; use codex_config::types::AuthCredentialsStoreMode; +use codex_features::Feature; use core_test_support::responses; use core_test_support::responses::strip_response_item_ids_from_json; use pretty_assertions::assert_eq; @@ -86,7 +84,15 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { .await; let codex_home = TempDir::new()?; - create_config_toml(codex_home.path(), &server.uri())?; + MockResponsesConfig::new(&server.uri()) + .with_model_provider("openai-custom") + .with_provider_name("OpenAI") + .with_provider_base_url(&format!("{}/api/codex", server.uri())) + .with_root_config(&format!("chatgpt_base_url = \"{}\"", server.uri())) + .enable_feature(Feature::StandaloneWebSearch) + .with_provider_config("supports_websockets = false") + .with_provider_config("requires_openai_auth = true") + .write(codex_home.path())?; write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("access-chatgpt"), @@ -96,9 +102,8 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let thread_req = mcp .send_thread_start_request_with_auto_env(ThreadStartParams { @@ -106,12 +111,8 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { ..Default::default() }) .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let ThreadStartResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(thread_req)).await??; let thread_id = thread.id.clone(); let turn_req = mcp @@ -126,12 +127,8 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { ..Default::default() }) .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let _turn: TurnStartResponse = to_response::(turn_resp)?; + let _turn: TurnStartResponse = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; let started = timeout(DEFAULT_READ_TIMEOUT, wait_for_web_search_started(&mut mcp)).await??; let completed = timeout( @@ -272,21 +269,16 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { let mut reloaded_mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) - .build() + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) .await?; - timeout(DEFAULT_READ_TIMEOUT, reloaded_mcp.initialize()).await??; let read_req = reloaded_mcp .send_thread_read_request(ThreadReadParams { thread_id, include_turns: true, }) .await?; - let read_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - reloaded_mcp.read_stream_until_response_message(RequestId::Integer(read_req)), - ) - .await??; - let ThreadReadResponse { thread, .. } = to_response::(read_resp)?; + let ThreadReadResponse { thread, .. } = + timeout(DEFAULT_READ_TIMEOUT, reloaded_mcp.read_response(read_req)).await??; let persisted_web_searches: Vec<&ThreadItem> = thread .turns .iter() @@ -300,14 +292,7 @@ async fn standalone_web_search_round_trips_output() -> Result<()> { async fn wait_for_web_search_started(mcp: &mut TestAppServer) -> Result { loop { - let notification = mcp - .read_stream_until_notification_message("item/started") - .await?; - let started: ItemStartedNotification = serde_json::from_value( - notification - .params - .context("item/started notification should include params")?, - )?; + let started: ItemStartedNotification = mcp.read_notification("item/started").await?; if matches!(&started.item, ThreadItem::WebSearch(_)) { return Ok(started); } @@ -318,14 +303,7 @@ async fn wait_for_web_search_completed( mcp: &mut TestAppServer, ) -> Result { loop { - let notification = mcp - .read_stream_until_notification_message("item/completed") - .await?; - let completed: ItemCompletedNotification = serde_json::from_value( - notification - .params - .context("item/completed notification should include params")?, - )?; + let completed: ItemCompletedNotification = mcp.read_notification("item/completed").await?; if matches!(&completed.item, ThreadItem::WebSearch(_)) { return Ok(completed); } @@ -372,30 +350,3 @@ async fn search_request(server: &MockServer) -> Result { .find(|request| request.url.path() == "/api/codex/alpha/search") .context("expected standalone search request") } - -fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { - std::fs::write( - codex_home.join("config.toml"), - format!( - r#" -model = "mock-model" -approval_policy = "never" -sandbox_mode = "read-only" -model_provider = "openai-custom" -chatgpt_base_url = "{server_uri}" - -[features] -standalone_web_search = true - -[model_providers.openai-custom] -name = "OpenAI" -base_url = "{server_uri}/api/codex" -wire_api = "responses" -request_max_retries = 0 -stream_max_retries = 0 -supports_websockets = false -requires_openai_auth = true -"# - ), - ) -}