fix(server): a stalled preview capture no longer times out snapshots and drops the agent's tab

apps/server/src/mcp/PreviewAutomationBroker.test.ts

@@ -1073,6 +1073,61 @@ it.effect("accepts responses only from the host that received the request", () =
   }),
 );
 
+it.effect.each([
+  { reconnectAfter: "4 seconds", carried: true },
+  { reconnectAfter: "6 seconds", carried: false },
+] as const)(
+  "keeps the current tab when a timed-out host re-registers within the grace period (%o)",
+  ({ reconnectAfter, carried }) =>
+    Effect.scoped(
+      Effect.gen(function* () {
+        const broker = yield* makeBroker;
+        const openedTabId = PreviewTabId.make("tab-before-timeout");
+        const firstRequests = requestsFrom(yield* broker.connect(makeHost()));
+        yield* Stream.runForEach(firstRequests, (request) =>
+          // The host stalls on the snapshot, as when a capture never gets a frame.
+          request.operation === "snapshot"
+            ? Effect.void
+            : broker.respond({
+                clientId: "client-1",
+                connectionId: request.connectionId,
+                requestId: request.requestId,
+                ok: true,
+                result: { tabId: openedTabId },
+              }),
+        ).pipe(Effect.forkScoped);
+        yield* Effect.yieldNow;
+
+        yield* broker.invoke({ scope, operation: "open", input: {} });
+        const stalled = yield* broker
+          .invoke<void>({ scope, operation: "snapshot", input: {} })
+          .pipe(Effect.flip, Effect.forkScoped);
+        yield* Effect.yieldNow;
+        yield* TestClock.adjust("16 seconds");
+        expect((yield* Fiber.join(stalled))._tag).toBe("PreviewAutomationTimeoutError");
+
+        yield* TestClock.adjust(reconnectAfter);
+        const routedRequests: RoutedRequest[] = [];
+        const replacementRequests = requestsFrom(yield* broker.connect(makeHost()));
+        yield* Stream.runForEach(replacementRequests, (request) => {
+          routedRequests.push(request);
+          return broker.respond({
+            clientId: "client-1",
+            connectionId: request.connectionId,
+            requestId: request.requestId,
+            ok: true,
+            result: {},
+          });
+        }).pipe(Effect.forkScoped);
+        yield* Effect.yieldNow;
+
+        yield* broker.invoke({ scope, operation: "status", input: {} });
+        expect(routedRequests.at(-1)?.tabId).toBe(carried ? openedTabId : undefined);
+        expect(routedRequests.at(-1)?.tabIdExplicit).toBe(false);
+      }),
+    ),
+);
+
 it.effect("keeps each thread's current tab when threads share a provider session", () =>
   Effect.gen(function* () {
     const broker = yield* makeBroker;

apps/server/src/mcp/PreviewAutomationBroker.ts

@@ -30,6 +30,7 @@ import {
 } from "@t3tools/contracts";
 import * as Context from "effect/Context";
 import type * as Cause from "effect/Cause";
+import * as Clock from "effect/Clock";
 import * as SecureRandom from "effect/crypto/SecureRandom";
 import * as Deferred from "effect/Deferred";
 import * as Effect from "effect/Effect";
@@ -43,6 +44,11 @@ import * as SynchronizedRef from "effect/SynchronizedRef";
 import * as McpInvocationContext from "./McpInvocationContext.ts";
 
 const HOST_RESPONSE_GRACE_MS = 1_000;
+/**
+ * A host evicted for a slow response is still running and re-registers within
+ * about a second. Its tabs survive, so its agents keep their current tab.
+ */
+const EVICTED_LEASE_GRACE_MS = 5_000;
 
 export interface PreviewAutomationInvokeInput {
   readonly scope: McpInvocationContext.McpThreadInvocationScope;
@@ -94,7 +100,10 @@ interface PendingRequest {
   readonly context: PreviewAutomationRequestErrorContext;
 }
 
-/** A lease belongs to this exact connectionId/queue; replacing the host invalidates it. */
+/**
+ * A lease lives as long as the connection it names; a host evicted for a timeout
+ * hands its leases to its own re-registration (see `EVICTED_LEASE_GRACE_MS`).
+ */
 interface HostAssignment {
   readonly clientId: ClientConnection["clientId"];
   readonly connectionId: ClientConnection["connectionId"];
@@ -103,6 +112,13 @@ interface HostAssignment {
   readonly tabSequence?: number;
 }
 
+interface ReleasedLease {
+  readonly clientId: ClientConnection["clientId"];
+  readonly tabId?: PreviewTabId;
+  readonly tabSequence?: number;
+  readonly releasedAt: number;
+}
+
 interface PreviewAutomationRequestErrorContext {
   readonly operation: PreviewAutomationOperation;
   readonly environmentId: McpInvocationContext.McpInvocationScope["environmentId"];
@@ -121,6 +137,8 @@ interface PreviewAutomationRequestErrorContext {
 interface BrokerState {
   readonly clients: ReadonlyMap<string, ClientConnection>;
   readonly assignments: ReadonlyMap<string, HostAssignment>;
+  /** Leases of hosts evicted for a timeout, kept until they re-register or the grace ends. */
+  readonly released: ReadonlyMap<string, ReleasedLease>;
   readonly pending: ReadonlyMap<string, PendingRequest>;
   readonly requestSequence: number;
   readonly focusSequence: number;
@@ -130,22 +148,32 @@ const removeConnectionFromState = (
   current: BrokerState,
   clientId: string,
   queue: ClientConnection["queue"],
+  releasedAt?: number,
 ): { readonly state: BrokerState; readonly disconnected: ReadonlyArray<PendingRequest> } => {
   const clients = new Map(current.clients);
   const assignments = new Map(current.assignments);
+  const released = new Map(current.released);
   const pending = new Map(current.pending);
   const disconnected: PendingRequest[] = [];
   if (current.clients.get(clientId)?.queue === queue) clients.delete(clientId);
   for (const [assignmentKey, assignment] of assignments) {
-    if (assignment.queue === queue) assignments.delete(assignmentKey);
+    if (assignment.queue !== queue) continue;
+    assignments.delete(assignmentKey);
+    if (releasedAt === undefined) continue;
+    released.set(assignmentKey, {
+      clientId: assignment.clientId,
+      ...(assignment.tabId === undefined ? {} : { tabId: assignment.tabId }),
+      ...(assignment.tabSequence === undefined ? {} : { tabSequence: assignment.tabSequence }),
+      releasedAt,
+    });
   }
   for (const [requestId, entry] of pending) {
     if (entry.queue !== queue) continue;
     pending.delete(requestId);
     disconnected.push(entry);
   }
   return {
-    state: { ...current, clients, assignments, pending },
+    state: { ...current, clients, assignments, released, pending },
     disconnected,
   };
 };
@@ -367,6 +395,7 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
   const state = yield* SynchronizedRef.make<BrokerState>({
     clients: new Map(),
     assignments: new Map(),
+    released: new Map(),
     pending: new Map(),
     requestSequence: 0,
     focusSequence: 0,
@@ -398,11 +427,14 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
     queue: ClientConnection["queue"],
     completeStream = false,
   ) {
+    // Only a timeout keeps the host's leases: a closed or replaced stream may
+    // belong to a runtime whose tabs are gone.
+    const releasedAt = completeStream ? yield* Clock.currentTimeMillis : undefined;
     yield* SynchronizedRef.modifyEffect(state, (current) => {
       if (current.clients.get(clientId)?.queue !== queue) {
         return Effect.succeed([undefined, current] as const);
       }
-      const removed = removeConnectionFromState(current, clientId, queue);
+      const removed = removeConnectionFromState(current, clientId, queue, releasedAt);
       return closeConnection(queue, removed.disconnected, completeStream).pipe(
         Effect.as([undefined, removed.state] as const),
       );
@@ -428,12 +460,28 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
       focusOrder: 0,
       queue,
     };
+    const now = yield* Clock.currentTimeMillis;
     const registration = yield* SynchronizedRef.modify(state, (current) => {
       const previousConnection = current.clients.get(clientId);
       const removed = previousConnection
         ? removeConnectionFromState(current, clientId, previousConnection.queue)
         : { state: current, disconnected: [] };
       const clients = new Map(removed.state.clients);
+      const assignments = new Map(removed.state.assignments);
+      const released = new Map(removed.state.released);
+      for (const [assignmentKey, lease] of released) {
+        const expired = now - lease.releasedAt > EVICTED_LEASE_GRACE_MS;
+        if (!expired && lease.clientId !== clientId) continue;
+        released.delete(assignmentKey);
+        if (expired || assignments.has(assignmentKey)) continue;
+        assignments.set(assignmentKey, {
+          clientId,
+          connectionId,
+          queue,
+          ...(lease.tabId === undefined ? {} : { tabId: lease.tabId }),
+          ...(lease.tabSequence === undefined ? {} : { tabSequence: lease.tabSequence }),
+        });
+      }
       const focusSequence = removed.state.focusSequence + 1;
       const registeredConnection = { ...connection, focusOrder: focusSequence };
       clients.set(clientId, registeredConnection);
@@ -443,7 +491,7 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
           disconnected: removed.disconnected,
           registeredConnection,
         },
-        { ...removed.state, clients, focusSequence },
+        { ...removed.state, clients, assignments, released, focusSequence },
       ] as const;
     });
     if (registration.previousConnection) {

apps/server/src/preview/CaptureLock.test.ts

@@ -0,0 +1,48 @@
+import { describe, expect, it } from "vite-plus/test";
+
+import { CaptureLock } from "./CaptureLock.ts";
+
+// oxlint-disable-next-line t3code/no-real-time-wait-in-tests -- a zero-delay macrotask boundary, not a wait for time
+const settle = () => new Promise<void>((resolve) => setImmediate(resolve));
+
+describe("CaptureLock", () => {
+  it("keeps the capture ahead exclusive when a waiting capture gives up", async () => {
+    const lock = new CaptureLock();
+    const ahead = Promise.withResolvers<void>();
+    const events: string[] = [];
+    const first = lock.run(async () => {
+      events.push("first started");
+      await ahead.promise;
+      events.push("first finished");
+    });
+    const giveUp = new AbortController();
+    const expired = lock.run(async () => {
+      events.push("expired started");
+    }, giveUp.signal);
+    const reason = new Error("expired");
+    giveUp.abort(reason);
+    await expect(expired).rejects.toBe(reason);
+
+    const next = lock.run(async () => {
+      events.push("next started");
+    });
+    await settle();
+    expect(events).toEqual(["first started"]);
+
+    ahead.resolve();
+    await Promise.all([first, next]);
+    expect(events).toEqual(["first started", "first finished", "next started"]);
+  });
+
+  it("frees the lock when a running capture that never settles gives up", async () => {
+    const lock = new CaptureLock();
+    const giveUp = new AbortController();
+    const stalled = lock.run(() => new Promise<never>(() => {}), giveUp.signal);
+    await settle();
+    const reason = new Error("expired");
+    giveUp.abort(reason);
+
+    await expect(stalled).rejects.toBe(reason);
+    await expect(lock.run(async () => "next")).resolves.toBe("next");
+  });
+});

apps/server/src/preview/CaptureLock.ts

@@ -0,0 +1,30 @@
+import { constVoid } from "effect/Function";
+
+const abandonOnAbort = <A>(promise: Promise<A>, signal: AbortSignal): Promise<A> =>
+  new Promise<A>((resolve, reject) => {
+    const abort = () => reject(signal.reason);
+    if (signal.aborted) abort();
+    signal.addEventListener("abort", abort, { once: true });
+    promise.then(resolve, reject).finally(() => signal.removeEventListener("abort", abort));
+  });
+
+/**
+ * Runs a tab's captures one at a time. A capture given a signal gives up when
+ * it aborts, waiting included, so one that never gets a frame cannot stall
+ * every capture behind it. One that gives up while still waiting keeps its
+ * place, so the capture ahead of it never overlaps the next.
+ */
+export class CaptureLock {
+  private tail: Promise<void> = Promise.resolve();
+
+  run<A>(operation: () => Promise<A>, signal?: AbortSignal): Promise<A> {
+    const previous = this.tail;
+    const run = previous.then(() => {
+      signal?.throwIfAborted();
+      return operation();
+    });
+    const result = signal ? abandonOnAbort(run, signal) : run;
+    this.tail = Promise.all([previous, result.then(constVoid, constVoid)]).then(constVoid);
+    return result;
+  }
+}