@workglow/task-graph 0.0.92 → 0.0.94

package/dist/bun.js

@@ -48,19 +48,26 @@ class Dataflow {
     if (!this.stream)
       return;
     const reader = this.stream.getReader();
-    let accumulatedText = "";
+    const accumulatedPorts = new Map;
     let lastSnapshotData = undefined;
     let finishData = undefined;
     let hasTextDelta = false;
+    let streamError;
     try {
       while (true) {
         const { done, value: event } = await reader.read();
         if (done)
           break;
         switch (event.type) {
-          case "text-delta":
+          case "text-delta": {
+            if (this.sourceTaskPortId !== DATAFLOW_ALL_PORTS && event.port !== this.sourceTaskPortId) {
+              break;
+            }
             hasTextDelta = true;
-            accumulatedText += event.textDelta;
+            accumulatedPorts.set(event.port, (accumulatedPorts.get(event.port) ?? "") + event.textDelta);
+            break;
+          }
+          case "object-delta":
             break;
           case "snapshot":
             lastSnapshotData = event.data;
@@ -69,6 +76,7 @@ class Dataflow {
             finishData = event.data;
             break;
           case "error":
+            streamError = event.error;
             break;
         }
       }
@@ -76,15 +84,25 @@ class Dataflow {
       reader.releaseLock();
       this.stream = undefined;
     }
+    if (streamError) {
+      this.error = streamError;
+      this.setStatus(TaskStatus.FAILED);
+      throw streamError;
+    }
     if (lastSnapshotData !== undefined) {
       this.setPortData(lastSnapshotData);
     } else if (finishData && Object.keys(finishData).length > 0) {
       this.setPortData(finishData);
     } else if (hasTextDelta) {
       if (this.sourceTaskPortId === DATAFLOW_ALL_PORTS) {
-        this.value = { text: accumulatedText };
+        const obj = {};
+        for (const [port, text] of accumulatedPorts) {
+          obj[port] = text;
+        }
+        this.value = obj;
       } else {
-        this.value = accumulatedText;
+        const text = accumulatedPorts.values().next().value ?? "";
+        this.value = text;
       }
     }
   }
@@ -466,29 +484,54 @@ function getPortStreamMode(schema, portId) {
     return xStream;
   return "none";
 }
-function getOutputStreamMode(outputSchema) {
-  if (typeof outputSchema === "boolean")
-    return "none";
-  const props = outputSchema.properties;
+function getStreamingPorts(schema) {
+  if (typeof schema === "boolean")
+    return [];
+  const props = schema.properties;
   if (!props)
-    return "none";
-  let found = "none";
-  for (const prop of Object.values(props)) {
+    return [];
+  const result = [];
+  for (const [name, prop] of Object.entries(props)) {
     if (!prop || typeof prop === "boolean")
       continue;
     const xStream = prop["x-stream"];
-    if (xStream === "append")
-      return "append";
-    if (xStream === "replace")
-      found = "replace";
+    if (xStream === "append" || xStream === "replace") {
+      result.push({ port: name, mode: xStream });
+    }
   }
-  return found;
+  return result;
+}
+function getOutputStreamMode(outputSchema) {
+  const ports = getStreamingPorts(outputSchema);
+  if (ports.length === 0)
+    return "none";
+  const mode = ports[0].mode;
+  for (let i = 1;i < ports.length; i++) {
+    if (ports[i].mode !== mode) {
+      throw new Error(`Mixed stream modes on a single task are not supported: ` + `port "${ports[0].port}" is "${mode}" but port "${ports[i].port}" is "${ports[i].mode}"`);
+    }
+  }
+  return mode;
 }
 function isTaskStreamable(task) {
   if (typeof task.executeStream !== "function")
     return false;
   return getOutputStreamMode(task.outputSchema()) !== "none";
 }
+function getAppendPortId(schema) {
+  if (typeof schema === "boolean")
+    return;
+  const props = schema.properties;
+  if (!props)
+    return;
+  for (const [name, prop] of Object.entries(props)) {
+    if (!prop || typeof prop === "boolean")
+      continue;
+    if (prop["x-stream"] === "append")
+      return name;
+  }
+  return;
+}
 function edgeNeedsAccumulation(sourceSchema, sourcePort, targetSchema, targetPort) {
   const sourceMode = getPortStreamMode(sourceSchema, sourcePort);
   if (sourceMode === "none")
@@ -505,6 +548,7 @@ class TaskRunner {
   abortController;
   outputCache;
   registry = globalServiceRegistry;
+  inputStreams;
   constructor(task) {
     this.task = task;
     this.own = this.own.bind(this);
@@ -607,7 +651,13 @@ class TaskRunner {
   }
   async executeStreamingTask(input) {
     const streamMode = getOutputStreamMode(this.task.outputSchema());
-    let accumulated = "";
+    if (streamMode === "append") {
+      const ports = getStreamingPorts(this.task.outputSchema());
+      if (ports.length === 0) {
+        throw new TaskError(`Task ${this.task.type} declares append streaming but no output port has x-stream: "append"`);
+      }
+    }
+    const accumulated = new Map;
     let chunkCount = 0;
     let finalOutput;
     this.task.emit("stream_start");
@@ -615,7 +665,8 @@ class TaskRunner {
       signal: this.abortController.signal,
       updateProgress: this.handleProgress.bind(this),
       own: this.own,
-      registry: this.registry
+      registry: this.registry,
+      inputStreams: this.inputStreams
     });
     for await (const event of stream) {
       chunkCount++;
@@ -629,20 +680,28 @@ class TaskRunner {
       this.task.emit("stream_chunk", event);
       switch (event.type) {
         case "text-delta": {
-          accumulated += event.textDelta;
-          const progress = Math.min(99, Math.round(100 * (1 - Math.exp(-0.02 * chunkCount))));
+          accumulated.set(event.port, (accumulated.get(event.port) ?? "") + event.textDelta);
+          const progress = Math.min(99, Math.round(100 * (1 - Math.exp(-0.05 * chunkCount))));
+          await this.handleProgress(progress);
+          break;
+        }
+        case "object-delta": {
+          const progress = Math.min(99, Math.round(100 * (1 - Math.exp(-0.05 * chunkCount))));
           await this.handleProgress(progress);
           break;
         }
         case "snapshot": {
-          const progress = Math.min(99, Math.round(100 * (1 - Math.exp(-0.02 * chunkCount))));
+          const progress = Math.min(99, Math.round(100 * (1 - Math.exp(-0.05 * chunkCount))));
           await this.handleProgress(progress);
           break;
         }
         case "finish": {
           if (streamMode === "append") {
-            const text = accumulated.length > 0 ? accumulated : event.data?.text ?? "";
-            finalOutput = { ...event.data || {}, text };
+            const merged = { ...event.data || {} };
+            for (const [port, text] of accumulated) {
+              merged[port] = text.length > 0 ? text : event.data?.[port] ?? "";
+            }
+            finalOutput = merged;
           } else if (streamMode === "replace") {
             finalOutput = event.data;
           }
@@ -1715,6 +1774,20 @@ class TaskGraphRunner {
   }
   async runTask(task, input) {
     const isStreamable = isTaskStreamable(task);
+    if (isStreamable) {
+      const dataflows = this.graph.getSourceDataflows(task.config.id);
+      const streamingEdges = dataflows.filter((df) => df.stream !== undefined);
+      if (streamingEdges.length > 0) {
+        const inputStreams = new Map;
+        for (const df of streamingEdges) {
+          const stream = df.stream;
+          const [forwardCopy, materializeCopy] = stream.tee();
+          inputStreams.set(df.targetTaskPortId, forwardCopy);
+          df.setStream(materializeCopy);
+        }
+        task.runner.inputStreams = inputStreams;
+      }
+    }
     await this.awaitStreamInputs(task);
     this.copyInputFromEdgesToNode(task);
     if (isStreamable) {
@@ -1782,14 +1855,17 @@ class TaskGraphRunner {
       task.off("stream_end", onStreamEnd);
     }
   }
-  pushStreamToEdges(task, streamMode) {
-    const targetDataflows = this.graph.getTargetDataflows(task.config.id);
-    if (targetDataflows.length === 0)
-      return;
-    const stream = new ReadableStream({
+  static isPortDelta(event) {
+    return event.type === "text-delta" || event.type === "object-delta";
+  }
+  createStreamFromTaskEvents(task, portId) {
+    return new ReadableStream({
       start: (controller) => {
         const onChunk = (event) => {
           try {
+            if (portId !== undefined && TaskGraphRunner.isPortDelta(event) && event.port !== portId) {
+              return;
+            }
             controller.enqueue(event);
           } catch {}
         };
@@ -1804,17 +1880,36 @@ class TaskGraphRunner {
         task.on("stream_end", onEnd);
       }
     });
-    if (targetDataflows.length === 1) {
-      targetDataflows[0].setStream(stream);
-    } else {
-      let currentStream = stream;
-      for (let i = 0;i < targetDataflows.length; i++) {
-        if (i === targetDataflows.length - 1) {
-          targetDataflows[i].setStream(currentStream);
-        } else {
-          const [s1, s2] = currentStream.tee();
-          targetDataflows[i].setStream(s1);
-          currentStream = s2;
+  }
+  pushStreamToEdges(task, streamMode) {
+    const targetDataflows = this.graph.getTargetDataflows(task.config.id);
+    if (targetDataflows.length === 0)
+      return;
+    const groups = new Map;
+    for (const df of targetDataflows) {
+      const key = df.sourceTaskPortId;
+      let group = groups.get(key);
+      if (!group) {
+        group = [];
+        groups.set(key, group);
+      }
+      group.push(df);
+    }
+    for (const [portKey, edges] of groups) {
+      const filterPort = portKey === DATAFLOW_ALL_PORTS ? undefined : portKey;
+      const stream = this.createStreamFromTaskEvents(task, filterPort);
+      if (edges.length === 1) {
+        edges[0].setStream(stream);
+      } else {
+        let currentStream = stream;
+        for (let i = 0;i < edges.length; i++) {
+          if (i === edges.length - 1) {
+            edges[i].setStream(currentStream);
+          } else {
+            const [s1, s2] = currentStream.tee();
+            edges[i].setStream(s1);
+            currentStream = s2;
+          }
         }
       }
     }
@@ -2145,6 +2240,69 @@ class GraphAsTask extends Task {
       });
     }
   }
+  async* executeStream(input, context) {
+    if (context.inputStreams) {
+      for (const [, stream] of context.inputStreams) {
+        const reader = stream.getReader();
+        try {
+          while (true) {
+            const { done, value } = await reader.read();
+            if (done)
+              break;
+            if (value.type === "finish")
+              continue;
+            yield value;
+          }
+        } finally {
+          reader.releaseLock();
+        }
+      }
+    }
+    if (this.hasChildren()) {
+      const endingNodeIds = new Set;
+      const tasks = this.subGraph.getTasks();
+      for (const task of tasks) {
+        if (this.subGraph.getTargetDataflows(task.config.id).length === 0) {
+          endingNodeIds.add(task.config.id);
+        }
+      }
+      const eventQueue = [];
+      let resolveWaiting;
+      let subgraphDone = false;
+      const unsub = this.subGraph.subscribeToTaskStreaming({
+        onStreamChunk: (taskId, event) => {
+          if (endingNodeIds.has(taskId) && event.type !== "finish") {
+            eventQueue.push(event);
+            resolveWaiting?.();
+          }
+        }
+      });
+      const runPromise = this.subGraph.run(input, { parentSignal: context.signal }).then((results2) => {
+        subgraphDone = true;
+        resolveWaiting?.();
+        return results2;
+      });
+      while (!subgraphDone) {
+        if (eventQueue.length === 0) {
+          await new Promise((resolve) => {
+            resolveWaiting = resolve;
+          });
+        }
+        while (eventQueue.length > 0) {
+          yield eventQueue.shift();
+        }
+      }
+      while (eventQueue.length > 0) {
+        yield eventQueue.shift();
+      }
+      unsub();
+      const results = await runPromise;
+      const mergedOutput = this.subGraph.mergeExecuteOutputsToRunOutput(results, this.compoundMerge);
+      yield { type: "finish", data: mergedOutput };
+    } else {
+      yield { type: "finish", data: input };
+    }
+  }
   regenerateGraph() {
     this._inputSchemaNode = undefined;
     this.events.emit("regenerate");
@@ -3437,6 +3595,9 @@ class IteratorTask extends GraphAsTask {
     }
     return this._runner;
   }
+  async* executeStream(input, _context) {
+    yield { type: "finish", data: input };
+  }
   set subGraph(subGraph) {
     super.subGraph = subGraph;
     this.invalidateIterationInputSchema();
@@ -3932,6 +4093,50 @@ class WhileTask extends GraphAsTask {
     }
     return currentOutput;
   }
+  async* executeStream(input, context) {
+    if (!this.hasChildren()) {
+      throw new TaskConfigurationError(`${this.type}: No subgraph set for while loop`);
+    }
+    const condition = this.condition ?? this.buildConditionFromExtras();
+    if (!condition) {
+      throw new TaskConfigurationError(`${this.type}: No condition function provided`);
+    }
+    const arrayAnalysis = this.analyzeArrayInputs(input);
+    this._currentIteration = 0;
+    let currentInput = { ...input };
+    let currentOutput = {};
+    const effectiveMax = arrayAnalysis ? Math.min(this.maxIterations, arrayAnalysis.iterationCount) : this.maxIterations;
+    while (this._currentIteration < effectiveMax) {
+      if (context.signal?.aborted)
+        break;
+      let iterationInput;
+      if (arrayAnalysis) {
+        iterationInput = {
+          ...this.buildIterationInput(currentInput, arrayAnalysis, this._currentIteration),
+          _iterationIndex: this._currentIteration
+        };
+      } else {
+        iterationInput = {
+          ...currentInput,
+          _iterationIndex: this._currentIteration
+        };
+      }
+      const results = await this.subGraph.run(iterationInput, {
+        parentSignal: context.signal
+      });
+      currentOutput = this.subGraph.mergeExecuteOutputsToRunOutput(results, this.compoundMerge);
+      if (!condition(currentOutput, this._currentIteration)) {
+        break;
+      }
+      if (this.chainIterations) {
+        currentInput = { ...currentInput, ...currentOutput };
+      }
+      this._currentIteration++;
+      const progress = Math.min(this._currentIteration / effectiveMax * 100, 99);
+      await context.updateProgress(progress, `Iteration ${this._currentIteration}`);
+    }
+    yield { type: "finish", data: currentOutput };
+  }
   getIterationContextSchema() {
     return this.constructor.getIterationContextSchema();
   }
@@ -4905,6 +5110,7 @@ export {
   hasVectorOutput,
   hasVectorLikeInput,
   getTaskQueueRegistry,
+  getStreamingPorts,
   getPortStreamMode,
   getOutputStreamMode,
   getNestedValue,
@@ -4912,6 +5118,7 @@ export {
   getJobQueueFactory,
   getIterationContextSchemaForType,
   getInputModeFromSchema,
+  getAppendPortId,
   findArrayPorts,
   filterIterationProperties,
   extractIterationProperties,
@@ -4981,4 +5188,4 @@ export {
   ConditionalTask
 };
 
-//# debugId=65C6DA6762951D4864756E2164756E21
+//# debugId=3CE7537068ABE3B064756E2164756E21