@neutrome/open-ai-router 0.4.2 → 0.5.0

package/src/router/execution-transforms.ts

@@ -0,0 +1,82 @@
+import { cloneProgram } from "@neutrome/lil-engine";
+import type { Program } from "@neutrome/lil-engine";
+import { createExecutionEvent } from "@neutrome/lilsdk-ts";
+import type { ExecutorContext } from "@neutrome/lilsdk-ts";
+import { ExecutionError } from "./errors.ts";
+import { executionTargetAuditRef, executionTargetId } from "./targets.ts";
+import type { ExecutionTarget, ProgramTransform } from "./execution-types.ts";
+
+export async function applyExecutionTransforms(options: {
+  program: Program;
+  target: ExecutionTarget;
+  transforms: ReadonlyMap<string, ProgramTransform>;
+  context: ExecutorContext;
+  requestId: string;
+  executionId: string;
+  parentExecutionId: string | undefined;
+  depth: number;
+  maxDepth: number;
+  observe: ExecutorContext["observe"];
+  validate(program: Program): void;
+}): Promise<Program> {
+  let current = cloneProgram(options.program);
+  for (const transformName of options.target.transforms ?? []) {
+    const transform = options.transforms.get(transformName);
+    if (!transform) {
+      throw new ExecutionError(
+        "transform",
+        `Unknown transform: ${transformName}`,
+        {
+          stage: "transform",
+          details: transformDetails(transformName, options),
+        },
+      );
+    }
+    try {
+      current = await transform.apply(current, options.context);
+    } catch (cause) {
+      throw new ExecutionError(
+        "transform",
+        `Transform ${transform.name} failed`,
+        {
+          stage: "transform",
+          details: transformDetails(transform.name, options),
+          cause,
+        },
+      );
+    }
+    options.validate(current);
+    options.observe(
+      createExecutionEvent({
+        kind: "transform.applied",
+        requestId: options.requestId,
+        executionId: options.executionId,
+        target: executionTargetAuditRef(options.target),
+        data: transformDetails(transform.name, options),
+        ...parent(options.parentExecutionId),
+      }),
+    );
+  }
+  return current;
+}
+
+function transformDetails(
+  transform: string,
+  options: Pick<
+    Parameters<typeof applyExecutionTransforms>[0],
+    "target" | "depth" | "maxDepth"
+  >,
+): Record<string, string | number> {
+  return {
+    transform,
+    target: executionTargetId(options.target),
+    depth: options.depth,
+    maxDepth: options.maxDepth,
+  };
+}
+
+function parent(parentExecutionId: string | undefined): {
+  parentExecutionId?: string;
+} {
+  return parentExecutionId === undefined ? {} : { parentExecutionId };
+}

package/src/router/execution-types.ts

@@ -1,6 +1,7 @@
 import type { Program, ValidationIssue } from "@neutrome/lil-engine";
 import type {
   ExecutionTarget,
+  ExecutionEvent,
   Executor,
   InvokeOptions,
   ProgramTransform,
@@ -11,33 +12,21 @@ import type {
   ExecutionErrorStage,
 } from "./errors.ts";
 
-export type AuditEvent = {
-  kind: string;
-  requestId: string;
-  executionId: string;
-  timestamp: string;
-  parentExecutionId?: string;
-  target?: { kind: "provider" | "executor"; id: string };
-  errorKind?: string;
-  data?: Record<string, unknown>;
-};
-
-export type AuditEventInput = Omit<AuditEvent, "timestamp"> & {
-  timestamp?: string;
-};
-
 export type ProviderInvocationContext = {
   requestId: string;
   executionId: string;
   parentExecutionId?: string;
   target: Extract<ExecutionTarget, { kind: "provider" }>;
   signal: AbortSignal;
-  observe(event: AuditEvent): void;
+  observe(event: ExecutionEvent): void;
 };
 
 export type ProviderInvoker = {
   execute(request: Program, ctx: ProviderInvocationContext): Promise<Program>;
-  stream(request: Program, ctx: ProviderInvocationContext): AsyncIterable<Program>;
+  stream(
+    request: Program,
+    ctx: ProviderInvocationContext,
+  ): AsyncIterable<Program>;
 };
 
 export type ResolveTarget = (
@@ -55,7 +44,7 @@ export type ExecutionRuntimeOptions = {
   transforms?: Readonly<Record<string, ProgramTransform>>;
   providerInvoker?: ProviderInvoker;
   resolveTarget?: ResolveTarget;
-  observe?: (event: AuditEvent) => void;
+  observe?: (event: ExecutionEvent) => void;
   signal?: AbortSignal;
   requestIdFactory?: () => string;
   executionIdFactory?: () => string;
@@ -73,8 +62,8 @@ export type {
   ExecutionErrorKind,
   ExecutionErrorStage,
   ExecutionTarget,
+  ExecutionEvent,
   Executor,
   InvokeOptions,
   ProgramTransform,
-  ValidationIssue,
 };

package/src/router/index.ts

@@ -1,4 +1,4 @@
-export { createAuditEvent, createInMemoryAuditSink } from "./audit.ts";
+export { createInMemoryAuditSink } from "./audit.ts";
 export { ExecutionError, isExecutionError } from "./errors.ts";
 export { createExecutionRuntime } from "./execution-runtime.ts";
 export { createRouterRuntime } from "./runtime.ts";
@@ -14,7 +14,11 @@ export {
   resolveRequestTarget,
   RouterError,
 } from "./execute.ts";
-export { listConfiguredModels, resolveInvocationTarget, resolveInvocationTargets } from "./resolve.ts";
+export {
+  listConfiguredModels,
+  resolveInvocationTarget,
+  resolveInvocationTargets,
+} from "./resolve.ts";
 export { executionTargetAuditRef, executionTargetId } from "./targets.ts";
 export {
   createRemoteMcpToolSet,
@@ -37,12 +41,9 @@ export type {
   ModelNamespaceRuntime,
   ModelRouteRuntime,
   ProviderRuntime,
-  RemoteMcpServerRuntime,
   RouterRuntime,
 } from "./runtime.ts";
 export type {
-  AuditEvent,
-  AuditEventInput,
   ExecutionErrorDetails,
   ExecutionErrorKind,
   ExecutionErrorStage,

package/src/router/mcp.ts

@@ -3,7 +3,7 @@ import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/
 import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import type { ExecutorContext, Tool } from "@neutrome/lilsdk-ts";
 import type { RemoteMcpServerRuntime, RouterRuntime } from "./runtime.ts";
-import { observeTimingEvent } from "../timing.ts";
+import { observeTimedExecution } from "./execution-events.ts";
 
 export type RemoteMcpToolCallClient = {
   callTool(name: string, args: Record<string, unknown>): Promise<unknown>;
@@ -26,7 +26,9 @@ export function createRemoteMcpToolSet(
   factory?: RemoteMcpClientFactory,
 ): RemoteMcpToolSet {
   const pool = new RemoteMcpSessionPool(
-    factory ?? ((server, signal) => createStreamableHttpMcpClient(server, signal, runtime.fetchImpl)),
+    factory ??
+      ((server, signal) =>
+        createStreamableHttpMcpClient(server, signal, runtime.fetchImpl)),
   );
   const tools: Tool[] = [];
   const usedNames = new Map<string, number>();
@@ -43,12 +45,14 @@ export function createRemoteMcpToolSet(
 
       tools.push({
         name,
-        description: snapshot.description || `Remote MCP tool ${server.name}/${snapshot.name}`,
+        description:
+          snapshot.description ||
+          `Remote MCP tool ${server.name}/${snapshot.name}`,
         schema: normalizeSchema(snapshot.inputSchema),
         async execute(args: Record<string, unknown>, ctx: ExecutorContext) {
           const connectStartedMs = Date.now();
           const client = await pool.get(server, ctx.signal);
-          observeTimingEvent({
+          observeTimedExecution({
             observe: ctx.observe,
             kind: "remote_mcp.connect",
             requestId: ctx.requestId,
@@ -63,7 +67,7 @@ export function createRemoteMcpToolSet(
           });
           const callStartedMs = Date.now();
           const result = await client.callTool(snapshot.name, args);
-          observeTimingEvent({
+          observeTimedExecution({
             observe: ctx.observe,
             kind: "remote_mcp.tool_call",
             requestId: ctx.requestId,
@@ -126,7 +130,9 @@ export function stringifyMcpToolResult(result: unknown): string {
   }
   const body = chunks.filter(Boolean).join("\n");
   if (result.isError === true) {
-    return body ? `MCP tool returned an error:\n${body}` : "MCP tool returned an error.";
+    return body
+      ? `MCP tool returned an error:\n${body}`
+      : "MCP tool returned an error.";
   }
   return body || JSON.stringify(result);
 }

package/src/router/provider-invoker.ts

@@ -0,0 +1,134 @@
+import {
+  emitProviderRequest,
+  isStreaming,
+  parseProviderResponse,
+  setModel,
+} from "@neutrome/lil-engine";
+import { toProviderError } from "./error-codec.ts";
+import {
+  createUpstreamSignal,
+  fetchProvider,
+  getProviderOrThrow,
+  normalizeUpstreamAbort,
+} from "./upstream-client.ts";
+import { emitProviderUsage, injectStreamUsage } from "./provider-telemetry.ts";
+import type { ProviderInvoker } from "./execution-types.ts";
+import type { RequestContext } from "../upstream-auth/types.ts";
+import { streamProvider } from "./provider-stream.ts";
+import type { RouterRuntime } from "./runtime.ts";
+import { observeTimedExecution } from "./execution-events.ts";
+
+const DEFAULT_UPSTREAM_TIMEOUT_MS = 120_000;
+
+export function createFetchProviderInvoker(
+  runtime: RouterRuntime,
+  ctx: RequestContext,
+  upstreamTimeoutMs = DEFAULT_UPSTREAM_TIMEOUT_MS,
+): ProviderInvoker {
+  return {
+    async execute(request, providerCtx) {
+      const provider = getProviderOrThrow(runtime, providerCtx.target.provider);
+      const timed = createUpstreamSignal(providerCtx.signal, upstreamTimeoutMs);
+
+      try {
+        const resolvedRequest = setModel(request, providerCtx.target.model);
+        const encodeStartedMs = Date.now();
+        let reqBody = emitProviderRequest(provider.style, resolvedRequest);
+        if (
+          provider.style === "chat-completions" &&
+          isStreaming(resolvedRequest)
+        ) {
+          reqBody = injectStreamUsage(reqBody);
+        }
+        observeTimedExecution({
+          observe: providerCtx.observe,
+          kind: "provider.request_encoded",
+          requestId: providerCtx.requestId,
+          executionId: providerCtx.executionId,
+          ...(providerCtx.parentExecutionId
+            ? { parentExecutionId: providerCtx.parentExecutionId }
+            : {}),
+          startedAtMs: encodeStartedMs,
+          finishedAtMs: Date.now(),
+          data: {
+            provider: provider.name,
+            model: providerCtx.target.model,
+            bytes: reqBody.byteLength,
+          },
+        });
+        const reqBytes = reqBody.byteLength;
+        const fetchStartedMs = Date.now();
+        const upstream = await fetchProvider(
+          runtime,
+          ctx,
+          provider,
+          reqBody,
+          timed.signal,
+          providerCtx,
+        );
+        observeTimedExecution({
+          observe: providerCtx.observe,
+          kind: "provider.fetch_headers",
+          requestId: providerCtx.requestId,
+          executionId: providerCtx.executionId,
+          ...(providerCtx.parentExecutionId
+            ? { parentExecutionId: providerCtx.parentExecutionId }
+            : {}),
+          startedAtMs: fetchStartedMs,
+          finishedAtMs: Date.now(),
+          data: {
+            provider: provider.name,
+            model: providerCtx.target.model,
+            status: upstream.status,
+          },
+        });
+        if (!upstream.ok) {
+          const providerError = await toProviderError(provider, upstream);
+          throw providerError;
+        }
+
+        const readStartedMs = Date.now();
+        const bytes = new Uint8Array(await upstream.arrayBuffer());
+        observeTimedExecution({
+          observe: providerCtx.observe,
+          kind: "provider.response_body_read",
+          requestId: providerCtx.requestId,
+          executionId: providerCtx.executionId,
+          ...(providerCtx.parentExecutionId
+            ? { parentExecutionId: providerCtx.parentExecutionId }
+            : {}),
+          startedAtMs: readStartedMs,
+          finishedAtMs: Date.now(),
+          data: {
+            provider: provider.name,
+            model: providerCtx.target.model,
+            bytes: bytes.byteLength,
+          },
+        });
+        const response = parseProviderResponse(provider.style, bytes);
+        emitProviderUsage(
+          providerCtx,
+          response,
+          provider.name,
+          reqBytes,
+          bytes.byteLength,
+        );
+        return response;
+      } catch (error) {
+        throw normalizeUpstreamAbort(error, timed.signal);
+      } finally {
+        timed.cleanup();
+      }
+    },
+
+    stream(request, providerCtx) {
+      return streamProvider(
+        runtime,
+        ctx,
+        upstreamTimeoutMs,
+        request,
+        providerCtx,
+      );
+    },
+  };
+}

package/src/router/provider-stream.ts

@@ -0,0 +1,192 @@
+import {
+  emitProviderRequest,
+  isStreaming,
+  parseProviderStreamChunk,
+  setModel,
+  type Program,
+  usageObject,
+} from "@neutrome/lil-engine";
+import { RouterError } from "./errors.ts";
+import { toProviderError } from "./error-codec.ts";
+import {
+  createUpstreamSignal,
+  fetchProvider,
+  getProviderOrThrow,
+  normalizeUpstreamAbort,
+} from "./upstream-client.ts";
+import { emitProviderUsage, injectStreamUsage } from "./provider-telemetry.ts";
+import type { ProviderInvocationContext } from "./execution-types.ts";
+import type { RequestContext } from "../upstream-auth/types.ts";
+import { isSse } from "../util/headers.ts";
+import { eventDataLines, splitSseEvents } from "../util/sse.ts";
+import type { RouterRuntime } from "./runtime.ts";
+import { observeTimedExecution } from "./execution-events.ts";
+
+const encoder = new TextEncoder();
+const decoder = new TextDecoder();
+
+export async function* streamProvider(
+  runtime: RouterRuntime,
+  ctx: RequestContext,
+  upstreamTimeoutMs: number,
+  request: Program,
+  providerCtx: ProviderInvocationContext,
+): AsyncGenerator<Program> {
+  const provider = getProviderOrThrow(runtime, providerCtx.target.provider);
+  const timed = createUpstreamSignal(providerCtx.signal, upstreamTimeoutMs);
+  let reader: ReadableStreamDefaultReader<Uint8Array> | null = null;
+  let buffer = "";
+
+  try {
+    const resolvedRequest = setModel(request, providerCtx.target.model);
+    const encodeStartedMs = Date.now();
+    let reqBody = emitProviderRequest(provider.style, resolvedRequest);
+    if (provider.style === "chat-completions" && isStreaming(resolvedRequest)) {
+      reqBody = injectStreamUsage(reqBody);
+    }
+    observeTimedExecution({
+      observe: providerCtx.observe,
+      kind: "provider.request_encoded",
+      requestId: providerCtx.requestId,
+      executionId: providerCtx.executionId,
+      ...(providerCtx.parentExecutionId
+        ? { parentExecutionId: providerCtx.parentExecutionId }
+        : {}),
+      startedAtMs: encodeStartedMs,
+      finishedAtMs: Date.now(),
+      data: {
+        provider: provider.name,
+        model: providerCtx.target.model,
+        bytes: reqBody.byteLength,
+      },
+    });
+    const reqBytes = reqBody.byteLength;
+    const fetchStartedMs = Date.now();
+    const upstream = await fetchProvider(
+      runtime,
+      ctx,
+      provider,
+      reqBody,
+      timed.signal,
+      providerCtx,
+    );
+    observeTimedExecution({
+      observe: providerCtx.observe,
+      kind: "provider.fetch_headers",
+      requestId: providerCtx.requestId,
+      executionId: providerCtx.executionId,
+      ...(providerCtx.parentExecutionId
+        ? { parentExecutionId: providerCtx.parentExecutionId }
+        : {}),
+      startedAtMs: fetchStartedMs,
+      finishedAtMs: Date.now(),
+      data: {
+        provider: provider.name,
+        model: providerCtx.target.model,
+        status: upstream.status,
+      },
+    });
+    if (!upstream.ok) {
+      const providerError = await toProviderError(provider, upstream);
+      throw providerError;
+    }
+    if (!isSse(upstream.headers) || !upstream.body) {
+      throw new RouterError(
+        `Provider ${provider.name} did not return an event stream`,
+        502,
+        "invalid_upstream_stream",
+      );
+    }
+
+    reader = upstream.body.getReader();
+    let firstByteObserved = false;
+    const streamReadStartedMs = Date.now();
+    let lastChunkWithUsage: Program | null = null;
+    let respBytes = 0;
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) {
+        break;
+      }
+      if (!firstByteObserved) {
+        firstByteObserved = true;
+        observeTimedExecution({
+          observe: providerCtx.observe,
+          kind: "provider.stream_first_byte",
+          requestId: providerCtx.requestId,
+          executionId: providerCtx.executionId,
+          ...(providerCtx.parentExecutionId
+            ? { parentExecutionId: providerCtx.parentExecutionId }
+            : {}),
+          startedAtMs: streamReadStartedMs,
+          finishedAtMs: Date.now(),
+          data: {
+            provider: provider.name,
+            model: providerCtx.target.model,
+            bytes: value.byteLength,
+          },
+        });
+      }
+      respBytes += value.byteLength;
+      buffer += decoder.decode(value, { stream: true });
+      const split = splitSseEvents(buffer);
+      buffer = split.remainder;
+      for (const event of split.events) {
+        for (const data of eventDataLines(event)) {
+          if (data === "[DONE]") {
+            emitProviderUsage(
+              providerCtx,
+              lastChunkWithUsage,
+              provider.name,
+              reqBytes,
+              respBytes,
+            );
+            return;
+          }
+          if (data[0] !== "{") continue;
+          const chunk = parseProviderStreamChunk(
+            provider.style,
+            encoder.encode(data),
+          );
+          if (usageObject(chunk) !== undefined) lastChunkWithUsage = chunk;
+          yield chunk;
+        }
+      }
+    }
+
+    if (buffer.length > 0) {
+      for (const data of eventDataLines(buffer)) {
+        if (data === "[DONE]") {
+          emitProviderUsage(
+            providerCtx,
+            lastChunkWithUsage,
+            provider.name,
+            reqBytes,
+            respBytes,
+          );
+          return;
+        }
+        if (data[0] !== "{") continue;
+        const chunk = parseProviderStreamChunk(
+          provider.style,
+          encoder.encode(data),
+        );
+        if (usageObject(chunk) !== undefined) lastChunkWithUsage = chunk;
+        yield chunk;
+      }
+    }
+    emitProviderUsage(
+      providerCtx,
+      lastChunkWithUsage,
+      provider.name,
+      reqBytes,
+      respBytes,
+    );
+  } catch (error) {
+    throw normalizeUpstreamAbort(error, timed.signal);
+  } finally {
+    await reader?.cancel().catch(() => {});
+    reader?.releaseLock();
+    timed.cleanup();
+  }
+}

package/src/router/provider-telemetry.ts

@@ -0,0 +1,59 @@
+import { usageObject, type Program } from "@neutrome/lil-engine";
+import { createExecutionEvent } from "@neutrome/lilsdk-ts";
+import type { ProviderInvocationContext } from "./execution-types.ts";
+
+const encoder = new TextEncoder();
+const decoder = new TextDecoder();
+
+export function emitProviderUsage(
+  providerCtx: ProviderInvocationContext,
+  response: Program | null,
+  provider: string,
+  requestBytes: number,
+  responseBytes: number,
+): void {
+  const usage = response ? usageObject(response) : undefined;
+  const u = (usage && typeof usage === "object" ? usage : {}) as Record<
+    string,
+    unknown
+  >;
+  const tokensInput = asNum(u.prompt_tokens) || asNum(u.input_tokens);
+  const tokensOutput = asNum(u.completion_tokens) || asNum(u.output_tokens);
+  const details = u.prompt_tokens_details as
+    | Record<string, unknown>
+    | undefined;
+  const cachedTokensInput = details ? asNum(details.cached_tokens) : 0;
+  providerCtx.observe(
+    createExecutionEvent({
+      kind: "provider.usage",
+      requestId: providerCtx.requestId,
+      executionId: providerCtx.executionId,
+      ...(providerCtx.parentExecutionId
+        ? { parentExecutionId: providerCtx.parentExecutionId }
+        : {}),
+      data: {
+        provider,
+        model: providerCtx.target.model,
+        tokensInput,
+        tokensOutput,
+        cachedTokensInput,
+        requestBytes,
+        responseBytes,
+      },
+    }),
+  );
+}
+
+function asNum(v: unknown): number {
+  return typeof v === "number" && v >= 0 ? v : 0;
+}
+
+export function injectStreamUsage(body: Uint8Array): Uint8Array {
+  try {
+    const obj = JSON.parse(decoder.decode(body)) as Record<string, unknown>;
+    obj.stream_options = { include_usage: true };
+    return encoder.encode(JSON.stringify(obj));
+  } catch {
+    return body;
+  }
+}

package/src/router/remote-mcp-execution.ts

@@ -0,0 +1,61 @@
+import { createTargetExecutor } from "@neutrome/lilsdk-ts/managed";
+import { connectTools } from "@neutrome/lilsdk-ts/tools";
+import type {
+  ExecutionTarget,
+  Executor,
+  ProgramTransform,
+} from "./execution-types.ts";
+import { createRemoteMcpToolSet, type RemoteMcpClientFactory } from "./mcp.ts";
+import type { ResolvedTarget } from "./resolve.ts";
+import type { RouterRuntime } from "./runtime.ts";
+
+type RemoteMcpExecutionOptions = {
+  executorImplementations?: Readonly<Record<string, Executor>>;
+  transforms?: Readonly<Record<string, ProgramTransform>>;
+  upstreamTimeoutMs?: number;
+  remoteMcpClientFactory?: RemoteMcpClientFactory;
+  incomingExecutionId?: string;
+};
+
+export function prepareRemoteMcpExecution(
+  runtime: RouterRuntime,
+  resolved: ResolvedTarget,
+  options: RemoteMcpExecutionOptions,
+): {
+  target: ExecutionTarget;
+  options: RemoteMcpExecutionOptions;
+  close: () => Promise<void>;
+} {
+  if (resolved.mcps.length === 0) {
+    return { target: resolved.target, options, close: async () => {} };
+  }
+
+  const remoteMcp = createRemoteMcpToolSet(
+    runtime,
+    resolved.mcps,
+    options.remoteMcpClientFactory,
+  );
+  if (remoteMcp.tools.length === 0) {
+    return { target: resolved.target, options, close: remoteMcp.close };
+  }
+
+  const executorName = `__remote_mcp_${crypto.randomUUID()}`;
+  return {
+    target: {
+      kind: "executor",
+      executorId: executorName,
+      alias: resolved.requestedModel,
+    },
+    options: {
+      ...options,
+      executorImplementations: {
+        ...(options.executorImplementations ?? {}),
+        [executorName]: connectTools(
+          remoteMcp.tools,
+          createTargetExecutor(resolved.target),
+        ),
+      },
+    },
+    close: remoteMcp.close,
+  };
+}

package/src/router/request-program.ts

@@ -0,0 +1,16 @@
+import {
+  parseChatCompletionsRequest,
+  parseProviderRequest,
+  type Program,
+  type ProviderStyle,
+} from "@neutrome/lil-engine";
+
+export async function parseRequestProgram(
+  style: ProviderStyle,
+  request: Request,
+): Promise<Program> {
+  const body = new Uint8Array(await request.arrayBuffer());
+  return style === "chat-completions"
+    ? parseChatCompletionsRequest(body)
+    : parseProviderRequest(style, body);
+}