diff --git a/src/server/middleware-body.spec.ts b/src/server/middleware-body.spec.ts index 524785b..00c5c3a 100644 --- a/src/server/middleware-body.spec.ts +++ b/src/server/middleware-body.spec.ts @@ -139,6 +139,30 @@ describe('middlewareBody', () => { expect(result.response.headers.get('Set-Cookie')).toContain('wos-session=new_value'); }); + it.each([ + ['organization switch', 'wos-session=org_b_r3; Path=/'], + ['sign out', 'wos-session=; Path=/; Max-Age=0'], + ])('keeps the %s cookie after the auto-refresh cookie', async (_, handlerCookie) => { + const refreshCookie = 'wos-session=org_a_r2; Path=/'; + mockAuthkit.withAuth.mockResolvedValue({ + auth: { user: { id: 'user_123' } }, + refreshedSessionData: 'refreshed_session', + }); + mockAuthkit.saveSession.mockResolvedValue({ + response: new Response(null, { headers: { 'Set-Cookie': refreshCookie } }), + }); + + const result = await middlewareBody({ + request: new Request('http://test.local'), + next: vi.fn(async ({ context }: any) => { + context.__setPendingHeader('Set-Cookie', handlerCookie); + return { response: new Response('OK') }; + }), + }); + + expect(result.response.headers.getSetCookie()).toEqual([refreshCookie, handlerCookie]); + }); + it('provides correct context shape to downstream handlers', async () => { const mockAuth = { user: { id: 'user_123' }, sessionId: 'session_123' }; mockAuthkit.withAuth.mockResolvedValue({ diff --git a/src/server/middleware-body.ts b/src/server/middleware-body.ts index 9aa9776..7b41de1 100644 --- a/src/server/middleware-body.ts +++ b/src/server/middleware-body.ts @@ -15,6 +15,14 @@ export async function middlewareBody(args: any, options?: AuthKitMiddlewareOptio const { auth, refreshedSessionData } = await authkit.withAuth(args.request); const pendingHeaders = new Headers(); + // Queue auto-refresh cookies first so downstream session writes take precedence. + if (refreshedSessionData) { + const { response: sessionResponse } = await authkit.saveSession(undefined, refreshedSessionData); + for (const cookie of sessionResponse?.headers.getSetCookie() ?? []) { + pendingHeaders.append('Set-Cookie', cookie); + } + } + const result = await args.next({ context: { auth: () => auth, @@ -30,13 +38,6 @@ export async function middlewareBody(args: any, options?: AuthKitMiddlewareOptio }, }); - if (refreshedSessionData) { - const { response: sessionResponse } = await authkit.saveSession(undefined, refreshedSessionData); - for (const cookie of sessionResponse?.headers.getSetCookie() ?? []) { - pendingHeaders.append('Set-Cookie', cookie); - } - } - const headerEntries = [...pendingHeaders]; if (headerEntries.length === 0) { return result;