fix(rmcp): budget fallback write lock waits

codex-rs/rmcp-client/src/oauth.rs

@@ -247,16 +247,23 @@ fn save_oauth_tokens_with_keyring<K: KeyringStore>(
     server_name: &str,
     tokens: &StoredOAuthTokens,
 ) -> Result<()> {
+    let key = save_oauth_tokens_to_keyring(keyring_store, server_name, tokens)?;
+    if let Err(error) = delete_oauth_tokens_from_file(&key) {
+        warn!("failed to remove OAuth tokens from fallback storage: {error:?}");
+    }
+    Ok(())
+}
+
+fn save_oauth_tokens_to_keyring<K: KeyringStore>(
+    keyring_store: &K,
+    server_name: &str,
+    tokens: &StoredOAuthTokens,
+) -> Result<String> {
     let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?;
 
     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(key),
         Err(error) => {
             let message = format!(
                 "failed to write OAuth tokens to keyring: {}",
@@ -284,6 +291,46 @@ fn save_oauth_tokens_with_keyring_with_fallback_to_file<K: KeyringStore>(
     }
 }
 
+async fn save_oauth_tokens_for_runtime(
+    server_name: &str,
+    tokens: &StoredOAuthTokens,
+    store_mode: OAuthCredentialsStoreMode,
+) -> Result<()> {
+    let keyring_store = DefaultKeyringStore;
+    match store_mode {
+        OAuthCredentialsStoreMode::Auto => {
+            match save_oauth_tokens_with_keyring_async(&keyring_store, server_name, tokens).await {
+                Ok(()) => Ok(()),
+                Err(error) => {
+                    let message = error.to_string();
+                    warn!("falling back to file storage for OAuth tokens: {message}");
+                    save_oauth_tokens_to_file_async(tokens)
+                        .await
+                        .with_context(|| {
+                            format!("failed to write OAuth tokens to keyring: {message}")
+                        })
+                }
+            }
+        }
+        OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file_async(tokens).await,
+        OAuthCredentialsStoreMode::Keyring => {
+            save_oauth_tokens_with_keyring_async(&keyring_store, server_name, tokens).await
+        }
+    }
+}
+
+async fn save_oauth_tokens_with_keyring_async<K: KeyringStore>(
+    keyring_store: &K,
+    server_name: &str,
+    tokens: &StoredOAuthTokens,
+) -> Result<()> {
+    let key = save_oauth_tokens_to_keyring(keyring_store, server_name, tokens)?;
+    if let Err(error) = delete_oauth_tokens_from_file_async(&key).await {
+        warn!("failed to remove OAuth tokens from fallback storage: {error:?}");
+    }
+    Ok(())
+}
+
 pub fn delete_oauth_tokens(
     server_name: &str,
     url: &str,
@@ -300,7 +347,29 @@ fn delete_oauth_tokens_from_keyring_and_file<K: KeyringStore>(
     url: &str,
 ) -> Result<bool> {
     let key = compute_store_key(server_name, url)?;
-    let keyring_result = keyring_store.delete(KEYRING_SERVICE, &key);
+    let keyring_removed = delete_oauth_tokens_from_keyring(keyring_store, store_mode, &key)?;
+    let file_removed = delete_oauth_tokens_from_file(&key)?;
+    Ok(keyring_removed || file_removed)
+}
+
+async fn delete_oauth_tokens_for_runtime(
+    server_name: &str,
+    url: &str,
+    store_mode: OAuthCredentialsStoreMode,
+) -> Result<bool> {
+    let keyring_store = DefaultKeyringStore;
+    let key = compute_store_key(server_name, url)?;
+    let keyring_removed = delete_oauth_tokens_from_keyring(&keyring_store, store_mode, &key)?;
+    let file_removed = delete_oauth_tokens_from_file_async(&key).await?;
+    Ok(keyring_removed || file_removed)
+}
+
+fn delete_oauth_tokens_from_keyring<K: KeyringStore>(
+    keyring_store: &K,
+    store_mode: OAuthCredentialsStoreMode,
+    key: &str,
+) -> Result<bool> {
+    let keyring_result = keyring_store.delete(KEYRING_SERVICE, key);
     let keyring_removed = match keyring_result {
         Ok(removed) => removed,
         Err(error) => {
@@ -315,9 +384,7 @@ fn delete_oauth_tokens_from_keyring_and_file<K: KeyringStore>(
             }
         }
     };
-
-    let file_removed = delete_oauth_tokens_from_file(&key)?;
-    Ok(keyring_removed || file_removed)
+    Ok(keyring_removed)
 }
 
 #[derive(Clone)]
@@ -389,18 +456,24 @@ impl OAuthPersistor {
                     expires_at,
                 };
                 if last_credentials.as_ref() != Some(&stored) {
-                    save_oauth_tokens(&self.inner.server_name, &stored, self.inner.store_mode)?;
+                    save_oauth_tokens_for_runtime(
+                        &self.inner.server_name,
+                        &stored,
+                        self.inner.store_mode,
+                    )
+                    .await?;
                     *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(
+                    && let Err(error) = delete_oauth_tokens_for_runtime(
                         &self.inner.server_name,
                         &self.inner.url,
                         self.inner.store_mode,
                     )
+                    .await
                 {
                     warn!(
                         "failed to remove OAuth tokens for server {}: {error}",
@@ -744,6 +817,27 @@ fn acquire_fallback_write_lock() -> Result<fs::File> {
     Ok(file)
 }
 
+async fn acquire_fallback_write_lock_async() -> Result<fs::File> {
+    let path = fallback_lock_path()?;
+    let file = open_oauth_lock_file(&path)?;
+    loop {
+        match file.try_lock() {
+            Ok(()) => return Ok(file),
+            Err(fs::TryLockError::WouldBlock) => {
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+            Err(error) => {
+                return Err(error).with_context(|| {
+                    format!(
+                        "failed to acquire OAuth fallback write lock at {}",
+                        path.display()
+                    )
+                });
+            }
+        }
+    }
+}
+
 fn acquire_fallback_read_lock() -> Result<fs::File> {
     let path = fallback_lock_path()?;
     let file = open_oauth_lock_file(&path)?;
@@ -931,6 +1025,15 @@ fn load_oauth_tokens_from_file_unlocked(
 
 fn save_oauth_tokens_to_file(tokens: &StoredOAuthTokens) -> Result<()> {
     let _write_lock = acquire_fallback_write_lock()?;
+    save_oauth_tokens_to_file_unlocked(tokens)
+}
+
+async fn save_oauth_tokens_to_file_async(tokens: &StoredOAuthTokens) -> Result<()> {
+    let _write_lock = acquire_fallback_write_lock_async().await?;
+    save_oauth_tokens_to_file_unlocked(tokens)
+}
+
+fn save_oauth_tokens_to_file_unlocked(tokens: &StoredOAuthTokens) -> Result<()> {
     let key = compute_store_key(&tokens.server_name, &tokens.url)?;
     let mut store = read_fallback_file_unlocked()?.unwrap_or_default();
 
@@ -964,6 +1067,11 @@ fn delete_oauth_tokens_from_file(key: &str) -> Result<bool> {
     delete_oauth_tokens_from_file_unlocked(key)
 }
 
+async fn delete_oauth_tokens_from_file_async(key: &str) -> Result<bool> {
+    let _write_lock = acquire_fallback_write_lock_async().await?;
+    delete_oauth_tokens_from_file_unlocked(key)
+}
+
 fn delete_oauth_tokens_from_file_unlocked(key: &str) -> Result<bool> {
     let mut store = match read_fallback_file_unlocked()? {
         Some(store) => store,

codex-rs/rmcp-client/tests/streamable_http_oauth_refresh_race.rs

@@ -272,6 +272,61 @@ async fn fallback_read_lock_wait_uses_startup_timeout_budget() -> anyhow::Result
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn fallback_write_lock_wait_uses_startup_timeout_budget() -> anyhow::Result<()> {
+    let server = MockServer::start().await;
+    mount_oauth_metadata(&server, /*expected_requests*/ 1).await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "access_token": NEW_ACCESS_TOKEN,
+            "token_type": "Bearer",
+            "expires_in": 7200,
+            "refresh_token": NEW_REFRESH_TOKEN,
+        })))
+        .expect(1)
+        .mount(&server)
+        .await;
+    mount_mcp_server(&server, /*expected_requests*/ 0).await;
+
+    let codex_home = TempDir::new()?;
+    let control_dir = TempDir::new()?;
+    let server_url = format!("{}/mcp", server.uri());
+    seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?;
+
+    let (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?;
+    let fallback_lock = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .create(true)
+        .truncate(false)
+        .open(lock_path)?;
+    fallback_lock.lock_shared()?;
+
+    let ready_path = control_dir.path().join("client.ready");
+    let go_path = control_dir.path().join("client.go");
+    let result_path = control_dir.path().join("client.result");
+    let child = spawn_client_child(
+        &codex_home,
+        &server_url,
+        &ready_path,
+        &go_path,
+        &result_path,
+    )?;
+    wait_for_paths(std::slice::from_ref(&ready_path)).await?;
+    fs::write(&go_path, "go")?;
+    wait_for_paths(std::slice::from_ref(&result_path)).await?;
+    assert_eq!(
+        fs::read_to_string(&result_path)?,
+        "error:timed out handshaking with MCP server after 5s"
+    );
+
+    drop(fallback_lock);
+    wait_for_children(vec![child]).await?;
+    server.verify().await;
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn fallback_read_lock_wait_uses_operation_timeout_budget() -> anyhow::Result<()> {
     let server = MockServer::start().await;
@@ -328,6 +383,62 @@ async fn fallback_read_lock_wait_uses_operation_timeout_budget() -> anyhow::Resu
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn fallback_write_lock_wait_uses_operation_timeout_budget() -> anyhow::Result<()> {
+    let server = MockServer::start().await;
+    mount_oauth_metadata(&server, /*expected_requests*/ 1).await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "access_token": NEW_ACCESS_TOKEN,
+            "token_type": "Bearer",
+            "expires_in": 31,
+            "refresh_token": NEW_REFRESH_TOKEN,
+        })))
+        .expect(2)
+        .mount(&server)
+        .await;
+    mount_mcp_server(&server, /*expected_requests*/ 2).await;
+
+    let codex_home = TempDir::new()?;
+    let control_dir = TempDir::new()?;
+    let server_url = format!("{}/mcp", server.uri());
+    seed_tokens(&codex_home, &server_url, /*expires_at*/ 0).await?;
+
+    let ready_path = control_dir.path().join("client.ready");
+    let go_path = control_dir.path().join("client.go");
+    let result_path = control_dir.path().join("client.result");
+    let child = spawn_operation_timeout_child(
+        &codex_home,
+        &server_url,
+        &ready_path,
+        &go_path,
+        &result_path,
+    )?;
+    wait_for_paths(std::slice::from_ref(&ready_path)).await?;
+    sleep(Duration::from_secs(2)).await;
+
+    let (lock_path, _) = fallback_lock_path_and_user_namespace(&codex_home)?;
+    let fallback_lock = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .create(true)
+        .truncate(false)
+        .open(lock_path)?;
+    fallback_lock.lock_shared()?;
+
+    fs::write(&go_path, "go")?;
+    wait_for_children(vec![child]).await?;
+    assert_eq!(
+        fs::read_to_string(result_path)?,
+        format!("error:timed out awaiting tools/list after {OPERATION_TIMEOUT:?}")
+    );
+
+    drop(fallback_lock);
+    server.verify().await;
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn refresh_provider_call_uses_operation_timeout_budget() -> anyhow::Result<()> {
     let server = MockServer::start().await;