allow configuring session storage on callback route (#22)

This PR addresses issue #63 where users deploying to serverless
environments (like AWS Lambda) experience errors in the callback route
during cold starts. The error occurs because the session storage
configuration happens in the root loader, but during cold starts the
callback route might be hit first, resulting in:

```
Error: SessionStorage was never configured. Did you forget to call
configureSessionStorage in your root loader?
```

## Changes

- Updated the `authLoader` function to accept the same `storage` and
`cookie` configuration options that `authkitLoader` already supports
- Modified `HandleAuthOptions` type to support these new configuration
options
- Changed the implementation to use `configureSessionStorage` directly
in the callback route

## How It Works

When a cold start happens in a serverless environment, the callback
route can now configure session storage itself through the same
interface already available in `authkitLoader`. This allows developers
to provide consistent session storage configuration to both routes,
ensuring the application works correctly regardless of which route is
hit first.

## Usage

For applications deployed to serverless environments, you can now use
this pattern:

```typescript
// Create a shared configuration function
function getAuthStorage() {
  return {
    storage: createCookieSessionStorage({...}),
    cookie: { name: "my-session" }
  };
}

// In your root loader
export const loader = (args) => authkitLoader(args, {
  ...getAuthStorage(),
  // Other options...
});

// In your callback route
export const loader = authLoader({
  ...getAuthStorage(),
  // Other options...
});
```

This ensures consistent session configuration across routes and prevents
cold start errors.

## README changes

This PR also adds a section describing sessoion storage configuration to
the README, and notes the important case to configure in both laoders
(when custom session storage is used).

Fixes #63.

Author: nicknisi

README.md

@@ -280,3 +280,75 @@ export const loader = (args: LoaderFunctionArgs) =>
     { debug: true },
   );
 ```
+
+## Customizing Session Storage
+
+By default, AuthKit for React Router uses cookie-based session storage with these settings:
+
+```typescript
+{
+  name: "wos-session", // Default or WORKOS_COOKIE_NAME if set
+  path: "/",
+  httpOnly: true,
+  secure: true, // When redirect URI uses HTTPS
+  sameSite: "lax",
+  maxAge: 34560000, // 400 days (configurable via WORKOS_COOKIE_MAX_AGE)
+  secrets: [/* your cookie password, configurable via WORKOS_COOKIE_PASSWORD */],
+}
+```
+
+### Custom Session Storage
+
+You can provide your own session storage implementation to both `authkitLoader` and `authLoader`:
+
+```typescript
+import { createMemorySessionStorage } from "@react-router/node";
+import { authkitLoader, authLoader } from "@workos-inc/authkit-react-router";
+
+// Create memory-based session storage
+const memoryStorage = createMemorySessionStorage({
+  cookie: {
+    name: "auth-session",
+    secrets: ["test-secret"],
+    sameSite: "lax",
+    path: "/",
+    httpOnly: true,
+    secure: false, // Use false for testing
+    maxAge: 60 * 60 * 24 // 1 day
+  }
+});
+
+// In your root loader
+export const loader = (args) => authkitLoader(args, {
+  storage: memoryStorage,
+  cookie: { name: "auth-session" }
+});
+
+// In your callback route
+export const loader = authLoader({
+  storage: memoryStorage,
+  cookie: { name: "auth-session" }
+});
+```
+
+For code reuse and consistency, consider using a shared function:
+
+```typescript
+// app/lib/session.ts
+export function getAuthStorage() {
+  const storage = createCookieSessionStorage({/* config */});
+  return { storage, cookie: { name: "my-custom-session" } };
+}
+
+// Then in your routes
+import { getAuthStorage } from "~/lib/session";
+export const loader = (args) => authkitLoader(args, {
+  ...getAuthStorage(),
+  // Other options...
+});
+```
+
+> [!NOTE]
+>When deploying to serverless environments like AWS Lambda, ensure you pass the same storage configuration to both your main routes and the callback route to handle cold starts properly.
+
+AuthKit works with any session storage that implements React Router's `SessionStorage` interface, including Redis-based or database-backed implementations.

src/authkit-callback-route.ts

@@ -2,13 +2,14 @@ import { LoaderFunctionArgs, data, redirect } from 'react-router';
 import { getConfig } from './config.js';
 import { HandleAuthOptions } from './interfaces.js';
 import { encryptSession } from './session.js';
-import { getSessionStorage } from './sessionStorage.js';
+import { configureSessionStorage } from './sessionStorage.js';
 import { getWorkOS } from './workos.js';
 
 export function authLoader(options: HandleAuthOptions = {}) {
   return async function loader({ request }: LoaderFunctionArgs) {
-    const { getSession, commitSession, cookieName } = await getSessionStorage();
-    const { returnPathname: returnPathnameOption = '/', onSuccess } = options;
+    const { storage, cookie, returnPathname: returnPathnameOption = '/', onSuccess } = options;
+    const cookieName = cookie?.name ?? getConfig('cookieName');
+    const { getSession, commitSession } = await configureSessionStorage({ storage, cookieName });
 
     const url = new URL(request.url);
 

src/interfaces.ts

@@ -3,10 +3,19 @@ import type { OauthTokens, User } from '@workos-inc/node';
 
 export type DataWithResponseInit<T> = ReturnType<typeof data<T>>;
 
-export interface HandleAuthOptions {
+export type HandleAuthOptions = {
   returnPathname?: string;
   onSuccess?: (data: AuthLoaderSuccessData) => void | Promise<void>;
-}
+} & (
+  | {
+      storage?: never;
+      cookie?: SessionIdStorageStrategy['cookie'];
+    }
+  | {
+      storage: SessionStorage;
+      cookie: SessionIdStorageStrategy['cookie'];
+    }
+);
 
 export interface AuthLoaderSuccessData {
   accessToken: string;