diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index 0b3b3bf6a..bef9edc6c 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -12,7 +12,8 @@ use reqwest::header::HeaderMap; use serde::Deserialize; use tracing::debug; -use crate::oauth::has_oauth_tokens; +use crate::oauth::StoredOAuthTokenStatus; +use crate::oauth::oauth_token_status; use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; @@ -44,8 +45,12 @@ pub async fn determine_streamable_http_auth_status( return Ok(McpAuthStatus::BearerToken); } - if has_oauth_tokens(server_name, url, store_mode)? { - return Ok(McpAuthStatus::OAuth); + match oauth_token_status(server_name, url, store_mode)? { + StoredOAuthTokenStatus::Usable => return Ok(McpAuthStatus::OAuth), + StoredOAuthTokenStatus::AuthorizationRequired => { + return Ok(McpAuthStatus::NotLoggedIn); + } + StoredOAuthTokenStatus::Missing => {} } match discover_streamable_http_oauth_with_headers(url, &default_headers).await { diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index c34846079..574266027 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -76,6 +76,13 @@ impl PartialEq for WrappedOAuthTokenResponse { } } +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum StoredOAuthTokenStatus { + Missing, + Usable, + AuthorizationRequired, +} + pub(crate) fn load_oauth_tokens( server_name: &str, url: &str, @@ -94,12 +101,33 @@ pub(crate) fn load_oauth_tokens( } } -pub(crate) fn has_oauth_tokens( +pub(crate) fn oauth_token_status( server_name: &str, url: &str, store_mode: OAuthCredentialsStoreMode, -) -> Result { - Ok(load_oauth_tokens(server_name, url, store_mode)?.is_some()) +) -> Result { + Ok( + match load_oauth_tokens(server_name, url, store_mode)?.as_ref() { + None => StoredOAuthTokenStatus::Missing, + Some(tokens) if oauth_tokens_are_usable(tokens) => StoredOAuthTokenStatus::Usable, + Some(_) => StoredOAuthTokenStatus::AuthorizationRequired, + }, + ) +} + +fn oauth_tokens_are_usable(tokens: &StoredOAuthTokens) -> bool { + if tokens.client_id.trim().is_empty() { + return false; + } + + let token_response = &tokens.token_response.0; + if token_needs_refresh(tokens.expires_at) { + return token_response + .refresh_token() + .is_some_and(|token| !token.secret().trim().is_empty()); + } + + !token_response.access_token().secret().trim().is_empty() } fn refresh_expires_in_from_timestamp(tokens: &mut StoredOAuthTokens) { @@ -852,6 +880,84 @@ mod tests { assert_eq!(tokens.token_response.0.expires_in(), Some(Duration::ZERO)); } + #[test] + fn oauth_tokens_are_usable_when_expiry_is_unknown() { + let mut tokens = sample_tokens(); + tokens.expires_at = None; + tokens.token_response.0.set_refresh_token(None); + + assert!(super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_usable_when_unexpired_without_refresh_token() { + let mut tokens = sample_tokens(); + tokens.token_response.0.set_refresh_token(None); + + assert!(super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_usable_when_expired_but_refreshable() { + let mut tokens = sample_tokens(); + tokens.expires_at = Some(0); + + assert!(super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_expired_and_unrefreshable() { + let mut tokens = sample_tokens(); + tokens.expires_at = Some(0); + tokens.token_response.0.set_refresh_token(None); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_near_expiry_and_unrefreshable() { + let mut tokens = sample_tokens(); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64; + tokens.expires_at = Some(now.saturating_add(REFRESH_SKEW_MILLIS - 1)); + tokens.token_response.0.set_refresh_token(None); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_client_id_is_blank() { + let mut tokens = sample_tokens(); + tokens.client_id = " ".to_string(); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_access_token_is_blank() { + let mut tokens = sample_tokens(); + tokens + .token_response + .0 + .set_access_token(AccessToken::new(" ".to_string())); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + + #[test] + fn oauth_tokens_are_not_usable_when_required_refresh_token_is_blank() { + let mut tokens = sample_tokens(); + tokens.expires_at = Some(0); + tokens + .token_response + .0 + .set_refresh_token(Some(RefreshToken::new(" ".to_string()))); + + assert!(!super::oauth_tokens_are_usable(&tokens)); + } + fn assert_tokens_match_without_expiry( actual: &StoredOAuthTokens, expected: &StoredOAuthTokens, 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 1c18f2c98..d7d8cf423 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -1,16 +1,21 @@ mod streamable_http_test_support; use std::time::Duration; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; +use codex_rmcp_client::McpAuthStatus; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::determine_streamable_http_auth_status; use codex_rmcp_client::save_oauth_tokens; use oauth2::AccessToken; use oauth2::RefreshToken; use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; use rmcp::transport::auth::OAuthTokenResponse; use rmcp::transport::auth::VendorExtraTokenFields; use serde_json::Value; @@ -33,6 +38,9 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token"; const REFRESH_TOKEN: &str = "valid-refresh-token"; const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-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"; +const REFRESHABLE_SERVER_URL: &str = "https://refreshable.example/mcp"; #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> { @@ -112,6 +120,101 @@ async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + + let status = Command::new(std::env::current_exe()?) + .args([ + "persisted_credentials_auth_status_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .status() + .await?; + + assert!( + status.success(), + "persisted credentials auth status child failed: {status}" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by reports_auth_status_for_persisted_credentials"] +async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { + let response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: UNREFRESHABLE_SERVER_URL.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + let status = auth_status(UNREFRESHABLE_SERVER_URL).await?; + assert_eq!(status, McpAuthStatus::NotLoggedIn); + + let response = OAuthTokenResponse::new( + AccessToken::new("unexpired-access-token".to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() as u64; + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + 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)), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + let status = auth_status(UNEXPIRED_SERVER_URL).await?; + assert_eq!(status, McpAuthStatus::OAuth); + + 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()))); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: REFRESHABLE_SERVER_URL.to_string(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?; + + let status = auth_status(REFRESHABLE_SERVER_URL).await?; + assert_eq!(status, McpAuthStatus::OAuth); + Ok(()) +} + +async fn auth_status(server_url: &str) -> anyhow::Result { + determine_streamable_http_auth_status( + SERVER_NAME, + server_url, + /*bearer_token_env_var*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + ) + .await +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"] async fn oauth_startup_child() -> anyhow::Result<()> {