diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index fd0b1a0e4fff..c498a2582cb5 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -29,7 +29,7 @@ use codex_mcp::oauth_login_support; use codex_mcp::resolve_oauth_scopes; use codex_mcp::should_retry_without_scopes; use codex_protocol::protocol::McpAuthStatus; -use codex_rmcp_client::delete_oauth_tokens; +use codex_rmcp_client::delete_oauth_tokens_locked; use codex_rmcp_client::perform_oauth_login; use codex_utils_cli::CliConfigOverrides; use codex_utils_cli::format_env_display; @@ -523,12 +523,14 @@ async fn run_logout(config_overrides: &CliConfigOverrides, logout_args: LogoutAr _ => bail!("OAuth logout is only supported for streamable_http transports."), }; - match delete_oauth_tokens( + match delete_oauth_tokens_locked( &name, &url, config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind(), - ) { + ) + .await + { Ok(true) => println!("Removed OAuth credentials for '{name}'."), Ok(false) => println!("No OAuth credentials stored for '{name}'."), Err(err) => return Err(anyhow!("failed to delete OAuth credentials: {err}")), diff --git a/codex-rs/cli/tests/mcp_oauth_logout.rs b/codex-rs/cli/tests/mcp_oauth_logout.rs new file mode 100644 index 000000000000..f43453937c30 --- /dev/null +++ b/codex-rs/cli/tests/mcp_oauth_logout.rs @@ -0,0 +1,80 @@ +use std::path::Path; + +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::save_oauth_tokens; +use predicates::str::contains; +use serde_json::json; +use tempfile::TempDir; + +const SERVER_NAME: &str = "oauth-server"; +const SERVER_URL: &str = "https://example.com/mcp"; + +fn codex_command(codex_home: &Path) -> Result { + let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?); + cmd.env("CODEX_HOME", codex_home); + Ok(cmd) +} + +#[tokio::test] +async fn mcp_logout_cli_removes_file_credentials() -> Result<()> { + let codex_home = TempDir::new()?; + let status = tokio::process::Command::new(std::env::current_exe()?) + .args([ + "mcp_logout_cli_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .status() + .await?; + anyhow::ensure!(status.success(), "MCP logout child failed: {status}"); + Ok(()) +} + +#[tokio::test] +#[ignore = "spawned by mcp_logout_cli_removes_file_credentials"] +async fn mcp_logout_cli_child() -> Result<()> { + let codex_home = std::env::var("CODEX_HOME")?; + let codex_home = Path::new(&codex_home); + std::fs::write( + codex_home.join("config.toml"), + format!( + "mcp_oauth_credentials_store = \"file\"\n\n[mcp_servers.{SERVER_NAME}]\nurl = \"{SERVER_URL}\"\n" + ), + )?; + + let tokens: StoredOAuthTokens = serde_json::from_value(json!({ + "server_name": SERVER_NAME, + "url": SERVER_URL, + "client_id": "test-client-id", + "token_response": { + "access_token": "access-token", + "token_type": "bearer", + "expires_in": 3600, + "refresh_token": "refresh-token", + }, + "expires_at": null, + }))?; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + assert!(codex_home.join(".credentials.json").exists()); + + let mut logout = codex_command(codex_home)?; + logout + .args(["mcp", "logout", SERVER_NAME]) + .assert() + .success() + .stdout(contains(format!( + "Removed OAuth credentials for '{SERVER_NAME}'." + ))); + assert!(!codex_home.join(".credentials.json").exists()); + Ok(()) +} diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index 5303fccd43b6..a66ff3f3a711 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -103,7 +103,9 @@ pub enum AuthCredentialsStoreMode { #[derive(Debug, Default, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "lowercase")] pub enum OAuthCredentialsStoreMode { - /// `Keyring` when available; otherwise, `File`. + /// Prefer `Keyring` and use `File` when keyring storage is unavailable. + /// Once an MCP client loads credentials from one store, that client keeps the resolved store + /// for its lifetime so refreshes cannot switch to a possibly stale credential source. /// Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access. #[default] Auto, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 0a48882b09c7..0bc23730af34 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2129,7 +2129,7 @@ "description": "Determine where Codex should store and read MCP credentials.", "oneOf": [ { - "description": "`Keyring` when available; otherwise, `File`. Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access.", + "description": "Prefer `Keyring` and use `File` when keyring storage is unavailable. Once an MCP client loads credentials from one store, that client keeps the resolved store for its lifetime so refreshes cannot switch to a possibly stale credential source. Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access.", "enum": [ "auto" ], diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index 517f7270bbe0..a8464f3054d9 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -84,5 +84,10 @@ keyring = { workspace = true, features = ["windows-native"] } [target.'cfg(any(target_os = "freebsd", target_os = "openbsd"))'.dependencies] keyring = { workspace = true, features = ["sync-secret-service"] } +# This test is compiled through `#[path]` inside the inline `oauth::tests` module. Cargo-shear +# cannot resolve that nested module path and otherwise reports the linked file as unlinked. +[package.metadata.cargo-shear] +ignored-paths = ["src/oauth/tests/persistor_tests.rs"] + [lib] doctest = false diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 19befb62355e..eccce8c03dfd 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -22,6 +22,7 @@ use codex_exec_server::HttpResponseBodyStream; use futures::StreamExt; use futures::stream; use futures::stream::BoxStream; +use oauth2::AccessToken; use reqwest::StatusCode; use reqwest::header::ACCEPT; use reqwest::header::AUTHORIZATION; @@ -55,12 +56,17 @@ pub(crate) struct StreamableHttpClientAdapter { http_client: Arc, default_headers: HeaderMap, auth_provider: Option, + attribute_rejected_access_token: bool, } #[derive(Debug, thiserror::Error)] pub(crate) enum StreamableHttpClientAdapterError { #[error("streamable HTTP session expired with 404 Not Found")] SessionExpired404, + #[error("MCP server rejected the access token with HTTP 401 Unauthorized")] + AccessTokenRejected { rejected_access_token: AccessToken }, + #[error("MCP OAuth operation failed: {0:#}")] + OAuth(#[source] anyhow::Error), #[error(transparent)] HttpRequest(#[from] ExecServerError), #[error("invalid HTTP header: {0}")] @@ -77,8 +83,15 @@ impl StreamableHttpClientAdapter { http_client, default_headers, auth_provider, + attribute_rejected_access_token: false, } } + + /// Preserves the access token associated with a 401 for Codex-owned OAuth recovery. + pub(crate) fn with_rejected_token_attribution(mut self) -> Self { + self.attribute_rejected_access_token = true; + self + } } impl StreamableHttpClient for StreamableHttpClientAdapter { @@ -109,7 +122,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { JSON_MIME_TYPE.to_string(), StreamableHttpClientAdapterError::Header, )?; - if let Some(auth_token) = auth_token { + if let Some(auth_token) = auth_token.as_deref() { insert_header( &mut headers, AUTHORIZATION, @@ -162,6 +175,12 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } + if self.attribute_rejected_access_token + && response.status == StatusCode::UNAUTHORIZED.as_u16() + && let Some(error) = access_token_rejected(auth_token.as_deref()) + { + return Err(error); + } if response.status == StatusCode::UNAUTHORIZED.as_u16() && let Some(header) = response_header(&response.headers, reqwest::header::WWW_AUTHENTICATE) @@ -240,7 +259,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { let mut headers = self.default_headers.clone(); headers.extend(custom_headers); self.add_auth_headers(&mut headers); - if let Some(auth_token) = auth_token { + if let Some(auth_token) = auth_token.as_deref() { insert_header( &mut headers, AUTHORIZATION, @@ -274,6 +293,12 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { if response.status == StatusCode::METHOD_NOT_ALLOWED.as_u16() { return Ok(()); } + if self.attribute_rejected_access_token + && response.status == StatusCode::UNAUTHORIZED.as_u16() + && let Some(error) = access_token_rejected(auth_token.as_deref()) + { + return Err(error); + } if !status_is_success(response.status) { return Err(StreamableHttpError::UnexpectedServerResponse( format!("DELETE returned HTTP {}", response.status).into(), @@ -316,7 +341,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::Header, )?; } - if let Some(auth_token) = auth_token { + if let Some(auth_token) = auth_token.as_deref() { insert_header( &mut headers, AUTHORIZATION, @@ -349,6 +374,12 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpClientAdapterError::SessionExpired404, )); } + if self.attribute_rejected_access_token + && response.status == StatusCode::UNAUTHORIZED.as_u16() + && let Some(error) = access_token_rejected(auth_token.as_deref()) + { + return Err(error); + } if !status_is_success(response.status) { return Err(StreamableHttpError::UnexpectedServerResponse( format!("GET returned HTTP {}", response.status).into(), @@ -371,6 +402,20 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } } +fn access_token_rejected( + auth_token: Option<&str>, +) -> Option> { + // Preserve the token associated with this response. Reading the current credential after a + // delayed 401 is racy: another concurrent request may already have refreshed A to B, in which + // case recovery must retry B rather than refresh B a second time. AccessToken's Debug + // implementation redacts the secret if this error is logged. + auth_token.map(|rejected_access_token| { + StreamableHttpError::Client(StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token: AccessToken::new(rejected_access_token.to_string()), + }) + }) +} + impl StreamableHttpClientAdapter { fn add_auth_headers(&self, headers: &mut HeaderMap) { if let Some(auth_provider) = &self.auth_provider { diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index e57eb3686c43..ecdce2c29b5f 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -6,6 +6,7 @@ mod in_process_transport; mod logging_client_handler; mod oauth; mod oauth_http_client; +mod oauth_transport; mod perform_oauth_login; mod program_resolver; mod rmcp_client; @@ -26,7 +27,7 @@ pub use in_process_transport::InProcessTransportFactory; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; -pub(crate) use oauth::load_oauth_tokens; +pub use oauth::delete_oauth_tokens_locked; pub use oauth::save_oauth_tokens; pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 46e6279a161e..11682a5fdc9d 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -16,6 +16,9 @@ //! //! If the keyring is not available or fails, we fall back to CODEX_HOME/.credentials.json which is consistent with other coding CLI agents. +mod refresh_lock; +mod refresh_transaction; +mod resolved_store; mod store_lock; #[cfg(test)] @@ -66,6 +69,12 @@ use tokio::sync::Mutex; use codex_utils_home_dir::find_codex_home; +use self::refresh_lock::RefreshCredentialLock; +pub(crate) use self::refresh_transaction::request_oauth_token_response; +pub(crate) use self::resolved_store::ResolvedOAuthCredentialStore; +pub(crate) use self::resolved_store::ResolvedOAuthTokens; +pub(crate) use self::resolved_store::resolve_oauth_tokens_from_store_policy; + const KEYRING_SERVICE: &str = "Codex MCP Credentials"; const MCP_OAUTH_SECRET_PREFIX: &str = "MCP_OAUTH"; const REFRESH_SKEW_MILLIS: u64 = 30_000; @@ -100,41 +109,24 @@ pub(crate) enum StoredOAuthTokenStatus { AuthorizationRequired, } -pub(crate) fn load_oauth_tokens( - server_name: &str, - url: &str, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, -) -> Result> { - let keyring_store = DefaultKeyringStore; - match store_mode { - OAuthCredentialsStoreMode::Auto => load_oauth_tokens_from_keyring_with_fallback_to_file( - &keyring_store, - keyring_backend_kind, - server_name, - url, - ), - OAuthCredentialsStoreMode::File => load_oauth_tokens_from_file(server_name, url), - OAuthCredentialsStoreMode::Keyring => { - load_oauth_tokens_from_keyring(&keyring_store, keyring_backend_kind, server_name, url) - .with_context(|| "failed to read OAuth tokens from keyring".to_string()) - } - } -} - pub(crate) fn oauth_token_status( server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> Result { - Ok( - match load_oauth_tokens(server_name, url, store_mode, keyring_backend_kind)?.as_ref() { - None => StoredOAuthTokenStatus::Missing, - Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, - Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, - }, - ) + let resolved = resolve_oauth_tokens_from_store_policy( + &DefaultKeyringStore, + server_name, + url, + store_mode, + keyring_backend_kind, + )?; + Ok(match resolved.as_ref().map(|resolved| &resolved.tokens) { + None => StoredOAuthTokenStatus::Missing, + Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, + Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, + }) } fn oauth_tokens_are_usable(tokens: &StoredOAuthTokens) -> bool { @@ -174,37 +166,16 @@ fn refresh_expires_in_from_timestamp(tokens: &mut StoredOAuthTokens) { } } -fn load_oauth_tokens_from_keyring_with_fallback_to_file( - keyring_store: &K, - keyring_backend_kind: AuthKeyringBackendKind, - server_name: &str, - url: &str, -) -> Result> { - match load_oauth_tokens_from_keyring(keyring_store, keyring_backend_kind, server_name, url) { - Ok(Some(tokens)) => Ok(Some(tokens)), - Ok(None) => load_oauth_tokens_from_file(server_name, url), - // A store lock failure means the configured aggregate authority could be changing, or - // that coordination itself is unavailable. It is not evidence that the keyring backend - // is unavailable, so consulting File here could replay credentials hidden behind a - // newer Secrets entry. This is the load-side counterpart of the save guard below. - Err(error) if error.downcast_ref::().is_some() => Err(error), - Err(error) => { - warn!("failed to read OAuth tokens from keyring: {error}"); - load_oauth_tokens_from_file(server_name, url) - .with_context(|| format!("failed to read OAuth tokens from keyring: {error}")) - } - } -} - fn load_oauth_tokens_from_keyring( keyring_store: &K, keyring_backend_kind: AuthKeyringBackendKind, server_name: &str, url: &str, -) -> Result> { +) -> std::result::Result, OAuthKeyringLoadError> { match keyring_backend_kind { AuthKeyringBackendKind::Direct => { load_oauth_tokens_from_direct_keyring(keyring_store, server_name, url) + .map_err(OAuthKeyringLoadError::Backend) } AuthKeyringBackendKind::Secrets => { load_oauth_tokens_from_secrets_keyring(keyring_store, server_name, url) @@ -234,9 +205,9 @@ fn load_oauth_tokens_from_secrets_keyring( keyring_store: &K, server_name: &str, url: &str, -) -> Result> { +) -> std::result::Result, OAuthKeyringLoadError> { let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; - let codex_home = find_codex_home()?; + let codex_home = find_codex_home().map_err(anyhow::Error::from)?; let manager = SecretsManager::new_with_keyring_store_and_namespace( codex_home.to_path_buf(), SecretsBackendKind::Local, @@ -258,6 +229,17 @@ fn load_oauth_tokens_from_secrets_keyring( } } +/// Classifies keyring load failures that affect Auto fallback policy. +#[derive(Debug, thiserror::Error)] +enum OAuthKeyringLoadError { + /// Store coordination failed, so consulting another authority would be unsafe. + #[error(transparent)] + StoreLock(#[from] OAuthStoreLockFailure), + /// The selected keyring backend itself was unavailable or its data was invalid. + #[error(transparent)] + Backend(#[from] anyhow::Error), +} + pub fn save_oauth_tokens( server_name: &str, tokens: &StoredOAuthTokens, @@ -265,16 +247,51 @@ pub fn save_oauth_tokens( keyring_backend_kind: AuthKeyringBackendKind, ) -> Result<()> { let keyring_store = DefaultKeyringStore; + save_oauth_tokens_with_keyring_store( + &keyring_store, + server_name, + tokens, + store_mode, + keyring_backend_kind, + ) +} + +pub(crate) async fn save_oauth_tokens_locked( + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result<()> { + // Login persistence shares the refresh transaction lock so a completed login always becomes + // authoritative: it either lands before refresh's reread or waits and overwrites the refresh + // result afterward. + let _lock = RefreshCredentialLock::acquire_for_server(server_name, &tokens.url).await?; + save_oauth_tokens_with_keyring_store( + &DefaultKeyringStore, + server_name, + tokens, + store_mode, + keyring_backend_kind, + ) +} + +fn save_oauth_tokens_with_keyring_store( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result<()> { match store_mode { OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file( - &keyring_store, + keyring_store, keyring_backend_kind, server_name, tokens, ), OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file(tokens), - OAuthCredentialsStoreMode::Keyring => save_oauth_tokens_with_keyring( - &keyring_store, + OAuthCredentialsStoreMode::Keyring => save_oauth_tokens_with_keyring_and_cleanup_file( + keyring_store, keyring_backend_kind, server_name, tokens, @@ -288,6 +305,8 @@ fn save_oauth_tokens_with_keyring( server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { + // This exact-store writer is used after a client resolves its authority. Only login-time + // policy resolution may clean up or update the non-selected store. match keyring_backend_kind { AuthKeyringBackendKind::Direct => { save_oauth_tokens_to_direct_keyring(keyring_store, server_name, tokens) @@ -307,12 +326,7 @@ fn save_oauth_tokens_to_direct_keyring( let key = compute_store_key(server_name, &tokens.url)?; match keyring_store.save(KEYRING_SERVICE, &key, &serialized) { - Ok(()) => { - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) - } + Ok(()) => Ok(()), Err(error) => { let message = format!( "failed to write OAuth tokens to keyring: {}", @@ -325,29 +339,19 @@ fn save_oauth_tokens_to_direct_keyring( } /// Saves one credential while holding the Secrets aggregate-store lock across the mutation. -/// -/// The Secrets lock is released before fallback File cleanup to preserve aggregate-lock ordering. fn save_oauth_tokens_to_secrets_keyring( keyring_store: &K, server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; - { - let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; - save_oauth_tokens_to_secrets_keyring_with_lock_held( - keyring_store, - server_name, - tokens, - &serialized, - )?; - } - - let key = compute_store_key(server_name, &tokens.url)?; - if let Err(error) = delete_oauth_tokens_from_file(&key) { - warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); - } - Ok(()) + let _store_lock = OAuthStoreLock::acquire(OAuthStore::Secrets)?; + save_oauth_tokens_to_secrets_keyring_with_lock_held( + keyring_store, + server_name, + tokens, + &serialized, + ) } /// Writes one credential to Secrets. The caller must hold the Secrets aggregate-store lock. @@ -370,13 +374,38 @@ fn save_oauth_tokens_to_secrets_keyring_with_lock_held( + keyring_store: &K, + keyring_backend_kind: AuthKeyringBackendKind, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result<()> { + save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens)?; + let key = compute_store_key(server_name, &tokens.url)?; + if let Err(error) = delete_oauth_tokens_from_file(&key) { + warn!( + server_name, + keyring_backend = ?keyring_backend_kind, + error = %error, + "failed to remove OAuth tokens from fallback storage" + ); + } + Ok(()) +} + fn save_oauth_tokens_with_keyring_with_fallback_to_file( keyring_store: &K, keyring_backend_kind: AuthKeyringBackendKind, server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { - match save_oauth_tokens_with_keyring(keyring_store, keyring_backend_kind, server_name, tokens) { + match save_oauth_tokens_with_keyring_and_cleanup_file( + keyring_store, + keyring_backend_kind, + server_name, + tokens, + ) { Ok(()) => Ok(()), // As on load, a store lock failure is a coordination failure rather than evidence that // the keyring backend is unavailable. Falling back could leave a newer File token hidden @@ -407,6 +436,24 @@ pub fn delete_oauth_tokens( ) } +pub async fn delete_oauth_tokens_locked( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result { + // Logout shares the refresh transaction lock so refresh cannot resurrect credentials after a + // completed delete: it either observes the deletion or finishes before logout removes it. + let _lock = RefreshCredentialLock::acquire_for_server(server_name, url).await?; + delete_oauth_tokens_from_keyring_and_file( + &DefaultKeyringStore, + store_mode, + keyring_backend_kind, + server_name, + url, + ) +} + fn delete_oauth_tokens_from_keyring_and_file( keyring_store: &K, store_mode: OAuthCredentialsStoreMode, @@ -495,8 +542,7 @@ struct OAuthPersistorInner { server_name: String, url: String, authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + credential_store: ResolvedOAuthCredentialStore, last_credentials: Mutex>, } @@ -505,8 +551,7 @@ impl OAuthPersistor { server_name: String, url: String, authorization_manager: Arc>, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + credential_store: ResolvedOAuthCredentialStore, initial_credentials: Option, ) -> Self { Self { @@ -514,106 +559,12 @@ impl OAuthPersistor { server_name, url, authorization_manager, - store_mode, - keyring_backend_kind, + credential_store, last_credentials: Mutex::new(initial_credentials), }), } } - - /// Persists the latest stored credentials if they have changed. - /// Deletes the credentials if they are no longer present. - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] - pub(crate) async fn persist_if_needed(&self) -> Result<()> { - let (client_id, maybe_credentials) = { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.get_credentials().await - }?; - - match maybe_credentials { - Some(credentials) => { - let mut last_credentials = self.inner.last_credentials.lock().await; - let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); - let same_token = last_credentials - .as_ref() - .map(|prev| prev.token_response == new_token_response) - .unwrap_or(false); - let expires_at = if same_token { - last_credentials.as_ref().and_then(|prev| prev.expires_at) - } else { - compute_expires_at_millis(&credentials) - }; - let stored = StoredOAuthTokens { - server_name: self.inner.server_name.clone(), - url: self.inner.url.clone(), - client_id, - token_response: new_token_response, - expires_at, - }; - if last_credentials.as_ref() != Some(&stored) { - save_oauth_tokens( - &self.inner.server_name, - &stored, - self.inner.store_mode, - self.inner.keyring_backend_kind, - )?; - *last_credentials = Some(stored); - } - } - None => { - let mut last_serialized = self.inner.last_credentials.lock().await; - if last_serialized.take().is_some() - && let Err(error) = delete_oauth_tokens( - &self.inner.server_name, - &self.inner.url, - self.inner.store_mode, - self.inner.keyring_backend_kind, - ) - { - warn!( - "failed to remove OAuth tokens for server {}: {error}", - self.inner.server_name - ); - } - } - } - - Ok(()) - } - - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] - pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - let expires_at = { - let guard = self.inner.last_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) - }; - - if !token_needs_refresh(expires_at) { - return Ok(()); - } - - { - 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 {}", - self.inner.server_name - ) - })?; - } - - self.persist_if_needed().await - } } - const FALLBACK_FILENAME: &str = ".credentials.json"; const MCP_SERVER_TYPE: &str = "http"; @@ -871,11 +822,13 @@ mod tests { use keyring::Error as KeyringError; use pretty_assertions::assert_eq; use std::sync::Arc; + #[path = "persistor_tests.rs"] + mod persistor_tests; use super::test_support::TempCodexHome; #[test] - fn load_oauth_tokens_reads_from_keyring_when_available() -> Result<()> { + fn resolve_oauth_tokens_from_store_policy_uses_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); let store = MockKeyringStore::default(); let tokens = sample_tokens(); @@ -884,14 +837,19 @@ mod tests { let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; store.save(KEYRING_SERVICE, &key, &serialized)?; - let loaded = super::load_oauth_tokens_from_keyring( + let resolved = super::resolve_oauth_tokens_from_store_policy( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from keyring"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_eq!( + resolved.store, + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) + ); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } @@ -904,14 +862,16 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + let resolved = super::resolve_oauth_tokens_from_store_policy( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); Ok(()) } @@ -926,14 +886,45 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + let resolved = super::resolve_oauth_tokens_from_store_policy( &store, - AuthKeyringBackendKind::Direct, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, )? .expect("tokens should load from fallback"); - assert_tokens_match_without_expiry(&loaded, &expected); + assert_eq!(resolved.store, ResolvedOAuthCredentialStore::File); + assert_tokens_match_without_expiry(&resolved.tokens, &expected); + Ok(()) + } + + #[test] + fn exact_store_operations_do_not_adopt_or_mutate_the_other_store() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockKeyringStore::default(); + let file_tokens = sample_tokens(); + let mut keyring_tokens = file_tokens.clone(); + keyring_tokens + .token_response + .0 + .set_access_token(AccessToken::new("keyring-access-token".to_string())); + + super::save_oauth_tokens_to_file(&file_tokens)?; + let fallback_path = super::fallback_file_path()?; + let fallback_before = fs::read(&fallback_path)?; + super::save_oauth_tokens_with_keyring( + &store, + AuthKeyringBackendKind::Direct, + &keyring_tokens.server_name, + &keyring_tokens, + )?; + + assert_eq!(fs::read(fallback_path)?, fallback_before); + let loaded = ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct) + .load(&store, &keyring_tokens.server_name, &keyring_tokens.url)? + .expect("tokens should load from the selected keyring store"); + assert_tokens_match_without_expiry(&loaded, &keyring_tokens); Ok(()) } diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs new file mode 100644 index 000000000000..9d1ecf9e60cc --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock.rs @@ -0,0 +1,104 @@ +//! Cross-process serialization for one MCP OAuth credential's refresh transaction. +//! +//! The guard is intentionally acquired before the authoritative credential reread and retained +//! through provider refresh and persistence. This prevents two processes from replaying the same +//! rotating refresh token or observing a partially persisted transaction. + +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use codex_utils_home_dir::find_codex_home; +use sha2::Digest; +use sha2::Sha256; +use std::fs; +use std::fs::File; +use std::fs::OpenOptions; +use std::path::Path; +use std::time::Duration; +use tokio::time::sleep; +use tokio::time::timeout; + +const REFRESH_LOCK_DIR: &str = "mcp-oauth-locks"; +const REFRESH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60); +const REFRESH_LOCK_RETRY_SLEEP: Duration = Duration::from_millis(/*millis*/ 50); +// Keep this internal target stable so diagnostics and cross-process tests can distinguish actual +// WouldBlock contention from a contender that merely started late and observed persisted tokens. +const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention"; + +pub(super) struct RefreshCredentialLock { + _file: File, +} + +impl RefreshCredentialLock { + pub(super) async fn acquire_for_server(server_name: &str, url: &str) -> Result { + let store_key = super::compute_store_key(server_name, url)?; + let codex_home = find_codex_home()?; + Self::acquire_in(&codex_home, &store_key, REFRESH_LOCK_ACQUIRE_TIMEOUT) + .await + .with_context(|| format!("failed to acquire OAuth credential lock for {server_name}")) + } + + async fn acquire_in( + codex_home: &Path, + store_key: &str, + acquire_timeout: Duration, + ) -> Result { + // Scope coordination to CODEX_HOME alongside File and Secrets state. Direct keyring + // coordination across homes needs a separate cross-platform rendezvous. + // TODO(stevenlee): define that rendezvous before expanding this lock's scope. + let mut hasher = Sha256::new(); + hasher.update(store_key.as_bytes()); + let path = codex_home + .join(REFRESH_LOCK_DIR) + .join(format!("{:x}.lock", hasher.finalize())); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&path) + .with_context(|| format!("failed to open 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. + let mut reported_contention = false; + timeout(acquire_timeout, async { + loop { + match file.try_lock() { + Ok(()) => return Ok(()), + Err(std::fs::TryLockError::WouldBlock) => { + if !reported_contention { + tracing::debug!( + target: LOCK_CONTENTION_EVENT_TARGET, + lock_path = %path.display(), + "waiting for another process to finish refreshing MCP OAuth credentials" + ); + reported_contention = true; + } + sleep(REFRESH_LOCK_RETRY_SLEEP).await; + } + Err(error) => return Err(std::io::Error::from(error)), + } + } + }) + .await + .map_err(|_| { + anyhow!( + "timed out after {acquire_timeout:?} waiting for OAuth refresh lock {}", + path.display() + ) + })? + .with_context(|| format!("failed to lock OAuth refresh lock {}", path.display()))?; + + Ok(Self { _file: file }) + } +} + +#[cfg(test)] +#[path = "refresh_lock_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs new file mode 100644 index 000000000000..6344bd036cbd --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_lock_tests.rs @@ -0,0 +1,40 @@ +use super::RefreshCredentialLock; +use anyhow::Result; +use std::time::Duration; +use tempfile::tempdir; + +#[tokio::test] +async fn acquisition_times_out_without_stealing() -> Result<()> { + let codex_home = tempdir()?; + let store_key = "test-store-key"; + let held_lock = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 100), + ) + .await?; + + let error = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 50), + ) + .await + .err() + .expect("contending lock acquisition should time out"); + assert!( + error + .to_string() + .contains("timed out after 50ms waiting for OAuth refresh lock"), + "unexpected error: {error:#}" + ); + + drop(held_lock); + let _reacquired = RefreshCredentialLock::acquire_in( + codex_home.path(), + store_key, + Duration::from_millis(/*millis*/ 100), + ) + .await?; + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs new file mode 100644 index 000000000000..9bd03cc401a9 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/refresh_transaction.rs @@ -0,0 +1,392 @@ +//! Serialized read-refresh-write transactions for MCP OAuth credentials. + +use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Context; +use anyhow::Result; +use codex_keyring_store::DefaultKeyringStore; +use codex_keyring_store::KeyringStore; +use oauth2::AccessToken; +use oauth2::TokenResponse; +use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::CredentialStore as _; +use rmcp::transport::auth::InMemoryCredentialStore; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::StoredCredentials; +use tokio::time::timeout; +use tracing::debug; +use tracing::warn; + +use super::OAuthPersistor; +use super::OAuthPersistorInner; +use super::StoredOAuthTokens; +use super::WrappedOAuthTokenResponse; +use super::compute_expires_at_millis; +use super::refresh_lock::RefreshCredentialLock; +use super::token_needs_refresh; + +const REFRESH_REQUEST_TIMEOUT: Duration = Duration::from_secs(45); + +impl OAuthPersistor { + pub(crate) async fn refresh_if_needed(&self) -> Result<()> { + self.refresh_in( + DefaultKeyringStore, + RefreshReason::Expiry, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + + pub(crate) async fn refresh_after_unauthorized( + &self, + rejected_access_token: AccessToken, + ) -> Result<()> { + self.refresh_in( + DefaultKeyringStore, + RefreshReason::Unauthorized { + rejected_access_token, + }, + REFRESH_REQUEST_TIMEOUT, + ) + .await + } + + /// Injects the credential backend and provider timeout for deterministic failure-path tests. + pub(super) async fn refresh_in( + &self, + keyring_store: K, + reason: RefreshReason, + refresh_request_timeout: Duration, + ) -> Result<()> { + if matches!(reason, RefreshReason::Expiry) { + let expires_at = { + let guard = self.inner.last_credentials.lock().await; + guard.as_ref().and_then(|tokens| tokens.expires_at) + }; + if !token_needs_refresh(expires_at) { + return Ok(()); + } + } + + let persistor = self.clone(); + // Once the provider can consume a rotating token, caller cancellation must not cancel + // persistence. The owned task continues with independently bounded lock and request waits. + // A provider timeout leaves the outcome unknown and permits a later serialized retry: + // provider grace may recover, otherwise reauthorization is unavoidable. This residual + // risk is preferred to holding the credential lock indefinitely. + let refresh_reason = reason.as_str(); + let transaction_task = tokio::spawn(async move { + let result = persistor + .refresh_transaction(&keyring_store, reason, refresh_request_timeout) + .await; + + // Keep this summary inside the owned task so caller cancellation cannot suppress it. + if let Err(error) = &result { + warn!( + server_name = %persistor.inner.server_name, + refresh_reason, + error = %error, + "MCP OAuth refresh transaction failed" + ); + } + + result + }); + transaction_task.await.with_context(|| { + format!( + "OAuth refresh task failed for server {}", + self.inner.server_name + ) + })? + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its Tokio mutex" + )] + #[tracing::instrument( + level = "debug", + skip_all, + fields( + server_name = %self.inner.server_name, + refresh_reason = reason.as_str(), + ), + err + )] + async fn refresh_transaction( + &self, + keyring_store: &K, + reason: RefreshReason, + refresh_request_timeout: Duration, + ) -> Result<()> { + debug!("waiting for the MCP OAuth credential transaction lock"); + let _lock = + RefreshCredentialLock::acquire_for_server(&self.inner.server_name, &self.inner.url) + .await?; + debug!("acquired the MCP OAuth credential transaction lock"); + + // Stay on the lifecycle-pinned store. A failure is surfaced rather than falling back and + // possibly replaying an older rotating refresh token from the other store. + debug!("rereading authoritative MCP OAuth credentials"); + let latest = self.inner.credential_store.load( + keyring_store, + &self.inner.server_name, + &self.inner.url, + )?; + + // The pre-lock snapshot is only a hint. This locked reread is authoritative, so adopt a + // winner from another process rather than refreshing its predecessor. + let Some(latest) = latest else { + let manager = self.inner.authorization_manager.clone(); + manager + .lock() + .await + .set_credential_store(InMemoryCredentialStore::new()); + *self.inner.last_credentials.lock().await = None; + return Err(AuthError::AuthorizationRequired).with_context(|| { + format!( + "OAuth tokens for server {} were removed before refresh; authorization required", + self.inner.server_name + ) + }); + }; + + let latest_access_token = latest.token_response.0.access_token().secret(); + // A 401 belongs to the token sent by that request. If another request already refreshed A + // to B, adopt B and retry rather than rotating B because a delayed response rejected A. + let should_adopt = !token_needs_refresh(latest.expires_at) + && match &reason { + RefreshReason::Expiry => true, + RefreshReason::Unauthorized { + rejected_access_token, + } => rejected_access_token.secret() != latest_access_token, + }; + if should_adopt { + debug!("adopting newer MCP OAuth credentials without contacting the provider"); + let manager = self.inner.authorization_manager.clone(); + let mut guard = manager.lock().await; + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await?; + *self.inner.last_credentials.lock().await = Some(latest); + return Ok(()); + } + + // Preserve RMCP's `AuthorizationRequired` marker only for credentials known to be + // unrefreshable. Network and provider failures below remain ordinary errors. + if latest + .token_response + .0 + .refresh_token() + .is_none_or(|refresh_token| refresh_token.secret().trim().is_empty()) + { + return Err(AuthError::AuthorizationRequired).with_context(|| { + format!( + "OAuth tokens for server {} cannot be refreshed; authorization required", + self.inner.server_name + ) + }); + } + + let manager = self.inner.authorization_manager.clone(); + // The provider uses a separate HTTP client and cannot re-enter `AuthClient`. Retain this + // async guard so requests cannot observe credentials while they are staged and committed. + let mut guard = manager.lock().await; + if let Err(error) = + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Refresh).await + { + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await + .context("failed to restore request-only OAuth credentials")?; + return Err(error).context("failed to stage OAuth credentials for refresh"); + } + // The owned task prevents caller deadlines from canceling after possible token rotation; + // this timeout independently bounds the provider request. + debug!( + timeout_ms = refresh_request_timeout.as_millis(), + "requesting refreshed MCP OAuth credentials from the provider" + ); + let refresh_result = match timeout(refresh_request_timeout, guard.refresh_token()).await { + Ok(Ok(token_response)) => { + debug!("received refreshed MCP OAuth credentials from the provider"); + Ok(refreshed_tokens(token_response, &latest, &self.inner)) + } + Ok(Err(error @ AuthError::TokenRefreshFailed(_))) => { + // RMCP 1.8 collapses definitive OAuth rejection (for example, + // `invalid_grant`) and transient token-endpoint failures into this string + // variant. Match RMCP's own request path for now so rejected refresh tokens + // prompt reauthorization instead of surfacing as generic MCP startup failures. + // This can also prompt reauthorization after a transient failure. + // TODO: When RMCP exposes a typed distinction for refresh-token rejection, + // map only that definitive rejection to `AuthorizationRequired` here. + warn!( + error = %error, + "MCP OAuth refresh failed; reauthorization required by RMCP compatibility policy" + ); + return Err(AuthError::AuthorizationRequired).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}: {error}", + self.inner.server_name + ) + }); + } + Ok(Err(error)) => { + warn!( + error = %error, + "MCP OAuth provider refresh failed" + ); + Err(error).with_context(|| { + format!( + "failed to refresh OAuth tokens for server {}", + self.inner.server_name + ) + }) + } + Err(_) => { + warn!( + timeout_ms = refresh_request_timeout.as_millis(), + "MCP OAuth provider refresh timed out; the outcome is unknown and a later serialized retry is permitted" + ); + Err(anyhow::anyhow!( + "timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}", + self.inner.server_name + )) + } + }; + let request_tokens = refresh_result.as_ref().unwrap_or(&latest); + install_tokens_in_manager_guard(&mut guard, request_tokens, CredentialExposure::Request) + .await + .context("failed to restore request-only OAuth credentials")?; + let refreshed = refresh_result?; + + // Persist to the pinned source before exposing the refreshed token. On failure, restore + // the prior in-process credential and return the error; serving an unpersisted token would + // hide the root cause until a later process restart. If the provider already consumed the + // prior token, the next refresh may require reauthorization. That is the deliberate + // fail-closed policy. + // TODO: Add a bounded persistence retry only if telemetry shows this is common; never + // silently switch stores or continue with an unpersisted credential. + debug!("persisting refreshed MCP OAuth credentials to the resolved store"); + if let Err(error) = + self.inner + .credential_store + .save(keyring_store, &self.inner.server_name, &refreshed) + { + warn!( + error = %error, + "failed to persist refreshed MCP OAuth credentials; returning the error and restoring the previous in-process credentials" + ); + install_tokens_in_manager_guard(&mut guard, &latest, CredentialExposure::Request) + .await + .context( + "failed to restore previous request-only OAuth credentials after refresh persistence failed", + )?; + return Err(error); + } + + // Commit the merged response, then expose only its request-safe form so RMCP cannot + // refresh independently outside this transaction. + install_tokens_in_manager_guard(&mut guard, &refreshed, CredentialExposure::Request) + .await + .context( + "refreshed OAuth tokens were persisted but could not be installed in the authorization manager", + )?; + *self.inner.last_credentials.lock().await = Some(refreshed); + drop(guard); + debug!("persisted refreshed MCP OAuth credentials and completed the transaction"); + Ok(()) + } +} + +pub(super) enum RefreshReason { + Expiry, + Unauthorized { rejected_access_token: AccessToken }, +} + +impl RefreshReason { + fn as_str(&self) -> &'static str { + match self { + Self::Expiry => "expiry", + Self::Unauthorized { .. } => "unauthorized", + } + } +} + +/// Ordinary requests receive neither refresh token nor expiry metadata, so RMCP cannot refresh +/// outside Codex's transaction. Full credentials are exposed only while both transaction locks +/// are held. +#[derive(Clone, Copy)] +enum CredentialExposure { + Request, + Refresh, +} + +async fn install_tokens_in_manager_guard( + authorization_manager: &mut AuthorizationManager, + tokens: &StoredOAuthTokens, + exposure: CredentialExposure, +) -> Result<()> { + let store = InMemoryCredentialStore::new(); + let token_response = match exposure { + CredentialExposure::Request => request_oauth_token_response(tokens), + CredentialExposure::Refresh => tokens.token_response.0.clone(), + }; + let granted_scopes = token_response + .scopes() + .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect()) + .unwrap_or_default(); + let token_received_at = match exposure { + CredentialExposure::Request => None, + CredentialExposure::Refresh => SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()), + }; + store + .save(StoredCredentials::new( + tokens.client_id.clone(), + Some(token_response), + granted_scopes, + token_received_at, + )) + .await + .context("failed to stage OAuth tokens for authorization manager")?; + + authorization_manager.set_credential_store(store); + // TODO(stevenlee): Add an RMCP adoption API that atomically updates credentials, client ID, + // and private `current_scopes`; this path cannot synchronize RMCP's scope-upgrade state. + authorization_manager + .initialize_from_store() + .await + .context("failed to adopt refreshed OAuth tokens")?; + Ok(()) +} + +pub(crate) fn request_oauth_token_response(tokens: &StoredOAuthTokens) -> OAuthTokenResponse { + let mut token_response = tokens.token_response.0.clone(); + token_response.set_refresh_token(None); + token_response.set_expires_in(None); + token_response +} + +fn refreshed_tokens( + mut token_response: OAuthTokenResponse, + previous: &StoredOAuthTokens, + inner: &OAuthPersistorInner, +) -> StoredOAuthTokens { + if token_response.refresh_token().is_none() { + token_response.set_refresh_token(previous.token_response.0.refresh_token().cloned()); + } + if token_response.scopes().is_none() { + token_response.set_scopes(previous.token_response.0.scopes().cloned()); + } + StoredOAuthTokens { + server_name: inner.server_name.clone(), + url: inner.url.clone(), + client_id: previous.client_id.clone(), + expires_at: compute_expires_at_millis(&token_response), + token_response: WrappedOAuthTokenResponse(token_response), + } +} diff --git a/codex-rs/rmcp-client/src/oauth/resolved_store.rs b/codex-rs/rmcp-client/src/oauth/resolved_store.rs new file mode 100644 index 000000000000..937fa51e6df2 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/resolved_store.rs @@ -0,0 +1,151 @@ +//! Resolves the configured MCP OAuth store and pins that concrete source for one client lifecycle. + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_keyring_store::KeyringStore; +use tracing::warn; + +use super::OAuthKeyringLoadError; +use super::StoredOAuthTokens; +use super::load_oauth_tokens_from_file; +use super::load_oauth_tokens_from_keyring; +use super::save_oauth_tokens_to_file; +use super::save_oauth_tokens_with_keyring; + +/// Concrete credential store resolved for one MCP OAuth client lifecycle. +/// +/// This is intentionally not durable. `Auto` may resolve differently in a later process, but a +/// client that loaded credentials from one store must reread, refresh, persist, and remove only +/// through that store. A mid-lifecycle backend failure is unexpected and must return an error +/// rather than falling back to another possibly stale refresh token. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResolvedOAuthCredentialStore { + File, + Keyring(AuthKeyringBackendKind), +} + +impl ResolvedOAuthCredentialStore { + /// Loads credentials only from this already-resolved authority. + /// + /// Unlike `resolve_oauth_tokens_from_store_policy`, this never evaluates configured + /// `Auto` fallback policy. + pub(crate) fn load( + self, + keyring_store: &K, + server_name: &str, + url: &str, + ) -> Result> { + match self { + Self::File => load_oauth_tokens_from_file(server_name, url) + .context("failed to reread OAuth tokens from resolved file storage"), + Self::Keyring(keyring_backend_kind) => load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) + .map_err(anyhow::Error::from) + .context( + "failed to reread OAuth tokens from resolved keyring storage; refusing file fallback", + ), + } + } + + /// Saves credentials only to this already-resolved authority. + pub(crate) fn save( + self, + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, + ) -> Result<()> { + match self { + Self::File => save_oauth_tokens_to_file(tokens), + Self::Keyring(keyring_backend_kind) => save_oauth_tokens_with_keyring( + keyring_store, + keyring_backend_kind, + server_name, + tokens, + ), + } + } +} + +#[derive(Debug)] +pub(crate) struct ResolvedOAuthTokens { + pub(crate) tokens: StoredOAuthTokens, + pub(crate) store: ResolvedOAuthCredentialStore, +} + +pub(crate) fn resolve_oauth_tokens_from_store_policy( + keyring_store: &K, + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { + match store_mode { + OAuthCredentialsStoreMode::Auto => { + // Auto remains keyring-first at lifecycle startup. The returned source is then pinned + // by the 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(ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + Ok(None) => Ok( + load_oauth_tokens_from_file(server_name, url)?.map(|tokens| { + ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + } + }), + ), + // Auto may fall back when the keyring backend is unavailable, but a Secrets + // aggregate-lock failure means authority may be changing. Consulting File in + // that state could replay credentials hidden behind a newer Secrets entry. + Err(OAuthKeyringLoadError::StoreLock(error)) => Err(error.into()), + Err(error) => { + warn!("failed to read OAuth tokens from keyring: {error}"); + Ok(load_oauth_tokens_from_file(server_name, url) + .with_context(|| { + format!("failed to read OAuth tokens from keyring: {error}") + })? + .map(|tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + })) + } + } + } + OAuthCredentialsStoreMode::File => Ok(load_oauth_tokens_from_file(server_name, url)?.map( + |tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::File, + }, + )), + OAuthCredentialsStoreMode::Keyring => Ok(load_oauth_tokens_from_keyring( + keyring_store, + keyring_backend_kind, + server_name, + url, + ) + .map_err(anyhow::Error::from) + .context("failed to read OAuth tokens from keyring")? + .map(|tokens| ResolvedOAuthTokens { + tokens, + store: ResolvedOAuthCredentialStore::Keyring(keyring_backend_kind), + })), + } +} diff --git a/codex-rs/rmcp-client/src/oauth/store_lock.rs b/codex-rs/rmcp-client/src/oauth/store_lock.rs index 5f83ebb8a63f..b053f4e5ae25 100644 --- a/codex-rs/rmcp-client/src/oauth/store_lock.rs +++ b/codex-rs/rmcp-client/src/oauth/store_lock.rs @@ -13,7 +13,6 @@ use std::path::PathBuf; use std::time::Duration; use std::time::Instant; -use anyhow::Result; use codex_utils_home_dir::find_codex_home; const OAUTH_LOCK_DIR: &str = "mcp-oauth-locks"; @@ -52,7 +51,7 @@ pub(super) struct OAuthStoreLock { } impl OAuthStoreLock { - pub(super) fn acquire(store: OAuthStore) -> Result { + pub(super) fn acquire(store: OAuthStore) -> Result { // This lock intentionally follows the existing local File/Secrets credential-store // authority. Those stores are CODEX_HOME-backed today: if CODEX_HOME is unset they use // the default home (`~/.codex`), and if an embedder has no local home/filesystem authority @@ -67,7 +66,7 @@ impl OAuthStoreLock { codex_home: &Path, store: OAuthStore, acquire_timeout: Duration, - ) -> Result { + ) -> Result { let path = oauth_store_lock_path(codex_home, store); if let Some(parent) = path.parent() { fs::create_dir_all(parent).map_err(|source| OAuthStoreLockFailure::CreateDir { @@ -99,8 +98,7 @@ impl OAuthStoreLock { store, path, acquire_timeout, - } - .into()); + }); } Err(std::fs::TryLockError::WouldBlock) => { if !reported_contention { @@ -119,16 +117,13 @@ impl OAuthStoreLock { store, path, source: io::Error::from(error), - } - .into()); + }); } } } } } -/// Marks aggregate-store coordination failures in an [`anyhow::Error`] chain. -/// /// Auto may fall back when the configured keyring backend is unavailable, but it must surface a /// lock failure. Falling back while another process owns the aggregate-store lock could leave the /// newer credential in File while a stale Secrets entry remains preferred. diff --git a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs new file mode 100644 index 000000000000..baaa6488a884 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs @@ -0,0 +1,655 @@ +use std::sync::Arc; +use std::sync::mpsc; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use keyring::Error as KeyringError; +use oauth2::AccessToken; +use oauth2::TokenResponse; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::AuthorizationManager; +use rmcp::transport::auth::OAuthState; +use tokio::sync::Mutex as TokioMutex; +use tracing::Event; +use tracing::Id; +use tracing::Metadata; +use tracing::Subscriber; +use tracing::span::Attributes; +use tracing::span::Record; +use tracing::subscriber::Interest; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::MockKeyringStore; +use super::TempCodexHome; +use super::assert_tokens_match_without_expiry; +use super::sample_tokens; +use crate::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::WrappedOAuthTokenResponse; +use crate::oauth::compute_store_key; +use crate::oauth::delete_oauth_tokens_locked; +use crate::oauth::load_oauth_tokens_from_file; +use crate::oauth::refresh_lock::RefreshCredentialLock; +use crate::oauth::refresh_transaction::RefreshReason; +use crate::oauth::save_oauth_tokens_to_file; +use crate::startup_error::is_authentication_required_error; + +const REFRESH_LOCK_CONTENTION_EVENT_TARGET: &str = + "codex_rmcp_client::oauth::refresh_lock::contention"; + +struct LockContentionSubscriber { + contended_tx: mpsc::Sender<()>, +} + +impl Subscriber for LockContentionSubscriber { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + metadata.target() == REFRESH_LOCK_CONTENTION_EVENT_TARGET + } + + fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { + if self.enabled(metadata) { + Interest::always() + } else { + Interest::never() + } + } + + fn max_level_hint(&self) -> Option { + Some(tracing::level_filters::LevelFilter::DEBUG) + } + + fn new_span(&self, _span: &Attributes<'_>) -> Id { + Id::from_u64(/*u*/ 1) + } + + fn record(&self, _span: &Id, _values: &Record<'_>) {} + + fn record_follows_from(&self, _span: &Id, _follows_from: &Id) {} + + fn event(&self, event: &Event<'_>) { + if self.enabled(event.metadata()) { + self.contended_tx + .send(()) + .expect("signal actual OAuth credential-lock contention"); + } + } + + fn enter(&self, _span: &Id) {} + + fn exit(&self, _span: &Id) {} +} + +#[tokio::test(flavor = "current_thread")] +async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> Result<()> { + let (_env, server, initial) = test_context().await?; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-access-token", + "token_type": "Bearer", + "expires_in": 3600, + }))) + .expect(1) + .mount(&server) + .await; + save_oauth_tokens_to_file(&initial)?; + + // Hold the real credential lock until both refresh transactions report WouldBlock. This makes + // the lock assertion independent of task scheduling and ensures removing transaction locking + // makes the test fail before either request can reach the provider. + let held_lock = + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url).await?; + let (contended_tx, contended_rx) = mpsc::channel(); + let _subscriber_guard = + tracing::subscriber::set_default(LockContentionSubscriber { contended_tx }); + + let first = persistor_for(&initial).await?; + let second = persistor_for(&initial).await?; + let first_task = tokio::spawn({ + let first = first.clone(); + async move { first.refresh_if_needed().await } + }); + let second_task = tokio::spawn({ + let second = second.clone(); + async move { second.refresh_if_needed().await } + }); + + wait_for_lock_contention(contended_rx, /*expected_count*/ 2).await?; + drop(held_lock); + first_task.await??; + second_task.await??; + server.verify().await; + + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("refreshed credentials should be stored"); + let mut expected_response = initial.token_response.0.clone(); + expected_response.set_access_token(AccessToken::new("refreshed-access-token".to_string())); + // File loads derive `expires_in` from stable `expires_at`, so it may tick down before this + // assertion. Normalize only that derived field and compare the complete token response so + // omitted refresh-token and scope carry-forward remain covered. + expected_response.set_expires_in(stored.token_response.0.expires_in().as_ref()); + assert_eq!( + stored.token_response, + WrappedOAuthTokenResponse(expected_response) + ); + Ok(()) +} + +#[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its Tokio mutex" +)] +#[tokio::test(flavor = "current_thread")] +async fn delayed_unauthorized_retries_adopt_the_winning_token() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-access-token", + "token_type": "Bearer", + "expires_in": 3600, + }))) + .expect(1) + .mount(&server) + .await; + let mut initial = expired_tokens(&format!("{}/mcp", server.uri())); + initial.expires_at = None; + initial.token_response.0.set_expires_in(None); + save_oauth_tokens_to_file(&initial)?; + + let first = persistor_for(&initial).await?; + let second_manager = authorization_manager_for(&initial).await?; + let second = OAuthPersistor::new( + initial.server_name.clone(), + initial.url.clone(), + Arc::clone(&second_manager), + ResolvedOAuthCredentialStore::File, + Some(initial.clone()), + ); + let rejected_access_token = initial.token_response.0.access_token().clone(); + + first + .refresh_after_unauthorized(rejected_access_token.clone()) + .await?; + // Both calls model requests that left with A. Once the first 401 rotates A to B, later 401s + // must adopt B and retry their requests instead of rotating B again. + first + .refresh_after_unauthorized(rejected_access_token.clone()) + .await?; + second + .refresh_after_unauthorized(rejected_access_token) + .await?; + + server.verify().await; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("the winning refresh should be persisted"); + assert_eq!( + stored.token_response.0.access_token().secret(), + "refreshed-access-token" + ); + let guard = second_manager.lock().await; + let (_client_id, adopted) = guard.get_credentials().await?; + let adopted = adopted.expect("second manager should adopt the winning token"); + assert_eq!(adopted.access_token().secret(), "refreshed-access-token"); + assert!(adopted.refresh_token().is_none()); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn oauth_callback_waits_for_refresh_and_then_becomes_authoritative() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "login-access-token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "login-refresh-token", + "scope": "scope-a scope-b", + }))) + .expect(1) + .mount(&server) + .await; + + let server_name = "callback-lock-test"; + let server_url = format!("{}/mcp", server.uri()); + let held_lock = RefreshCredentialLock::acquire_for_server(server_name, &server_url).await?; + let (contended_tx, contended_rx) = mpsc::channel(); + let _subscriber_guard = + tracing::subscriber::set_default(LockContentionSubscriber { contended_tx }); + let handle = crate::perform_oauth_login_return_url_with_http_client( + server_name, + &server_url, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + /*http_headers*/ None, + /*env_http_headers*/ None, + &["scope-a".to_string(), "scope-b".to_string()], + Some("test-client-id"), + /*oauth_resource*/ None, + Some(/*timeout_secs*/ 5), + /*callback_port*/ None, + /*callback_url*/ None, + Environment::default_for_tests().get_http_client(), + ) + .await?; + let authorization_url = reqwest::Url::parse(handle.authorization_url())?; + let query = authorization_url + .query_pairs() + .collect::>(); + let redirect_uri = query + .get("redirect_uri") + .context("authorization URL omitted redirect_uri")?; + let state = query + .get("state") + .context("authorization URL omitted state")?; + let mut callback_url = reqwest::Url::parse(redirect_uri)?; + callback_url + .query_pairs_mut() + .append_pair("code", "authorization-code") + .append_pair("state", state); + let (_authorization_url, completion) = handle.into_parts(); + reqwest::Client::new() + .get(callback_url) + .send() + .await? + .error_for_status()?; + + // This event is emitted only after the callback's real persistence path observes WouldBlock. + // Writing while the lock is held models a refresh that started first; login must overwrite it + // after the transaction finishes. + wait_for_signal(contended_rx).await?; + let mut refresh_winner = sample_tokens(); + refresh_winner.server_name = server_name.to_string(); + refresh_winner.url.clone_from(&server_url); + refresh_winner + .token_response + .0 + .set_access_token(AccessToken::new("refresh-winner".to_string())); + save_oauth_tokens_to_file(&refresh_winner)?; + drop(held_lock); + completion + .await + .context("OAuth login task was cancelled")??; + + let stored = load_oauth_tokens_from_file(server_name, &server_url)? + .expect("callback credentials should be persisted"); + assert_eq!( + stored.token_response.0.access_token().secret(), + "login-access-token" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn locked_logout_waits_for_refresh_and_removes_its_result() -> Result<()> { + let _env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let (request_received_tx, request_received_rx) = mpsc::channel(); + let (release_response_tx, release_response_rx) = mpsc::channel(); + let release_response_rx = Arc::new(std::sync::Mutex::new(release_response_rx)); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(move |_request: &wiremock::Request| { + request_received_tx + .send(()) + .expect("signal OAuth refresh request"); + release_response_rx + .lock() + .expect("lock OAuth refresh response gate") + .recv() + .expect("wait to release OAuth refresh response"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "refreshed-before-logout", + "token_type": "Bearer", + "expires_in": 3600, + })) + }) + .expect(1) + .mount(&server) + .await; + let initial = expired_tokens(&format!("{}/mcp", server.uri())); + save_oauth_tokens_to_file(&initial)?; + let persistor = persistor_for(&initial).await?; + let refresh_task = tokio::spawn(async move { persistor.refresh_if_needed().await }); + wait_for_signal(request_received_rx).await?; + + let (contended_tx, contended_rx) = mpsc::channel(); + let server_name = initial.server_name.clone(); + let url = initial.url.clone(); + let logout_thread = std::thread::spawn(move || -> Result { + let _subscriber_guard = + tracing::subscriber::set_default(LockContentionSubscriber { contended_tx }); + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()? + .block_on(delete_oauth_tokens_locked( + &server_name, + &url, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )) + }); + + // Do not release the provider response until the real logout path reports WouldBlock on the + // credential lock held by refresh. + wait_for_signal(contended_rx).await?; + release_response_tx + .send(()) + .context("release OAuth refresh response")?; + refresh_task.await??; + let deleted = tokio::task::spawn_blocking(move || logout_thread.join()) + .await? + .map_err(|_| anyhow::anyhow!("logout thread panicked"))??; + assert!(deleted); + assert!(load_oauth_tokens_from_file(&initial.server_name, &initial.url)?.is_none()); + server.verify().await; + Ok(()) +} + +#[expect( + clippy::await_holding_invalid_type, + reason = "AuthorizationManager async access must be serialized through its Tokio mutex" +)] +#[tokio::test(flavor = "current_thread")] +async fn resolved_keyring_read_error_preserves_in_memory_credentials() -> Result<()> { + let (_env, _server, initial) = test_context().await?; + let keyring_store = MockKeyringStore::default(); + let key = compute_store_key(&initial.server_name, &initial.url)?; + keyring_store.set_error(&key, KeyringError::Invalid("error".into(), "load".into())); + let manager = authorization_manager_for(&initial).await?; + let persistor = OAuthPersistor::new( + initial.server_name.clone(), + initial.url.clone(), + Arc::clone(&manager), + ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct), + Some(initial.clone()), + ); + + let error = persistor + .refresh_in( + keyring_store, + RefreshReason::Expiry, + Duration::from_secs(/*secs*/ 45), + ) + .await + .expect_err("the resolved keyring read error should abort refresh"); + assert!( + error + .to_string() + .contains("failed to reread OAuth tokens from resolved keyring storage"), + "unexpected error: {error:#}" + ); + let guard = manager.lock().await; + let (_client_id, token_response) = guard.get_credentials().await?; + assert_eq!( + WrappedOAuthTokenResponse(token_response.expect("manager should retain credentials")), + initial.token_response + ); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn missing_authoritative_credentials_require_reauthorization() -> Result<()> { + let (_env, _server, initial) = test_context().await?; + let persistor = persistor_for(&initial).await?; + + let error = persistor + .refresh_if_needed() + .await + .expect_err("a removed authoritative credential should abort refresh"); + assert!(error.chain().any(|source| matches!( + source.downcast_ref::(), + Some(AuthError::AuthorizationRequired) + ))); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn rejected_refresh_token_requires_reauthorization() -> Result<()> { + let (_env, server, initial) = test_context().await?; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(ResponseTemplate::new(400).set_body_json(serde_json::json!({ + "error": "invalid_grant", + "error_description": "refresh token expired or revoked", + }))) + .expect(1) + .mount(&server) + .await; + save_oauth_tokens_to_file(&initial)?; + let persistor = persistor_for(&initial).await?; + + let error = persistor + .refresh_if_needed() + .await + .expect_err("a provider-rejected refresh token should require reauthorization"); + assert!(is_authentication_required_error(&error)); + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("rejected refresh must preserve the durable credentials"); + assert_tokens_match_without_expiry(&stored, &initial); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn caller_cancellation_does_not_cancel_refresh_persistence() -> Result<()> { + let (_env, server, initial) = test_context().await?; + let (request_received_tx, request_received_rx) = mpsc::channel(); + let (release_response_tx, release_response_rx) = mpsc::channel(); + let release_response_rx = Arc::new(std::sync::Mutex::new(release_response_rx)); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with(move |_request: &wiremock::Request| { + request_received_tx + .send(()) + .expect("signal OAuth refresh request"); + release_response_rx + .lock() + .expect("lock OAuth refresh response gate") + .recv() + .expect("wait to release OAuth refresh response"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": "cancel-safe-access-token", + "token_type": "Bearer", + "expires_in": 3600, + })) + }) + .expect(1) + .mount(&server) + .await; + save_oauth_tokens_to_file(&initial)?; + let persistor = persistor_for(&initial).await?; + let caller = tokio::spawn({ + let persistor = persistor.clone(); + async move { persistor.refresh_if_needed().await } + }); + + tokio::task::spawn_blocking(move || { + request_received_rx + .recv_timeout(Duration::from_secs(/*secs*/ 5)) + .context("timed out waiting for OAuth refresh request") + }) + .await??; + caller.abort(); + assert!( + caller + .await + .expect_err("caller should be cancelled") + .is_cancelled() + ); + + release_response_tx + .send(()) + .context("release OAuth refresh response")?; + + // Reacquiring the same credential lock waits for the detached refresh task to persist and + // release it, avoiding a scheduler-sensitive sleep after cancellation. + let _lock = tokio::time::timeout( + Duration::from_secs(/*secs*/ 2), + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url), + ) + .await + .context("detached refresh did not release its credential lock")??; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("detached refresh should persist credentials"); + assert_eq!( + stored.token_response.0.access_token().secret(), + "cancel-safe-access-token" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn provider_timeout_releases_lock_and_preserves_durable_credentials() -> Result<()> { + let (_env, server, initial) = test_context().await?; + mount_delayed_refresh(&server, "late-access-token").await; + save_oauth_tokens_to_file(&initial)?; + let persistor = persistor_for(&initial).await?; + + let error = persistor + .refresh_in( + MockKeyringStore::default(), + RefreshReason::Expiry, + Duration::from_millis(/*millis*/ 50), + ) + .await + .expect_err("provider request should reach its explicit timeout"); + assert!(error.to_string().contains("timed out after 50ms")); + + let _lock = tokio::time::timeout( + Duration::from_millis(/*millis*/ 100), + RefreshCredentialLock::acquire_for_server(&initial.server_name, &initial.url), + ) + .await + .context("provider timeout did not release the credential lock")??; + let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? + .expect("timed-out refresh must leave durable credentials present"); + assert_tokens_match_without_expiry(&stored, &initial); + server.verify().await; + Ok(()) +} + +async fn persistor_for(tokens: &StoredOAuthTokens) -> Result { + Ok(OAuthPersistor::new( + tokens.server_name.clone(), + tokens.url.clone(), + authorization_manager_for(tokens).await?, + ResolvedOAuthCredentialStore::File, + Some(tokens.clone()), + )) +} + +async fn test_context() -> Result<(TempCodexHome, MockServer, StoredOAuthTokens)> { + let env = TempCodexHome::new(); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + let tokens = expired_tokens(&format!("{}/mcp", server.uri())); + Ok((env, server, tokens)) +} + +async fn authorization_manager_for( + tokens: &StoredOAuthTokens, +) -> Result>> { + let mut state = OAuthState::new(tokens.url.clone(), Some(reqwest::Client::new())).await?; + state + .set_credentials(&tokens.client_id, tokens.token_response.0.clone()) + .await?; + let manager = match state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, + OAuthState::Session(_) | OAuthState::AuthorizedHttpClient(_) => { + anyhow::bail!("unexpected OAuth state") + } + _ => anyhow::bail!("unexpected OAuth state"), + }; + Ok(Arc::new(TokioMutex::new(manager))) +} + +async fn mount_oauth_metadata(server: &MockServer) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a", "scope-b"], + }))) + .mount(server) + .await; +} + +async fn mount_delayed_refresh(server: &MockServer, response_access_token: &str) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains("refresh_token=refresh-token")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(/*millis*/ 200)) + .set_body_json(serde_json::json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + })), + ) + .expect(1) + .mount(server) + .await; +} + +async fn wait_for_lock_contention(rx: mpsc::Receiver<()>, expected_count: usize) -> Result<()> { + tokio::task::spawn_blocking(move || { + for _ in 0..expected_count { + rx.recv_timeout(Duration::from_secs(/*secs*/ 5)) + .context("timed out waiting for lock contention")?; + } + Ok(()) + }) + .await? +} + +async fn wait_for_signal(rx: mpsc::Receiver<()>) -> Result<()> { + tokio::task::spawn_blocking(move || { + rx.recv_timeout(Duration::from_secs(/*secs*/ 5)) + .context("timed out waiting for lock contention") + }) + .await? +} + +fn expired_tokens(url: &str) -> StoredOAuthTokens { + let mut tokens = sample_tokens(); + tokens.url = url.to_string(); + tokens.expires_at = Some(0); + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::ZERO)); + tokens +} diff --git a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs index aafc362c204b..9ecbf2fff6fd 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/store_lock_tests.rs @@ -31,13 +31,14 @@ use crate::oauth::WrappedOAuthTokenResponse; use crate::oauth::fallback_file_path; use crate::oauth::load_oauth_tokens_from_file; use crate::oauth::load_oauth_tokens_from_keyring; -use crate::oauth::load_oauth_tokens_from_keyring_with_fallback_to_file; +use crate::oauth::resolve_oauth_tokens_from_store_policy; use crate::oauth::save_oauth_tokens_to_file; use crate::oauth::save_oauth_tokens_to_file_with_lock_held; use crate::oauth::save_oauth_tokens_to_secrets_keyring_with_lock_held; use crate::oauth::save_oauth_tokens_with_keyring; use crate::oauth::save_oauth_tokens_with_keyring_with_fallback_to_file; use crate::oauth::test_support::TempCodexHome; +use codex_config::types::OAuthCredentialsStoreMode; const STORE_LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::store_lock::contention"; @@ -136,10 +137,7 @@ fn store_lock_is_released_when_holder_process_exits() -> Result<()> { } Err(error) => error, }; - assert!(matches!( - error.downcast_ref::(), - Some(OAuthStoreLockFailure::Timeout { .. }) - )); + assert!(matches!(error, OAuthStoreLockFailure::Timeout { .. })); child .kill() @@ -216,11 +214,12 @@ fn auto_load_secrets_lock_failure_does_not_fall_back_to_file() -> Result<()> { let lock_dir = env.path().join("mcp-oauth-locks"); std::fs::create_dir(lock_dir.join("secrets-store.lock"))?; - let error = load_oauth_tokens_from_keyring_with_fallback_to_file( + let error = resolve_oauth_tokens_from_store_policy( &keyring_store, - AuthKeyringBackendKind::Secrets, &tokens.server_name, &tokens.url, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Secrets, ) .expect_err("aggregate-store lock failure must abort Auto resolution"); @@ -445,12 +444,12 @@ fn secrets_store_load_and_delete_observe_aggregate_lock() -> Result<()> { let url = tokens.url.clone(); let loaded = complete_after_store_lock_contention(env.path(), OAuthStore::Secrets, move || { - load_oauth_tokens_from_keyring( + Ok(load_oauth_tokens_from_keyring( &store_for_load, AuthKeyringBackendKind::Secrets, &server_name, &url, - ) + )?) })? .expect("encrypted credentials should remain readable after contention"); assert_tokens_match_without_expiry(&loaded, &tokens); diff --git a/codex-rs/rmcp-client/src/oauth_transport.rs b/codex-rs/rmcp-client/src/oauth_transport.rs new file mode 100644 index 000000000000..e5ea3fc2f1f5 --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_transport.rs @@ -0,0 +1,254 @@ +//! Codex-owned OAuth policy for RMCP Streamable HTTP traffic. +//! +//! RMCP remains responsible for transport mechanics and bearer-token injection. Codex owns the +//! credential lifecycle: every POST, SSE GET/reconnect, and session DELETE receives proactive +//! refresh from its owning Codex layer, and each path has at most one 401 recovery. The +//! authorization manager only receives request-safe credentials, so it cannot independently +//! refresh outside Codex's serialized transaction. +//! +//! POST recovery is split at an intentional ownership boundary. Client-originated requests and +//! notifications retain their outer `RmcpClient` recovery, which knows the startup/tool deadline +//! and can avoid replaying a request after its caller timed out. RMCP-owned responses to +//! server-initiated requests have no such outer operation, so they recover here. GET/reconnect and +//! DELETE are always RMCP-owned and also recover here. + +use std::collections::HashMap; +use std::sync::Arc; + +use reqwest::header::HeaderName; +use reqwest::header::HeaderValue; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::JsonRpcMessage; +use rmcp::transport::auth::AuthClient; +use rmcp::transport::auth::AuthError; +use rmcp::transport::streamable_http_client::StreamableHttpClient; +use rmcp::transport::streamable_http_client::StreamableHttpError; +use rmcp::transport::streamable_http_client::StreamableHttpPostResponse; +use tracing::debug; + +use crate::http_client_adapter::StreamableHttpClientAdapter; +use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::oauth::OAuthPersistor; + +type TransportResult = + std::result::Result>; + +#[derive(Clone)] +pub(crate) struct OAuthTransportClient { + auth_client: AuthClient, + persistor: OAuthPersistor, +} + +impl OAuthTransportClient { + pub(crate) fn new( + auth_client: AuthClient, + persistor: OAuthPersistor, + ) -> Self { + Self { + auth_client, + persistor, + } + } + + pub(crate) fn persistor(&self) -> OAuthPersistor { + self.persistor.clone() + } + + async fn preflight(&self, operation: &'static str) -> TransportResult<()> { + debug!( + operation, + "checking MCP OAuth credentials before transport request" + ); + self.persistor + .refresh_if_needed() + .await + .map_err(oauth_transport_error) + } + + async fn recover_after_unauthorized( + &self, + operation: &'static str, + rejected_access_token: Option, + ) -> TransportResult { + let Some(rejected_access_token) = rejected_access_token else { + return Ok(false); + }; + + debug!( + operation, + "recovering once after MCP transport rejected an OAuth access token" + ); + self.persistor + .refresh_after_unauthorized(rejected_access_token) + .await + .map_err(oauth_transport_error)?; + Ok(true) + } +} + +impl StreamableHttpClient for OAuthTransportClient { + type Error = StreamableHttpClientAdapterError; + + async fn post_message( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_token: Option, + custom_headers: HashMap, + ) -> TransportResult { + let is_rmcp_owned_response = matches!( + message, + JsonRpcMessage::Response(_) | JsonRpcMessage::Error(_) + ); + if is_rmcp_owned_response { + self.preflight("post_message").await?; + } + let result = self + .auth_client + .post_message( + Arc::clone(&uri), + message.clone(), + session_id.clone(), + auth_token.clone(), + custom_headers.clone(), + ) + .await; + + // RMCP queues client-originated requests independently of the caller waiting on them. If + // recovery happened here, a timed-out public tool call could still be replayed after its + // refresh finished. The outer RmcpClient path owns those deadlines. Responses to + // server-initiated requests have no outer operation and therefore recover here. + if !is_rmcp_owned_response { + return result; + } + let rejected_access_token = result.as_ref().err().and_then(rejected_access_token); + if self + .recover_after_unauthorized("post_message", rejected_access_token) + .await? + { + authorization_required_after_retry( + self.auth_client + .post_message(uri, message, session_id, auth_token, custom_headers) + .await, + ) + } else { + result + } + } + + async fn delete_session( + &self, + uri: Arc, + session_id: Arc, + auth_token: Option, + custom_headers: HashMap, + ) -> TransportResult<()> { + self.preflight("delete_session").await?; + let result = self + .auth_client + .delete_session( + Arc::clone(&uri), + Arc::clone(&session_id), + auth_token.clone(), + custom_headers.clone(), + ) + .await; + let rejected_access_token = result.as_ref().err().and_then(rejected_access_token); + if self + .recover_after_unauthorized("delete_session", rejected_access_token) + .await? + { + authorization_required_after_retry( + self.auth_client + .delete_session(uri, session_id, auth_token, custom_headers) + .await, + ) + } else { + result + } + } + + async fn get_stream( + &self, + uri: Arc, + session_id: Arc, + last_event_id: Option, + auth_token: Option, + custom_headers: HashMap, + ) -> TransportResult< + futures::stream::BoxStream<'static, Result>, + > { + self.preflight("get_stream").await?; + let result = self + .auth_client + .get_stream( + Arc::clone(&uri), + Arc::clone(&session_id), + last_event_id.clone(), + auth_token.clone(), + custom_headers.clone(), + ) + .await; + let rejected_access_token = result.as_ref().err().and_then(rejected_access_token); + if self + .recover_after_unauthorized("get_stream", rejected_access_token) + .await? + { + authorization_required_after_retry( + self.auth_client + .get_stream(uri, session_id, last_event_id, auth_token, custom_headers) + .await, + ) + } else { + result + } + } +} + +fn authorization_required_after_retry(result: TransportResult) -> TransportResult { + match result { + // The first 401 carries the token that was actually rejected so concurrent recovery can + // distinguish A from a newer B. Once the single retry also rejects B, attribution is no + // longer useful: surface the existing reauthentication marker instead of leaking the + // adapter-only error past the Codex-owned recovery boundary. + Err(StreamableHttpError::Client( + StreamableHttpClientAdapterError::AccessTokenRejected { .. }, + )) => Err(StreamableHttpError::Auth(AuthError::AuthorizationRequired)), + result => result, + } +} + +fn rejected_access_token( + error: &StreamableHttpError, +) -> Option { + match error { + StreamableHttpError::Client(StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token, + }) => Some(rejected_access_token.clone()), + _ => None, + } +} + +fn oauth_transport_error( + error: anyhow::Error, +) -> StreamableHttpError { + if let Some(auth_error) = + error + .chain() + .find_map(|source| match source.downcast_ref::() { + Some(AuthError::AuthorizationRequired) => Some(AuthError::AuthorizationRequired), + Some(AuthError::TokenExpired) => Some(AuthError::TokenExpired), + _ => None, + }) + { + // Preserve RMCP's established reauthentication variants across Codex's transport policy + // boundary. Other OAuth failures retain their context-rich adapter error. + return StreamableHttpError::Auth(auth_error); + } + StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuth(error)) +} + +#[cfg(test)] +#[path = "oauth_transport_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/oauth_transport_tests.rs b/codex-rs/rmcp-client/src/oauth_transport_tests.rs new file mode 100644 index 000000000000..acbca36c10bc --- /dev/null +++ b/codex-rs/rmcp-client/src/oauth_transport_tests.rs @@ -0,0 +1,220 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use reqwest::header::HeaderMap; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::transport::auth::AuthClient; +use rmcp::transport::auth::AuthError; +use rmcp::transport::auth::OAuthState; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use rmcp::transport::streamable_http_client::StreamableHttpClient; +use rmcp::transport::streamable_http_client::StreamableHttpError; +use rmcp::transport::streamable_http_client::StreamableHttpPostResponse; +use serde_json::json; +use tempfile::TempDir; +use tokio::process::Command; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::OAuthTransportClient; +use super::authorization_required_after_retry; +use super::oauth_transport_error; +use crate::http_client_adapter::StreamableHttpClientAdapter; +use crate::http_client_adapter::StreamableHttpClientAdapterError; +use crate::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::WrappedOAuthTokenResponse; +use crate::oauth::request_oauth_token_response; +use crate::oauth::save_oauth_tokens; +use crate::oauth_http_client::OAuthHttpClientAdapter; + +const SERVER_NAME: &str = "oauth-transport-response-test"; +const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_RESPONSE_SERVER_URL"; +const ACCESS_TOKEN_A: &str = "response-access-a"; +const REFRESH_TOKEN_A: &str = "response-refresh-a"; +const ACCESS_TOKEN_B: &str = "response-access-b"; +const REFRESH_TOKEN_B: &str = "response-refresh-b"; + +#[test] +fn exhausted_transport_retry_requires_reauthentication() { + let result = authorization_required_after_retry::<()>(Err(StreamableHttpError::Client( + StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token: AccessToken::new(ACCESS_TOKEN_B.to_string()), + }, + ))); + + assert!(matches!( + result, + Err(StreamableHttpError::Auth(AuthError::AuthorizationRequired)) + )); +} + +#[test] +fn oauth_transport_preserves_reauthentication_errors() { + let error = anyhow::Error::new(AuthError::AuthorizationRequired) + .context("refreshing rejected MCP access token"); + + assert!(matches!( + oauth_transport_error(error), + StreamableHttpError::Auth(AuthError::AuthorizationRequired) + )); +} + +#[tokio::test] +async fn server_response_post_receives_one_shot_oauth_recovery() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN_A}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": ACCESS_TOKEN_B, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": REFRESH_TOKEN_B, + "scope": "scope-a", + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + .respond_with(ResponseTemplate::new(202)) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_transport::tests::server_response_post_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + anyhow::ensure!(status.success(), "OAuth response child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test] +#[ignore = "spawned by server_response_post_receives_one_shot_oauth_recovery"] +async fn server_response_post_child() -> anyhow::Result<()> { + let server_url = std::env::var(SERVER_URL_ENV)?; + let initial_tokens = initial_tokens(&server_url); + save_oauth_tokens( + SERVER_NAME, + &initial_tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let http_client = Environment::default_for_tests().get_http_client(); + let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( + Arc::clone(&http_client), + HeaderMap::new(), + )); + let mut oauth_state = + OAuthState::new_with_oauth_http_client(server_url.clone(), oauth_http_client).await?; + oauth_state + .set_credentials( + &initial_tokens.client_id, + request_oauth_token_response(&initial_tokens), + ) + .await?; + let manager = match oauth_state { + OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => manager, + _ => anyhow::bail!("unexpected OAuth state during response test setup"), + }; + let auth_client = AuthClient::new( + StreamableHttpClientAdapter::new( + Arc::clone(&http_client), + HeaderMap::new(), + /*auth_provider*/ None, + ) + .with_rejected_token_attribution(), + manager, + ); + let persistor = OAuthPersistor::new( + SERVER_NAME.to_string(), + server_url.clone(), + Arc::clone(&auth_client.auth_manager), + ResolvedOAuthCredentialStore::File, + Some(initial_tokens), + ); + let client = OAuthTransportClient::new(auth_client, persistor); + let response_message: ClientJsonRpcMessage = serde_json::from_value(json!({ + "jsonrpc": "2.0", + "id": "server-request-1", + "result": { + "action": "accept", + "content": { "confirmed": true } + } + }))?; + + let response = client + .post_message( + Arc::from(server_url), + response_message, + Some(Arc::from("response-session")), + /*auth_token*/ None, + HashMap::new(), + ) + .await?; + + assert!(matches!(response, StreamableHttpPostResponse::Accepted)); + Ok(()) +} + +fn initial_tokens(server_url: &str) -> StoredOAuthTokens { + let mut response = OAuthTokenResponse::new( + AccessToken::new(ACCESS_TOKEN_A.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN_A.to_string()))); + response.set_expires_in(None); + StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + } +} diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 7a9d242bc465..cd2b961b8103 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -28,8 +28,8 @@ use urlencoding::decode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; use crate::oauth::compute_expires_at_millis; +use crate::oauth::save_oauth_tokens_locked; use crate::oauth_http_client::OAuthHttpClientAdapter; -use crate::save_oauth_tokens; use crate::utils::build_default_headers; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; @@ -623,12 +623,13 @@ impl OauthLoginFlow { token_response: WrappedOAuthTokenResponse(credentials), expires_at, }; - save_oauth_tokens( + save_oauth_tokens_locked( &self.server_name, &stored, self.store_mode, self.keyring_backend_kind, - )?; + ) + .await?; Ok(()) } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index c6527990fac0..a15611bd3b1e 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; @@ -14,8 +15,10 @@ use codex_api::SharedAuthProvider; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::McpServerEnvVar; use codex_exec_server::HttpClient; +use codex_keyring_store::DefaultKeyringStore; use futures::FutureExt; use futures::future::BoxFuture; +use oauth2::AccessToken; use oauth2::TokenResponse; use reqwest::header::AUTHORIZATION; use reqwest::header::HeaderMap; @@ -64,10 +67,14 @@ use crate::elicitation_client_service::ElicitationClientService; use crate::http_client_adapter::StreamableHttpClientAdapter; use crate::http_client_adapter::StreamableHttpClientAdapterError; use crate::in_process_transport::InProcessTransportFactory; -use crate::load_oauth_tokens; use crate::oauth::OAuthPersistor; +use crate::oauth::ResolvedOAuthCredentialStore; +use crate::oauth::ResolvedOAuthTokens; use crate::oauth::StoredOAuthTokens; +use crate::oauth::request_oauth_token_response; +use crate::oauth::resolve_oauth_tokens_from_store_policy; use crate::oauth_http_client::OAuthHttpClientAdapter; +use crate::oauth_transport::OAuthTransportClient; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerProcessHandle; @@ -93,7 +100,7 @@ enum PendingTransport { transport: StreamableHttpClientTransport, }, StreamableHttpWithOAuth { - transport: StreamableHttpClientTransport>, + transport: StreamableHttpClientTransport, oauth_persistor: OAuthPersistor, }, } @@ -126,6 +133,8 @@ enum TransportRecipe { env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, + pinned_credential_store: Arc>, + oauth_client: Arc>, http_client: Arc, auth_provider: Option, }, @@ -402,6 +411,8 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + pinned_credential_store: Arc::new(OnceLock::new()), + oauth_client: Arc::new(OnceLock::new()), http_client, auth_provider, }; @@ -445,7 +456,7 @@ impl RmcpClient { }; let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, client_service.clone(), timeout, @@ -477,12 +488,6 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after initialize: {error}"); - } - Ok(initialize_result) } @@ -491,14 +496,13 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); async move { service.list_tools(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -508,7 +512,7 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("tools/list", timeout, move |service| { let params = params.clone(); @@ -533,7 +537,6 @@ impl RmcpClient { }) }) .collect::>>()?; - self.persist_oauth_tokens().await; Ok(ListToolsWithConnectorIdResult { next_cursor: result.next_cursor, tools, @@ -553,14 +556,13 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/list", timeout, move |service| { let params = params.clone(); async move { service.list_resources(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -569,14 +571,13 @@ impl RmcpClient { params: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/templates/list", timeout, move |service| { let params = params.clone(); async move { service.list_resource_templates(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -585,14 +586,13 @@ impl RmcpClient { params: ReadResourceRequestParams, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); async move { service.read_resource(params).await }.boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -603,7 +603,7 @@ impl RmcpClient { meta: Option, timeout: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let arguments = match arguments { Some(Value::Object(map)) => Some(map), Some(other) => { @@ -650,7 +650,6 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(result) } @@ -659,7 +658,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result<()> { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; self.run_service_operation( "notifications/custom", /*timeout*/ None, @@ -680,7 +679,6 @@ impl RmcpClient { }, ) .await?; - self.persist_oauth_tokens().await; Ok(()) } @@ -689,7 +687,7 @@ impl RmcpClient { method: &str, params: Option, ) -> Result { - self.refresh_oauth_if_needed().await; + self.refresh_oauth_if_needed().await?; let response = self .run_service_operation("requests/custom", /*timeout*/ None, move |service| { let params = params.clone(); @@ -703,14 +701,18 @@ impl RmcpClient { .boxed() }) .await?; - self.persist_oauth_tokens().await; Ok(response) } - async fn service(&self) -> Result>> { + async fn service_and_oauth_persistor( + &self, + ) -> Result<( + Arc>, + Option, + )> { let guard = self.state.lock().await; match &*guard { - ClientState::Ready { service, .. } => Ok(Arc::clone(service)), + ClientState::Ready { service, oauth } => Ok((Arc::clone(service), oauth.clone())), ClientState::Connecting { .. } => Err(anyhow!("MCP client not initialized")), ClientState::Closed => Err(anyhow!("MCP client is shut down")), } @@ -743,22 +745,11 @@ impl RmcpClient { drop(previous_state); } - /// This should be called after every tool call so that if a given tool call triggered - /// a refresh of the OAuth tokens, they are persisted. - async fn persist_oauth_tokens(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens: {error}"); - } - } - - async fn refresh_oauth_if_needed(&self) { - if let Some(runtime) = self.oauth_persistor().await - && let Err(error) = runtime.refresh_if_needed().await - { - warn!("failed to refresh OAuth tokens: {error}"); + async fn refresh_oauth_if_needed(&self) -> Result<()> { + if let Some(runtime) = self.oauth_persistor().await { + runtime.refresh_if_needed().await?; } + Ok(()) } async fn create_pending_transport( @@ -781,6 +772,8 @@ impl RmcpClient { env_http_headers, store_mode, keyring_backend_kind, + pinned_credential_store, + oauth_client, http_client, auth_provider, } => { @@ -793,34 +786,91 @@ impl RmcpClient { auth_provider.clone() }; - let initial_oauth_tokens = if bearer_token.is_none() + // Reuse one OAuth manager and persistor across initialize retries and session + // reconstruction. This preserves the lifecycle-pinned store and keeps each failed + // request paired with the manager snapshot that supplied its access token. + if let Some(oauth_client) = oauth_client.get() { + let runtime = oauth_client.persistor(); + let transport = StreamableHttpClientTransport::with_client( + oauth_client.clone(), + StreamableHttpClientTransportConfig::with_uri(url.clone()), + ); + return Ok(PendingTransport::StreamableHttpWithOAuth { + transport, + oauth_persistor: runtime, + }); + } + + let resolved_oauth_tokens = if bearer_token.is_none() && auth_provider.is_none() && !default_headers.contains_key(AUTHORIZATION) { - match load_oauth_tokens(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) = pinned_credential_store.get().copied() { + // Rebuilds reread the source selected during first construction. Only the + // initial construction below evaluates configured store policy. + store + .load(&DefaultKeyringStore, server_name, url)? + .map(|tokens| ResolvedOAuthTokens { tokens, store }) + } else { + match resolve_oauth_tokens_from_store_policy( + &DefaultKeyringStore, + server_name, + url, + *store_mode, + *keyring_backend_kind, + ) { + Ok(tokens) => { + if let Some(resolved) = 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. + pinned_credential_store.set(resolved.store).map_err(|_| { + anyhow!( + "OAuth credential store pinned concurrently for MCP server `{server_name}`" + ) + })?; + } + tokens + } + Err(err) => { + warn!("failed to read tokens for server `{server_name}`: {err}"); + None + } } } } else { None }; - if let Some(initial_tokens) = initial_oauth_tokens.clone() { - match create_oauth_transport_and_runtime( + if let Some(ResolvedOAuthTokens { + tokens: initial_tokens, + store: credential_store, + }) = resolved_oauth_tokens + { + match create_oauth_transport_client( server_name, url, initial_tokens.clone(), - *store_mode, - *keyring_backend_kind, + credential_store, default_headers.clone(), Arc::clone(http_client), ) .await { - Ok((transport, oauth_persistor)) => { + Ok(resolved_oauth_client) => { + oauth_client + .set(resolved_oauth_client.clone()) + .map_err(|_| { + anyhow!( + "OAuth client resolved concurrently for MCP server `{server_name}`" + ) + })?; + let oauth_persistor = resolved_oauth_client.persistor(); + let transport = StreamableHttpClientTransport::with_client( + resolved_oauth_client, + StreamableHttpClientTransportConfig::with_uri(url.clone()), + ); Ok(PendingTransport::StreamableHttpWithOAuth { transport, oauth_persistor, @@ -919,19 +969,7 @@ impl RmcpClient { .await .map_err(|source| anyhow::Error::from(HandshakeError { source })), }; - let service = match service_result { - Ok(service) => service, - Err(error) => { - if let Some(runtime) = oauth_persistor.as_ref() - && let Err(persist_error) = runtime.persist_if_needed().await - { - warn!( - "failed to persist OAuth tokens after failed initialize: {persist_error}" - ); - } - return Err(error); - } - }; + let service = service_result?; Ok((Arc::new(service), oauth_persistor)) } @@ -946,31 +984,83 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { - let service = self.service().await?; - match Self::run_service_operation_with_transient_retries( - Arc::clone(&service), - label, - timeout, - self.elicitation_pause_state.clone(), - &operation, - ) - .await - { - Ok(result) => Ok(result), - Err(error) if Self::is_session_expired_404(&error) => { + let deadline = timeout.map(|duration| Instant::now() + duration); + // Keep the OAuth persistor paired with the service that performs this operation. Session + // recovery can replace both while the request is in flight; rereading only the persistor + // after a 401 could refresh credentials owned by a different transport lifecycle. + let (mut service, mut oauth_persistor) = self.service_and_oauth_persistor().await?; + let mut oauth_recovered = false; + let mut session_recovered = false; + + loop { + let result = Self::run_service_operation_with_transient_retries( + Arc::clone(&service), + label, + timeout, + deadline, + self.elicitation_pause_state.clone(), + &operation, + ) + .await; + + if let Some(rejected_access_token) = result + .as_ref() + .err() + .and_then(Self::rejected_access_token_from_operation_error) + { + if oauth_recovered { + // The rejected token is needed only to attribute the first 401. A second 401 + // after the one allowed refresh means this lifecycle needs reauthentication. + return Err(AuthError::AuthorizationRequired.into()); + } + let Some(oauth_persistor) = oauth_persistor.as_ref() else { + return result.map_err(Into::into); + }; + + // Public request/notification recovery stays here rather than in the transport + // wrapper because this layer owns the caller deadline. RMCP can continue + // processing a queued transport message after the caller times out; retrying it + // inside the wrapper could replay a timed-out tool call. + let remaining = remaining_operation_timeout(label, timeout, deadline)?; + let refresh = oauth_persistor.refresh_after_unauthorized(rejected_access_token); + let refresh_result = match remaining { + Some(remaining) => match time::timeout(remaining, refresh).await { + Ok(result) => result, + Err(_) => { + // The owned transaction keeps running after this caller stops waiting, + // but the rejected operation is not replayed after its deadline. + return Err(ClientOperationError::Timeout { + label: label.to_string(), + duration: timeout.unwrap_or(remaining), + } + .into()); + } + }, + None => refresh.await, + }; + if let Err(error) = refresh_result { + if let Err(timeout_error) = + remaining_operation_timeout(label, timeout, deadline) + { + return Err(timeout_error.into()); + } + return Err(error); + } + oauth_recovered = true; + continue; + } + + if !session_recovered && result.as_ref().is_err_and(Self::is_session_expired_404) { + // OAuth and session recovery are each one-shot, but either error may arrive first. + // Re-entering this loop lets 404 -> 401 compose just like the existing 401 -> 404 + // path without allowing either recovery to repeat indefinitely. self.reinitialize_after_session_expiry(&service).await?; - let recovered_service = self.service().await?; - Self::run_service_operation_with_transient_retries( - recovered_service, - label, - timeout, - self.elicitation_pause_state.clone(), - &operation, - ) - .await - .map_err(Into::into) + (service, oauth_persistor) = self.service_and_oauth_persistor().await?; + session_recovered = true; + continue; } - Err(error) => Err(error.into()), + + return result.map_err(Into::into); } } @@ -978,6 +1068,7 @@ impl RmcpClient { service: Arc>, label: &str, timeout: Option, + retry_deadline: Option, pause_state: ElicitationPauseState, operation: &F, ) -> std::result::Result @@ -985,7 +1076,6 @@ impl RmcpClient { F: Fn(Arc>) -> Fut, Fut: std::future::Future>, { - let retry_deadline = timeout.map(|duration| Instant::now() + duration); for (attempt, retry_delay_ms) in STREAMABLE_HTTP_RETRY_DELAYS_MS .iter() .copied() @@ -1091,6 +1181,34 @@ impl RmcpClient { }) } + fn rejected_access_token_from_operation_error( + error: &ClientOperationError, + ) -> Option { + let ClientOperationError::Service(rmcp::service::ServiceError::TransportSend(error)) = + error + else { + return None; + }; + + error + .error + .downcast_ref::>() + .and_then(Self::rejected_access_token) + } + + pub(super) fn rejected_access_token( + error: &StreamableHttpError, + ) -> Option { + match error { + StreamableHttpError::Client( + StreamableHttpClientAdapterError::AccessTokenRejected { + rejected_access_token, + }, + ) => Some(rejected_access_token.clone()), + _ => None, + } + } + async fn reinitialize_after_session_expiry( &self, failed_service: &Arc>, @@ -1125,7 +1243,7 @@ impl RmcpClient { .ok_or_else(|| anyhow!("MCP client cannot recover before initialize succeeds"))?; let pending_transport = Self::create_pending_transport(&self.transport_recipe).await?; let (service, oauth_persistor) = self - .connect_pending_transport_with_initialize_retries( + .connect_pending_transport_with_oauth_recovery( pending_transport, initialize_context.client_service, initialize_context.timeout, @@ -1143,28 +1261,18 @@ impl RmcpClient { }; } - if let Some(runtime) = oauth_persistor - && let Err(error) = runtime.persist_if_needed().await - { - warn!("failed to persist OAuth tokens after session recovery: {error}"); - } - Ok(()) } } -async fn create_oauth_transport_and_runtime( +async fn create_oauth_transport_client( server_name: &str, url: &str, initial_tokens: StoredOAuthTokens, - credentials_store: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + credential_store: ResolvedOAuthCredentialStore, default_headers: HeaderMap, http_client: Arc, -) -> Result<( - StreamableHttpClientTransport>, - OAuthPersistor, -)> { +) -> Result { let oauth_http_client = Arc::new(OAuthHttpClientAdapter::new( http_client.clone(), default_headers.clone(), @@ -1175,7 +1283,7 @@ async fn create_oauth_transport_and_runtime( oauth_state .set_credentials( &initial_tokens.client_id, - initial_tokens.token_response.0.clone(), + request_oauth_token_response(&initial_tokens), ) .await?; @@ -1188,26 +1296,21 @@ async fn create_oauth_transport_and_runtime( }; let auth_client = AuthClient::new( - StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None), + StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None) + .with_rejected_token_attribution(), manager, ); let auth_manager = auth_client.auth_manager.clone(); - let transport = StreamableHttpClientTransport::with_client( - auth_client, - StreamableHttpClientTransportConfig::with_uri(url.to_string()), - ); - let runtime = OAuthPersistor::new( server_name.to_string(), url.to_string(), auth_manager, - credentials_store, - keyring_backend_kind, + credential_store, Some(initial_tokens), ); - Ok((transport, runtime)) + Ok(OAuthTransportClient::new(auth_client, runtime)) } #[cfg(test)] diff --git a/codex-rs/rmcp-client/src/startup_error.rs b/codex-rs/rmcp-client/src/startup_error.rs index c74f637e3d26..d10c8cd0ed9c 100644 --- a/codex-rs/rmcp-client/src/startup_error.rs +++ b/codex-rs/rmcp-client/src/startup_error.rs @@ -34,7 +34,7 @@ fn client_initialize_error_requires_authentication(error: &ClientInitializeError error, StreamableHttpError::Auth(auth_error) if auth_error_requires_authentication(auth_error) - ) + ) || matches!(error, StreamableHttpError::AuthRequired(_)) }) } diff --git a/codex-rs/rmcp-client/src/streamable_http_retry.rs b/codex-rs/rmcp-client/src/streamable_http_retry.rs index 73da95de58ec..9bff1d548ad4 100644 --- a/codex-rs/rmcp-client/src/streamable_http_retry.rs +++ b/codex-rs/rmcp-client/src/streamable_http_retry.rs @@ -5,9 +5,11 @@ use std::time::Instant; use anyhow::Result; use anyhow::anyhow; use codex_exec_server::ExecServerError; +use oauth2::AccessToken; use reqwest::StatusCode; use rmcp::service::RoleClient; use rmcp::service::RunningService; +use rmcp::transport::auth::AuthError; use rmcp::transport::streamable_http_client::StreamableHttpError; use tokio::time; use tracing::warn; @@ -22,8 +24,13 @@ use super::RmcpClient; const JSON_RPC_INTERNAL_ERROR_CODE: i64 = -32603; pub(super) const STREAMABLE_HTTP_RETRY_DELAYS_MS: [u64; 2] = [250, 1_000]; +#[derive(Default)] +struct InitializeAttemptContext { + oauth_persistor: Option, +} + impl RmcpClient { - pub(super) async fn connect_pending_transport_with_initialize_retries( + pub(super) async fn connect_pending_transport_with_oauth_recovery( &self, initial_transport: PendingTransport, client_service: ElicitationClientService, @@ -31,13 +38,95 @@ impl RmcpClient { ) -> Result<( Arc>, Option, + )> { + let mut initialize_deadline = timeout.map(|duration| Instant::now() + duration); + let mut attempt_context = InitializeAttemptContext::default(); + match self + .connect_pending_transport_with_initialize_retries( + initial_transport, + client_service.clone(), + timeout, + &mut initialize_deadline, + &mut attempt_context, + ) + .await + { + Ok(result) => Ok(result), + Err(error) => { + let Some(rejected_access_token) = + Self::rejected_access_token_from_initialize_error(&error) + else { + return Err(error); + }; + let Some(oauth_persistor) = attempt_context.oauth_persistor else { + return Err(error); + }; + // Initialization gets one OAuth refresh and one reconstructed transport. Reusing + // this wrapper for the retry would turn persistent 401s into a refresh loop. The + // startup deadline gates whether recovery starts and bounds transport setup plus + // the retry handshake, but the refresh transaction has its own bounds and is + // deliberately excluded from the startup budget. + remaining_initialize_timeout(timeout, initialize_deadline)?; + let refresh_started_at = Instant::now(); + let refresh_result = oauth_persistor + .refresh_after_unauthorized(rejected_access_token) + .await; + if let Some(deadline) = initialize_deadline.as_mut() { + *deadline += refresh_started_at.elapsed(); + } + refresh_result?; + let remaining = remaining_initialize_timeout(timeout, initialize_deadline)?; + let transport = match remaining { + Some(remaining) => time::timeout( + remaining, + Self::create_pending_transport(&self.transport_recipe), + ) + .await + .map_err(|_| initialize_timeout_error(timeout, remaining))??, + None => Self::create_pending_transport(&self.transport_recipe).await?, + }; + let mut retry_context = InitializeAttemptContext::default(); + let result = self + .connect_pending_transport_with_initialize_retries( + transport, + client_service, + timeout, + &mut initialize_deadline, + &mut retry_context, + ) + .await; + if result + .as_ref() + .err() + .and_then(Self::rejected_access_token_from_initialize_error) + .is_some() + { + // The first 401 identifies which access token failed. If the reconstructed + // transport still rejects the refreshed token, preserve Codex's established + // signal that the user must authenticate again. + return Err(AuthError::AuthorizationRequired.into()); + } + result + } + } + } + + async fn connect_pending_transport_with_initialize_retries( + &self, + initial_transport: PendingTransport, + client_service: ElicitationClientService, + timeout: Option, + initialize_deadline: &mut Option, + attempt_context: &mut InitializeAttemptContext, + ) -> Result<( + Arc>, + Option, )> { let should_retry = match &initial_transport { PendingTransport::InProcess { .. } | PendingTransport::Stdio { .. } => false, 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 +139,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,8 +151,30 @@ impl RmcpClient { } } }; - let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?; - + if let PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } = &transport + { + // OAuth has independent bounds; pause the MCP handshake budget until refreshed + // credentials are durably committed. + let refresh_started_at = Instant::now(); + oauth_persistor.refresh_if_needed().await?; + if let Some(deadline) = initialize_deadline.as_mut() { + *deadline += refresh_started_at.elapsed(); + } + } + // Keep the persistor paired with the transport attempt that returned 401. Rebuilt + // transports reuse the recipe's lifecycle-pinned credential source, and this pairing + // also keeps the authorization manager and snapshot aligned with the failed attempt. + attempt_context.oauth_persistor = match &transport { + PendingTransport::StreamableHttpWithOAuth { + oauth_persistor, .. + } => Some(oauth_persistor.clone()), + PendingTransport::InProcess { .. } + | PendingTransport::Stdio { .. } + | PendingTransport::StreamableHttp { .. } => None, + }; + let attempt_timeout = remaining_initialize_timeout(timeout, *initialize_deadline)?; match Self::connect_pending_transport( transport, client_service.clone(), @@ -84,7 +195,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:?}" @@ -109,6 +220,33 @@ impl RmcpClient { }) } + fn rejected_access_token_from_initialize_error(error: &anyhow::Error) -> Option { + error.chain().find_map(|source| { + source + .downcast_ref::() + .and_then(|error| { + Self::rejected_access_token_from_client_initialize_error(&error.source) + }) + .or_else(|| { + source + .downcast_ref::() + .and_then(Self::rejected_access_token_from_client_initialize_error) + }) + }) + } + + fn rejected_access_token_from_client_initialize_error( + error: &rmcp::service::ClientInitializeError, + ) -> Option { + match error { + rmcp::service::ClientInitializeError::TransportError { error, .. } => error + .error + .downcast_ref::>() + .and_then(Self::rejected_access_token), + _ => None, + } + } + fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool { match error { rmcp::service::ClientInitializeError::TransportError { error, context } @@ -159,6 +297,10 @@ impl RmcpClient { | StreamableHttpError::ServerDoesNotSupportSse | StreamableHttpError::Deserialize(_) | StreamableHttpError::Client(StreamableHttpClientAdapterError::SessionExpired404) + | StreamableHttpError::Client( + StreamableHttpClientAdapterError::AccessTokenRejected { .. }, + ) + | StreamableHttpError::Client(StreamableHttpClientAdapterError::OAuth(_)) | StreamableHttpError::Client(StreamableHttpClientAdapterError::Header(_)) => false, _ => false, } diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs new file mode 100644 index 000000000000..fc40872a15c7 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_internal.rs @@ -0,0 +1,250 @@ +mod streamable_http_test_support; + +use std::path::PathBuf; +use std::time::Duration; + +use anyhow::Context; +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +use oauth2::AccessToken; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use serde_json::Value; +use serde_json::json; +use tempfile::TempDir; +use tokio::process::Command; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_string_contains; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use streamable_http_test_support::initialize_client; + +const SERVER_NAME: &str = "test-streamable-http-oauth-internal"; +const SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_SERVER_URL"; +const GET_RETRY_MARKER_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_GET_RETRY_MARKER"; +const DELETE_RETRY_MARKER_ENV: &str = "MCP_TEST_OAUTH_INTERNAL_DELETE_RETRY_MARKER"; +const ACCESS_TOKEN_A: &str = "internal-access-a"; +const REFRESH_TOKEN_A: &str = "internal-refresh-a"; +const ACCESS_TOKEN_B: &str = "internal-access-b"; +const REFRESH_TOKEN_B: &str = "internal-refresh-b"; +const ACCESS_TOKEN_C: &str = "internal-access-c"; +const REFRESH_TOKEN_C: &str = "internal-refresh-c"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rmcp_owned_get_and_delete_receive_oauth_recovery() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let get_retry_marker = codex_home.path().join("get-retry-observed"); + let delete_retry_marker = codex_home.path().join("delete-retry-observed"); + let server = MockServer::start().await; + mount_oauth_metadata(&server).await; + mount_refresh(&server, REFRESH_TOKEN_A, ACCESS_TOKEN_B, REFRESH_TOKEN_B).await; + mount_refresh(&server, REFRESH_TOKEN_B, ACCESS_TOKEN_C, REFRESH_TOKEN_C).await; + + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_A}"))) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + // A 405 tells RMCP that the optional common SSE stream is unsupported. Reaching this + // response proves that the wrapper retried the RMCP-owned GET with B after refreshing A. + .respond_with({ + let get_retry_marker = get_retry_marker.clone(); + move |_request: &Request| { + std::fs::write(&get_retry_marker, b"observed") + .expect("record retried RMCP-owned GET"); + ResponseTemplate::new(405) + } + }) + .expect(1) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_B}"))) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/mcp")) + .and(header("authorization", format!("Bearer {ACCESS_TOKEN_C}"))) + .respond_with({ + let delete_retry_marker = delete_retry_marker.clone(); + move |_request: &Request| { + std::fs::write(&delete_retry_marker, b"observed") + .expect("record retried RMCP-owned DELETE"); + ResponseTemplate::new(204) + } + }) + .expect(1) + .mount(&server) + .await; + + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_internal_get_delete_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .env(GET_RETRY_MARKER_ENV, &get_retry_marker) + .env(DELETE_RETRY_MARKER_ENV, &delete_retry_marker) + .status() + .await?; + anyhow::ensure!(status.success(), "OAuth internal child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by rmcp_owned_get_and_delete_receive_oauth_recovery"] +async fn oauth_internal_get_delete_child() -> anyhow::Result<()> { + let client = create_oauth_client().await?; + initialize_client(&client).await?; + wait_for_marker(GET_RETRY_MARKER_ENV).await?; + client.shutdown().await; + wait_for_marker(DELETE_RETRY_MARKER_ENV).await?; + Ok(()) +} + +async fn wait_for_marker(env_name: &str) -> anyhow::Result<()> { + let path = PathBuf::from( + std::env::var_os(env_name).with_context(|| format!("missing {env_name} environment"))?, + ); + tokio::time::timeout(Duration::from_secs(/*secs*/ 5), async { + while !path.exists() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await + .with_context(|| format!("timed out waiting for marker {}", path.display())) +} + +async fn create_oauth_client() -> anyhow::Result { + let server_url = std::env::var(SERVER_URL_ENV)?; + save_initial_tokens(&server_url)?; + RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await +} + +fn save_initial_tokens(server_url: &str) -> anyhow::Result<()> { + let mut response = OAuthTokenResponse::new( + AccessToken::new(ACCESS_TOKEN_A.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN_A.to_string()))); + response.set_expires_in(None); + save_oauth_tokens( + SERVER_NAME, + &StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) +} + +async fn mount_oauth_metadata(server: &MockServer) { + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .mount(server) + .await; +} + +async fn mount_refresh( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, +) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": response_refresh_token, + "scope": "scope-a", + }))) + .expect(1) + .mount(server) + .await; +} + +fn initialize_response(body: &Value) -> ResponseTemplate { + ResponseTemplate::new(200) + .insert_header("mcp-session-id", "oauth-internal-session") + .set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-internal-test", + "version": "0.0.0-test" + } + } + })) +} 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 32f90bc41e8e..e01c7bb36c1b 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,8 @@ mod streamable_http_test_support; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -40,6 +43,11 @@ 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 FINAL_ACCESS_TOKEN: &str = "final-access-token"; +const FINAL_REFRESH_TOKEN: &str = "final-refresh-token"; +const REJECTED_RETRY_ACCESS_TOKEN: &str = "rejected-retry-access-token"; +const REJECTED_RETRY_REFRESH_TOKEN: &str = "rejected-retry-refresh-token"; 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"; @@ -123,6 +131,290 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn recovers_initialization_and_operation_401_once() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .expect(1) + .mount(&server) + .await; + mount_refresh( + &server, + REFRESH_TOKEN, + REFRESHED_ACCESS_TOKEN, + ROTATED_REFRESH_TOKEN, + ) + .await; + mount_refresh( + &server, + ROTATED_REFRESH_TOKEN, + FINAL_ACCESS_TOKEN, + FINAL_REFRESH_TOKEN, + ) + .await; + mount_refresh( + &server, + FINAL_REFRESH_TOKEN, + REJECTED_RETRY_ACCESS_TOKEN, + REJECTED_RETRY_REFRESH_TOKEN, + ) + .await; + + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(401), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {FINAL_ACCESS_TOKEN}"), + )) + .respond_with({ + let resource_attempts = Arc::new(AtomicUsize::new(0)); + move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })), + Some("resources/list") + if resource_attempts.fetch_add(1, Ordering::SeqCst) == 0 => + { + ResponseTemplate::new(404) + } + Some("resources/list") => ResponseTemplate::new(401) + .insert_header("www-authenticate", "Bearer realm=\"mcp\""), + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + } + }) + .expect(5) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REJECTED_RETRY_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=\"mcp\""), + ) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_401_recovery_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + assert!(status.success(), "OAuth recovery child failed: {status}"); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn rejected_initialize_retry_requires_reauthentication() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .expect(1) + .mount(&server) + .await; + mount_refresh( + &server, + REFRESH_TOKEN, + REFRESHED_ACCESS_TOKEN, + ROTATED_REFRESH_TOKEN, + ) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=\"mcp\""), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with( + ResponseTemplate::new(401).insert_header("www-authenticate", "Bearer realm=\"mcp\""), + ) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_rejected_initialize_retry_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + assert!( + status.success(), + "OAuth rejected-retry child failed: {status}" + ); + server.verify().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn operation_timeout_bounds_unauthorized_refresh_wait() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + "scopes_supported": ["scope-a"], + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={REFRESH_TOKEN}" + ))) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(/*millis*/ 500)) + .set_body_json(json!({ + "access_token": REFRESHED_ACCESS_TOKEN, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": ROTATED_REFRESH_TOKEN, + "scope": "scope-a", + })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {EXPIRED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body.get("method").and_then(Value::as_str) { + Some("initialize") => initialize_response(&body), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => ResponseTemplate::new(401), + method => ResponseTemplate::new(400) + .set_body_string(format!("unexpected JSON-RPC method: {method:?}")), + } + }) + // Three requests always use A: initialize, initialized, and the first tools/list. The next + // operation may wait for the in-flight authorization-manager guard and send B directly, + // or it may send A once, receive 401, and join the same refresh transaction. Both + // interleavings are valid. The exact provider and B expectations below prove that neither + // path performs a second refresh or sends more than one successful retry. + .expect(3..=4) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/mcp")) + .and(header( + "authorization", + format!("Bearer {REFRESHED_ACCESS_TOKEN}"), + )) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { "tools": [] }, + })) + }) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "oauth_401_timeout_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + assert!(status.success(), "OAuth timeout child failed: {status}"); + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { let codex_home = TempDir::new()?; @@ -236,7 +528,9 @@ async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { url: UNEXPIRED_SERVER_URL.to_string(), client_id: "test-client-id".to_string(), token_response: WrappedOAuthTokenResponse(response), - expires_at: Some(now.saturating_add(/*rhs*/ 60_000)), + // Keep this outside the 60-second proactive refresh guard band. The test is checking a + // healthy persisted access token, not the boundary where a refresh becomes necessary. + expires_at: Some(now.saturating_add(/*rhs*/ 120_000)), }; save_oauth_tokens( SERVER_NAME, @@ -376,3 +670,142 @@ async fn expired_unrefreshable_startup_child() -> anyhow::Result<()> { assert!(is_authentication_required_error(&error)); Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by recovers_initialization_and_operation_401_once"] +async fn oauth_401_recovery_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let client = refreshable_oauth_client(&server_url).await?; + initialize_client(&client).await?; + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(/*secs*/ 5))) + .await?; + assert!(tools.tools.is_empty()); + + let error = client + .list_resources(/*params*/ None, Some(Duration::from_secs(/*secs*/ 5))) + .await + .expect_err("a rejected one-shot OAuth retry should require reauthentication"); + assert!(is_authentication_required_error(&error)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by rejected_initialize_retry_requires_reauthentication"] +async fn oauth_rejected_initialize_retry_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let client = refreshable_oauth_client(&server_url).await?; + let error = initialize_client(&client) + .await + .expect_err("a rejected initialize retry should require reauthentication"); + assert!(is_authentication_required_error(&error)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by operation_timeout_bounds_unauthorized_refresh_wait"] +async fn oauth_401_timeout_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let client = refreshable_oauth_client(&server_url).await?; + initialize_client(&client).await?; + + let error = client + .list_tools( + /*params*/ None, + Some(Duration::from_millis(/*millis*/ 50)), + ) + .await + .expect_err("operation deadline should expire before the delayed refresh"); + assert!( + error.to_string().contains("timed out awaiting tools/list"), + "unexpected operation error: {error:#}" + ); + // The caller stopped waiting, but the owned refresh transaction still holds the credential + // lock. Starting the next operation immediately makes its preflight wait for that transaction, + // then adopt the persisted token without another provider request. + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(/*secs*/ 5))) + .await?; + assert!(tools.tools.is_empty()); + Ok(()) +} + +async fn refreshable_oauth_client(server_url: &str) -> anyhow::Result { + let mut response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string()))); + response.set_expires_in(None); + save_oauth_tokens( + SERVER_NAME, + &StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: None, + }, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + Ok(client) +} + +async fn mount_refresh( + server: &MockServer, + request_refresh_token: &str, + response_access_token: &str, + response_refresh_token: &str, +) { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(body_string_contains("grant_type=refresh_token")) + .and(body_string_contains(format!( + "refresh_token={request_refresh_token}" + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": response_access_token, + "token_type": "Bearer", + "expires_in": 7200, + "refresh_token": response_refresh_token, + "scope": "scope-a", + }))) + .expect(1) + .mount(server) + .await; +} + +fn initialize_response(body: &Value) -> ResponseTemplate { + ResponseTemplate::new(200) + .insert_header("mcp-session-id", "oauth-recovery-session") + .set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "protocolVersion": body + .pointer("/params/protocolVersion") + .cloned() + .unwrap_or_else(|| json!("2025-06-18")), + "capabilities": {}, + "serverInfo": { + "name": "oauth-401-recovery-test", + "version": "0.0.0-test", + }, + }, + })) +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs new file mode 100644 index 000000000000..d92df03dd253 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_store_pinning.rs @@ -0,0 +1,296 @@ +mod streamable_http_test_support; + +use std::any::Any; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::save_oauth_tokens; +use futures::future::BoxFuture; +use keyring::credential::Credential; +use keyring::credential::CredentialApi; +use keyring::credential::CredentialBuilderApi; +use keyring::credential::CredentialPersistence; +use oauth2::AccessToken; +use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::OAuthTokenResponse; +use rmcp::transport::auth::VendorExtraTokenFields; +use tempfile::TempDir; +use tokio::process::Command; + +use streamable_http_test_support::arm_session_post_failure; +use streamable_http_test_support::call_echo_tool; +use streamable_http_test_support::expected_echo_result; +use streamable_http_test_support::initialize_client; +use streamable_http_test_support::spawn_streamable_http_server; + +const SERVER_NAME: &str = "test-streamable-http-oauth-store-pinning"; +const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_PINNED_STORE_SERVER_URL"; +const KEYRING_ACCESS_TOKEN: &str = "keyring-access-token"; +const FILE_ACCESS_TOKEN: &str = "stale-file-access-token"; + +#[derive(Clone)] +struct RecordingHttpClient { + inner: Arc, + bearer_tokens: Arc>>, +} + +impl RecordingHttpClient { + fn new(inner: Arc) -> Self { + Self { + inner, + bearer_tokens: Arc::new(Mutex::new(Vec::new())), + } + } + + fn record_bearer_token(&self, params: &HttpRequestParams) { + let Some(header) = params + .headers + .iter() + .find(|header| header.name.eq_ignore_ascii_case("authorization")) + else { + return; + }; + self.bearer_tokens + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(header.value.clone()); + } + + fn bearer_tokens(&self) -> Vec { + self.bearer_tokens + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } +} + +impl HttpClient for RecordingHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + self.record_bearer_token(¶ms); + self.inner.http_request(params) + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + self.record_bearer_token(¶ms); + self.inner.http_request_stream(params) + } +} + +#[derive(Debug, Default)] +struct TestKeyringState { + secret: Mutex>>, + fail_reads: AtomicBool, +} + +#[derive(Clone, Debug)] +struct TestCredential { + state: Arc, +} + +impl CredentialApi for TestCredential { + fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> { + *self + .state + .secret + .lock() + .unwrap_or_else(PoisonError::into_inner) = Some(secret.to_vec()); + Ok(()) + } + + fn get_secret(&self) -> keyring::Result> { + if self.state.fail_reads.load(Ordering::SeqCst) { + return Err(keyring::Error::Invalid( + "simulated keyring read failure".to_string(), + "load".to_string(), + )); + } + + self.state + .secret + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + .ok_or(keyring::Error::NoEntry) + } + + fn delete_credential(&self) -> keyring::Result<()> { + self.state + .secret + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .map(|_| ()) + .ok_or(keyring::Error::NoEntry) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +#[derive(Debug)] +struct TestCredentialBuilder { + state: Arc, +} + +impl CredentialBuilderApi for TestCredentialBuilder { + fn build( + &self, + _target: Option<&str>, + _service: &str, + _user: &str, + ) -> keyring::Result> { + Ok(Box::new(TestCredential { + state: Arc::clone(&self.state), + })) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn persistence(&self) -> CredentialPersistence { + CredentialPersistence::ProcessOnly + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn auto_store_remains_pinned_across_session_recovery() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + let codex_home = TempDir::new()?; + + let status = Command::new(std::env::current_exe()?) + .args([ + "auto_store_remains_pinned_across_session_recovery_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, &base_url) + .status() + .await?; + + assert!( + status.success(), + "OAuth store-pinning child failed: {status}" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by auto_store_remains_pinned_across_session_recovery"] +async fn auto_store_remains_pinned_across_session_recovery_child() -> anyhow::Result<()> { + let state = Arc::new(TestKeyringState::default()); + keyring::set_default_credential_builder(Box::new(TestCredentialBuilder { + state: Arc::clone(&state), + })); + + let base_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let server_url = format!("{base_url}/mcp"); + let keyring_tokens = stored_tokens(&server_url, KEYRING_ACCESS_TOKEN); + save_oauth_tokens( + SERVER_NAME, + &keyring_tokens, + OAuthCredentialsStoreMode::Keyring, + AuthKeyringBackendKind::Direct, + )?; + let file_tokens = stored_tokens(&server_url, FILE_ACCESS_TOKEN); + save_oauth_tokens( + SERVER_NAME, + &file_tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )?; + let http_client = RecordingHttpClient::new(Environment::default_for_tests().get_http_client()); + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::Auto, + AuthKeyringBackendKind::Direct, + Arc::new(http_client.clone()), + /*auth_provider*/ None, + ) + .await?; + initialize_client(&client).await?; + assert_eq!( + call_echo_tool(&client, "warmup").await?, + expected_echo_result("warmup") + ); + + arm_session_post_failure( + &base_url, + /*status*/ 404, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; + // The selected keyring becomes unavailable only after initial construction. If recovery + // reevaluates Auto, it adopts the stale File token and this operation incorrectly succeeds. + state.fail_reads.store(true, Ordering::SeqCst); + + match call_echo_tool(&client, "recovery-must-not-fallback").await { + Ok(result) => assert_eq!(result, expected_echo_result("recovery-must-not-fallback")), + Err(error) => { + let error_chain = format!("{error:#}"); + assert!( + error_chain.contains("failed to reread OAuth tokens from resolved keyring storage"), + "unexpected recovery error: {error_chain}" + ); + } + } + + let bearer_tokens = http_client.bearer_tokens(); + assert!( + bearer_tokens + .iter() + .any(|token| token == &format!("Bearer {KEYRING_ACCESS_TOKEN}")), + "expected requests authenticated by the keyring token: {bearer_tokens:?}" + ); + assert!( + bearer_tokens + .iter() + .all(|token| token != &format!("Bearer {FILE_ACCESS_TOKEN}")), + "stale File token must never be sent during recovery: {bearer_tokens:?}" + ); + Ok(()) +} + +fn stored_tokens(server_url: &str, access_token: &str) -> StoredOAuthTokens { + StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(OAuthTokenResponse::new( + AccessToken::new(access_token.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + )), + expires_at: None, + } +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index 29d2404ede9e..a53d0cb08109 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -11,6 +11,8 @@ use codex_exec_server::HttpClient; use codex_exec_server::HttpRequestParams; use codex_exec_server::HttpRequestResponse; use codex_exec_server::HttpResponseBodyStream; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::is_authentication_required_error; use futures::FutureExt as _; use futures::future::BoxFuture; use pretty_assertions::assert_eq; @@ -124,7 +126,13 @@ async fn streamable_http_initialize_retries_remote_no_response_error() -> anyhow async fn streamable_http_initialize_retries_transient_http_status() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; - arm_initialize_post_failure(&base_url, /*status*/ 502, /*remaining*/ 1).await?; + arm_initialize_post_failure( + &base_url, + /*status*/ 502, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[], + ) + .await?; let client = create_client(&base_url).await?; let result = call_echo_tool(&client, "after-status-retry").await?; @@ -301,6 +309,37 @@ async fn streamable_http_401_does_not_trigger_recovery() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn static_bearer_initialize_401_requires_authentication() -> anyhow::Result<()> { + let (_server, base_url) = spawn_streamable_http_server().await?; + arm_initialize_post_failure( + &base_url, + /*status*/ 401, + /*remaining*/ 1, + /*www_authenticate_headers*/ &[r#"Bearer realm="mcp""#], + ) + .await?; + + let client = RmcpClient::new_streamable_http_client( + "test-static-bearer-401", + &format!("{base_url}/mcp"), + Some("test-bearer".to_string()), + /*http_headers*/ None, + /*env_http_headers*/ None, + codex_config::types::OAuthCredentialsStoreMode::File, + codex_config::types::AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + let error = streamable_http_test_support::initialize_client(&client) + .await + .expect_err("a challenged static bearer token should require authentication"); + + assert!(is_authentication_required_error(&error)); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn streamable_http_403_scope_challenge_returns_insufficient_scope() -> anyhow::Result<()> { let (_server, base_url) = spawn_streamable_http_server().await?; 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 87a52f55dff6..30d3039f5bb7 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -258,12 +258,14 @@ pub(crate) async fn arm_initialize_post_failure( base_url: &str, status: u16, remaining: usize, + www_authenticate_headers: &[&str], ) -> anyhow::Result<()> { let response = reqwest::Client::new() .post(format!("{base_url}{INITIALIZE_POST_FAILURE_CONTROL_PATH}")) .json(&json!({ "status": status, "remaining": remaining, + "www_authenticate_headers": www_authenticate_headers, })) .send() .await?;