fix(rmcp-client): close OAuth external reload gaps

2026-10-28 18:58:46 +00:00 · 2026-06-05 06:47:55 +00:00
parent 59813d06c9
commit e6319c8c7f
3 changed files with 459 additions and 99 deletions
@@ -48,9 +48,8 @@ use tracing::warn;
 use codex_keyring_store::DefaultKeyringStore;
 use codex_keyring_store::KeyringStore;
 use rmcp::transport::auth::AuthorizationManager;
-use rmcp::transport::auth::CredentialStore;
 use rmcp::transport::auth::InMemoryCredentialStore;
-use rmcp::transport::auth::StoredCredentials;
+use rmcp::transport::auth::OAuthState;
 use tokio::sync::Mutex;
 use tokio::sync::OwnedMutexGuard;

@@ -140,8 +139,14 @@ fn load_oauth_tokens_from_keyring_with_fallback_to_file<K: KeyringStore>(
        Ok(None) => load_oauth_tokens_from_file(server_name, url),
        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}"))
+            match load_oauth_tokens_from_file(server_name, url) {
+                Ok(Some(tokens)) => Ok(Some(tokens)),
+                Ok(None) => Err(error).context(
+                    "failed to read OAuth tokens from keyring and no fallback credentials exist",
+                ),
+                Err(file_error) => Err(file_error)
+                    .with_context(|| format!("failed to read OAuth tokens from keyring: {error}")),
+            }
        }
    }
 }
@@ -297,6 +302,7 @@ struct OAuthPersistorInner {
    server_name: String,
    url: String,
    authorization_manager: Arc<Mutex<AuthorizationManager>>,
+    oauth_metadata_client: reqwest::Client,
    store_mode: OAuthCredentialsStoreMode,
    last_credentials: Mutex<Option<StoredOAuthTokens>>,
 }
@@ -307,11 +313,30 @@ enum CredentialReload {
    Removed,
 }

+pub(crate) async fn authorization_manager_from_tokens(
+    url: &str,
+    oauth_metadata_client: reqwest::Client,
+    tokens: &StoredOAuthTokens,
+) -> Result<AuthorizationManager> {
+    let mut oauth_state = OAuthState::new(url.to_string(), Some(oauth_metadata_client)).await?;
+    oauth_state
+        .set_credentials(&tokens.client_id, tokens.token_response.0.clone())
+        .await?;
+
+    match oauth_state {
+        OAuthState::Authorized(manager) | OAuthState::Unauthorized(manager) => Ok(manager),
+        _ => Err(anyhow::anyhow!(
+            "unexpected OAuth state while replacing credentials"
+        )),
+    }
+}
+
 impl OAuthPersistor {
    pub(crate) fn new(
        server_name: String,
        url: String,
        authorization_manager: Arc<Mutex<AuthorizationManager>>,
+        oauth_metadata_client: reqwest::Client,
        store_mode: OAuthCredentialsStoreMode,
        initial_credentials: Option<StoredOAuthTokens>,
    ) -> Self {
@@ -320,6 +345,7 @@ impl OAuthPersistor {
                server_name,
                url,
                authorization_manager,
+                oauth_metadata_client,
                store_mode,
                last_credentials: Mutex::new(initial_credentials),
            }),
@@ -481,29 +507,15 @@ impl OAuthPersistor {
    }

    async fn replace_manager_credentials(&self, tokens: &StoredOAuthTokens) -> Result<()> {
-        let token_response = tokens.token_response.0.clone();
-        let store = InMemoryCredentialStore::new();
-        store
-            .save(StoredCredentials::new(
-                tokens.client_id.clone(),
-                Some(token_response.clone()),
-                token_response
-                    .scopes()
-                    .map(|scopes| scopes.iter().map(|scope| scope.to_string()).collect())
-                    .unwrap_or_default(),
-                Some(
-                    SystemTime::now()
-                        .duration_since(UNIX_EPOCH)
-                        .unwrap_or_else(|_| Duration::from_secs(0))
-                        .as_secs(),
-                ),
-            ))
-            .await?;
-
+        let replacement = authorization_manager_from_tokens(
+            &self.inner.url,
+            self.inner.oauth_metadata_client.clone(),
+            tokens,
+        )
+        .await?;
        let manager = self.inner.authorization_manager.clone();
        let mut guard = manager.lock().await;
-        guard.configure_client_id(&tokens.client_id)?;
-        guard.set_credential_store(store);
+        *guard = replacement;
        Ok(())
    }

@@ -944,17 +956,28 @@ mod tests {
    use keyring::Error as KeyringError;
    use pretty_assertions::assert_eq;
    use std::ffi::OsString;
-    use std::sync::Barrier;
+    use std::process::Child;
+    use std::process::Command;
    use std::sync::Mutex;
    use std::sync::MutexGuard;
    use std::sync::OnceLock;
    use std::sync::PoisonError;
-    use std::sync::mpsc;
    use std::thread;
+    use std::time::Instant;
    use tempfile::tempdir;
+    use wiremock::Mock;
+    use wiremock::MockServer;
+    use wiremock::ResponseTemplate;
+    use wiremock::matchers::body_string_contains;
+    use wiremock::matchers::method;
+    use wiremock::matchers::path;

    use codex_keyring_store::tests::MockKeyringStore;

+    const FALLBACK_WRITER_VARIANT_ENV: &str = "CODEX_TEST_FALLBACK_WRITER_VARIANT";
+    const FALLBACK_WRITER_ATTEMPT_ENV: &str = "CODEX_TEST_FALLBACK_WRITER_ATTEMPT";
+    const FALLBACK_WRITER_DONE_ENV: &str = "CODEX_TEST_FALLBACK_WRITER_DONE";
+
    struct TempCodexHome {
        _guard: MutexGuard<'static, ()>,
        _dir: tempfile::TempDir,
@@ -1071,6 +1094,99 @@ mod tests {
        Ok(())
    }

+    #[test]
+    fn load_oauth_tokens_preserves_keyring_error_without_fallback() -> Result<()> {
+        let _env = TempCodexHome::new();
+        let store = MockKeyringStore::default();
+        let tokens = sample_tokens();
+        let key = super::compute_store_key(&tokens.server_name, &tokens.url)?;
+        store.set_error(&key, KeyringError::Invalid("error".into(), "load".into()));
+
+        let error = super::load_oauth_tokens_from_keyring_with_fallback_to_file(
+            &store,
+            &tokens.server_name,
+            &tokens.url,
+        )
+        .expect_err("missing fallback must not turn a keyring failure into logout");
+
+        assert!(error.to_string().contains(
+            "failed to read OAuth tokens from keyring and no fallback credentials exist"
+        ));
+        Ok(())
+    }
+
+    #[tokio::test]
+    #[expect(
+        clippy::await_holding_invalid_type,
+        reason = "the test must inspect and refresh the same authorization manager instance"
+    )]
+    async fn replacement_updates_current_scopes_used_by_scope_less_refresh() -> Result<()> {
+        let server = MockServer::start().await;
+        let server_url = format!("{}/mcp", server.uri());
+        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": ["external-scope"],
+            })))
+            .expect(2)
+            .mount(&server)
+            .await;
+        Mock::given(method("POST"))
+            .and(path("/oauth/token"))
+            .and(body_string_contains("refresh_token=refresh-token"))
+            .and(body_string_contains("scope=external-scope"))
+            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
+                "access_token": "refreshed-access-token",
+                "token_type": "Bearer",
+                "expires_in": 7200,
+            })))
+            .expect(1)
+            .mount(&server)
+            .await;
+
+        let mut initial = sample_tokens();
+        initial.url = server_url.clone();
+        let oauth_metadata_client = reqwest::Client::new();
+        let manager = Arc::new(tokio::sync::Mutex::new(
+            super::authorization_manager_from_tokens(
+                &server_url,
+                oauth_metadata_client.clone(),
+                &initial,
+            )
+            .await?,
+        ));
+        let runtime = OAuthPersistor::new(
+            initial.server_name.clone(),
+            server_url,
+            Arc::clone(&manager),
+            oauth_metadata_client,
+            OAuthCredentialsStoreMode::File,
+            Some(initial.clone()),
+        );
+        let mut replacement = initial;
+        replacement
+            .token_response
+            .0
+            .set_scopes(Some(vec![Scope::new("external-scope".to_string())]));
+
+        runtime.replace_manager_credentials(&replacement).await?;
+        let guard = manager.lock().await;
+        assert_eq!(
+            guard.get_current_scopes().await,
+            vec!["external-scope".to_string()]
+        );
+        guard.refresh_token().await?;
+        assert_eq!(
+            guard.get_current_scopes().await,
+            vec!["external-scope".to_string()]
+        );
+        drop(guard);
+        server.verify().await;
+        Ok(())
+    }
+
    #[test]
    fn save_oauth_tokens_prefers_keyring_when_available() -> Result<()> {
        let _env = TempCodexHome::new();
@@ -1319,7 +1435,7 @@ mod tests {
    }

    #[test]
-    fn fallback_store_serializes_writes_for_different_servers() -> Result<()> {
+    fn fallback_store_serializes_cross_process_writes_for_different_servers() -> Result<()> {
        let _env = TempCodexHome::new();
        let first = sample_tokens();
        let mut second = sample_tokens();
@@ -1332,32 +1448,45 @@ mod tests {
        );

        let fallback_lock = super::acquire_fallback_store_lock()?;
-        let barrier = Arc::new(Barrier::new(3));
-        let (tx, rx) = mpsc::channel();
-        let handles = [first.clone(), second.clone()]
+        let codex_home = find_codex_home()?;
+        let markers = tempdir()?;
+        let mut writers = ["first", "second"]
            .into_iter()
-            .map(|tokens| {
-                let barrier = Arc::clone(&barrier);
-                let tx = tx.clone();
-                thread::spawn(move || {
-                    barrier.wait();
-                    tx.send(super::save_oauth_tokens_to_file(&tokens))
-                        .expect("send save result");
-                })
+            .map(|variant| {
+                let attempt = markers.path().join(format!("{variant}-attempt"));
+                let done = markers.path().join(format!("{variant}-done"));
+                let child = Command::new(std::env::current_exe()?)
+                    .args(["fallback_store_writer_child", "--ignored", "--nocapture"])
+                    .env("CODEX_HOME", codex_home.as_path())
+                    .env(FALLBACK_WRITER_VARIANT_ENV, variant)
+                    .env(FALLBACK_WRITER_ATTEMPT_ENV, &attempt)
+                    .env(FALLBACK_WRITER_DONE_ENV, &done)
+                    .spawn()?;
+                Ok((child, attempt, done))
            })
-            .collect::<Vec<_>>();
-        drop(tx);
-
-        barrier.wait();
-        thread::sleep(Duration::from_millis(50));
-        assert!(matches!(rx.try_recv(), Err(mpsc::TryRecvError::Empty)));
+            .collect::<Result<Vec<_>>>()?;
+        wait_for_paths(
+            &writers
+                .iter()
+                .map(|(_, attempt, _)| attempt.clone())
+                .collect::<Vec<_>>(),
+            Duration::from_secs(5),
+        )?;
+        thread::sleep(Duration::from_millis(100));
+        for (child, _, done) in &mut writers {
+            assert!(
+                child.try_wait()?.is_none(),
+                "fallback writer should block on the cross-process file lock"
+            );
+            assert!(
+                !done.exists(),
+                "fallback writer completed while another process held the file lock"
+            );
+        }

        drop(fallback_lock);
-        for _ in 0..2 {
-            rx.recv_timeout(Duration::from_secs(5))??;
-        }
-        for handle in handles {
-            handle.join().expect("fallback writer should not panic");
+        for (child, _, _) in &mut writers {
+            wait_for_child_success(child, Duration::from_secs(5))?;
        }

        let store = super::read_fallback_file()?.expect("fallback file should load");
@@ -1373,6 +1502,38 @@ mod tests {
        Ok(())
    }

+    #[test]
+    #[ignore = "spawned by fallback_store_serializes_cross_process_writes_for_different_servers"]
+    fn fallback_store_writer_child() -> Result<()> {
+        let mut tokens = sample_tokens();
+        match std::env::var(FALLBACK_WRITER_VARIANT_ENV)?.as_str() {
+            "first" => {}
+            "second" => {
+                tokens.server_name = "other-server".to_string();
+                tokens.url = "https://other.example.test".to_string();
+                tokens.token_response.0 = OAuthTokenResponse::new(
+                    AccessToken::new("other-access-token".to_string()),
+                    BasicTokenType::Bearer,
+                    VendorExtraTokenFields::default(),
+                );
+            }
+            variant => anyhow::bail!("unexpected fallback writer variant: {variant}"),
+        }
+
+        let attempt = PathBuf::from(
+            std::env::var_os(FALLBACK_WRITER_ATTEMPT_ENV)
+                .context("fallback writer attempt path should be set")?,
+        );
+        let done = PathBuf::from(
+            std::env::var_os(FALLBACK_WRITER_DONE_ENV)
+                .context("fallback writer done path should be set")?,
+        );
+        fs::write(attempt, b"attempting")?;
+        super::save_oauth_tokens_to_file(&tokens)?;
+        fs::write(done, b"done")?;
+        Ok(())
+    }
+
    #[cfg(any(unix, windows))]
    #[test]
    fn server_generation_lock_is_stable_across_read_only_homes_and_tmpdirs() -> Result<()> {
@@ -1460,6 +1621,31 @@ mod tests {
        );
    }

+    fn wait_for_paths(paths: &[PathBuf], timeout: Duration) -> Result<()> {
+        let deadline = Instant::now() + timeout;
+        while !paths.iter().all(|path| path.exists()) {
+            if Instant::now() >= deadline {
+                anyhow::bail!("timed out waiting for fallback writer readiness");
+            }
+            thread::sleep(Duration::from_millis(10));
+        }
+        Ok(())
+    }
+
+    fn wait_for_child_success(child: &mut Child, timeout: Duration) -> Result<()> {
+        let deadline = Instant::now() + timeout;
+        loop {
+            if let Some(status) = child.try_wait()? {
+                anyhow::ensure!(status.success(), "fallback writer failed: {status}");
+                return Ok(());
+            }
+            if Instant::now() >= deadline {
+                anyhow::bail!("timed out waiting for fallback writer");
+            }
+            thread::sleep(Duration::from_millis(10));
+        }
+    }
+
    fn sample_tokens() -> StoredOAuthTokens {
        let mut response = OAuthTokenResponse::new(
            AccessToken::new("access-token".to_string()),
@@ -47,7 +47,6 @@ use rmcp::service::{self};
 use rmcp::transport::StreamableHttpClientTransport;
 use rmcp::transport::auth::AuthClient;
 use rmcp::transport::auth::AuthError;
-use rmcp::transport::auth::OAuthState;
 use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
 use rmcp::transport::streamable_http_client::StreamableHttpError;
 use serde::Deserialize;
@@ -66,6 +65,7 @@ use crate::in_process_transport::InProcessTransportFactory;
 use crate::load_oauth_tokens;
 use crate::oauth::OAuthPersistor;
 use crate::oauth::StoredOAuthTokens;
+use crate::oauth::authorization_manager_from_tokens;
 use crate::stdio_server_launcher::StdioServerCommand;
 use crate::stdio_server_launcher::StdioServerLauncher;
 use crate::stdio_server_launcher::StdioServerProcessHandle;
@@ -439,7 +439,6 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListToolsResult> {
-        self.refresh_oauth_if_needed().await;
        let result = self
            .run_service_operation("tools/list", timeout, move |service| {
                let params = params.clone();
@@ -455,7 +454,6 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListToolsWithConnectorIdResult> {
-        self.refresh_oauth_if_needed().await;
        let result = self
            .run_service_operation("tools/list", timeout, move |service| {
                let params = params.clone();
@@ -500,7 +498,6 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListResourcesResult> {
-        self.refresh_oauth_if_needed().await;
        let result = self
            .run_service_operation("resources/list", timeout, move |service| {
                let params = params.clone();
@@ -516,7 +513,6 @@ impl RmcpClient {
        params: Option<PaginatedRequestParams>,
        timeout: Option<Duration>,
    ) -> Result<ListResourceTemplatesResult> {
-        self.refresh_oauth_if_needed().await;
        let result = self
            .run_service_operation("resources/templates/list", timeout, move |service| {
                let params = params.clone();
@@ -532,7 +528,6 @@ impl RmcpClient {
        params: ReadResourceRequestParams,
        timeout: Option<Duration>,
    ) -> Result<ReadResourceResult> {
-        self.refresh_oauth_if_needed().await;
        let result = self
            .run_service_operation("resources/read", timeout, move |service| {
                let params = params.clone();
@@ -550,7 +545,6 @@ impl RmcpClient {
        meta: Option<serde_json::Value>,
        timeout: Option<Duration>,
    ) -> Result<CallToolResult> {
-        self.refresh_oauth_if_needed().await;
        let arguments = match arguments {
            Some(Value::Object(map)) => Some(map),
            Some(other) => {
@@ -606,7 +600,6 @@ impl RmcpClient {
        method: &str,
        params: Option<serde_json::Value>,
    ) -> Result<()> {
-        self.refresh_oauth_if_needed().await;
        self.run_service_operation(
            "notifications/custom",
            /*timeout*/ None,
@@ -636,7 +629,6 @@ impl RmcpClient {
        method: &str,
        params: Option<serde_json::Value>,
    ) -> Result<ServerResult> {
-        self.refresh_oauth_if_needed().await;
        let response = self
            .run_service_operation("requests/custom", /*timeout*/ None, move |service| {
                let params = params.clone();
@@ -869,20 +861,21 @@ impl RmcpClient {
        Fut: std::future::Future<Output = std::result::Result<T, rmcp::service::ServiceError>>,
    {
        let service = self.service().await?;
-        match Self::run_service_operation_once(
-            Arc::clone(&service),
-            label,
-            timeout,
-            self.elicitation_pause_state.clone(),
-            &operation,
-        )
-        .await
+        match self
+            .run_service_operation_once(
+                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) => {
                self.reinitialize_after_session_expiry(&service).await?;
                let recovered_service = self.service().await?;
-                Self::run_service_operation_once(
+                self.run_service_operation_once(
                    recovered_service,
                    label,
                    timeout,
@@ -897,6 +890,7 @@ impl RmcpClient {
    }

    async fn run_service_operation_once<T, F, Fut>(
+        &self,
        service: Arc<RunningService<RoleClient, ElicitationClientService>>,
        label: &str,
        timeout: Option<Duration>,
@@ -907,17 +901,19 @@ impl RmcpClient {
        F: Fn(Arc<RunningService<RoleClient, ElicitationClientService>>) -> Fut,
        Fut: std::future::Future<Output = std::result::Result<T, rmcp::service::ServiceError>>,
    {
+        let operation = async {
+            self.refresh_oauth_if_needed().await;
+            operation(service).await
+        };
        match timeout {
-            Some(duration) => {
-                active_time_timeout(duration, pause_state.subscribe(), operation(service))
-                    .await
-                    .map_err(|_| ClientOperationError::Timeout {
-                        label: label.to_string(),
-                        duration,
-                    })?
-                    .map_err(ClientOperationError::from)
-            }
-            None => operation(service).await.map_err(ClientOperationError::from),
+            Some(duration) => active_time_timeout(duration, pause_state.subscribe(), operation)
+                .await
+                .map_err(|_| ClientOperationError::Timeout {
+                    label: label.to_string(),
+                    duration,
+                })?
+                .map_err(ClientOperationError::from),
+            None => operation.await.map_err(ClientOperationError::from),
        }
    }

@@ -1021,23 +1017,9 @@ async fn create_oauth_transport_and_runtime(
    // TODO(aibrahim): teach OAuth bootstrap and refresh to use the same
    // shared HTTP client abstraction instead of always creating the local
    // reqwest metadata client here.
-    let mut oauth_state =
-        OAuthState::new(url.to_string(), Some(oauth_metadata_client.clone())).await?;
-
-    oauth_state
-        .set_credentials(
-            &initial_tokens.client_id,
-            initial_tokens.token_response.0.clone(),
-        )
-        .await?;
-
-    let manager = match oauth_state {
-        OAuthState::Authorized(manager) => manager,
-        OAuthState::Unauthorized(manager) => manager,
-        _ => {
-            return Err(anyhow!("unexpected OAuth state during client setup"));
-        }
-    };
+    let manager =
+        authorization_manager_from_tokens(url, oauth_metadata_client.clone(), &initial_tokens)
+            .await?;

    let auth_client = AuthClient::new(
        StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None),
@@ -1054,6 +1036,7 @@ async fn create_oauth_transport_and_runtime(
        server_name.to_string(),
        url.to_string(),
        auth_manager,
+        oauth_metadata_client,
        credentials_store,
        Some(initial_tokens),
    );
@@ -1,6 +1,7 @@
 mod streamable_http_test_support;

 use std::collections::BTreeMap;
+use std::sync::Arc;
 use std::time::Duration;
 use std::time::SystemTime;
 use std::time::UNIX_EPOCH;
@@ -28,28 +29,218 @@ 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-startup";
+const STARTUP_REFRESH_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_REFRESH_SERVER_URL";
 const EXTERNAL_UPDATE_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_EXTERNAL_UPDATE_SERVER_URL";
+const OPERATION_TIMEOUT_CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_TIMEOUT_SERVER_URL";
+const OPERATION_TIMEOUT_REFRESH_STARTED_ENV: &str = "MCP_TEST_OAUTH_TIMEOUT_REFRESH_STARTED";
 const REFRESH_CONTENTION_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_SERVER_URL";
 const REFRESH_CONTENTION_ATTEMPT_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_ATTEMPT";
 const REFRESH_CONTENTION_DONE_ENV: &str = "MCP_TEST_OAUTH_REFRESH_CONTENTION_DONE";
+const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token";
+const VALID_REFRESH_TOKEN: &str = "valid-refresh-token";
+const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token";
 const OLD_ACCESS_TOKEN: &str = "old-access-token";
 const OLD_REFRESH_TOKEN: &str = "old-refresh-token";
 const EXTERNAL_ACCESS_TOKEN: &str = "external-access-token";
 const EXTERNAL_REFRESH_TOKEN: &str = "external-refresh-token";
 const EXTERNAL_EXPIRES_AT: u64 = u64::MAX;
+const TIMEOUT_REFRESHED_ACCESS_TOKEN: &str = "timeout-refreshed-access-token";
 const STALE_REFRESHED_ACCESS_TOKEN: &str = "stale-refreshed-access-token";
 const STALE_ROTATED_REFRESH_TOKEN: &str = "stale-rotated-refresh-token";

+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> {
+    let server = MockServer::start().await;
+    mount_oauth_metadata(&server, /*expected_calls*/ 1).await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .and(body_string_contains("grant_type=refresh_token"))
+        .and(body_string_contains(format!(
+            "refresh_token={VALID_REFRESH_TOKEN}"
+        )))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "access_token": REFRESHED_ACCESS_TOKEN,
+            "token_type": "Bearer",
+            "expires_in": 7200,
+            "refresh_token": VALID_REFRESH_TOKEN,
+        })))
+        .expect(1)
+        .mount(&server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path("/mcp"))
+        .and(header(
+            "authorization",
+            format!("Bearer {REFRESHED_ACCESS_TOKEN}"),
+        ))
+        .respond_with(|request: &Request| oauth_mcp_response(request, REFRESHED_ACCESS_TOKEN))
+        .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([
+            "oauth_startup_refresh_child",
+            "--exact",
+            "--ignored",
+            "--nocapture",
+        ])
+        .env("CODEX_HOME", codex_home.path())
+        .env(STARTUP_REFRESH_CHILD_SERVER_URL_ENV, server_url)
+        .status()
+        .await?;
+    assert!(status.success(), "OAuth startup child failed: {status}");
+    server.verify().await;
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+#[ignore = "spawned by refreshes_expired_persisted_token_before_initialize"]
+async fn oauth_startup_refresh_child() -> anyhow::Result<()> {
+    let server_url = std::env::var(STARTUP_REFRESH_CHILD_SERVER_URL_ENV)?;
+    save_test_tokens(
+        &server_url,
+        EXPIRED_ACCESS_TOKEN,
+        VALID_REFRESH_TOKEN,
+        /*expires_at*/ 0,
+    )?;
+    initialized_oauth_client(&server_url).await?;
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn operation_timeout_includes_oauth_generation_lock_wait() -> anyhow::Result<()> {
+    let server = MockServer::start().await;
+    mount_oauth_metadata(&server, /*expected_calls*/ 3).await;
+    let codex_home = TempDir::new()?;
+    let refresh_started = codex_home.path().join("refresh-started");
+    let refresh_started_for_mock = refresh_started.clone();
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .and(body_string_contains("grant_type=refresh_token"))
+        .and(body_string_contains(format!(
+            "refresh_token={EXTERNAL_REFRESH_TOKEN}"
+        )))
+        .respond_with(move |_request: &Request| {
+            std::fs::write(&refresh_started_for_mock, b"started").expect("write refresh marker");
+            ResponseTemplate::new(200)
+                .set_delay(Duration::from_secs(1))
+                .set_body_json(json!({
+                    "access_token": TIMEOUT_REFRESHED_ACCESS_TOKEN,
+                    "token_type": "Bearer",
+                    "expires_in": 7200,
+                    "refresh_token": EXTERNAL_REFRESH_TOKEN,
+                }))
+        })
+        .expect(1)
+        .mount(&server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path("/mcp"))
+        .respond_with(|request: &Request| {
+            let Ok(body) = request.body_json::<Value>() else {
+                return ResponseTemplate::new(400);
+            };
+            let expected_access_token = match (
+                body.get("method").and_then(Value::as_str),
+                body.pointer("/params/name").and_then(Value::as_str),
+            ) {
+                (Some("tools/call"), Some("hold-generation-lock")) => {
+                    TIMEOUT_REFRESHED_ACCESS_TOKEN
+                }
+                _ => OLD_ACCESS_TOKEN,
+            };
+            oauth_mcp_response(request, expected_access_token)
+        })
+        .expect(5)
+        .mount(&server)
+        .await;
+
+    let server_url = format!("{}/mcp", server.uri());
+    let status = Command::new(std::env::current_exe()?)
+        .args([
+            "oauth_operation_timeout_child",
+            "--exact",
+            "--ignored",
+            "--nocapture",
+        ])
+        .env("CODEX_HOME", codex_home.path())
+        .env(OPERATION_TIMEOUT_CHILD_SERVER_URL_ENV, server_url)
+        .env(OPERATION_TIMEOUT_REFRESH_STARTED_ENV, refresh_started)
+        .status()
+        .await?;
+    assert!(
+        status.success(),
+        "OAuth operation timeout child failed: {status}"
+    );
+    server.verify().await;
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+#[ignore = "spawned by operation_timeout_includes_oauth_generation_lock_wait"]
+async fn oauth_operation_timeout_child() -> anyhow::Result<()> {
+    let server_url = std::env::var(OPERATION_TIMEOUT_CHILD_SERVER_URL_ENV)?;
+    let refresh_started =
+        std::path::PathBuf::from(std::env::var(OPERATION_TIMEOUT_REFRESH_STARTED_ENV)?);
+    save_test_tokens(
+        &server_url,
+        OLD_ACCESS_TOKEN,
+        OLD_REFRESH_TOKEN,
+        future_expiry()?,
+    )?;
+    let lock_holder = Arc::new(initialized_oauth_client(&server_url).await?);
+    let timed_client = initialized_oauth_client(&server_url).await?;
+    save_test_tokens(
+        &server_url,
+        EXTERNAL_ACCESS_TOKEN,
+        EXTERNAL_REFRESH_TOKEN,
+        /*expires_at*/ 0,
+    )?;
+
+    let lock_holder_call = tokio::spawn(async move {
+        lock_holder
+            .call_tool(
+                "hold-generation-lock".to_string(),
+                /*arguments*/ None,
+                /*meta*/ None,
+                Some(Duration::from_secs(5)),
+            )
+            .await
+    });
+    wait_for_path(&refresh_started).await?;
+    let error = timed_client
+        .call_tool(
+            "must-time-out-on-generation-lock".to_string(),
+            /*arguments*/ None,
+            /*meta*/ None,
+            Some(Duration::from_millis(100)),
+        )
+        .await
+        .expect_err("generation lock wait must consume the operation timeout");
+    assert!(
+        error
+            .to_string()
+            .contains("timed out awaiting tools/call after 100ms"),
+        "unexpected timeout error: {error}"
+    );
+    lock_holder_call.await??;
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn running_client_adopts_external_oauth_update_and_deletion() -> anyhow::Result<()> {
    let server = MockServer::start().await;
-    mount_oauth_metadata(&server).await;
+    mount_oauth_metadata(&server, /*expected_calls*/ 2).await;

    Mock::given(method("POST"))
        .and(path("/mcp"))
@@ -168,7 +359,7 @@ async fn oauth_external_delete_writer_child() -> anyhow::Result<()> {
 async fn refresh_holds_generation_lock_against_external_oauth_update() -> anyhow::Result<()> {
    let server = MockServer::start().await;
    let server_url = format!("{}/mcp", server.uri());
-    mount_oauth_metadata(&server).await;
+    mount_oauth_metadata(&server, /*expected_calls*/ 3).await;
    Mock::given(method("POST"))
        .and(path("/oauth/token"))
        .and(body_string_contains("grant_type=refresh_token"))
@@ -322,7 +513,7 @@ async fn oauth_external_update_writer_child() -> anyhow::Result<()> {
    Ok(())
 }

-async fn mount_oauth_metadata(server: &MockServer) {
+async fn mount_oauth_metadata(server: &MockServer, expected_calls: u64) {
    Mock::given(method("GET"))
        .and(path("/.well-known/oauth-authorization-server/mcp"))
        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
@@ -330,7 +521,7 @@ async fn mount_oauth_metadata(server: &MockServer) {
            "token_endpoint": format!("{}/oauth/token", server.uri()),
            "scopes_supported": [""],
        })))
-        .expect(1)
+        .expect(expected_calls)
        .mount(server)
        .await;
 }