fix: export getAuthKitContext as the public auth context accessor

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,14 +125,36 @@ 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',
         'X-Custom': 'value',
       });
     });
   });
+
+  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<null>().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);
+    });
+  });
 });

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<User>;
+  /** 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

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 {

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.',

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 }),

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<Request, Response> {
@@ -22,7 +22,7 @@ export class TanStackStartCookieSessionStorage extends CookieSessionStorage<Requ
     response: Response | undefined,
     headers: Record<string, string>,
   ): Promise<{ response: Response }> {
-    const ctx = getAuthKitContextOrNull();
+    const ctx = getInternalAuthKitContextOrNull();
 
     // When middleware context is available, use it exclusively
     if (ctx?.__setPendingHeader) {

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();