From 0d24353fd31ab68f3c6cdebf0745730db540f832 Mon Sep 17 00:00:00 2001 From: rayan-mazouz Date: Thu, 8 Oct 2026 03:13:09 +0200 Subject: [PATCH 1/3] refactor(rotation): extract stored access token decryption Both the fast path and the locked path decrypted the stored access token and flagged the account for re-authentication on failure. Share one helper so later changes to the locked path stay small. --- backend/app/utils/rotation.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/backend/app/utils/rotation.py b/backend/app/utils/rotation.py index 8f50802..01b0688 100644 --- a/backend/app/utils/rotation.py +++ b/backend/app/utils/rotation.py @@ -167,6 +167,14 @@ 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, @@ -191,22 +199,14 @@ def ensure_fresh_token( 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 + 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 + return _stored_access_token(db, locked) try: refresh_plain = crypto.decrypt(locked.refresh_token_enc) From 587f878b35e6890731b1b95821e5c0c3d8d1eccb Mon Sep 17 00:00:00 2001 From: rayan-mazouz Date: Thu, 8 Oct 2026 03:13:31 +0200 Subject: [PATCH 2/3] fix(rotation): re-read the account row under the refresh lock ensure_fresh_token locked the row with SELECT ... FOR UPDATE and then re-checked the expiry, but the query returned the identity-mapped object with its previously loaded attributes. A worker that loaded the account before another worker rotated its tokens therefore saw the old expiry and spent the already-rotated refresh token. The provider rejected it and the healthy account was marked as needing re-authentication. Re-read the row with populate_existing() (flushing pending changes first so they are not discarded), and stop if the account was marked for re-authentication in the meantime. --- backend/app/utils/rotation.py | 10 +++++-- backend/tests/test_rotation.py | 48 ++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/backend/app/utils/rotation.py b/backend/app/utils/rotation.py index 01b0688..508f759 100644 --- a/backend/app/utils/rotation.py +++ b/backend/app/utils/rotation.py @@ -188,7 +188,9 @@ def ensure_fresh_token( 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. + concurrently rotated token/expiry. 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() @@ -201,7 +203,11 @@ def ensure_fresh_token( if not force_refresh and expires_at - leeway > now: return _stored_access_token(db, account) - locked = db.query(AccountDb).filter(AccountDb.id == account.id).with_for_update().one() + # 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() locked_expiry = locked.expires_at if locked_expiry.tzinfo is None: locked_expiry = locked_expiry.replace(tzinfo=timezone.utc) diff --git a/backend/tests/test_rotation.py b/backend/tests/test_rotation.py index 49007d8..3ee298f 100644 --- a/backend/tests/test_rotation.py +++ b/backend/tests/test_rotation.py @@ -78,6 +78,54 @@ def test_force_refresh_does_not_require_a_synthetic_expiry(seed_account, monkeyp 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_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) From d45e365e445ba1da5216ca90168b882f2e0752c1 Mon Sep 17 00:00:00 2001 From: rayan-mazouz Date: Thu, 8 Oct 2026 03:13:39 +0200 Subject: [PATCH 3/3] fix(rotation): reuse a concurrent replacement after an upstream 401 When several in-flight requests hit a 401 on the same revoked token, each forced its own refresh, spending one refresh token per request. Pass the rejected access token instead of a force flag: under the row lock, a token that no longer matches it was already replaced by another worker and is returned as-is, so only the first request refreshes. --- backend/app/routes/proxy.py | 2 +- backend/app/utils/rotation.py | 32 ++++++++++++++++++-------------- backend/tests/test_rotation.py | 22 ++++++++++++++++++---- 3 files changed, 37 insertions(+), 19 deletions(-) diff --git a/backend/app/routes/proxy.py b/backend/app/routes/proxy.py index dfc898d..f06edda 100644 --- a/backend/app/routes/proxy.py +++ b/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, diff --git a/backend/app/utils/rotation.py b/backend/app/utils/rotation.py index 508f759..0bcc83d 100644 --- a/backend/app/utils/rotation.py +++ b/backend/app/utils/rotation.py @@ -179,18 +179,17 @@ 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. 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. + ``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() @@ -200,7 +199,7 @@ 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: + if rejected_token is None and expires_at - leeway > now: return _stored_access_token(db, account) # populate_existing() discards unflushed changes on the identity-mapped row, so flush them first. @@ -208,11 +207,16 @@ def ensure_fresh_token( 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() - 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: - return _stored_access_token(db, locked) + 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) diff --git a/backend/tests/test_rotation.py b/backend/tests/test_rotation.py index 3ee298f..a7bc078 100644 --- a/backend/tests/test_rotation.py +++ b/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,7 +71,7 @@ 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 @@ -105,6 +103,22 @@ def test_expiring_token_reuses_a_concurrent_rotation(seed_account, monkeypatch): 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