fix: use subscribable as the source of truth (#93)

.changeset/loud-cooks-attack.md

@@ -0,0 +1,5 @@
+---
+"@llama-flow/core": patch
+---
+
+fix: use `subscribable` as the source of truth

packages/core/jsr.json

@@ -4,10 +4,6 @@
   "exports": {
     ".": "./src/core/index.ts",
     "./async-context": "./src/async-context/index.ts",
-    "./interrupter/hono": "./src/interrupter/hono.ts",
-    "./interrupter/mcp": "./src/interrupter/mcp.ts",
-    "./interrupter/next": "./src/interrupter/next.ts",
-    "./interrupter/promise": "./src/interrupter/promise.ts",
     "./middleware/store": "./src/middleware/store.ts",
     "./middleware/trace-events": "./src/middleware/trace-events.ts",
     "./middleware/validation": "./src/middleware/validation.ts",

packages/core/src/core/context.ts

@@ -198,36 +198,23 @@ export const createContext = ({
     handlerContext: HandlerContext,
   ): WorkflowContext => ({
     get stream() {
-      let unsubscribe: () => void;
-      const stream = new ReadableStream({
-        start: async (controller) => {
-          unsubscribe =
-            rootWorkflowContext.__internal__call_send_event.subscribe(
-              (newEvent: WorkflowEventData<any>) => {
-                let currentEventContext = eventContextWeakMap.get(newEvent);
-                while (currentEventContext) {
-                  if (currentEventContext === handlerContext) {
-                    controller.enqueue(newEvent);
-                    break;
-                  }
-                  currentEventContext = currentEventContext.prev;
-                }
-              },
-            );
-        },
-        cancel: () => {
-          if (unsubscribe) {
-            unsubscribe();
+      const subscribable = createSubscribable<
+        [event: WorkflowEventData<any>],
+        void
+      >();
+      rootWorkflowContext.__internal__call_send_event.subscribe(
+        (newEvent: WorkflowEventData<any>) => {
+          let currentEventContext = eventContextWeakMap.get(newEvent);
+          while (currentEventContext) {
+            if (currentEventContext === handlerContext) {
+              subscribable.publish(newEvent);
+              break;
+            }
+            currentEventContext = currentEventContext.prev;
           }
         },
-      });
-      return new WorkflowStream(
-        rootWorkflowContext.__internal__call_send_event as unknown as Subscribable<
-          [event: WorkflowEventData<any>],
-          void
-        >,
-        stream,
       );
+      return new WorkflowStream(subscribable, null);
     },
     get signal() {
       return handlerContext.abortController.signal;

packages/core/src/core/stream.ts

@@ -27,10 +27,42 @@ export class WorkflowStream<R = any>
 
   constructor(
     subscribable: Subscribable<[event: WorkflowEventData<any>], void>,
-    stream: ReadableStream<WorkflowEventData<any>>,
+    rootStream: ReadableStream<WorkflowEventData<any>> | null,
   ) {
     this.#subscribable = subscribable;
-    this.#stream = stream;
+    let unsubscribe: () => void;
+    this.#stream =
+      rootStream ??
+      new ReadableStream<WorkflowEventData<any>>({
+        start: (controller) => {
+          unsubscribe = subscribable.subscribe((event) => {
+            controller.enqueue(event);
+          });
+        },
+        cancel: () => {
+          unsubscribe();
+        },
+      });
+  }
+
+  static fromReadableStream(
+    stream: ReadableStream<WorkflowEventData<any>>,
+  ): WorkflowStream {
+    const subscribable = createSubscribable<
+      [event: WorkflowEventData<any>],
+      void
+    >();
+    return new WorkflowStream(
+      subscribable,
+      stream.pipeThrough(
+        new TransformStream<WorkflowEventData<any>>({
+          transform: (event, controller) => {
+            subscribable.publish(event);
+            controller.enqueue(event);
+          },
+        }),
+      ),
+    );
   }
 
   static fromResponse(
@@ -46,29 +78,31 @@ export class WorkflowStream<R = any>
       void
     >();
     const events = Object.values(eventMap);
-    const stream = body
-      .pipeThrough(new TextDecoderStream())
-      .pipeThrough<WorkflowEventData<any>>(
-        new TransformStream({
-          transform: (chunk, controller) => {
-            const eventData = JSON.parse(chunk) as {
-              data: ReturnType<WorkflowEvent<any>["with"]>;
-              uniqueId: string;
-            };
-            const targetEvent = events.find(
-              (e) => e.uniqueId === eventData.uniqueId,
-            );
-            if (targetEvent) {
-              const ev = targetEvent.with(
-                eventData.data,
-              ) as WorkflowEventData<any>;
-              subscribable.publish(ev);
-              controller.enqueue(ev);
-            }
-          },
-        }),
-      );
-    return new WorkflowStream(subscribable, stream);
+    return new WorkflowStream(
+      subscribable,
+      body
+        .pipeThrough(new TextDecoderStream())
+        .pipeThrough<WorkflowEventData<any>>(
+          new TransformStream({
+            transform: (chunk, controller) => {
+              const eventData = JSON.parse(chunk) as {
+                data: ReturnType<WorkflowEvent<any>["with"]>;
+                uniqueId: string;
+              };
+              const targetEvent = events.find(
+                (e) => e.uniqueId === eventData.uniqueId,
+              );
+              if (targetEvent) {
+                const ev = targetEvent.with(
+                  eventData.data,
+                ) as WorkflowEventData<any>;
+                subscribable.publish(ev);
+                controller.enqueue(ev);
+              }
+            },
+          }),
+        ),
+    );
   }
 
   toResponse(

packages/core/src/stream/filter.ts

@@ -8,9 +8,9 @@ export function filter<
   Event extends WorkflowEventData<any>,
   Final extends Event,
 >(
-  stream: ReadableStream<Event> | WorkflowStream,
+  stream: ReadableStream<Event> | WorkflowStream<Event>,
   cond: (event: Event) => event is Final,
-): ReadableStream<Final> {
+): ReadableStream<Final> | WorkflowStream<Final> {
   return stream.pipeThrough(
     new TransformStream<Event, Final>({
       transform(event, controller) {
@@ -19,5 +19,5 @@ export function filter<
         }
       },
     }),
-  ) as ReadableStream<Final>;
+  );
 }

packages/core/src/stream/until.ts

@@ -11,42 +11,44 @@ const isWorkflowEvent = (value: unknown): value is WorkflowEvent<any> =>
   "include" in value;
 
 export function until(
-  stream: ReadableStream<WorkflowEventData<any>> | WorkflowStream,
+  stream: WorkflowStream | ReadableStream<WorkflowEventData<any>>,
   cond: (event: WorkflowEventData<any>) => boolean | Promise<boolean>,
-): ReadableStream<WorkflowEventData<any>>;
+): WorkflowStream;
 export function until<Stop>(
-  stream: ReadableStream<WorkflowEventData<any>> | WorkflowStream,
+  stream: WorkflowStream | ReadableStream<WorkflowEventData<any>>,
   cond: WorkflowEvent<Stop>,
-): ReadableStream<WorkflowEventData<any> | WorkflowEventData<Stop>>;
+): WorkflowStream;
 export function until(
-  stream: ReadableStream<WorkflowEventData<any>> | WorkflowStream,
+  stream: WorkflowStream | ReadableStream<WorkflowEventData<any>>,
   cond:
     | ((event: WorkflowEventData<any>) => boolean | Promise<boolean>)
     | WorkflowEvent<any>,
-): ReadableStream<WorkflowEventData<any>> {
+): WorkflowStream<WorkflowEventData<any>> {
   let reader: ReadableStreamDefaultReader<WorkflowEventData<any>> | null = null;
-  return new ReadableStream<WorkflowEventData<any>>({
-    start: () => {
-      reader = stream.getReader();
-    },
-    pull: async (controller) => {
-      const { done, value } = await reader!.read();
-      if (value) {
-        controller.enqueue(value);
-      }
-      if (done) {
-        reader!.releaseLock();
-        reader = null;
-        controller.close();
-      } else {
-        if (isWorkflowEvent(cond) && cond.include(value)) {
-          reader!.releaseLock();
-          controller.close();
-        } else if (typeof cond === "function" && (await cond(value))) {
-          reader!.releaseLock();
-          controller.close();
+  return WorkflowStream.fromReadableStream(
+    new ReadableStream<WorkflowEventData<any>>({
+      start: () => {
+        reader = stream.getReader();
+      },
+      pull: async (controller) => {
+        const { done, value } = await reader!.read();
+        if (value) {
+          controller.enqueue(value);
         }
-      }
-    },
-  });
+        if (done) {
+          reader!.releaseLock();
+          reader = null;
+          controller.close();
+        } else {
+          if (isWorkflowEvent(cond) && cond.include(value)) {
+            reader!.releaseLock();
+            controller.close();
+          } else if (typeof cond === "function" && (await cond(value))) {
+            reader!.releaseLock();
+            controller.close();
+          }
+        }
+      },
+    }),
+  );
 }

packages/core/tests/integration/stream-library.spec.ts

@@ -57,7 +57,8 @@ describe("stream-chain", () => {
     sendEvent(startEvent.with());
     const outputs: WorkflowEventData<any>[] = [];
     const pipeline = chain([
-      until(stream, stopEvent),
+      // fixme: upstream should be treat it as a stream
+      until(stream, stopEvent)[Symbol.asyncIterator],
       new TransformStream({
         transform: (event: WorkflowEventData<any>, controller) => {
           if (messageEvent.include(event)) {