@hebo-ai/gateway 0.4.0-beta.3 → 0.4.0

package/src/errors/openai.ts

@@ -22,28 +22,36 @@ export class OpenAIError {
   }
 }
 
+const mapType = (status: number) => (status < 500 ? "invalid_request_error" : "server_error");
+
+const maybeMaskMessage = (meta: ReturnType<typeof getErrorMeta>, requestId?: string) => {
+  // FUTURE: consider masking all upstream errors, also 4xx
+  if (!(isProduction() && meta.status >= 500)) {
+    return meta.message;
+  }
+  // FUTURE: always attach requestId to errors (masked and unmasked)
+  return `${STATUS_CODE(meta.status)} (${requestId ?? "see requestId in response headers"})`;
+};
+
 export function toOpenAIError(error: unknown): OpenAIError {
   const meta = getErrorMeta(error);
-  return new OpenAIError(meta.message, meta.type, meta.code);
+
+  return new OpenAIError(maybeMaskMessage(meta), mapType(meta.status), meta.code);
 }
 
 export function toOpenAIErrorResponse(error: unknown, responseInit?: ResponseInit) {
   const meta = getErrorMeta(error);
 
-  const shouldMask = isProduction() && (meta.status >= 500 || meta.code.includes("UPSTREAM"));
-
-  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;
-  }
-
-  return toResponse(new OpenAIError(message, meta.type, meta.code), {
-    ...responseInit,
-    status: meta.status,
-    statusText: meta.code,
-  });
+  return toResponse(
+    new OpenAIError(
+      maybeMaskMessage(meta, resolveRequestId(responseInit)),
+      mapType(meta.status),
+      meta.code,
+    ),
+    {
+      ...responseInit,
+      status: meta.status,
+      statusText: meta.code,
+    },
+  );
 }

package/src/errors/utils.ts

@@ -11,6 +11,7 @@ export const STATUS_CODES = {
   409: "CONFLICT",
   422: "UNPROCESSABLE_ENTITY",
   429: "TOO_MANY_REQUESTS",
+  499: "CLIENT_CLOSED_REQUEST",
   500: "INTERNAL_SERVER_ERROR",
   502: "BAD_GATEWAY",
   503: "SERVICE_UNAVAILABLE",
@@ -23,26 +24,24 @@ export const STATUS_CODE = (status: number) => {
   return status >= 400 && status < 500 ? STATUS_CODES[400] : STATUS_CODES[500];
 };
 
+// FUTURE: always return a wrapped GatewayError?
 export function getErrorMeta(error: unknown) {
   const message = error instanceof Error ? error.message : String(error);
 
-  let code: string;
   let status: number;
-  let param = "";
+  let code: string;
 
   if (error instanceof GatewayError) {
-    ({ code, status } = error);
+    ({ status, code } = error);
   } else {
     const normalized = normalizeAiSdkError(error);
     if (normalized) {
-      ({ code, status } = normalized);
+      ({ status, code } = normalized);
     } else {
       status = 500;
       code = STATUS_CODE(status);
     }
   }
 
-  const type = status < 500 ? "invalid_request_error" : "server_error";
-
-  return { code, status, param, type, message };
+  return { status, code, message };
 }

package/src/gateway.ts

@@ -30,7 +30,7 @@ export function gateway(config: GatewayConfig) {
       pathname = pathname.slice(basePath.length);
     }
 
-    logger.debug(`[gateway] ${req.method} ${pathname}`);
+    logger.info(`[gateway] ${req.method} ${pathname}`);
     for (const [route, endpoint] of routeEntries) {
       if (pathname === route || pathname.startsWith(route + "/")) {
         return endpoint.handler(req, state);

package/src/lifecycle.ts

@@ -6,10 +6,15 @@ import type {
 } from "./types";
 
 import { parseConfig } from "./config";
+import { GatewayError } from "./errors/gateway";
 import { toOpenAIErrorResponse } from "./errors/openai";
 import { logger } from "./logger";
-import { withOtel } from "./telemetry/otel";
-import { addSpanEvent, recordSpanError } from "./telemetry/span";
+import { getBaggageAttributes } from "./telemetry/baggage";
+import { initFetch } from "./telemetry/fetch";
+import { getRequestAttributes, getResponseAttributes } from "./telemetry/http";
+import { recordV8jsMemory } from "./telemetry/memory";
+import { addSpanEvent, setSpanEventsEnabled, setSpanTracer, startSpan } from "./telemetry/span";
+import { wrapStream } from "./telemetry/stream";
 import { resolveRequestId } from "./utils/headers";
 import { maybeApplyRequestPatch, prepareRequestHeaders } from "./utils/request";
 import { prepareResponseInit, toResponse } from "./utils/response";
@@ -20,23 +25,81 @@ export const winterCgHandler = (
 ) => {
   const parsedConfig = parseConfig(config);
 
-  const core = async (ctx: GatewayContext): Promise<void> => {
+  if (parsedConfig.telemetry?.enabled) {
+    setSpanTracer(parsedConfig.telemetry?.tracer);
+    setSpanEventsEnabled(parsedConfig.telemetry?.signals?.hebo);
+    initFetch(parsedConfig.telemetry?.signals?.hebo);
+  }
+
+  return async (request: Request, state?: Record<string, unknown>): Promise<Response> => {
+    const ctx: GatewayContext = {
+      request,
+      state: state ?? {},
+      providers: parsedConfig.providers,
+      models: parsedConfig.models,
+    };
+
+    const headers = prepareRequestHeaders(ctx.request);
+    if (headers) ctx.request = new Request(ctx.request, { headers });
+
+    const span = startSpan(ctx.request.url);
+    span.setAttributes(getBaggageAttributes(ctx.request));
+    if (!span.isExisting) {
+      span.setAttributes(getRequestAttributes(ctx.request, parsedConfig.telemetry?.signals?.http));
+    }
+
+    const finalize = (status: number, reason?: unknown) => {
+      if (ctx.operation) {
+        span.updateName(`${ctx.operation}${ctx.modelId ? ` ${ctx.modelId}` : ""}`);
+      }
+
+      if (!span.isExisting) {
+        // FUTURE add http.server.request.duration
+        span.setAttributes(
+          getResponseAttributes(ctx.response!, parsedConfig.telemetry?.signals?.http),
+        );
+      }
+
+      let realStatus = status;
+      if (ctx.request.signal.aborted) realStatus = 499;
+      else if (status === 200 && ctx.response?.status) realStatus = ctx.response.status;
+
+      if (realStatus !== 200) {
+        (realStatus >= 500 ? logger.error : logger.warn)({
+          requestId: resolveRequestId(ctx.request),
+          err: reason ?? ctx.request.signal.reason,
+        });
+
+        if (realStatus >= 500) span.recordError(reason);
+      }
+      span.setAttributes({ "http.response.status_code_effective": realStatus });
+
+      recordV8jsMemory(parsedConfig.telemetry?.signals?.hebo);
+
+      span.finish();
+    };
+
     try {
       if (parsedConfig.hooks?.onRequest) {
         const onRequest = await parsedConfig.hooks.onRequest(ctx as OnRequestHookContext);
         addSpanEvent("hebo.hooks.on_request.completed");
 
-        if (onRequest) {
-          if (onRequest instanceof Response) {
-            ctx.response = onRequest;
-            return;
-          }
+        if (onRequest instanceof Response) {
+          ctx.response = onRequest;
+        } else if (onRequest) {
           ctx.request = maybeApplyRequestPatch(ctx.request, onRequest);
         }
       }
 
-      ctx.result = (await run(ctx)) as typeof ctx.result;
-      ctx.response = toResponse(ctx.result!, prepareResponseInit(ctx.request));
+      if (!ctx.response) {
+        ctx.result = (await span.runWithContext(() => run(ctx))) as typeof ctx.result;
+
+        if (ctx.result instanceof ReadableStream) {
+          ctx.result = wrapStream(ctx.result, { onDone: finalize });
+        }
+
+        ctx.response = toResponse(ctx.result!, prepareResponseInit(ctx.request));
+      }
 
       if (parsedConfig.hooks?.onResponse) {
         const onResponse = await parsedConfig.hooks.onResponse(ctx as OnResponseHookContext);
@@ -45,30 +108,20 @@ export const winterCgHandler = (
           ctx.response = onResponse;
         }
       }
+
+      // FUTURE: this can leak if onResponse removed wrapper from response.body
+      if (!(ctx.result instanceof ReadableStream)) {
+        finalize(ctx.response.status);
+      }
     } catch (error) {
-      recordSpanError(error);
-      logger.error({
-        requestId: resolveRequestId(ctx.request),
-        err: error instanceof Error ? error : new Error(String(error)),
-      });
-      ctx.response = toOpenAIErrorResponse(error, prepareResponseInit(ctx.request));
+      ctx.response = toOpenAIErrorResponse(
+        ctx.request.signal.aborted
+          ? new GatewayError(error ?? ctx.request.signal.reason, 499)
+          : error,
+        prepareResponseInit(ctx.request),
+      );
+      finalize(ctx.response.status, error);
     }
-  };
-
-  const handler = parsedConfig.telemetry?.enabled ? withOtel(core, parsedConfig) : core;
-
-  return async (request: Request, state?: Record<string, unknown>): Promise<Response> => {
-    const ctx: GatewayContext = {
-      request,
-      state: state ?? {},
-      providers: parsedConfig.providers,
-      models: parsedConfig.models,
-    };
-
-    const headers = prepareRequestHeaders(ctx.request);
-    if (headers) ctx.request = new Request(ctx.request, { headers });
-
-    await handler(ctx);
 
     return ctx.response ?? new Response("Internal Server Error", { status: 500 });
   };

package/src/middleware/matcher.ts

@@ -118,7 +118,7 @@ class ModelMiddlewareMatcher {
         if (--n === 0) break;
       }
       logger.warn(`[middleware] cache eviction`);
-      addSpanEvent("hebo.middelware.cache.evicted");
+      addSpanEvent("hebo.middleware.cache.evicted");
     }
 
     this.cache.set(key, out);

package/src/models/amazon/presets.ts

@@ -24,7 +24,7 @@ const NOVA_MULTIMODAL_BASE = {
 const NOVA_EMBEDDINGS_BASE = {
   modalities: {
     input: ["text", "image", "audio", "video", "pdf"] as const,
-    output: ["embeddings"] as const,
+    output: ["embedding"] as const,
   },
   providers: ["bedrock"] as const satisfies readonly CanonicalProviderId[],
 } satisfies DeepPartial<CatalogModel>;

package/src/models/cohere/presets.ts

@@ -26,7 +26,7 @@ const COMMAND_VISION_BASE = {
 const EMBED_V3_BASE = {
   modalities: {
     input: ["text", "image"] as const,
-    output: ["embeddings"] as const,
+    output: ["embedding"] as const,
   },
   providers: ["cohere"] as const satisfies readonly CanonicalProviderId[],
 } satisfies DeepPartial<CatalogModel>;
@@ -34,7 +34,7 @@ const EMBED_V3_BASE = {
 const EMBED_V4_BASE = {
   modalities: {
     input: ["text", "image", "pdf"] as const,
-    output: ["embeddings"] as const,
+    output: ["embedding"] as const,
   },
   providers: ["cohere"] as const satisfies readonly CanonicalProviderId[],
 } satisfies DeepPartial<CatalogModel>;

package/src/models/google/presets.ts

@@ -22,7 +22,7 @@ const GEMINI_BASE = {
 const GEMINI_EMBEDDINGS_BASE = {
   modalities: {
     input: ["text"] as const,
-    output: ["embeddings"] as const,
+    output: ["embedding"] as const,
   },
   providers: ["vertex"] as const satisfies readonly CanonicalProviderId[],
 } satisfies DeepPartial<CatalogModel>;

package/src/models/openai/presets.ts

@@ -52,7 +52,7 @@ const GPT_PRO_BASE = {
 const EMBEDDINGS_BASE = {
   modalities: {
     input: ["text"] as const,
-    output: ["embeddings"] as const,
+    output: ["embedding"] as const,
   },
   providers: ["openai", "azure"] as const satisfies readonly CanonicalProviderId[],
 } satisfies DeepPartial<CatalogModel>;

package/src/models/types.ts

@@ -93,7 +93,7 @@ export type CatalogModel = {
   knowledge?: string;
   modalities?: {
     input: readonly ("text" | "image" | "file" | "audio" | "video" | "pdf")[];
-    output: readonly ("text" | "image" | "audio" | "video" | "embeddings")[];
+    output: readonly ("text" | "image" | "audio" | "video" | "embedding")[];
   };
   context?: number;
   capabilities?: readonly (

package/src/models/voyage/presets.ts

@@ -6,7 +6,7 @@ import { presetFor, type DeepPartial } from "../../utils/preset";
 const VOYAGE_BASE = {
   modalities: {
     input: ["text"] as const,
-    output: ["embeddings"] as const,
+    output: ["embedding"] as const,
   },
   providers: ["voyage"] as const satisfies readonly CanonicalProviderId[],
 } satisfies DeepPartial<CatalogModel>;

package/src/providers/registry.ts

@@ -22,10 +22,10 @@ export const resolveProvider = (args: {
     throw new GatewayError(`Model '${modelId}' not found in catalog`, 422, "MODEL_NOT_FOUND");
   }
 
-  const modality = operation === "embeddings" ? "embeddings" : "text";
+  const modality = operation === "embeddings" ? "embedding" : "text";
   if (catalogModel.modalities && !catalogModel.modalities.output.includes(modality)) {
     throw new GatewayError(
-      `Model '${modelId}' does not support '${operation}' output`,
+      `Model '${modelId}' does not support '${modality}' output`,
       422,
       "MODEL_UNSUPPORTED_OPERATION",
     );

package/src/telemetry/baggage.ts

@@ -0,0 +1,27 @@
+const HEBO_BAGGAGE_PREFIX = "hebo.";
+
+export const getBaggageAttributes = (request?: Request) => {
+  const h = request?.headers.get("baggage");
+  if (!h) return {};
+
+  const attrs: Record<string, string> = {};
+
+  for (const part of h.split(",")) {
+    const [k, v] = part.trim().split("=", 2);
+    if (!k || !v) continue;
+
+    const [rawValue] = v.split(";", 1);
+    if (!rawValue) continue;
+
+    let value = rawValue;
+    try {
+      value = decodeURIComponent(rawValue);
+    } catch {}
+
+    if (k.startsWith(HEBO_BAGGAGE_PREFIX)) {
+      attrs[k.slice(HEBO_BAGGAGE_PREFIX.length)] = value;
+    }
+  }
+
+  return attrs;
+};

package/src/telemetry/fetch.ts

@@ -1,5 +1,7 @@
 import { SpanKind } from "@opentelemetry/api";
 
+import type { TelemetrySignalLevel } from "../types";
+
 import { withSpan } from "./span";
 
 const ORIGINAL_FETCH_KEY = Symbol.for("@hebo/fetch/original-fetch");
@@ -9,15 +11,25 @@ type GlobalFetchState = typeof globalThis & {
 };
 
 const g = globalThis as GlobalFetchState;
+let fetchTracingEnabled = false;
+
+const shouldTraceFetch = (init?: RequestInit): boolean =>
+  typeof (init?.headers as any)?.["user-agent"] === "string" &&
+  (init!.headers as any)["user-agent"].indexOf("ai-sdk/provider-utils") !== -1;
 
-const perfFetch = (input: RequestInfo | URL, init?: RequestInit) => {
+const otelFetch = (input: RequestInfo | URL, init?: RequestInit) => {
   const original = g[ORIGINAL_FETCH_KEY]!;
+
+  if (!fetchTracingEnabled) return original(input, init);
+  if (!shouldTraceFetch(init)) return original(input, init);
   return withSpan("fetch", () => original(input, init), { kind: SpanKind.CLIENT });
 };
 
-export const initFetch = () => {
+export const initFetch = (level?: TelemetrySignalLevel) => {
+  fetchTracingEnabled = level === "full";
+  if (!fetchTracingEnabled) return;
   if (g[ORIGINAL_FETCH_KEY]) return;
 
   g[ORIGINAL_FETCH_KEY] = globalThis.fetch.bind(globalThis);
-  globalThis.fetch = perfFetch as typeof fetch;
+  globalThis.fetch = otelFetch as typeof fetch;
 };

package/src/telemetry/gen-ai.ts

@@ -0,0 +1,88 @@
+import { metrics, type Attributes } from "@opentelemetry/api";
+
+import type { TelemetrySignalLevel } from "../types";
+
+const meter = metrics.getMeter("@hebo/gateway");
+
+const requestDurationHistogram = meter.createHistogram("gen_ai.server.request.duration", {
+  description: "End-to-end gateway request duration",
+  unit: "s",
+  advice: {
+    explicitBucketBoundaries: [
+      0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 120, 240,
+    ],
+  },
+});
+
+const timePerOutputTokenHistogram = meter.createHistogram("gen_ai.server.time_per_output_token", {
+  description: "End-to-end gateway request duration per output token",
+  unit: "s",
+  advice: {
+    explicitBucketBoundaries: [
+      0.01, 0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.75, 1.0, 2.5,
+    ],
+  },
+});
+
+const tokenUsageHistogram = meter.createHistogram("gen_ai.client.token.usage", {
+  description: "Token usage reported by upstream model responses",
+  unit: "{token}",
+  advice: {
+    explicitBucketBoundaries: [
+      1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144,
+      524288, 1048576,
+    ],
+  },
+});
+
+// FUTURE: record unsuccessful calls
+export const recordRequestDuration = (
+  start: number,
+  attrs: Attributes,
+  signalLevel?: TelemetrySignalLevel,
+) => {
+  if (!signalLevel || signalLevel === "off") return;
+
+  requestDurationHistogram.record((performance.now() - start) / 1000, attrs);
+};
+
+// FUTURE: record unsuccessful calls
+export const recordTimePerOutputToken = (
+  start: number,
+  tokenAttrs: Attributes,
+  metricAttrs: Attributes,
+  signalLevel?: TelemetrySignalLevel,
+) => {
+  if (!signalLevel || (signalLevel !== "recommended" && signalLevel !== "full")) return;
+
+  const outputTokens = tokenAttrs["gen_ai.usage.output_tokens"];
+  if (typeof outputTokens !== "number" || outputTokens <= 0) return;
+
+  timePerOutputTokenHistogram.record(
+    (performance.now() - start) / 1000 / outputTokens,
+    metricAttrs,
+  );
+};
+
+// FUTURE: record unsuccessful calls
+export const recordTokenUsage = (
+  tokenAttrs: Attributes,
+  metricAttrs: Attributes,
+  signalLevel?: TelemetrySignalLevel,
+) => {
+  if (!signalLevel || (signalLevel !== "recommended" && signalLevel !== "full")) return;
+
+  const record = (value: unknown, tokenType: string) => {
+    if (typeof value !== "number") return;
+    tokenUsageHistogram.record(
+      value,
+      Object.assign({}, metricAttrs, { "gen_ai.token.type": tokenType }),
+    );
+  };
+
+  record(tokenAttrs["gen_ai.usage.input_tokens"], "input");
+  record(tokenAttrs["gen_ai.usage.output_tokens"], "output");
+  record(tokenAttrs["gen_ai.usage.total_tokens"], "total");
+  record(tokenAttrs["gen_ai.usage.cached_tokens"], "cached");
+  record(tokenAttrs["gen_ai.usage.reasoning_tokens"], "reasoning");
+};

package/src/telemetry/http.ts

@@ -0,0 +1,65 @@
+import { type TelemetrySignalLevel } from "../types";
+import { resolveRequestId } from "../utils/headers";
+
+const headerArr = (h: Headers, k: string) => (h.has(k) ? [h.get(k)!] : undefined);
+
+export const getRequestAttributes = (request: Request, signalLevel?: TelemetrySignalLevel) => {
+  if (!signalLevel || signalLevel === "off") return {};
+
+  let url;
+  try {
+    // FUTURE: reuse URL from lifecycle
+    url = new URL(request.url);
+  } catch {}
+
+  const attrs = {
+    "http.request.method": request.method,
+    "url.full": request.url,
+    "url.path": url?.pathname,
+    "url.scheme": url?.protocol.replace(":", ""),
+    "server.address": url?.hostname,
+    "server.port": url
+      ? url.port
+        ? Number(url.port)
+        : url.protocol === "https:"
+          ? 443
+          : 80
+      : undefined,
+  };
+
+  if (signalLevel !== "required") {
+    Object.assign(attrs, {
+      // FUTURE: does ElysiaJS and other frameworks attach request id?
+      "http.request.id": resolveRequestId(request),
+      "user_agent.original": request.headers.get("user-agent") ?? undefined,
+    });
+  }
+
+  if (signalLevel === "full") {
+    Object.assign(attrs, {
+      // FUTURE: "url.query"
+      "http.request.header.content-type": headerArr(request.headers, "content-type"),
+      "http.request.header.content-length": headerArr(request.headers, "content-length"),
+      // FUTURE: "client.address"
+    });
+  }
+
+  return attrs;
+};
+
+export const getResponseAttributes = (response: Response, signalLevel?: TelemetrySignalLevel) => {
+  if (!signalLevel || signalLevel === "off") return {};
+
+  const attrs = {
+    "http.response.status_code": response.status,
+  };
+
+  if (signalLevel === "full") {
+    Object.assign(attrs, {
+      "http.response.header.content-type": [headerArr(response.headers, "content-type")],
+      "http.response.header.content-length": [headerArr(response.headers, "content-length")],
+    });
+  }
+
+  return attrs;
+};

package/src/telemetry/memory.ts

@@ -0,0 +1,36 @@
+import { metrics } from "@opentelemetry/api";
+
+import type { TelemetrySignalLevel } from "../types";
+
+const meter = metrics.getMeter("@hebo/gateway");
+const defaultHeapSpaceAttrs = { "v8js.heap.space.name": "total" } as const;
+
+const heapUsedCounter = meter.createUpDownCounter("v8js.memory.heap.used", {
+  description: "Used bytes in the V8 heap",
+  unit: "By",
+});
+
+const heapSpacePhysicalSizeCounter = meter.createUpDownCounter(
+  "v8js.memory.heap.space.physical_size",
+  {
+    description: "Physical bytes allocated for the V8 heap space",
+    unit: "By",
+  },
+);
+
+const isEnabled = (level?: TelemetrySignalLevel) => level === "recommended" || level === "full";
+
+export const recordV8jsMemory = (level?: TelemetrySignalLevel) => {
+  if (!isEnabled(level)) return;
+
+  let usage;
+  try {
+    usage = globalThis.process?.memoryUsage?.();
+  } catch {
+    return;
+  }
+  if (!usage) return;
+
+  heapUsedCounter.add(usage.heapUsed, defaultHeapSpaceAttrs);
+  heapSpacePhysicalSizeCounter.add(usage.rss, defaultHeapSpaceAttrs);
+};