Lazy instantiate WorkOS instance on first use (#201)

__tests__/actions.spec.ts

@@ -6,19 +6,20 @@ import {
   refreshAuthAction,
 } from '../src/actions.js';
 import { signOut } from '../src/auth.js';
-import { workos } from '../src/workos.js';
+import { getWorkOS } from '../src/workos.js';
 import { withAuth, refreshSession } from '../src/session.js';
 
 jest.mock('../src/auth.js', () => ({
   signOut: jest.fn().mockResolvedValue(true),
 }));
 
-jest.mock('../src/workos.js', () => ({
-  workos: {
-    organizations: {
-      getOrganization: jest.fn().mockResolvedValue({ id: 'org_123', name: 'Test Org' }),
-    },
+const fakeWorkosInstance = {
+  organizations: {
+    getOrganization: jest.fn().mockResolvedValue({ id: 'org_123', name: 'Test Org' }),
   },
+};
+jest.mock('../src/workos.js', () => ({
+  getWorkOS: jest.fn(() => fakeWorkosInstance),
 }));
 
 jest.mock('../src/session.js', () => ({
@@ -27,6 +28,7 @@ jest.mock('../src/session.js', () => ({
 }));
 
 describe('actions', () => {
+  const workos = getWorkOS();
   describe('checkSessionAction', () => {
     it('should return true for authenticated users', async () => {
       const result = await checkSessionAction();

__tests__/authkit-callback-route.spec.ts

@@ -1,21 +1,24 @@
-import { workos } from '../src/workos.js';
+import { getWorkOS } from '../src/workos.js';
 import { handleAuth } from '../src/authkit-callback-route.js';
 import { NextRequest, NextResponse } from 'next/server';
 
 // Mocked in jest.setup.ts
 import { cookies, headers } from 'next/headers';
 
 // Mock dependencies
-jest.mock('../src/workos', () => ({
-  workos: {
-    userManagement: {
-      authenticateWithCode: jest.fn(),
-      getJwksUrl: jest.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
-    },
+const fakeWorkosInstance = {
+  userManagement: {
+    authenticateWithCode: jest.fn(),
+    getJwksUrl: jest.fn(() => 'https://api.workos.com/sso/jwks/client_1234567890'),
   },
+};
+
+jest.mock('../src/workos', () => ({
+  getWorkOS: jest.fn(() => fakeWorkosInstance),
 }));
 
 describe('authkit-callback-route', () => {
+  const workos = getWorkOS();
   const mockAuthResponse = {
     accessToken: 'access123',
     refreshToken: 'refresh123',

__tests__/get-authorization-url.spec.ts

@@ -1,12 +1,23 @@
 import { describe, it, expect, beforeEach, jest } from '@jest/globals';
 import { getAuthorizationUrl } from '../src/get-authorization-url.js';
 import { headers } from 'next/headers';
-import { workos } from '../src/workos.js';
+import { getWorkOS } from '../src/workos.js';
 
 jest.mock('next/headers');
-jest.mock('../src/workos.js');
+
+// Mock dependencies
+const fakeWorkosInstance = {
+  userManagement: {
+    getAuthorizationUrl: jest.fn(),
+  },
+};
+
+jest.mock('../src/workos', () => ({
+  getWorkOS: jest.fn(() => fakeWorkosInstance),
+}));
 
 describe('getAuthorizationUrl', () => {
+  const workos = getWorkOS();
   beforeEach(() => {
     jest.clearAllMocks();
   });

__tests__/session.spec.ts

@@ -3,7 +3,7 @@ import { cookies, headers } from 'next/headers';
 import { redirect } from 'next/navigation';
 import { generateTestToken } from './test-helpers.js';
 import { withAuth, updateSession, refreshSession, terminateSession, updateSessionMiddleware } from '../src/session.js';
-import { workos } from '../src/workos.js';
+import { getWorkOS } from '../src/workos.js';
 import * as envVariables from '../src/env-variables.js';
 
 import { jwtVerify } from 'jose';
@@ -20,6 +20,8 @@ jest.mock('jose', () => ({
 // logging is disabled by default, flip this to true to still have logs in the console
 const DEBUG = false;
 
+const workos = getWorkOS();
+
 describe('session.ts', () => {
   const mockSession = {
     accessToken: 'access-token',

__tests__/workos.spec.ts

@@ -1,7 +1,8 @@
 import { WorkOS } from '@workos-inc/node';
-import { workos, VERSION } from '../src/workos.js';
+import { getWorkOS, VERSION } from '../src/workos.js';
 
 describe('workos', () => {
+  const workos = getWorkOS();
   beforeEach(() => {
     jest.clearAllMocks();
   });
@@ -44,23 +45,23 @@ describe('workos', () => {
 
     it('uses custom API hostname when provided', async () => {
       process.env.WORKOS_API_HOSTNAME = 'custom.workos.com';
-      const { workos: customWorkos } = await import('../src/workos.js');
+      const { getWorkOS: customWorkos } = await import('../src/workos.js');
 
-      expect(customWorkos.options.apiHostname).toEqual('custom.workos.com');
+      expect(customWorkos().options.apiHostname).toEqual('custom.workos.com');
     });
 
     it('uses custom HTTPS setting when provided', async () => {
       process.env.WORKOS_API_HTTPS = 'false';
-      const { workos: customWorkos } = await import('../src/workos.js');
+      const { getWorkOS: customWorkos } = await import('../src/workos.js');
 
-      expect(customWorkos.options.https).toEqual(false);
+      expect(customWorkos().options.https).toEqual(false);
     });
 
     it('uses custom port when provided', async () => {
       process.env.WORKOS_API_PORT = '8080';
-      const { workos: customWorkos } = await import('../src/workos.js');
+      const { getWorkOS: customWorkos } = await import('../src/workos.js');
 
-      expect(customWorkos.options.port).toEqual(8080);
+      expect(customWorkos().options.port).toEqual(8080);
     });
   });
 });

src/actions.ts

@@ -2,7 +2,7 @@
 
 import { signOut } from './auth.js';
 import { refreshSession, withAuth } from './session.js';
-import { workos } from './workos.js';
+import { getWorkOS } from './workos.js';
 
 /**
  * This action is only accessible to authenticated users,
@@ -18,7 +18,7 @@ export const handleSignOutAction = async ({ returnTo }: { returnTo?: string } =
 };
 
 export const getOrganizationAction = async (organizationId: string) => {
-  return await workos.organizations.getOrganization(organizationId);
+  return await getWorkOS().organizations.getOrganization(organizationId);
 };
 
 export const getAuthAction = async (options?: { ensureSignedIn?: boolean }) => {

src/authkit-callback-route.ts

@@ -5,7 +5,7 @@ import { WORKOS_CLIENT_ID, WORKOS_COOKIE_NAME } from './env-variables.js';
 import { HandleAuthOptions } from './interfaces.js';
 import { encryptSession } from './session.js';
 import { errorResponseWithFallback, redirectWithFallback } from './utils.js';
-import { workos } from './workos.js';
+import { getWorkOS } from './workos.js';
 
 export function handleAuth(options: HandleAuthOptions = {}) {
   const { returnPathname: returnPathnameOption = '/', baseURL, onSuccess, onError } = options;
@@ -28,7 +28,7 @@ export function handleAuth(options: HandleAuthOptions = {}) {
       try {
         // Use the code returned to us by AuthKit and authenticate the user with WorkOS
         const { accessToken, refreshToken, user, impersonator, oauthTokens } =
-          await workos.userManagement.authenticateWithCode({
+          await getWorkOS().userManagement.authenticateWithCode({
             clientId: WORKOS_CLIENT_ID,
             code,
           });

src/get-authorization-url.ts

@@ -1,4 +1,4 @@
-import { workos } from './workos.js';
+import { getWorkOS } from './workos.js';
 import { WORKOS_CLIENT_ID, WORKOS_REDIRECT_URI } from './env-variables.js';
 import { GetAuthURLOptions } from './interfaces.js';
 import { headers } from 'next/headers';
@@ -7,7 +7,7 @@ async function getAuthorizationUrl(options: GetAuthURLOptions = {}) {
   const headersList = await headers();
   const { returnPathname, screenHint, organizationId, redirectUri = headersList.get('x-redirect-uri') } = options;
 
-  return workos.userManagement.getAuthorizationUrl({
+  return getWorkOS().userManagement.getAuthorizationUrl({
     provider: 'authkit',
     clientId: WORKOS_CLIENT_ID,
     redirectUri: redirectUri ?? WORKOS_REDIRECT_URI,

src/interfaces.ts

@@ -1,4 +1,4 @@
-import { OauthTokens, User } from '@workos-inc/node';
+import type { OauthTokens, User } from '@workos-inc/node';
 import { type NextRequest } from 'next/server';
 
 export interface HandleAuthOptions {

src/session.ts

@@ -6,7 +6,7 @@ import { NextRequest, NextResponse } from 'next/server';
 import { jwtVerify, createRemoteJWKSet, decodeJwt } from 'jose';
 import { sealData, unsealData } from 'iron-session';
 import { getCookieOptions } from './cookie.js';
-import { workos } from './workos.js';
+import { getWorkOS } from './workos.js';
 import { WORKOS_CLIENT_ID, WORKOS_COOKIE_PASSWORD, WORKOS_COOKIE_NAME, WORKOS_REDIRECT_URI } from './env-variables.js';
 import { getAuthorizationUrl } from './get-authorization-url.js';
 import {
@@ -21,13 +21,13 @@ import {
 } from './interfaces.js';
 
 import { parse, tokensToRegexp } from 'path-to-regexp';
-import { redirectWithFallback } from './utils.js';
+import { lazy, redirectWithFallback } from './utils.js';
 
 const sessionHeaderName = 'x-workos-session';
 const middlewareHeaderName = 'x-workos-middleware';
 const signUpPathsHeaderName = 'x-sign-up-paths';
 
-const JWKS = createRemoteJWKSet(new URL(workos.userManagement.getJwksUrl(WORKOS_CLIENT_ID)));
+const JWKS = lazy(() => createRemoteJWKSet(new URL(getWorkOS().userManagement.getJwksUrl(WORKOS_CLIENT_ID))));
 
 async function encryptSession(session: Session) {
   return sealData(session, {
@@ -184,11 +184,12 @@ async function updateSession(
 
     const { org_id: organizationIdFromAccessToken } = decodeJwt<AccessToken>(session.accessToken);
 
-    const { accessToken, refreshToken, user, impersonator } = await workos.userManagement.authenticateWithRefreshToken({
-      clientId: WORKOS_CLIENT_ID,
-      refreshToken: session.refreshToken,
-      organizationId: organizationIdFromAccessToken,
-    });
+    const { accessToken, refreshToken, user, impersonator } =
+      await getWorkOS().userManagement.authenticateWithRefreshToken({
+        clientId: WORKOS_CLIENT_ID,
+        refreshToken: session.refreshToken,
+        organizationId: organizationIdFromAccessToken,
+      });
 
     if (options.debug) {
       console.log('Session successfully refreshed');
@@ -270,7 +271,7 @@ async function refreshSession({
   let refreshResult;
 
   try {
-    refreshResult = await workos.userManagement.authenticateWithRefreshToken({
+    refreshResult = await getWorkOS().userManagement.authenticateWithRefreshToken({
       clientId: WORKOS_CLIENT_ID,
       refreshToken: session.refreshToken,
       organizationId: nextOrganizationId ?? organizationIdFromAccessToken,
@@ -389,15 +390,15 @@ async function withAuth(options?: { ensureSignedIn?: boolean }): Promise<UserInf
 async function terminateSession({ returnTo }: { returnTo?: string } = {}) {
   const { sessionId } = await withAuth();
   if (sessionId) {
-    redirect(workos.userManagement.getLogoutUrl({ sessionId, returnTo }));
+    redirect(getWorkOS().userManagement.getLogoutUrl({ sessionId, returnTo }));
   } else {
     redirect(returnTo ?? '/');
   }
 }
 
 async function verifyAccessToken(accessToken: string) {
   try {
-    await jwtVerify(accessToken, JWKS);
+    await jwtVerify(accessToken, JWKS());
     return true;
   } catch {
     return false;

src/utils.ts

@@ -21,3 +21,22 @@ export function errorResponseWithFallback(errorBody: { error: { message: string;
         headers: { 'Content-Type': 'application/json' },
       });
 }
+
+/**
+ * Returns a function that can only be called once.
+ * Subsequent calls will return the result of the first call.
+ * This is useful for lazy initialization.
+ * @param fn - The function to be called once.
+ * @returns A function that can only be called once.
+ */
+export function lazy<T>(fn: () => T): () => T {
+  let called = false;
+  let result: T;
+  return () => {
+    if (!called) {
+      result = fn();
+      called = true;
+    }
+    return result;
+  };
+}

src/workos.ts

@@ -1,5 +1,6 @@
 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 = '1.4.0';
 
@@ -13,7 +14,9 @@ const options = {
   },
 };
 
-// Initialize the WorkOS client
-const workos = new WorkOS(WORKOS_API_KEY, options);
-
-export { workos };
+/**
+ * Create a WorkOS instance with the provided API key and options.
+ * If an instance already exists, it returns the existing instance.
+ * @returns The WorkOS instance.
+ */
+export const getWorkOS = lazy(() => new WorkOS(WORKOS_API_KEY, options));