Allow onSuccess callback to update session

src/authkit-callback-route.spec.ts

@@ -1,5 +1,6 @@
 import { getWorkOS } from './workos.js';
 import { handleAuth } from './authkit-callback-route.js';
+import { getSessionFromCookie, saveSession } from './session.js';
 import { NextRequest, NextResponse } from 'next/server';
 
 // Mocked in jest.setup.ts
@@ -253,6 +254,26 @@ describe('authkit-callback-route', () => {
       await handler(request);
 
       expect(onSuccess).toHaveBeenCalledWith(mockAuthResponse);
+      const session = await getSessionFromCookie();
+      expect(session?.accessToken).toBe(mockAuthResponse.accessToken);
+    });
+
+    it('should allow onSuccess to update session', async () => {
+      const newAccessToken = 'new-access-token';
+      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+      // Set up request with code
+      request.nextUrl.searchParams.set('code', 'test-code');
+
+      const handler = handleAuth({
+        onSuccess: async (data) => {
+          await saveSession({ ...data, accessToken: newAccessToken }, request);
+        },
+      });
+      await handler(request);
+
+      const session = await getSessionFromCookie();
+      expect(session?.accessToken).toBe(newAccessToken);
     });
   });
 });

src/authkit-callback-route.ts

@@ -61,6 +61,8 @@ export function handleAuth(options: HandleAuthOptions = {}) {
 
         if (!accessToken || !refreshToken) throw new Error('response is missing tokens');
 
+        await saveSession({ accessToken, refreshToken, user, impersonator }, request);
+
         if (onSuccess) {
           await onSuccess({
             accessToken,
@@ -73,8 +75,6 @@ export function handleAuth(options: HandleAuthOptions = {}) {
           });
         }
 
-        await saveSession({ accessToken, refreshToken, user, impersonator }, request);
-
         return response;
       } catch (error) {
         const errorRes = {

src/components/useAccessToken.ts

@@ -1,3 +1,5 @@
+'use client';
+
 import { useCallback, useEffect, useRef, useState, useSyncExternalStore } from 'react';
 import { useAuth } from './authkit-provider.js';
 import { tokenStore } from './tokenStore.js';