@hebo-ai/gateway 0.4.0-beta.2 → 0.4.0-beta.4

package/src/lifecycle.ts

@@ -8,8 +8,11 @@ import type {
 import { parseConfig } from "./config";
 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 { 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 +23,77 @@ 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),
+        );
+      }
+
+      const realStatus = status === 200 ? (ctx.response?.status ?? status) : status;
+      if (realStatus !== 200) {
+        // FUTURE: in-stream errors are redacted in prod
+        (realStatus >= 500 ? logger.error : logger.warn)({
+          requestId: resolveRequestId(ctx.request),
+          err: reason,
+        });
+
+        if (realStatus >= 500) span.recordError(reason);
+      }
+      span.setAttributes({ "http.response.status_code_effective": realStatus });
+
+      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.request.signal);
+        }
+
+        ctx.response = toResponse(ctx.result!, prepareResponseInit(ctx.request));
+      }
 
       if (parsedConfig.hooks?.onResponse) {
         const onResponse = await parsedConfig.hooks.onResponse(ctx as OnResponseHookContext);
@@ -45,30 +102,15 @@ 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));
+      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,60 @@
+import { metrics, type Attributes } from "@opentelemetry/api";
+
+import type { TelemetrySignalLevel } from "../types";
+
+const meter = metrics.getMeter("@hebo-ai/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 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 = (
+  duration: number,
+  attrs: Attributes,
+  signalLevel?: TelemetrySignalLevel,
+) => {
+  if (!signalLevel || signalLevel === "off") return;
+
+  requestDurationHistogram.record(duration / 1000, attrs);
+};
+
+// 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/span.ts

@@ -1,28 +1,13 @@
-import type { Attributes, Span, SpanOptions, Tracer } from "@opentelemetry/api";
+import type { Attributes, SpanOptions, Tracer } from "@opentelemetry/api";
 
 import { INVALID_SPAN_CONTEXT, SpanKind, SpanStatusCode, context, trace } from "@opentelemetry/api";
 
-const DEFAULT_TRACER_NAME = "@hebo-ai/gateway";
-const mem = () => process?.memoryUsage?.();
-
-const toError = (error: unknown) => (error instanceof Error ? error : new Error(String(error)));
+import type { TelemetrySignalLevel } from "../types";
 
-const maybeSetDynamicAttributes = (span: Span, getAttributes: () => Attributes) => {
-  const attrs = getAttributes();
-  if (Object.keys(attrs).length === 0) return;
-  span.setAttributes(attrs);
-};
-
-const getMemoryAttributes = (): Attributes => {
-  const memory = mem();
-  if (!memory) return {};
+const DEFAULT_TRACER_NAME = "@hebo-ai/gateway";
 
-  return {
-    "process.memory.usage": memory.rss,
-    "process.memory.heap.used": memory.heapUsed,
-    "process.memory.heap.total": memory.heapTotal,
-  };
-};
+let spanTracer: Tracer | undefined;
+let spanEventsEnabled = false;
 
 const NOOP_SPAN = {
   runWithContext: <T>(fn: () => Promise<T> | T) => fn(),
@@ -31,35 +16,38 @@ const NOOP_SPAN = {
   isExisting: true,
 };
 
-export const startSpan = (name: string, options?: SpanOptions, customTracer?: Tracer) => {
-  const tracer = customTracer ?? trace.getTracer(DEFAULT_TRACER_NAME);
+export const setSpanTracer = (tracer?: Tracer) => {
+  spanTracer = tracer ?? trace.getTracer(DEFAULT_TRACER_NAME);
+};
+
+export const setSpanEventsEnabled = (level?: TelemetrySignalLevel) => {
+  spanEventsEnabled = level === "recommended" || level === "full";
+};
+
+export const startSpan = (name: string, options?: SpanOptions) => {
+  if (!spanTracer) {
+    return Object.assign(trace.wrapSpanContext(INVALID_SPAN_CONTEXT), NOOP_SPAN);
+  }
 
   const parentContext = context.active();
   const activeSpan = trace.getActiveSpan();
 
-  const span = tracer.startSpan(
+  const span = spanTracer.startSpan(
     name,
     { kind: activeSpan ? SpanKind.INTERNAL : SpanKind.SERVER, ...options },
     parentContext,
   );
 
-  if (!span.isRecording()) {
-    return Object.assign(trace.wrapSpanContext(INVALID_SPAN_CONTEXT), NOOP_SPAN);
-  }
-
-  maybeSetDynamicAttributes(span, getMemoryAttributes);
-
   const runWithContext = <T>(fn: () => Promise<T> | T) =>
     context.with(trace.setSpan(parentContext, span), fn);
 
   const recordError = (error: unknown) => {
-    const err = toError(error);
+    const err = error instanceof Error ? error : new Error(String(error));
     span.recordException(err);
     span.setStatus({ code: SpanStatusCode.ERROR, message: err.message });
   };
 
   const finish = () => {
-    maybeSetDynamicAttributes(span, getMemoryAttributes);
     span.end();
   };
 
@@ -71,6 +59,10 @@ export const withSpan = async <T>(
   run: () => Promise<T> | T,
   options?: SpanOptions,
 ): Promise<T> => {
+  if (!spanTracer) {
+    return await run();
+  }
+
   const started = startSpan(name, options);
   try {
     return await started.runWithContext(run);
@@ -83,15 +75,11 @@ export const withSpan = async <T>(
 };
 
 export const addSpanEvent = (name: string, attributes?: Attributes) => {
-  const allAttributes = Object.assign(attributes ?? {}, getMemoryAttributes());
-  trace.getActiveSpan()?.addEvent(name, allAttributes);
+  if (!spanEventsEnabled) return;
+  trace.getActiveSpan()?.addEvent(name, attributes);
 };
 
-export const recordSpanError = (error: unknown) => {
-  const span = trace.getActiveSpan();
-  if (!span) return;
-
-  const err = toError(error);
-  span.recordException(err);
-  span.setStatus({ code: SpanStatusCode.ERROR, message: err.message });
+export const setSpanAttributes = (attributes?: Attributes) => {
+  if (!attributes) return;
+  trace.getActiveSpan()?.setAttributes(attributes);
 };

package/src/telemetry/stream.ts

@@ -1,39 +1,33 @@
-export type InstrumentStreamHooks = {
-  onComplete?: (status: number, stats: { bytes: number }) => void;
-  onError?: (error: unknown, status: number) => void;
-};
+const isErrorChunk = (v: unknown) => !!(v as any)?.error;
 
-export const instrumentStream = (
-  src: ReadableStream<Uint8Array>,
-  hooks: InstrumentStreamHooks,
+export const wrapStream = (
+  src: ReadableStream,
+  hooks: { onDone?: (status: number, reason: unknown) => void },
   signal?: AbortSignal,
-): ReadableStream<Uint8Array> => {
-  const stats = { bytes: 0 };
-  let done = false;
+): ReadableStream => {
+  let finishOnce = false;
 
   const finish = (status: number, reason?: unknown) => {
-    if (done) return;
-    done = true;
-
-    if (!reason) reason = signal?.reason;
-
-    if (status >= 400) {
-      hooks.onError?.(reason, status);
-    }
+    if (finishOnce) return;
+    finishOnce = true;
 
-    hooks.onComplete?.(status, stats);
+    hooks.onDone?.(status, reason ?? signal?.reason);
   };
 
-  return new ReadableStream<Uint8Array>({
+  return new ReadableStream({
     async start(controller) {
       const reader = src.getReader();
 
+      const close = (status: number, reason?: unknown) => {
+        finish(status, reason);
+        reader.cancel(reason).catch(() => {});
+        controller.close();
+      };
+
       try {
         for (;;) {
           if (signal?.aborted) {
-            finish(499, signal.reason);
-            reader.cancel(signal.reason).catch(() => {});
-            controller.close();
+            close(499, signal.reason);
             return;
           }
 
@@ -41,18 +35,20 @@ export const instrumentStream = (
           const { value, done } = await reader.read();
           if (done) break;
 
-          stats.bytes += value!.byteLength;
-          controller.enqueue(value!);
+          controller.enqueue(value);
+
+          if (isErrorChunk(value)) {
+            const status = value.error.type === "invalid_request_error" ? 422 : 502;
+            close(status, value.error.message);
+            return;
+          }
         }
 
         finish(200);
         controller.close();
       } catch (err) {
         const status = signal?.aborted ? 499 : (err as any)?.name === "AbortError" ? 503 : 502;
-
-        finish(status, err);
-        reader.cancel(err).catch(() => {});
-        controller.close();
+        close(status, err);
       } finally {
         try {
           reader.releaseLock();
@@ -60,7 +56,7 @@ export const instrumentStream = (
       }
     },
 
-    cancel(reason) {
+    cancel(reason?: unknown) {
       finish(499, reason);
       src.cancel(reason).catch(() => {});
     },