@workos-inc/authkit-nextjs 2.14.0 → 2.16.0

package/src/auth.spec.ts

@@ -29,6 +29,13 @@ const fakeWorkosInstance = {
     getJwksUrl: vi.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
     getLogoutUrl: vi.fn(),
   },
+  pkce: {
+    generate: vi.fn().mockResolvedValue({
+      codeVerifier: 'test-code-verifier',
+      codeChallenge: 'test-code-challenge',
+      codeChallengeMethod: 'S256' as const,
+    }),
+  },
 };
 
 const revalidatePath = vi.mocked(cache.revalidatePath);
@@ -74,6 +81,15 @@ describe('auth.ts', () => {
       expect(url).toBeDefined();
       expect(() => new URL(url)).not.toThrow();
     });
+
+    it('should include returnTo as returnPathname in the state parameter', async () => {
+      const url = await getSignInUrl({ returnTo: '/dashboard' });
+      const parsedUrl = new URL(url);
+      const state = parsedUrl.searchParams.get('state');
+      expect(state).toBeDefined();
+      const decoded = JSON.parse(atob(state!.replace(/-/g, '+').replace(/_/g, '/')));
+      expect(decoded.returnPathname).toBe('/dashboard');
+    });
   });
 
   it('should not include prompt when not specified for getSignInUrl', async () => {
@@ -101,6 +117,15 @@ describe('auth.ts', () => {
       const url = await getSignUpUrl({ prompt: 'consent' });
       expect(url).toContain('prompt=consent');
     });
+
+    it('should include returnTo as returnPathname in the state parameter', async () => {
+      const url = await getSignUpUrl({ returnTo: '/welcome' });
+      const parsedUrl = new URL(url);
+      const state = parsedUrl.searchParams.get('state');
+      expect(state).toBeDefined();
+      const decoded = JSON.parse(atob(state!.replace(/-/g, '+').replace(/_/g, '/')));
+      expect(decoded.returnPathname).toBe('/welcome');
+    });
   });
 
   describe('switchToOrganization', () => {
@@ -133,9 +158,16 @@ describe('auth.ts', () => {
         nextHeaders.set('x-url', 'http://localhost/test');
         await generateSession();
 
+        fakeWorkosInstance.pkce.generate.mockResolvedValue({
+          codeVerifier: 'test-code-verifier',
+          codeChallenge: 'test-code-challenge',
+          codeChallengeMethod: 'S256' as const,
+        });
+
         // Create a WorkOS-like object that matches what our tests need
         const mockWorkOS = {
           userManagement: fakeWorkosInstance.userManagement,
+          pkce: fakeWorkosInstance.pkce,
           // Add minimal properties to satisfy TypeScript
           createHttpClient: vi.fn(),
           createWebhookClient: vi.fn(),

package/src/auth.ts

@@ -7,7 +7,8 @@ import { redirect } from 'next/navigation';
 import { WORKOS_COOKIE_NAME } from './env-variables.js';
 import { getCookieOptions } from './cookie.js';
 import { getAuthorizationUrl } from './get-authorization-url.js';
-import type { AccessToken, SwitchToOrganizationOptions, UserInfo } from './interfaces.js';
+import type { AccessToken, GetAuthURLOptions, SwitchToOrganizationOptions, UserInfo } from './interfaces.js';
+import { setPKCECookie } from './pkce.js';
 import { getSessionFromCookie, refreshSession, withAuth } from './session.js';
 import { getWorkOS } from './workos.js';
 
@@ -20,20 +21,36 @@ function revalidateTagCompat(tag: string): void {
   return fn(tag, 'max');
 }
 
+async function getAuthURLAndSetPKCECookie(options: GetAuthURLOptions): Promise<string> {
+  const { url, pkceCookieValue } = await getAuthorizationUrl(options);
+  await setPKCECookie(pkceCookieValue);
+  return url;
+}
+
 export async function getSignInUrl({
   organizationId,
   loginHint,
   redirectUri,
   prompt,
   state,
+  returnTo,
 }: {
   organizationId?: string;
   loginHint?: string;
   redirectUri?: string;
   prompt?: 'consent';
   state?: string;
+  returnTo?: string;
 } = {}) {
-  return getAuthorizationUrl({ organizationId, screenHint: 'sign-in', loginHint, redirectUri, prompt, state });
+  return getAuthURLAndSetPKCECookie({
+    organizationId,
+    screenHint: 'sign-in',
+    loginHint,
+    redirectUri,
+    prompt,
+    state,
+    returnPathname: returnTo,
+  });
 }
 
 export async function getSignUpUrl({
@@ -42,14 +59,24 @@ export async function getSignUpUrl({
   redirectUri,
   prompt,
   state,
+  returnTo,
 }: {
   organizationId?: string;
   loginHint?: string;
   redirectUri?: string;
   prompt?: 'consent';
   state?: string;
+  returnTo?: string;
 } = {}) {
-  return getAuthorizationUrl({ organizationId, screenHint: 'sign-up', loginHint, redirectUri, prompt, state });
+  return getAuthURLAndSetPKCECookie({
+    organizationId,
+    screenHint: 'sign-up',
+    loginHint,
+    redirectUri,
+    prompt,
+    state,
+    returnPathname: returnTo,
+  });
 }
 
 /**
@@ -77,7 +104,12 @@ export async function signOut({ returnTo }: { returnTo?: string } = {}) {
     const nextCookies = await cookies();
     const cookieName = WORKOS_COOKIE_NAME || 'wos-session';
     const { domain, path, sameSite, secure } = getCookieOptions();
-    nextCookies.delete({ name: cookieName, domain, path, sameSite, secure });
+    try {
+      nextCookies.delete({ name: cookieName, domain, path, sameSite, secure });
+    } catch {
+      // Some environments (e.g., vinext) only accept a string cookie name
+      nextCookies.delete(cookieName);
+    }
 
     if (sessionId) {
       redirect(getWorkOS().userManagement.getLogoutUrl({ sessionId, returnTo }));
@@ -108,22 +140,25 @@ export async function switchToOrganization(
       redirect(cause.rawData.authkit_redirect_url);
     } else {
       if (cause?.error === 'sso_required' || cause?.error === 'mfa_enrollment') {
-        const url = await getAuthorizationUrl({ organizationId });
-        return redirect(url);
+        return redirect(await getAuthURLAndSetPKCECookie({ organizationId }));
       }
       throw error;
     }
   }
 
-  switch (revalidationStrategy) {
-    case 'path':
-      revalidatePath(pathname);
-      break;
-    case 'tag':
-      for (const tag of revalidationTags) {
-        revalidateTagCompat(tag);
-      }
-      break;
+  try {
+    switch (revalidationStrategy) {
+      case 'path':
+        revalidatePath(pathname);
+        break;
+      case 'tag':
+        for (const tag of revalidationTags) {
+          revalidateTagCompat(tag);
+        }
+        break;
+    }
+  } catch {
+    // revalidatePath/revalidateTag may not be available in non-Next.js environments (e.g., vinext)
   }
   if (revalidationStrategy !== 'none') {
     redirect(pathname);

package/src/authkit-callback-route.spec.ts

@@ -2,6 +2,7 @@ import { getWorkOS } from './workos.js';
 import { handleAuth } from './authkit-callback-route.js';
 import { getSessionFromCookie, saveSession } from './session.js';
 import { NextRequest, NextResponse } from 'next/server';
+import { sealData } from 'iron-session';
 
 // Mocked in vitest.setup.ts
 import { cookies, headers } from 'next/headers';
@@ -13,6 +14,9 @@ const { fakeWorkosInstance } = vi.hoisted(() => ({
       authenticateWithCode: vi.fn(),
       getJwksUrl: vi.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
     },
+    pkce: {
+      generate: vi.fn(),
+    },
   },
 }));
 
@@ -361,5 +365,86 @@ describe('authkit-callback-route', () => {
       // Should still redirect correctly
       expect(response.headers.get('Location')).toContain('/old-path');
     });
+
+    describe('PKCE', () => {
+      it('should pass codeVerifier from cookie to authenticateWithCode', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        // Seal a verifier into a cookie value
+        const sealedVerifier = await sealData(
+          { codeVerifier: 'test-verifier-123' },
+          { password: process.env.WORKOS_COOKIE_PASSWORD! },
+        );
+
+        // Set the PKCE cookie on the request
+        request.cookies.set('wos-pkce-verifier', sealedVerifier);
+        request.nextUrl.searchParams.set('code', 'test-code');
+
+        const handler = handleAuth();
+        await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith(
+          expect.objectContaining({
+            code: 'test-code',
+            codeVerifier: 'test-verifier-123',
+          }),
+        );
+      });
+
+      it('should proceed without codeVerifier when PKCE cookie is missing', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        request.nextUrl.searchParams.set('code', 'test-code');
+
+        const handler = handleAuth();
+        await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith(
+          expect.objectContaining({
+            code: 'test-code',
+            codeVerifier: undefined,
+          }),
+        );
+      });
+
+      it('should proceed without codeVerifier when PKCE cookie is corrupted', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        // Set a corrupted cookie
+        request.cookies.set('wos-pkce-verifier', 'not-a-valid-sealed-value');
+        request.nextUrl.searchParams.set('code', 'test-code');
+
+        const handler = handleAuth();
+        await handler(request);
+
+        expect(workos.userManagement.authenticateWithCode).toHaveBeenCalledWith(
+          expect.objectContaining({
+            code: 'test-code',
+            codeVerifier: undefined,
+          }),
+        );
+      });
+
+      it('should delete PKCE cookie after successful authentication', async () => {
+        vi.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+        const sealedVerifier = await sealData(
+          { codeVerifier: 'test-verifier-123' },
+          { password: process.env.WORKOS_COOKIE_PASSWORD! },
+        );
+
+        request.cookies.set('wos-pkce-verifier', sealedVerifier);
+        request.nextUrl.searchParams.set('code', 'test-code');
+
+        const handler = handleAuth();
+        const response = await handler(request);
+
+        // The response should have a Set-Cookie header to delete the PKCE cookie
+        const setCookieHeaders = response.headers.getSetCookie();
+        const pkceDeletionCookie = setCookieHeaders.find((c: string) => c.startsWith('wos-pkce-verifier='));
+        expect(pkceDeletionCookie).toBeDefined();
+        expect(pkceDeletionCookie).toContain('Max-Age=0');
+      });
+    });
   });
 });

package/src/authkit-callback-route.ts

@@ -1,6 +1,8 @@
 import { NextRequest } from 'next/server';
+import { getCookieOptions } from './cookie.js';
 import { WORKOS_CLIENT_ID } from './env-variables.js';
 import { HandleAuthOptions } from './interfaces.js';
+import { PKCE_COOKIE_NAME, getPKCECodeVerifier } from './pkce.js';
 import { saveSession } from './session.js';
 import { errorResponseWithFallback, redirectWithFallback, setCachePreventionHeaders } from './utils.js';
 import { getWorkOS } from './workos.js';
@@ -54,24 +56,30 @@ export function handleAuth(options: HandleAuthOptions = {}) {
   }
 
   return async function GET(request: NextRequest) {
-    const code = request.nextUrl.searchParams.get('code');
-    const state = request.nextUrl.searchParams.get('state');
+    // Fall back to standard URL parsing when nextUrl is not available (e.g., vinext)
+    const requestUrl = request.nextUrl ?? new URL(request.url);
+    const code = requestUrl.searchParams.get('code');
+    const state = requestUrl.searchParams.get('state');
 
     const { state: customState, returnPathname: returnPathnameState } = handleState(state);
 
     if (code) {
       try {
+        const pkceCookie = request.cookies.get(PKCE_COOKIE_NAME);
+        const codeVerifier = await getPKCECodeVerifier(pkceCookie?.value);
+
         // Use the code returned to us by AuthKit and authenticate the user with WorkOS
         const { accessToken, refreshToken, user, impersonator, oauthTokens, authenticationMethod, organizationId } =
           await getWorkOS().userManagement.authenticateWithCode({
             clientId: WORKOS_CLIENT_ID,
             code,
+            codeVerifier,
           });
 
         // If baseURL is provided, use it instead of request.nextUrl
         // This is useful if the app is being run in a container like docker where
         // the hostname can be different from the one in the request
-        const url = baseURL ? new URL(baseURL) : request.nextUrl.clone();
+        const url = baseURL ? new URL(baseURL) : new URL(requestUrl.toString());
 
         // Cleanup params
         url.searchParams.delete('code');
@@ -90,6 +98,10 @@ export function handleAuth(options: HandleAuthOptions = {}) {
         const response = redirectWithFallback(url.toString());
         preventCaching(response.headers);
 
+        if (pkceCookie) {
+          response.headers.append('Set-Cookie', `${PKCE_COOKIE_NAME}=; ${getCookieOptions(request.url, true, true)}`);
+        }
+
         if (!accessToken || !refreshToken) throw new Error('response is missing tokens');
 
         await saveSession({ accessToken, refreshToken, user, impersonator }, request);

package/src/components/authkit-provider.tsx

@@ -71,7 +71,7 @@ export const AuthKitProvider = ({ children, onSessionExpired, initialAuth }: Aut
       setEntitlements(auth.entitlements);
       setFeatureFlags(auth.featureFlags);
       setImpersonator(auth.impersonator);
-    } catch (error) {
+    } catch {
       setUser(null);
       setSessionId(undefined);
       setOrganizationId(undefined);

package/src/components/tokenStore.spec.ts

@@ -530,7 +530,7 @@ describe('tokenStore', () => {
 
       try {
         await tokenStore.refreshToken();
-      } catch (e) {
+      } catch {
         // Expected to throw
       }
 
@@ -547,7 +547,7 @@ describe('tokenStore', () => {
 
       try {
         await tokenStore.refreshToken();
-      } catch (e) {
+      } catch {
         // Expected to throw
       }
 
@@ -699,7 +699,7 @@ describe('tokenStore', () => {
 
       try {
         await tokenStore.refreshToken();
-      } catch (e) {
+      } catch {
         // Expected to throw
       }
 

package/src/env-variables.ts

@@ -12,6 +12,7 @@ const WORKOS_COOKIE_DOMAIN = getEnvVariable('WORKOS_COOKIE_DOMAIN');
 const WORKOS_COOKIE_MAX_AGE = getEnvVariable('WORKOS_COOKIE_MAX_AGE');
 const WORKOS_COOKIE_NAME = getEnvVariable('WORKOS_COOKIE_NAME');
 const WORKOS_COOKIE_SAMESITE = getEnvVariable('WORKOS_COOKIE_SAMESITE') as 'lax' | 'strict' | 'none' | undefined;
+const WORKOS_DISABLE_PKCE = getEnvVariable('WORKOS_DISABLE_PKCE');
 
 // Required env variables
 const WORKOS_API_KEY = getEnvVariable('WORKOS_API_KEY') ?? '';
@@ -31,4 +32,5 @@ export {
   WORKOS_COOKIE_PASSWORD,
   WORKOS_REDIRECT_URI,
   WORKOS_COOKIE_SAMESITE,
+  WORKOS_DISABLE_PKCE,
 };

package/src/get-authorization-url.spec.ts

@@ -8,6 +8,13 @@ const { fakeWorkosInstance } = vi.hoisted(() => ({
     userManagement: {
       getAuthorizationUrl: vi.fn(),
     },
+    pkce: {
+      generate: vi.fn().mockResolvedValue({
+        codeVerifier: 'test-code-verifier',
+        codeChallenge: 'test-code-challenge',
+        codeChallengeMethod: 'S256' as const,
+      }),
+    },
   },
 }));
 
@@ -19,6 +26,12 @@ describe('getAuthorizationUrl', () => {
   const workos = getWorkOS();
   beforeEach(() => {
     vi.clearAllMocks();
+    delete process.env.WORKOS_DISABLE_PKCE;
+    fakeWorkosInstance.pkce.generate.mockResolvedValue({
+      codeVerifier: 'test-code-verifier',
+      codeChallenge: 'test-code-challenge',
+      codeChallengeMethod: 'S256' as const,
+    });
   });
 
   it('uses x-redirect-uri header when redirectUri option is not provided', async () => {
@@ -56,4 +69,53 @@ describe('getAuthorizationUrl', () => {
       }),
     );
   });
+
+  describe('PKCE', () => {
+    it('generates PKCE pair and includes codeChallenge in authorization URL', async () => {
+      vi.mocked(workos.userManagement.getAuthorizationUrl).mockReturnValue('mock-url');
+
+      const result = await getAuthorizationUrl({});
+
+      expect(fakeWorkosInstance.pkce.generate).toHaveBeenCalled();
+      expect(workos.userManagement.getAuthorizationUrl).toHaveBeenCalledWith(
+        expect.objectContaining({
+          codeChallenge: 'test-code-challenge',
+          codeChallengeMethod: 'S256',
+        }),
+      );
+      expect(result.url).toBe('mock-url');
+      expect(result.pkceCookieValue).toBeDefined();
+      expect(result.pkceCookieValue).not.toBe('');
+    });
+
+    it('returns sealed cookie value for the verifier', async () => {
+      vi.mocked(workos.userManagement.getAuthorizationUrl).mockReturnValue('mock-url');
+
+      const result = await getAuthorizationUrl({});
+
+      // pkceCookieValue should be a sealed (encrypted) string
+      expect(typeof result.pkceCookieValue).toBe('string');
+      expect(result.pkceCookieValue!.length).toBeGreaterThan(0);
+    });
+
+    it('skips PKCE when WORKOS_DISABLE_PKCE is set to true', async () => {
+      process.env.WORKOS_DISABLE_PKCE = 'true';
+
+      // Re-import to pick up the new env var
+      vi.resetModules();
+      const { getAuthorizationUrl: freshGetAuthorizationUrl } = await import('./get-authorization-url.js');
+
+      vi.mocked(workos.userManagement.getAuthorizationUrl).mockReturnValue('mock-url');
+
+      const result = await freshGetAuthorizationUrl({});
+
+      expect(fakeWorkosInstance.pkce.generate).not.toHaveBeenCalled();
+      expect(workos.userManagement.getAuthorizationUrl).toHaveBeenCalledWith(
+        expect.not.objectContaining({
+          codeChallenge: expect.any(String),
+        }),
+      );
+      expect(result.pkceCookieValue).toBeUndefined();
+    });
+  });
 });

package/src/get-authorization-url.ts

@@ -1,9 +1,10 @@
-import { getWorkOS } from './workos.js';
-import { WORKOS_CLIENT_ID, WORKOS_REDIRECT_URI } from './env-variables.js';
-import { GetAuthURLOptions } from './interfaces.js';
+import { sealData } from 'iron-session';
 import { headers } from 'next/headers';
+import { WORKOS_CLIENT_ID, WORKOS_COOKIE_PASSWORD, WORKOS_DISABLE_PKCE, WORKOS_REDIRECT_URI } from './env-variables.js';
+import { GetAuthURLOptions, GetAuthURLResult } from './interfaces.js';
+import { getWorkOS } from './workos.js';
 
-async function getAuthorizationUrl(options: GetAuthURLOptions = {}) {
+async function getAuthorizationUrl(options: GetAuthURLOptions = {}): Promise<GetAuthURLResult> {
   const { returnPathname, screenHint, organizationId, loginHint, prompt, state: customState } = options;
   let redirectUri = options.redirectUri;
   if (!redirectUri) {
@@ -18,8 +19,8 @@ async function getAuthorizationUrl(options: GetAuthURLOptions = {}) {
   const finalState =
     internalState && customState ? `${internalState}.${customState}` : internalState || customState || undefined;
 
-  return getWorkOS().userManagement.getAuthorizationUrl({
-    provider: 'authkit',
+  const baseOptions = {
+    provider: 'authkit' as const,
     clientId: WORKOS_CLIENT_ID,
     redirectUri: redirectUri ?? WORKOS_REDIRECT_URI,
     state: finalState,
@@ -27,7 +28,24 @@ async function getAuthorizationUrl(options: GetAuthURLOptions = {}) {
     organizationId,
     loginHint,
     prompt,
+  };
+
+  if (WORKOS_DISABLE_PKCE === 'true') {
+    return { url: getWorkOS().userManagement.getAuthorizationUrl(baseOptions), pkceCookieValue: undefined };
+  }
+
+  const pkce = await getWorkOS().pkce.generate();
+  const pkceCookieValue = await sealData(
+    { codeVerifier: pkce.codeVerifier },
+    { password: WORKOS_COOKIE_PASSWORD, ttl: 600 },
+  );
+  const url = getWorkOS().userManagement.getAuthorizationUrl({
+    ...baseOptions,
+    codeChallenge: pkce.codeChallenge,
+    codeChallengeMethod: pkce.codeChallengeMethod,
   });
+
+  return { url, pkceCookieValue };
 }
 
 export { getAuthorizationUrl };

package/src/interfaces.ts

@@ -61,6 +61,11 @@ export interface AccessToken {
   feature_flags?: string[];
 }
 
+export interface GetAuthURLResult {
+  url: string;
+  pkceCookieValue?: string;
+}
+
 export interface GetAuthURLOptions {
   screenHint?: 'sign-up' | 'sign-in';
   returnPathname?: string;

package/src/pkce.ts

@@ -0,0 +1,42 @@
+import { unsealData } from 'iron-session';
+import { cookies } from 'next/headers';
+import { getCookieOptions } from './cookie.js';
+import { WORKOS_COOKIE_PASSWORD } from './env-variables.js';
+
+export const PKCE_COOKIE_NAME = 'wos-pkce-verifier';
+const PKCE_COOKIE_MAX_AGE = 600; // 10 minutes
+
+/**
+ * Set the PKCE verifier cookie in server action context.
+ * In middleware context, callers must set the cookie via Set-Cookie headers instead.
+ */
+export async function setPKCECookie(pkceCookieValue: string | undefined): Promise<void> {
+  if (!pkceCookieValue) return;
+  const nextCookies = await cookies();
+  const { domain, path, sameSite, secure } = getCookieOptions();
+  nextCookies.set(PKCE_COOKIE_NAME, pkceCookieValue, {
+    domain,
+    path,
+    sameSite,
+    secure,
+    httpOnly: true,
+    maxAge: PKCE_COOKIE_MAX_AGE,
+  });
+}
+
+/**
+ * Read and unseal the PKCE code verifier from the cookie.
+ * Returns undefined if the cookie is missing or corrupted.
+ */
+export async function getPKCECodeVerifier(cookieValue: string | undefined): Promise<string | undefined> {
+  if (!cookieValue) return undefined;
+  try {
+    const unsealed = await unsealData<{ codeVerifier: string }>(cookieValue, {
+      password: WORKOS_COOKIE_PASSWORD,
+    });
+    return unsealed.codeVerifier;
+  } catch {
+    // Cookie corrupted or expired — caller will proceed without PKCE
+    return undefined;
+  }
+}

package/src/session.ts

@@ -18,6 +18,7 @@ import {
   Session,
   UserInfo,
 } from './interfaces.js';
+import { PKCE_COOKIE_NAME, setPKCECookie } from './pkce.js';
 import { getWorkOS } from './workos.js';
 
 import type { AuthenticationResponse } from '@workos-inc/node';
@@ -25,6 +26,12 @@ import { parse, tokensToRegexp } from 'path-to-regexp';
 import { handleAuthkitHeaders } from './middleware-helpers.js';
 import { lazy, setCachePreventionHeaders } from './utils.js';
 
+function appendPKCESetCookieHeader(headers: Headers, pkceCookieValue: string | undefined, requestUrl: string): void {
+  if (pkceCookieValue) {
+    headers.append('Set-Cookie', `${PKCE_COOKIE_NAME}=${pkceCookieValue}; ${getCookieOptions(requestUrl, true)}`);
+  }
+}
+
 const sessionHeaderName = 'x-workos-session';
 const middlewareHeaderName = 'x-workos-middleware';
 const signUpPathsHeaderName = 'x-sign-up-paths';
@@ -202,14 +209,18 @@ async function updateSession(
       console.log('No session found from cookie');
     }
 
+    const { url: authorizationUrl, pkceCookieValue } = await getAuthorizationUrl({
+      returnPathname: getReturnPathname(request.url),
+      redirectUri: options.redirectUri || WORKOS_REDIRECT_URI,
+      screenHint: options.screenHint,
+    });
+
+    appendPKCESetCookieHeader(newRequestHeaders, pkceCookieValue, request.url);
+
     return {
       session: { user: null },
       headers: newRequestHeaders,
-      authorizationUrl: await getAuthorizationUrl({
-        returnPathname: getReturnPathname(request.url),
-        redirectUri: options.redirectUri || WORKOS_REDIRECT_URI,
-        screenHint: options.screenHint,
-      }),
+      authorizationUrl,
     };
   }
 
@@ -340,13 +351,17 @@ async function updateSession(
 
     options.onSessionRefreshError?.({ error: e, request });
 
+    const { url: authorizationUrl, pkceCookieValue } = await getAuthorizationUrl({
+      returnPathname: getReturnPathname(request.url),
+      redirectUri: options.redirectUri || WORKOS_REDIRECT_URI,
+    });
+
+    appendPKCESetCookieHeader(newRequestHeaders, pkceCookieValue, request.url);
+
     return {
       session: { user: null },
       headers: newRequestHeaders,
-      authorizationUrl: await getAuthorizationUrl({
-        returnPathname: getReturnPathname(request.url),
-        redirectUri: options.redirectUri || WORKOS_REDIRECT_URI,
-      }),
+      authorizationUrl,
     };
   }
 }
@@ -453,7 +468,9 @@ async function redirectToSignIn() {
 
   const returnPathname = getReturnPathname(url);
 
-  redirect(await getAuthorizationUrl({ returnPathname, screenHint }));
+  const { url: authkitUrl, pkceCookieValue } = await getAuthorizationUrl({ returnPathname, screenHint });
+  await setPKCECookie(pkceCookieValue);
+  redirect(authkitUrl);
 }
 
 export async function getTokenClaims<T = Record<string, unknown>>(