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

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;
     },
@@ -344,11 +356,22 @@ const gw = gateway({
       // - 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.js

@@ -302,6 +302,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.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.request = maybeApplyRequestPatch(ctx.request, onRequest);
             }
             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.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

@@ -5,7 +5,7 @@ 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 = {
     /**
@@ -77,10 +77,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 +92,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 +110,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 | object | ReadableStream<Uint8Array> | Promise<void | object | ReadableStream<Uint8Array>>;
+    /**
+     * 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>;
 };
 /**
  * Main configuration object for the gateway.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@hebo-ai/gateway",
-  "version": "0.3.0",
+  "version": "0.4.0-alpha.0",
   "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

@@ -476,6 +476,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";
@@ -16,25 +21,20 @@ export const winterCgHandler = (
 
   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);
-
-      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

@@ -7,7 +7,7 @@ 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 = {
   /**
@@ -83,7 +83,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 +92,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 +103,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,7 +131,7 @@ 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,
@@ -127,6 +140,11 @@ export type GatewayHooks = {
     | object
     | ReadableStream<Uint8Array>
     | Promise<void | object | ReadableStream<Uint8Array>>;
+  /**
+   * 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>;
 };
 
 /**