support custom auth api domain (#52)

Author: erquhart

README.md

@@ -88,6 +88,17 @@ authKit.registerRoutes(http);
 export default http;
 ```
 
+## Custom domain
+
+If you've configured a [custom WorkOS authentication API domain](https://workos.com/docs/custom-domains/auth-api), pass it as `apiHostname` when constructing the AuthKit client. This routes the WorkOS SDK through your domain and updates the JWT issuer and JWKS URLs to match.
+
+```ts
+// convex/auth.ts
+export const authKit = new AuthKit<DataModel>(components.workOSAuthKit, {
+  apiHostname: "auth.example.com",
+});
+```
+
 ## Usage
 
 User create/update/delete in WorkOS will be automatically synced by the

eslint.config.js

@@ -10,8 +10,6 @@ export default [
     ignores: [
       "dist/**",
       "example/dist/**",
-      "*.config.{js,mjs,cjs,ts,tsx}",
-      "example/**/*.config.{js,mjs,cjs,ts,tsx}",
       "**/_generated/",
       "initTemplate.mjs",
     ],

src/client/index.test.ts

@@ -0,0 +1,109 @@
+/// <reference types="vite/client" />
+import { describe, test, expect, vi, beforeEach, afterEach } from "vitest";
+import { WorkOS } from "@workos-inc/node";
+import { AuthKit } from "./index.js";
+import type { ComponentApi } from "../component/_generated/component.js";
+
+vi.mock("@workos-inc/node", () => {
+  return {
+    WorkOS: vi.fn().mockImplementation(function () {
+      return {};
+    }),
+  };
+});
+
+const requiredEnv = {
+  WORKOS_CLIENT_ID: "client_test",
+  WORKOS_API_KEY: "sk_test",
+  WORKOS_WEBHOOK_SECRET: "whsec_test",
+};
+
+const fakeComponent = {} as ComponentApi;
+
+describe("AuthKit constructor", () => {
+  beforeEach(() => {
+    for (const [k, v] of Object.entries(requiredEnv)) {
+      process.env[k] = v;
+    }
+  });
+
+  afterEach(() => {
+    for (const k of Object.keys(requiredEnv)) {
+      delete process.env[k];
+    }
+    vi.clearAllMocks();
+  });
+
+  describe("apiHostname", () => {
+    test("forwards apiHostname option to the WorkOS SDK", () => {
+      new AuthKit(fakeComponent, { apiHostname: "auth.example.com" });
+      expect(vi.mocked(WorkOS)).toHaveBeenCalledWith(
+        "sk_test",
+        expect.objectContaining({ apiHostname: "auth.example.com" })
+      );
+    });
+
+    test("forwards undefined when option is not set", () => {
+      new AuthKit(fakeComponent);
+      expect(vi.mocked(WorkOS)).toHaveBeenCalledWith(
+        "sk_test",
+        expect.objectContaining({ apiHostname: undefined })
+      );
+    });
+  });
+
+  test("clientId is forwarded to the WorkOS SDK", () => {
+    new AuthKit(fakeComponent);
+    expect(vi.mocked(WorkOS)).toHaveBeenCalledWith(
+      "sk_test",
+      expect.objectContaining({ clientId: "client_test" })
+    );
+  });
+});
+
+describe("AuthKit.getAuthConfigProviders", () => {
+  beforeEach(() => {
+    for (const [k, v] of Object.entries(requiredEnv)) {
+      process.env[k] = v;
+    }
+  });
+
+  afterEach(() => {
+    for (const k of Object.keys(requiredEnv)) {
+      delete process.env[k];
+    }
+    vi.clearAllMocks();
+  });
+
+  test("falls back to api.workos.com when no custom hostname is set", () => {
+    const authKit = new AuthKit(fakeComponent);
+    const providers = authKit.getAuthConfigProviders();
+    expect(providers[0].issuer).toBe("https://api.workos.com/");
+    expect(providers[0].jwks).toBe(
+      "https://api.workos.com/sso/jwks/client_test"
+    );
+    expect(providers[1].issuer).toBe(
+      "https://api.workos.com/user_management/client_test"
+    );
+    expect(providers[1].jwks).toBe(
+      "https://api.workos.com/sso/jwks/client_test"
+    );
+  });
+
+  test("custom hostname rewrites issuer but not jwks", () => {
+    const authKit = new AuthKit(fakeComponent, {
+      apiHostname: "auth.example.com",
+    });
+    const providers = authKit.getAuthConfigProviders();
+    expect(providers[0].issuer).toBe("https://auth.example.com/");
+    expect(providers[0].jwks).toBe(
+      "https://api.workos.com/sso/jwks/client_test"
+    );
+    expect(providers[1].issuer).toBe(
+      "https://auth.example.com/user_management/client_test"
+    );
+    expect(providers[1].jwks).toBe(
+      "https://api.workos.com/sso/jwks/client_test"
+    );
+  });
+});

src/client/index.ts

@@ -6,7 +6,6 @@ import {
   type HttpRouter,
   createFunctionHandle,
   httpActionGeneric,
-  internalActionGeneric,
   internalMutationGeneric,
 } from "convex/server";
 import type { RunQueryCtx } from "./types.js";
@@ -31,6 +30,7 @@ type Options = {
   authFunctions?: AuthFunctions;
   clientId?: string;
   apiKey?: string;
+  apiHostname?: string;
   webhookSecret?: string;
   webhookPath?: string;
   additionalEventTypes?: WorkOSEvent["event"][];
@@ -97,25 +97,30 @@ export class AuthKit<DataModel extends GenericDataModel> {
       actionSecret: options?.actionSecret ?? process.env.WORKOS_ACTION_SECRET,
       webhookPath: options?.webhookPath ?? "/workos/webhook",
     };
-    this.workos = new WorkOS(this.config.apiKey);
+    this.workos = new WorkOS(this.config.apiKey, {
+      clientId: this.config.clientId,
+      apiHostname: this.config.apiHostname,
+    });
   }
 
-  getAuthConfigProviders = () =>
-    [
+  getAuthConfigProviders = () => {
+    const apiBaseUrl = `https://${this.config.apiHostname ?? "api.workos.com"}`;
+    return [
       {
         type: "customJwt",
-        issuer: `https://api.workos.com/`,
+        issuer: `${apiBaseUrl}/`,
         algorithm: "RS256",
         jwks: `https://api.workos.com/sso/jwks/${this.config.clientId}`,
         applicationID: this.config.clientId,
       },
       {
         type: "customJwt",
-        issuer: `https://api.workos.com/user_management/${this.config.clientId}`,
+        issuer: `${apiBaseUrl}/user_management/${this.config.clientId}`,
         algorithm: "RS256",
         jwks: `https://api.workos.com/sso/jwks/${this.config.clientId}`,
       },
     ] satisfies AuthConfig["providers"];
+  };
 
   async getAuthUser(ctx: RunQueryCtx) {
     const identity = await ctx.auth.getUserIdentity();