Fix concurrent token refresh marking healthy accounts for re-authentication

backend/app/routes/proxy.py

@@ -888,7 +888,7 @@ async def proxy_messages(
             if candidate.status_code == 401:
                 await candidate.aclose()
                 connection = egress.get_pool().acquire(account.egress_target_id, priority=user.priority)
-                access_token = rotation.ensure_fresh_token(db, account, force_refresh=True, egress_target=connection.target)
+                access_token = rotation.ensure_fresh_token(db, account, rejected_token=access_token, egress_target=connection.target)
                 candidate = await _send_with_account_limit(
                     connection,
                     request.method,

backend/app/utils/rotation.py

@@ -167,20 +167,29 @@ def select_account(
     return chosen
 
 
+def _stored_access_token(db: Session, account: AccountDb) -> str:
+    try:
+        return crypto.decrypt(account.access_token_enc)
+    except Exception as exc:
+        provider_health.persist_failure(db, account.id, exc, context="credential_decrypt")
+        raise provider_health.reauthentication_error() from exc
+
+
 def ensure_fresh_token(
     db: Session,
     account: AccountDb,
     *,
-    force_refresh: bool = False,
+    rejected_token: Optional[str] = None,
     egress_target=None,
 ) -> str:
     """Return a usable access token, serializing refresh-token rotation in the database.
 
-    ``force_refresh`` is used after an upstream 401.  The provider can revoke an
-    access token before its JWT expiry, so callers must be able to refresh it
-    without writing a synthetic expiry into the account row.  Keeping that
-    decision inside the row lock prevents a stale request from clobbering a
-    concurrently rotated token/expiry.
+    ``rejected_token`` is the access token an upstream 401 just rejected. The
+    provider can revoke an access token before its JWT expiry, so that token is
+    refreshed regardless of expiry -- unless a concurrent request already
+    replaced it, in which case the replacement is returned. The row is re-read
+    from the database under the lock, so a refresh always spends the latest
+    refresh token rather than one another worker has already rotated.
     """
     if account.provider_health == ProviderHealth.REAUTH_REQUIRED:
         raise provider_health.reauthentication_error()
@@ -190,23 +199,24 @@ def ensure_fresh_token(
     expires_at = account.expires_at
     if expires_at.tzinfo is None:
         expires_at = expires_at.replace(tzinfo=timezone.utc)
-    if not force_refresh and expires_at - leeway > now:
-        try:
-            return crypto.decrypt(account.access_token_enc)
-        except Exception as exc:
-            provider_health.persist_failure(db, account.id, exc, context="credential_decrypt")
-            raise provider_health.reauthentication_error() from exc
+    if rejected_token is None and expires_at - leeway > now:
+        return _stored_access_token(db, account)
 
-    locked = db.query(AccountDb).filter(AccountDb.id == account.id).with_for_update().one()
-    locked_expiry = locked.expires_at
-    if locked_expiry.tzinfo is None:
-        locked_expiry = locked_expiry.replace(tzinfo=timezone.utc)
-    if not force_refresh and locked_expiry - leeway > now:
-        try:
-            return crypto.decrypt(locked.access_token_enc)
-        except Exception as exc:
-            provider_health.persist_failure(db, account.id, exc, context="credential_decrypt")
-            raise provider_health.reauthentication_error() from exc
+    # populate_existing() discards unflushed changes on the identity-mapped row, so flush them first.
+    db.flush()
+    locked = db.query(AccountDb).filter(AccountDb.id == account.id).with_for_update().populate_existing().one()
+    if locked.provider_health == ProviderHealth.REAUTH_REQUIRED:
+        raise provider_health.reauthentication_error()
+    if rejected_token is None:
+        locked_expiry = locked.expires_at
+        if locked_expiry.tzinfo is None:
+            locked_expiry = locked_expiry.replace(tzinfo=timezone.utc)
+        if locked_expiry - leeway > now:
+            return _stored_access_token(db, locked)
+    else:
+        current = _stored_access_token(db, locked)
+        if current != rejected_token:
+            return current
 
     try:
         refresh_plain = crypto.decrypt(locked.refresh_token_enc)

backend/tests/test_rotation.py

@@ -57,11 +57,9 @@ def test_new_account_gets_default_rotation_policy(seed_account):
         assert acct.max_failover_attempts == config.DEFAULT_MAX_FAILOVER_ATTEMPTS
 
 
-def test_force_refresh_does_not_require_a_synthetic_expiry(seed_account, monkeypatch):
+def test_rejected_token_refreshes_without_a_synthetic_expiry(seed_account, monkeypatch):
     """A provider 401 must refresh the current row without clobbering its expiry."""
     from app.utils import oauth
-    from app.utils.postgres import AccountDb
-    from app.utils.postgres.base import SessionFactory
 
     account_id = seed_account("force-refresh", expires_in_hours=1)
     refreshed_expiry = datetime.now(timezone.utc) + timedelta(hours=2)
@@ -73,11 +71,75 @@ def test_force_refresh_does_not_require_a_synthetic_expiry(seed_account, monkeyp
 
     with SessionFactory() as db:
         account = db.get(AccountDb, account_id)
-        assert rotation.ensure_fresh_token(db, account, force_refresh=True) == "rotated-access"
+        assert rotation.ensure_fresh_token(db, account, rejected_token="upstream-access-token") == "rotated-access"
         persisted = db.get(AccountDb, account_id)
         assert persisted.expires_at == refreshed_expiry
 
 
+def _rotate_in_other_session(account_id, access_token: str, refresh_token: str, expires_at: datetime) -> None:
+    from app.utils import crypto
+
+    with SessionFactory() as other:
+        row = other.get(AccountDb, account_id)
+        row.access_token_enc = crypto.encrypt(access_token)
+        row.refresh_token_enc = crypto.encrypt(refresh_token)
+        row.expires_at = expires_at
+        other.commit()
+
+
+def test_expiring_token_reuses_a_concurrent_rotation(seed_account, monkeypatch):
+    """A worker holding a stale row must not spend a refresh token another worker already rotated."""
+    from app.utils import oauth
+
+    account_id = seed_account("concurrent-expiry", expires_in_hours=0)
+    calls = []
+    monkeypatch.setattr(oauth, "refresh_access_token", lambda refresh: calls.append(refresh))
+
+    with SessionFactory() as db:
+        account = db.get(AccountDb, account_id)
+        assert account.expires_at is not None  # loaded before the other worker rotates
+        _rotate_in_other_session(account_id, "other-access", "other-refresh", datetime.now(timezone.utc) + timedelta(hours=8))
+        assert rotation.ensure_fresh_token(db, account) == "other-access"
+    assert calls == []
+
+
+def test_rejected_token_reuses_a_concurrent_rotation(seed_account, monkeypatch):
+    """Concurrent 401s on one token trigger one refresh; later callers get the replacement."""
+    from app.utils import oauth
+
+    account_id = seed_account("concurrent-401", expires_in_hours=1)
+    calls = []
+    monkeypatch.setattr(oauth, "refresh_access_token", lambda refresh: calls.append(refresh))
+
+    with SessionFactory() as db:
+        account = db.get(AccountDb, account_id)
+        assert account.expires_at is not None
+        _rotate_in_other_session(account_id, "other-access", "other-refresh", datetime.now(timezone.utc) + timedelta(hours=8))
+        assert rotation.ensure_fresh_token(db, account, rejected_token="upstream-access-token") == "other-access"
+    assert calls == []
+
+
+def test_refresh_spends_the_latest_refresh_token(seed_account, monkeypatch):
+    from app.utils import oauth
+
+    account_id = seed_account("latest-refresh", expires_in_hours=0)
+    spent = []
+
+    def refresh(refresh_token):
+        spent.append(refresh_token)
+        return {"access_token": "new-access", "refresh_token": "new-refresh", "expires_at": datetime.now(timezone.utc) + timedelta(hours=8)}
+
+    monkeypatch.setattr(oauth, "refresh_access_token", refresh)
+
+    with SessionFactory() as db:
+        account = db.get(AccountDb, account_id)
+        assert account.expires_at is not None
+        # Another worker rotated, but its new access token is already near expiry, so this worker must refresh again.
+        _rotate_in_other_session(account_id, "other-access", "other-refresh", datetime.now(timezone.utc))
+        assert rotation.ensure_fresh_token(db, account) == "new-access"
+    assert spent == ["other-refresh"]
+
+
 def test_is_available_uses_per_account_threshold(seed_account):
     # The same 0.50 utilization is benched under a 0.40 per-account threshold but available under the default 1.0.
     strict = seed_account("strict", session_used_pct=0.50, rotation_threshold=0.40)