From 70f1d02e595a346ac751fe0d4c2f4ae7eca3b195 Mon Sep 17 00:00:00 2001 From: Rohit Sabu <13933510+rohitsabu@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:32:49 +0400 Subject: [PATCH] fix(auth): preserve desktop sessions across refresh and restart --- apps/desktop/electron/main.ts | 6 +- apps/desktop/electron/native-oauth.test.ts | 69 ++++++++++++++ apps/desktop/electron/native-oauth.ts | 37 ++++++++ hermes_cli/dashboard_auth/middleware.py | 95 ++++++++++++++++++- .../test_dashboard_auth_401_reauth.py | 51 +++++++++- 5 files changed, 254 insertions(+), 4 deletions(-) diff --git a/apps/desktop/electron/main.ts b/apps/desktop/electron/main.ts index 5f6ba957ff8ce..b87762ade8905 100644 --- a/apps/desktop/electron/main.ts +++ b/apps/desktop/electron/main.ts @@ -130,8 +130,10 @@ import { resolveReadinessProbeAuth } from './native-auth-decisions' import { + mergeRefreshedNativeTokenSet, nativeRefreshUrl, type NativeTokenSet, + parseStoredNativeTokenSet, parseTokenResponse, resolveLoginStrategy, tokenNeedsRefresh @@ -6043,7 +6045,7 @@ function _loadNativeTokens(baseUrl: string): NativeTokenSet | null { return null } - const tokens = parseTokenResponse(JSON.parse(plaintext)) + const tokens = parseStoredNativeTokenSet(JSON.parse(plaintext)) _nativeTokens.set(baseUrl, tokens) return tokens @@ -6108,7 +6110,7 @@ async function ensureNativeAccessToken(baseUrl: string): Promise { timeoutMs: 10_000 } ) - const rotated = parseTokenResponse(body) + const rotated = mergeRefreshedNativeTokenSet(tokens, body) _storeNativeTokens(baseUrl, rotated) return rotated.accessToken diff --git a/apps/desktop/electron/native-oauth.test.ts b/apps/desktop/electron/native-oauth.test.ts index 58d863c4b1a60..a0063892456a6 100644 --- a/apps/desktop/electron/native-oauth.test.ts +++ b/apps/desktop/electron/native-oauth.test.ts @@ -16,10 +16,12 @@ import { buildNativeAuthorizeUrl, generatePkcePair, generateState, + mergeRefreshedNativeTokenSet, NATIVE_FLOW_ID, nativeRefreshUrl, nativeTokenUrl, parseLoopbackCallback, + parseStoredNativeTokenSet, parseTokenResponse, resolveLoginStrategy, statusSupportsNativeFlow, @@ -176,6 +178,73 @@ test('parseTokenResponse tolerates an absent refresh token / expiry', () => { assert.equal(t.expiresAt, 0) }) +// --- persisted token-set restoration --- + +test('parseStoredNativeTokenSet restores the JSON shape encrypted by the desktop', () => { + const stored = { + accessToken: 'AT-persisted', + refreshToken: 'RT-persisted', + expiresAt: 1_893_456_000, + provider: 'nous', + userId: 'u-persisted' + } + + const persistedPlaintext = JSON.stringify(stored) + + assert.deepEqual(parseStoredNativeTokenSet(JSON.parse(persistedPlaintext)), stored) +}) + +test('parseStoredNativeTokenSet rejects a persisted record without an access token', () => { + assert.throws( + () => + parseStoredNativeTokenSet({ + refreshToken: 'RT-persisted', + expiresAt: 1_893_456_000, + provider: 'nous', + userId: 'u-persisted' + }), + /missing accessToken/i + ) +}) + +test('mergeRefreshedNativeTokenSet preserves the prior refresh token when rotation omits one', () => { + const previous = { + accessToken: 'AT-old', + refreshToken: 'RT-still-live', + expiresAt: 1_800_000_000, + provider: 'nous', + userId: 'u-1' + } + const refreshed = mergeRefreshedNativeTokenSet(previous, { + access_token: 'AT-new', + expires_at: 1_900_000_000, + provider: 'nous', + user_id: 'u-1' + }) + + assert.equal(refreshed.accessToken, 'AT-new') + assert.equal(refreshed.refreshToken, 'RT-still-live') +}) + +test('mergeRefreshedNativeTokenSet accepts a rotated refresh token', () => { + const previous = { + accessToken: 'AT-old', + refreshToken: 'RT-old', + expiresAt: 1_800_000_000, + provider: 'nous', + userId: 'u-1' + } + const refreshed = mergeRefreshedNativeTokenSet(previous, { + access_token: 'AT-new', + refresh_token: 'RT-new', + expires_at: 1_900_000_000, + provider: 'nous', + user_id: 'u-1' + }) + + assert.equal(refreshed.refreshToken, 'RT-new') +}) + // --- refresh timing --- test('tokenNeedsRefresh respects the skew window', () => { diff --git a/apps/desktop/electron/native-oauth.ts b/apps/desktop/electron/native-oauth.ts index 691cd32caca69..ea34283a17cd0 100644 --- a/apps/desktop/electron/native-oauth.ts +++ b/apps/desktop/electron/native-oauth.ts @@ -193,6 +193,43 @@ export function parseTokenResponse(body: any): NativeTokenSet { } } +/** + * Validate and restore the camelCase token shape encrypted by the desktop. + * + * This is deliberately separate from parseTokenResponse: the gateway wire + * format is snake_case, while JSON.stringify(NativeTokenSet) persists the + * normalized camelCase object. Feeding the latter back through the wire parser + * makes every post-restart load look like a missing access token. + */ +export function parseStoredNativeTokenSet(body: any): NativeTokenSet { + const accessToken = String(body?.accessToken || '') + + if (!accessToken) { + throw new Error('Stored native token set missing accessToken') + } + + const expiresAt = Number(body?.expiresAt) + + return { + accessToken, + refreshToken: String(body?.refreshToken || ''), + expiresAt: Number.isFinite(expiresAt) ? expiresAt : 0, + provider: String(body?.provider || ''), + userId: String(body?.userId || '') + } +} + +/** + * Normalize a refresh response without discarding a still-valid refresh token + * when an OAuth provider omits refresh_token. Providers that rotate return a + * replacement; providers that do not rotate leave the previous token live. + */ +export function mergeRefreshedNativeTokenSet(previous: NativeTokenSet, body: any): NativeTokenSet { + const refreshed = parseTokenResponse(body) + + return refreshed.refreshToken ? refreshed : { ...refreshed, refreshToken: previous.refreshToken } +} + /** * True when a stored token set is at/near expiry and should be refreshed * before use. `skewSeconds` refreshes slightly early to avoid a race where diff --git a/hermes_cli/dashboard_auth/middleware.py b/hermes_cli/dashboard_auth/middleware.py index 5b11e98cf2b16..3d0387a7c9f66 100644 --- a/hermes_cli/dashboard_auth/middleware.py +++ b/hermes_cli/dashboard_auth/middleware.py @@ -16,7 +16,11 @@ """ from __future__ import annotations +import hashlib import logging +import threading +import time +from collections import OrderedDict from typing import Awaitable, Callable from fastapi import Request @@ -41,6 +45,26 @@ _log = logging.getLogger(__name__) +# A browser can have several authenticated requests in flight when its access +# token expires. They all carry the same old refresh-token cookie because the +# first response has not reached the browser yet. Providers such as Nous rotate +# refresh tokens with reuse detection, so calling the provider once per stale +# request can turn a healthy rotation into a replay failure whose 401 response +# clears the newly-issued cookies. +# +# Keep the just-rotated Session briefly, keyed by a non-reversible token digest, +# client address, and provider-registry identity. A stale parallel request then +# receives the same rotated cookies without replaying the old token upstream. +# The cache is process-local, bounded, never persisted or logged, and much +# shorter-lived than either credential. +_REFRESH_REPLAY_GRACE_SECONDS = 15.0 +_REFRESH_REPLAY_CACHE_MAX_ENTRIES = 256 +_refresh_replay_cache: OrderedDict[ + tuple[bytes, str, tuple[tuple[str, int], ...]], + tuple[float, tuple[object, str]], +] = OrderedDict() +_refresh_replay_cache_lock = threading.Lock() + # Prefixes that bypass the auth gate. Match via ``path == prefix`` or # ``path.startswith(prefix)`` — so ``/assets/`` (with trailing slash) # matches ``/assets/foo.css`` but not ``/assetsleak``. Auth-bootstrap @@ -454,7 +478,7 @@ async def gated_auth_middleware( # serve the request transparently; only after every provider rejects # the RT do we fall through to clear-and-relogin. try: - refreshed = _attempt_refresh( + refreshed = _attempt_refresh_with_replay_grace( request, refresh_token=_rt, provider_hint=provider_hint, @@ -589,3 +613,72 @@ def _attempt_refresh(request: Request, *, refresh_token, provider_hint: str | No if unavailable_provider is not None: raise ProviderError(unavailable_provider) return None + + +def _refresh_replay_key( + request: Request, + *, + refresh_token: str, + provider_hint: str | None, +) -> tuple[bytes, str, tuple[tuple[str, int], ...]]: + providers = _ordered_session_providers(provider_hint) + provider_scope = tuple((provider.name, id(provider)) for provider in providers) + digest = hashlib.sha256(refresh_token.encode("utf-8")).digest() + return digest, _client_ip(request), provider_scope + + +def _get_recent_refresh(key): + now = time.monotonic() + with _refresh_replay_cache_lock: + while _refresh_replay_cache: + oldest_key, (expires_at, _) = next(iter(_refresh_replay_cache.items())) + if expires_at > now: + break + _refresh_replay_cache.pop(oldest_key, None) + + cached = _refresh_replay_cache.get(key) + if cached is None or cached[0] <= now: + _refresh_replay_cache.pop(key, None) + return None + + _refresh_replay_cache.move_to_end(key) + return cached[1] + + +def _remember_recent_refresh(key, refreshed) -> None: + with _refresh_replay_cache_lock: + _refresh_replay_cache[key] = ( + time.monotonic() + _REFRESH_REPLAY_GRACE_SECONDS, + refreshed, + ) + _refresh_replay_cache.move_to_end(key) + while len(_refresh_replay_cache) > _REFRESH_REPLAY_CACHE_MAX_ENTRIES: + _refresh_replay_cache.popitem(last=False) + + +def _attempt_refresh_with_replay_grace( + request: Request, + *, + refresh_token, + provider_hint: str | None = None, +): + if not refresh_token: + return None + + key = _refresh_replay_key( + request, + refresh_token=refresh_token, + provider_hint=provider_hint, + ) + cached = _get_recent_refresh(key) + if cached is not None: + return cached + + refreshed = _attempt_refresh( + request, + refresh_token=refresh_token, + provider_hint=provider_hint, + ) + if refreshed is not None: + _remember_recent_refresh(key, refreshed) + return refreshed diff --git a/tests/hermes_cli/test_dashboard_auth_401_reauth.py b/tests/hermes_cli/test_dashboard_auth_401_reauth.py index 63ab76a9a8338..1073aacd77ac6 100644 --- a/tests/hermes_cli/test_dashboard_auth_401_reauth.py +++ b/tests/hermes_cli/test_dashboard_auth_401_reauth.py @@ -33,7 +33,7 @@ from hermes_cli import web_server from hermes_cli.dashboard_auth import clear_providers, register_provider -from hermes_cli.dashboard_auth.base import ProviderError, RefreshExpiredError +from hermes_cli.dashboard_auth.base import ProviderError, RefreshExpiredError, Session from hermes_cli.dashboard_auth.cookies import ( SESSION_AT_COOKIE, SESSION_PROVIDER_COOKIE, @@ -254,6 +254,55 @@ def test_at_evicted_rt_present_refreshes_transparently(self, gated_app): for c in set_cookies ), f"no rotated RT cookie in {set_cookies!r}" + def test_stale_parallel_request_reuses_the_just_rotated_session(self, gated_app): + """Two requests can leave the browser with the same rotating RT. + + The first response cannot update the second in-flight request's Cookie + header. Replaying that old RT against a reuse-detecting provider must + reuse the first rotation result instead of calling the provider again + and clearing an otherwise healthy browser session. + """ + + class ReuseDetectingProvider(StubAuthProvider): + def __init__(self): + super().__init__(default_ttl=900) + self.refresh_calls = 0 + + def refresh_session(self, *, refresh_token: str): + self.refresh_calls += 1 + if self.refresh_calls > 1: + raise RefreshExpiredError("rotated refresh token replayed") + + return Session( + user_id="stub-user-1", + email="stub@example.test", + display_name="Stub User", + org_id="stub-org-1", + provider=self.name, + expires_at=1_900_000_000, + access_token="rotated-access-token", + refresh_token="rotated-refresh-token", + ) + + provider = ReuseDetectingProvider() + clear_providers() + register_provider(provider) + first = TestClient(web_server.app, base_url="https://fly-app.fly.dev") + stale_parallel = TestClient(web_server.app, base_url="https://fly-app.fly.dev") + first.cookies.set(SESSION_RT_COOKIE, "same-pre-rotation-refresh-token") + stale_parallel.cookies.set(SESSION_RT_COOKIE, "same-pre-rotation-refresh-token") + + first_response = first.get("/api/sessions", follow_redirects=False) + stale_response = stale_parallel.get("/api/sessions", follow_redirects=False) + + assert first_response.status_code == 200 + assert stale_response.status_code == 200 + assert provider.refresh_calls == 1 + assert any( + SESSION_RT_COOKIE in cookie and "rotated-refresh-token" in cookie + for cookie in stale_response.headers.get_list("set-cookie") + ) + def test_provider_hint_routes_refresh_to_token_owner(self, gated_app): """A Nous-style RT must not be rejected by Basic just because Basic was registered first. The non-secret provider hint routes directly to