Adding streaming support

Author: ericallam

packages/core/src/v3/apiClient/index.ts

@@ -61,14 +61,13 @@ import {
   SubscribeToRunsQueryParams,
   UpdateEnvironmentVariableParams,
 } from "./types.js";
-import type { AsyncIterableStream } from "./stream.js";
+import { AsyncIterableStream } from "../streams/asyncIterableStream.js";
 
 export type {
   CreateEnvironmentVariableParams,
   ImportEnvironmentVariablesParams,
   SubscribeToRunsQueryParams,
   UpdateEnvironmentVariableParams,
-  AsyncIterableStream,
 };
 
 export type ClientTriggerOptions = {

packages/core/src/v3/apiClient/runStream.ts

@@ -16,12 +16,11 @@ import {
 } from "../utils/ioSerialization.js";
 import { ApiError } from "./errors.js";
 import { ApiClient } from "./index.js";
+import { LineTransformStream, zodShapeStream } from "./stream.js";
 import {
   AsyncIterableStream,
   createAsyncIterableReadable,
-  LineTransformStream,
-  zodShapeStream,
-} from "./stream.js";
+} from "../streams/asyncIterableStream.js";
 
 export type RunShape<TRunTypes extends AnyRunTypes> = TRunTypes extends AnyRunTypes
   ? {

packages/core/src/v3/apiClient/stream.ts

@@ -9,6 +9,7 @@ import {
   type Row,
   type ShapeStreamInterface,
 } from "@electric-sql/client";
+import { AsyncIterableStream, createAsyncIterableStream } from "../streams/asyncIterableStream.js";
 
 export type ZodShapeStreamOptions = {
   headers?: Record<string, string>;
@@ -82,57 +83,6 @@ export function zodShapeStream<TShapeSchema extends z.ZodTypeAny>(
   };
 }
 
-export type AsyncIterableStream<T> = AsyncIterable<T> & ReadableStream<T>;
-
-export function createAsyncIterableStream<S, T>(
-  source: ReadableStream<S>,
-  transformer: Transformer<S, T>
-): AsyncIterableStream<T> {
-  const transformedStream: any = source.pipeThrough(new TransformStream(transformer));
-
-  transformedStream[Symbol.asyncIterator] = () => {
-    const reader = transformedStream.getReader();
-    return {
-      async next(): Promise<IteratorResult<string>> {
-        const { done, value } = await reader.read();
-        return done ? { done: true, value: undefined } : { done: false, value };
-      },
-    };
-  };
-
-  return transformedStream;
-}
-
-export function createAsyncIterableReadable<S, T>(
-  source: ReadableStream<S>,
-  transformer: Transformer<S, T>,
-  signal: AbortSignal
-): AsyncIterableStream<T> {
-  return new ReadableStream<T>({
-    async start(controller) {
-      const transformedStream = source.pipeThrough(new TransformStream(transformer));
-      const reader = transformedStream.getReader();
-
-      signal.addEventListener("abort", () => {
-        queueMicrotask(() => {
-          reader.cancel();
-          controller.close();
-        });
-      });
-
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) {
-          controller.close();
-          break;
-        }
-
-        controller.enqueue(value);
-      }
-    },
-  }) as AsyncIterableStream<T>;
-}
-
 class ReadableShapeStream<T extends Row<unknown> = Row> {
   readonly #stream: ShapeStreamInterface<T>;
   readonly #currentState: Map<string, T> = new Map();

packages/core/src/v3/index.ts

@@ -21,6 +21,7 @@ export * from "./types/index.js";
 export { links } from "./links.js";
 export * from "./jwt.js";
 export * from "./idempotencyKeys.js";
+export * from "./streams/asyncIterableStream.js";
 export {
   formatDuration,
   formatDurationInDays,

packages/core/src/v3/logger/index.ts

@@ -51,6 +51,10 @@ export class LoggerAPI implements TaskLogger {
     return this.#getTaskLogger().trace(name, fn, options);
   }
 
+  public startSpan(name: string, options?: SpanOptions): Span {
+    return this.#getTaskLogger().startSpan(name, options);
+  }
+
   #getTaskLogger(): TaskLogger {
     return getGlobal(API_NAME) ?? NOOP_TASK_LOGGER;
   }

packages/core/src/v3/logger/taskLogger.ts

@@ -24,6 +24,7 @@ export interface TaskLogger {
   warn(message: string, properties?: Record<string, unknown>): void;
   error(message: string, properties?: Record<string, unknown>): void;
   trace<T>(name: string, fn: (span: Span) => Promise<T>, options?: SpanOptions): Promise<T>;
+  startSpan(name: string, options?: SpanOptions): Span;
 }
 
 export class OtelTaskLogger implements TaskLogger {
@@ -90,6 +91,10 @@ export class OtelTaskLogger implements TaskLogger {
     return this._config.tracer.startActiveSpan(name, fn, options);
   }
 
+  startSpan(name: string, options?: SpanOptions): Span {
+    return this._config.tracer.startSpan(name, options);
+  }
+
   #getTimestampInHrTime(): ClockTime {
     return clock.preciseNow();
   }
@@ -104,6 +109,9 @@ export class NoopTaskLogger implements TaskLogger {
   trace<T>(name: string, fn: (span: Span) => Promise<T>): Promise<T> {
     return fn({} as Span);
   }
+  startSpan(): Span {
+    return {} as Span;
+  }
 }
 
 function safeJsonProcess(value?: Record<string, unknown>): Record<string, unknown> | undefined {

packages/core/src/v3/runMetadata/index.ts

@@ -1,5 +1,5 @@
 import { DeserializedJson } from "../../schemas/json.js";
-import { AsyncIterableStream } from "../apiClient/stream.js";
+import { AsyncIterableStream } from "../streams/asyncIterableStream.js";
 import { getGlobal, registerGlobal } from "../utils/globals.js";
 import { ApiRequestOptions } from "../zodfetch.js";
 import { NoopRunMetadataManager } from "./noopManager.js";

packages/core/src/v3/runMetadata/manager.ts

@@ -1,12 +1,12 @@
 import { dequal } from "dequal/lite";
 import { DeserializedJson } from "../../schemas/json.js";
 import { ApiClient } from "../apiClient/index.js";
-import { AsyncIterableStream } from "../apiClient/stream.js";
 import { FlushedRunMetadata, RunMetadataChangeOperation } from "../schemas/common.js";
 import { ApiRequestOptions } from "../zodfetch.js";
 import { MetadataStream } from "./metadataStream.js";
 import { applyMetadataOperations } from "./operations.js";
 import { RunMetadataManager, RunMetadataUpdater } from "./types.js";
+import { AsyncIterableStream } from "../streams/asyncIterableStream.js";
 
 const MAXIMUM_ACTIVE_STREAMS = 5;
 const MAXIMUM_TOTAL_STREAMS = 10;

packages/core/src/v3/runMetadata/noopManager.ts

@@ -1,5 +1,5 @@
 import { DeserializedJson } from "../../schemas/json.js";
-import { AsyncIterableStream } from "../apiClient/stream.js";
+import { AsyncIterableStream } from "../streams/asyncIterableStream.js";
 import { ApiRequestOptions } from "../zodfetch.js";
 import type { RunMetadataManager, RunMetadataUpdater } from "./types.js";
 

packages/core/src/v3/runMetadata/types.ts

@@ -1,5 +1,5 @@
 import { DeserializedJson } from "../../schemas/json.js";
-import { AsyncIterableStream } from "../apiClient/stream.js";
+import { AsyncIterableStream } from "../streams/asyncIterableStream.js";
 import { ApiRequestOptions } from "../zodfetch.js";
 
 export interface RunMetadataUpdater {

packages/core/src/v3/streams/asyncIterableStream.ts

@@ -0,0 +1,97 @@
+export type AsyncIterableStream<T> = AsyncIterable<T> & ReadableStream<T>;
+
+export function createAsyncIterableStream<S, T>(
+  source: ReadableStream<S>,
+  transformer: Transformer<S, T>
+): AsyncIterableStream<T> {
+  const transformedStream: any = source.pipeThrough(new TransformStream(transformer));
+
+  transformedStream[Symbol.asyncIterator] = () => {
+    const reader = transformedStream.getReader();
+    return {
+      async next(): Promise<IteratorResult<string>> {
+        const { done, value } = await reader.read();
+        return done ? { done: true, value: undefined } : { done: false, value };
+      },
+    };
+  };
+
+  return transformedStream;
+}
+
+export function createAsyncIterableReadable<S, T>(
+  source: ReadableStream<S>,
+  transformer: Transformer<S, T>,
+  signal: AbortSignal
+): AsyncIterableStream<T> {
+  return new ReadableStream<T>({
+    async start(controller) {
+      const transformedStream = source.pipeThrough(new TransformStream(transformer));
+      const reader = transformedStream.getReader();
+
+      signal.addEventListener("abort", () => {
+        queueMicrotask(() => {
+          reader.cancel();
+          controller.close();
+        });
+      });
+
+      while (true) {
+        const { done, value } = await reader.read();
+        if (done) {
+          controller.close();
+          break;
+        }
+
+        controller.enqueue(value);
+      }
+    },
+  }) as AsyncIterableStream<T>;
+}
+
+export function createAsyncIterableStreamFromAsyncIterable<T>(
+  asyncIterable: AsyncIterable<T>,
+  transformer: Transformer<T, T>,
+  signal?: AbortSignal
+): AsyncIterableStream<T> {
+  const stream = new ReadableStream<T>({
+    async start(controller) {
+      try {
+        if (signal) {
+          signal.addEventListener("abort", () => {
+            controller.close();
+          });
+        }
+
+        const iterator = asyncIterable[Symbol.asyncIterator]();
+
+        while (true) {
+          if (signal?.aborted) {
+            break;
+          }
+
+          const { done, value } = await iterator.next();
+
+          if (done) {
+            controller.close();
+            break;
+          }
+
+          controller.enqueue(value);
+        }
+      } catch (error) {
+        controller.error(error);
+      }
+    },
+    cancel() {
+      // If the stream is a tinyexec process with a kill method, kill it
+      if ("kill" in asyncIterable) {
+        (asyncIterable as any).kill();
+      }
+    },
+  });
+
+  const transformedStream = stream.pipeThrough(new TransformStream(transformer));
+
+  return transformedStream as AsyncIterableStream<T>;
+}

packages/python/README.md

@@ -76,6 +76,22 @@ export const myScript = task({
     return result.stdout;
   },
 });
+
+export const myStreamingScript = task({
+  id: "my-streaming-python-script",
+  run: async () => {
+    // You can also stream the output of the script
+    const result = python.stream.runScript("my_script.py", ["hello", "world"]);
+
+    // result is an async iterable/readable stream
+    for await (const chunk of streamingResult) {
+      logger.debug("convert-url-to-markdown", {
+        url: payload.url,
+        chunk,
+      });
+    }
+  },
+});
 ```
 
 ### Running Inline Python Code