diff --git a/src/server/context.spec.ts b/src/server/context.spec.ts index 71b3fbd..7610dbf 100644 --- a/src/server/context.spec.ts +++ b/src/server/context.spec.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { describe, it, expect, expectTypeOf, vi, beforeEach } from 'vitest'; // Store original mock state let mockContext: any = undefined; @@ -13,7 +13,7 @@ vi.mock('@tanstack/react-start', () => ({ }, })); -import { getAuthKitContext, getAuthKitContextOrNull } from './context'; +import { getAuthKitContext, getAuthKitContextOrNull, getInternalAuthKitContextOrNull } from './context'; describe('Context Functions', () => { beforeEach(() => { @@ -125,9 +125,9 @@ describe('Context Functions', () => { __setPendingHeader: setPendingHeader, }; - const ctx = getAuthKitContext(); - ctx.__setPendingHeader('Set-Cookie', 'session=abc123'); - ctx.__setPendingHeader('X-Custom', 'value'); + const ctx = getInternalAuthKitContextOrNull(); + ctx?.__setPendingHeader('Set-Cookie', 'session=abc123'); + ctx?.__setPendingHeader('X-Custom', 'value'); expect(pendingHeaders).toEqual({ 'Set-Cookie': 'session=abc123', @@ -135,4 +135,26 @@ describe('Context Functions', () => { }); }); }); + + describe('public type narrowing', () => { + it('hides __setPendingHeader from the public context types', () => { + expectTypeOf(getAuthKitContext).returns.not.toHaveProperty('__setPendingHeader'); + expectTypeOf(getAuthKitContext).returns.toHaveProperty('auth'); + expectTypeOf(getAuthKitContext).returns.toHaveProperty('request'); + expectTypeOf(getAuthKitContext).returns.toHaveProperty('redirectUri'); + expectTypeOf(getInternalAuthKitContextOrNull).returns.exclude().toHaveProperty('__setPendingHeader'); + }); + + it('returns the same underlying context object from public and internal accessors', () => { + mockContext = { + auth: () => ({ user: { id: 'user_123' } }), + request: new Request('http://test.local'), + __setPendingHeader: vi.fn(), + }; + + expect(getAuthKitContext()).toBe(mockContext); + expect(getAuthKitContextOrNull()).toBe(mockContext); + expect(getInternalAuthKitContextOrNull()).toBe(mockContext); + }); + }); }); diff --git a/src/server/context.ts b/src/server/context.ts index a2f3824..65c42c2 100644 --- a/src/server/context.ts +++ b/src/server/context.ts @@ -3,12 +3,24 @@ import type { AuthResult } from '@workos/authkit-session'; import type { User } from '../types.js'; /** - * Internal context shape set by authkitMiddleware. + * Auth context provided by `authkitMiddleware`, available in server + * functions and any middleware that runs after it. */ -export interface AuthKitServerContext { +export interface AuthKitContext { + /** Returns the auth result for the current request. */ auth: () => AuthResult; + /** The original incoming request. */ request: Request; + /** The redirect URI configured on the middleware, if any. */ redirectUri?: string; +} + +/** + * Internal context shape set by authkitMiddleware. Extends the public + * context with the pending-header channel used to flush Set-Cookie + * headers onto the outgoing response. Not part of the public API. + */ +export interface AuthKitServerContext extends AuthKitContext { __setPendingHeader: (key: string, value: string) => void; } @@ -29,8 +41,22 @@ See the documentation for more details: https://github.com/workos/authkit-tansta /** * Gets the AuthKit context from TanStack's global context. * Throws if middleware is not configured. + * + * Use this in your own server functions or middleware to access the + * auth result that `authkitMiddleware` already resolved for the request: + * + * @example + * ```typescript + * import { getAuthKitContext } from '@workos/authkit-tanstack-react-start'; + * + * const myServerFn = createServerFn().handler(async () => { + * const { auth } = getAuthKitContext(); + * const { user } = auth(); + * // ... + * }); + * ``` */ -export function getAuthKitContext(): AuthKitServerContext { +export function getAuthKitContext(): AuthKitContext { const ctx = getGlobalStartContext() as AuthKitServerContext | undefined; // Validate that both auth and request are present (ensures middleware ran correctly) @@ -46,7 +72,17 @@ export function getAuthKitContext(): AuthKitServerContext { * Gracefully handles the case where TanStack Start context is not available * (e.g., when called after args.next() returns in middleware). */ -export function getAuthKitContextOrNull(): AuthKitServerContext | null { +export function getAuthKitContextOrNull(): AuthKitContext | null { + return getInternalAuthKitContextOrNull(); +} + +/** + * Full middleware context including the pending-header channel. + * For sibling server modules only — not exported from the package. + * + * @internal + */ +export function getInternalAuthKitContextOrNull(): AuthKitServerContext | null { try { const ctx = getGlobalStartContext() as AuthKitServerContext | undefined; // Validate that both auth and request are present diff --git a/src/server/index.ts b/src/server/index.ts index e6c09bb..6fe1a3d 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -15,6 +15,8 @@ export type { HandleCallbackOptions, HandleAuthSuccessData, OauthTokens } from ' export { authkitMiddleware, type AuthKitMiddlewareOptions } from './middleware.js'; +export { getAuthKitContext, getAuthKitContextOrNull, type AuthKitContext } from './context.js'; + export { getAuthkit, type AuthService } from './authkit-loader.js'; export { diff --git a/src/server/server-fn-bodies.ts b/src/server/server-fn-bodies.ts index 1691275..bf07ac8 100644 --- a/src/server/server-fn-bodies.ts +++ b/src/server/server-fn-bodies.ts @@ -1,7 +1,7 @@ import type { GetAuthorizationUrlOptions as GetAuthURLOptions, HeadersBag } from '@workos/authkit-session'; import { getRawAuthFromContext, mapAuthToBaseInfo, refreshSession, getRedirectUriFromContext } from './auth-helpers.js'; import { getAuthkit } from './authkit-loader.js'; -import { getAuthKitContextOrNull } from './context.js'; +import { getInternalAuthKitContextOrNull } from './context.js'; import { emitHeadersFrom, forEachHeaderBagEntry } from './headers-bag.js'; import type { NoUserInfo, UserInfo } from './server-functions.js'; @@ -19,7 +19,7 @@ type AuthorizationResult = { * emissions survive as distinct HTTP headers. */ function forwardAuthorizationCookies(result: AuthorizationResult): string { - const ctx = getAuthKitContextOrNull(); + const ctx = getInternalAuthKitContextOrNull(); if (!ctx?.__setPendingHeader) { throw new Error( '[authkit-tanstack-react-start] PKCE cookie could not be set: middleware context unavailable. Ensure authkitMiddleware is registered in your request middleware stack.', diff --git a/src/server/storage.spec.ts b/src/server/storage.spec.ts index 4fb95fa..726920c 100644 --- a/src/server/storage.spec.ts +++ b/src/server/storage.spec.ts @@ -9,7 +9,7 @@ const mockSetPendingHeader = vi.fn(); let mockContextAvailable = true; vi.mock('./context', () => ({ - getAuthKitContextOrNull: () => { + getInternalAuthKitContextOrNull: () => { if (!mockContextAvailable) return null; return { auth: () => ({ user: null }), diff --git a/src/server/storage.ts b/src/server/storage.ts index 07cc881..c670bff 100644 --- a/src/server/storage.ts +++ b/src/server/storage.ts @@ -1,5 +1,5 @@ import { CookieSessionStorage } from '@workos/authkit-session'; -import { getAuthKitContextOrNull } from './context.js'; +import { getInternalAuthKitContextOrNull } from './context.js'; import { parseCookies } from './cookie-utils.js'; export class TanStackStartCookieSessionStorage extends CookieSessionStorage { @@ -22,7 +22,7 @@ export class TanStackStartCookieSessionStorage extends CookieSessionStorage, ): Promise<{ response: Response }> { - const ctx = getAuthKitContextOrNull(); + const ctx = getInternalAuthKitContextOrNull(); // When middleware context is available, use it exclusively if (ctx?.__setPendingHeader) { diff --git a/tests/exports.spec.ts b/tests/exports.spec.ts index 5b0a93f..cb7a195 100644 --- a/tests/exports.spec.ts +++ b/tests/exports.spec.ts @@ -17,6 +17,11 @@ describe('SDK exports', () => { // Middleware expect(exports.authkitMiddleware).toBeDefined(); + // Auth context accessors (public replacement for the middleware's + // inferred downstream context type, lost in the lazy-shell refactor) + expect(exports.getAuthKitContext).toBeDefined(); + expect(exports.getAuthKitContextOrNull).toBeDefined(); + // Error classes re-exported from authkit-session for adopter error handling expect(exports.OAuthStateMismatchError).toBeDefined(); expect(exports.PKCECookieMissingError).toBeDefined();