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

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

@@ -18,6 +18,7 @@ import * as Deferred from "effect/Deferred";
 import * as Fiber from "effect/Fiber";
 import * as Result from "effect/Result";
 import * as Stream from "effect/Stream";
+import * as TestClock from "effect/testing/TestClock";
 
 import * as PreviewAutomationBroker from "./PreviewAutomationBroker.ts";
 
@@ -925,6 +926,61 @@ it.effect("does not carry a tab id across a replacement automation stream", () =
   ),
 );
 
+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("fails requests assigned to the stream that is replaced", () =>
   Effect.scoped(
     Effect.gen(function* () {

apps/server/src/mcp/PreviewAutomationBroker.ts

@@ -29,6 +29,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 Crypto from "effect/Crypto";
 import * as Deferred from "effect/Deferred";
 import * as Effect from "effect/Effect";
@@ -42,6 +43,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;
@@ -93,7 +99,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"];
@@ -102,6 +111,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"];
@@ -120,6 +136,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;
@@ -129,22 +147,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,
   };
 };
@@ -351,6 +379,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,
@@ -382,11 +411,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),
       );
@@ -412,12 +444,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);
@@ -427,7 +475,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;
+  }
+}

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

@@ -15,6 +15,7 @@ import {
   type PreviewAutomationSnapshot,
   type PreviewAutomationStatus,
 } from "@t3tools/contracts";
+import * as Deferred from "effect/Deferred";
 import * as Effect from "effect/Effect";
 import * as Fiber from "effect/Fiber";
 import * as FileSystem from "effect/FileSystem";
@@ -93,6 +94,11 @@ function makeSession(frame?: ReturnType<typeof makeFrame>) {
       if (method === "Page.getNavigationHistory") return { currentIndex: 0, entries: [{}] };
       if (method === "Page.getLayoutMetrics") return { cssVisualViewport: { pageX: 0, pageY: 0 } };
       if (method === "Page.captureScreenshot") {
+        // A guest that is hidden mid-transition never paints the frame a still waits for.
+        if ((input as { format: string }).format === "jpeg" && stallStills) {
+          onStalledStill?.();
+          return new Promise<never>(() => {});
+        }
         // A snapshot sizes its screenshot from the PNG header; viewers get JPEG frames as-is.
         return { data: (input as { format: string }).format === "png" ? PNG_HEADER : "ZnJhbWU=" };
       }
@@ -268,6 +274,8 @@ let focusedFrame: ReturnType<typeof makeFrame> | null = null;
 let contextFailure: Error | null = null;
 let initialNavigation: (() => Promise<void>) | null = null;
 let desktopRendersNext = false;
+let stallStills = false;
+let onStalledStill: (() => void) | undefined;
 const desktopDetaches = new NodeEvents.EventEmitter();
 const desktopTabs = new Set<string>();
 const desktopRenders = (tabId: string) => {
@@ -380,6 +388,8 @@ beforeEach(() => {
   focusedFrame = null;
   desktopTabs.clear();
   desktopRendersNext = false;
+  stallStills = false;
+  onStalledStill = undefined;
   releasedDesktopTabs.length = 0;
   desktopConnections.length = 0;
 });
@@ -627,6 +637,44 @@ it.live("enforces provider ownership and explicit targets when a session has mul
   ).pipe(Effect.provide(layer)),
 );
 
+// The viewer still's deadline is a signal the test aborts; later deadlines never fire.
+const viewerStill = new AbortController();
+const neverAborts = new AbortController().signal;
+
+it.effect("a viewer still that never paints does not hold up the agent's snapshot", () =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const { browser, broker, tabId } = yield* ready;
+      const timeout = vi
+        .spyOn(AbortSignal, "timeout")
+        .mockImplementationOnce(() => viewerStill.signal)
+        .mockImplementation(() => neverAborts);
+      yield* Effect.addFinalizer(() => Effect.sync(() => timeout.mockRestore()));
+      const stillStarted = yield* Deferred.make<void>();
+      onStalledStill = () => Deferred.doneUnsafe(stillStarted, Effect.void);
+      stallStills = true;
+      const attaching = yield* browser
+        .attachViewer(viewerInput(tabId, false))
+        .pipe(Effect.forkScoped);
+      yield* Deferred.await(stillStarted);
+
+      const snapshotting = yield* broker
+        .invoke<PreviewAutomationSnapshot>({
+          scope,
+          operation: "snapshot",
+          input: {},
+          tabId,
+        })
+        .pipe(Effect.forkScoped);
+      viewerStill.abort(new DOMException("timed out", "TimeoutError"));
+      const snapshot = yield* Fiber.join(snapshotting);
+
+      expect(snapshot.accessibilityTree).toContain('button "delete"');
+      yield* Fiber.join(attaching);
+    }),
+  ).pipe(Effect.provide(layer)),
+);
+
 it.live("streams to a read-only viewer without allowing takeover, input, or viewport changes", () =>
   Effect.scoped(
     Effect.gen(function* () {