@neutrome/open-ai-router 0.1.12 → 0.1.13

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@neutrome/open-ai-router",
-  "version": "0.1.12",
+  "version": "0.1.13",
   "type": "module",
   "exports": {
     ".": "./src/index.ts"

package/src/app/factory.test.ts

@@ -0,0 +1,156 @@
+import { describe, expect, it, vi } from "vitest";
+import type { RouterConfig } from "../router/index.ts";
+import { createRouterApp } from "./factory.ts";
+
+const encoder = new TextEncoder();
+
+const config: RouterConfig = {
+  providers: {
+    openrouter: {
+      api_base_url: "https://openrouter.ai/api/v1",
+      style: "chat-completions",
+      exports: ["*"],
+    },
+  },
+  executors: {
+    default: {
+      exports: ["*"],
+      models: {
+        "gpt-4o": { provider: "openrouter", model: "gpt-4o" },
+      },
+    },
+  },
+};
+
+describe("createRouterApp trace finalization", () => {
+  it("keeps non-stream trace hooks alive with waitUntil", async () => {
+    const waited: Promise<unknown>[] = [];
+    const onTrace = vi.fn(async () => {});
+    const app = createRouterApp({
+      config,
+      executors: {},
+      auth: {
+        async authenticate() {
+          return {};
+        },
+        canAccess() {
+          return true;
+        },
+      },
+      hooks: { onTrace },
+      fetchImpl: vi.fn(async () =>
+        new Response(
+          JSON.stringify({
+            id: "resp_1",
+            model: "gpt-4o",
+            choices: [
+              {
+                index: 0,
+                message: { role: "assistant", content: "hello" },
+                finish_reason: "stop",
+              },
+            ],
+          }),
+          { headers: { "content-type": "application/json; charset=utf-8" } },
+        )),
+    });
+
+    const response = await app.fetch(
+      new Request("http://local/v1/chat/completions", {
+        method: "POST",
+        headers: { "content-type": "application/json" },
+        body: JSON.stringify({
+          model: "gpt-4o",
+          messages: [{ role: "user", content: "hi" }],
+        }),
+      }),
+      {},
+      createExecutionContext(waited) as never,
+    );
+
+    expect(response.status).toBe(200);
+    expect(waited).toHaveLength(1);
+    await waited[0];
+    expect(onTrace).toHaveBeenCalledTimes(1);
+  });
+
+  it("keeps stream trace hooks alive until the stream closes", async () => {
+    const waited: Promise<unknown>[] = [];
+    const onTrace = vi.fn(async () => {});
+    const app = createRouterApp({
+      config,
+      executors: {},
+      auth: {
+        async authenticate() {
+          return {};
+        },
+        canAccess() {
+          return true;
+        },
+      },
+      hooks: { onTrace },
+      fetchImpl: vi.fn(async () => {
+        const stream = new ReadableStream<Uint8Array>({
+          start(controller) {
+            controller.enqueue(
+              encoder.encode(
+                `data: ${JSON.stringify({
+                  id: "chunk_1",
+                  model: "gpt-4o",
+                  choices: [{ index: 0, delta: { role: "assistant", content: "Hi" } }],
+                })}\n\n`,
+              ),
+            );
+            controller.enqueue(
+              encoder.encode(
+                `data: ${JSON.stringify({
+                  id: "chunk_1",
+                  model: "gpt-4o",
+                  choices: [{ index: 0, delta: {}, finish_reason: "stop" }],
+                })}\n\n`,
+              ),
+            );
+            controller.enqueue(encoder.encode("data: [DONE]\n\n"));
+            controller.close();
+          },
+        });
+
+        return new Response(stream, {
+          headers: { "content-type": "text/event-stream; charset=utf-8" },
+        });
+      }),
+    });
+
+    const response = await app.fetch(
+      new Request("http://local/v1/chat/completions", {
+        method: "POST",
+        headers: { "content-type": "application/json" },
+        body: JSON.stringify({
+          model: "gpt-4o",
+          stream: true,
+          messages: [{ role: "user", content: "hi" }],
+        }),
+      }),
+      {},
+      createExecutionContext(waited) as never,
+    );
+
+    expect(response.headers.get("content-type")).toBe("text/event-stream; charset=utf-8");
+    expect(waited).toHaveLength(0);
+
+    await response.text();
+
+    expect(waited).toHaveLength(1);
+    await waited[0];
+    expect(onTrace).toHaveBeenCalledTimes(1);
+  });
+});
+
+function createExecutionContext(waited: Promise<unknown>[]) {
+  return {
+    waitUntil(promise: Promise<unknown>) {
+      waited.push(promise);
+    },
+    passThroughOnException() {},
+  };
+}

package/src/app/factory.ts

@@ -112,6 +112,7 @@ export function createRouterApp(options: RouterAppOptions) {
     }
 
     const collector = createTraceCollector();
+    const onTrace = options.hooks?.onTrace;
     const observe = (event: Parameters<NonNullable<ExecutionRuntimeOptions["observe"]>>[0]) => {
       collector.observe(event);
       options.observe?.(event);
@@ -120,9 +121,12 @@ export function createRouterApp(options: RouterAppOptions) {
     const finalizeTrace = () => {
       if (traceFinalized) return;
       traceFinalized = true;
-      Promise.resolve(options.hooks?.onTrace?.(collector.getTrace(), env)).catch((error) => {
-        console.error("[open-ai-router] onTrace hook failed:", error);
-      });
+      if (!onTrace) return;
+      runBackgroundTask(
+        c,
+        () => onTrace(collector.getTrace(), env),
+        "onTrace hook",
+      );
     };
 
     try {
@@ -208,6 +212,20 @@ function isEventStreamResponse(response: Response): boolean {
   return contentType.startsWith("text/event-stream");
 }
 
+function runBackgroundTask(c: Context, task: () => unknown, label: string): void {
+  const tracked = Promise.resolve()
+    .then(task)
+    .catch((error) => {
+      console.error(`[open-ai-router] ${label} failed:`, error);
+    });
+  const executionCtx = c.executionCtx as { waitUntil?(promise: Promise<unknown>): void } | undefined;
+  if (executionCtx?.waitUntil) {
+    executionCtx.waitUntil(tracked);
+    return;
+  }
+  void tracked;
+}
+
 function wrapResponseWithFinalize(response: Response, finalize: () => void): Response {
   const reader = response.body!.getReader();
   const body = new ReadableStream<Uint8Array>({