Reread MCP OAuth credentials before refresh

2026-11-05 20:10:33 +00:00 · 2026-06-25 18:24:35 +00:00
parent 09c6d6f00e
commit 8b83b05278
4 changed files with 374 additions and 41 deletions
@@ -1042,15 +1042,7 @@ mod tests {
    async fn resolved_keyring_write_failure_never_falls_back_to_file() -> Result<()> {
        let _env = TempCodexHome::new();
        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", "scope-b"],
-            })))
-            .mount(&server)
-            .await;
+        mount_oauth_metadata(&server).await;
        let store = MockKeyringStore::default();
        let mut initial_tokens = sample_tokens();
        initial_tokens.url = format!("{}/mcp", server.uri());
@@ -1086,6 +1078,192 @@ mod tests {
        Ok(())
    }

+    #[tokio::test]
+    async fn refresh_transaction_adopts_valid_reread_without_provider_refresh() -> Result<()> {
+        let _env = TempCodexHome::new();
+        let server = MockServer::start().await;
+        mount_oauth_metadata(&server).await;
+        let store = MockKeyringStore::default();
+        let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
+        initial_tokens
+            .token_response
+            .0
+            .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string())));
+
+        let mut latest_tokens = sample_tokens();
+        latest_tokens.url.clone_from(&initial_tokens.url);
+        latest_tokens
+            .token_response
+            .0
+            .set_access_token(AccessToken::new(
+                "already-refreshed-access-token".to_string(),
+            ));
+        latest_tokens
+            .token_response
+            .0
+            .set_refresh_token(Some(RefreshToken::new(
+                "already-rotated-refresh-token".to_string(),
+            )));
+
+        super::save_oauth_tokens_with_keyring_store(
+            &store,
+            &latest_tokens.server_name,
+            &latest_tokens,
+            OAuthCredentialsStoreMode::Keyring,
+            AuthKeyringBackendKind::Direct,
+        )?;
+
+        let manager = authorization_manager_for(&initial_tokens).await?;
+        let persistor = OAuthPersistor::new(
+            initial_tokens.server_name.clone(),
+            initial_tokens.url.clone(),
+            manager.clone(),
+            ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
+            Some(initial_tokens),
+        );
+
+        persistor
+            .refresh_if_needed_with_keyring_store(&store)
+            .await?;
+
+        let manager_tokens = tokens_from_manager(&manager).await?;
+        assert_eq!(
+            access_token(&manager_tokens),
+            "already-refreshed-access-token"
+        );
+        assert_eq!(
+            refresh_token(&manager_tokens),
+            Some("already-rotated-refresh-token".to_string())
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn refresh_transaction_refreshes_when_only_derived_expires_in_drifted() -> Result<()> {
+        let _env = TempCodexHome::new();
+        let server = MockServer::start().await;
+        mount_oauth_metadata(&server).await;
+        mount_refresh_response(
+            &server,
+            "refresh-token",
+            "refreshed-after-expiry-drift",
+            "rotated-after-expiry-drift",
+        )
+        .await;
+
+        let store = MockKeyringStore::default();
+        let mut initial_tokens = sample_tokens();
+        initial_tokens.url = format!("{}/mcp", server.uri());
+        initial_tokens.expires_at = Some(now_millis().saturating_add(5_000));
+        initial_tokens
+            .token_response
+            .0
+            .set_expires_in(Some(&Duration::from_secs(3600)));
+        super::save_oauth_tokens_with_keyring_store(
+            &store,
+            &initial_tokens.server_name,
+            &initial_tokens,
+            OAuthCredentialsStoreMode::Keyring,
+            AuthKeyringBackendKind::Direct,
+        )?;
+
+        let manager = authorization_manager_for(&initial_tokens).await?;
+        let persistor = OAuthPersistor::new(
+            initial_tokens.server_name.clone(),
+            initial_tokens.url.clone(),
+            manager,
+            ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
+            Some(initial_tokens.clone()),
+        );
+
+        persistor
+            .refresh_if_needed_with_keyring_store(&store)
+            .await?;
+
+        server.verify().await;
+        let stored = super::load_oauth_tokens_with_source_and_keyring_store(
+            &store,
+            &initial_tokens.server_name,
+            &initial_tokens.url,
+            OAuthCredentialsStoreMode::Keyring,
+            AuthKeyringBackendKind::Direct,
+        )?
+        .expect("refreshed tokens should be persisted");
+        assert_eq!(access_token(&stored.tokens), "refreshed-after-expiry-drift");
+        assert_eq!(
+            refresh_token(&stored.tokens),
+            Some("rotated-after-expiry-drift".to_string())
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn refresh_transaction_uses_latest_refresh_token_when_reread_is_expired() -> Result<()> {
+        let _env = TempCodexHome::new();
+        let server = MockServer::start().await;
+        mount_oauth_metadata(&server).await;
+        mount_refresh_response(
+            &server,
+            "latest-refresh-token",
+            "refreshed-from-latest-token",
+            "rotated-from-latest-token",
+        )
+        .await;
+
+        let store = MockKeyringStore::default();
+        let mut initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
+        initial_tokens
+            .token_response
+            .0
+            .set_refresh_token(Some(RefreshToken::new("stale-refresh-token".to_string())));
+
+        let mut latest_tokens = initial_tokens.clone();
+        latest_tokens
+            .token_response
+            .0
+            .set_access_token(AccessToken::new("latest-expired-access-token".to_string()));
+        latest_tokens
+            .token_response
+            .0
+            .set_refresh_token(Some(RefreshToken::new("latest-refresh-token".to_string())));
+        super::save_oauth_tokens_with_keyring_store(
+            &store,
+            &latest_tokens.server_name,
+            &latest_tokens,
+            OAuthCredentialsStoreMode::Keyring,
+            AuthKeyringBackendKind::Direct,
+        )?;
+
+        let manager = authorization_manager_for(&initial_tokens).await?;
+        let persistor = OAuthPersistor::new(
+            initial_tokens.server_name.clone(),
+            initial_tokens.url.clone(),
+            manager,
+            ResolvedOAuthCredentialStore::Keyring(AuthKeyringBackendKind::Direct),
+            Some(initial_tokens.clone()),
+        );
+
+        persistor
+            .refresh_if_needed_with_keyring_store(&store)
+            .await?;
+
+        server.verify().await;
+        let stored = super::load_oauth_tokens_with_source_and_keyring_store(
+            &store,
+            &initial_tokens.server_name,
+            &initial_tokens.url,
+            OAuthCredentialsStoreMode::Keyring,
+            AuthKeyringBackendKind::Direct,
+        )?
+        .expect("refreshed tokens should be persisted");
+        assert_eq!(access_token(&stored.tokens), "refreshed-from-latest-token");
+        assert_eq!(
+            refresh_token(&stored.tokens),
+            Some("rotated-from-latest-token".to_string())
+        );
+        Ok(())
+    }
+
    #[tokio::test]
    async fn provider_refresh_timeout_permits_a_later_serialized_retry() -> Result<()> {
        let _env = TempCodexHome::new();
@@ -1547,6 +1725,30 @@ mod tests {
            .await;
    }

+    async fn mount_refresh_response(
+        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 scope-b",
+            })))
+            .expect(1)
+            .mount(server)
+            .await;
+    }
+
    async fn mount_refresh_response_with_signal(
        server: &MockServer,
        request_refresh_token: &str,
@@ -1633,6 +1835,13 @@ mod tests {
        tokens
    }

+    fn now_millis() -> u64 {
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap_or_else(|_| Duration::from_secs(0))
+            .as_millis() as u64
+    }
+
    fn sample_tokens() -> StoredOAuthTokens {
        let mut response = OAuthTokenResponse::new(
            AccessToken::new("access-token".to_string()),
@@ -241,10 +241,6 @@ impl OAuthPersistor {
        let transaction_started_at = Instant::now();
        let lock_started_at = Instant::now();
        debug!("waiting for the MCP OAuth credential transaction lock");
-        let snapshot = {
-            let guard = self.inner.last_credentials.lock().await;
-            guard.clone()
-        };
        let key = compute_store_key(&self.inner.server_name, &self.inner.url)?;
        let _lock = RefreshCredentialLock::acquire(&key).await?;
        debug!(
@@ -273,7 +269,10 @@ impl OAuthPersistor {
            }
        };

-        if latest.is_none() && snapshot.is_some() {
+        // The pre-lock snapshot only decides whether a refresh transaction might be needed. Once
+        // the lock is held, this reread is authoritative: adopt it before deciding whether to
+        // refresh so this process never sends a refresh token superseded by another process.
+        let Some(latest) = latest else {
            self.clear_manager_credentials().await;
            let mut last_credentials = self.inner.last_credentials.lock().await;
            *last_credentials = None;
@@ -281,22 +280,14 @@ impl OAuthPersistor {
                "OAuth tokens for server {} were removed before refresh; authorization required",
                self.inner.server_name
            );
+        };
+
+        if !token_needs_refresh(latest.expires_at) {
+            self.adopt_credentials(latest).await?;
+            return Ok(());
        }

-        if latest != snapshot {
-            if let Some(latest) = latest {
-                let needs_refresh = token_needs_refresh(latest.expires_at);
-                self.adopt_credentials(latest).await?;
-                // `expires_in` is derived from `expires_at` on each load and can drift without a
-                // persisted change. Even for a real concurrent update, keep going when the
-                // authoritative token is still inside the refresh window.
-                if !needs_refresh {
-                    return Ok(());
-                }
-            } else {
-                return Ok(());
-            }
-        }
+        self.adopt_credentials(latest).await?;

        {
            let manager = self.inner.authorization_manager.clone();
@@ -500,7 +500,7 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListToolsResult> {
-        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();
@@ -517,7 +517,7 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListToolsWithConnectorIdResult> {
-        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();
@@ -562,7 +562,7 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListResourcesResult> {
-        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();
@@ -578,7 +578,7 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListResourceTemplatesResult> {
-        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();
@@ -594,7 +594,7 @@ impl RmcpClient {
        params: ReadResourceRequestParams,
        timeout: Option<Duration>,
    ) -> Result<ReadResourceResult> {
-        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();
@@ -612,7 +612,7 @@ impl RmcpClient {
        meta: Option<serde_json::Value>,
        timeout: Option<Duration>,
    ) -> Result<CallToolResult> {
-        self.refresh_oauth_if_needed().await;
+        self.refresh_oauth_if_needed().await?;
        let arguments = match arguments {
            Some(Value::Object(map)) => Some(map),
            Some(other) => {
@@ -668,7 +668,7 @@ impl RmcpClient {
        method: &str,
        params: Option<serde_json::Value>,
    ) -> Result<()> {
-        self.refresh_oauth_if_needed().await;
+        self.refresh_oauth_if_needed().await?;
        self.run_service_operation(
            "notifications/custom",
            /*timeout*/ None,
@@ -698,7 +698,7 @@ impl RmcpClient {
        method: &str,
        params: Option<serde_json::Value>,
    ) -> Result<ServerResult> {
-        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();
@@ -762,12 +762,11 @@ impl RmcpClient {
        }
    }

-    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(
@@ -64,6 +64,7 @@ const ROTATED_REFRESH_TOKEN: &str = "rotated-refresh-token";
 const CHILD_CONTENTION_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_CONTENTION_FILE";
 const CHILD_READY_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_READY_FILE";
 const CHILD_RELEASE_FILE_ENV: &str = "MCP_TEST_OAUTH_STARTUP_RELEASE_FILE";
+const PREFLIGHT_REFRESH_ERROR: &str = "preflight refresh failed distinctly";
 const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL";
 // This mirrors the private event target in oauth::refresh_lock without exposing test-only crate API.
 const LOCK_CONTENTION_EVENT_TARGET: &str = "codex_rmcp_client::oauth::refresh_lock::contention";
@@ -430,6 +431,101 @@ async fn concurrent_file_mode_startup_refreshes_once() -> anyhow::Result<()> {
    Ok(())
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn operation_preflight_refresh_failure_blocks_rmcp_request() -> 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": [""],
+        })))
+        .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_body_json(json!({
+            "access_token": REFRESHED_ACCESS_TOKEN,
+            "token_type": "Bearer",
+            "expires_in": 31,
+            "refresh_token": ROTATED_REFRESH_TOKEN,
+        })))
+        .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={ROTATED_REFRESH_TOKEN}"
+        )))
+        .respond_with(ResponseTemplate::new(400).set_body_json(json!({
+            "error": "invalid_grant",
+            "error_description": PREFLIGHT_REFRESH_ERROR,
+        })))
+        .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") => ResponseTemplate::new(200).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-preflight-test",
+                            "version": "0.0.0-test",
+                        },
+                    },
+                })),
+                Some("notifications/initialized") => ResponseTemplate::new(202),
+                method => ResponseTemplate::new(400)
+                    .set_body_string(format!("unexpected JSON-RPC method: {method:?}")),
+            }
+        })
+        .expect(2)
+        .mount(&server)
+        .await;
+
+    let codex_home = TempDir::new()?;
+    let server_url = format!("{}/mcp", server.uri());
+    let status = Command::new(std::env::current_exe()?)
+        .args([
+            "operation_preflight_refresh_failure_child",
+            "--exact",
+            "--ignored",
+            "--nocapture",
+        ])
+        .env("CODEX_HOME", codex_home.path())
+        .env(CHILD_SERVER_URL_ENV, &server_url)
+        .status()
+        .await?;
+    assert!(
+        status.success(),
+        "OAuth preflight failure 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()?;
@@ -589,6 +685,44 @@ async fn oauth_startup_child() -> anyhow::Result<()> {
    Ok(())
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+#[ignore = "spawned by operation_preflight_refresh_failure_blocks_rmcp_request"]
+async fn operation_preflight_refresh_failure_child() -> anyhow::Result<()> {
+    let server_url = std::env::var(CHILD_SERVER_URL_ENV)?;
+    save_expired_file_mode_tokens(&server_url)?;
+
+    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?;
+
+    initialize_client(&client).await?;
+
+    tokio::time::sleep(Duration::from_millis(1_200)).await;
+    let error = client
+        .list_tools(/*params*/ None, Some(Duration::from_secs(5)))
+        .await
+        .expect_err("preflight refresh failure should abort the operation");
+    let message = format!("{error:#}");
+    assert!(
+        message.contains("failed to refresh OAuth tokens for server"),
+        "unexpected error: {message}"
+    );
+    assert!(
+        message.contains(PREFLIGHT_REFRESH_ERROR),
+        "unexpected error: {message}"
+    );
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 #[ignore = "spawned by concurrent_file_mode_startup_refreshes_once"]
 async fn oauth_concurrency_seed_child() -> anyhow::Result<()> {