@hebo-ai/gateway 0.3.0 → 0.4.0-alpha.1

package/README.md

@@ -288,11 +288,23 @@ const gw = gateway({
      * @returns Optional RequestPatch to merge into headers / override body.
      * Returning a Response stops execution of the endpoint.
      */
-    before: async (ctx: { request: Request }): Promise<RequestPatch | Response | void> => {
+    onRequest: async (ctx: { request: Request }): Promise<RequestPatch | Response | void> => {
       // Example Use Cases:
-      // - Transform request body
       // - Verify authentication
       // - Enforce rate limits
+      return undefined;
+    },
+    /**
+     * Runs after body is parsed & validated.
+     * @param ctx.body Parsed request body.
+     * @returns Replacement parsed body, or undefined to keep original body unchanged.
+     */
+    before: async (ctx: {
+      body: ChatCompletionsBody | EmbeddingsBody;
+      operation: "text" | "embeddings";
+    }): Promise<ChatCompletionsBody | EmbeddingsBody | void> => {
+      // Example Use Cases:
+      // - Transform request body
       // - Observability integration
       return undefined;
     },
@@ -337,18 +349,29 @@ const gw = gateway({
      * @returns Modified result, or undefined to keep original.
      */
     after: async (ctx: {
-      result: object | ReadableStream<Uint8Array>
-    }): Promise<object | ReadableStream<Uint8Array> | void> => {
+      result: ChatCompletions  | ReadableStream<ChatCompletionsChunk | OpenAIError> | Embeddings | object
+    }): Promise<ChatCompletions  | ReadableStream<ChatCompletionsChunk | OpenAIError> | Embeddings | object | void> => {
       // Example Use Cases:
       // - Transform result
       // - Result logging
       return undefined;
     },
+    /**
+     * Runs after the gateway has produced the final Response.
+     * @param ctx.response Response object returned by the lifecycle.
+     * @returns Replacement response, or undefined to keep original.
+     */
+    onResponse: async (ctx: { response: Response }): Promise<Response | void> => {
+      // Example Use Cases:
+      // - Add response headers
+      // - Replace or redact response payload
+      return undefined;
+    },
   },
 });
 ```
 
-The `ctx` object is **readonly for core fields**. Use return values to override request / result and to provide modelId / provider instances.
+The `ctx` object is **readonly for core fields**. Use return values to override request / parsed body / result / response and to provide modelId / provider instances.
 
 > [!TIP]
 > To pass data between hooks, use `ctx.state`. It’s a per-request mutable bag in which you can stash things like auth info, routing decisions, timers, or trace IDs and read them later again in any of the other hooks.

package/dist/endpoints/chat-completions/converters.d.ts

@@ -25,7 +25,7 @@ export declare const convertToToolSet: (tools: ChatCompletionsTool[] | undefined
 export declare const convertToToolChoice: (toolChoice: ChatCompletionsToolChoice | undefined) => ToolChoice<ToolSet> | undefined;
 export declare function toChatCompletions(result: GenerateTextResult<ToolSet, Output.Output>, model: string): ChatCompletions;
 export declare function toChatCompletionsResponse(result: GenerateTextResult<ToolSet, Output.Output>, model: string, responseInit?: ResponseInit): Response;
-export declare function toChatCompletionsStream(result: StreamTextResult<ToolSet, Output.Output>, model: string): ReadableStream<Uint8Array>;
+export declare function toChatCompletionsStream(result: StreamTextResult<ToolSet, Output.Output>, model: string): ReadableStream<ChatCompletionsChunk | OpenAIError>;
 export declare function toChatCompletionsStreamResponse(result: StreamTextResult<ToolSet, Output.Output>, model: string, responseInit?: ResponseInit): Response;
 export declare class ChatCompletionsStream extends TransformStream<TextStreamPart<ToolSet>, ChatCompletionsChunk | OpenAIError> {
     constructor(model: string);

package/dist/endpoints/chat-completions/converters.js

@@ -1,5 +1,5 @@
 import { convertBase64ToUint8Array } from "@ai-sdk/provider-utils";
-import { jsonSchema, JsonToSseTransformStream, tool } from "ai";
+import { jsonSchema, tool } from "ai";
 import { GatewayError } from "../../errors/gateway";
 import { OpenAIError, toOpenAIError } from "../../errors/openai";
 import { toResponse } from "../../utils/response";
@@ -241,10 +241,7 @@ export function toChatCompletionsResponse(result, model, responseInit) {
     return toResponse(toChatCompletions(result, model), responseInit);
 }
 export function toChatCompletionsStream(result, model) {
-    return result.fullStream
-        .pipeThrough(new ChatCompletionsStream(model))
-        .pipeThrough(new JsonToSseTransformStream())
-        .pipeThrough(new TextEncoderStream());
+    return result.fullStream.pipeThrough(new ChatCompletionsStream(model));
 }
 export function toChatCompletionsStreamResponse(result, model, responseInit) {
     return toResponse(toChatCompletionsStream(result, model), responseInit);
@@ -302,6 +299,7 @@ export class ChatCompletionsStream extends TransformStream {
                     }
                     case "error": {
                         const error = part.error;
+                        // FUTURE mask in production mode and return responseID
                         controller.enqueue(toOpenAIError(error));
                         break;
                     }

package/dist/endpoints/chat-completions/handler.js

@@ -30,13 +30,14 @@ export const chatCompletions = (config) => {
             throw new GatewayError(z.prettifyError(parsed.error), 400);
         }
         ctx.body = parsed.data;
+        ctx.operation = "text";
+        ctx.body = (await hooks?.before?.(ctx)) ?? ctx.body;
         // Resolve model + provider (hooks may override defaults).
         let inputs, stream;
-        ({ model: ctx.modelId, stream, ...inputs } = parsed.data);
+        ({ model: ctx.modelId, stream, ...inputs } = ctx.body);
         ctx.resolvedModelId =
             (await hooks?.resolveModelId?.(ctx)) ?? ctx.modelId;
         logger.debug(`[chat] resolved ${ctx.modelId} to ${ctx.resolvedModelId}`);
-        ctx.operation = "text";
         const override = await hooks?.resolveProvider?.(ctx);
         ctx.provider =
             override ??
@@ -79,7 +80,7 @@ export const chatCompletions = (config) => {
                     throw new DOMException("Upstream failed", "AbortError");
                 },
                 timeout: {
-                    chunkMs: 5 * 60 * 1000,
+                    totalMs: 5 * 60 * 1000,
                 },
                 experimental_include: {
                     requestBody: false,
@@ -88,7 +89,8 @@ export const chatCompletions = (config) => {
                 ...textOptions,
             });
             markPerf(ctx.request, "aiSdkEnd");
-            return toChatCompletionsStream(result, ctx.modelId);
+            ctx.result = toChatCompletionsStream(result, ctx.modelId);
+            return (await hooks?.after?.(ctx)) ?? ctx.result;
         }
         const result = await generateText({
             model: languageModelWithMiddleware,
@@ -104,7 +106,8 @@ export const chatCompletions = (config) => {
         });
         markPerf(ctx.request, "aiSdkEnd");
         logger.trace({ requestId: resolveRequestId(ctx.request), result }, "[chat] AI SDK result");
-        return toChatCompletions(result, ctx.modelId);
+        ctx.result = toChatCompletions(result, ctx.modelId);
+        return (await hooks?.after?.(ctx)) ?? ctx.result;
     };
     return { handler: winterCgHandler(handler, config) };
 };

package/dist/endpoints/embeddings/handler.js

@@ -30,13 +30,14 @@ export const embeddings = (config) => {
             throw new GatewayError(z.prettifyError(parsed.error), 400);
         }
         ctx.body = parsed.data;
+        ctx.operation = "embeddings";
+        ctx.body = (await hooks?.before?.(ctx)) ?? ctx.body;
         // Resolve model + provider (hooks may override defaults).
         let inputs;
-        ({ model: ctx.modelId, ...inputs } = parsed.data);
+        ({ model: ctx.modelId, ...inputs } = ctx.body);
         ctx.resolvedModelId =
             (await hooks?.resolveModelId?.(ctx)) ?? ctx.modelId;
         logger.debug(`[embeddings] resolved ${ctx.modelId} to ${ctx.resolvedModelId}`);
-        ctx.operation = "embeddings";
         const override = await hooks?.resolveProvider?.(ctx);
         ctx.provider =
             override ??
@@ -67,7 +68,8 @@ export const embeddings = (config) => {
         });
         markPerf(ctx.request, "aiSdkEnd");
         logger.trace({ requestId: resolveRequestId(ctx.request), result }, "[embeddings] AI SDK result");
-        return toEmbeddings(result, ctx.modelId);
+        ctx.result = toEmbeddings(result, ctx.modelId);
+        return (await hooks?.after?.(ctx)) ?? ctx.result;
     };
     return { handler: winterCgHandler(handler, config) };
 };

package/dist/errors/openai.js

@@ -27,6 +27,7 @@ export function toOpenAIErrorResponse(error, responseInit) {
     let message;
     if (shouldMask) {
         const requestId = resolveRequestId(responseInit);
+        // FUTURE: always attach requestId to errors (masked and unmasked)
         message = `${STATUS_CODE(meta.status)} (${requestId})`;
     }
     else {

package/dist/lifecycle.d.ts

@@ -1,2 +1,2 @@
 import type { GatewayConfig, GatewayContext } from "./types";
-export declare const winterCgHandler: (run: (ctx: GatewayContext) => Promise<object | ReadableStream<Uint8Array>>, config: GatewayConfig) => (request: Request, state?: Record<string, unknown>) => Promise<Response>;
+export declare const winterCgHandler: (run: (ctx: GatewayContext) => Promise<object | ReadableStream<object>>, config: GatewayConfig) => (request: Request, state?: Record<string, unknown>) => Promise<Response>;

package/dist/lifecycle.js

@@ -9,23 +9,19 @@ export const winterCgHandler = (run, config) => {
     const parsedConfig = parseConfig(config);
     const core = async (ctx) => {
         try {
-            const before = await parsedConfig.hooks?.before?.(ctx);
-            if (before) {
-                if (before instanceof Response) {
-                    ctx.response = before;
+            const onRequest = await parsedConfig.hooks?.onRequest?.(ctx);
+            if (onRequest) {
+                if (onRequest instanceof Response) {
+                    ctx.response = onRequest;
                     return;
                 }
-                ctx.request = maybeApplyRequestPatch(ctx.request, before);
-            }
-            ctx.result = await run(ctx);
-            const after = await parsedConfig.hooks?.after?.(ctx);
-            if (after)
-                ctx.result = after;
-            if (ctx.result instanceof Response) {
-                ctx.response = ctx.result;
-                return;
+                ctx.request = maybeApplyRequestPatch(ctx.request, onRequest);
             }
+            ctx.result = (await run(ctx));
             ctx.response = toResponse(ctx.result, prepareResponseInit(ctx.request));
+            const onResponse = await parsedConfig.hooks?.onResponse?.(ctx);
+            if (onResponse)
+                ctx.response = onResponse;
         }
         catch (error) {
             logger.error({

package/dist/types.d.ts

@@ -1,11 +1,12 @@
 import type { ProviderV3 } from "@ai-sdk/provider";
-import type { ChatCompletionsBody } from "./endpoints/chat-completions/schema";
-import type { EmbeddingsBody } from "./endpoints/embeddings/schema";
+import type { ChatCompletions, ChatCompletionsBody, ChatCompletionsChunk } from "./endpoints/chat-completions/schema";
+import type { Embeddings, EmbeddingsBody } from "./endpoints/embeddings/schema";
+import type { OpenAIError } from "./errors/openai";
 import type { Logger, LoggerConfig } from "./logger";
 import type { ModelCatalog, ModelId } from "./models/types";
 import type { ProviderId, ProviderRegistry } from "./providers/types";
 /**
- * Request overrides returned from the `before` hook.
+ * Request overrides returned from the `onRequest` hook.
  */
 export type RequestPatch = {
     /**
@@ -64,7 +65,7 @@ export type GatewayContext = {
     /**
      * Result returned by the handler (pre-response).
      */
-    result?: object | ReadableStream<Uint8Array>;
+    result?: ChatCompletions | ReadableStream<ChatCompletionsChunk | OpenAIError> | Embeddings | object;
     /**
      * Final response returned by the lifecycle.
      */
@@ -77,10 +78,12 @@ export type HookContext = Omit<Readonly<GatewayContext>, "state"> & {
     state: GatewayContext["state"];
 };
 type RequiredHookContext<K extends keyof GatewayContext> = Omit<HookContext, K> & Required<Pick<HookContext, K>>;
-export type BeforeHookContext = RequiredHookContext<"request">;
+export type OnRequestHookContext = RequiredHookContext<"request">;
+export type BeforeHookContext = RequiredHookContext<"request" | "body" | "operation">;
 export type ResolveModelHookContext = RequiredHookContext<"request" | "body" | "modelId">;
 export type ResolveProviderHookContext = RequiredHookContext<"request" | "body" | "modelId" | "resolvedModelId" | "operation">;
 export type AfterHookContext = RequiredHookContext<"request" | "result" | "provider" | "resolvedModelId" | "operation">;
+export type OnResponseHookContext = RequiredHookContext<"request" | "response">;
 /**
  * Hooks to plugin to the gateway lifecycle.
  */
@@ -90,7 +93,12 @@ export type GatewayHooks = {
      * @returns Optional RequestPatch to merge into headers / override body,
      * or Response to short-circuit the request.
      */
-    before?: (ctx: BeforeHookContext) => void | RequestPatch | Response | Promise<void | RequestPatch | Response>;
+    onRequest?: (ctx: OnRequestHookContext) => void | RequestPatch | Response | Promise<void | RequestPatch | Response>;
+    /**
+     * Runs after request JSON is parsed and validated for chat completions / embeddings.
+     * @returns Replacement parsed body, or undefined to keep original.
+     */
+    before?: (ctx: BeforeHookContext) => void | ChatCompletionsBody | EmbeddingsBody | Promise<void | ChatCompletionsBody | EmbeddingsBody>;
     /**
      * Maps a user-provided model ID or alias to a canonical ID.
      * @returns Canonical model ID or undefined to keep original.
@@ -103,9 +111,14 @@ export type GatewayHooks = {
     resolveProvider?: (ctx: ResolveProviderHookContext) => ProviderV3 | void | Promise<ProviderV3 | void>;
     /**
      * Runs after the endpoint handler.
-     * @returns Response to replace, or undefined to keep original.
+     * @returns Result to replace, or undefined to keep original.
+     */
+    after?: (ctx: AfterHookContext) => void | ChatCompletions | ReadableStream<ChatCompletionsChunk | OpenAIError> | Embeddings | Promise<void | ChatCompletions | ReadableStream<ChatCompletionsChunk | OpenAIError> | Embeddings>;
+    /**
+     * Runs after the lifecycle has produced the final Response.
+     * @returns Replacement Response, or undefined to keep original.
      */
-    after?: (ctx: AfterHookContext) => void | object | ReadableStream<Uint8Array> | Promise<void | object | ReadableStream<Uint8Array>>;
+    onResponse?: (ctx: OnResponseHookContext) => void | Response | Promise<void | Response>;
 };
 /**
  * Main configuration object for the gateway.

package/dist/utils/response.d.ts

@@ -1,3 +1,3 @@
 export declare const prepareResponseInit: (request: Request) => ResponseInit;
 export declare const mergeResponseInit: (defaultHeaders: HeadersInit, responseInit?: ResponseInit) => ResponseInit;
-export declare const toResponse: (result: ReadableStream<Uint8Array> | Uint8Array<ArrayBuffer> | object | string, responseInit?: ResponseInit) => Response;
+export declare const toResponse: (result: ReadableStream | Uint8Array<ArrayBuffer> | object | string, responseInit?: ResponseInit) => Response;

package/dist/utils/response.js

@@ -1,5 +1,17 @@
 import { REQUEST_ID_HEADER, resolveRequestId } from "./headers";
 const TEXT_ENCODER = new TextEncoder();
+class JsonToSseTransformStream extends TransformStream {
+    constructor() {
+        super({
+            transform(part, controller) {
+                controller.enqueue(`data: ${JSON.stringify(part)}\n\n`);
+            },
+            flush(controller) {
+                controller.enqueue("data: [DONE]\n\n");
+            },
+        });
+    }
+}
 export const prepareResponseInit = (request) => ({
     headers: { [REQUEST_ID_HEADER]: resolveRequestId(request.headers) },
 });
@@ -20,7 +32,10 @@ export const mergeResponseInit = (defaultHeaders, responseInit) => {
 export const toResponse = (result, responseInit) => {
     let body;
     const isStream = result instanceof ReadableStream;
-    if (isStream || result instanceof Uint8Array) {
+    if (isStream) {
+        body = result.pipeThrough(new JsonToSseTransformStream()).pipeThrough(new TextEncoderStream());
+    }
+    else if (result instanceof Uint8Array) {
         body = result;
     }
     else if (typeof result === "string") {
@@ -44,7 +59,7 @@ export const toResponse = (result, responseInit) => {
         ? {
             "content-type": "text/event-stream",
             "cache-control": "no-cache",
-            Connection: "keep-alive",
+            connection: "keep-alive",
         }
         : {
             "content-type": "application/json",

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@hebo-ai/gateway",
-  "version": "0.3.0",
+  "version": "0.4.0-alpha.1",
   "description": "AI gateway as a framework. For full control over models, routing & lifecycle. OpenAI-compatible /chat/completions, /embeddings & /models.",
   "keywords": [
     "ai",

package/src/endpoints/chat-completions/converters.ts

@@ -18,7 +18,7 @@ import type {
 } from "ai";
 
 import { convertBase64ToUint8Array } from "@ai-sdk/provider-utils";
-import { jsonSchema, JsonToSseTransformStream, tool } from "ai";
+import { jsonSchema, tool } from "ai";
 
 import type {
   ChatCompletionsToolCall,
@@ -368,11 +368,8 @@ export function toChatCompletionsResponse(
 export function toChatCompletionsStream(
   result: StreamTextResult<ToolSet, Output.Output>,
   model: string,
-): ReadableStream<Uint8Array> {
-  return result.fullStream
-    .pipeThrough(new ChatCompletionsStream(model))
-    .pipeThrough(new JsonToSseTransformStream())
-    .pipeThrough(new TextEncoderStream());
+): ReadableStream<ChatCompletionsChunk | OpenAIError> {
+  return result.fullStream.pipeThrough(new ChatCompletionsStream(model));
 }
 
 export function toChatCompletionsStreamResponse(
@@ -476,6 +473,7 @@ export class ChatCompletionsStream extends TransformStream<
 
           case "error": {
             const error = part.error;
+            // FUTURE mask in production mode and return responseID
             controller.enqueue(toOpenAIError(error));
             break;
           }

package/src/endpoints/chat-completions/handler.ts

@@ -2,6 +2,8 @@ import { generateText, streamText, wrapLanguageModel } from "ai";
 import * as z from "zod/mini";
 
 import type {
+  AfterHookContext,
+  BeforeHookContext,
   GatewayConfig,
   Endpoint,
   GatewayContext,
@@ -43,15 +45,17 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
     }
     ctx.body = parsed.data;
 
+    ctx.operation = "text";
+    ctx.body = (await hooks?.before?.(ctx as BeforeHookContext)) ?? ctx.body;
+
     // Resolve model + provider (hooks may override defaults).
     let inputs, stream;
-    ({ model: ctx.modelId, stream, ...inputs } = parsed.data);
+    ({ model: ctx.modelId, stream, ...inputs } = ctx.body);
 
     ctx.resolvedModelId =
       (await hooks?.resolveModelId?.(ctx as ResolveModelHookContext)) ?? ctx.modelId;
     logger.debug(`[chat] resolved ${ctx.modelId} to ${ctx.resolvedModelId}`);
 
-    ctx.operation = "text";
     const override = await hooks?.resolveProvider?.(ctx as ResolveProviderHookContext);
     ctx.provider =
       override ??
@@ -101,7 +105,7 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
           throw new DOMException("Upstream failed", "AbortError");
         },
         timeout: {
-          chunkMs: 5 * 60 * 1000,
+          totalMs: 5 * 60 * 1000,
         },
         experimental_include: {
           requestBody: false,
@@ -111,7 +115,9 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
       });
       markPerf(ctx.request, "aiSdkEnd");
 
-      return toChatCompletionsStream(result, ctx.modelId);
+      ctx.result = toChatCompletionsStream(result, ctx.modelId);
+
+      return (await hooks?.after?.(ctx as AfterHookContext)) ?? ctx.result;
     }
 
     const result = await generateText({
@@ -130,7 +136,9 @@ export const chatCompletions = (config: GatewayConfig): Endpoint => {
 
     logger.trace({ requestId: resolveRequestId(ctx.request), result }, "[chat] AI SDK result");
 
-    return toChatCompletions(result, ctx.modelId);
+    ctx.result = toChatCompletions(result, ctx.modelId);
+
+    return (await hooks?.after?.(ctx as AfterHookContext)) ?? ctx.result;
   };
 
   return { handler: winterCgHandler(handler, config) };

package/src/endpoints/embeddings/handler.ts

@@ -2,6 +2,8 @@ import { embedMany, wrapEmbeddingModel } from "ai";
 import * as z from "zod/mini";
 
 import type {
+  AfterHookContext,
+  BeforeHookContext,
   GatewayConfig,
   Endpoint,
   GatewayContext,
@@ -43,15 +45,17 @@ export const embeddings = (config: GatewayConfig): Endpoint => {
     }
     ctx.body = parsed.data;
 
+    ctx.operation = "embeddings";
+    ctx.body = (await hooks?.before?.(ctx as BeforeHookContext)) ?? ctx.body;
+
     // Resolve model + provider (hooks may override defaults).
     let inputs;
-    ({ model: ctx.modelId, ...inputs } = parsed.data);
+    ({ model: ctx.modelId, ...inputs } = ctx.body);
 
     ctx.resolvedModelId =
       (await hooks?.resolveModelId?.(ctx as ResolveModelHookContext)) ?? ctx.modelId;
     logger.debug(`[embeddings] resolved ${ctx.modelId} to ${ctx.resolvedModelId}`);
 
-    ctx.operation = "embeddings";
     const override = await hooks?.resolveProvider?.(ctx as ResolveProviderHookContext);
     ctx.provider =
       override ??
@@ -94,7 +98,9 @@ export const embeddings = (config: GatewayConfig): Endpoint => {
       "[embeddings] AI SDK result",
     );
 
-    return toEmbeddings(result, ctx.modelId);
+    ctx.result = toEmbeddings(result, ctx.modelId);
+
+    return (await hooks?.after?.(ctx as AfterHookContext)) ?? ctx.result;
   };
 
   return { handler: winterCgHandler(handler, config) };

package/src/errors/openai.ts

@@ -35,6 +35,7 @@ export function toOpenAIErrorResponse(error: unknown, responseInit?: ResponseIni
   let message;
   if (shouldMask) {
     const requestId = resolveRequestId(responseInit);
+    // FUTURE: always attach requestId to errors (masked and unmasked)
     message = `${STATUS_CODE(meta.status)} (${requestId})`;
   } else {
     message = meta.message;

package/src/lifecycle.ts

@@ -1,4 +1,9 @@
-import type { AfterHookContext, BeforeHookContext, GatewayConfig, GatewayContext } from "./types";
+import type {
+  GatewayConfig,
+  GatewayContext,
+  OnRequestHookContext,
+  OnResponseHookContext,
+} from "./types";
 
 import { parseConfig } from "./config";
 import { toOpenAIErrorResponse } from "./errors/openai";
@@ -9,32 +14,27 @@ import { maybeApplyRequestPatch, prepareRequestHeaders } from "./utils/request";
 import { prepareResponseInit, toResponse } from "./utils/response";
 
 export const winterCgHandler = (
-  run: (ctx: GatewayContext) => Promise<object | ReadableStream<Uint8Array>>,
+  run: (ctx: GatewayContext) => Promise<object | ReadableStream<object>>,
   config: GatewayConfig,
 ) => {
   const parsedConfig = parseConfig(config);
 
   const core = async (ctx: GatewayContext): Promise<void> => {
     try {
-      const before = await parsedConfig.hooks?.before?.(ctx as BeforeHookContext);
-      if (before) {
-        if (before instanceof Response) {
-          ctx.response = before;
+      const onRequest = await parsedConfig.hooks?.onRequest?.(ctx as OnRequestHookContext);
+      if (onRequest) {
+        if (onRequest instanceof Response) {
+          ctx.response = onRequest;
           return;
         }
-        ctx.request = maybeApplyRequestPatch(ctx.request, before);
+        ctx.request = maybeApplyRequestPatch(ctx.request, onRequest);
       }
 
-      ctx.result = await run(ctx);
+      ctx.result = (await run(ctx)) as typeof ctx.result;
+      ctx.response = toResponse(ctx.result!, prepareResponseInit(ctx.request));
 
-      const after = await parsedConfig.hooks?.after?.(ctx as AfterHookContext);
-      if (after) ctx.result = after;
-
-      if (ctx.result instanceof Response) {
-        ctx.response = ctx.result;
-        return;
-      }
-      ctx.response = toResponse(ctx.result, prepareResponseInit(ctx.request));
+      const onResponse = await parsedConfig.hooks?.onResponse?.(ctx as OnResponseHookContext);
+      if (onResponse) ctx.response = onResponse;
     } catch (error) {
       logger.error({
         requestId: resolveRequestId(ctx.request)!,

package/src/types.ts

@@ -1,13 +1,18 @@
 import type { ProviderV3 } from "@ai-sdk/provider";
 
-import type { ChatCompletionsBody } from "./endpoints/chat-completions/schema";
-import type { EmbeddingsBody } from "./endpoints/embeddings/schema";
+import type {
+  ChatCompletions,
+  ChatCompletionsBody,
+  ChatCompletionsChunk,
+} from "./endpoints/chat-completions/schema";
+import type { Embeddings, EmbeddingsBody } from "./endpoints/embeddings/schema";
+import type { OpenAIError } from "./errors/openai";
 import type { Logger, LoggerConfig } from "./logger";
 import type { ModelCatalog, ModelId } from "./models/types";
 import type { ProviderId, ProviderRegistry } from "./providers/types";
 
 /**
- * Request overrides returned from the `before` hook.
+ * Request overrides returned from the `onRequest` hook.
  */
 export type RequestPatch = {
   /**
@@ -67,7 +72,11 @@ export type GatewayContext = {
   /**
    * Result returned by the handler (pre-response).
    */
-  result?: object | ReadableStream<Uint8Array>;
+  result?:
+    | ChatCompletions
+    | ReadableStream<ChatCompletionsChunk | OpenAIError>
+    | Embeddings
+    | object;
   /**
    * Final response returned by the lifecycle.
    */
@@ -83,7 +92,8 @@ export type HookContext = Omit<Readonly<GatewayContext>, "state"> & {
 
 type RequiredHookContext<K extends keyof GatewayContext> = Omit<HookContext, K> &
   Required<Pick<HookContext, K>>;
-export type BeforeHookContext = RequiredHookContext<"request">;
+export type OnRequestHookContext = RequiredHookContext<"request">;
+export type BeforeHookContext = RequiredHookContext<"request" | "body" | "operation">;
 export type ResolveModelHookContext = RequiredHookContext<"request" | "body" | "modelId">;
 export type ResolveProviderHookContext = RequiredHookContext<
   "request" | "body" | "modelId" | "resolvedModelId" | "operation"
@@ -91,6 +101,7 @@ export type ResolveProviderHookContext = RequiredHookContext<
 export type AfterHookContext = RequiredHookContext<
   "request" | "result" | "provider" | "resolvedModelId" | "operation"
 >;
+export type OnResponseHookContext = RequiredHookContext<"request" | "response">;
 
 /**
  * Hooks to plugin to the gateway lifecycle.
@@ -101,9 +112,20 @@ export type GatewayHooks = {
    * @returns Optional RequestPatch to merge into headers / override body,
    * or Response to short-circuit the request.
    */
+  onRequest?: (
+    ctx: OnRequestHookContext,
+  ) => void | RequestPatch | Response | Promise<void | RequestPatch | Response>;
+  /**
+   * Runs after request JSON is parsed and validated for chat completions / embeddings.
+   * @returns Replacement parsed body, or undefined to keep original.
+   */
   before?: (
     ctx: BeforeHookContext,
-  ) => void | RequestPatch | Response | Promise<void | RequestPatch | Response>;
+  ) =>
+    | void
+    | ChatCompletionsBody
+    | EmbeddingsBody
+    | Promise<void | ChatCompletionsBody | EmbeddingsBody>;
   /**
    * Maps a user-provided model ID or alias to a canonical ID.
    * @returns Canonical model ID or undefined to keep original.
@@ -118,15 +140,23 @@ export type GatewayHooks = {
   ) => ProviderV3 | void | Promise<ProviderV3 | void>;
   /**
    * Runs after the endpoint handler.
-   * @returns Response to replace, or undefined to keep original.
+   * @returns Result to replace, or undefined to keep original.
    */
   after?: (
     ctx: AfterHookContext,
   ) =>
     | void
-    | object
-    | ReadableStream<Uint8Array>
-    | Promise<void | object | ReadableStream<Uint8Array>>;
+    | ChatCompletions
+    | ReadableStream<ChatCompletionsChunk | OpenAIError>
+    | Embeddings
+    | Promise<
+        void | ChatCompletions | ReadableStream<ChatCompletionsChunk | OpenAIError> | Embeddings
+      >;
+  /**
+   * Runs after the lifecycle has produced the final Response.
+   * @returns Replacement Response, or undefined to keep original.
+   */
+  onResponse?: (ctx: OnResponseHookContext) => void | Response | Promise<void | Response>;
 };
 
 /**

package/src/utils/response.ts

@@ -2,6 +2,19 @@ import { REQUEST_ID_HEADER, resolveRequestId } from "./headers";
 
 const TEXT_ENCODER = new TextEncoder();
 
+class JsonToSseTransformStream extends TransformStream<unknown, string> {
+  constructor() {
+    super({
+      transform(part, controller) {
+        controller.enqueue(`data: ${JSON.stringify(part)}\n\n`);
+      },
+      flush(controller) {
+        controller.enqueue("data: [DONE]\n\n");
+      },
+    });
+  }
+}
+
 export const prepareResponseInit = (request: Request): ResponseInit => ({
   headers: { [REQUEST_ID_HEADER]: resolveRequestId(request.headers)! },
 });
@@ -25,13 +38,15 @@ export const mergeResponseInit = (
 };
 
 export const toResponse = (
-  result: ReadableStream<Uint8Array> | Uint8Array<ArrayBuffer> | object | string,
+  result: ReadableStream | Uint8Array<ArrayBuffer> | object | string,
   responseInit?: ResponseInit,
 ): Response => {
   let body: BodyInit;
 
   const isStream = result instanceof ReadableStream;
-  if (isStream || result instanceof Uint8Array) {
+  if (isStream) {
+    body = result.pipeThrough(new JsonToSseTransformStream()).pipeThrough(new TextEncoderStream());
+  } else if (result instanceof Uint8Array) {
     body = result;
   } else if (typeof result === "string") {
     body = TEXT_ENCODER.encode(result);
@@ -57,7 +72,7 @@ export const toResponse = (
       ? {
           "content-type": "text/event-stream",
           "cache-control": "no-cache",
-          Connection: "keep-alive",
+          connection: "keep-alive",
         }
       : {
           "content-type": "application/json",