Use eventsource-parser to clean up the implemenation

Author: aron

index.js

@@ -293,10 +293,9 @@ class Replicate {
 
     if (prediction.urls && prediction.urls.stream) {
       const { signal } = options;
-      const stream = await createReadableStream({
+      const stream = createReadableStream({
         url: prediction.urls.stream,
-        EventSource: this.EventSource,
-        ReadableStream: this.ReadableStream,
+        fetch: this.fetch,
         options: { signal },
       });
       yield* stream;

index.test.ts

@@ -7,7 +7,6 @@ import Replicate, {
   parseProgressFromLogs,
 } from "replicate";
 import nock from "nock";
-import fetch from "cross-fetch";
 import { createReadableStream } from "./lib/stream";
 import { PassThrough } from "node:stream";
 
@@ -23,7 +22,6 @@ describe("Replicate client", () => {
 
   beforeEach(() => {
     client = new Replicate({ auth: "test-token" });
-    client.fetch = fetch;
 
     unmatched = [];
     nock.emitter.on("no match", handleNoMatch);
@@ -1199,8 +1197,9 @@ describe("Replicate client", () => {
         .matchHeader("Accept", "text/event-stream")
         .reply(status, body);
 
-      return await createReadableStream({
+      return createReadableStream({
         url: `${streamEndpoint}/fake_stream`,
+        fetch: fetch,
       });
     }
 
@@ -1484,7 +1483,9 @@ describe("Replicate client", () => {
         done: false,
         value: { event: "output", id: "EVENT_1", data: "hello world" },
       });
-      await expect(iterator.next()).rejects.toThrowError("Unexpected Error");
+      await expect(iterator.next()).rejects.toThrowError(
+        "An unexpected error occurred"
+      );
       expect(await iterator.next()).toEqual({ done: true });
     });
 

lib/stream.js

@@ -1,6 +1,7 @@
 // Attempt to use readable-stream if available, attempt to use the built-in stream module.
 
 const ApiError = require("./error");
+const { EventSourceParserStream } = require("eventsource-parser/stream");
 
 /**
  * A server-sent event.
@@ -38,87 +39,49 @@ class ServerSentEvent {
  *
  * @param {object} config
  * @param {string} config.url The URL to connect to.
- * @param {any} [config.EventSource] A standards compliant EventSource implementation.
- * @param {any} [config.ReadableStream] A standards compliant ReadableStream implementation.
+ * @param {typeof fetch} [config.fetch] The URL to connect to.
  * @param {object} [config.options] The EventSource options.
  * @returns {Promise<ReadableStream & AsyncIterable<ServerSentEvent>>}
  */
-async function createReadableStream({
-  url,
-  EventSource = globalThis.EventSource,
-  ReadableStream = globalThis.ReadableStream,
-  options = {},
-}) {
-  const EventSourceClass = EventSource
-    ? EventSource
-    : (await import("eventsource")).default;
-  const source = new EventSourceClass(url, options);
-
-  const stream = new ReadableStream({
-    cancel() {
-      source.close();
-    },
-
+function createReadableStream({ url, fetch, options = {} }) {
+  return new ReadableStream({
     async start(controller) {
-      return new Promise((resolve) => {
-        source.addEventListener("output", (evt) => {
-          const entry = new ServerSentEvent(
-            evt.type,
-            evt.data,
-            evt.lastEventId
-          );
-          controller.enqueue(entry);
-        });
-
-        source.addEventListener("done", (evt) => {
-          const entry = new ServerSentEvent(
-            evt.type,
-            evt.data,
-            evt.lastEventId
-          );
-          controller.enqueue(entry);
-          source.close();
-          controller.close();
-        });
+      const init = {
+        ...options,
+        headers: {
+          ...options.headers,
+          Accept: "text/event-stream",
+        },
+      };
+      const response = await fetch(url, init);
 
-        source.addEventListener("open", (_evt) => {
-          resolve();
-        });
+      if (!response.ok) {
+        const text = await response.text();
+        const request = new Request(url, init);
+        controller.error(
+          new ApiError(
+            `Request to ${url} failed with status ${response.status}`,
+            request,
+            response
+          )
+        );
+      }
 
-        source.addEventListener("error", (evt) => {
-          // HTTP Error
-          if (typeof evt.status === "number") {
-            source.close();
-            controller.error(
-              new ApiError(
-                `Request to ${url} failed with status ${evt.status} ${
-                  evt.message ?? ""
-                }`.trim()
-              )
-            );
-            return;
-          }
-
-          // Connection closed
-          if (!evt.message && source.readyState === 0) {
-            controller.close();
-            source.close();
-            return;
-          }
-
-          // Other
-          source.close();
-          controller.error(new Error(evt.message ?? "Unexpected Error"));
-        });
-
-        options.signal?.addEventListener("abort", () => {
-          source.close();
-        });
-      });
+      const stream = response.body
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough(new EventSourceParserStream());
+      for await (const event of stream) {
+        if (event.event === "error") {
+          controller.error(new Error(event.data));
+        } else {
+          controller.enqueue(
+            new ServerSentEvent(event.event, event.data, event.id)
+          );
+        }
+      }
+      controller.close();
     },
   });
-
-  return stream;
 }
 
 module.exports = {