Add saveSession for manual auth flows (#52)

Author: chaance

src/index.ts

@@ -1,7 +1,7 @@
 import { getSignInUrl, getSignUpUrl, signOut, switchToOrganization, withAuth } from './auth.js';
 import { authLoader } from './authkit-callback-route.js';
 import { configure, getConfig } from './config.js';
-import { authkitLoader, refreshSession } from './session.js';
+import { authkitLoader, refreshSession, saveSession } from './session.js';
 import { getWorkOS } from './workos.js';
 
 export {
@@ -14,6 +14,7 @@ export {
   getSignUpUrl,
   getWorkOS,
   refreshSession,
+  saveSession,
   signOut,
   switchToOrganization,
 };

src/session.spec.ts

@@ -7,7 +7,7 @@ import {
   getSessionStorage as getSessionStorageMock,
 } from './sessionStorage.js';
 import { Session } from './interfaces.js';
-import { authkitLoader, encryptSession, terminateSession, refreshSession } from './session.js';
+import { authkitLoader, encryptSession, terminateSession, refreshSession, saveSession } from './session.js';
 import { assertIsResponse } from './test-utils/test-helpers.js';
 import { getWorkOS } from './workos.js';
 import { getConfig } from './config.js';
@@ -818,4 +818,96 @@ describe('session', () => {
       );
     });
   });
+
+  describe('saveSession', () => {
+    const sessionData = {
+      accessToken: 'new.valid.token',
+      refreshToken: 'new.refresh.token',
+      user: {
+        object: 'user',
+        id: 'user-1',
+        email: '[email protected]',
+        emailVerified: true,
+        profilePictureUrl: null,
+        firstName: 'Test',
+        lastName: 'User',
+        lastSignInAt: '2021-01-01T00:00:00Z',
+        createdAt: '2021-01-01T00:00:00Z',
+        updatedAt: '2021-01-01T00:00:00Z',
+        externalId: null,
+        locale: null,
+        metadata: {},
+      } satisfies User,
+      impersonator: undefined,
+      headers: {},
+    } satisfies Session;
+
+    const createMockRequest = (cookie = 'test-cookie', url = 'http://example.com./some-path') =>
+      new Request(url, {
+        headers: new Headers({
+          Cookie: cookie,
+        }),
+      });
+
+    let getSession: jest.Mock;
+    let destroySession: jest.Mock;
+    let commitSession: jest.Mock;
+    let mockSession: ReactRouterSession;
+
+    beforeEach(() => {
+      getSession = jest.fn();
+      destroySession = jest.fn().mockResolvedValue('destroyed-session-cookie');
+      commitSession = jest.fn().mockResolvedValue('new-session-cookie');
+
+      mockSession = createMockSession({
+        has: jest.fn().mockReturnValue(true),
+        get: jest.fn().mockReturnValue('encrypted-jwt'),
+        set: jest.fn(),
+      });
+
+      getSessionStorage.mockResolvedValue({
+        cookieName: 'wos-cookie',
+        getSession,
+        destroySession,
+        commitSession,
+      });
+
+      getSession.mockResolvedValue(mockSession);
+
+      const validSessionData = {
+        accessToken: 'valid.token',
+        refreshToken: 'refresh.token',
+        user: {
+          id: 'user-1',
+          email: '[email protected]',
+          firstName: 'Test',
+          lastName: 'User',
+          object: 'user',
+        },
+        impersonator: null,
+      };
+      unsealData.mockResolvedValue(validSessionData);
+      sealData.mockResolvedValue('new-encrypted-jwt');
+
+      authenticateWithRefreshToken.mockResolvedValue(sessionData);
+
+      // Mock JWT decoding
+      (jose.decodeJwt as jest.Mock).mockReturnValue({
+        sid: 'new-session-id',
+        org_id: 'org-123',
+        role: 'user',
+        roles: ['user'],
+        permissions: ['read'],
+        entitlements: ['basic'],
+        feature_flags: ['flag-1'],
+      });
+    });
+
+    it('should save the session to the cookie', async () => {
+      await saveSession(sessionData, createMockRequest());
+      expect(getSessionStorage).toHaveBeenCalled();
+      expect(commitSession).toHaveBeenCalledWith(mockSession);
+      expect(mockSession.set).toHaveBeenCalledWith('jwt', 'new-encrypted-jwt');
+    });
+  });
 });

src/session.ts

@@ -16,6 +16,7 @@ import { createRemoteJWKSet, decodeJwt, jwtVerify } from 'jose';
 import { getConfig } from './config.js';
 import { configureSessionStorage, getSessionStorage } from './sessionStorage.js';
 import { isDataWithResponseInit, isJsonResponse, isRedirect, isResponse } from './utils.js';
+import type { AuthenticationResponse } from '@workos-inc/node';
 
 // must be a type since this is a subtype of response
 // interfaces must conform to the types they extend
@@ -37,37 +38,24 @@ export class SessionRefreshError extends Error {
  * @param options - Optional configuration options
  * @returns A promise that resolves to the new session object
  */
-export async function refreshSession(request: Request, { organizationId }: { organizationId?: string } = {}) {
-  const { getSession, commitSession } = await getSessionStorage();
-  const session = await getSessionFromCookie(request.headers.get('Cookie') as string);
-
+export async function refreshSession(request: Request, options: { organizationId?: string } = {}) {
+  const { organizationId } = options;
+  const { getSession } = await getSessionStorage();
+  const cookie = request.headers.get('Cookie');
+  const session = cookie ? await getSessionFromCookie(cookie) : null;
   if (!session) {
     throw redirect(await getAuthorizationUrl());
   }
 
   try {
-    const { accessToken, refreshToken, user, impersonator } =
-      await getWorkOS().userManagement.authenticateWithRefreshToken({
-        clientId: getConfig('clientId'),
-        refreshToken: session.refreshToken,
-        organizationId,
-      });
-
-    const newSession = {
-      accessToken,
-      refreshToken,
-      user,
-      impersonator,
-      headers: {} as Record<string, string>,
-    };
-
-    const cookieSession = await getSession(request.headers.get('Cookie'));
-    cookieSession.set('jwt', await encryptSession(newSession));
-    const cookie = await commitSession(cookieSession);
-
-    newSession.headers = {
-      'Set-Cookie': cookie,
-    };
+    const refreshResult = await getWorkOS().userManagement.authenticateWithRefreshToken({
+      clientId: getConfig('clientId'),
+      refreshToken: session.refreshToken,
+      organizationId,
+    });
+    const { headers } = await saveSession(refreshResult, request);
+    const cookieSession = await getSession(cookie);
+    const { accessToken, user, impersonator } = refreshResult;
 
     const {
       sessionId,
@@ -91,7 +79,7 @@ export async function refreshSession(request: Request, { organizationId }: { org
       featureFlags,
       impersonator: impersonator ?? null,
       sealedSession: cookieSession.get('jwt'),
-      headers: newSession.headers,
+      headers,
     };
   } catch (error) {
     throw new Error(`Failed to refresh session: ${error instanceof Error ? error.message : String(error)}`, {
@@ -100,7 +88,54 @@ export async function refreshSession(request: Request, { organizationId }: { org
   }
 }
 
-async function updateSession(request: Request, debug: boolean) {
+/**
+ * Saves a WorkOS session to a cookie for use with AuthKit.
+ *
+ * This function is intended for advanced use cases where you need to manually
+ * manage sessions, such as custom authentication flows (email verification,
+ * etc.) that don't use the standard AuthKit authentication flow.
+ *
+ * @param sessionOrResponse The WorkOS session or AuthenticationResponse
+ * containing access token, refresh token, and user information.
+ * @param request A Request object, used to determine cookie settings.
+ *
+ * @example
+ * import { saveSession } from '@workos-inc/authkit-react-router';
+ *
+ * async function handleEmailVerification(req: Request) {
+ *   const { code } = await req.json();
+ *   const authResponse = await workos.userManagement.authenticateWithEmailVerification({
+ *     clientId: process.env.WORKOS_CLIENT_ID,
+ *     code,
+ *   });
+ *
+ *   await saveSession(authResponse, req);
+ * }
+ */
+export async function saveSession(
+  sessionOrResponse: Session | AuthenticationResponse,
+  request: Request,
+): Promise<Session> {
+  const { getSession, commitSession } = await getSessionStorage();
+  const { accessToken, refreshToken, user, impersonator } = sessionOrResponse;
+  const newSession: Session = {
+    accessToken,
+    refreshToken,
+    user,
+    impersonator,
+    headers: {},
+  };
+  const cookieSession = await getSession(request.headers.get('Cookie'));
+  cookieSession.set('jwt', await encryptSession(newSession));
+  const cookie = await commitSession(cookieSession);
+  newSession.headers = {
+    'Set-Cookie': cookie,
+  };
+
+  return newSession;
+}
+
+async function updateSession(request: Request, debug: boolean): Promise<Session | null> {
   const session = await getSessionFromCookie(request.headers.get('Cookie') as string);
   const { commitSession, getSession } = await getSessionStorage();
 
@@ -158,7 +193,7 @@ async function updateSession(request: Request, debug: boolean) {
   }
 }
 
-export async function encryptSession(session: Session) {
+export async function encryptSession(session: Session | AuthenticationResponse) {
   return sealData(session, {
     password: getConfig('cookiePassword'),
     ttl: 0,