feat: add redirectUri option to middleware (#29)

* implement

* docs: document redirectUri middleware option

* fix: align redirectUri injection pattern across URL functions

Make getAuthorizationUrl use the same conditional pattern as
getSignInUrl and getSignUpUrl for injecting context redirectUri.

Author: nicknisi

README.md

@@ -510,11 +510,25 @@ function ClaimsDisplay() {
 
 ### Middleware
 
-#### `authkitMiddleware()`
+#### `authkitMiddleware(options?)`
 
 Processes authentication on every request. Validates tokens, refreshes sessions, and provides auth context to server functions.
 
-Already shown in setup, but can be imported separately if needed.
+```typescript
+import { authkitMiddleware } from '@workos/authkit-tanstack-react-start';
+
+// Basic usage
+authkitMiddleware();
+
+// With custom redirect URI (e.g., for Vercel preview deployments)
+authkitMiddleware({
+  redirectUri: 'https://preview-123.example.com/api/auth/callback',
+});
+```
+
+**Options:**
+
+- `redirectUri` - Override the default redirect URI from `WORKOS_REDIRECT_URI`. Useful for dynamic environments like preview deployments.
 
 ## TypeScript
 

src/server/auth-helpers.ts

@@ -19,6 +19,14 @@ export function isAuthConfigured(): boolean {
   return getAuthKitContextOrNull() !== null;
 }
 
+/**
+ * Gets the redirect URI from middleware context if configured.
+ */
+export function getRedirectUriFromContext(): string | undefined {
+  const ctx = getAuthKitContextOrNull();
+  return ctx?.redirectUri;
+}
+
 /**
  * Gets the session with refresh token from the current request.
  * Returns null if no valid session exists.

src/server/context.ts

@@ -8,6 +8,7 @@ import type { User } from '../types.js';
 export interface AuthKitServerContext {
   auth: () => AuthResult<User>;
   request: Request;
+  redirectUri?: string;
   __setPendingHeader: (key: string, value: string) => void;
 }
 

src/server/index.ts

@@ -13,7 +13,7 @@ export {
 export { handleCallbackRoute } from './server.js';
 export type { HandleCallbackOptions, HandleAuthSuccessData, OauthTokens } from './types.js';
 
-export { authkitMiddleware } from './middleware.js';
+export { authkitMiddleware, type AuthKitMiddlewareOptions } from './middleware.js';
 
 export { getAuthkit, type AuthService } from './authkit-loader.js';
 

src/server/middleware.spec.ts

@@ -220,4 +220,56 @@ describe('authkitMiddleware', () => {
       expect(result.response.headers.get('X-New')).toBe('added');
     });
   });
+
+  describe('redirectUri option', () => {
+    it('passes redirectUri to context when provided', async () => {
+      mockAuthkit.withAuth.mockResolvedValue({
+        auth: { user: null },
+        refreshedSessionData: null,
+      });
+
+      authkitMiddleware({ redirectUri: 'https://custom.example.com/callback' });
+
+      const mockRequest = new Request('http://test.local');
+      const mockResponse = new Response('OK', { status: 200 });
+
+      let capturedContext: any = null;
+      const args = {
+        request: mockRequest,
+        next: vi.fn(async ({ context }: any) => {
+          capturedContext = context;
+          return { response: mockResponse };
+        }),
+      };
+
+      await middlewareServerCallback(args);
+
+      expect(capturedContext.redirectUri).toBe('https://custom.example.com/callback');
+    });
+
+    it('passes undefined redirectUri when not provided', async () => {
+      mockAuthkit.withAuth.mockResolvedValue({
+        auth: { user: null },
+        refreshedSessionData: null,
+      });
+
+      authkitMiddleware();
+
+      const mockRequest = new Request('http://test.local');
+      const mockResponse = new Response('OK', { status: 200 });
+
+      let capturedContext: any = null;
+      const args = {
+        request: mockRequest,
+        next: vi.fn(async ({ context }: any) => {
+          capturedContext = context;
+          return { response: mockResponse };
+        }),
+      };
+
+      await middlewareServerCallback(args);
+
+      expect(capturedContext.redirectUri).toBeUndefined();
+    });
+  });
 });

src/server/middleware.ts

@@ -3,6 +3,17 @@ import { getAuthkit, validateConfig } from './authkit-loader.js';
 
 let configValidated = false;
 
+/**
+ * Options for AuthKit middleware.
+ */
+export interface AuthKitMiddlewareOptions {
+  /**
+   * Override the default redirect URI for OAuth callbacks.
+   * Useful for dynamic environments like Vercel preview deployments.
+   */
+  redirectUri?: string;
+}
+
 /**
  * AuthKit middleware for TanStack Start.
  * Validates/refreshes sessions and provides auth context to downstream handlers.
@@ -16,8 +27,16 @@ let configValidated = false;
  *   requestMiddleware: [authkitMiddleware()],
  * }));
  * ```
+ *
+ * @example
+ * ```typescript
+ * // With custom redirect URI
+ * authkitMiddleware({
+ *   redirectUri: 'https://preview.example.com/callback',
+ * })
+ * ```
  */
-export const authkitMiddleware = () => {
+export const authkitMiddleware = (options?: AuthKitMiddlewareOptions) => {
   return createMiddleware().server(async (args) => {
     const authkit = await getAuthkit();
 
@@ -33,6 +52,7 @@ export const authkitMiddleware = () => {
       context: {
         auth: () => auth,
         request: args.request,
+        redirectUri: options?.redirectUri,
         __setPendingHeader: (key: string, value: string) => {
           // Use append for Set-Cookie to support multiple cookies
           if (key.toLowerCase() === 'set-cookie') {

src/server/server-functions.ts

@@ -1,7 +1,7 @@
 import { redirect } from '@tanstack/react-router';
 import { createServerFn } from '@tanstack/react-start';
 import type { Impersonator, User } from '../types.js';
-import { getRawAuthFromContext, refreshSession } from './auth-helpers.js';
+import { getRawAuthFromContext, refreshSession, getRedirectUriFromContext } from './auth-helpers.js';
 import { getAuthkit } from './authkit-loader.js';
 
 // Type-only import - safe for bundling
@@ -159,6 +159,16 @@ export const getAuthorizationUrl = createServerFn({ method: 'GET' })
   .inputValidator((options?: GetAuthURLOptions) => options)
   .handler(async ({ data: options = {} }) => {
     const authkit = await getAuthkit();
+    const contextRedirectUri = getRedirectUriFromContext();
+
+    // Only inject context redirectUri if it exists and user didn't provide one
+    if (contextRedirectUri && !options.redirectUri) {
+      return authkit.getAuthorizationUrl({
+        ...options,
+        redirectUri: contextRedirectUri,
+      });
+    }
+
     return authkit.getAuthorizationUrl(options);
   });
 
@@ -185,7 +195,17 @@ export const getSignInUrl = createServerFn({ method: 'GET' })
   .inputValidator((data?: string | SignInUrlOptions) => data)
   .handler(async ({ data }) => {
     const options = typeof data === 'string' ? { returnPathname: data } : data;
+    const contextRedirectUri = getRedirectUriFromContext();
     const authkit = await getAuthkit();
+
+    // Only inject context redirectUri if it exists and user didn't provide one
+    if (contextRedirectUri && !options?.redirectUri) {
+      return authkit.getSignInUrl({
+        ...options,
+        redirectUri: contextRedirectUri,
+      });
+    }
+
     return authkit.getSignInUrl(options);
   });
 
@@ -209,7 +229,17 @@ export const getSignUpUrl = createServerFn({ method: 'GET' })
   .inputValidator((data?: string | SignInUrlOptions) => data)
   .handler(async ({ data }) => {
     const options = typeof data === 'string' ? { returnPathname: data } : data;
+    const contextRedirectUri = getRedirectUriFromContext();
     const authkit = await getAuthkit();
+
+    // Only inject context redirectUri if it exists and user didn't provide one
+    if (contextRedirectUri && !options?.redirectUri) {
+      return authkit.getSignUpUrl({
+        ...options,
+        redirectUri: contextRedirectUri,
+      });
+    }
+
     return authkit.getSignUpUrl(options);
   });