automatically cleanup input stream on handlers and tail SSE requests when a task execution cleansup

Author: ericallam

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

@@ -51,6 +51,10 @@ export class InputStreamsAPI implements InputStreamManager {
     return this.#getManager().lastSeqNum(streamId);
   }
 
+  public clearHandlers(): void {
+    this.#getManager().clearHandlers();
+  }
+
   public reset(): void {
     this.#getManager().reset();
   }

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

@@ -145,6 +145,19 @@ export class StandardInputStreamManager implements InputStreamManager {
     return undefined;
   }
 
+  clearHandlers(): void {
+    this.handlers.clear();
+
+    // Abort tails that no longer have any once waiters either
+    for (const [streamId, tail] of this.tails) {
+      const hasWaiters = this.onceWaiters.has(streamId) && this.onceWaiters.get(streamId)!.length > 0;
+      if (!hasWaiters) {
+        tail.abortController.abort();
+        this.tails.delete(streamId);
+      }
+    }
+  }
+
   connectTail(runId: string, _fromSeq?: number): void {
     // No-op: tails are now created per-stream lazily
   }

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

@@ -22,6 +22,7 @@ export class NoopInputStreamManager implements InputStreamManager {
     return undefined;
   }
 
+  clearHandlers(): void {}
   reset(): void {}
   disconnect(): void {}
   connectTail(_runId: string, _fromSeq?: number): void {}

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

@@ -47,6 +47,8 @@ export interface InputStreamManager {
 
   /**
    * Register a handler that fires every time data arrives on the given input stream.
+   * Handlers are automatically cleaned up when the task run completes.
+   * Returns `{ off }` for early unsubscription if needed.
    */
   on(streamId: string, handler: (data: unknown) => void | Promise<void>): { off: () => void };
 
@@ -68,6 +70,12 @@ export interface InputStreamManager {
    */
   lastSeqNum(streamId: string): number | undefined;
 
+  /**
+   * Clear all persistent `.on()` handlers and abort tails that have no remaining once waiters.
+   * Called automatically when a task run completes.
+   */
+  clearHandlers(): void;
+
   /**
    * Reset state between task executions.
    */

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

@@ -160,7 +160,8 @@ export type RealtimeDefinedInputStream<TData> = {
   id: string;
   /**
    * Register a handler that fires every time data arrives on this input stream.
-   * Returns a subscription object with an `.off()` method to unsubscribe.
+   * Handlers are automatically cleaned up when the task run completes, so calling
+   * `.off()` is optional. Returns a subscription with `.off()` for early unsubscription.
    */
   on: (handler: (data: TData) => void | Promise<void>) => InputStreamSubscription;
   /**

packages/core/src/v3/workers/taskExecutor.ts

@@ -13,6 +13,7 @@ import {
   accessoryAttributes,
   attemptKey,
   flattenAttributes,
+  inputStreams,
   lifecycleHooks,
   OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT,
   runMetadata,
@@ -1046,6 +1047,7 @@ export class TaskExecutor {
     signal: AbortSignal
   ) {
     await this.#callCleanupFunctions(payload, ctx, initOutput, signal);
+    inputStreams.clearHandlers();
     await this.#blockForWaitUntil();
   }
 

references/hello-world/src/trigger/inputStreams.ts

@@ -119,7 +119,7 @@ export const inputStreamOn = task({
 
     logger.info("Subscribing to messages via .on()", { expected });
 
-    const { off } = messageStream.on((data) => {
+    messageStream.on((data) => {
       logger.info("Received message", { data });
       received.push(data);
     });
@@ -128,7 +128,6 @@ export const inputStreamOn = task({
       await wait.for({ seconds: 1 });
     }
 
-    off();
     logger.info("Done receiving messages", { count: received.length });
     return { messages: received };
   },