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

package/src/telemetry/stream.ts

@@ -1,54 +1,51 @@
-const isErrorChunk = (v: unknown) => !!(v as any)?.error;
+import { toOpenAIError } from "../errors/openai";
+
+const isErrorChunk = (v: unknown) => v instanceof Error || !!(v as any)?.error;
 
 export const wrapStream = (
   src: ReadableStream,
   hooks: { onDone?: (status: number, reason: unknown) => void },
-  signal?: AbortSignal,
 ): ReadableStream => {
-  let finishOnce = false;
-
-  const finish = (status: number, reason?: unknown) => {
-    if (finishOnce) return;
-    finishOnce = true;
+  let finished = false;
 
-    hooks.onDone?.(status, reason ?? signal?.reason);
+  const done = (
+    reader: ReadableStreamDefaultReader,
+    controller: ReadableStreamDefaultController,
+    status: number,
+    reason?: unknown,
+  ) => {
+    if (!finished) {
+      finished = true;
+      hooks.onDone?.(status, reason);
+    }
+    reader.cancel(reason).catch(() => {});
+    controller.close();
   };
 
   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) {
-            close(499, signal.reason);
-            return;
-          }
-
           // eslint-disable-next-line no-await-in-loop
-          const { value, done } = await reader.read();
-          if (done) break;
+          const { value, done: eof } = await reader.read();
+          if (eof) break;
 
-          controller.enqueue(value);
+          const out = isErrorChunk(value) ? toOpenAIError(value) : value;
+          controller.enqueue(out);
 
-          if (isErrorChunk(value)) {
-            const status = value.error.type === "invalid_request_error" ? 422 : 502;
-            close(status, value.error.message);
+          if (out !== value) {
+            const status = out.error?.type === "invalid_request_error" ? 422 : 502;
+            done(reader, controller, status, value);
             return;
           }
         }
 
-        finish(200);
-        controller.close();
+        done(reader, controller, 200);
       } catch (err) {
-        const status = signal?.aborted ? 499 : (err as any)?.name === "AbortError" ? 503 : 502;
-        close(status, err);
+        controller.enqueue(toOpenAIError(err));
+        done(reader, controller, 502, err);
       } finally {
         try {
           reader.releaseLock();
@@ -56,8 +53,11 @@ export const wrapStream = (
       }
     },
 
-    cancel(reason?: unknown) {
-      finish(499, reason);
+    cancel(reason) {
+      if (!finished) {
+        finished = true;
+        hooks.onDone?.(499, reason);
+      }
       src.cancel(reason).catch(() => {});
     },
   });

package/src/types.ts

@@ -8,7 +8,6 @@ import type {
 } from "./endpoints/chat-completions/schema";
 import type { Embeddings, EmbeddingsBody } from "./endpoints/embeddings/schema";
 import type { Model, ModelList } from "./endpoints/models";
-import type { OpenAIError } from "./errors/openai";
 import type { Logger, LoggerConfig } from "./logger";
 import type { ModelCatalog, ModelId } from "./models/types";
 import type { ProviderId, ProviderRegistry } from "./providers/types";
@@ -76,7 +75,7 @@ export type GatewayContext = {
    */
   result?:
     | ChatCompletions
-    | ReadableStream<ChatCompletionsChunk | OpenAIError>
+    | ReadableStream<ChatCompletionsChunk | Error>
     | Embeddings
     | Model
     | ModelList;
@@ -150,11 +149,9 @@ export type GatewayHooks = {
   ) =>
     | void
     | ChatCompletions
-    | ReadableStream<ChatCompletionsChunk | OpenAIError>
+    | ReadableStream<ChatCompletionsChunk | Error>
     | Embeddings
-    | Promise<
-        void | ChatCompletions | ReadableStream<ChatCompletionsChunk | OpenAIError> | Embeddings
-      >;
+    | Promise<void | ChatCompletions | ReadableStream<ChatCompletionsChunk | Error> | Embeddings>;
   /**
    * Runs after the lifecycle has produced the final Response.
    * @returns Replacement Response, or undefined to keep original.