fix(core): tighten stream-manager lifecycles

Author: ericallam

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

@@ -14,6 +14,13 @@ type OnceWaiter = {
   resolve: (result: InputStreamOnceResult<unknown>) => void;
   reject: (error: Error) => void;
   timeoutHandle?: ReturnType<typeof setTimeout>;
+  // The abort signal and its handler are tracked on the waiter so any
+  // resolution path (dispatch / timeout / explicit removal) can detach
+  // the listener. Without this, a long-lived `AbortSignal` reused across
+  // many `once()` calls accumulates listeners — `{ once: true }` only
+  // self-clears if the signal actually aborts.
+  signal?: AbortSignal;
+  abortHandler?: () => void;
 };
 
 
@@ -36,6 +43,13 @@ export class StandardInputStreamManager implements InputStreamManager {
   // reconnect in a tight loop. Reset to 0 by `#dispatch` whenever a
   // record flows through.
   private reconnectAttempts = new Map<string, number>();
+  // Stream IDs that were explicitly torn down by `disconnectStream`. The
+  // tail's `.finally` reconnect path consults this set so a deliberate
+  // teardown isn't immediately undone by the auto-reconnect when
+  // handlers or once-waiters are still registered. Cleared on the next
+  // explicit `on()` / `once()` (those are the only legitimate reasons to
+  // bring the tail back up).
+  private explicitlyDisconnected = new Set<string>();
 
   constructor(
     private apiClient: ApiClient,
@@ -75,6 +89,10 @@ export class StandardInputStreamManager implements InputStreamManager {
   on(streamId: string, handler: InputStreamHandler): { off: () => void } {
     this.#requireV2Streams();
 
+    // A fresh attach is a legitimate reason to bring the tail back up;
+    // clear any prior explicit-disconnect flag.
+    this.explicitlyDisconnected.delete(streamId);
+
     let handlerSet = this.handlers.get(streamId);
     if (!handlerSet) {
       handlerSet = new Set();
@@ -107,6 +125,10 @@ export class StandardInputStreamManager implements InputStreamManager {
   once(streamId: string, options?: InputStreamOnceOptions): InputStreamOncePromise<unknown> {
     this.#requireV2Streams();
 
+    // A fresh waiter is a legitimate reason to bring the tail back up;
+    // clear any prior explicit-disconnect flag.
+    this.explicitlyDisconnected.delete(streamId);
+
     // Lazily connect a tail for this stream
     this.#ensureStreamTailConnected(streamId);
 
@@ -131,17 +153,16 @@ export class StandardInputStreamManager implements InputStreamManager {
           reject(new Error("Aborted"));
           return;
         }
-        options.signal.addEventListener(
-          "abort",
-          () => {
-            if (waiter.timeoutHandle) {
-              clearTimeout(waiter.timeoutHandle);
-            }
-            this.#removeOnceWaiter(streamId, waiter);
-            reject(new Error("Aborted"));
-          },
-          { once: true }
-        );
+        const abortHandler = () => {
+          if (waiter.timeoutHandle) {
+            clearTimeout(waiter.timeoutHandle);
+          }
+          this.#removeOnceWaiter(streamId, waiter);
+          reject(new Error("Aborted"));
+        };
+        waiter.signal = options.signal;
+        waiter.abortHandler = abortHandler;
+        options.signal.addEventListener("abort", abortHandler, { once: true });
       }
 
       // Handle timeout — resolve with error result instead of rejecting
@@ -186,6 +207,14 @@ export class StandardInputStreamManager implements InputStreamManager {
   }
 
   disconnectStream(streamId: string): void {
+    // Mark as explicitly disconnected BEFORE we abort, so the tail's
+    // `.finally` reconnect path sees the flag when it runs (which can be
+    // synchronous in the AbortError catch). Without this, an in-flight
+    // `.on(...)` or pending `.once()` would immediately resurrect the
+    // tail and negate the disconnect — defeating the
+    // "drop-the-duplicate before .wait() suspends" contract. Cleared on
+    // the next explicit `on()` / `once()`.
+    this.explicitlyDisconnected.add(streamId);
     const tail = this.tails.get(streamId);
     if (tail) {
       tail.abortController.abort();
@@ -225,13 +254,17 @@ export class StandardInputStreamManager implements InputStreamManager {
     this.seqNums.clear();
     this.handlers.clear();
     this.reconnectAttempts.clear();
+    this.explicitlyDisconnected.clear();
 
     // Reject all pending once waiters
     for (const [, waiters] of this.onceWaiters) {
       for (const waiter of waiters) {
         if (waiter.timeoutHandle) {
           clearTimeout(waiter.timeoutHandle);
         }
+        if (waiter.signal && waiter.abortHandler) {
+          waiter.signal.removeEventListener("abort", waiter.abortHandler);
+        }
         waiter.reject(new Error("Input stream manager reset"));
       }
     }
@@ -259,6 +292,13 @@ export class StandardInputStreamManager implements InputStreamManager {
         .finally(() => {
           this.tails.delete(streamId);
 
+          // If the tail was torn down explicitly via `disconnectStream`,
+          // don't auto-reconnect — that's the whole point of the
+          // disconnect call. The next `on()` / `once()` clears the flag.
+          if (this.explicitlyDisconnected.has(streamId)) {
+            return;
+          }
+
           // Auto-reconnect with exponential backoff if there are still
           // active handlers or waiters. Without backoff a persistent
           // failure (auth rejected, 5xx, DNS) would reconnect in a tight
@@ -273,6 +313,7 @@ export class StandardInputStreamManager implements InputStreamManager {
             this.reconnectAttempts.set(streamId, attempt + 1);
             const delayMs = computeReconnectDelayMs(attempt);
             setTimeout(() => {
+              if (this.explicitlyDisconnected.has(streamId)) return;
               if (this.tails.has(streamId)) return;
               const stillHasHandlers =
                 this.handlers.has(streamId) && this.handlers.get(streamId)!.size > 0;
@@ -361,6 +402,9 @@ export class StandardInputStreamManager implements InputStreamManager {
       if (waiter.timeoutHandle) {
         clearTimeout(waiter.timeoutHandle);
       }
+      if (waiter.signal && waiter.abortHandler) {
+        waiter.signal.removeEventListener("abort", waiter.abortHandler);
+      }
       waiter.resolve({ ok: true, output: data });
       // Also invoke persistent handlers
       this.#invokeHandlers(streamId, data);
@@ -410,6 +454,13 @@ export class StandardInputStreamManager implements InputStreamManager {
   }
 
   #removeOnceWaiter(streamId: string, waiter: OnceWaiter): void {
+    // Centralized cleanup — both timeout and explicit abort paths funnel
+    // through here, so detach the abort listener once instead of at every
+    // callsite. The dispatch path doesn't go through this method (the
+    // waiter is shifted off inline), so it detaches the listener there.
+    if (waiter.signal && waiter.abortHandler) {
+      waiter.signal.removeEventListener("abort", waiter.abortHandler);
+    }
     const waiters = this.onceWaiters.get(streamId);
     if (!waiters) return;
     const index = waiters.indexOf(waiter);

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

@@ -14,6 +14,13 @@ type OnceWaiter = {
   resolve: (result: InputStreamOnceResult<unknown>) => void;
   reject: (error: Error) => void;
   timeoutHandle?: ReturnType<typeof setTimeout>;
+  // The abort signal and its handler are tracked on the waiter so any
+  // resolution path (dispatch / timeout / explicit removal) can detach
+  // the listener. Without this, a long-lived `AbortSignal` reused across
+  // many `once()` calls accumulates listeners — `{ once: true }` only
+  // self-clears if the signal actually aborts.
+  signal?: AbortSignal;
+  abortHandler?: () => void;
 };
 
 type TailState = {
@@ -131,15 +138,14 @@ export class StandardSessionStreamManager implements SessionStreamManager {
           reject(new Error("Aborted"));
           return;
         }
-        options.signal.addEventListener(
-          "abort",
-          () => {
-            if (waiter.timeoutHandle) clearTimeout(waiter.timeoutHandle);
-            this.#removeOnceWaiter(key, waiter);
-            reject(new Error("Aborted"));
-          },
-          { once: true }
-        );
+        const abortHandler = () => {
+          if (waiter.timeoutHandle) clearTimeout(waiter.timeoutHandle);
+          this.#removeOnceWaiter(key, waiter);
+          reject(new Error("Aborted"));
+        };
+        waiter.signal = options.signal;
+        waiter.abortHandler = abortHandler;
+        options.signal.addEventListener("abort", abortHandler, { once: true });
       }
 
       if (options?.timeoutMs) {
@@ -261,6 +267,9 @@ export class StandardSessionStreamManager implements SessionStreamManager {
     for (const [, waiters] of this.onceWaiters) {
       for (const waiter of waiters) {
         if (waiter.timeoutHandle) clearTimeout(waiter.timeoutHandle);
+        if (waiter.signal && waiter.abortHandler) {
+          waiter.signal.removeEventListener("abort", waiter.abortHandler);
+        }
         waiter.reject(new Error("Session stream manager reset"));
       }
     }
@@ -403,6 +412,9 @@ export class StandardSessionStreamManager implements SessionStreamManager {
       const waiter = waiters.shift()!;
       if (waiters.length === 0) this.onceWaiters.delete(key);
       if (waiter.timeoutHandle) clearTimeout(waiter.timeoutHandle);
+      if (waiter.signal && waiter.abortHandler) {
+        waiter.signal.removeEventListener("abort", waiter.abortHandler);
+      }
       waiter.resolve({ ok: true, output: data });
       this.#invokeHandlers(key, data);
       return;
@@ -450,6 +462,13 @@ export class StandardSessionStreamManager implements SessionStreamManager {
   }
 
   #removeOnceWaiter(key: string, waiter: OnceWaiter): void {
+    // Centralized cleanup — both timeout and explicit abort paths funnel
+    // through here, so detach the abort listener once instead of at every
+    // callsite. The dispatch path doesn't go through this method (the
+    // waiter is shifted off inline), so it detaches the listener there.
+    if (waiter.signal && waiter.abortHandler) {
+      waiter.signal.removeEventListener("abort", waiter.abortHandler);
+    }
     const waiters = this.onceWaiters.get(key);
     if (!waiters) return;
     const index = waiters.indexOf(waiter);

packages/core/src/v3/test/test-realtime-streams-manager.ts

@@ -158,6 +158,7 @@ export class TestRealtimeStreamsManager implements RealtimeStreamsManager {
   reset(): void {
     this.buffers.clear();
     this.pipeWaits.clear();
+    this.writeListeners.clear();
   }
 
   private getBuffer(key: string): unknown[] {

packages/core/src/v3/test/test-session-stream-manager.ts

@@ -221,8 +221,18 @@ export class TestSessionStreamManager implements SessionStreamManager {
 
     const handlers = this.handlers.get(key);
     if (handlers && handlers.size > 0) {
+      // Awaited so test code can rely on handlers having completed by the
+      // time `__sendFromTest` resolves. Wrapped per-handler so a
+      // throwing/rejecting handler doesn't poison Promise.all and break
+      // unrelated test state.
       await Promise.all(
-        Array.from(handlers).map((h) => Promise.resolve().then(() => h(data)))
+        Array.from(handlers).map(async (h) => {
+          try {
+            await h(data);
+          } catch {
+            // Never let a handler error break test state
+          }
+        })
       );
     }
 
@@ -267,17 +277,6 @@ export class TestSessionStreamManager implements SessionStreamManager {
     }
   }
 
-  private invoke(handler: Handler, data: unknown): void {
-    try {
-      const result = handler(data);
-      if (result && typeof result === "object" && "catch" in result) {
-        (result as Promise<void>).catch(() => {});
-      }
-    } catch {
-      // Never let a handler error break test state
-    }
-  }
-
   private removeWaiter(key: string, waiter: OnceWaiter): void {
     const waiters = this.onceWaiters.get(key);
     if (!waiters) return;