@workos-inc/authkit-nextjs 2.9.0 → 2.11.0

package/src/auth.ts

@@ -10,13 +10,30 @@ import { getAuthorizationUrl } from './get-authorization-url.js';
 import type { AccessToken, SwitchToOrganizationOptions, UserInfo } from './interfaces.js';
 import { getSessionFromCookie, refreshSession, withAuth } from './session.js';
 import { getWorkOS } from './workos.js';
+
+/**
+ * A wrapper around revalidateTag to provide compatibility with previous versions.
+ * @param tag The tag to revalidate.
+ */
+function revalidateTagCompat(tag: string): void {
+  const fn = revalidateTag as (tag: string, profile: string) => void;
+  return fn(tag, 'max');
+}
+
 export async function getSignInUrl({
   organizationId,
   loginHint,
   redirectUri,
   prompt,
-}: { organizationId?: string; loginHint?: string; redirectUri?: string; prompt?: 'consent' } = {}) {
-  return getAuthorizationUrl({ organizationId, screenHint: 'sign-in', loginHint, redirectUri, prompt });
+  state,
+}: {
+  organizationId?: string;
+  loginHint?: string;
+  redirectUri?: string;
+  prompt?: 'consent';
+  state?: string;
+} = {}) {
+  return getAuthorizationUrl({ organizationId, screenHint: 'sign-in', loginHint, redirectUri, prompt, state });
 }
 
 export async function getSignUpUrl({
@@ -24,8 +41,15 @@ export async function getSignUpUrl({
   loginHint,
   redirectUri,
   prompt,
-}: { organizationId?: string; loginHint?: string; redirectUri?: string; prompt?: 'consent' } = {}) {
-  return getAuthorizationUrl({ organizationId, screenHint: 'sign-up', loginHint, redirectUri, prompt });
+  state,
+}: {
+  organizationId?: string;
+  loginHint?: string;
+  redirectUri?: string;
+  prompt?: 'consent';
+  state?: string;
+} = {}) {
+  return getAuthorizationUrl({ organizationId, screenHint: 'sign-up', loginHint, redirectUri, prompt, state });
 }
 
 /**
@@ -97,7 +121,7 @@ export async function switchToOrganization(
       break;
     case 'tag':
       for (const tag of revalidationTags) {
-        revalidateTag(tag);
+        revalidateTagCompat(tag);
       }
       break;
   }

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

@@ -36,6 +36,7 @@ describe('authkit-callback-route', () => {
       lastSignInAt: '2024-01-01T00:00:00Z',
       externalId: null,
       metadata: {},
+      locale: null,
     },
     oauthTokens: {
       accessToken: 'access123',
@@ -275,5 +276,73 @@ describe('authkit-callback-route', () => {
       const session = await getSessionFromCookie();
       expect(session?.accessToken).toBe(newAccessToken);
     });
+
+    it('should pass custom state data to onSuccess callback', async () => {
+      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+      // Create state with new format: internal.user
+      const internalState = btoa(JSON.stringify({ returnPathname: '/dashboard' }))
+        .replace(/\+/g, '-')
+        .replace(/\//g, '_');
+      const userState = 'custom-user-state-string';
+      const state = `${internalState}.${userState}`;
+
+      request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', state);
+
+      const onSuccess = jest.fn();
+      const handler = handleAuth({ onSuccess });
+      await handler(request);
+
+      // Verify onSuccess was called with the custom state string
+      expect(onSuccess).toHaveBeenCalledWith(
+        expect.objectContaining({
+          ...mockAuthResponse,
+          state: 'custom-user-state-string',
+        }),
+      );
+
+      // Verify the redirect went to the correct path
+      const response = await handler(request);
+      expect(response.headers.get('Location')).toContain('/dashboard');
+    });
+
+    it('should handle state without custom data', async () => {
+      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+      // State with only returnPathname
+      const state = btoa(JSON.stringify({ returnPathname: '/profile' }));
+
+      request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', state);
+
+      const onSuccess = jest.fn();
+      const handler = handleAuth({ onSuccess });
+      await handler(request);
+
+      // Verify onSuccess was called without state property when no custom data exists
+      expect(onSuccess).toHaveBeenCalledWith(
+        expect.objectContaining({
+          ...mockAuthResponse,
+          state: undefined,
+        }),
+      );
+    });
+
+    it('should handle backward compatibility with old state format', async () => {
+      jest.mocked(workos.userManagement.authenticateWithCode).mockResolvedValue(mockAuthResponse);
+
+      // Old format: just returnPathname
+      const state = btoa(JSON.stringify({ returnPathname: '/old-path' }));
+
+      request.nextUrl.searchParams.set('code', 'test-code');
+      request.nextUrl.searchParams.set('state', state);
+
+      const handler = handleAuth();
+      const response = await handler(request);
+
+      // Should still redirect correctly
+      expect(response.headers.get('Location')).toContain('/old-path');
+    });
   });
 });

package/src/authkit-callback-route.ts

@@ -5,6 +5,37 @@ import { saveSession } from './session.js';
 import { errorResponseWithFallback, redirectWithFallback } from './utils.js';
 import { getWorkOS } from './workos.js';
 
+function handleState(state: string | null) {
+  let returnPathname: string | undefined = undefined;
+  let userState: string | undefined;
+  if (state?.includes('.')) {
+    const [internal, ...rest] = state.split('.');
+    userState = rest.join('.');
+    try {
+      // Reverse URL-safe base64 encoding
+      const decoded = internal.replace(/-/g, '+').replace(/_/g, '/');
+      returnPathname = JSON.parse(atob(decoded)).returnPathname;
+    } catch {
+      // Malformed internal part, ignore it
+    }
+  } else if (state) {
+    try {
+      const decoded = JSON.parse(atob(state));
+      if (decoded.returnPathname) {
+        returnPathname = decoded.returnPathname;
+      } else {
+        userState = state;
+      }
+    } catch {
+      userState = state;
+    }
+  }
+  return {
+    returnPathname,
+    state: userState,
+  };
+}
+
 export function handleAuth(options: HandleAuthOptions = {}) {
   const { returnPathname: returnPathnameOption = '/', baseURL, onSuccess, onError } = options;
 
@@ -20,7 +51,8 @@ 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');
-    let returnPathname = state && state !== 'null' ? JSON.parse(atob(state)).returnPathname : null;
+
+    const { state: customState, returnPathname: returnPathnameState } = handleState(state);
 
     if (code) {
       try {
@@ -41,7 +73,7 @@ export function handleAuth(options: HandleAuthOptions = {}) {
         url.searchParams.delete('state');
 
         // Redirect to the requested path and store the session
-        returnPathname = returnPathname ?? returnPathnameOption;
+        const returnPathname = returnPathnameState ?? returnPathnameOption;
 
         // Extract the search params if they are present
         if (returnPathname.includes('?')) {
@@ -72,6 +104,7 @@ export function handleAuth(options: HandleAuthOptions = {}) {
             oauthTokens,
             authenticationMethod,
             organizationId,
+            state: customState,
           });
         }
 

package/src/components/impersonation.spec.tsx

@@ -27,19 +27,6 @@ describe('Impersonation', () => {
       impersonator: null,
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
-    });
-
-    const { container } = render(<Impersonation />);
-    expect(container).toBeEmptyDOMElement();
-  });
-
-  it('should return null if loading', () => {
-    (useAuth as jest.Mock).mockReturnValue({
-      impersonator: { email: 'admin@example.com' },
-      user: { id: '123', email: 'user@example.com' },
-      organizationId: null,
-      loading: true,
     });
 
     const { container } = render(<Impersonation />);
@@ -51,7 +38,6 @@ describe('Impersonation', () => {
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const { container } = render(<Impersonation />);
@@ -63,7 +49,6 @@ describe('Impersonation', () => {
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: 'org_123',
-      loading: false,
     });
 
     (getOrganizationAction as jest.Mock).mockResolvedValue({
@@ -83,7 +68,6 @@ describe('Impersonation', () => {
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const { container } = render(<Impersonation />);
@@ -96,7 +80,6 @@ describe('Impersonation', () => {
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const { container } = render(<Impersonation side="top" />);
@@ -109,7 +92,6 @@ describe('Impersonation', () => {
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     const customStyle = { backgroundColor: 'red' };
@@ -123,12 +105,146 @@ describe('Impersonation', () => {
       impersonator: { email: 'admin@example.com' },
       user: { id: '123', email: 'user@example.com' },
       organizationId: null,
-      loading: false,
     });
 
     render(<Impersonation />);
     const stopButton = await screen.findByText('Stop');
     stopButton.click();
-    expect(handleSignOutAction).toHaveBeenCalled();
+    expect(handleSignOutAction).toHaveBeenCalledWith({});
+  });
+
+  it('should pass returnTo prop to handleSignOutAction when provided', async () => {
+    (useAuth as jest.Mock).mockReturnValue({
+      impersonator: { email: 'admin@example.com' },
+      user: { id: '123', email: 'user@example.com' },
+      organizationId: null,
+      loading: false,
+    });
+
+    const returnTo = '/dashboard';
+    render(<Impersonation returnTo={returnTo} />);
+    const stopButton = await screen.findByText('Stop');
+    stopButton.click();
+    expect(handleSignOutAction).toHaveBeenCalledWith({ returnTo });
+  });
+
+  it('should not call getOrganizationAction when organizationId is not provided', () => {
+    (useAuth as jest.Mock).mockReturnValue({
+      impersonator: { email: 'admin@example.com' },
+      user: { id: '123', email: 'user@example.com' },
+      organizationId: null,
+    });
+
+    render(<Impersonation />);
+    expect(getOrganizationAction).not.toHaveBeenCalled();
+  });
+
+  it('should not call getOrganizationAction when impersonator is not present', () => {
+    (useAuth as jest.Mock).mockReturnValue({
+      impersonator: null,
+      user: { id: '123', email: 'user@example.com' },
+      organizationId: 'org_123',
+    });
+
+    render(<Impersonation />);
+    expect(getOrganizationAction).not.toHaveBeenCalled();
+  });
+
+  it('should not call getOrganizationAction when user is not present', () => {
+    (useAuth as jest.Mock).mockReturnValue({
+      impersonator: { email: 'admin@example.com' },
+      user: null,
+      organizationId: 'org_123',
+    });
+
+    render(<Impersonation />);
+    expect(getOrganizationAction).not.toHaveBeenCalled();
+  });
+
+  it('should not call getOrganizationAction again when organization is already loaded with same ID', async () => {
+    const mockOrg = {
+      id: 'org_123',
+      name: 'Test Org',
+    };
+
+    (getOrganizationAction as jest.Mock).mockResolvedValue(mockOrg);
+
+    const { rerender } = await act(async () => {
+      (useAuth as jest.Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_123',
+      });
+
+      return render(<Impersonation />);
+    });
+
+    // Wait for the initial call to complete
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(getOrganizationAction).toHaveBeenCalledTimes(1);
+
+    // Rerender with the same organizationId
+    await act(async () => {
+      (useAuth as jest.Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_123',
+      });
+
+      rerender(<Impersonation />);
+    });
+
+    // Should still be called only once
+    expect(getOrganizationAction).toHaveBeenCalledTimes(1);
+  });
+
+  it('should call getOrganizationAction again when organizationId changes', async () => {
+    const mockOrg1 = {
+      id: 'org_123',
+      name: 'Test Org 1',
+    };
+
+    const mockOrg2 = {
+      id: 'org_456',
+      name: 'Test Org 2',
+    };
+
+    (getOrganizationAction as jest.Mock).mockResolvedValueOnce(mockOrg1).mockResolvedValueOnce(mockOrg2);
+
+    const { rerender } = await act(async () => {
+      (useAuth as jest.Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_123',
+      });
+
+      return render(<Impersonation />);
+    });
+
+    // Wait for the initial call to complete
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(getOrganizationAction).toHaveBeenCalledTimes(1);
+    expect(getOrganizationAction).toHaveBeenCalledWith('org_123');
+
+    // Rerender with a different organizationId
+    await act(async () => {
+      (useAuth as jest.Mock).mockReturnValue({
+        impersonator: { email: 'admin@example.com' },
+        user: { id: '123', email: 'user@example.com' },
+        organizationId: 'org_456',
+      });
+
+      rerender(<Impersonation />);
+    });
+
+    // Should be called again with the new ID
+    expect(getOrganizationAction).toHaveBeenCalledTimes(2);
+    expect(getOrganizationAction).toHaveBeenCalledWith('org_456');
   });
 });

package/src/components/impersonation.tsx

@@ -9,19 +9,21 @@ import { useAuth } from './authkit-provider.js';
 
 interface ImpersonationProps extends React.ComponentPropsWithoutRef<'div'> {
   side?: 'top' | 'bottom';
+  returnTo?: string;
 }
 
-export function Impersonation({ side = 'bottom', ...props }: ImpersonationProps) {
-  const { user, impersonator, organizationId, loading } = useAuth();
+export function Impersonation({ side = 'bottom', returnTo, ...props }: ImpersonationProps) {
+  const { user, impersonator, organizationId } = useAuth();
 
   const [organization, setOrganization] = React.useState<Organization | null>(null);
 
   React.useEffect(() => {
-    if (!organizationId) return;
+    if (!organizationId || !impersonator || !user) return;
+    if (organization && organization.id === organizationId) return;
     getOrganizationAction(organizationId).then(setOrganization);
-  }, [organizationId]);
+  }, [organizationId, impersonator, user]);
 
-  if (loading || !impersonator || !user) return null;
+  if (!impersonator || !user) return null;
 
   return (
     <div
@@ -78,7 +80,7 @@ export function Impersonation({ side = 'bottom', ...props }: ImpersonationProps)
         <form
           onSubmit={async (event) => {
             event.preventDefault();
-            await handleSignOutAction();
+            await handleSignOutAction({ returnTo });
           }}
           style={{
             display: 'flex',

package/src/get-authorization-url.ts

@@ -12,13 +12,21 @@ async function getAuthorizationUrl(options: GetAuthURLOptions = {}) {
     redirectUri = headersList.get('x-redirect-uri'),
     loginHint,
     prompt,
+    state: customState,
   } = options;
 
+  const internalState = returnPathname
+    ? btoa(JSON.stringify({ returnPathname })).replace(/\+/g, '-').replace(/\//g, '_')
+    : null;
+
+  const finalState =
+    internalState && customState ? `${internalState}.${customState}` : internalState || customState || undefined;
+
   return getWorkOS().userManagement.getAuthorizationUrl({
     provider: 'authkit',
     clientId: WORKOS_CLIENT_ID,
     redirectUri: redirectUri ?? WORKOS_REDIRECT_URI,
-    state: returnPathname ? btoa(JSON.stringify({ returnPathname })) : undefined,
+    state: finalState,
     screenHint,
     organizationId,
     loginHint,

package/src/index.ts

@@ -2,6 +2,7 @@ import { getSignInUrl, getSignUpUrl, signOut, switchToOrganization } from './aut
 import { handleAuth } from './authkit-callback-route.js';
 import { authkit, authkitMiddleware } from './middleware.js';
 import { getTokenClaims, refreshSession, saveSession, withAuth } from './session.js';
+import { validateApiKey } from './validate-api-key.js';
 import { getWorkOS } from './workos.js';
 
 export * from './interfaces.js';
@@ -19,4 +20,5 @@ export {
   switchToOrganization,
   withAuth,
   getTokenClaims,
+  validateApiKey,
 };

package/src/interfaces.ts

@@ -12,6 +12,7 @@ export interface HandleAuthSuccessData extends Session {
   oauthTokens?: OauthTokens;
   organizationId?: string;
   authenticationMethod?: AuthenticationResponse['authenticationMethod'];
+  state?: string | undefined;
 }
 
 export interface Impersonator {
@@ -67,6 +68,7 @@ export interface GetAuthURLOptions {
   redirectUri?: string;
   loginHint?: string;
   prompt?: 'consent';
+  state?: string;
 }
 
 export interface AuthkitMiddlewareAuth {

package/src/session.spec.ts

@@ -146,7 +146,12 @@ describe('session.ts', () => {
 
       await withAuth({ ensureSignedIn: true });
 
-      const pathname = encodeURIComponent(btoa(JSON.stringify({ returnPathname: '/protected?test=123' })));
+      // URL-safe base64 encoding
+      const pathname = encodeURIComponent(
+        btoa(JSON.stringify({ returnPathname: '/protected?test=123' }))
+          .replace(/\+/g, '-')
+          .replace(/\//g, '_'),
+      );
 
       expect(redirect).toHaveBeenCalledWith(expect.stringContaining(pathname));
     });

package/src/test-helpers.ts

@@ -45,6 +45,7 @@ export async function generateSession(overrides: Partial<User> = {}) {
     lastSignInAt: '2024-01-01T00:00:00Z',
     externalId: null,
     metadata: {},
+    locale: null,
     ...overrides,
   } satisfies User;
 

package/src/validate-api-key.spec.ts

@@ -0,0 +1,113 @@
+import { describe, it, expect, beforeEach, jest } from '@jest/globals';
+
+import { validateApiKey } from './validate-api-key.js';
+import { getWorkOS } from './workos.js';
+
+// These are mocked in jest.setup.ts
+import { headers } from 'next/headers';
+
+const workos = getWorkOS();
+
+describe('validate-api-key.ts', () => {
+  beforeEach(async () => {
+    // Clear all mocks between tests
+    jest.clearAllMocks();
+
+    const nextHeaders = await headers();
+    // @ts-expect-error - _reset is part of the mock
+    nextHeaders._reset();
+  });
+
+  describe('validateApiKey', () => {
+    it('should return valid API key when Bearer token is present and valid', async () => {
+      const mockApiKeyResponse = {
+        apiKey: {
+          id: 'api_key_123',
+          object: 'api_key' as const,
+          name: 'Test API Key',
+          obfuscatedValue: 'sk_…7890',
+          createdAt: '2024-01-01T00:00:00Z',
+          updatedAt: '2024-01-01T00:00:00Z',
+          lastUsedAt: '2024-01-01T00:00:00Z',
+          permissions: [],
+          owner: { type: 'organization' as const, id: 'org_123' },
+        },
+      };
+
+      jest.spyOn(workos.apiKeys, 'validateApiKey').mockResolvedValue(mockApiKeyResponse);
+
+      const nextHeaders = await headers();
+      nextHeaders.set('authorization', 'Bearer sk_test_1234567890');
+
+      const result = await validateApiKey();
+
+      expect(workos.apiKeys.validateApiKey).toHaveBeenCalledWith({
+        value: 'sk_test_1234567890',
+      });
+      expect(result).toEqual(mockApiKeyResponse);
+    });
+
+    it('should return { apiKey: null } when no authorization header is present', async () => {
+      // Don't set any authorization header
+      const result = await validateApiKey();
+
+      expect(workos.apiKeys.validateApiKey).not.toHaveBeenCalled();
+      expect(result).toEqual({ apiKey: null });
+    });
+
+    it('should return { apiKey: null } when authorization header is empty', async () => {
+      const nextHeaders = await headers();
+      nextHeaders.set('authorization', '');
+
+      const result = await validateApiKey();
+
+      expect(workos.apiKeys.validateApiKey).not.toHaveBeenCalled();
+      expect(result).toEqual({ apiKey: null });
+    });
+
+    it('should return { apiKey: null } when authorization header does not start with Bearer', async () => {
+      const nextHeaders = await headers();
+      nextHeaders.set('authorization', 'Basic dXNlcjpwYXNz');
+
+      const result = await validateApiKey();
+
+      expect(workos.apiKeys.validateApiKey).not.toHaveBeenCalled();
+      expect(result).toEqual({ apiKey: null });
+    });
+
+    it('should return { apiKey: null } when Bearer token is missing', async () => {
+      const nextHeaders = await headers();
+      nextHeaders.set('authorization', 'Bearer');
+
+      const result = await validateApiKey();
+
+      expect(workos.apiKeys.validateApiKey).not.toHaveBeenCalled();
+      expect(result).toEqual({ apiKey: null });
+    });
+
+    it('should return { apiKey: null } when Bearer token is only whitespace', async () => {
+      const nextHeaders = await headers();
+      nextHeaders.set('authorization', 'Bearer   ');
+
+      const result = await validateApiKey();
+
+      expect(workos.apiKeys.validateApiKey).not.toHaveBeenCalled();
+      expect(result).toEqual({ apiKey: null });
+    });
+
+    it('should return { apiKey: null } when WorkOS validation fails', async () => {
+      const mockResponse = { apiKey: null };
+      jest.spyOn(workos.apiKeys, 'validateApiKey').mockResolvedValue(mockResponse);
+
+      const nextHeaders = await headers();
+      nextHeaders.set('authorization', 'Bearer invalid_key');
+
+      const result = await validateApiKey();
+
+      expect(workos.apiKeys.validateApiKey).toHaveBeenCalledWith({
+        value: 'invalid_key',
+      });
+      expect(result).toEqual({ apiKey: null });
+    });
+  });
+});

package/src/validate-api-key.ts

@@ -0,0 +1,19 @@
+'use server';
+
+import { getWorkOS } from './workos.js';
+import { headers } from 'next/headers';
+
+export async function validateApiKey() {
+  const headersList = await headers();
+  const authorizationHeader = headersList.get('authorization');
+  if (!authorizationHeader) {
+    return { apiKey: null };
+  }
+
+  const value = authorizationHeader.match(/Bearer\s+(.*)/i)?.[1];
+  if (!value) {
+    return { apiKey: null };
+  }
+
+  return getWorkOS().apiKeys.validateApiKey({ value });
+}

package/src/workos.ts

@@ -2,7 +2,7 @@ import { WorkOS } from '@workos-inc/node';
 import { WORKOS_API_HOSTNAME, WORKOS_API_KEY, WORKOS_API_HTTPS, WORKOS_API_PORT } from './env-variables.js';
 import { lazy } from './utils.js';
 
-export const VERSION = '2.9.0';
+export const VERSION = '2.11.0';
 
 const options = {
   apiHostname: WORKOS_API_HOSTNAME,