@hebo-ai/gateway 0.1.2 → 0.2.0

package/src/endpoints/embeddings/handler.ts

@@ -1,32 +1,31 @@
-import { embedMany } from "ai";
+import { embedMany, wrapEmbeddingModel } from "ai";
 import * as z from "zod/mini";
 
-import type { GatewayConfig, Endpoint } from "../../types";
+import type { GatewayConfig, Endpoint, GatewayContext } from "../../types";
 
-import { parseConfig } from "../../config";
+import { withLifecycle } from "../../lifecycle";
+import { modelMiddlewareMatcher } from "../../middleware/matcher";
 import { resolveProvider } from "../../providers/registry";
 import { createErrorResponse } from "../../utils/errors";
-import { withHooks } from "../../utils/hooks";
-import { transformEmbeddingsInputs, createEmbeddingsResponse } from "./converters";
+import { convertToEmbedCallOptions, createEmbeddingsResponse } from "./converters";
 import { EmbeddingsBodySchema } from "./schema";
 
 export const embeddings = (config: GatewayConfig): Endpoint => {
-  const { providers, models, hooks } = parseConfig(config);
+  const hooks = config.hooks;
 
-  const handler = async (req: Request): Promise<Response> => {
-    if (req.method !== "POST") {
+  const handler = async (ctx: GatewayContext): Promise<Response> => {
+    if (!ctx.request || ctx.request.method !== "POST") {
       return createErrorResponse("METHOD_NOT_ALLOWED", "Method Not Allowed", 405);
     }
 
     let body;
     try {
-      body = await req.json();
+      body = await ctx.request.json();
     } catch {
       return createErrorResponse("BAD_REQUEST", "Invalid JSON", 400);
     }
 
     const parsed = EmbeddingsBodySchema.safeParse(body);
-
     if (!parsed.success) {
       return createErrorResponse(
         "UNPROCESSABLE_ENTITY",
@@ -35,51 +34,58 @@ export const embeddings = (config: GatewayConfig): Endpoint => {
         z.prettifyError(parsed.error),
       );
     }
+    ctx.body = parsed.data;
 
-    const { model: modelId, ...inputs } = parsed.data;
+    let inputs;
+    ({ model: ctx.modelId, ...inputs } = parsed.data);
 
-    let resolvedModelId;
     try {
-      resolvedModelId = (await hooks?.resolveModelId?.({ modelId })) ?? modelId;
+      ctx.resolvedModelId = (await hooks?.resolveModelId?.(ctx)) ?? ctx.modelId;
     } catch (error) {
       return createErrorResponse("BAD_REQUEST", error, 400);
     }
 
-    let embedOptions;
+    ctx.operation = "embeddings";
     try {
-      embedOptions = transformEmbeddingsInputs(inputs);
+      const override = await hooks?.resolveProvider?.(ctx);
+      ctx.provider =
+        override ??
+        resolveProvider({
+          providers: ctx.providers,
+          models: ctx.models,
+          modelId: ctx.resolvedModelId,
+          operation: ctx.operation,
+        });
     } catch (error) {
       return createErrorResponse("BAD_REQUEST", error, 400);
     }
 
-    let provider;
+    const embeddingModel = ctx.provider.embeddingModel(ctx.resolvedModelId);
+
+    let embedOptions;
     try {
-      const args = {
-        providers,
-        models,
-        modelId: resolvedModelId,
-        operation: "embeddings" as const,
-      };
-      const override = await hooks?.resolveProvider?.(args);
-      provider = override ?? resolveProvider(args);
+      embedOptions = convertToEmbedCallOptions(inputs);
     } catch (error) {
       return createErrorResponse("BAD_REQUEST", error, 400);
     }
 
-    const embeddingModel = provider.embeddingModel(resolvedModelId);
+    const embeddingModelWithMiddleware = wrapEmbeddingModel({
+      model: embeddingModel,
+      middleware: modelMiddlewareMatcher.forEmbedding(ctx.resolvedModelId, embeddingModel.provider),
+    });
 
     let result;
     try {
       result = await embedMany({
-        model: embeddingModel,
+        model: embeddingModelWithMiddleware,
         ...embedOptions,
       });
     } catch (error) {
       return createErrorResponse("INTERNAL_SERVER_ERROR", error, 500);
     }
 
-    return createEmbeddingsResponse(result, modelId);
+    return createEmbeddingsResponse(result, ctx.modelId);
   };
 
-  return { handler: withHooks(hooks, handler) };
+  return { handler: withLifecycle(handler, config) };
 };

package/src/endpoints/embeddings/schema.ts

@@ -1,28 +1,27 @@
-import * as z from "zod/mini";
+import * as z from "zod";
 
 export const EmbeddingsInputsSchema = z.object({
   input: z.union([z.string(), z.array(z.string())]),
-  encoding_format: z.optional(z.enum(["float", "base64"])),
-  dimensions: z.optional(z.number()),
-  user: z.optional(z.string()),
+  dimensions: z.int().nonnegative().max(65536).optional(),
 });
 export type EmbeddingsInputs = z.infer<typeof EmbeddingsInputsSchema>;
 
-export const EmbeddingsBodySchema = z.extend(EmbeddingsInputsSchema, {
+export const EmbeddingsBodySchema = z.looseObject({
   model: z.string(),
+  ...EmbeddingsInputsSchema.shape,
 });
 export type EmbeddingsBody = z.infer<typeof EmbeddingsBodySchema>;
 
 export const EmbeddingsDataSchema = z.object({
   object: z.literal("embedding"),
   embedding: z.array(z.number()),
-  index: z.number(),
+  index: z.int().nonnegative(),
 });
 export type EmbeddingsData = z.infer<typeof EmbeddingsDataSchema>;
 
 export const EmbeddingsUsageSchema = z.object({
-  prompt_tokens: z.number(),
-  total_tokens: z.number(),
+  prompt_tokens: z.int().nonnegative().optional(),
+  total_tokens: z.int().nonnegative().optional(),
 });
 export type EmbeddingsUsage = z.infer<typeof EmbeddingsUsageSchema>;
 
@@ -30,7 +29,8 @@ export const EmbeddingsSchema = z.object({
   object: z.literal("list"),
   data: z.array(EmbeddingsDataSchema),
   model: z.string(),
-  usage: EmbeddingsUsageSchema,
-  providerMetadata: z.optional(z.any()),
+  usage: EmbeddingsUsageSchema.nullable(),
+  // Extensions
+  provider_metadata: z.optional(z.any()).meta({ extension: true }),
 });
 export type Embeddings = z.infer<typeof EmbeddingsSchema>;

package/src/endpoints/models/converters.ts

@@ -1,6 +1,8 @@
 import type { ModelCatalog, CatalogModel } from "../../models/types";
 import type { ModelList, Model } from "./schema";
 
+import { mergeResponseInit } from "../../utils/response";
+
 export function toModel(id: string, catalogModel: CatalogModel): Model {
   const { created, providers, modalities, additionalProperties, ...rest } = catalogModel;
   let createdTimestamp = Math.floor(Date.now() / 1000);
@@ -11,18 +13,18 @@ export function toModel(id: string, catalogModel: CatalogModel): Model {
     }
   }
 
-  const model: Model = {
+  const model = {
     id,
-    object: "model",
+    object: "model" as const,
     created: createdTimestamp,
     owned_by: providers?.[0] || "system",
     architecture: {
-      input_modalities: modalities.input || [],
+      input_modalities: modalities?.input || [],
       modality:
-        modalities.input &&
-        modalities.output &&
+        modalities?.input &&
+        modalities?.output &&
         `${modalities.input?.[0]}->${modalities.output?.[0]}`,
-      output_modalities: modalities.output || [],
+      output_modalities: modalities?.output || [],
     },
     endpoints:
       providers?.map((provider) => ({
@@ -41,14 +43,20 @@ export function toModels(models: ModelCatalog): ModelList {
     data: Object.entries(models).map(([id, catalogModel]) => toModel(id, catalogModel!)),
   };
 }
-export function createModelsResponse(models: ModelCatalog): Response {
-  return new Response(JSON.stringify(toModels(models)), {
-    headers: { "Content-Type": "application/json" },
-  });
+export function createModelsResponse(models: ModelCatalog, responseInit?: ResponseInit): Response {
+  return new Response(
+    JSON.stringify(toModels(models)),
+    mergeResponseInit({ "Content-Type": "application/json" }, responseInit),
+  );
 }
 
-export function createModelResponse(id: string, catalogModel: CatalogModel): Response {
-  return new Response(JSON.stringify(toModel(id, catalogModel)), {
-    headers: { "Content-Type": "application/json" },
-  });
+export function createModelResponse(
+  id: string,
+  catalogModel: CatalogModel,
+  responseInit?: ResponseInit,
+): Response {
+  return new Response(
+    JSON.stringify(toModel(id, catalogModel)),
+    mergeResponseInit({ "Content-Type": "application/json" }, responseInit),
+  );
 }

package/src/endpoints/models/handler.test.ts

@@ -1,46 +1,43 @@
-import { createProviderRegistry } from "ai";
 import { MockProviderV3 } from "ai/test";
 import { describe, expect, test } from "bun:test";
 
 import { parseResponse } from "../../../test/helpers/http";
-import { createModelCatalog } from "../../models/catalog";
 import { models } from "./handler";
 
 const baseUrl = "http://localhost/models";
 
 describe("Models Handler", () => {
-  const registry = createProviderRegistry({
-    anthropic: new MockProviderV3(),
-    google: new MockProviderV3(),
-  });
-
-  const catalog = createModelCatalog({
-    "anthropic/claude-opus-4.5": {
-      name: "Claude Opus 4.5",
-      created: "2025-09-29T10:00:00.000Z",
-      knowledge: "2025-07",
-      modalities: {
-        input: ["text", "image"],
-        output: ["text"],
-      },
-      context: 200000,
-      capabilities: ["reasoning", "tool_call"],
-      providers: ["anthropic"],
+  const endpoint = models({
+    providers: {
+      anthropic: new MockProviderV3(),
+      google: new MockProviderV3(),
     },
-    "google/gemini-3-flash": {
-      name: "Gemini 3 Flash",
-      created: "2025-10-01T08:30:00.000Z",
-      modalities: {
-        input: ["text", "video"],
-        output: ["text"],
+    models: {
+      "anthropic/claude-opus-4.5": {
+        name: "Claude Opus 4.5",
+        created: "2025-09-29T10:00:00.000Z",
+        knowledge: "2025-07",
+        modalities: {
+          input: ["text", "image"],
+          output: ["text"],
+        },
+        context: 200000,
+        capabilities: ["reasoning", "tool_call"],
+        providers: ["anthropic"],
+      },
+      "google/gemini-3-flash": {
+        name: "Gemini 3 Flash",
+        created: "2025-10-01T08:30:00.000Z",
+        modalities: {
+          input: ["text", "video"],
+          output: ["text"],
+        },
+        context: 128000,
+        providers: ["google"],
       },
-      context: 128000,
-      providers: ["google"],
     },
   });
 
-  const endpoint = models({ providers: registry, models: catalog });
-
   test("should list models via GET request with realistic data (exact match)", async () => {
     const request = new Request(baseUrl, { method: "GET" });
 

package/src/endpoints/models/handler.ts

@@ -1,23 +1,21 @@
-import type { GatewayConfig, Endpoint } from "../../types";
+import type { GatewayConfig, Endpoint, GatewayContext } from "../../types";
 
-import { parseConfig } from "../../config";
+import { withLifecycle } from "../../lifecycle";
 import { createErrorResponse } from "../../utils/errors";
-import { withHooks } from "../../utils/hooks";
 import { createModelsResponse, createModelResponse } from "./converters";
 
 export const models = (config: GatewayConfig): Endpoint => {
-  const { models, hooks } = parseConfig(config);
-
   // eslint-disable-next-line require-await
-  const handler = async (req: Request): Promise<Response> => {
-    if (req.method !== "GET") {
+  const handler = async (ctx: GatewayContext): Promise<Response> => {
+    const request = ctx.request;
+
+    if (!request || request.method !== "GET") {
       return createErrorResponse("METHOD_NOT_ALLOWED", "Method Not Allowed", 405);
     }
 
-    const rawId = req.url.split("/models/", 2)[1]?.split("?", 1)[0];
-
+    const rawId = request.url.split("/models/", 2)[1]?.split("?", 1)[0];
     if (!rawId) {
-      return createModelsResponse(models);
+      return createModelsResponse(ctx.models);
     }
 
     let modelId = rawId;
@@ -27,7 +25,7 @@ export const models = (config: GatewayConfig): Endpoint => {
       return createErrorResponse("BAD_REQUEST", "Invalid model ID", 400);
     }
 
-    const model = models[modelId];
+    const model = ctx.models[modelId];
     if (!model) {
       return createErrorResponse("NOT_FOUND", `Model '${modelId}' not found`, 404);
     }
@@ -35,5 +33,5 @@ export const models = (config: GatewayConfig): Endpoint => {
     return createModelResponse(modelId, model);
   };
 
-  return { handler: withHooks(hooks, handler) };
+  return { handler: withLifecycle(handler, config) };
 };

package/src/endpoints/models/schema.ts

@@ -1,27 +1,33 @@
-import * as z from "zod/mini";
+import * as z from "zod";
 
-export const ModelSchema = z.catchall(
-  z.object({
-    id: z.string(),
-    object: z.literal("model"),
-    created: z.number(),
-    owned_by: z.string(),
-    description: z.optional(z.string()),
-    architecture: z.optional(
-      z.object({
-        modality: z.optional(z.string()),
-        input_modalities: z.readonly(z.array(z.string())),
-        output_modalities: z.readonly(z.array(z.string())),
-      }),
-    ),
-    endpoints: z.array(
+export const ModelSchema = z.looseObject({
+  // Core
+  id: z.string(),
+  object: z.literal("model"),
+  created: z.int().nonnegative(),
+  owned_by: z.string(),
+  // Extensions
+  name: z.string().optional().meta({ extension: true }),
+  knowledge: z.string().optional().meta({ extension: true }),
+  context: z.int().nonnegative().optional().meta({ extension: true }),
+  architecture: z
+    .object({
+      modality: z.string().optional(),
+      input_modalities: z.array(z.string()).readonly().optional(),
+      output_modalities: z.array(z.string()).readonly().optional(),
+    })
+    .optional()
+    .meta({ extension: true }),
+  endpoints: z
+    .array(
       z.object({
         tag: z.string(),
       }),
-    ),
-  }),
-  z.unknown(),
-);
+    )
+    .optional()
+    .meta({ extension: true }),
+  capabilities: z.array(z.string()).readonly().optional().meta({ extension: true }),
+});
 export type Model = z.infer<typeof ModelSchema>;
 
 export const ModelListSchema = z.object({

package/src/gateway.ts

@@ -1,35 +1,21 @@
-import type {
-  Endpoint,
-  GatewayConfig,
-  GatewayConfigBase,
-  GatewayConfigRegistry,
-  HeboGateway,
-} from "./types";
-
-import { parseConfig } from "./config";
+import type { Endpoint, GatewayConfig, HeboGateway } from "./types";
+
 import { chatCompletions } from "./endpoints/chat-completions/handler";
 import { embeddings } from "./endpoints/embeddings/handler";
 import { models } from "./endpoints/models/handler";
 
-const buildRoutes = (config: GatewayConfig) =>
-  ({
+export function gateway(config: GatewayConfig) {
+  const basePath = (config.basePath ?? "").replace(/\/+$/, "");
+
+  const routes = {
     ["/chat/completions"]: chatCompletions(config),
     ["/embeddings"]: embeddings(config),
     ["/models"]: models(config),
-  }) as const satisfies Record<string, Endpoint>;
+  } as const satisfies Record<string, Endpoint>;
 
-type GatewayRoutes = ReturnType<typeof buildRoutes>;
-
-export function gateway(config: GatewayConfigBase): HeboGateway<GatewayRoutes>;
-export function gateway(config: GatewayConfigRegistry): HeboGateway<GatewayRoutes>;
-export function gateway(config: GatewayConfig): HeboGateway<GatewayRoutes> {
-  const parsedConfig = parseConfig(config);
-
-  const basePath = (config.basePath ?? "").replace(/\/+$/, "");
-  const routes = buildRoutes(parsedConfig);
   const routeEntries = Object.entries(routes);
 
-  const handler = (req: Request): Promise<Response> => {
+  const handler = (req: Request, state?: Record<string, unknown>): Promise<Response> => {
     let pathname = new URL(req.url).pathname;
     if (basePath && pathname.startsWith(basePath)) {
       pathname = pathname.slice(basePath.length);
@@ -37,12 +23,12 @@ export function gateway(config: GatewayConfig): HeboGateway<GatewayRoutes> {
 
     for (const [route, endpoint] of routeEntries) {
       if (pathname === route || pathname.startsWith(route + "/")) {
-        return endpoint.handler(req);
+        return endpoint.handler(req, state);
       }
     }
 
     return Promise.resolve(new Response("Not Found", { status: 404 }));
   };
 
-  return { handler, routes };
+  return { handler, routes } satisfies HeboGateway<typeof routes>;
 }

package/src/index.ts

@@ -1,6 +1,9 @@
 export * from "./gateway";
 export type * from "./types";
 
+export * from "./middleware/common";
+export * from "./middleware/matcher";
+
 export * from "./endpoints/chat-completions";
 export * from "./endpoints/embeddings";
 export * from "./endpoints/models";

package/src/{utils/hooks.ts → lifecycle.ts}

@@ -1,6 +1,7 @@
-import type { GatewayHooks, RequestPatch } from "../types";
+import type { GatewayConfig, GatewayContext, RequestPatch } from "./types";
 
-import { createErrorResponse } from "./errors";
+import { parseConfig } from "./config";
+import { createErrorResponse } from "./utils/errors";
 
 const maybeApplyRequestPatch = (request: Request, patch: RequestPatch) => {
   if (!patch.headers && patch.body === undefined) return request;
@@ -21,30 +22,39 @@ const maybeApplyRequestPatch = (request: Request, patch: RequestPatch) => {
   return new Request(request, init);
 };
 
-export const withHooks = (
-  hooks: GatewayHooks | undefined,
-  run: (request: Request) => Promise<Response>,
+export const withLifecycle = (
+  run: (ctx: GatewayContext) => Promise<Response>,
+  config: GatewayConfig,
 ) => {
-  const handler = async (request: Request): Promise<Response> => {
+  const parsedConfig = parseConfig(config);
+
+  const handler = async (request: Request, state?: Record<string, unknown>): Promise<Response> => {
+    const context: GatewayContext = {
+      request,
+      state: state ?? {},
+      providers: parsedConfig.providers,
+      models: parsedConfig.models,
+    };
+
     let beforeResult;
     try {
-      beforeResult = await hooks?.before?.({ request });
+      beforeResult = await parsedConfig.hooks?.before?.(context);
     } catch (error) {
       return createErrorResponse("INTERNAL_SERVER_ERROR", error, 500);
     }
     if (beforeResult instanceof Response) return beforeResult;
 
-    const nextRequest = beforeResult ? maybeApplyRequestPatch(request, beforeResult) : request;
+    context.request = beforeResult ? maybeApplyRequestPatch(request, beforeResult) : request;
 
-    const response = await run(nextRequest);
+    context.response = await run(context);
 
     let after;
     try {
-      after = await hooks?.after?.({ response });
+      after = await parsedConfig.hooks?.after?.(context);
     } catch (error) {
       return createErrorResponse("INTERNAL_SERVER_ERROR", error, 500);
     }
-    return after ?? response;
+    return after ?? context.response;
   };
 
   return handler;

package/src/middleware/common.ts

@@ -0,0 +1,68 @@
+import type { JSONObject } from "@ai-sdk/provider";
+import type { EmbeddingModelMiddleware, LanguageModelMiddleware } from "ai";
+
+import type { ProviderId } from "../providers/types";
+
+function snakeToCamel(key: string): string {
+  return key.replaceAll(/_([a-z])/g, (_, c: string) => c.toUpperCase());
+}
+
+function camelizeKeysDeep(value: unknown): unknown {
+  if (value === null || typeof value !== "object") return value;
+
+  if (Array.isArray(value)) {
+    return value.map((v) => camelizeKeysDeep(v));
+  }
+
+  const out: Record<string, unknown> = {};
+  for (const [k, v] of Object.entries(value)) {
+    if (v === undefined || v === null) continue;
+    out[snakeToCamel(k)] = camelizeKeysDeep(v);
+  }
+  return out;
+}
+
+/**
+ * Converts snake_case params in providerOptions to camelCase
+ * and moves all of them into providerOptions[providerName].
+ */
+type Kind = "embedding" | "language";
+type MiddlewareFor<K extends Kind> = K extends "embedding"
+  ? EmbeddingModelMiddleware
+  : LanguageModelMiddleware;
+type TransformOptsFor<K extends Kind> = Parameters<
+  NonNullable<MiddlewareFor<K>["transformParams"]>
+>[0];
+function forwardParamsForMiddleware<K extends Kind>(
+  _kind: K,
+  providerName: ProviderId,
+): MiddlewareFor<K> {
+  return {
+    specificationVersion: "v3" as const,
+    // eslint-disable-next-line require-await
+    transformParams: async (options: TransformOptsFor<K>) => {
+      const { params } = options;
+      const providerOptions = params.providerOptions;
+      if (!providerOptions) return params;
+
+      const target = (providerOptions[providerName] ??= {});
+      for (const key in providerOptions) {
+        if (key === providerName) continue;
+        Object.assign(target, camelizeKeysDeep(providerOptions[key]) as Record<string, JSONObject>);
+        if (key === "unknown") delete providerOptions[key];
+      }
+
+      return params;
+    },
+  } as MiddlewareFor<K>;
+}
+
+export function forwardParamsMiddleware(providerName: ProviderId): LanguageModelMiddleware {
+  return forwardParamsForMiddleware("language", providerName);
+}
+
+export function forwardParamsEmbeddingMiddleware(
+  providerName: ProviderId,
+): EmbeddingModelMiddleware {
+  return forwardParamsForMiddleware("embedding", providerName);
+}