From 08a1f262aabe9beff3ec1326858ef44a5bda1a67 Mon Sep 17 00:00:00 2001 From: MOLW <85048713+MatrixAdventurer@users.noreply.github.com> Date: Sun, 25 Oct 2026 17:13:54 +0800 Subject: [PATCH 1/7] fix(sdk): coordinate concurrent token refresh MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 由 LW 授权的 ChatGPT Personal 智能体提交。 --- packages/sdk/README.md | 6 + .../token-refresh-concurrency.test.ts | 129 ++++++++++++++++++ packages/sdk/src/client.ts | 48 +++++-- 3 files changed, 172 insertions(+), 11 deletions(-) create mode 100644 packages/sdk/src/__tests__/token-refresh-concurrency.test.ts diff --git a/packages/sdk/README.md b/packages/sdk/README.md index 16be074fbc..0ea837b01b 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -161,6 +161,12 @@ already accepted frame to other subscribed callbacks. ### Browser apps and the BFF pattern +Within one `AstraClient`, concurrent REST and buffered run-event requests share +one token refresh. A delayed 401 reuses the rotated token. Refresh responses +started before `setTokens`, successful login/registration, or logout cannot +restore the previous credentials or retry an old request under a new account. +This coordination is per client instance, not across separate BFF requests. + Do not expose long-lived **refresh tokens** in browser-accessible JavaScript. The dashboard (`web/`) keeps tokens in **httpOnly cookies** and routes browser requests through explicit **Next.js Route Handlers** such as `/api/chats`, diff --git a/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts new file mode 100644 index 0000000000..a34eef6e19 --- /dev/null +++ b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts @@ -0,0 +1,129 @@ +import { AstraClient } from '../client'; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} + +const rotated = { access_token: 'new-access', refresh_token: 'new-refresh' }; +const unauthorized = () => Response.json({ detail: 'expired' }, { status: 401 }); +const authorization = (init?: RequestInit) => new Headers(init?.headers).get('Authorization'); + +afterEach(() => vi.unstubAllGlobals()); + +test('REST and buffered SSE share one refresh of a single-use token', async () => { + const refresh = deferred(); + const onTokenRefresh = vi.fn(); + let refreshCount = 0; + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshCount++; + expect(JSON.parse(String(init?.body))).toEqual({ refresh_token: 'old-refresh' }); + // The runtime revokes the token on its first successful refresh. + return refreshCount === 1 ? refresh.promise : unauthorized(); + } + if (authorization(init) === 'Bearer old-access') return unauthorized(); + expect(authorization(init)).toBe('Bearer new-access'); + return url.includes('/stream') + ? new Response('data: {"type":"run_finished"}\n\n') + : Response.json({ user_id: 'alice' }); + }); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = Promise.allSettled([client.getMe(), client.getRunEvents('run-1')]); + await vi.waitFor(() => expect(refreshCount).toBeGreaterThan(0)); + refresh.resolve(Response.json(rotated)); + expect((await results).map((result) => result.status)).toEqual(['fulfilled', 'fulfilled']); + expect(refreshCount).toBe(1); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); +}); + +test('a late 401 reuses the token already rotated by another request', async () => { + const late = deferred(); + let initialRequests = 0; + let refreshCount = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshCount++; + return Response.json(rotated); + } + if (authorization(init) === 'Bearer old-access') { + return ++initialRequests === 1 ? late.promise : unauthorized(); + } + return Response.json({ user_id: 'alice' }); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + }); + const first = client.getMe(); + await client.getMe(); + late.resolve(unauthorized()); + await expect(first).resolves.toEqual({ user_id: 'alice' }); + expect(refreshCount).toBe(1); +}); + +test.each(['http', 'network'] as const)('failed %s refresh is shared but a later attempt can recover', async (failure) => { + const refresh = deferred(); + let refreshCount = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshCount++; + if (refreshCount > 1) return Response.json(rotated); + await refresh.promise; + if (failure === 'network') throw new TypeError('connection lost'); + return unauthorized(); + } + return authorization(init) === 'Bearer old-access' + ? unauthorized() : Response.json({ user_id: 'alice' }); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + }); + const results = Promise.allSettled([client.getMe(), client.getMe()]); + await vi.waitFor(() => expect(refreshCount).toBeGreaterThan(0)); + refresh.resolve(unauthorized()); + expect((await results).map((result) => result.status)).toEqual(['rejected', 'rejected']); + expect(refreshCount).toBe(1); + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); + expect(refreshCount).toBe(2); +}); + +test.each(['setTokens', 'login', 'register', 'logout'] as const)( + 'a refresh started before %s cannot restore the previous credentials', async (action) => { + const refresh = deferred(); + const onTokenRefresh = vi.fn(); + let refreshCount = 0; + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { refreshCount++; return refresh.promise; } + if (url.endsWith('/auth/logout')) return new Response(null, { status: 204 }); + if (url.endsWith('/auth/login') || url.endsWith('/auth/register')) { + return Response.json({ access_token: 'bob-access', refresh_token: 'bob-refresh' }); + } + return authorization(init) === 'Bearer old-access' + ? unauthorized() : Response.json({ user_id: 'bob' }); + }); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const pending = Promise.allSettled([client.getMe()]); + await vi.waitFor(() => expect(refreshCount).toBe(1)); + if (action === 'setTokens') client.setTokens('bob-access', 'bob-refresh'); + else if (action === 'logout') await client.logout(); + else await client[action]('bob', 'password'); + const callsBeforeCompletion = fetchMock.mock.calls.length; + refresh.resolve(Response.json(rotated)); + expect((await pending)[0].status).toBe('rejected'); + expect(fetchMock).toHaveBeenCalledTimes(callsBeforeCompletion); + expect(onTokenRefresh).not.toHaveBeenCalled(); + await client.getMe(); + expect(authorization(fetchMock.mock.lastCall?.[1])).toBe( + action === 'logout' ? null : 'Bearer bob-access', + ); + }, +); diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index 02b78c224f..1ed937010f 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -627,6 +627,8 @@ export class AstraClient { private config: AstraClientConfig; private accessToken: string | null; private refreshTokenValue: string | null; + private authGeneration = 0; + private refreshAttempt: { generation: number; promise: Promise } | null = null; constructor(config: AstraClientConfig) { this.config = config; @@ -654,8 +656,7 @@ export class AstraClient { const body: Record = { username, email, password }; if (options?.displayName) body.display_name = options.displayName; const result = await this.post(PATH_AUTH_REGISTER, body); - this.accessToken = result.access_token; - this.refreshTokenValue = result.refresh_token; + this.setTokens(result.access_token, result.refresh_token); return result; } @@ -665,13 +666,13 @@ export class AstraClient { username, password, }); - this.accessToken = result.access_token; - this.refreshTokenValue = result.refresh_token; + this.setTokens(result.access_token, result.refresh_token); return result; } /** Log out and clear stored tokens. Requires refresh token in client state. */ async logout(): Promise { + this.authGeneration++; try { if (this.refreshTokenValue) { await this.post(PATH_AUTH_LOGOUT, { @@ -679,6 +680,7 @@ export class AstraClient { }); } } finally { + this.authGeneration++; this.accessToken = null; this.refreshTokenValue = null; } @@ -743,27 +745,47 @@ export class AstraClient { } setTokens(accessToken: string, refreshToken?: string): void { + this.authGeneration++; this.accessToken = accessToken; if (refreshToken) this.refreshTokenValue = refreshToken; } - private async tryRefreshToken(): Promise { + private async tryRefreshToken(generation: number, rejectedToken: string | null): Promise { + if (generation !== this.authGeneration) return false; + if (this.refreshAttempt?.generation === generation) return this.refreshAttempt.promise; + // A delayed 401 may belong to a token another request already rotated. + if (rejectedToken !== this.accessToken) return true; if (!this.refreshTokenValue) return false; + const attempt = { + generation, + promise: this.refreshTokens(generation, this.refreshTokenValue), + }; + this.refreshAttempt = attempt; + try { + return await attempt.promise; + } finally { + // An older account's completion must not clear a newer refresh attempt. + if (this.refreshAttempt === attempt) this.refreshAttempt = null; + } + } + + private async refreshTokens(generation: number, refreshToken: string): Promise { try { const res = await fetch(this.apiPath(PATH_AUTH_REFRESH), { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ refresh_token: this.refreshTokenValue }), + body: JSON.stringify({ refresh_token: refreshToken }), }); if (!res.ok) return false; const data = (await res.json()) as AuthResult; + if (generation !== this.authGeneration) return false; this.accessToken = data.access_token; this.refreshTokenValue = data.refresh_token; await this.config.onTokenRefresh?.({ accessToken: data.access_token, refreshToken: data.refresh_token, }); - return true; + return generation === this.authGeneration; } catch { return false; } @@ -795,14 +817,16 @@ export class AstraClient { private async request(path: string, init?: RequestInit): Promise { const url = this.apiPath(path); + const generation = this.authGeneration; + const accessToken = this.accessToken; let res = await fetch(url, { ...init, headers: headersInitToRecord(this.buildHeaders(init), init?.headers), }); if (res.status === 401) { - const refreshed = await this.tryRefreshToken(); - if (refreshed) { + const refreshed = await this.tryRefreshToken(generation, accessToken); + if (refreshed && generation === this.authGeneration) { res = await fetch(url, { ...init, headers: headersInitToRecord(this.buildHeaders(init), init?.headers), @@ -3061,14 +3085,16 @@ export class AstraClient { async getRunEvents(runId: string, startIndex = 0): Promise { const path = `${chatRunStreamPath(runId)}?last_index=${encodeURIComponent(String(startIndex))}`; const url = this.apiPath(path); + const generation = this.authGeneration; + const accessToken = this.accessToken; const streamHeaders = this.buildHeaders(undefined, { Accept: "text/event-stream", }); let res = await fetch(url, { headers: streamHeaders }); if (res.status === 401) { - const refreshed = await this.tryRefreshToken(); - if (refreshed) { + const refreshed = await this.tryRefreshToken(generation, accessToken); + if (refreshed && generation === this.authGeneration) { res = await fetch(url, { headers: this.buildHeaders(undefined, { Accept: "text/event-stream", From 5261c1545723cdf80e1ef5ac63d37dba60ef138e Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sun, 25 Oct 2026 21:20:14 +0800 Subject: [PATCH 2/7] fix(sdk): allow token refresh callbacks to reenter client --- packages/sdk/README.md | 8 + packages/sdk/src/__tests__/client.test.ts | 23 ++- .../token-refresh-concurrency.test.ts | 158 ++++++++++++++++++ packages/sdk/src/client.ts | 27 +-- .../lib/runtime-client-refresh.test.ts | 50 ++++++ 5 files changed, 241 insertions(+), 25 deletions(-) create mode 100644 web/__tests__/lib/runtime-client-refresh.test.ts diff --git a/packages/sdk/README.md b/packages/sdk/README.md index 0ea837b01b..faade2bc7a 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -167,6 +167,14 @@ started before `setTokens`, successful login/registration, or logout cannot restore the previous credentials or retry an old request under a new account. This coordination is per client instance, not across separate BFF requests. +`onTokenRefresh` may be asynchronous and call the same client. Requests already +awaiting its rotation still wait for the callback. Callback failure preserves +the rotated credentials and returns those callers' original authentication +error. Notifications from successive rotations may overlap. + +Buffered `getRunEvents` uses the same authenticated request and structured +`AstraApiError` decoding as REST, including error codes and action hints. + Do not expose long-lived **refresh tokens** in browser-accessible JavaScript. The dashboard (`web/`) keeps tokens in **httpOnly cookies** and routes browser requests through explicit **Next.js Route Handlers** such as `/api/chats`, diff --git a/packages/sdk/src/__tests__/client.test.ts b/packages/sdk/src/__tests__/client.test.ts index 944c9c70dd..ebcaa4f1c7 100644 --- a/packages/sdk/src/__tests__/client.test.ts +++ b/packages/sdk/src/__tests__/client.test.ts @@ -428,11 +428,30 @@ describe("AstraClient — Runs", () => { 'data: {"type":"turn_complete"}\n\n'; globalThis.fetch = mockFetch(200, sse); - const events = await createClient().getRunEvents("r1", 5); + const client = new AstraClient({ + baseUrl: "http://localhost:8000", pathPrefix: "/api", accessToken: "current-token", + headers: { accept: "application/json", AUTHORIZATION: "Bearer stale-token", "X-Custom": "preserved" }, + }); + const events = await client.getRunEvents("r1", 5); expect(events).toHaveLength(2); expect( (globalThis.fetch as ReturnType).mock.calls[0][0], - ).toContain("last_index=5"); + ).toBe("http://localhost:8000/api/chat/runs/r1/stream?last_index=5"); + const headers = new Headers((globalThis.fetch as ReturnType).mock.calls[0][1].headers); + expect(headers.get("Accept")).toBe("text/event-stream"); + expect(headers.get("Authorization")).toBe("Bearer current-token"); + expect(headers.get("X-Custom")).toBe("preserved"); + }); + + test("getRunEvents preserves structured API errors", async () => { + globalThis.fetch = mockFetch(429, { + detail: "Try later", code: "rate_limited", category: "capacity", + retryable: true, action_hints: ["retry_later"], + }, { "Content-Type": "application/json" }); + await expect(createClient().getRunEvents("r1")).rejects.toMatchObject({ + status: 429, body: "Try later", path: "/chat/runs/r1/stream?last_index=0", + code: "rate_limited", category: "capacity", retryable: true, actionHints: ["retry_later"], + }); }); test("getSessionConversation preserves typed input identity and raw cursors", async () => { diff --git a/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts index a34eef6e19..d65febe23e 100644 --- a/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts +++ b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts @@ -127,3 +127,161 @@ test.each(['setTokens', 'login', 'register', 'logout'] as const)( ); }, ); + +test('an async notification can receive 401 without blocking its callers or later recovery', async () => { + const refreshTokens: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshTokens.push(JSON.parse(String(init?.body)).refresh_token); + if (refreshTokens.length === 2) return Response.json({}, { status: 503 }); + return Response.json(refreshTokens.length === 1 ? rotated + : { access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + return authorization(init) === 'Bearer recovered-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const onTokenRefresh = vi.fn(async () => { await client.getMe(); }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = await Promise.allSettled([client.getMe(), client.getRunEvents('run-1')]); + expect(results.map(result => result.status)).toEqual(['rejected', 'rejected']); + expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); + expect(refreshTokens).toEqual(['old-refresh', 'new-refresh', 'new-refresh']); + expect(onTokenRefresh).toHaveBeenCalledTimes(2); +}); + +test('a successful nested refresh still permits only one replay of the buffered SSE request', async () => { + const refreshTokens: string[] = []; + const callbackResults: unknown[] = []; + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshTokens.push(JSON.parse(String(init?.body)).refresh_token); + return Response.json(refreshTokens.length === 1 ? rotated + : { access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + return url.endsWith('/auth/me') && authorization(init) === 'Bearer recovered-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + }); + vi.stubGlobal('fetch', fetchMock); + const onTokenRefresh = vi.fn(async () => { callbackResults.push(await client.getMe()); }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + await expect(client.getRunEvents('run-1')).rejects.toMatchObject({ status: 401 }); + expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); + expect(callbackResults).toEqual([{ user_id: 'alice' }, { user_id: 'alice' }]); + expect(onTokenRefresh).toHaveBeenCalledTimes(2); + const streamRequests = fetchMock.mock.calls.filter(([url]) => url.includes('/stream')); + expect(streamRequests.map(([, init]) => authorization(init))).toEqual([ + 'Bearer old-access', 'Bearer recovered-access', + ]); + expect(streamRequests.every(([, init]) => new Headers(init?.headers).get('Accept') === 'text/event-stream')).toBe(true); +}); + +test('already joined callers await an async notification that successfully reenters the client', async () => { + const notification = deferred(); + const entered = deferred(); + const settled = vi.fn(); + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(rotated); + if (authorization(init) === 'Bearer old-access') return unauthorized(); + return url.includes('/stream') ? new Response('data: {"type":"run_finished"}\n\n') + : Response.json({ user_id: 'alice' }); + })); + const onTokenRefresh = vi.fn(async () => { + await client.getMe(); + entered.resolve(); + await notification.promise; + }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = Promise.allSettled([client.getMe(), client.getRunEvents('run-1')] + .map(request => request.then(settled, settled))); + await entered.promise; + expect(settled).not.toHaveBeenCalled(); + notification.resolve(); + await results; + expect(settled).toHaveBeenCalledTimes(2); + expect(settled).toHaveBeenCalledWith({ user_id: 'alice' }); + expect(settled).toHaveBeenCalledWith([{ type: 'run_finished' }]); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); +}); + +test.each(['throw', 'reject'] as const)('a notification %s settles joined callers and preserves rotated credentials', async failure => { + const refresh = deferred(); + const refreshTokens: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + refreshTokens.push(JSON.parse(String(init?.body)).refresh_token); + return refreshTokens.length === 1 ? refresh.promise + : Response.json({ access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + return authorization(init) === 'Bearer recovered-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const onTokenRefresh = vi.fn(() => { + if (onTokenRefresh.mock.calls.length > 1) return; + const error = new Error('notification failed'); + if (failure === 'throw') throw error; + return Promise.reject(error); + }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const results = Promise.allSettled([client.getMe(), client.getRunEvents('run-1')]); + await vi.waitFor(() => expect(refreshTokens).toHaveLength(1)); + refresh.resolve(Response.json(rotated)); + expect((await results).map(result => result.status)).toEqual(['rejected', 'rejected']); + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); + expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); + expect(onTokenRefresh).toHaveBeenCalledTimes(2); +}); + +test('an older notification cannot replay across accounts or release a newer refresh', async () => { + const notification = deferred(); + const notified = deferred(); + const bobRefresh = deferred(); + const bobRefreshing = deferred(); + const refreshTokens: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + const token = JSON.parse(String(init?.body)).refresh_token; + refreshTokens.push(token); + if (token === 'old-refresh') return Response.json(rotated); + bobRefreshing.resolve(); + return bobRefresh.promise; + } + return authorization(init) === 'Bearer bob-access' + ? Response.json({ user_id: 'bob' }) : unauthorized(); + })); + const onTokenRefresh = vi.fn(async ({ accessToken }: { accessToken: string }) => { + if (accessToken !== 'new-access') return; + notified.resolve(); + await notification.promise; + }); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', + refreshToken: 'old-refresh', onTokenRefresh, + }); + const alice = Promise.allSettled([client.getMe()]); + await notified.promise; + client.setTokens('bob-expired', 'bob-refresh'); + const firstBob = client.getMe(); + await bobRefreshing.promise; + notification.resolve(); + expect((await alice)[0].status).toBe('rejected'); + const secondBob = client.getMe(); + await new Promise(resolve => setImmediate(resolve)); + expect(refreshTokens).toEqual(['old-refresh', 'bob-refresh']); + bobRefresh.resolve(Response.json({ access_token: 'bob-access', refresh_token: 'bob-new-refresh' })); + expect(await Promise.all([firstBob, secondBob])).toEqual([{ user_id: 'bob' }, { user_id: 'bob' }]); + expect(onTokenRefresh).toHaveBeenCalledTimes(2); +}); diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index 1ed937010f..b8f14c2c14 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -242,7 +242,6 @@ import { SSEClient, parseSseDataEvents } from "./sse-client"; import { headersInitToRecord, readAstraError, - readAstraErrorDetail, } from "./http"; import { modelSelectionToWire } from "./wire"; import { @@ -781,6 +780,8 @@ export class AstraClient { if (generation !== this.authGeneration) return false; this.accessToken = data.access_token; this.refreshTokenValue = data.refresh_token; + // Notification may reenter the client and need another refresh. + this.refreshAttempt = null; await this.config.onTokenRefresh?.({ accessToken: data.access_token, refreshToken: data.refresh_token, @@ -3084,29 +3085,9 @@ export class AstraClient { */ async getRunEvents(runId: string, startIndex = 0): Promise { const path = `${chatRunStreamPath(runId)}?last_index=${encodeURIComponent(String(startIndex))}`; - const url = this.apiPath(path); - const generation = this.authGeneration; - const accessToken = this.accessToken; - const streamHeaders = this.buildHeaders(undefined, { - Accept: "text/event-stream", + const res = await this.request(path, { + headers: { Accept: "text/event-stream" }, }); - let res = await fetch(url, { headers: streamHeaders }); - - if (res.status === 401) { - const refreshed = await this.tryRefreshToken(generation, accessToken); - if (refreshed && generation === this.authGeneration) { - res = await fetch(url, { - headers: this.buildHeaders(undefined, { - Accept: "text/event-stream", - }), - }); - } - } - - if (!res.ok) { - const body = await readAstraErrorDetail(res); - throw new AstraApiError(res.status, body, path); - } // Response.text() strips a BOM before the shared parser can apply its rule. const text = new TextDecoder("utf-8", { ignoreBOM: true }).decode(await res.arrayBuffer()); diff --git a/web/__tests__/lib/runtime-client-refresh.test.ts b/web/__tests__/lib/runtime-client-refresh.test.ts new file mode 100644 index 0000000000..1b500e7819 --- /dev/null +++ b/web/__tests__/lib/runtime-client-refresh.test.ts @@ -0,0 +1,50 @@ +// @vitest-environment node +vi.mock("next/headers", () => ({ cookies: vi.fn() })); + +import { cookies } from "next/headers"; +import { WebRuntimeClient } from "@/lib/runtime-client/server"; +import { ACCESS_TOKEN_COOKIE, REFRESH_TOKEN_COOKIE } from "@/lib/runtime-config"; + +afterEach(() => vi.unstubAllGlobals()); + +it("persists refreshed httpOnly cookies before the initiating SDK request completes", async () => { + let release!: () => void; + const available = new Promise(resolve => { release = resolve; }); + const set = vi.fn(); + vi.mocked(cookies).mockImplementation(async () => { + await available; + // This consumer only uses the cookie store's setter. + return { set } as unknown as Awaited>; + }); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + return Response.json({ access_token: "new-access", refresh_token: "new-refresh" }); + } + return new Headers(init?.headers).get("Authorization") === "Bearer new-access" + ? Response.json({ user_id: "alice" }) : Response.json({}, { status: 401 }); + }); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient({ + mode: "live", source: "cookie", apiUrl: "https://example.test", + accessToken: "old-access", refreshToken: "old-refresh", + demoMode: false, hasAccessToken: true, hasRefreshToken: true, message: "", + }); + const settled = vi.fn(); + const pending = runtime.sdk.getMe(); + void pending.then(settled, settled); + await vi.waitFor(() => expect(cookies).toHaveBeenCalledTimes(1)); + expect(settled).not.toHaveBeenCalled(); + expect(set).not.toHaveBeenCalled(); + await expect(runtime.get("/auth/me", { auth: "required" })).resolves.toEqual({ user_id: "alice" }); + expect(settled).not.toHaveBeenCalled(); + release(); + await expect(pending).resolves.toEqual({ user_id: "alice" }); + const options = { httpOnly: true, sameSite: "lax", path: "/" }; + expect(set.mock.calls).toEqual([ + [ACCESS_TOKEN_COOKIE, "new-access", options], + [REFRESH_TOKEN_COOKIE, "new-refresh", options], + ]); + expect(fetchMock.mock.calls.map(([, init]) => new Headers(init?.headers).get("Authorization"))).toEqual([ + "Bearer old-access", null, "Bearer new-access", "Bearer new-access", + ]); +}); From 2df60a07550732f5c2e0a904b9fcd8fee311eeb1 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sun, 25 Oct 2026 22:01:21 +0800 Subject: [PATCH 3/7] fix(web): discard stale refresh cookie writes --- packages/sdk/README.md | 2 + .../lib/runtime-client-refresh.test.ts | 78 +++++++++++++++++-- web/lib/runtime-client/server.ts | 4 + 3 files changed, 78 insertions(+), 6 deletions(-) diff --git a/packages/sdk/README.md b/packages/sdk/README.md index faade2bc7a..e8dd610508 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -171,6 +171,8 @@ This coordination is per client instance, not across separate BFF requests. awaiting its rotation still wait for the callback. Callback failure preserves the rotated credentials and returns those callers' original authentication error. Notifications from successive rotations may overlap. +The Web runtime client persists a captured pair only if it still matches that +instance's current credentials after asynchronous cookie acquisition. Buffered `getRunEvents` uses the same authenticated request and structured `AstraApiError` decoding as REST, including error codes and action hints. diff --git a/web/__tests__/lib/runtime-client-refresh.test.ts b/web/__tests__/lib/runtime-client-refresh.test.ts index 1b500e7819..e8547a0eab 100644 --- a/web/__tests__/lib/runtime-client-refresh.test.ts +++ b/web/__tests__/lib/runtime-client-refresh.test.ts @@ -3,10 +3,19 @@ vi.mock("next/headers", () => ({ cookies: vi.fn() })); import { cookies } from "next/headers"; import { WebRuntimeClient } from "@/lib/runtime-client/server"; -import { ACCESS_TOKEN_COOKIE, REFRESH_TOKEN_COOKIE } from "@/lib/runtime-config"; +import { ACCESS_TOKEN_COOKIE, REFRESH_TOKEN_COOKIE, type RuntimeConfig } from "@/lib/runtime-config"; +beforeEach(() => vi.mocked(cookies).mockReset()); afterEach(() => vi.unstubAllGlobals()); +function runtimeConfig(accessToken: string, refreshToken: string): RuntimeConfig { + return { + mode: "live", source: "cookie", apiUrl: "https://example.test", + accessToken, refreshToken, demoMode: false, + hasAccessToken: true, hasRefreshToken: true, message: "", + }; +} + it("persists refreshed httpOnly cookies before the initiating SDK request completes", async () => { let release!: () => void; const available = new Promise(resolve => { release = resolve; }); @@ -24,11 +33,7 @@ it("persists refreshed httpOnly cookies before the initiating SDK request comple ? Response.json({ user_id: "alice" }) : Response.json({}, { status: 401 }); }); vi.stubGlobal("fetch", fetchMock); - const runtime = new WebRuntimeClient({ - mode: "live", source: "cookie", apiUrl: "https://example.test", - accessToken: "old-access", refreshToken: "old-refresh", - demoMode: false, hasAccessToken: true, hasRefreshToken: true, message: "", - }); + const runtime = new WebRuntimeClient(runtimeConfig("old-access", "old-refresh")); const settled = vi.fn(); const pending = runtime.sdk.getMe(); void pending.then(settled, settled); @@ -48,3 +53,64 @@ it("persists refreshed httpOnly cookies before the initiating SDK request comple "Bearer old-access", null, "Bearer new-access", "Bearer new-access", ]); }); + +it.each([ + { name: "SDK notification", sdkRefresh: true }, + { name: "Web transport", sdkRefresh: false }, +])("keeps the newest usable cookies when the $name refresh finishes before an older notification", async ({ sdkRefresh }) => { + let release!: () => void; + const available = new Promise(resolve => { release = resolve; }); + const jar = new Map(); + const set = vi.fn((name: string, value: string) => { jar.set(name, value); }); + vi.mocked(cookies).mockImplementation(async () => { + if (vi.mocked(cookies).mock.calls.length === 1) await available; + return { set } as unknown as Awaited>; + }); + + const issued = new Set(["access-0"]); + const expired = new Set(["access-0"]); + const refreshRequests: string[] = []; + let rotation = 0; + let validRefresh = "refresh-0"; + const unauthorized = () => Response.json({ detail: "Unauthorized" }, { status: 401 }); + vi.stubGlobal("fetch", vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + const token = JSON.parse(String(init?.body)).refresh_token; + refreshRequests.push(token); + if (token !== validRefresh) return unauthorized(); + rotation++; + validRefresh = `refresh-${rotation}`; + const access = `access-${rotation}`; + issued.add(access); + return Response.json({ access_token: access, refresh_token: validRefresh }); + } + const token = new Headers(init?.headers).get("Authorization")?.replace("Bearer ", ""); + return token && issued.has(token) && !expired.has(token) + ? Response.json({ user_id: "alice" }) : unauthorized(); + })); + + const runtime = new WebRuntimeClient(runtimeConfig("access-0", "refresh-0")); + const first = Promise.allSettled([runtime.sdk.getMe()]); + await vi.waitFor(() => expect(cookies).toHaveBeenCalledTimes(1)); + expired.add("access-1"); + if (sdkRefresh) { + await expect(runtime.sdk.getMe()).resolves.toEqual({ user_id: "alice" }); + } else { + await expect(runtime.get("/auth/me", { auth: "required" })).resolves.toEqual({ user_id: "alice" }); + } + release(); + // The Web transport's existing setTokens handoff fences the older SDK call. + expect(await first).toMatchObject([sdkRefresh + ? { status: "fulfilled", value: { user_id: "alice" } } + : { status: "rejected", reason: { status: 401 } }]); + + const committed = set.mock.calls.map(([name, value]) => [name, value]); + const next = new WebRuntimeClient(runtimeConfig(jar.get(ACCESS_TOKEN_COOKIE)!, jar.get(REFRESH_TOKEN_COOKIE)!)); + await expect(next.sdk.getMe()).resolves.toEqual({ user_id: "alice" }); + expect(committed).toEqual([ + [ACCESS_TOKEN_COOKIE, "access-2"], + [REFRESH_TOKEN_COOKIE, "refresh-2"], + ]); + expect(jar.get(REFRESH_TOKEN_COOKIE)).toBe(validRefresh); + expect(refreshRequests).toEqual(["refresh-0", "refresh-1"]); +}); diff --git a/web/lib/runtime-client/server.ts b/web/lib/runtime-client/server.ts index 7a0923d9f8..1023300f25 100644 --- a/web/lib/runtime-client/server.ts +++ b/web/lib/runtime-client/server.ts @@ -243,6 +243,10 @@ export class WebRuntimeClient { ): Promise { const { cookies } = await import("next/headers"); const cookieStore = await cookies(); + // A newer refresh may have replaced this pair while cookie acquisition waited. + if (this.accessToken !== accessToken || this.refreshToken !== refreshToken) { + return; + } cookieStore.set(ACCESS_TOKEN_COOKIE, accessToken, { httpOnly: true, sameSite: "lax", From c49bf9d480564a3aa8b8230fc8a229767ee6583e Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Sun, 25 Oct 2026 22:58:38 +0800 Subject: [PATCH 4/7] fix(web): share SDK automatic refresh owner --- packages/sdk/README.md | 13 +++ .../sdk/src/__tests__/raw-response.test.ts | 44 +++++++++ packages/sdk/src/client.ts | 71 +++++++------- .../lib/runtime-client-refresh.test.ts | 94 ++++++++++++++++++- .../lib/web-store-model-cache.test.ts | 12 ++- web/lib/api/web-store.ts | 2 +- web/lib/runtime-client/server.ts | 82 +++------------- 7 files changed, 209 insertions(+), 109 deletions(-) create mode 100644 packages/sdk/src/__tests__/raw-response.test.ts diff --git a/packages/sdk/README.md b/packages/sdk/README.md index e8dd610508..b8cb668033 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -174,6 +174,19 @@ error. Notifications from successive rotations may overlap. The Web runtime client persists a captured pair only if it still matches that instance's current credentials after asynchronous cookie acquisition. +The Web runtime's raw HTTP requests share this SDK refresh owner. Its API URL +is origin-relative: a pathname suffix is not a gateway prefix. Standalone SDK +`baseUrl` and `pathPrefix` routing are unchanged. Web refresh or notification +failure returns the original 401, which its typed adapter converts to a runtime +error; initial transport errors still propagate. + +`fetchResponse(path, init, { auth })` returns an unconsumed response, including +non-2xx responses, without inferring a JSON content type. Its default auth is +`"optional"`; `"none"` suppresses default authorization and automatic refresh. +Explicit request headers still override defaults. `getTokens()` returns a +detached snapshot of current credentials, not the initial configuration. +Bodies must remain replayable for the single automatic retry. + Buffered `getRunEvents` uses the same authenticated request and structured `AstraApiError` decoding as REST, including error codes and action hints. diff --git a/packages/sdk/src/__tests__/raw-response.test.ts b/packages/sdk/src/__tests__/raw-response.test.ts new file mode 100644 index 0000000000..231d74170f --- /dev/null +++ b/packages/sdk/src/__tests__/raw-response.test.ts @@ -0,0 +1,44 @@ +import { AstraClient } from '../client'; + +afterEach(() => vi.unstubAllGlobals()); + +test('raw transport leaves error bodies and FormData untouched; typed fetch keeps JSON defaults', async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response('upstream detail', { status: 403 })); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ baseUrl: 'https://example.test/base', pathPrefix: '/api' }); + const body = new FormData(); + body.set('file', new Blob(['contents']), 'test.txt'); + const signal = new AbortController().signal; + const response = await client.fetchResponse('/upload', { method: 'POST', body, signal }); + expect(response.status).toBe(403); + expect(response.bodyUsed).toBe(false); + expect(await response.text()).toBe('upstream detail'); + expect(fetchMock.mock.calls[0]).toEqual(['https://example.test/base/api/upload', expect.objectContaining({ body, signal })]); + expect(new Headers(fetchMock.mock.calls[0][1]?.headers).has('Content-Type')).toBe(false); + await expect(client.fetch('/typed', { method: 'POST', body: '{}' })).rejects.toMatchObject({ status: 403 }); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get('Content-Type')).toBe('application/json'); +}); + +test('anonymous raw requests suppress configured auth and refresh but allow explicit headers', async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response('unauthorized', { status: 401 })); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'access', refreshToken: 'refresh', headers: { authorization: 'Bearer configured' } }); + await client.fetchResponse('/public', {}, { auth: 'none' }); + await client.fetchResponse('/public', { headers: { Authorization: 'Bearer explicit' } }, { auth: 'none' }); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(new Headers(fetchMock.mock.calls[0][1]?.headers).has('Authorization')).toBe(false); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get('Authorization')).toBe('Bearer explicit'); +}); + +test.each([null, {}, { access_token: 'new', refresh_token: 1 }])('malformed refresh does not publish credentials or notify: %j', async payload => { + const notify = vi.fn(); + const fetchMock = vi.fn(async (url: string) => url.endsWith('/auth/refresh') + ? Response.json(payload) : new Response('expired', { status: 401 })); + vi.stubGlobal('fetch', fetchMock); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'refresh', onTokenRefresh: notify }); + const snapshot = client.getTokens(); + await expect(client.getMe()).rejects.toMatchObject({ status: 401 }); + expect(client.getTokens()).toEqual(snapshot); + expect(notify).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(2); +}); diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index b8f14c2c14..688f3ba76d 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -749,6 +749,11 @@ export class AstraClient { if (refreshToken) this.refreshTokenValue = refreshToken; } + /** A detached snapshot of the credentials currently owned by this client. */ + getTokens(): Readonly<{ accessToken: string | null; refreshToken: string | null }> { + return { accessToken: this.accessToken, refreshToken: this.refreshTokenValue }; + } + private async tryRefreshToken(generation: number, rejectedToken: string | null): Promise { if (generation !== this.authGeneration) return false; if (this.refreshAttempt?.generation === generation) return this.refreshAttempt.promise; @@ -774,10 +779,16 @@ export class AstraClient { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ refresh_token: refreshToken }), + cache: "no-store", }); if (!res.ok) return false; const data = (await res.json()) as AuthResult; - if (generation !== this.authGeneration) return false; + if ( + generation !== this.authGeneration || + !data || + typeof data.access_token !== "string" || + typeof data.refresh_token !== "string" + ) return false; this.accessToken = data.access_token; this.refreshTokenValue = data.refresh_token; // Notification may reenter the client and need another refresh. @@ -792,48 +803,44 @@ export class AstraClient { } } - private buildHeaders( - init?: RequestInit, - extra?: Record, - ): Record { - const headers = headersInitToRecord(this.config.headers ?? {}, extra); - const generated: Record = {}; - const method = (init?.method ?? "GET").toUpperCase(); - if ( - init?.body != null && - (method === "POST" || - method === "PUT" || - method === "PATCH" || - method === "DELETE") - ) { - generated["Content-Type"] = "application/json"; - } - if (this.accessToken) { - generated["Authorization"] = `Bearer ${this.accessToken}`; - } - return headersInitToRecord(headers, generated); + private buildHeaders(auth: "optional" | "none"): Record { + const headers = new Headers(headersInitToRecord(this.config.headers ?? {})); + if (auth === "none") headers.delete("Authorization"); + else if (this.accessToken) headers.set("Authorization", `Bearer ${this.accessToken}`); + return headersInitToRecord({}, headers); } // ─── HTTP helpers ────────────────────────────────────────────────── - private async request(path: string, init?: RequestInit): Promise { + /** Raw, body-neutral transport; returns the final response without consuming it. */ + async fetchResponse( + path: string, + init?: RequestInit, + options?: { auth?: "optional" | "none" }, + ): Promise { + const auth = options?.auth ?? "optional"; const url = this.apiPath(path); const generation = this.authGeneration; const accessToken = this.accessToken; - let res = await fetch(url, { + const send = () => fetch(url, { ...init, - headers: headersInitToRecord(this.buildHeaders(init), init?.headers), + headers: headersInitToRecord(this.buildHeaders(auth), init?.headers), }); - - if (res.status === 401) { + let res = await send(); + if (res.status === 401 && auth !== "none") { const refreshed = await this.tryRefreshToken(generation, accessToken); - if (refreshed && generation === this.authGeneration) { - res = await fetch(url, { - ...init, - headers: headersInitToRecord(this.buildHeaders(init), init?.headers), - }); - } + if (refreshed && generation === this.authGeneration) res = await send(); } + return res; + } + + private async request(path: string, init?: RequestInit): Promise { + const method = (init?.method ?? "GET").toUpperCase(); + const json = init?.body != null && ["POST", "PUT", "PATCH", "DELETE"].includes(method); + const res = await this.fetchResponse(path, { + ...init, + headers: headersInitToRecord(json ? { "Content-Type": "application/json" } : {}, init?.headers), + }); if (!res.ok) { const error = await readAstraError(res); diff --git a/web/__tests__/lib/runtime-client-refresh.test.ts b/web/__tests__/lib/runtime-client-refresh.test.ts index e8547a0eab..2f60c909e6 100644 --- a/web/__tests__/lib/runtime-client-refresh.test.ts +++ b/web/__tests__/lib/runtime-client-refresh.test.ts @@ -99,10 +99,7 @@ it.each([ await expect(runtime.get("/auth/me", { auth: "required" })).resolves.toEqual({ user_id: "alice" }); } release(); - // The Web transport's existing setTokens handoff fences the older SDK call. - expect(await first).toMatchObject([sdkRefresh - ? { status: "fulfilled", value: { user_id: "alice" } } - : { status: "rejected", reason: { status: 401 } }]); + expect(await first).toMatchObject([{ status: "fulfilled", value: { user_id: "alice" } }]); const committed = set.mock.calls.map(([name, value]) => [name, value]); const next = new WebRuntimeClient(runtimeConfig(jar.get(ACCESS_TOKEN_COOKIE)!, jar.get(REFRESH_TOKEN_COOKIE)!)); @@ -114,3 +111,92 @@ it.each([ expect(jar.get(REFRESH_TOKEN_COOKIE)).toBe(validRefresh); expect(refreshRequests).toEqual(["refresh-0", "refresh-1"]); }); + +it("shares one in-flight refresh between SDK and Web requests", async () => { + let release!: () => void; + const refresh = new Promise(resolve => { release = resolve; }); + const set = vi.fn(); + vi.mocked(cookies).mockResolvedValue({ set } as unknown as Awaited>); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + await refresh; + return Response.json({ access_token: "new", refresh_token: "rotated" }); + } + return new Headers(init?.headers).get("Authorization") === "Bearer new" + ? Response.json({ user_id: "alice" }) : new Response("expired", { status: 401 }); + }); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig("old", "refresh")); + const pending = Promise.all([runtime.sdk.getMe(), runtime.get("/auth/me", { auth: "required" })]); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(3)); + release(); + expect(await pending).toEqual([{ user_id: "alice" }, { user_id: "alice" }]); + expect(fetchMock.mock.calls.filter(([url]) => url.endsWith("/auth/refresh"))).toHaveLength(1); + expect(cookies).toHaveBeenCalledTimes(1); + expect(set).toHaveBeenCalledTimes(2); +}); + +it("keeps Web raw responses, body, cancellation and origin routing intact", async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response("data: event\n\n", { status: 403 })); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient({ ...runtimeConfig("access", "refresh"), apiUrl: "https://example.test/ignored?query=ignored" }); + const body = new FormData(); + body.set("upload", "contents"); + const signal = new AbortController().signal; + const response = await runtime.fetchResponse("/stream", { method: "POST", body, signal, headers: { Accept: "text/event-stream" } }); + expect(response.bodyUsed).toBe(false); + expect(await response.text()).toBe("data: event\n\n"); + const [url, init] = fetchMock.mock.calls[0]; + expect(url).toBe("https://example.test/stream"); + expect(init).toMatchObject({ method: "POST", body, signal, cache: "no-store" }); + expect(new Headers(init?.headers).get("Accept")).toBe("text/event-stream"); + expect(new Headers(init?.headers).has("Content-Type")).toBe(false); + await expect(runtime.get("/error", { operation: "read item" })).rejects.toMatchObject({ status: 403, operation: "read item" }); + await expect(runtime.sdk.getMe()).rejects.toMatchObject({ status: 403 }); + expect(fetchMock.mock.calls[2][0]).toBe("https://example.test/auth/me"); +}); + +it("uses current SDK credentials for Web auth and keeps old requests fenced on identity replacement", async () => { + let release!: () => void; + const refresh = new Promise(resolve => { release = resolve; }); + const set = vi.fn(); + vi.mocked(cookies).mockResolvedValue({ set } as unknown as Awaited>); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith("/auth/refresh")) { + await refresh; + return Response.json({ access_token: "stale", refresh_token: "stale-refresh" }); + } + if (url.endsWith("/auth/logout")) return new Response(null, { status: 204 }); + return new Headers(init?.headers).get("Authorization") === "Bearer bob" + ? Response.json({ user_id: "bob" }) : new Response("expired", { status: 401 }); + }); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig("alice", "refresh")); + const old = Promise.allSettled([runtime.get("/auth/me", { auth: "required" })]); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + runtime.sdk.setTokens("bob", "bob-refresh"); + release(); + expect(await old).toMatchObject([{ status: "rejected", reason: { status: 401 } }]); + await expect(runtime.get("/auth/me", { auth: "required" })).resolves.toEqual({ user_id: "bob" }); + expect(set).not.toHaveBeenCalled(); + await runtime.sdk.logout(); + const count = fetchMock.mock.calls.length; + await expect(runtime.get("/auth/me", { auth: "required" })).rejects.toMatchObject({ status: 401 }); + expect(fetchMock).toHaveBeenCalledTimes(count); +}); + +it("does not authenticate or refresh anonymous Web calls and preserves explicit overrides", async () => { + const fetchMock = vi.fn(async (_url: string, _init?: RequestInit) => new Response("unauthorized", { status: 401 })); + vi.stubGlobal("fetch", fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig("access", "refresh")); + await runtime.fetchResponse("/auth/login", { auth: "none", method: "POST", json: { username: "alice" } }); + await runtime.fetchResponse("/auth/login", { auth: "none", headers: { Authorization: "Bearer explicit", "X-User-Id": "hint" } }); + expect(fetchMock).toHaveBeenCalledTimes(2); + const headers = new Headers(fetchMock.mock.calls[0][1]?.headers); + expect(headers.has("Authorization")).toBe(false); + expect(headers.has("X-User-Id")).toBe(false); + expect(headers.get("Content-Type")).toBe("application/json"); + expect(fetchMock.mock.calls[0][1]?.body).toBe('{"username":"alice"}'); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get("Authorization")).toBe("Bearer explicit"); + expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get("X-User-Id")).toBe("hint"); +}); diff --git a/web/__tests__/lib/web-store-model-cache.test.ts b/web/__tests__/lib/web-store-model-cache.test.ts index 32fed15bfc..906be0aa66 100644 --- a/web/__tests__/lib/web-store-model-cache.test.ts +++ b/web/__tests__/lib/web-store-model-cache.test.ts @@ -30,15 +30,16 @@ function makeRuntime(overrides: Partial<{ accessToken: string | null }> = {}) { tokenExpiresAtMs: null as number | null, }; + const getTokens = vi.fn((): { accessToken: string | null; refreshToken: string | null } => ({ accessToken: mockConfig.accessToken, refreshToken: mockConfig.refreshToken })); MockedWRC.mockImplementation(function () { return { config: mockConfig, - sdk: { listModels: mockListModels }, + sdk: { listModels: mockListModels, getTokens }, }; }); const client = new WebRuntimeClient(mockConfig as any); - return { client, mockListModels }; + return { client, mockListModels, getTokens }; } describe("resolveModelOfferingSelection", () => { @@ -82,6 +83,13 @@ describe("resolveModelOfferingSelection", () => { expect(mockListModels).not.toHaveBeenCalled(); }); + it("checks live credentials after logout rather than initialization config", async () => { + const { client, getTokens } = makeRuntime(); + getTokens.mockReturnValue({ accessToken: null, refreshToken: null }); + await expect(resolveModelOfferingSelection(client, "offering")) + .rejects.toMatchObject({ code: "authentication_required" }); + }); + it("rejects unauthenticated selection without querying the catalog", async () => { const { client, mockListModels } = makeRuntime({ accessToken: null }); diff --git a/web/lib/api/web-store.ts b/web/lib/api/web-store.ts index 2fcf9c8e91..f326381ade 100644 --- a/web/lib/api/web-store.ts +++ b/web/lib/api/web-store.ts @@ -2347,7 +2347,7 @@ export async function resolveModelOfferingSelection( runtime instanceof WebRuntimeClient ? runtime : new WebRuntimeClient(runtime); - const accessToken = client.config.accessToken; + const { accessToken } = client.sdk.getTokens(); if (!accessToken) { throw new ModelOfferingSelectionError( "authentication_required", diff --git a/web/lib/runtime-client/server.ts b/web/lib/runtime-client/server.ts index 1023300f25..02696cce34 100644 --- a/web/lib/runtime-client/server.ts +++ b/web/lib/runtime-client/server.ts @@ -1,9 +1,7 @@ import { AstraClient, - PATH_AUTH_REFRESH, extractJwtSubject, headersInitToRecord, - joinApiPath, methodCanHaveJson, } from "@astra/sdk"; import { @@ -34,8 +32,6 @@ type RuntimeClientOptions = { export class WebRuntimeClient { readonly config: RuntimeConfig; readonly sdk: AstraClient; - private accessToken?: string; - private refreshToken?: string; constructor(config: RuntimeConfig) { if (config.mode !== "live" || !config.apiUrl) { @@ -48,25 +44,16 @@ export class WebRuntimeClient { } this.config = config; - this.accessToken = config.accessToken; - this.refreshToken = config.refreshToken; this.sdk = new AstraClient({ - baseUrl: config.apiUrl, + baseUrl: new URL("/", config.apiUrl).toString(), accessToken: config.accessToken, refreshToken: config.refreshToken, - headers: this.baseAuthHeaders("optional"), onTokenRefresh: async (tokens) => { - this.accessToken = tokens.accessToken; - this.refreshToken = tokens.refreshToken; await this.persistTokens(tokens.accessToken, tokens.refreshToken); }, }); } - get apiUrl(): string { - return this.config.apiUrl ?? DEFAULT_API_URL; - } - async get(path: string, init?: RuntimeRequestInit): Promise { return this.request(path, { ...init, method: "GET" }); } @@ -130,24 +117,11 @@ export class WebRuntimeClient { ): Promise { const operation = init.operation ?? `${init.method ?? "GET"} ${path}`; const auth = init.auth ?? "optional"; - const url = this.url(path); - const request = this.toRequestInit(init, auth, operation, path); - - let response = await fetch(url, request); - if (response.status === 401 && auth !== "none" && this.refreshToken) { - const refreshed = await this.refreshTokens(); - if (refreshed) { - response = await fetch( - url, - this.toRequestInit(init, auth, operation, path), - ); - } - } - return response; - } - - url(path: string): string { - return new URL(joinApiPath(undefined, path), this.apiUrl).toString(); + return this.sdk.fetchResponse( + path, + this.toRequestInit(init, auth, operation, path), + { auth: auth === "none" ? "none" : "optional" }, + ); } private toRequestInit( @@ -167,7 +141,7 @@ export class WebRuntimeClient { init.headers, ); - if (auth === "required" && !this.accessToken) { + if (auth === "required" && !this.sdk.getTokens().accessToken) { throw new RuntimeClientError({ operation, path, @@ -194,9 +168,9 @@ export class WebRuntimeClient { if (request?.hasJsonBody && methodCanHaveJson(request.method)) { headers["Content-Type"] = "application/json"; } - if (auth !== "none" && this.accessToken) { - headers.Authorization = `Bearer ${this.accessToken}`; - const userId = extractJwtSubject(this.accessToken); + const { accessToken } = this.sdk.getTokens(); + if (auth !== "none" && accessToken) { + const userId = extractJwtSubject(accessToken); if (userId) { headers["X-User-Id"] = userId; } @@ -204,39 +178,6 @@ export class WebRuntimeClient { return headers; } - private async refreshTokens(): Promise { - if (!this.refreshToken) { - return false; - } - - const response = await fetch(this.url(PATH_AUTH_REFRESH), { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ refresh_token: this.refreshToken }), - cache: "no-store", - }); - if (!response.ok) { - return false; - } - - const payload = (await response.json()) as { - access_token?: unknown; - refresh_token?: unknown; - }; - if ( - typeof payload.access_token !== "string" || - typeof payload.refresh_token !== "string" - ) { - return false; - } - - this.accessToken = payload.access_token; - this.refreshToken = payload.refresh_token; - this.sdk.setTokens(payload.access_token, payload.refresh_token); - await this.persistTokens(payload.access_token, payload.refresh_token); - return true; - } - private async persistTokens( accessToken: string, refreshToken: string, @@ -244,7 +185,8 @@ export class WebRuntimeClient { const { cookies } = await import("next/headers"); const cookieStore = await cookies(); // A newer refresh may have replaced this pair while cookie acquisition waited. - if (this.accessToken !== accessToken || this.refreshToken !== refreshToken) { + const current = this.sdk.getTokens(); + if (current.accessToken !== accessToken || current.refreshToken !== refreshToken) { return; } cookieStore.set(ACCESS_TOKEN_COOKIE, accessToken, { From 2bbde77e59b1f594bcb01ed7ff1dab838d026bcc Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Tue, 27 Oct 2026 16:08:38 +0800 Subject: [PATCH 5/7] fix(auth): fence credential transitions and stream refresh --- crates/services/src/auth/mod.rs | 47 +++-- .../tests/auth_logout_rotation_db_it.rs | 158 ++++++++++++++++ docs/design/authentication.md | 5 + docs/reference/api-reference.md | 5 + packages/sdk/README.md | 11 +- .../auth-lifecycle-regression.test.ts | 173 ++++++++++++++++++ .../src/__tests__/client-edge-stream.test.ts | 4 +- packages/sdk/src/client.ts | 91 ++++++--- packages/sdk/src/sse-client.ts | 2 +- packages/sdk/src/types.ts | 2 + web/__tests__/lib/chat-cold-recovery.test.ts | 1 + .../lib/runtime-client-refresh.test.ts | 17 ++ web/app/api/chats/[chatId]/stream/route.ts | 1 + .../branches/[branchId]/turns/route.ts | 1 + web/lib/runtime-client/server.ts | 7 +- 15 files changed, 475 insertions(+), 50 deletions(-) create mode 100644 crates/services/tests/auth_logout_rotation_db_it.rs create mode 100644 packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts diff --git a/crates/services/src/auth/mod.rs b/crates/services/src/auth/mod.rs index ceb28dc373..4e7a3b508e 100644 --- a/crates/services/src/auth/mod.rs +++ b/crates/services/src/auth/mod.rs @@ -1867,12 +1867,37 @@ impl AuthService for DatabaseAuthService { .await .map_err(|e| map_auth_sqlx(e, "auth.get_pool", None))?; let refresh_token_hash = sha256_hex(&request.refresh_token); - let stored = self - .fetch_refresh_token(&pool, &refresh_token_hash) + let _invalidation = self.authentication_cache.mutation(); + let mut tx = pool + .begin() .await - .map_err(|e| map_auth_sqlx(e, "logout.fetch_refresh_token", Some(&pool)))? - .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Token expired or revoked"))?; - + .map_err(|e| map_auth_sqlx(e, "logout.begin_tx", Some(&pool)))?; + let account = query("SELECT user_id FROM auth_users WHERE user_id = ? FOR UPDATE") + .bind(&user_id) + .fetch_optional(&mut *tx) + .await + .map_err(|e| map_auth_sqlx(e, "logout.lock_user", Some(&pool)))?; + if account.is_none() { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Token expired or revoked", + )); + } + // A consumed token cannot refresh again, but can still terminate the + // session it originally owned. The shared user lock serializes this + // revocation with refresh's compare-and-revoke and successor insert. + let stored = query( + "SELECT user_id, DATE_FORMAT(expires_at, '%Y-%m-%dT%H:%i:%s') AS expires_at, session_id \ + FROM auth_refresh_tokens WHERE token_hash = ? LIMIT 1", + ) + .bind(&refresh_token_hash) + .fetch_optional(&mut *tx) + .await + .map_err(|e| map_auth_sqlx(e, "logout.fetch_refresh_token", Some(&pool)))? + .map(|row| decode_refresh_token_row(&row)) + .transpose() + .map_err(|e| map_auth_sqlx(e, "logout.decode_refresh_token", Some(&pool)))? + .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Token expired or revoked"))?; if stored.0 != user_id || stored.2.as_deref() != Some(session_id.as_str()) { return Err(error_response( StatusCode::UNAUTHORIZED, @@ -1885,18 +1910,12 @@ impl AuthService for DatabaseAuthService { "Token expired or revoked", )); } - - let _invalidation = self.authentication_cache.mutation(); - let mut tx = pool - .begin() - .await - .map_err(|e| map_auth_sqlx(e, "logout.begin_tx", Some(&pool)))?; - query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ? AND user_id = ?") - .bind(&refresh_token_hash) + query("UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE user_id = ? AND session_id = ? AND is_revoked = 0") .bind(&user_id) + .bind(&session_id) .execute(&mut *tx) .await - .map_err(|e| map_auth_sqlx(e, "logout.revoke_submitted_token", Some(&pool)))?; + .map_err(|e| map_auth_sqlx(e, "logout.revoke_session", Some(&pool)))?; tx.commit() .await diff --git a/crates/services/tests/auth_logout_rotation_db_it.rs b/crates/services/tests/auth_logout_rotation_db_it.rs new file mode 100644 index 0000000000..be676e142d --- /dev/null +++ b/crates/services/tests/auth_logout_rotation_db_it.rs @@ -0,0 +1,158 @@ +mod common; + +use astra_core::JwtSettings; +use astra_services::{ + AuthLoginRequestData, AuthRefreshRequestData, AuthRegisterRequestData, AuthService, + DatabaseAuthService, +}; +use axum::http::{HeaderMap, HeaderValue, StatusCode}; + +fn headers(token: &str) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {token}")).unwrap(), + ); + headers +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires live MatrixOne: ASTRA_TEST_DB_IT=1"] +async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_other_sessions() { + let (pool, settings) = common::setup_pool_and_settings().await; + let service = DatabaseAuthService::new( + settings, + JwtSettings { + secret_key: "isolated-logout-rotation-test-key".into(), + algorithm: "HS256".into(), + access_token_expire_minutes: 5, + refresh_token_expire_days: 1, + }, + ) + .with_pool(pool.clone()); + let username = format!("logout_{}", uuid::Uuid::new_v4().simple()); + let password = "correct-horse-battery-staple"; + let user = service + .register(AuthRegisterRequestData { + username: username.clone(), + email: format!("{username}@example.test"), + password: password.into(), + display_name: None, + }) + .await + .unwrap(); + let login = || AuthLoginRequestData { + username: username.clone(), + password: password.into(), + }; + let original = service.login(login()).await.unwrap(); + let other = service.login(login()).await.unwrap(); + for invalid in [original.access_token.clone(), "invalid.token".to_string()] { + assert_eq!( + service + .logout(AuthRefreshRequestData { + refresh_token: invalid, + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + } + let rotated = service + .refresh(AuthRefreshRequestData { + refresh_token: original.refresh_token.clone(), + }) + .await + .unwrap(); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: original.refresh_token.clone() + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + service + .logout(AuthRefreshRequestData { + refresh_token: original.refresh_token.clone(), + }) + .await + .unwrap(); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: rotated.refresh_token + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + service + .current_principal(&headers(&rotated.access_token)) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + service + .current_principal(&headers(&other.access_token)) + .await + .expect("another session remains active"); + + // Exercise simultaneous refresh and logout with three independently issued sessions. + // Whichever transaction wins, no successor from that session may stay usable. + for _ in 0..3 { + let token = service.login(login()).await.unwrap(); + let (refreshed, logged_out) = tokio::join!( + service.refresh(AuthRefreshRequestData { + refresh_token: token.refresh_token.clone() + }), + service.logout(AuthRefreshRequestData { + refresh_token: token.refresh_token.clone() + }), + ); + logged_out.unwrap(); + match refreshed { + Ok(next) => { + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: next.refresh_token + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + service + .current_principal(&headers(&next.access_token)) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + } + Err(error) => assert_eq!(error.0, StatusCode::UNAUTHORIZED), + } + service + .current_principal(&headers(&other.access_token)) + .await + .expect("another session remains active"); + } + for table in ["auth_refresh_tokens", "auth_user_roles", "auth_users"] { + sqlx::query(sqlx::AssertSqlSafe(format!( + "DELETE FROM {table} WHERE user_id = ?" + ))) + .bind(&user.user_id) + .execute(pool.get()) + .await + .unwrap(); + } + pool.get().close().await; +} diff --git a/docs/design/authentication.md b/docs/design/authentication.md index f50b9442f4..7b5358fd66 100644 --- a/docs/design/authentication.md +++ b/docs/design/authentication.md @@ -67,6 +67,11 @@ unrelated principal. Absolute token, observed Session and Memoria access expiry are checked again after awaited validation. A request starting after a local logout cannot reuse an invalidated observation. +Refresh and logout serialize on the account lock. Logout validates the +recorded token and revokes its entire session, so a concurrent rotation +cannot leave a usable successor. A previously consumed, unexpired token may +terminate its own session but never regain refresh authority. + This cache observes identity, not operation permissions. Memory operations continue resolving scoped credentials through the canonical credential resolver; provider-request replay checks, sensitive reauthentication and execution-handoff diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 527aac18fe..849443d8a7 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -10,6 +10,11 @@ All protected endpoints require a JWT. The public authentication exceptions are `/live`, `/ready`, `/health`, `/auth/register`, `/auth/login`, `GET /auth/methods`, `POST /auth/memoria`, `/auth/refresh`, and `/auth/logout` (refresh/logout authenticate the supplied refresh token): +Logout revokes the token's session, including a successor issued by a +concurrent refresh. A signed, unexpired refresh token recorded for that +session can terminate it even after rotation; it cannot be used to refresh +again. Other sessions are unaffected. + ``` Authorization: Bearer ``` diff --git a/packages/sdk/README.md b/packages/sdk/README.md index b8cb668033..ad9b6e65fa 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -161,10 +161,13 @@ already accepted frame to other subscribed callbacks. ### Browser apps and the BFF pattern -Within one `AstraClient`, concurrent REST and buffered run-event requests share +Within one `AstraClient`, concurrent REST, live streams and buffered run-event requests share one token refresh. A delayed 401 reuses the rotated token. Refresh responses started before `setTokens`, successful login/registration, or logout cannot restore the previous credentials or retry an old request under a new account. +`setTokens(accessToken)` clears any previous refresh token; pass both tokens +when replacing a credential pair. Pending login and logout operations cannot +overwrite a newer authentication decision. This coordination is per client instance, not across separate BFF requests. `onTokenRefresh` may be asynchronous and call the same client. Requests already @@ -183,7 +186,11 @@ error; initial transport errors still propagate. `fetchResponse(path, init, { auth })` returns an unconsumed response, including non-2xx responses, without inferring a JSON content type. Its default auth is `"optional"`; `"none"` suppresses default authorization and automatic refresh. -Explicit request headers still override defaults. `getTokens()` returns a +Explicit request headers still override defaults and are not automatically +refreshed using the client's credentials. Raw mutation requests are not retried +unless `retryUnauthorized` is explicitly enabled for an endpoint that rejects +authentication before admitting the operation. Streams and one-shot bodies +are never replayed after acceptance. `getTokens()` returns a detached snapshot of current credentials, not the initial configuration. Bodies must remain replayable for the single automatic retry. diff --git a/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts new file mode 100644 index 0000000000..d9ce9f722a --- /dev/null +++ b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts @@ -0,0 +1,173 @@ +import { AstraClient } from '../client'; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} +const unauthorized = () => Response.json({ detail: 'expired' }, { status: 401 }); +const tokens = { access_token: 'new-access', refresh_token: 'new-refresh' }; +const bearer = (init?: RequestInit) => new Headers(init?.headers).get('authorization'); +afterEach(() => vi.unstubAllGlobals()); + +test('replacing only the access token cannot refresh as the previous account', async () => { + const fetch = vi.fn(async () => unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'alice', refreshToken: 'alice-refresh' }); + client.setTokens('bob'); + await expect(client.getMe()).rejects.toMatchObject({ status: 401 }); + expect(client.getTokens()).toEqual({ accessToken: 'bob', refreshToken: null }); + expect(fetch).toHaveBeenCalledTimes(1); +}); + +test.each([204, 500])('late logout response %i cannot erase a newer login', async (status) => { + const logout = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string) => { + if (url.endsWith('/auth/logout')) return logout.promise; + if (url.endsWith('/auth/login')) return Response.json(tokens); + throw new Error(url); + })); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const pending = client.logout().then(() => 'ok', () => 'error'); + expect(client.getTokens()).toEqual({ accessToken: null, refreshToken: null }); + await client.login('bob', 'password'); + logout.resolve(status === 204 ? new Response(null, { status }) : new Response('failed', { status })); + expect(await pending).toBe(status === 204 ? 'ok' : 'error'); + expect(client.getTokens()).toEqual({ accessToken: 'new-access', refreshToken: 'new-refresh' }); +}); + +test.each(['login', 'register'] as const)('pending %s cannot restore a session after logout', async (method) => { + const result = deferred(); + vi.stubGlobal('fetch', vi.fn(async () => result.promise)); + const client = new AstraClient({ baseUrl: 'https://example.test' }); + const pending = client[method]('bob', 'password'); + await client.logout(); + result.resolve(Response.json(tokens)); + await expect(pending).rejects.toThrow('superseded'); + expect(client.getTokens()).toEqual({ accessToken: null, refreshToken: null }); +}); + +test('a slower login cannot overwrite a later successful login', async () => { + const first = deferred(); + let calls = 0; + vi.stubGlobal('fetch', vi.fn(async () => ++calls === 1 ? first.promise : Response.json(tokens))); + const client = new AstraClient({ baseUrl: 'https://example.test' }); + const oldLogin = client.login('alice', 'password'); + await client.login('bob', 'password'); + first.resolve(Response.json({ access_token: 'alice', refresh_token: 'alice-refresh' })); + await expect(oldLogin).rejects.toThrow('superseded'); + expect(client.getTokens()).toEqual({ accessToken: 'new-access', refreshToken: 'new-refresh' }); +}); + +test.each(['login', 'register'] as const)('failed %s does not discard an in-flight refresh of the current account', async (method) => { + const refresh = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return refresh.promise; + if (url.endsWith('/auth/login') || url.endsWith('/auth/register')) return unauthorized(); + return bearer(init) === 'Bearer new-access' ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const me = client.getMe(); + await vi.waitFor(() => expect(vi.mocked(fetch)).toHaveBeenCalledTimes(2)); + await expect(client[method]('bob', 'bad-password')).rejects.toMatchObject({ status: 401 }); + refresh.resolve(Response.json(tokens)); + await expect(me).resolves.toEqual({ user_id: 'alice' }); + expect(client.getTokens()).toEqual({ accessToken: 'new-access', refreshToken: 'new-refresh' }); +}); + +test.each([ + { authorization: 'Bearer external' }, + new Headers({ AUTHORIZATION: 'Bearer external' }), + { forEach: (callback: (value: string, key: string) => void) => callback('Bearer external', 'authorization') } as Headers, +])('explicit authorization is not refreshed using client credentials', async (headers) => { + const fetch = vi.fn(async () => unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const response = await client.fetchResponse('/resource', { headers }); + expect(response.status).toBe(401); + expect(fetch).toHaveBeenCalledTimes(1); + expect(client.getTokens()).toEqual({ accessToken: 'old', refreshToken: 'old-refresh' }); +}); + +test('raw mutations require explicit retry and a replayable body', async () => { + let refreshes = 0; + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { refreshes++; return Response.json(tokens); } + return bearer(init) === 'Bearer new-access' ? new Response('ok') : unauthorized(); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + expect((await client.fetchResponse('/write', { method: 'POST', body: '{}' })).status).toBe(401); + expect(refreshes).toBe(0); + expect((await client.fetchResponse('/write', { method: 'POST', body: '{}' }, { retryUnauthorized: true })).status).toBe(200); + expect(refreshes).toBe(1); + client.setTokens('old', 'old-refresh'); + const body = new ReadableStream({ start(controller) { controller.close(); } }); + expect((await client.fetchResponse('/write', { method: 'POST', body }, { retryUnauthorized: true })).status).toBe(401); + expect(refreshes).toBe(1); +}); + +test.each(['chat', 'work'] as const)('%s stream and REST share refresh without replaying an accepted stream', async (kind) => { + const refresh = deferred(); + let refreshes = 0; + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { refreshes++; return refresh.promise; } + if (bearer(init) === 'Bearer old') return unauthorized(); + if (url.endsWith('/auth/me')) return Response.json({ user_id: 'alice' }); + return new Response('data: {"type":"run_finished","run_id":"run-1","status":"completed"}\n\n'); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const events: unknown[] = []; + if (kind === 'chat') client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: event => events.push(event) }); + else client.continueWorkBranch('work-1', 'branch-1', { requestId: 'request-1', attachmentId: 'attachment-1', message: 'hi' }, { onEvent: event => events.push(event) }); + const rest = client.getMe(); + await vi.waitFor(() => expect(refreshes).toBe(1)); + refresh.resolve(Response.json(tokens)); + await expect(rest).resolves.toEqual({ user_id: 'alice' }); + await vi.waitFor(() => expect(events).toHaveLength(1)); + expect(events[0]).toMatchObject({ type: 'run_finished' }); + expect(refreshes).toBe(1); + expect(fetch.mock.calls.filter(([url]) => url.includes(kind === 'chat' ? '/chat/stream' : '/turns'))).toHaveLength(2); +}); + +test('closing a stream while refresh is pending prevents a second POST', async () => { + const refresh = deferred(); + const fetch = vi.fn(async (url: string) => url.endsWith('/auth/refresh') ? refresh.promise : unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const stream = client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: () => {} }); + await vi.waitFor(() => expect(fetch).toHaveBeenCalledTimes(2)); + stream.close(); + refresh.resolve(Response.json(tokens)); + await vi.waitFor(() => expect(client.getTokens().accessToken).toBe('new-access')); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/chat/stream'))).toHaveLength(1); +}); + +test('reconnecting an old stream cannot issue a request under a replacement account', async () => { + const fetch = vi.fn(async () => new Response('data: {"type":"run_finished","status":"completed"}\n\n')); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'alice' }); + const events: unknown[] = []; + const stream = client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: event => events.push(event) }); + await vi.waitFor(() => expect(events).toHaveLength(1)); + client.setTokens('bob'); + await stream.connect(); + expect(fetch).toHaveBeenCalledTimes(1); + expect(events[1]).toMatchObject({ type: 'error', retryable: false }); +}); + +test('stream requests preserve configured headers while using the current owned bearer', async () => { + const fetch = vi.fn(async (_url: string, _init?: RequestInit) => new Response('data: {"type":"run_finished","status":"completed"}\n\n')); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'current', headers: { + Authorization: 'Bearer configured', Accept: 'application/custom', 'Content-Type': 'application/custom+json', + } }); + const events: unknown[] = []; + client.streamChat({ message: 'hi', modelSelection: { offeringId: 'offer' } }, { onEvent: event => events.push(event) }); + await vi.waitFor(() => expect(events).toHaveLength(1)); + const headers = new Headers(fetch.mock.calls[0][1]?.headers); + expect(headers.get('authorization')).toBe('Bearer current'); + expect(headers.get('accept')).toBe('application/custom'); + expect(headers.get('content-type')).toBe('application/custom+json'); +}); diff --git a/packages/sdk/src/__tests__/client-edge-stream.test.ts b/packages/sdk/src/__tests__/client-edge-stream.test.ts index cbf59ddaad..7c2cd2ff11 100644 --- a/packages/sdk/src/__tests__/client-edge-stream.test.ts +++ b/packages/sdk/src/__tests__/client-edge-stream.test.ts @@ -63,8 +63,8 @@ describe('AstraClient — streamChat', () => { const [url, init] = (globalThis.fetch as ReturnType).mock.calls[0] as [string, RequestInit]; expect(String(url).endsWith(PATH_CHAT_STREAM)).toBe(true); expect(init.method).toBe('POST'); - const headers = init.headers as Record; - expect(headers['Authorization']).toBe('Bearer tok-1'); + const headers = new Headers(init.headers); + expect(headers.get('Authorization')).toBe('Bearer tok-1'); expect(init.body).toBe( JSON.stringify( chatRequestToWire({ diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index 688f3ba76d..8c304a5d36 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -627,6 +627,7 @@ export class AstraClient { private accessToken: string | null; private refreshTokenValue: string | null; private authGeneration = 0; + private authAttempt: object | null = null; private refreshAttempt: { generation: number; promise: Promise } | null = null; constructor(config: AstraClientConfig) { @@ -654,34 +655,48 @@ export class AstraClient { const email = options?.email ?? `${username}@users.local.astra`; const body: Record = { username, email, password }; if (options?.displayName) body.display_name = options.displayName; - const result = await this.post(PATH_AUTH_REGISTER, body); - this.setTokens(result.access_token, result.refresh_token); - return result; + return this.authenticate(PATH_AUTH_REGISTER, body); } /** Log in with username/password. Stores tokens automatically. */ async login(username: string, password: string): Promise { - const result = await this.post(PATH_AUTH_LOGIN, { + return this.authenticate(PATH_AUTH_LOGIN, { username, password, }); - this.setTokens(result.access_token, result.refresh_token); - return result; } - /** Log out and clear stored tokens. Requires refresh token in client state. */ - async logout(): Promise { - this.authGeneration++; + private async authenticate(path: string, body: unknown): Promise { + const attempt = {}; + this.authAttempt = attempt; try { - if (this.refreshTokenValue) { - await this.post(PATH_AUTH_LOGOUT, { - refresh_token: this.refreshTokenValue, - }); + const response = await this.request(path, { + method: "POST", + body: JSON.stringify(body), + }, { auth: "none" }); + const result = await response.json() as AuthResult; + if (this.authAttempt !== attempt) { + throw new Error("Authentication attempt was superseded"); } + this.setTokens(result.access_token, result.refresh_token); + return result; } finally { - this.authGeneration++; - this.accessToken = null; - this.refreshTokenValue = null; + if (this.authAttempt === attempt) this.authAttempt = null; + } + } + + /** Log out and clear stored tokens. Requires refresh token in client state. */ + async logout(): Promise { + this.authGeneration++; + this.authAttempt = null; + const refreshToken = this.refreshTokenValue; + this.accessToken = null; + this.refreshTokenValue = null; + if (refreshToken) { + await this.request(PATH_AUTH_LOGOUT, { + method: "POST", + body: JSON.stringify({ refresh_token: refreshToken }), + }, { auth: "none" }); } } @@ -745,8 +760,9 @@ export class AstraClient { setTokens(accessToken: string, refreshToken?: string): void { this.authGeneration++; + this.authAttempt = null; this.accessToken = accessToken; - if (refreshToken) this.refreshTokenValue = refreshToken; + this.refreshTokenValue = refreshToken || null; } /** A detached snapshot of the credentials currently owned by this client. */ @@ -810,37 +826,49 @@ export class AstraClient { return headersInitToRecord({}, headers); } + private streamHeaders(extra?: Record): Record { + const headers = new Headers(headersInitToRecord(this.config.headers ?? {})); + headers.delete("Authorization"); + return headersInitToRecord(headersInitToRecord({}, headers), extra); + } + // ─── HTTP helpers ────────────────────────────────────────────────── /** Raw, body-neutral transport; returns the final response without consuming it. */ async fetchResponse( path: string, init?: RequestInit, - options?: { auth?: "optional" | "none" }, + options?: { auth?: "optional" | "none"; retryUnauthorized?: boolean }, ): Promise { const auth = options?.auth ?? "optional"; const url = this.apiPath(path); const generation = this.authGeneration; const accessToken = this.accessToken; + const requestHeaders = headersInitToRecord({}, init?.headers); + const explicitAuthorization = new Headers(requestHeaders).has("Authorization"); + const method = (init?.method ?? "GET").toUpperCase(); + const replayable = !(typeof ReadableStream !== "undefined" && init?.body instanceof ReadableStream); + const mayRetry = options?.retryUnauthorized ?? (method === "GET" || method === "HEAD"); const send = () => fetch(url, { ...init, - headers: headersInitToRecord(this.buildHeaders(auth), init?.headers), + headers: headersInitToRecord(this.buildHeaders(auth), requestHeaders), }); let res = await send(); - if (res.status === 401 && auth !== "none") { + if (res.status === 401 && auth !== "none" && accessToken && !explicitAuthorization && mayRetry && replayable && !init?.signal?.aborted) { const refreshed = await this.tryRefreshToken(generation, accessToken); + if (init?.signal?.aborted) throw new DOMException("The operation was aborted", "AbortError"); if (refreshed && generation === this.authGeneration) res = await send(); } return res; } - private async request(path: string, init?: RequestInit): Promise { + private async request(path: string, init?: RequestInit, options?: { auth?: "optional" | "none" }): Promise { const method = (init?.method ?? "GET").toUpperCase(); const json = init?.body != null && ["POST", "PUT", "PATCH", "DELETE"].includes(method); const res = await this.fetchResponse(path, { ...init, headers: headersInitToRecord(json ? { "Content-Type": "application/json" } : {}, init?.headers), - }); + }, { ...options, retryUnauthorized: true }); if (!res.ok) { const error = await readAstraError(res); @@ -2697,6 +2725,7 @@ export class AstraClient { ); } const url = this.apiPath(workBranchTurnsPath(workId, branchId)); + const generation = this.authGeneration; let client: SSEClient; client = new SSEClient({ url, @@ -2706,11 +2735,13 @@ export class AstraClient { attachment_id: input.attachmentId, message: input.message, }), - token: this.accessToken ?? undefined, - headers: { - ...this.config.headers, - [ASTRA_WORK_API_MAJOR_HEADER]: ASTRA_WORK_API_MAJOR, + transport: (_url, init) => { + if (generation !== this.authGeneration) throw new Error("Authentication changed before stream connection"); + return this.fetchResponse(workBranchTurnsPath(workId, branchId), init, { retryUnauthorized: true }); }, + headers: this.streamHeaders({ + [ASTRA_WORK_API_MAJOR_HEADER]: ASTRA_WORK_API_MAJOR, + }), onEvent: (event) => { try { callbacks.onEvent(decodeWorkTurnStreamEventV1(event)); @@ -3412,12 +3443,16 @@ export class AstraClient { }, ): SSEClient { const url = this.apiPath(PATH_CHAT_STREAM); + const generation = this.authGeneration; const client = new SSEClient({ url, method: "POST", body: JSON.stringify(chatRequestToWire(request)), - token: this.accessToken ?? undefined, - headers: this.config.headers, + transport: (_url, init) => { + if (generation !== this.authGeneration) throw new Error("Authentication changed before stream connection"); + return this.fetchResponse(PATH_CHAT_STREAM, init, { retryUnauthorized: true }); + }, + headers: this.streamHeaders(), onEvent: callbacks.onEvent, onStateChange: callbacks.onStateChange, onRawLine: callbacks.onRawLine, diff --git a/packages/sdk/src/sse-client.ts b/packages/sdk/src/sse-client.ts index ba2d68692f..b4ee134931 100644 --- a/packages/sdk/src/sse-client.ts +++ b/packages/sdk/src/sse-client.ts @@ -187,7 +187,7 @@ export class SSEClient { headers = headersInitToRecord(headers, { 'Content-Type': 'application/json' }); } - const response = await fetch(this.options.url, { + const response = await (this.options.transport ?? fetch)(this.options.url, { method: this.options.method ?? 'GET', headers, body: this.options.body, diff --git a/packages/sdk/src/types.ts b/packages/sdk/src/types.ts index 593b31626e..f93bca4eaa 100644 --- a/packages/sdk/src/types.ts +++ b/packages/sdk/src/types.ts @@ -1209,6 +1209,8 @@ export type AstraClientConfig = { export type SSEClientOptions = { url: string; + /** Optional request transport for clients that own authentication. */ + transport?: (url: string, init: RequestInit) => Promise; token?: string; headers?: Record; onEvent: (event: StreamEvent) => void; diff --git a/web/__tests__/lib/chat-cold-recovery.test.ts b/web/__tests__/lib/chat-cold-recovery.test.ts index c9bdc72889..37dc2ce3eb 100644 --- a/web/__tests__/lib/chat-cold-recovery.test.ts +++ b/web/__tests__/lib/chat-cold-recovery.test.ts @@ -28,6 +28,7 @@ function setup() { globalThis.__astraWebStores = undefined; vi.mocked(requireRuntimeUser).mockResolvedValue({ user: { user_id: "owner", display_name: null }, response: null } as never); const sdk = { + getTokens: vi.fn(() => ({ accessToken: "synthetic-test", refreshToken: null })), listRuntimeSessions: vi.fn().mockResolvedValue({ sessions: [session] }), listRuns: vi.fn().mockResolvedValue({ runs: [{ runId: "run", sessionId: "session", parentRunId: null, depth: 0, status: "running" }] }), getSessionConversation: vi.fn().mockResolvedValue(conversation()), diff --git a/web/__tests__/lib/runtime-client-refresh.test.ts b/web/__tests__/lib/runtime-client-refresh.test.ts index 2f60c909e6..ba35a53ef0 100644 --- a/web/__tests__/lib/runtime-client-refresh.test.ts +++ b/web/__tests__/lib/runtime-client-refresh.test.ts @@ -200,3 +200,20 @@ it("does not authenticate or refresh anonymous Web calls and preserves explicit expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get("Authorization")).toBe("Bearer explicit"); expect(new Headers(fetchMock.mock.calls[1][1]?.headers).get("X-User-Id")).toBe("hint"); }); + +it("only retries a raw mutation when its caller opts in", async () => { + vi.mocked(cookies).mockResolvedValue({ set: vi.fn() } as unknown as Awaited>); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json({ access_token: 'new', refresh_token: 'rotated' }); + return new Headers(init?.headers).get('Authorization') === 'Bearer new' + ? new Response('ok') : new Response('expired', { status: 401 }); + }); + vi.stubGlobal('fetch', fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig('old', 'refresh')); + expect((await runtime.fetchResponse('/write', { method: 'POST', json: { value: 1 } })).status).toBe(401); + expect(fetchMock).toHaveBeenCalledTimes(1); + await expect(runtime.post('/write', { value: 1 }, { retryUnauthorized: false })).rejects.toMatchObject({ status: 401 }); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect((await runtime.fetchResponse('/write', { method: 'POST', json: { value: 1 }, retryUnauthorized: true })).status).toBe(200); + expect(fetchMock).toHaveBeenCalledTimes(5); +}); diff --git a/web/app/api/chats/[chatId]/stream/route.ts b/web/app/api/chats/[chatId]/stream/route.ts index e47cc52be8..e76ab11839 100644 --- a/web/app/api/chats/[chatId]/stream/route.ts +++ b/web/app/api/chats/[chatId]/stream/route.ts @@ -702,6 +702,7 @@ export async function POST( return runtime.fetchResponse(PATH_CHAT_STREAM, { method: "POST", auth: "required", + retryUnauthorized: true, operation: "stream web chat turn", signal: backendAbortController.signal, json: { diff --git a/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts b/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts index 2eb600c952..4c28cab38e 100644 --- a/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts +++ b/web/app/api/works/[workId]/branches/[branchId]/turns/route.ts @@ -148,6 +148,7 @@ export async function POST( const backend = await runtime.fetchResponse(path, { method: "POST", auth: "required", + retryUnauthorized: true, body: JSON.stringify(input), signal: request.signal, headers: { diff --git a/web/lib/runtime-client/server.ts b/web/lib/runtime-client/server.ts index 02696cce34..418cbd85c8 100644 --- a/web/lib/runtime-client/server.ts +++ b/web/lib/runtime-client/server.ts @@ -19,6 +19,7 @@ export type RuntimeAuthMode = "required" | "optional" | "none"; export type RuntimeRequestInit = Omit & { auth?: RuntimeAuthMode; + retryUnauthorized?: boolean; json?: unknown; operation?: string; body?: BodyInit | null; @@ -79,7 +80,7 @@ export class WebRuntimeClient { } async request(path: string, init: RuntimeRequestInit = {}): Promise { - const response = await this.fetchResponse(path, init); + const response = await this.fetchResponse(path, { ...init, retryUnauthorized: init.retryUnauthorized ?? true }); if (!response.ok) { throw new RuntimeClientError({ operation: init.operation ?? `${init.method ?? "GET"} ${path}`, @@ -120,7 +121,7 @@ export class WebRuntimeClient { return this.sdk.fetchResponse( path, this.toRequestInit(init, auth, operation, path), - { auth: auth === "none" ? "none" : "optional" }, + { auth: auth === "none" ? "none" : "optional", retryUnauthorized: init.retryUnauthorized }, ); } @@ -150,7 +151,7 @@ export class WebRuntimeClient { }); } - const { auth: _auth, json: _json, operation: _operation, ...rest } = init; + const { auth: _auth, json: _json, operation: _operation, retryUnauthorized: _retryUnauthorized, ...rest } = init; return { ...rest, method, From ba3550d3f7ea0839404b08ef1f7471407abe6e52 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Tue, 27 Oct 2026 16:50:33 +0800 Subject: [PATCH 6/7] fix(auth): retain logout proof and cancel refresh waiters --- crates/services/src/storage.rs | 8 +-- .../tests/auth_logout_rotation_db_it.rs | 42 ++++++++++++- docs/design/authentication.md | 4 +- .../auth-lifecycle-regression.test.ts | 60 +++++++++++++++++++ packages/sdk/src/client.ts | 56 +++++++++++------ 5 files changed, 145 insertions(+), 25 deletions(-) diff --git a/crates/services/src/storage.rs b/crates/services/src/storage.rs index 60a1fa7f7d..6264ab19f0 100644 --- a/crates/services/src/storage.rs +++ b/crates/services/src/storage.rs @@ -6867,7 +6867,7 @@ pub struct CleanupResult { /// Configuration for expiry policies of authentication and operational data. pub struct RetentionPolicy { - /// Max age in days for expired/revoked refresh tokens (default: 7) + /// Minimum age in days for expired refresh token records (default: 7) pub refresh_token_days: u32, /// Max age in days for inactive auth tokens (default: 30) pub auth_token_days: u32, @@ -6909,12 +6909,12 @@ pub async fn cleanup_expired_data( const AUTH_AUDIT_LOG_BATCH_LIMIT: u32 = 1000; let mut results = Vec::new(); - // 1. Expired + revoked refresh tokens + // 1. Expired refresh tokens past the minimum retention age let deleted = sqlx::query( "DELETE FROM auth_refresh_tokens \ - WHERE (expires_at < NOW(6) OR is_revoked = 1) \ + WHERE expires_at < NOW(6) \ AND created_at < DATE_SUB(NOW(6), INTERVAL ? DAY) \ - ORDER BY created_at ASC, token_id ASC \ + ORDER BY expires_at ASC, token_id ASC \ LIMIT ?", ) .bind(policy.refresh_token_days) diff --git a/crates/services/tests/auth_logout_rotation_db_it.rs b/crates/services/tests/auth_logout_rotation_db_it.rs index be676e142d..2d8eb717b8 100644 --- a/crates/services/tests/auth_logout_rotation_db_it.rs +++ b/crates/services/tests/auth_logout_rotation_db_it.rs @@ -3,7 +3,7 @@ mod common; use astra_core::JwtSettings; use astra_services::{ AuthLoginRequestData, AuthRefreshRequestData, AuthRegisterRequestData, AuthService, - DatabaseAuthService, + DatabaseAuthService, RetentionPolicy, cleanup_expired_data, }; use axum::http::{HeaderMap, HeaderValue, StatusCode}; @@ -46,6 +46,12 @@ async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_oth password: password.into(), }; let original = service.login(login()).await.unwrap(); + // Model a still-valid token issued more than the cleanup retention age ago. + sqlx::query("UPDATE auth_refresh_tokens SET created_at = DATE_SUB(NOW(6), INTERVAL 8 DAY) WHERE user_id = ?") + .bind(&user.user_id) + .execute(pool.get()) + .await + .unwrap(); let other = service.login(login()).await.unwrap(); for invalid in [original.access_token.clone(), "invalid.token".to_string()] { assert_eq!( @@ -75,12 +81,46 @@ async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_oth .0, StatusCode::UNAUTHORIZED ); + // The old token's revoked row may be removed even though its JWT is still valid. + let cleanup = cleanup_expired_data(pool.get(), &RetentionPolicy::default()) + .await + .unwrap(); + assert!( + cleanup + .iter() + .any(|result| result.table == "auth_refresh_tokens") + ); + let retained: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM auth_refresh_tokens WHERE user_id = ? AND is_revoked = 1", + ) + .bind(&user.user_id) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!( + retained, 1, + "cleanup must retain an unexpired rotated token" + ); service .logout(AuthRefreshRequestData { refresh_token: original.refresh_token.clone(), }) .await .unwrap(); + // Once the expired record is beyond retention, cleanup must remove it. + sqlx::query("UPDATE auth_refresh_tokens SET expires_at = DATE_SUB(NOW(6), INTERVAL 8 DAY) WHERE user_id = ? AND is_revoked = 1 AND created_at < DATE_SUB(NOW(6), INTERVAL 7 DAY)") + .bind(&user.user_id) + .execute(pool.get()) + .await + .unwrap(); + let cleanup = cleanup_expired_data(pool.get(), &RetentionPolicy::default()) + .await + .unwrap(); + assert!( + cleanup + .iter() + .any(|result| result.table == "auth_refresh_tokens" && result.rows_deleted > 0) + ); assert_eq!( service .refresh(AuthRefreshRequestData { diff --git a/docs/design/authentication.md b/docs/design/authentication.md index 7b5358fd66..9a20093f64 100644 --- a/docs/design/authentication.md +++ b/docs/design/authentication.md @@ -70,7 +70,9 @@ logout cannot reuse an invalidated observation. Refresh and logout serialize on the account lock. Logout validates the recorded token and revokes its entire session, so a concurrent rotation cannot leave a usable successor. A previously consumed, unexpired token may -terminate its own session but never regain refresh authority. +terminate its own session but never regain refresh authority. Refresh-token +records remain until they are expired and older than the configured minimum +retention age, so cleanup cannot remove this logout proof early. This cache observes identity, not operation permissions. Memory operations continue resolving scoped credentials through the canonical credential resolver; diff --git a/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts index d9ce9f722a..fa83c6a7fb 100644 --- a/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts +++ b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts @@ -131,6 +131,66 @@ test.each(['chat', 'work'] as const)('%s stream and REST share refresh without r expect(fetch.mock.calls.filter(([url]) => url.includes(kind === 'chat' ? '/chat/stream' : '/turns'))).toHaveLength(2); }); +test('an aborted refresh waiter settles while other callers keep sharing the refresh', async () => { + const refresh = deferred(); + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return refresh.promise; + return bearer(init) === 'Bearer new-access' ? new Response('ok') : unauthorized(); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const controller = new AbortController(); + const aborted = client.fetchResponse('/cancelled', { signal: controller.signal }); + const other = client.fetchResponse('/other'); + await vi.waitFor(() => expect(fetch.mock.calls.filter(([url]) => url.endsWith('/auth/refresh'))).toHaveLength(1)); + controller.abort(); + // This must settle without releasing the pending shared refresh. + await expect(aborted).rejects.toMatchObject({ name: 'AbortError' }); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/cancelled'))).toHaveLength(1); + refresh.resolve(Response.json(tokens)); + expect((await other).status).toBe(200); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/auth/refresh'))).toHaveLength(1); +}); + +test('an abort between refresh completion and replay preserves the caller reason', async () => { + const controller = new AbortController(); + const reason = new Error('caller stopped'); + const fetch = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(tokens); + if (bearer(init) === 'Bearer new-access') { + if (url.endsWith('/first')) controller.abort(reason); + return new Response('ok'); + } + return unauthorized(); + }); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const first = client.fetchResponse('/first'); + const second = client.fetchResponse('/second', { signal: controller.signal }); + await expect(first).resolves.toMatchObject({ status: 200 }); + await expect(second).rejects.toBe(reason); + expect(fetch.mock.calls.filter(([url]) => url.endsWith('/second'))).toHaveLength(1); +}); + +test('aborted waiters release their registrations even when refresh never settles', async () => { + const refresh = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string) => url.endsWith('/auth/refresh') ? refresh.promise : unauthorized())); + const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old', refreshToken: 'old-refresh' }); + const reason = new Error('caller stopped'); + for (let i = 0; i < 20; i++) { + const controller = new AbortController(); + const removed = vi.spyOn(controller.signal, 'removeEventListener'); + const request = client.fetchResponse('/resource', { signal: controller.signal }); + await vi.waitFor(() => expect(vi.mocked(fetch)).toHaveBeenCalledTimes(i + 2)); + controller.abort(reason); + await expect(request).rejects.toBe(reason); + expect(removed).toHaveBeenCalledWith('abort', expect.any(Function)); + } + const attempt = (client as unknown as { refreshAttempt: { waiters: Set } }).refreshAttempt; + expect(attempt.waiters.size).toBe(0); + refresh.resolve(Response.json(tokens)); +}); + test('closing a stream while refresh is pending prevents a second POST', async () => { const refresh = deferred(); const fetch = vi.fn(async (url: string) => url.endsWith('/auth/refresh') ? refresh.promise : unauthorized()); diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index 8c304a5d36..ed9d3a42b3 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -628,7 +628,7 @@ export class AstraClient { private refreshTokenValue: string | null; private authGeneration = 0; private authAttempt: object | null = null; - private refreshAttempt: { generation: number; promise: Promise } | null = null; + private refreshAttempt: { generation: number; waiters: Set<(result: boolean) => void> } | null = null; constructor(config: AstraClientConfig) { this.config = config; @@ -770,23 +770,41 @@ export class AstraClient { return { accessToken: this.accessToken, refreshToken: this.refreshTokenValue }; } - private async tryRefreshToken(generation: number, rejectedToken: string | null): Promise { - if (generation !== this.authGeneration) return false; - if (this.refreshAttempt?.generation === generation) return this.refreshAttempt.promise; - // A delayed 401 may belong to a token another request already rotated. - if (rejectedToken !== this.accessToken) return true; - if (!this.refreshTokenValue) return false; - const attempt = { - generation, - promise: this.refreshTokens(generation, this.refreshTokenValue), - }; - this.refreshAttempt = attempt; - try { - return await attempt.promise; - } finally { - // An older account's completion must not clear a newer refresh attempt. - if (this.refreshAttempt === attempt) this.refreshAttempt = null; + private tryRefreshToken(generation: number, rejectedToken: string | null, signal?: AbortSignal | null): Promise { + if (signal?.aborted) return Promise.reject(signal.reason); + if (generation !== this.authGeneration) return Promise.resolve(false); + let attempt = this.refreshAttempt; + if (attempt?.generation !== generation) { + // A delayed 401 may belong to a token another request already rotated. + if (rejectedToken !== this.accessToken) return Promise.resolve(true); + if (!this.refreshTokenValue) return Promise.resolve(false); + attempt = { generation, waiters: new Set() }; + this.refreshAttempt = attempt; + const current = attempt; + // Observe the operation once; cancelled waiters can be removed even if it hangs. + void this.refreshTokens(generation, this.refreshTokenValue).then( + result => { + if (this.refreshAttempt === current) this.refreshAttempt = null; + for (const settle of current.waiters) settle(result); + }, + () => { + if (this.refreshAttempt === current) this.refreshAttempt = null; + for (const settle of current.waiters) settle(false); + }, + ); } + const current = attempt; + return new Promise((resolve, reject) => { + const cleanup = () => { + current.waiters.delete(settle); + signal?.removeEventListener("abort", abort); + }; + const settle = (result: boolean) => { cleanup(); resolve(result); }; + const abort = () => { cleanup(); reject(signal?.reason); }; + current.waiters.add(settle); + signal?.addEventListener("abort", abort, { once: true }); + if (signal?.aborted) abort(); + }); } private async refreshTokens(generation: number, refreshToken: string): Promise { @@ -855,8 +873,8 @@ export class AstraClient { }); let res = await send(); if (res.status === 401 && auth !== "none" && accessToken && !explicitAuthorization && mayRetry && replayable && !init?.signal?.aborted) { - const refreshed = await this.tryRefreshToken(generation, accessToken); - if (init?.signal?.aborted) throw new DOMException("The operation was aborted", "AbortError"); + const refreshed = await this.tryRefreshToken(generation, accessToken, init?.signal); + if (init?.signal?.aborted) throw init.signal.reason; if (refreshed && generation === this.authGeneration) res = await send(); } return res; From 0eabb13f86c68fa51a792f1ab4545a077e6df3a1 Mon Sep 17 00:00:00 2001 From: XuPeng-SH Date: Wed, 28 Oct 2026 00:22:46 +0800 Subject: [PATCH 7/7] fix(auth): bound refresh coordination and session revocation --- crates/services/src/auth/cache.rs | 151 +++++++++- crates/services/src/auth/mod.rs | 66 ++--- crates/services/src/storage.rs | 8 +- crates/services/tests/auth_cache_db_it.rs | 30 +- .../tests/auth_logout_rotation_db_it.rs | 269 ++++++++++++++---- docs/design/authentication.md | 33 ++- docs/reference/api-reference.md | 6 +- packages/sdk/README.md | 17 +- .../auth-lifecycle-regression.test.ts | 6 +- .../token-refresh-concurrency.test.ts | 212 +++++++++++++- packages/sdk/src/client.ts | 121 +++++--- scripts/schema/schema_inventory.py | 2 +- .../lib/runtime-client-refresh.test.ts | 84 +++++- web/lib/runtime-client/server.ts | 4 +- 14 files changed, 815 insertions(+), 194 deletions(-) diff --git a/crates/services/src/auth/cache.rs b/crates/services/src/auth/cache.rs index e097fc129f..a3c10f309d 100644 --- a/crates/services/src/auth/cache.rs +++ b/crates/services/src/auth/cache.rs @@ -13,10 +13,11 @@ const TTL: Duration = Duration::from_secs(30); const MAX_ENTRIES: usize = 1_024; const MAX_BYTES: usize = 4 * 1024 * 1024; const MAX_ENTRY_BYTES: usize = 16 * 1024; -// Remove/insert churn can retain 4096 buckets for 1024 live entries. As in -// the model cache, reserve keys, values and structural overhead for that -// allocation, including storage retained by invalidate(). -const CONTAINER_BYTES: usize = MAX_ENTRIES * 4 * (std::mem::size_of::<([u8; 32], Entry)>() + 64); +// Reserve entry and session-fence buckets, including allocations retained +// after invalidation. Churn can retain four buckets per live entry. +const CONTAINER_BYTES: usize = MAX_ENTRIES + * 4 + * (std::mem::size_of::<([u8; 32], Entry)>() + std::mem::size_of::<([u8; 32], u64)>() + 128); struct Entry { principal: AuthPrincipal, @@ -34,6 +35,8 @@ impl Entry { #[derive(Default)] struct State { generation: u64, + global_invalidation: u64, + session_invalidations: HashMap<[u8; 32], u64>, bytes: usize, entries: HashMap<[u8; 32], Entry>, } @@ -46,6 +49,14 @@ impl AuthenticationCache { Sha256::digest(token.as_bytes()).into() } + fn session_key(user_id: &str, session_id: &str) -> [u8; 32] { + let mut hash = Sha256::new(); + hash.update((user_id.len() as u64).to_be_bytes()); + hash.update(user_id.as_bytes()); + hash.update(session_id.as_bytes()); + hash.finalize().into() + } + pub(super) fn lookup(&self, key: &[u8; 32]) -> (Option, u64) { let mut state = self .0 @@ -80,10 +91,18 @@ impl AuthenticationCache { .unwrap_or_else(|poisoned| poisoned.into_inner()); let now = Instant::now(); let now_unix = chrono::Utc::now().timestamp(); - if state.generation != generation - || now_unix >= expires_at_unix - || now.duration_since(checked_at) >= TTL - { + if state.generation != generation { + let session_invalidated = principal.session_id.as_deref().is_some_and(|session| { + state + .session_invalidations + .get(&Self::session_key(&principal.user.user_id, session)) + .is_some_and(|last| *last > generation) + }); + if generation < state.global_invalidation || session_invalidated { + return false; + } + } + if now_unix >= expires_at_unix || now.duration_since(checked_at) >= TTL { return false; } let bytes = principal_bytes(principal); @@ -125,28 +144,71 @@ impl AuthenticationCache { true } - fn invalidate(&self) { + fn invalidate(&self, session: Option<(&str, &str)>) { let mut state = self .0 .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); state.generation = state.generation.wrapping_add(1); - state.entries.clear(); - state.bytes = 0; + if let Some((user_id, session_id)) = session { + let key = Self::session_key(user_id, session_id); + // Bound the fence metadata. On overflow, fence all older fills + // before discarding their session-specific invalidations. + if !state.session_invalidations.contains_key(&key) + && state.session_invalidations.len() >= MAX_ENTRIES + { + state.global_invalidation = state.generation; + state.session_invalidations.clear(); + } + let generation = state.generation; + state.session_invalidations.insert(key, generation); + state.entries.retain(|_, entry| { + entry.principal.user.user_id != user_id + || entry.principal.session_id.as_deref() != Some(session_id) + }); + state.bytes = state.entries.values().map(|entry| entry.bytes).sum(); + } else { + state.global_invalidation = state.generation; + state.session_invalidations.clear(); + state.entries.clear(); + state.bytes = 0; + } } pub(super) fn mutation(self: &Arc) -> InvalidateOnDrop { - self.invalidate(); - InvalidateOnDrop(Arc::clone(self)) + self.invalidate(None); + InvalidateOnDrop { + cache: Arc::clone(self), + session: None, + } + } + + pub(super) fn session_mutation( + self: &Arc, + user_id: &str, + session_id: &str, + ) -> InvalidateOnDrop { + self.invalidate(Some((user_id, session_id))); + InvalidateOnDrop { + cache: Arc::clone(self), + session: Some((user_id.to_owned(), session_id.to_owned())), + } } } -pub(super) struct InvalidateOnDrop(Arc); +pub(super) struct InvalidateOnDrop { + cache: Arc, + session: Option<(String, String)>, +} impl Drop for InvalidateOnDrop { fn drop(&mut self) { // Includes errors, cancellation, and unknown commit acknowledgements. - self.0.invalidate(); + self.cache.invalidate( + self.session + .as_ref() + .map(|(user, session)| (user.as_str(), session.as_str())), + ); } } @@ -216,6 +278,65 @@ mod tests { assert!(cache.lookup(&key).0.is_none()); } + #[test] + fn session_invalidation_preserves_other_sessions_and_fences_inflight_fills() { + let cache = Arc::new(AuthenticationCache::default()); + let first_key = AuthenticationCache::key("first"); + let second_key = AuthenticationCache::key("second"); + let other_key = AuthenticationCache::key("other"); + let other_fill = AuthenticationCache::key("other-cold-fill"); + let mut first = principal(); + first.session_id = Some("first-session".into()); + let mut second = first.clone(); + second.session_id = Some("second-session".into()); + let mut other = first.clone(); + other.user.user_id = "other-user".into(); + let now = Instant::now(); + let expires = chrono::Utc::now().timestamp() + 60; + for (key, value) in [ + (first_key, &first), + (second_key, &second), + (other_key, &other), + ] { + assert!(cache.publish(key, 0, now, expires, value)); + } + let mutation = cache.session_mutation("owner", "first-session"); + assert!(cache.lookup(&first_key).0.is_none()); + assert_eq!(cache.lookup(&second_key).0, Some(second.clone())); + assert_eq!(cache.lookup(&other_key).0, Some(other.clone())); + assert!(!cache.publish(first_key, 0, now, expires, &first)); + let (_, during) = cache.lookup(&first_key); + drop(mutation); + assert!(!cache.publish(first_key, during, now, expires, &first)); + for _ in 0..100 { + drop(cache.session_mutation("owner", "first-session")); + } + assert!( + cache.publish(other_fill, 0, now, expires, &other), + "a different session's in-flight database observation stays cacheable" + ); + assert_eq!(cache.lookup(&second_key).0, Some(second)); + assert_eq!(cache.lookup(&other_key).0, Some(other)); + } + + #[test] + fn session_fence_overflow_is_bounded_and_fails_closed() { + let cache = Arc::new(AuthenticationCache::default()); + let now = Instant::now(); + let expires = chrono::Utc::now().timestamp() + 60; + for index in 0..=MAX_ENTRIES { + drop(cache.session_mutation("owner", &format!("session-{index}"))); + } + let state = cache.0.lock().unwrap(); + assert!(state.session_invalidations.len() <= MAX_ENTRIES); + assert!(state.global_invalidation > 0); + drop(state); + let key = AuthenticationCache::key("unrelated"); + assert!(!cache.publish(key, 0, now, expires, &principal())); + let (_, current) = cache.lookup(&key); + assert!(cache.publish(key, current, now, expires, &principal())); + } + #[test] fn cache_retention_is_bounded_and_oversized_principals_stay_local() { for content_bytes in [0, MAX_ENTRY_BYTES / 2] { diff --git a/crates/services/src/auth/mod.rs b/crates/services/src/auth/mod.rs index 4e7a3b508e..2363822843 100644 --- a/crates/services/src/auth/mod.rs +++ b/crates/services/src/auth/mod.rs @@ -1799,33 +1799,26 @@ impl AuthService for DatabaseAuthService { .fetch_one(&mut *tx) .await .map_err(internal_error)?; - // A concurrent disconnect revokes this row under the same user lock; - // the compare-and-revoke prevents refresh from resurrecting it. - let revoked = query( - "UPDATE auth_refresh_tokens SET is_revoked = 1 WHERE token_hash = ? AND is_revoked = 0", + // Replace the current token under the account lock. The old hash no + // longer matches, so refresh and disconnect still have one winner. + let rotated = query( + "UPDATE auth_refresh_tokens SET token_hash = ?, expires_at = ?, created_at = NOW(6), token_prefix = NULL \ + WHERE token_hash = ? AND user_id = ? AND session_id = ? AND is_revoked = 0", ) + .bind(&new_refresh_token_hash) + .bind(expires_at) .bind(&refresh_token_hash) + .bind(&user_id) + .bind(&session_id) .execute(&mut *tx) .await - .map_err(|e| map_auth_sqlx(e, "refresh.revoke_old_token", Some(&pool)))?; - if revoked.rows_affected() != 1 { + .map_err(|e| map_auth_sqlx(e, "refresh.rotate_token", Some(&pool)))?; + if rotated.rows_affected() != 1 { return Err(error_response( StatusCode::UNAUTHORIZED, "Token expired or revoked", )); } - query( - "INSERT INTO auth_refresh_tokens (token_id, user_id, session_id, token_hash, expires_at, is_revoked) \ - VALUES (?, ?, ?, ?, ?, 0)", - ) - .bind(Uuid::new_v4().to_string()) - .bind(&user.user_id) - .bind(&session_id) - .bind(&new_refresh_token_hash) - .bind(expires_at) - .execute(&mut *tx) - .await - .map_err(|e| map_auth_sqlx(e, "refresh.insert_new_refresh_token", Some(&pool)))?; tx.commit() .await .map_err(|e| map_auth_sqlx(e, "refresh.commit_tx", Some(&pool)))?; @@ -1861,13 +1854,21 @@ impl AuthService for DatabaseAuthService { .clone() .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Invalid token"))?; let (session_id, _, _) = self.parse_token_session(&claims)?; + // JWT validation allows clock-skew leeway; logout proofs do not. + if claims.exp <= Utc::now().timestamp() { + return Err(error_response( + StatusCode::UNAUTHORIZED, + "Token expired or revoked", + )); + } let pool = self .get_pool() .await .map_err(|e| map_auth_sqlx(e, "auth.get_pool", None))?; - let refresh_token_hash = sha256_hex(&request.refresh_token); - let _invalidation = self.authentication_cache.mutation(); + let _invalidation = self + .authentication_cache + .session_mutation(&user_id, &session_id); let mut tx = pool .begin() .await @@ -1883,28 +1884,9 @@ impl AuthService for DatabaseAuthService { "Token expired or revoked", )); } - // A consumed token cannot refresh again, but can still terminate the - // session it originally owned. The shared user lock serializes this - // revocation with refresh's compare-and-revoke and successor insert. - let stored = query( - "SELECT user_id, DATE_FORMAT(expires_at, '%Y-%m-%dT%H:%i:%s') AS expires_at, session_id \ - FROM auth_refresh_tokens WHERE token_hash = ? LIMIT 1", - ) - .bind(&refresh_token_hash) - .fetch_optional(&mut *tx) - .await - .map_err(|e| map_auth_sqlx(e, "logout.fetch_refresh_token", Some(&pool)))? - .map(|row| decode_refresh_token_row(&row)) - .transpose() - .map_err(|e| map_auth_sqlx(e, "logout.decode_refresh_token", Some(&pool)))? - .ok_or_else(|| error_response(StatusCode::UNAUTHORIZED, "Token expired or revoked"))?; - if stored.0 != user_id || stored.2.as_deref() != Some(session_id.as_str()) { - return Err(error_response( - StatusCode::UNAUTHORIZED, - "Token session mismatch", - )); - } - if stored.1 < Utc::now().format("%Y-%m-%dT%H:%M:%S").to_string() { + // Recheck after waiting for the lock. The signed session claim is the + // proof even when a consumed token's old hash has been replaced. + if claims.exp <= Utc::now().timestamp() { return Err(error_response( StatusCode::UNAUTHORIZED, "Token expired or revoked", diff --git a/crates/services/src/storage.rs b/crates/services/src/storage.rs index 6264ab19f0..60a1fa7f7d 100644 --- a/crates/services/src/storage.rs +++ b/crates/services/src/storage.rs @@ -6867,7 +6867,7 @@ pub struct CleanupResult { /// Configuration for expiry policies of authentication and operational data. pub struct RetentionPolicy { - /// Minimum age in days for expired refresh token records (default: 7) + /// Max age in days for expired/revoked refresh tokens (default: 7) pub refresh_token_days: u32, /// Max age in days for inactive auth tokens (default: 30) pub auth_token_days: u32, @@ -6909,12 +6909,12 @@ pub async fn cleanup_expired_data( const AUTH_AUDIT_LOG_BATCH_LIMIT: u32 = 1000; let mut results = Vec::new(); - // 1. Expired refresh tokens past the minimum retention age + // 1. Expired + revoked refresh tokens let deleted = sqlx::query( "DELETE FROM auth_refresh_tokens \ - WHERE expires_at < NOW(6) \ + WHERE (expires_at < NOW(6) OR is_revoked = 1) \ AND created_at < DATE_SUB(NOW(6), INTERVAL ? DAY) \ - ORDER BY expires_at ASC, token_id ASC \ + ORDER BY created_at ASC, token_id ASC \ LIMIT ?", ) .bind(policy.refresh_token_days) diff --git a/crates/services/tests/auth_cache_db_it.rs b/crates/services/tests/auth_cache_db_it.rs index 1c3f011610..7e38561510 100644 --- a/crates/services/tests/auth_cache_db_it.rs +++ b/crates/services/tests/auth_cache_db_it.rs @@ -29,7 +29,6 @@ struct SnapshotBarrier { impl tracing_subscriber::Layer for Queries { fn on_event(&self, event: &tracing::Event<'_>, _: tracing_subscriber::layer::Context<'_, S>) { if event.metadata().target() == "sqlx::query" { - self.count.fetch_add(1, Ordering::SeqCst); #[derive(Default)] struct Fields(String); impl tracing::field::Visit for Fields { @@ -39,6 +38,11 @@ impl tracing_subscriber::Layer for Queries { } let mut fields = Fields::default(); event.record(&mut fields); + // Pool expansion can initialize a new connection alongside an + // auth query; count application statements, not session setup. + if !fields.0.contains("SET sql_mode=") { + self.count.fetch_add(1, Ordering::SeqCst); + } if fields.0.contains("SELECT u.user_id, u.username") && self.snapshot.armed.swap(false, Ordering::SeqCst) { @@ -175,6 +179,26 @@ async fn jwt_authentication_uses_one_cold_query_and_zero_warm_queries_with_local .current_principal(&second) .await .expect("logout does not revoke another device"); + assert_eq!( + count.swap(0, Ordering::SeqCst), + 0, + "another session stays warm" + ); + for _ in 0..3 { + service + .logout(AuthRefreshRequestData { + refresh_token: tokens[0].refresh_token.clone(), + }) + .await + .unwrap(); + count.store(0, Ordering::SeqCst); + service.current_principal(&second).await.unwrap(); + assert_eq!( + count.swap(0, Ordering::SeqCst), + 0, + "replayed logout must not flush unrelated sessions" + ); + } sqlx::query("UPDATE auth_users SET is_active = 0 WHERE user_id = ?") .bind(&user.user_id) @@ -246,8 +270,8 @@ async fn jwt_authentication_uses_one_cold_query_and_zero_warm_queries_with_local let subsequent = observer.current_principal(&second).await; assert_eq!( count.swap(0, Ordering::SeqCst), - 1, - "invalidated fill cannot re-enter cache" + usize::from(same_session), + "only the revoked session's in-flight fill must be fenced" ); if same_session { assert_eq!( diff --git a/crates/services/tests/auth_logout_rotation_db_it.rs b/crates/services/tests/auth_logout_rotation_db_it.rs index 2d8eb717b8..0040d0503f 100644 --- a/crates/services/tests/auth_logout_rotation_db_it.rs +++ b/crates/services/tests/auth_logout_rotation_db_it.rs @@ -6,6 +6,72 @@ use astra_services::{ DatabaseAuthService, RetentionPolicy, cleanup_expired_data, }; use axum::http::{HeaderMap, HeaderValue, StatusCode}; +use std::sync::{ + Arc, Condvar, Mutex, + atomic::{AtomicBool, Ordering}, +}; +use tracing_subscriber::prelude::*; + +#[derive(Default)] +struct AccountLockGate { + armed: AtomicBool, + acquired: tokio::sync::Notify, + released: Mutex, + wake: Condvar, +} + +impl AccountLockGate { + fn arm(&self) { + *self.released.lock().unwrap() = false; + self.armed.store(true, Ordering::SeqCst); + } + + fn release(&self) { + *self.released.lock().unwrap() = true; + self.wake.notify_all(); + } +} + +struct PauseAccountLock(Arc); + +impl tracing_subscriber::Layer for PauseAccountLock { + fn on_event(&self, event: &tracing::Event<'_>, _: tracing_subscriber::layer::Context<'_, S>) { + if event.metadata().target() != "sqlx::query" || !self.0.armed.load(Ordering::SeqCst) { + return; + } + #[derive(Default)] + struct Fields(String); + impl tracing::field::Visit for Fields { + fn record_debug(&mut self, _: &tracing::field::Field, value: &dyn std::fmt::Debug) { + self.0.push_str(&format!("{value:?}")); + } + } + let mut fields = Fields::default(); + event.record(&mut fields); + if fields + .0 + .contains("SELECT user_id FROM auth_users WHERE user_id = ? FOR UPDATE") + && self.0.armed.swap(false, Ordering::SeqCst) + { + self.0.acquired.notify_one(); + tokio::task::block_in_place(|| { + let (released, timeout) = self + .0 + .wake + .wait_timeout_while( + self.0.released.lock().unwrap(), + std::time::Duration::from_secs(20), + |released| !*released, + ) + .unwrap(); + assert!( + *released && !timeout.timed_out(), + "account lock barrier timed out" + ); + }); + } + } +} fn headers(token: &str) -> HeaderMap { let mut headers = HeaderMap::new(); @@ -19,6 +85,11 @@ fn headers(token: &str) -> HeaderMap { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[ignore = "requires live MatrixOne: ASTRA_TEST_DB_IT=1"] async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_other_sessions() { + let gate = Arc::new(AccountLockGate::default()); + tracing::subscriber::set_global_default( + tracing_subscriber::registry().with(PauseAccountLock(gate.clone())), + ) + .expect("this test binary owns its tracing subscriber"); let (pool, settings) = common::setup_pool_and_settings().await; let service = DatabaseAuthService::new( settings, @@ -53,7 +124,27 @@ async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_oth .await .unwrap(); let other = service.login(login()).await.unwrap(); - for invalid in [original.access_token.clone(), "invalid.token".to_string()] { + let key = b"isolated-logout-rotation-test-key"; + let mut expired_claims = jsonwebtoken::decode::( + &original.refresh_token, + &jsonwebtoken::DecodingKey::from_secret(key), + &jsonwebtoken::Validation::default(), + ) + .unwrap() + .claims; + // This is inside the JWT library's usual expiry leeway. + expired_claims["exp"] = (chrono::Utc::now().timestamp() - 1).into(); + let expired = jsonwebtoken::encode( + &jsonwebtoken::Header::default(), + &expired_claims, + &jsonwebtoken::EncodingKey::from_secret(key), + ) + .unwrap(); + for invalid in [ + original.access_token.clone(), + "invalid.token".to_string(), + expired, + ] { assert_eq!( service .logout(AuthRefreshRequestData { @@ -65,7 +156,7 @@ async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_oth StatusCode::UNAUTHORIZED ); } - let rotated = service + let mut rotated = service .refresh(AuthRefreshRequestData { refresh_token: original.refresh_token.clone(), }) @@ -81,7 +172,16 @@ async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_oth .0, StatusCode::UNAUTHORIZED ); - // The old token's revoked row may be removed even though its JWT is still valid. + // Repeated rotations replace the same session row rather than retaining a + // historical row for every consumed token. + for _ in 0..3 { + rotated = service + .refresh(AuthRefreshRequestData { + refresh_token: rotated.refresh_token, + }) + .await + .unwrap(); + } let cleanup = cleanup_expired_data(pool.get(), &RetentionPolicy::default()) .await .unwrap(); @@ -90,25 +190,22 @@ async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_oth .iter() .any(|result| result.table == "auth_refresh_tokens") ); - let retained: i64 = sqlx::query_scalar( - "SELECT COUNT(*) FROM auth_refresh_tokens WHERE user_id = ? AND is_revoked = 1", - ) - .bind(&user.user_id) - .fetch_one(pool.get()) - .await - .unwrap(); - assert_eq!( - retained, 1, - "cleanup must retain an unexpired rotated token" - ); + let retained: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM auth_refresh_tokens WHERE user_id = ?") + .bind(&user.user_id) + .fetch_one(pool.get()) + .await + .unwrap(); + assert_eq!(retained, 2, "two sessions retain only one refresh row each"); service .logout(AuthRefreshRequestData { refresh_token: original.refresh_token.clone(), }) .await .unwrap(); - // Once the expired record is beyond retention, cleanup must remove it. - sqlx::query("UPDATE auth_refresh_tokens SET expires_at = DATE_SUB(NOW(6), INTERVAL 8 DAY) WHERE user_id = ? AND is_revoked = 1 AND created_at < DATE_SUB(NOW(6), INTERVAL 7 DAY)") + // An old signed proof still revokes the current row even though its own + // hash disappeared on rotation. Revoked rows are eligible after retention. + sqlx::query("UPDATE auth_refresh_tokens SET created_at = DATE_SUB(NOW(6), INTERVAL 8 DAY) WHERE user_id = ? AND is_revoked = 1") .bind(&user.user_id) .execute(pool.get()) .await @@ -144,41 +241,115 @@ async fn logout_revokes_the_session_across_refresh_rotation_without_revoking_oth .await .expect("another session remains active"); - // Exercise simultaneous refresh and logout with three independently issued sessions. - // Whichever transaction wins, no successor from that session may stay usable. - for _ in 0..3 { + let concurrent = service.login(login()).await.unwrap(); + let (left, right) = tokio::join!( + service.refresh(AuthRefreshRequestData { + refresh_token: concurrent.refresh_token.clone() + }), + service.refresh(AuthRefreshRequestData { + refresh_token: concurrent.refresh_token.clone() + }), + ); + let next = match (left, right) { + (Ok(next), Err(error)) | (Err(error), Ok(next)) => { + assert_eq!(error.0, StatusCode::UNAUTHORIZED); + next + } + _ => panic!("a single-use refresh token must have exactly one winner"), + }; + service + .logout(AuthRefreshRequestData { + refresh_token: concurrent.refresh_token, + }) + .await + .unwrap(); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: next.refresh_token + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + + // Pause the winner after it acquires the real account lock so both orders + // are exercised without depending on task scheduling. + for refresh_first in [true, false] { let token = service.login(login()).await.unwrap(); - let (refreshed, logged_out) = tokio::join!( - service.refresh(AuthRefreshRequestData { - refresh_token: token.refresh_token.clone() - }), - service.logout(AuthRefreshRequestData { - refresh_token: token.refresh_token.clone() - }), - ); - logged_out.unwrap(); - match refreshed { - Ok(next) => { - assert_eq!( - service - .refresh(AuthRefreshRequestData { - refresh_token: next.refresh_token - }) - .await - .unwrap_err() - .0, - StatusCode::UNAUTHORIZED - ); - assert_eq!( - service - .current_principal(&headers(&next.access_token)) - .await - .unwrap_err() - .0, - StatusCode::UNAUTHORIZED - ); + gate.arm(); + let winner_service = service.clone(); + let winner_token = token.refresh_token.clone(); + let winner = tokio::spawn(async move { + if refresh_first { + winner_service + .refresh(AuthRefreshRequestData { + refresh_token: winner_token, + }) + .await + .map(Some) + } else { + winner_service + .logout(AuthRefreshRequestData { + refresh_token: winner_token, + }) + .await + .map(|()| None) } - Err(error) => assert_eq!(error.0, StatusCode::UNAUTHORIZED), + }); + tokio::time::timeout(std::time::Duration::from_secs(15), gate.acquired.notified()) + .await + .unwrap(); + let other_service = service.clone(); + let other_token = token.refresh_token.clone(); + let mut loser = tokio::spawn(async move { + if refresh_first { + other_service + .logout(AuthRefreshRequestData { + refresh_token: other_token, + }) + .await + .map(|()| None) + } else { + other_service + .refresh(AuthRefreshRequestData { + refresh_token: other_token, + }) + .await + .map(Some) + } + }); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), &mut loser) + .await + .is_err() + ); + gate.release(); + let first = winner.await.unwrap().unwrap(); + let second = loser.await.unwrap(); + if let Some(next) = first { + second.unwrap(); + assert_eq!( + service + .refresh(AuthRefreshRequestData { + refresh_token: next.refresh_token + }) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + service + .current_principal(&headers(&next.access_token)) + .await + .unwrap_err() + .0, + StatusCode::UNAUTHORIZED + ); + } else { + assert_eq!(second.unwrap_err().0, StatusCode::UNAUTHORIZED); } service .current_principal(&headers(&other.access_token)) diff --git a/docs/design/authentication.md b/docs/design/authentication.md index 9a20093f64..5c17b9f988 100644 --- a/docs/design/authentication.md +++ b/docs/design/authentication.md @@ -55,9 +55,11 @@ is retained in the cache. The cache is shared by service clones and bounded to 1,024 entries, 4 MiB of accounted storage and 16 KiB per entry. Reconfiguring a service's pool, issuer or -encryptor starts a new namespace. Logout, disconnect and Memoria binding changes -invalidate local observations and fence in-flight fills, including uncertain -commits. Other processes and out-of-band account/session mutations can take up +encryptor starts a new namespace. Logout invalidates only its session's local +observations and fences matching in-flight fills; account-wide changes still +invalidate all observations, including on uncertain commits. Session fences are +bounded within the cache; overflow conservatively fences older in-flight fills. +Other processes and out-of-band account/session mutations can take up to 30 seconds to affect ordinary request admission. Expired observations require a successful fresh lookup; database errors do not enable stale fallback. @@ -67,12 +69,25 @@ unrelated principal. Absolute token, observed Session and Memoria access expiry are checked again after awaited validation. A request starting after a local logout cannot reuse an invalidated observation. -Refresh and logout serialize on the account lock. Logout validates the -recorded token and revokes its entire session, so a concurrent rotation -cannot leave a usable successor. A previously consumed, unexpired token may -terminate its own session but never regain refresh authority. Refresh-token -records remain until they are expired and older than the configured minimum -retention age, so cleanup cannot remove this logout proof early. +Refresh and logout serialize on the account lock. Refresh atomically replaces +the active token hash in its existing session row; the consumed hash cannot +refresh again. Logout uses a signed, strictly unexpired refresh token's user and +session claims to revoke that session, even after the old hash is replaced or +cleaned up. A valid proof with no remaining session is an idempotent success; +it cannot revoke another session. This relies on unique session IDs and trusted +JWT signing keys: removing an old hash alone no longer revokes its logout proof. +Expired or revoked rows remain eligible for the existing bounded cleanup. +Do not mix this version with older servers: an old server may report logout +success after a new server has already replaced the hash, leaving the session +active. For upgrade or rollback, stop admission to all API replicas, drain or +terminate their in-flight requests, and wait until every old process has exited +before starting any replica of the other version. This cutover interrupts +service; ordinary rolling updates and session affinity do not provide the +required ordering. Keep all API replicas on the same side of the cutover. +After rollback, an older server cannot use a token already consumed by the +newer version to revoke its successor; it rejects that logout. Draining +requests does not restore replaced hashes, so clients must use their current +refresh tokens or the affected sessions must be revoked separately. This cache observes identity, not operation permissions. Memory operations continue resolving scoped credentials through the canonical credential resolver; diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 849443d8a7..368db7659a 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -11,9 +11,9 @@ are `/live`, `/ready`, `/health`, `/auth/register`, `/auth/login`, `GET /auth/me `/auth/logout` (refresh/logout authenticate the supplied refresh token): Logout revokes the token's session, including a successor issued by a -concurrent refresh. A signed, unexpired refresh token recorded for that -session can terminate it even after rotation; it cannot be used to refresh -again. Other sessions are unaffected. +concurrent refresh. A signed, unexpired refresh token can terminate its own +session even after rotation; it cannot be used to refresh again. A valid proof +for a session already gone is an idempotent success. Other sessions are unaffected. ``` Authorization: Bearer diff --git a/packages/sdk/README.md b/packages/sdk/README.md index ad9b6e65fa..d385bf1160 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -171,9 +171,20 @@ overwrite a newer authentication decision. This coordination is per client instance, not across separate BFF requests. `onTokenRefresh` may be asynchronous and call the same client. Requests already -awaiting its rotation still wait for the callback. Callback failure preserves -the rotated credentials and returns those callers' original authentication -error. Notifications from successive rotations may overlap. +awaiting its rotation still wait for the callback, so the callback must not +await a request that has already joined that rotation. A delayed old-token 401 +returns its original error while that notification is pending or has failed +without a recovery in progress, so it cannot use credentials whose persistence +is uncertain. Callback failure +preserves the rotated credentials in memory. Requests made during notification +can use the current token, but a 401 does not start another refresh until the +notification finishes. This prevents overlapping persistence callbacks from +storing a superseded refresh token. A delayed 401 joins a newer in-flight +rotation when no notification is active. +If credentials can be replaced while a callback writes to external storage, +that storage must use an atomic conditional write or another ordering mechanism +to prevent an older callback from overwriting newer credentials. Checking +`getTokens()` before an asynchronous write alone is not sufficient. The Web runtime client persists a captured pair only if it still matches that instance's current credentials after asynchronous cookie acquisition. diff --git a/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts index fa83c6a7fb..14ba17e8cd 100644 --- a/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts +++ b/packages/sdk/src/__tests__/auth-lifecycle-regression.test.ts @@ -186,9 +186,11 @@ test('aborted waiters release their registrations even when refresh never settle await expect(request).rejects.toBe(reason); expect(removed).toHaveBeenCalledWith('abort', expect.any(Function)); } - const attempt = (client as unknown as { refreshAttempt: { waiters: Set } }).refreshAttempt; - expect(attempt.waiters.size).toBe(0); + const remaining = client.fetchResponse('/remaining'); + await vi.waitFor(() => expect(vi.mocked(fetch)).toHaveBeenCalledTimes(22)); refresh.resolve(Response.json(tokens)); + expect((await remaining).status).toBe(401); + expect(vi.mocked(fetch).mock.calls.filter(([url]) => String(url).endsWith('/auth/refresh'))).toHaveLength(1); }); test('closing a stream while refresh is pending prevents a second POST', async () => { diff --git a/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts index d65febe23e..a5f4bee567 100644 --- a/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts +++ b/packages/sdk/src/__tests__/token-refresh-concurrency.test.ts @@ -12,6 +12,75 @@ const authorization = (init?: RequestInit) => new Headers(init?.headers).get('Au afterEach(() => vi.unstubAllGlobals()); +test('an asynchronous reentrant notification cannot rotate forever on persistent 401', async () => { + let refreshes = 0; + let recovered = false; + vi.stubGlobal('fetch', vi.fn(async (url: string) => { + if (url.endsWith('/auth/refresh')) { + refreshes++; + // Keep a broken implementation bounded during the regression test. + return refreshes > 4 ? new Response('stop', { status: 503 }) + : Response.json({ access_token: `access-${refreshes}`, refresh_token: `refresh-${refreshes}` }); + } + return recovered ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { await Promise.resolve(); await client.getMe(); }, + }); + await expect(client.getMe()).rejects.toMatchObject({ status: 401 }); + expect(refreshes).toBe(1); + recovered = true; + await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); +}); + +test.each(['pending', 'failed'] as const)('a delayed old-token 401 cannot bypass a %s notification', async (when) => { + const late401 = deferred(); + const notification = deferred(); + const entered = deferred(); + const calls: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(rotated); + calls.push(`${url.split('/').pop()}:${authorization(init)}`); + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late401.promise; + return authorization(init) === 'Bearer old-access' ? unauthorized() : new Response('ok'); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { entered.resolve(); await notification.promise; throw new Error('persistence failed'); }, + }); + const late = client.fetchResponse('/late'); + const first = client.fetchResponse('/first'); + await entered.promise; + if (when === 'failed') notification.resolve(); + if (when === 'failed') expect((await first).status).toBe(401); + late401.resolve(unauthorized()); + expect((await late).status).toBe(401); + expect(calls.filter(call => call.startsWith('late:'))).toEqual(['late:Bearer old-access']); + if (when === 'pending') notification.resolve(); + expect((await first).status).toBe(401); +}); + +test('a callback awaiting a delayed old-token response does not deadlock', async () => { + const late401 = deferred(); + const entered = deferred(); + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json(rotated); + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late401.promise; + return authorization(init) === 'Bearer old-access' ? unauthorized() : new Response('ok'); + })); + let late!: Promise; + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { entered.resolve(); expect((await late).status).toBe(401); }, + }); + late = client.fetchResponse('/late'); + const first = client.fetchResponse('/first'); + await entered.promise; + late401.resolve(unauthorized()); + expect((await first).status).toBe(200); +}); + test('REST and buffered SSE share one refresh of a single-use token', async () => { const refresh = deferred(); const onTokenRefresh = vi.fn(); @@ -42,6 +111,38 @@ test('REST and buffered SSE share one refresh of a single-use token', async () = expect(onTokenRefresh).toHaveBeenCalledTimes(1); }); +test('refresh-token-only clients resume a session with one shared rotation', async () => { + const refresh = deferred(); + let rotations = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { rotations++; return refresh.promise; } + return authorization(init) === 'Bearer new-access' + ? Response.json({ user_id: 'alice' }) : unauthorized(); + })); + const client = new AstraClient({ baseUrl: 'https://example.test', refreshToken: 'old-refresh' }); + const results = Promise.all([client.getMe(), client.getMe()]); + await vi.waitFor(() => expect(rotations).toBe(1)); + refresh.resolve(Response.json(rotated)); + expect(await results).toEqual([{ user_id: 'alice' }, { user_id: 'alice' }]); + expect(rotations).toBe(1); +}); + +test.each(['anonymous', 'request authorization', 'configured authorization'] as const)( + 'a refresh-token-only client does not refresh %s requests', async (mode) => { + const fetch = vi.fn(async () => unauthorized()); + vi.stubGlobal('fetch', fetch); + const client = new AstraClient({ + baseUrl: 'https://example.test', refreshToken: 'old-refresh', + headers: mode === 'configured authorization' ? { Authorization: 'Bearer external' } : undefined, + }); + const response = await client.fetchResponse('/resource', + mode === 'request authorization' ? { headers: { Authorization: 'Bearer external' } } : undefined, + mode === 'anonymous' ? { auth: 'none' } : undefined); + expect(response.status).toBe(401); + expect(fetch).toHaveBeenCalledTimes(1); + }, +); + test('a late 401 reuses the token already rotated by another request', async () => { const late = deferred(); let initialRequests = 0; @@ -66,6 +167,100 @@ test('a late 401 reuses the token already rotated by another request', async () expect(refreshCount).toBe(1); }); +test('a late 401 joins a newer refresh instead of retrying its rejected token', async () => { + const late = deferred(); + const secondRefresh = deferred(); + const secondStarted = deferred(); + const lateTokens: string[] = []; + let refreshCount = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + if (++refreshCount === 1) return Response.json(rotated); + secondStarted.resolve(); + return secondRefresh.promise; + } + const token = authorization(init); + if (url.endsWith('/late')) { + lateTokens.push(token ?? ''); + if (token === 'Bearer old-access') return late.promise; + } + return token === 'Bearer final-access' || (url.endsWith('/first') && token === 'Bearer new-access') + ? new Response('ok') : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + }); + const delayed = client.fetchResponse('/late'); + expect((await client.fetchResponse('/first')).status).toBe(200); + const next = client.fetchResponse('/next'); + await secondStarted.promise; + late.resolve(unauthorized()); + await new Promise(resolve => setTimeout(resolve, 0)); + secondRefresh.resolve(Response.json({ access_token: 'final-access', refresh_token: 'final-refresh' })); + expect((await delayed).status).toBe(200); + expect((await next).status).toBe(200); + expect(lateTokens).toEqual(['Bearer old-access', 'Bearer final-access']); + expect(refreshCount).toBe(2); +}); + +test('a delayed 401 can use a later successfully persisted recovery', async () => { + const late = deferred(); + let rotations = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + rotations++; + return Response.json(rotations === 1 ? rotated + : { access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); + } + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late.promise; + return authorization(init) === 'Bearer recovered-access' ? new Response('ok') : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { if (rotations === 1) throw new Error('persistence failed'); }, + }); + const delayed = client.fetchResponse('/late'); + expect((await client.fetchResponse('/first')).status).toBe(401); + expect((await client.fetchResponse('/recovery')).status).toBe(200); + late.resolve(unauthorized()); + expect((await delayed).status).toBe(200); + expect(rotations).toBe(2); +}); + +test.each(['success', 'failure'] as const)( + 'a delayed 401 joins a %s recovery after an earlier notification failure', async (result) => { + const late = deferred(); + const recovery = deferred(); + const recoveryStarted = deferred(); + let rotations = 0; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + if (++rotations === 1) return Response.json(rotated); + recoveryStarted.resolve(); + return recovery.promise; + } + if (url.endsWith('/late') && authorization(init) === 'Bearer old-access') return late.promise; + return authorization(init) === 'Bearer recovered-access' ? new Response('ok') : unauthorized(); + })); + const client = new AstraClient({ + baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', + onTokenRefresh: async () => { if (rotations === 1) throw new Error('persistence failed'); }, + }); + const delayed = client.fetchResponse('/late'); + expect((await client.fetchResponse('/first')).status).toBe(401); + const recovering = client.fetchResponse('/recovery'); + await recoveryStarted.promise; + late.resolve(unauthorized()); + await new Promise(resolve => setTimeout(resolve, 0)); + recovery.resolve(result === 'success' + ? Response.json({ access_token: 'recovered-access', refresh_token: 'recovered-refresh' }) + : new Response('unavailable', { status: 503 })); + expect((await recovering).status).toBe(result === 'success' ? 200 : 401); + expect((await delayed).status).toBe(result === 'success' ? 200 : 401); + expect(rotations).toBe(2); + }, +); + test.each(['http', 'network'] as const)('failed %s refresh is shared but a later attempt can recover', async (failure) => { const refresh = deferred(); let refreshCount = 0; @@ -133,7 +328,6 @@ test('an async notification can receive 401 without blocking its callers or late vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { if (url.endsWith('/auth/refresh')) { refreshTokens.push(JSON.parse(String(init?.body)).refresh_token); - if (refreshTokens.length === 2) return Response.json({}, { status: 503 }); return Response.json(refreshTokens.length === 1 ? rotated : { access_token: 'recovered-access', refresh_token: 'recovered-refresh' }); } @@ -147,14 +341,14 @@ test('an async notification can receive 401 without blocking its callers or late }); const results = await Promise.allSettled([client.getMe(), client.getRunEvents('run-1')]); expect(results.map(result => result.status)).toEqual(['rejected', 'rejected']); - expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); + expect(refreshTokens).toEqual(['old-refresh']); expect(onTokenRefresh).toHaveBeenCalledTimes(1); await expect(client.getMe()).resolves.toEqual({ user_id: 'alice' }); - expect(refreshTokens).toEqual(['old-refresh', 'new-refresh', 'new-refresh']); + expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); expect(onTokenRefresh).toHaveBeenCalledTimes(2); }); -test('a successful nested refresh still permits only one replay of the buffered SSE request', async () => { +test('a callback 401 does not start an overlapping rotation', async () => { const refreshTokens: string[] = []; const callbackResults: unknown[] = []; const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { @@ -167,18 +361,18 @@ test('a successful nested refresh still permits only one replay of the buffered ? Response.json({ user_id: 'alice' }) : unauthorized(); }); vi.stubGlobal('fetch', fetchMock); - const onTokenRefresh = vi.fn(async () => { callbackResults.push(await client.getMe()); }); + const onTokenRefresh = vi.fn(async () => { callbackResults.push(await client.getMe().catch(error => error.status)); }); const client = new AstraClient({ baseUrl: 'https://example.test', accessToken: 'old-access', refreshToken: 'old-refresh', onTokenRefresh, }); await expect(client.getRunEvents('run-1')).rejects.toMatchObject({ status: 401 }); - expect(refreshTokens).toEqual(['old-refresh', 'new-refresh']); - expect(callbackResults).toEqual([{ user_id: 'alice' }, { user_id: 'alice' }]); - expect(onTokenRefresh).toHaveBeenCalledTimes(2); + expect(refreshTokens).toEqual(['old-refresh']); + expect(callbackResults).toEqual([401]); + expect(onTokenRefresh).toHaveBeenCalledTimes(1); const streamRequests = fetchMock.mock.calls.filter(([url]) => url.includes('/stream')); expect(streamRequests.map(([, init]) => authorization(init))).toEqual([ - 'Bearer old-access', 'Bearer recovered-access', + 'Bearer old-access', 'Bearer new-access', ]); expect(streamRequests.every(([, init]) => new Headers(init?.headers).get('Accept') === 'text/event-stream')).toBe(true); }); diff --git a/packages/sdk/src/client.ts b/packages/sdk/src/client.ts index ed9d3a42b3..5efeb6f6e6 100644 --- a/packages/sdk/src/client.ts +++ b/packages/sdk/src/client.ts @@ -622,18 +622,32 @@ export function chatRequestToWire(req: ChatRequest): Record { * same layout as `astra-thin-client` / `astra-server`; set `pathPrefix` when a * gateway mounts the API under a prefix. */ +type RefreshAttempt = { + phase: "refreshing" | "notifying" | "done"; + result: boolean; + waiters: Set<(result: boolean) => void>; +}; + +type TokenRevision = { + accessToken: string | null; + refreshToken: string | null; + notification: "pending" | "succeeded" | "failed"; + attempt: RefreshAttempt | null; +}; + +function tokenRevision(accessToken: string | null, refreshToken: string | null): TokenRevision { + return { accessToken, refreshToken, notification: "succeeded", attempt: null }; +} + export class AstraClient { private config: AstraClientConfig; - private accessToken: string | null; - private refreshTokenValue: string | null; + private tokens: TokenRevision; private authGeneration = 0; private authAttempt: object | null = null; - private refreshAttempt: { generation: number; waiters: Set<(result: boolean) => void> } | null = null; constructor(config: AstraClientConfig) { this.config = config; - this.accessToken = config.accessToken ?? null; - this.refreshTokenValue = config.refreshToken ?? null; + this.tokens = tokenRevision(config.accessToken ?? null, config.refreshToken ?? null); } private apiPath(path: string): string { @@ -689,9 +703,8 @@ export class AstraClient { async logout(): Promise { this.authGeneration++; this.authAttempt = null; - const refreshToken = this.refreshTokenValue; - this.accessToken = null; - this.refreshTokenValue = null; + const refreshToken = this.tokens.refreshToken; + this.tokens = tokenRevision(null, null); if (refreshToken) { await this.request(PATH_AUTH_LOGOUT, { method: "POST", @@ -761,35 +774,52 @@ export class AstraClient { setTokens(accessToken: string, refreshToken?: string): void { this.authGeneration++; this.authAttempt = null; - this.accessToken = accessToken; - this.refreshTokenValue = refreshToken || null; + this.tokens = tokenRevision(accessToken, refreshToken || null); } /** A detached snapshot of the credentials currently owned by this client. */ getTokens(): Readonly<{ accessToken: string | null; refreshToken: string | null }> { - return { accessToken: this.accessToken, refreshToken: this.refreshTokenValue }; + return { accessToken: this.tokens.accessToken, refreshToken: this.tokens.refreshToken }; } - private tryRefreshToken(generation: number, rejectedToken: string | null, signal?: AbortSignal | null): Promise { + private tryRefreshToken(generation: number, revision: TokenRevision, signal?: AbortSignal | null): Promise { if (signal?.aborted) return Promise.reject(signal.reason); if (generation !== this.authGeneration) return Promise.resolve(false); - let attempt = this.refreshAttempt; - if (attempt?.generation !== generation) { - // A delayed 401 may belong to a token another request already rotated. - if (rejectedToken !== this.accessToken) return Promise.resolve(true); - if (!this.refreshTokenValue) return Promise.resolve(false); - attempt = { generation, waiters: new Set() }; - this.refreshAttempt = attempt; + // Do not overlap notifications: an older async persistence callback could + // otherwise overwrite the successor refresh token after rotation. + if (this.tokens.notification === "pending") return Promise.resolve(false); + let attempt = revision.attempt; + if (attempt?.phase === "notifying") return Promise.resolve(false); + if (attempt?.phase === "done") { + const newerAttempt = this.tokens.attempt; + if (newerAttempt?.phase === "refreshing") { + // Join the newer rotation instead of retrying a token that another + // request has already rejected. + attempt = newerAttempt; + } else { + return Promise.resolve( + (attempt.result || revision !== this.tokens) && this.tokens.notification === "succeeded", + ); + } + } + if (!attempt) { + if (revision !== this.tokens || !revision.refreshToken) return Promise.resolve(false); + attempt = { phase: "refreshing", result: false, waiters: new Set() }; + revision.attempt = attempt; const current = attempt; - // Observe the operation once; cancelled waiters can be removed even if it hangs. - void this.refreshTokens(generation, this.refreshTokenValue).then( + const finish = (result: boolean) => { + current.phase = "done"; + current.result = result; + // Failed network/HTTP refreshes may be tried again later. + if (revision === this.tokens) revision.attempt = null; + for (const settle of current.waiters) settle(result); + }; + void this.refreshTokens(generation, revision, current).then( result => { - if (this.refreshAttempt === current) this.refreshAttempt = null; - for (const settle of current.waiters) settle(result); + finish(result); }, () => { - if (this.refreshAttempt === current) this.refreshAttempt = null; - for (const settle of current.waiters) settle(false); + finish(false); }, ); } @@ -807,12 +837,12 @@ export class AstraClient { }); } - private async refreshTokens(generation: number, refreshToken: string): Promise { + private async refreshTokens(generation: number, revision: TokenRevision, attempt: RefreshAttempt): Promise { try { const res = await fetch(this.apiPath(PATH_AUTH_REFRESH), { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ refresh_token: refreshToken }), + body: JSON.stringify({ refresh_token: revision.refreshToken }), cache: "no-store", }); if (!res.ok) return false; @@ -823,15 +853,20 @@ export class AstraClient { typeof data.access_token !== "string" || typeof data.refresh_token !== "string" ) return false; - this.accessToken = data.access_token; - this.refreshTokenValue = data.refresh_token; - // Notification may reenter the client and need another refresh. - this.refreshAttempt = null; - await this.config.onTokenRefresh?.({ - accessToken: data.access_token, - refreshToken: data.refresh_token, - }); - return generation === this.authGeneration; + const next = tokenRevision(data.access_token, data.refresh_token); + this.tokens = next; + if (this.config.onTokenRefresh) { + next.notification = "pending"; + attempt.phase = "notifying"; + try { + await this.config.onTokenRefresh({ accessToken: data.access_token, refreshToken: data.refresh_token }); + next.notification = "succeeded"; + } catch { + next.notification = "failed"; + return false; + } + } + return generation === this.authGeneration && this.tokens.notification === "succeeded"; } catch { return false; } @@ -840,7 +875,7 @@ export class AstraClient { private buildHeaders(auth: "optional" | "none"): Record { const headers = new Headers(headersInitToRecord(this.config.headers ?? {})); if (auth === "none") headers.delete("Authorization"); - else if (this.accessToken) headers.set("Authorization", `Bearer ${this.accessToken}`); + else if (this.tokens.accessToken) headers.set("Authorization", `Bearer ${this.tokens.accessToken}`); return headersInitToRecord({}, headers); } @@ -861,9 +896,11 @@ export class AstraClient { const auth = options?.auth ?? "optional"; const url = this.apiPath(path); const generation = this.authGeneration; - const accessToken = this.accessToken; + const revision = this.tokens; + const accessToken = revision.accessToken; const requestHeaders = headersInitToRecord({}, init?.headers); - const explicitAuthorization = new Headers(requestHeaders).has("Authorization"); + const explicitAuthorization = new Headers(requestHeaders).has("Authorization") + || (!accessToken && new Headers(headersInitToRecord(this.config.headers ?? {})).has("Authorization")); const method = (init?.method ?? "GET").toUpperCase(); const replayable = !(typeof ReadableStream !== "undefined" && init?.body instanceof ReadableStream); const mayRetry = options?.retryUnauthorized ?? (method === "GET" || method === "HEAD"); @@ -872,10 +909,10 @@ export class AstraClient { headers: headersInitToRecord(this.buildHeaders(auth), requestHeaders), }); let res = await send(); - if (res.status === 401 && auth !== "none" && accessToken && !explicitAuthorization && mayRetry && replayable && !init?.signal?.aborted) { - const refreshed = await this.tryRefreshToken(generation, accessToken, init?.signal); + if (res.status === 401 && auth !== "none" && (accessToken || revision.refreshToken) && !explicitAuthorization && mayRetry && replayable && !init?.signal?.aborted) { + const refreshed = await this.tryRefreshToken(generation, revision, init?.signal); if (init?.signal?.aborted) throw init.signal.reason; - if (refreshed && generation === this.authGeneration) res = await send(); + if (refreshed && generation === this.authGeneration && this.tokens.notification === "succeeded") res = await send(); } return res; } diff --git a/scripts/schema/schema_inventory.py b/scripts/schema/schema_inventory.py index de16f64dde..137fdc094f 100755 --- a/scripts/schema/schema_inventory.py +++ b/scripts/schema/schema_inventory.py @@ -886,7 +886,7 @@ class AutoIncrementMetadata: "auth_refresh_tokens": TableMetadata( semantic_owner="astra_services::auth", state_class="durable refresh-session credential fact", - primary_query="refresh/logout lookup by token_hash, session_id, user_id, is_revoked, and expires_at", + primary_query="refresh compare-and-replace by token_hash and session; logout revokes by signed user/session claims", retention_policy="cleanup_expired_data prunes expired or revoked rows after refresh_token_days default 7 in ordered bounded batches; logout and refresh revoke specific token hashes", rebuildability="not rebuildable after token_hash/session binding is lost because refresh/logout idempotency and external session linkage depend on it", merge_guidance="keep separate from auth_users; refresh credentials are high-churn secrets with TTL/revocation lifecycle", diff --git a/web/__tests__/lib/runtime-client-refresh.test.ts b/web/__tests__/lib/runtime-client-refresh.test.ts index ba35a53ef0..c8b5d91fa4 100644 --- a/web/__tests__/lib/runtime-client-refresh.test.ts +++ b/web/__tests__/lib/runtime-client-refresh.test.ts @@ -46,18 +46,83 @@ it("persists refreshed httpOnly cookies before the initiating SDK request comple await expect(pending).resolves.toEqual({ user_id: "alice" }); const options = { httpOnly: true, sameSite: "lax", path: "/" }; expect(set.mock.calls).toEqual([ - [ACCESS_TOKEN_COOKIE, "new-access", options], [REFRESH_TOKEN_COOKIE, "new-refresh", options], + [ACCESS_TOKEN_COOKIE, "new-access", options], ]); expect(fetchMock.mock.calls.map(([, init]) => new Headers(init?.headers).get("Authorization"))).toEqual([ "Bearer old-access", null, "Bearer new-access", "Bearer new-access", ]); }); +it("does not replay a delayed old-token request while cookie persistence fails", async () => { + let release401!: (response: Response) => void; + let releaseCookies!: () => void; + const late401 = new Promise(resolve => { release401 = resolve; }); + const cookiesReady = new Promise(resolve => { releaseCookies = resolve; }); + vi.mocked(cookies).mockImplementation(async () => { + await cookiesReady; + return { set: () => { throw new Error("cookie write unavailable"); } } as unknown as Awaited>; + }); + const expired = () => new Response("expired", { status: 401 }); + const fetchMock = vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) return Response.json({ access_token: 'new', refresh_token: 'rotated' }); + const auth = new Headers(init?.headers).get('Authorization'); + if (url.endsWith('/late') && auth === 'Bearer old') return late401; + return auth === 'Bearer old' ? expired() : new Response('ok'); + }); + vi.stubGlobal('fetch', fetchMock); + const runtime = new WebRuntimeClient(runtimeConfig('old', 'old-refresh')); + const late = runtime.fetchResponse('/late', { auth: 'required' }); + const first = runtime.fetchResponse('/first', { auth: 'required' }); + await vi.waitFor(() => expect(cookies).toHaveBeenCalledTimes(1)); + release401(expired()); + expect((await late).status).toBe(401); + releaseCookies(); + expect((await first).status).toBe(401); + expect(fetchMock.mock.calls.filter(([url]) => url.endsWith('/late'))).toHaveLength(1); + expect(fetchMock.mock.calls.filter(([url]) => url.endsWith('/first'))).toHaveLength(1); +}); + +it("can recover when writing the access cookie fails after the rotated refresh cookie", async () => { + const jar = new Map([[ACCESS_TOKEN_COOKIE, 'old'], [REFRESH_TOKEN_COOKIE, 'old-refresh']]); + let failAccess = true; + vi.mocked(cookies).mockResolvedValue({ + set(name: string, value: string) { + if (name === ACCESS_TOKEN_COOKIE && failAccess) { + failAccess = false; + throw new Error('access cookie write failed'); + } + jar.set(name, value); + }, + } as unknown as Awaited>); + let validRefresh = 'old-refresh'; + let rotations = 0; + const seenRefresh: string[] = []; + vi.stubGlobal('fetch', vi.fn(async (url: string, init?: RequestInit) => { + if (url.endsWith('/auth/refresh')) { + const refresh = JSON.parse(String(init?.body)).refresh_token; + seenRefresh.push(refresh); + if (refresh !== validRefresh) return new Response('expired', { status: 401 }); + rotations++; + validRefresh = `refresh-${rotations}`; + return Response.json({ access_token: `access-${rotations}`, refresh_token: validRefresh }); + } + return new Headers(init?.headers).get('Authorization') === `Bearer access-${rotations}` + ? Response.json({ user_id: 'alice' }) : new Response('expired', { status: 401 }); + })); + const first = new WebRuntimeClient(runtimeConfig('old', 'old-refresh')); + expect((await first.fetchResponse('/resource', { auth: 'required' })).status).toBe(401); + expect(jar.get(ACCESS_TOKEN_COOKIE)).toBe('old'); + expect(jar.get(REFRESH_TOKEN_COOKIE)).toBe('refresh-1'); + const retry = new WebRuntimeClient(runtimeConfig(jar.get(ACCESS_TOKEN_COOKIE)!, jar.get(REFRESH_TOKEN_COOKIE)!)); + await expect(retry.get('/resource', { auth: 'required' })).resolves.toEqual({ user_id: 'alice' }); + expect(seenRefresh).toEqual(['old-refresh', 'refresh-1']); +}); + it.each([ { name: "SDK notification", sdkRefresh: true }, { name: "Web transport", sdkRefresh: false }, -])("keeps the newest usable cookies when the $name refresh finishes before an older notification", async ({ sdkRefresh }) => { +])("does not rotate during an outstanding $name notification and recovers afterward", async ({ sdkRefresh }) => { let release!: () => void; const available = new Promise(resolve => { release = resolve; }); const jar = new Map(); @@ -94,20 +159,19 @@ it.each([ await vi.waitFor(() => expect(cookies).toHaveBeenCalledTimes(1)); expired.add("access-1"); if (sdkRefresh) { - await expect(runtime.sdk.getMe()).resolves.toEqual({ user_id: "alice" }); + await expect(runtime.sdk.getMe()).rejects.toMatchObject({ status: 401 }); } else { - await expect(runtime.get("/auth/me", { auth: "required" })).resolves.toEqual({ user_id: "alice" }); + await expect(runtime.get("/auth/me", { auth: "required" })).rejects.toMatchObject({ status: 401 }); } release(); - expect(await first).toMatchObject([{ status: "fulfilled", value: { user_id: "alice" } }]); + expect(await first).toMatchObject([{ status: "rejected", reason: { status: 401 } }]); - const committed = set.mock.calls.map(([name, value]) => [name, value]); + expect(set.mock.calls.map(([name, value]) => [name, value])).toEqual([ + [REFRESH_TOKEN_COOKIE, "refresh-1"], + [ACCESS_TOKEN_COOKIE, "access-1"], + ]); const next = new WebRuntimeClient(runtimeConfig(jar.get(ACCESS_TOKEN_COOKIE)!, jar.get(REFRESH_TOKEN_COOKIE)!)); await expect(next.sdk.getMe()).resolves.toEqual({ user_id: "alice" }); - expect(committed).toEqual([ - [ACCESS_TOKEN_COOKIE, "access-2"], - [REFRESH_TOKEN_COOKIE, "refresh-2"], - ]); expect(jar.get(REFRESH_TOKEN_COOKIE)).toBe(validRefresh); expect(refreshRequests).toEqual(["refresh-0", "refresh-1"]); }); diff --git a/web/lib/runtime-client/server.ts b/web/lib/runtime-client/server.ts index 418cbd85c8..a771a39f17 100644 --- a/web/lib/runtime-client/server.ts +++ b/web/lib/runtime-client/server.ts @@ -190,12 +190,12 @@ export class WebRuntimeClient { if (current.accessToken !== accessToken || current.refreshToken !== refreshToken) { return; } - cookieStore.set(ACCESS_TOKEN_COOKIE, accessToken, { + cookieStore.set(REFRESH_TOKEN_COOKIE, refreshToken, { httpOnly: true, sameSite: "lax", path: "/", }); - cookieStore.set(REFRESH_TOKEN_COOKIE, refreshToken, { + cookieStore.set(ACCESS_TOKEN_COOKIE, accessToken, { httpOnly: true, sameSite: "lax", path: "/",