diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index e1dfe1ef3b..d3e442de8b 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -59,6 +59,7 @@ use rmcp::transport::auth::InMemoryCredentialStore; use rmcp::transport::auth::StoredCredentials; use tokio::sync::Mutex; use tokio::time::sleep; +use tokio::time::timeout; use codex_utils_home_dir::find_codex_home; @@ -66,7 +67,9 @@ const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; const REFRESH_SKEW_MILLIS: u64 = 30_000; const REFRESH_LOCK_DIR: &str = "mcp-oauth-refresh-locks"; +const REFRESH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(60); const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(50); +const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct StoredOAuthTokens { @@ -135,7 +138,7 @@ pub(crate) fn load_oauth_tokens_with_source( keyring_backend_kind: AuthKeyringBackendKind, ) -> Result> { let keyring_store = DefaultKeyringStore; - load_oauth_tokens_with_keyring_store( + load_oauth_tokens_with_source_and_keyring_store( &keyring_store, server_name, url, @@ -144,7 +147,7 @@ pub(crate) fn load_oauth_tokens_with_source( ) } -fn load_oauth_tokens_with_keyring_store( +fn load_oauth_tokens_with_source_and_keyring_store( keyring_store: &K, server_name: &str, url: &str, @@ -178,6 +181,28 @@ fn load_oauth_tokens_with_keyring_store( } } +pub(crate) fn load_oauth_tokens_from_resolved_store( + server_name: &str, + url: &str, + store: ResolvedOAuthCredentialStore, +) -> Result> { + match store { + ResolvedOAuthCredentialStore::File => load_oauth_tokens_from_file(server_name, url) + .context("failed to read OAuth tokens from resolved file storage"), + ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind) => { + load_oauth_tokens_from_keyring( + &DefaultKeyringStore, + keyring_backend_kind, + server_name, + url, + ) + .context( + "failed to read OAuth tokens from resolved keyring storage; refusing file fallback", + ) + } + } +} + pub(crate) fn oauth_token_status( server_name: &str, url: &str, @@ -237,7 +262,11 @@ fn load_oauth_tokens_from_keyring_with_fallback_to_file Result> { // Auto remains keyring-first at lifecycle startup. The returned source is then pinned by the - // per-server OAuth persistor so later refresh work cannot hot-switch stores. + // client transport recipe and OAuth persistor so retries, recovery, and refresh work cannot + // hot-switch stores. + // TODO(stevenlee): Different processes can still resolve Auto to different stores when + // keyring availability differs. Solving that safely requires durable backend selection or + // reconciliation of legacy entries and is intentionally outside this stack. match load_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url) { Ok(Some(tokens)) => Ok(Some(LoadedOAuthTokens { tokens, @@ -378,6 +407,10 @@ fn save_oauth_tokens_with_keyring( } } +/// Saves to the selected keyring backend, then best-effort removes the fallback file entry. +/// +/// A cleanup failure does not change the current client's selected authority, but it can leave +/// legacy residue that a different `Auto` process may discover if keyring availability changes. fn save_oauth_tokens_with_keyring_and_cleanup_file( keyring_store: &K, keyring_backend_kind: AuthKeyringBackendKind, @@ -678,13 +711,25 @@ impl OAuthPersistor { .await } + async fn refresh_if_needed_with_keyring_store( + &self, + keyring_store: &K, + ) -> Result<()> { + self.refresh_if_needed_with_keyring_store_and_timeout( + keyring_store, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + #[expect( clippy::await_holding_invalid_type, reason = "AuthorizationManager async access must be serialized through its mutex" )] - async fn refresh_if_needed_with_keyring_store( + async fn refresh_if_needed_with_keyring_store_and_timeout( &self, keyring_store: &K, + refresh_request_timeout: Duration, ) -> Result<()> { let expires_at = { let guard = self.inner.last_credentials.lock().await; @@ -751,14 +796,24 @@ impl OAuthPersistor { { let manager = self.inner.authorization_manager.clone(); let guard = manager.lock().await; - guard.refresh_token().await.with_context(|| { - format!( - "failed to refresh OAuth tokens for server {}", + match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(result) => { + result.with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", + self.inner.server_name + ) + })?; + } + Err(_) => anyhow::bail!( + "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", self.inner.server_name - ) - })?; + ), + } } + // Once the provider returns a rotated token, persistence must finish before the credential + // lock is released. In particular, caller startup deadlines must not cancel this step. self.persist_if_needed_with_keyring_store(keyring_store) .await } @@ -783,6 +838,10 @@ struct RefreshCredentialLock { impl RefreshCredentialLock { async fn acquire(store_key: &str) -> Result { + Self::acquire_with_timeout(store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT).await + } + + async fn acquire_with_timeout(store_key: &str, acquire_timeout: Duration) -> Result { let path = refresh_lock_path(store_key)?; if let Some(parent) = path.parent() { fs::create_dir_all(parent)?; @@ -796,18 +855,32 @@ impl RefreshCredentialLock { .open(&path) .with_context(|| format!("failed to open OAuth refresh lock {}", path.display()))?; - loop { - match file.try_lock() { - Ok(()) => break, - Err(std::fs::TryLockError::WouldBlock) => { - sleep(REFRESH_LOCK_RETRY_SLEEP).await; - } - Err(error) => { - return Err(std::io::Error::from(error)).with_context(|| { - format!("failed to lock OAuth refresh lock {}", path.display()) - }); + // Bound every contender, but keep the acquired lock for the full provider request and + // persistence transaction. Releasing it while awaiting the provider would allow concurrent + // use of a rotating refresh token. + match timeout(acquire_timeout, async { + loop { + match file.try_lock() { + Ok(()) => return Ok(()), + Err(std::fs::TryLockError::WouldBlock) => { + sleep(REFRESH_LOCK_RETRY_SLEEP).await; + } + Err(error) => return Err(std::io::Error::from(error)), } } + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + return Err(error).with_context(|| { + format!("failed to lock OAuth refresh lock {}", path.display()) + }); + } + Err(_) => anyhow::bail!( + "timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}", + path.display() + ), } Ok(Self { _file: file }) @@ -831,6 +904,9 @@ async fn install_tokens_in_manager( let manager = authorization_manager.clone(); let mut guard = manager.lock().await; guard.set_credential_store(store); + // TODO(stevenlee): RMCP's `initialize_from_store` updates the credential store and client ID + // but not its private `current_scopes`. Credential adoption can therefore leave scope-upgrade + // state stale until RMCP exposes an adoption API that synchronizes both. guard .initialize_from_store() .await @@ -1039,6 +1115,8 @@ fn refresh_lock_path(store_key: &str) -> Result { // Credential coordination is deliberately scoped to the active CODEX_HOME, alongside File // and Secrets state. Coordinating the process-global Direct keyring across distinct homes // would require a separately defined global lock namespace and is outside this transaction. + // TODO(stevenlee): define a safe per-user, cross-platform rendezvous before extending Direct + // keyring coordination across distinct CODEX_HOME values. let mut hasher = Sha256::new(); hasher.update(store_key.as_bytes()); let digest = hasher.finalize(); @@ -1250,7 +1328,7 @@ mod tests { &keyring_tokens, )?; - let loaded = super::load_oauth_tokens_with_keyring_store( + let loaded = super::load_oauth_tokens_with_source_and_keyring_store( &store, &keyring_tokens.server_name, &keyring_tokens.url, @@ -1401,6 +1479,35 @@ mod tests { Ok(()) } + #[tokio::test] + async fn refresh_lock_acquisition_times_out_without_stealing() -> Result<()> { + let _env = TempCodexHome::new(); + let store_key = "test-store-key"; + let held_lock = + RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(100)) + .await?; + + let error = + match RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(50)) + .await + { + Ok(_) => panic!("contending lock acquisition should time out"), + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("timed out after 50ms waiting for OAuth refresh lock"), + "unexpected error: {error:#}" + ); + + drop(held_lock); + let _reacquired = + RefreshCredentialLock::acquire_with_timeout(store_key, Duration::from_millis(100)) + .await?; + Ok(()) + } + #[tokio::test] async fn refresh_transaction_preserves_credentials_when_resolved_keyring_reread_fails() -> Result<()> { diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 39083ffb86..bb85b9c2e7 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -3,6 +3,7 @@ use std::ffi::OsString; use std::future::Future; use std::io; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -68,6 +69,7 @@ use crate::oauth::LoadedOAuthTokens; use crate::oauth::OAuthPersistor; use crate::oauth::ResolvedOAuthCredentialStore; use crate::oauth::StoredOAuthTokens; +use crate::oauth::load_oauth_tokens_from_resolved_store; use crate::oauth::load_oauth_tokens_with_source; use crate::oauth_http_client::OAuthHttpClientAdapter; use crate::stdio_server_launcher::StdioServerCommand; @@ -82,7 +84,6 @@ mod streamable_http_retry; use self::streamable_http_retry::HandshakeError; use self::streamable_http_retry::STREAMABLE_HTTP_RETRY_DELAYS_MS; -use self::streamable_http_retry::initialize_timeout_error; use self::streamable_http_retry::remaining_initialize_timeout; use self::streamable_http_retry::sleep_with_retry_deadline; @@ -130,6 +131,7 @@ enum TransportRecipe { env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, + resolved_store: Arc>, http_client: Arc, auth_provider: Option, }, @@ -406,6 +408,7 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + resolved_store: Arc::new(OnceLock::new()), http_client, auth_provider, }; @@ -448,11 +451,13 @@ impl RmcpClient { } }; + let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration); let (service, oauth_persistor) = self .connect_pending_transport_with_initialize_retries( pending_transport, client_service.clone(), timeout, + &mut initialize_deadline, ) .await?; @@ -785,6 +790,7 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + resolved_store, http_client, auth_provider, } => { @@ -801,16 +807,34 @@ impl RmcpClient { && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) { - match load_oauth_tokens_with_source( - server_name, - url, - *store_mode, - *keyring_backend_kind, - ) { - Ok(tokens) => tokens, - Err(err) => { - warn!("failed to read tokens for server `{server_name}`: {err}"); - None + if let Some(store) = resolved_store.get().copied() { + load_oauth_tokens_from_resolved_store(server_name, url, store)? + .map(|tokens| LoadedOAuthTokens { tokens, store }) + } else { + match load_oauth_tokens_with_source( + server_name, + url, + *store_mode, + *keyring_backend_kind, + ) { + Ok(tokens) => { + if let Some(loaded) = tokens.as_ref() { + // Transport retries and session recovery are part of the same + // client lifecycle. Pin the first concrete source in memory so + // rebuilding a transport never re-evaluates Auto and adopts a + // possibly stale credential from another store. + resolved_store.set(loaded.store).map_err(|_| { + anyhow!( + "OAuth credential store resolved concurrently for MCP server `{server_name}`" + ) + })?; + } + tokens + } + Err(err) => { + warn!("failed to read tokens for server `{server_name}`: {err}"); + None + } } } } else { @@ -892,11 +916,11 @@ impl RmcpClient { pending_transport: PendingTransport, client_service: ElicitationClientService, timeout: Option, + initialize_deadline: &mut Option, ) -> Result<( Arc>, Option, )> { - let deadline = timeout.map(|duration| Instant::now() + duration); let (transport, oauth_persistor) = match pending_transport { PendingTransport::InProcess { transport } => ( service::serve_client(client_service, transport).boxed(), @@ -914,14 +938,16 @@ impl RmcpClient { transport, oauth_persistor, } => { - match remaining_initialize_timeout(timeout, deadline)? { - Some(remaining) => { - time::timeout(remaining, oauth_persistor.refresh_if_needed()) - .await - .map_err(|_| initialize_timeout_error(timeout, remaining))??; - } - None => oauth_persistor.refresh_if_needed().await?, + // `startup_timeout_sec` bounds MCP transport setup, retry delays, and the + // initialization handshake. OAuth refresh has independent lock and provider + // request bounds, and persistence must finish after a successful response, so the + // complete refresh transaction is deliberately excluded from that deadline. + let refresh_started_at = Instant::now(); + let refresh_result = oauth_persistor.refresh_if_needed().await; + if let Some(deadline) = initialize_deadline.as_mut() { + *deadline += refresh_started_at.elapsed(); } + refresh_result?; ( service::serve_client(client_service, transport).boxed(), Some(oauth_persistor), @@ -929,7 +955,7 @@ impl RmcpClient { } }; - let handshake_timeout = remaining_initialize_timeout(timeout, deadline)?; + let handshake_timeout = remaining_initialize_timeout(timeout, *initialize_deadline)?; let service_result = match handshake_timeout { Some(duration) => match time::timeout(duration, transport).await { Ok(result) => { @@ -1148,11 +1174,15 @@ impl RmcpClient { .clone() .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; + let mut initialize_deadline = initialize_context + .timeout + .map(|duration| Instant::now() + duration); let (service, oauth_persistor) = self .connect_pending_transport_with_initialize_retries( pending_transport, initialize_context.client_service, initialize_context.timeout, + &mut initialize_deadline, ) .await?; diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 8df3358e7d..2794489b2a 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -28,6 +28,7 @@ impl RmcpClient { initial_transport: PendingTransport, client_service: ElicitationClientService, timeout: Option, + initialize_deadline: &mut Option, ) -> Result<( Arc>, Option, @@ -37,7 +38,6 @@ impl RmcpClient { PendingTransport::StreamableHttp { .. } | PendingTransport::StreamableHttpWithOAuth { .. } => true, }; - let retry_deadline = timeout.map(|duration| Instant::now() + duration); let mut pending_transport = Some(initial_transport); for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS @@ -50,7 +50,7 @@ impl RmcpClient { let transport = match pending_transport.take() { Some(transport) => transport, None => { - let remaining = remaining_initialize_timeout(timeout, retry_deadline)?; + let remaining = remaining_initialize_timeout(timeout, *initialize_deadline)?; match remaining { Some(remaining) => time::timeout( remaining, @@ -62,12 +62,11 @@ impl RmcpClient { } } }; - let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; - match Self::connect_pending_transport( transport, client_service.clone(), - attempt_timeout, + timeout, + initialize_deadline, ) .await { @@ -84,7 +83,7 @@ impl RmcpClient { error = %error, "streamable HTTP MCP initialize failed with a retryable error; retrying" ); - if !sleep_with_retry_deadline(delay, retry_deadline).await { + if !sleep_with_retry_deadline(delay, *initialize_deadline).await { let duration = timeout.unwrap_or(delay); return Err(anyhow!( "timed out handshaking with MCP server after {duration:?}" diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index 755af0823f..02866823f2 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,5 +1,7 @@ mod streamable_http_test_support; +use std::fs; +use std::path::PathBuf; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -33,19 +35,22 @@ use wiremock::matchers::method; use wiremock::matchers::path; use streamable_http_test_support::initialize_client; +use streamable_http_test_support::initialize_client_with_timeout; const SERVER_NAME: &str = "test-streamable-http-oauth-startup"; const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token"; const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token"; +const CHILD_READY_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_READY_FILE"; +const CHILD_RELEASE_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_RELEASE_FILE"; const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL"; const UNREFRESHABLE_SERVER_URL: &str = "https://unrefreshable.example/mcp"; const UNEXPIRED_SERVER_URL: &str = "https://unexpired.example/mcp"; const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { +async fn startup_refresh_does_not_consume_handshake_timeout() -> anyhow::Result<()> { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/.well-known/oauth-authorization-server/mcp")) @@ -63,12 +68,18 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result .and(body_string_contains(format!( "refresh_token={REFRESH_TOKEN}" ))) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "access_token": REFRESHED_ACCESS_TOKEN, - "token_type": "Bearer", - "expires_in": 7200, - "refresh_token": REFRESH_TOKEN, - }))) + // The provider takes longer than the configured MCP handshake timeout. Refresh has its own + // bound, so this delay must not leave the subsequent handshake with an expired budget. + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(1_500)) + .set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": REFRESH_TOKEN, + })), + ) .expect(1) .mount(&server) .await; @@ -150,6 +161,8 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { "refresh_token": ROTATED_REFRESH_TOKEN, })), ) + // This exact expectation proves that the two child processes issue one serialized provider + // refresh rather than both replaying the same rotating refresh token. .expect(1) .mount(&server) .await; @@ -199,17 +212,36 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> { "OAuth concurrency seed child failed: {seed_status}" ); - let first_status = Command::new(std::env::current_exe()?) + let release_file = codex_home.path().join("oauth-clients.release"); + let first_ready_file = codex_home.path().join("oauth-client-first.ready"); + let second_ready_file = codex_home.path().join("oauth-client-second.ready"); + let mut first_child = Command::new(std::env::current_exe()?) .args(["oauth_concurrency_client_child", "--exact", "--ignored"]) .env("CODEX_HOME", codex_home.path()) .env(CHILD_SERVER_URL_ENV, &server_url) - .status(); - let second_status = Command::new(std::env::current_exe()?) + .env(CHILD_READY_FILE_ENV, &first_ready_file) + .env(CHILD_RELEASE_FILE_ENV, &release_file) + .kill_on_drop(true) + .spawn()?; + let mut second_child = Command::new(std::env::current_exe()?) .args(["oauth_concurrency_client_child", "--exact", "--ignored"]) .env("CODEX_HOME", codex_home.path()) .env(CHILD_SERVER_URL_ENV, &server_url) - .status(); - let (first_status, second_status) = tokio::try_join!(first_status, second_status)?; + .env(CHILD_READY_FILE_ENV, &second_ready_file) + .env(CHILD_RELEASE_FILE_ENV, &release_file) + .kill_on_drop(true) + .spawn()?; + + tokio::time::timeout(Duration::from_secs(10), async { + while !first_ready_file.exists() || !second_ready_file.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .map_err(|_| anyhow::anyhow!("OAuth concurrency children did not become ready"))?; + fs::write(&release_file, b"release")?; + + let (first_status, second_status) = tokio::try_join!(first_child.wait(), second_child.wait())?; assert!( first_status.success(), "first OAuth concurrency child failed: {first_status}" @@ -335,7 +367,7 @@ async fn auth_status(server_url: &str) -> anyhow::Result { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] +#[ignore = "spawned by startup_refresh_does_not_consume_handshake_timeout"] async fn oauth_startup_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; @@ -378,7 +410,7 @@ async fn oauth_startup_child() -> anyhow::Result<()> { ) .await?; - initialize_client(&client).await?; + initialize_client_with_timeout(&client, Duration::from_secs(1)).await?; Ok(()) } @@ -394,6 +426,8 @@ async fn oauth_concurrency_seed_child() -> anyhow::Result<()> { #[ignore = "spawned by concurrent_file_mode_startup_refreshes_once"] async fn oauth_concurrency_client_child() -> anyhow::Result<()> { let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let ready_file = PathBuf::from(std::env::var(CHILD_READY_FILE_ENV)?); + let release_file = PathBuf::from(std::env::var(CHILD_RELEASE_FILE_ENV)?); let client = RmcpClient::new_streamable_http_client( SERVER_NAME, &server_url, @@ -407,6 +441,12 @@ async fn oauth_concurrency_client_child() -> anyhow::Result<()> { ) .await?; + // Both processes must construct their OAuth client from the same expired snapshot before + // either begins refresh. This removes process scheduling from the contention setup. + fs::write(ready_file, b"ready")?; + while !release_file.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } initialize_client(&client).await?; Ok(()) } diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index 87a52f55df..532861de7d 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -106,10 +106,17 @@ pub(crate) async fn create_client_with_http_client( } pub(crate) async fn initialize_client(client: &RmcpClient) -> anyhow::Result<()> { + initialize_client_with_timeout(client, Duration::from_secs(5)).await +} + +pub(crate) async fn initialize_client_with_timeout( + client: &RmcpClient, + timeout: Duration, +) -> anyhow::Result<()> { client .initialize( init_params(), - Some(Duration::from_secs(5)), + Some(timeout), Box::new(|_, _| { async { Ok(ElicitationResponse {