Reread MCP OAuth credentials before refresh

2026-10-26 18:38:47 +00:00 · 2026-06-19 01:12:52 +00:00
parent 1328e80255
commit 3817102633
3 changed files with 405 additions and 46 deletions
@@ -611,10 +611,6 @@ impl OAuthPersistor {
            return Ok(());
        }

-        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?;
        let latest = load_oauth_tokens_with_keyring_store(
@@ -625,7 +621,7 @@ impl OAuthPersistor {
            self.inner.keyring_backend_kind,
        )?;

-        if latest.is_none() && snapshot.is_some() {
+        let Some(latest) = latest else {
            self.clear_manager_credentials().await;
            let mut last_credentials = self.inner.last_credentials.lock().await;
            *last_credentials = None;
@@ -633,15 +629,15 @@ impl OAuthPersistor {
                "OAuth tokens for server {} were removed before refresh; authorization required",
                self.inner.server_name
            );
-        }
+        };

-        if latest != snapshot {
-            if let Some(latest) = latest {
-                self.adopt_credentials(latest).await?;
-            }
+        if !token_needs_refresh(latest.expires_at) {
+            self.adopt_credentials(latest).await?;
            return Ok(());
        }

+        self.adopt_credentials(latest).await?;
+
        {
            let manager = self.inner.authorization_manager.clone();
            let guard = manager.lock().await;
@@ -1261,29 +1257,14 @@ mod tests {
    -> 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;
-        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(json!({
-                "access_token": "refreshed-access-token",
-                "token_type": "Bearer",
-                "expires_in": 3600,
-                "refresh_token": "rotated-refresh-token",
-                "scope": "scope-a scope-b",
-            })))
-            .expect(1)
-            .mount(&server)
-            .await;
+        mount_oauth_metadata(&server).await;
+        mount_refresh_response(
+            &server,
+            "refresh-token",
+            "refreshed-access-token",
+            "rotated-refresh-token",
+        )
+        .await;

        let store = MockKeyringStore::default();
        let initial_tokens = expired_sample_tokens(&format!("{}/mcp", server.uri()));
@@ -1341,6 +1322,195 @@ 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(),
+            OAuthCredentialsStoreMode::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,
+            OAuthCredentialsStoreMode::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_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), "refreshed-after-expiry-drift");
+        assert_eq!(
+            refresh_token(&stored),
+            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,
+            OAuthCredentialsStoreMode::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_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), "refreshed-from-latest-token");
+        assert_eq!(
+            refresh_token(&stored),
+            Some("rotated-from-latest-token".to_string())
+        );
+        Ok(())
+    }
+
    #[test]
    fn save_oauth_tokens_with_secrets_backend_falls_back_to_file_when_keyring_fails() -> Result<()>
    {
@@ -1644,6 +1814,42 @@ mod tests {
        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(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_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;
+    }
+
    #[expect(
        clippy::await_holding_invalid_type,
        reason = "AuthorizationManager async access must be serialized through its mutex"
@@ -1686,6 +1892,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()),
@@ -491,7 +491,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();
@@ -507,7 +507,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();
@@ -552,7 +552,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();
@@ -568,7 +568,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();
@@ -584,7 +584,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();
@@ -602,7 +602,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) => {
@@ -658,7 +658,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,
@@ -688,7 +688,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();
@@ -752,12 +752,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(
@@ -39,6 +39,7 @@ 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 PREFLIGHT_REFRESH_ERROR: &str = "preflight refresh failed distinctly";
 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";
@@ -233,6 +234,86 @@ 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(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 {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") => 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()?;
@@ -392,6 +473,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_near_expiry_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<()> {
@@ -444,3 +563,31 @@ fn save_expired_file_mode_tokens(server_url: &str) -> anyhow::Result<()> {
    )?;
    Ok(())
 }
+
+fn save_near_expiry_file_mode_tokens(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(Some(&Duration::from_secs(31)));
+    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: server_url.to_string(),
+        client_id: "test-client-id".to_string(),
+        token_response: WrappedOAuthTokenResponse(response),
+        expires_at: Some(now.saturating_add(/*rhs*/ 31_000)),
+    };
+    save_oauth_tokens(
+        SERVER_NAME,
+        &tokens,
+        OAuthCredentialsStoreMode::File,
+        AuthKeyringBackendKind::default(),
+    )?;
+    Ok(())
+}