diff --git a/src/session.spec.ts b/src/session.spec.ts index d8a1d48..456e11c 100644 --- a/src/session.spec.ts +++ b/src/session.spec.ts @@ -11,12 +11,17 @@ import { authkitLoader, encryptSession, terminateSession, refreshSession } from import { assertIsResponse } from './test-utils/test-helpers.js'; import { getWorkOS } from './workos.js'; import { getConfig } from './config.js'; +import { getAuthorizationUrl } from './get-authorization-url.js'; jest.mock('./sessionStorage.js', () => ({ configureSessionStorage: jest.fn(), getSessionStorage: jest.fn(), })); +jest.mock('./get-authorization-url.js', () => ({ + getAuthorizationUrl: jest.fn(), +})); + // Mock dependencies const fakeWorkosInstance = { userManagement: { @@ -39,6 +44,7 @@ const authenticateWithRefreshToken = jest.mocked(workos.userManagement.authentic const getSessionStorage = jest.mocked(getSessionStorageMock); const configureSessionStorage = jest.mocked(configureSessionStorageMock); const jwtVerify = jest.mocked(jose.jwtVerify); +const getAuthorizationUrlMock = jest.mocked(getAuthorizationUrl); function getHeaderValue(headers: HeadersInit | undefined, name: string): string | null { if (!headers) { @@ -113,6 +119,10 @@ describe('session', () => { destroySession, commitSession, }); + + // Reset getAuthorizationUrl mock + getAuthorizationUrlMock.mockReset(); + getAuthorizationUrlMock.mockResolvedValue('https://auth.workos.com/oauth/authorize'); }); describe('encryptSession', () => { @@ -594,17 +604,26 @@ describe('session', () => { expect(getHeaderValue(init?.headers, 'Set-Cookie')).toBe('new-session-cookie'); }); - it('should redirect to root when refresh fails', async () => { + it('should redirect to authorization URL preserving returnPathname when refresh fails', async () => { authenticateWithRefreshToken.mockRejectedValue(new Error('Refresh token invalid')); + // Setup the mock to return a URL with state parameter + getAuthorizationUrlMock.mockResolvedValue('https://auth.workos.com/oauth/authorize?state=abc123'); + try { - await authkitLoader(createLoaderArgs(createMockRequest())); + const mockRequest = createMockRequest('test-cookie', 'https://app.example.com/dashboard/settings'); + await authkitLoader(createLoaderArgs(mockRequest)); fail('Expected redirect response to be thrown'); } catch (response: unknown) { assertIsResponse(response); expect(response.status).toBe(302); - expect(response.headers.get('Location')).toBe('/'); + expect(response.headers.get('Location')).toBe('https://auth.workos.com/oauth/authorize?state=abc123'); expect(response.headers.get('Set-Cookie')).toBe('destroyed-session-cookie'); + + // Verify getAuthorizationUrl was called with the correct returnPathname + expect(getAuthorizationUrlMock).toHaveBeenCalledWith({ + returnPathname: '/dashboard/settings', + }); } }); diff --git a/src/session.ts b/src/session.ts index 9879779..e4bf4f6 100644 --- a/src/session.ts +++ b/src/session.ts @@ -393,7 +393,8 @@ export async function authkitLoader( } } - throw redirect('/', { + const returnPathname = getReturnPathname(request.url); + throw redirect(await getAuthorizationUrl({ returnPathname }), { headers: { 'Set-Cookie': await destroySession(cookieSession), },