fix(core): end runs.subscribeToRun after the run completes

`runs.subscribeToRun` now ends after the run completes. The `for await`
loop exits right after the update with the final status and the
subscription stops its connection, so a script that subscribes to a run
exits on its own. Pass `stopOnCompletion: false` to keep the
subscription open. `withStreams()` subscriptions keep their current
behaviour.

Mono-RevId: 04b515ab1c1e80b3cb328b5f80075d9086760c54

Author: ericallam

.changeset/subscribe-to-run-ends.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/core": patch
+---
+
+`runs.subscribeToRun` now ends after the run completes. The `for await` loop exits right after the update with the final status, and the subscription stops its connection, so a script that subscribes to a run exits on its own. Pass `stopOnCompletion: false` to keep the subscription open after completion.

packages/core/src/v3/apiClient/runStream-cancel.test.ts

@@ -139,6 +139,64 @@ describe("RunSubscription cancellation", () => {
     expect(pendingTimers()).toBeLessThanOrEqual(baseline);
   });
 
+  it("ends iteration and stops polling once the run completes", async () => {
+    runFields = {
+      status: "COMPLETED_SUCCESSFULLY",
+      completedAt: new Date().toISOString().slice(0, -1),
+    };
+    const client = new ApiClient(baseUrl, "tr_test");
+    const shared = new AbortController();
+    const baseline = pendingTimers();
+    const statuses: string[] = [];
+
+    const iterate = (async () => {
+      for await (const run of client.subscribeToRun("run_1", { signal: shared.signal })) {
+        statuses.push(run.status);
+      }
+      return "ended" as const;
+    })();
+
+    const outcome = await Promise.race([iterate, sleep(2_000).then(() => "timeout" as const)]);
+    teardown.abort();
+    if (outcome === "timeout") {
+      shared.abort();
+    }
+
+    expect(outcome).toBe("ended");
+    expect(statuses).toEqual(["COMPLETED"]);
+    const { settled, later } = await requestsAfterSettling();
+    expect(later).toBe(settled);
+    expect(pendingTimers()).toBeLessThanOrEqual(baseline);
+    expect(getEventListeners(shared.signal, "abort")).toHaveLength(0);
+  });
+
+  it("keeps iterating after the run completes when closeOnComplete is false", async () => {
+    runFields = {
+      status: "COMPLETED_SUCCESSFULLY",
+      completedAt: new Date().toISOString().slice(0, -1),
+    };
+    const client = new ApiClient(baseUrl, "tr_test");
+    const subscription = client.subscribeToRun("run_1", {
+      signal: teardown.signal,
+      closeOnComplete: false,
+    });
+
+    const iterate = (async () => {
+      for await (const _ of subscription) {
+      }
+      return "ended" as const;
+    })();
+
+    const outcome = await Promise.race([iterate, sleep(500).then(() => "open" as const)]);
+    const before = requests;
+    await sleep(200);
+
+    expect(outcome).toBe("open");
+    expect(requests).toBeGreaterThan(before);
+    subscription.unsubscribe();
+    await iterate;
+  });
+
   it("removes its listener from a shared signal when the consumer breaks", async () => {
     const client = new ApiClient(baseUrl, "tr_test");
     const shared = new AbortController();

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

@@ -153,7 +153,7 @@ export function runShapeStream<TRunTypes extends AnyRunTypes>(
 
   const $options: RunSubscriptionOptions = {
     runShapeStream: runStreamInstance.stream,
-    stopRunShapeStream: () => runStreamInstance.stop(30 * 1000),
+    stopRunShapeStream: (delay) => runStreamInstance.stop(delay),
     streamFactory: streamFactory,
     abortController,
     ...options,
@@ -760,7 +760,7 @@ export class SSEStreamSubscriptionFactory implements StreamSubscriptionFactory {
 
 export type RunSubscriptionOptions = RunShapeStreamOptions & {
   runShapeStream: ReadableStream<SubscribeRunRawShape>;
-  stopRunShapeStream: () => void;
+  stopRunShapeStream: (delay?: number) => void;
   streamFactory: StreamSubscriptionFactory;
   abortController: AbortController;
 };
@@ -773,6 +773,7 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
   >();
   private _closeOnComplete: boolean;
   private _isRunComplete = false;
+  private _withStreams = false;
 
   constructor(private options: RunSubscriptionOptions) {
     this._closeOnComplete =
@@ -794,12 +795,18 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
             this._isRunComplete &&
             !this.options.abortController.signal.aborted
           ) {
-            this.options.stopRunShapeStream();
+            if (this._withStreams) {
+              this.options.stopRunShapeStream(30 * 1000);
+            } else {
+              this.options.stopRunShapeStream();
+              controller.terminate();
+            }
           }
         },
       },
       this.options.abortController.signal,
       // A consumer `break` cancels the stream; unsubscribe so the shape stream stops polling
+      () => this.unsubscribe(),
       () => this.unsubscribe()
     );
   }
@@ -822,6 +829,8 @@ export class RunSubscription<TRunTypes extends AnyRunTypes> {
   withStreams<TStreams extends Record<string, any>>(): AsyncIterableStream<
     RunWithStreamsResult<RunShape<TRunTypes>, TStreams>
   > {
+    this._withStreams = true;
+
     // Keep track of which streams we've already subscribed to
     const activeStreams = new Set<string>();
 

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

@@ -23,7 +23,8 @@ export function createAsyncIterableReadable<S, T>(
   source: ReadableStream<S>,
   transformer: Transformer<S, T>,
   signal: AbortSignal,
-  onCancel?: () => void
+  onCancel?: () => void,
+  onClose?: () => void
 ): AsyncIterableStream<T> {
   // Set once the stream is closed, errored or cancelled; the controller must not be touched after
   let closed = false;
@@ -70,6 +71,7 @@ export function createAsyncIterableReadable<S, T>(
         }
       } finally {
         signal.removeEventListener("abort", onAbort);
+        onClose?.();
       }
     },
     cancel(reason) {