From 1115032e62e08b22e1804e7fea5a388d36918416 Mon Sep 17 00:00:00 2001 From: Mikiyas Date: Tue, 27 Oct 2026 12:44:03 +0300 Subject: [PATCH] fix(server): a stalled preview capture no longer times out snapshots and drops the agent's tab --- .../src/mcp/PreviewAutomationBroker.test.ts | 56 +++++++++ .../server/src/mcp/PreviewAutomationBroker.ts | 58 ++++++++- apps/server/src/preview/CaptureLock.test.ts | 48 ++++++++ apps/server/src/preview/CaptureLock.ts | 30 +++++ apps/server/src/preview/ServerBrowser.test.ts | 48 ++++++++ apps/server/src/preview/ServerBrowser.ts | 116 ++++++++++++------ apps/server/src/preview/ServerBrowserPage.ts | 2 +- 7 files changed, 312 insertions(+), 46 deletions(-) create mode 100644 apps/server/src/preview/CaptureLock.test.ts create mode 100644 apps/server/src/preview/CaptureLock.ts diff --git a/apps/server/src/mcp/PreviewAutomationBroker.test.ts b/apps/server/src/mcp/PreviewAutomationBroker.test.ts index f9d618f0520e3..d7393157aba52 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.test.ts +++ b/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({ 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* () { diff --git a/apps/server/src/mcp/PreviewAutomationBroker.ts b/apps/server/src/mcp/PreviewAutomationBroker.ts index 53622f3054775..6399406b21a4b 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.ts +++ b/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; readonly assignments: ReadonlyMap; + /** Leases of hosts evicted for a timeout, kept until they re-register or the grace ends. */ + readonly released: ReadonlyMap; readonly pending: ReadonlyMap; readonly requestSequence: number; readonly focusSequence: number; @@ -129,14 +147,24 @@ const removeConnectionFromState = ( current: BrokerState, clientId: string, queue: ClientConnection["queue"], + releasedAt?: number, ): { readonly state: BrokerState; readonly disconnected: ReadonlyArray } => { 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; @@ -144,7 +172,7 @@ const removeConnectionFromState = ( 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({ 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) { diff --git a/apps/server/src/preview/CaptureLock.test.ts b/apps/server/src/preview/CaptureLock.test.ts new file mode 100644 index 0000000000000..668971e7b5740 --- /dev/null +++ b/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((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(); + 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(() => {}), 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"); + }); +}); diff --git a/apps/server/src/preview/CaptureLock.ts b/apps/server/src/preview/CaptureLock.ts new file mode 100644 index 0000000000000..b64d32853d480 --- /dev/null +++ b/apps/server/src/preview/CaptureLock.ts @@ -0,0 +1,30 @@ +import { constVoid } from "effect/Function"; + +const abandonOnAbort = (promise: Promise, signal: AbortSignal): Promise => + new Promise((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 = Promise.resolve(); + + run(operation: () => Promise, signal?: AbortSignal): Promise { + 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; + } +} diff --git a/apps/server/src/preview/ServerBrowser.test.ts b/apps/server/src/preview/ServerBrowser.test.ts index b87bc8d9ca4d8..30db0a3b8468b 100644 --- a/apps/server/src/preview/ServerBrowser.test.ts +++ b/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) { 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(() => {}); + } // 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 | null = null; let contextFailure: Error | null = null; let initialNavigation: (() => Promise) | null = null; let desktopRendersNext = false; +let stallStills = false; +let onStalledStill: (() => void) | undefined; const desktopDetaches = new NodeEvents.EventEmitter(); const desktopTabs = new Set(); 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(); + 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({ + 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* () { diff --git a/apps/server/src/preview/ServerBrowser.ts b/apps/server/src/preview/ServerBrowser.ts index a2485b30e9e36..05e9052523263 100644 --- a/apps/server/src/preview/ServerBrowser.ts +++ b/apps/server/src/preview/ServerBrowser.ts @@ -83,6 +83,7 @@ import * as ServerBrowserPage from "./ServerBrowserPage.ts"; import * as PreviewBrowser from "./PreviewBrowser.ts"; import * as PreviewBrowserHost from "./PreviewBrowserHost.ts"; import { presentAsChrome, ServerBrowserContexts } from "./ServerBrowserContexts.ts"; +import { CaptureLock } from "./CaptureLock.ts"; import { BrowserControlInterrupted, SessionControl } from "./SessionControl.ts"; import { installPdfDocumentNavigation, isPdfAssetUrl } from "./PdfDocumentNavigation.ts"; @@ -94,6 +95,8 @@ const SCREENCAST_SETTLE_MS = 200; const SCREENCAST_MOTION_FRAMES = 4; const SCREENCAST_MOTION_WINDOW_MS = 300; const SCREENCAST_MOTION_QUALITY = 50; +// A still is a nicety; a guest that is hidden mid-transition may never paint it. +const VIEWER_STILL_TIMEOUT_MS = 1_000; const HOST_RECONNECT_DELAY = "1 second"; const DESKTOP_ATTACH_TIMEOUT = "10 seconds"; const VIEWER_OUTPUT_LIMIT = 64; @@ -321,7 +324,7 @@ interface ServerTab { } | null; recordingStart: Promise | null; - captureLock: Promise; + captureLock: CaptureLock; capturing: number; } @@ -774,7 +777,7 @@ const make = Effect.gen(function* () { initialNavigation: null, navigationGeneration: 0, abortedNavigation: null, - captureLock: Promise.resolve(), + captureLock: new CaptureLock(), capturing: 0, }; if (!desktop) { @@ -1450,13 +1453,27 @@ const make = Effect.gen(function* () { }); }; - const withCaptureLock = (tab: ServerTab, operation: () => Promise): Promise => { - const run = tab.captureLock.then(() => { - if (tab.closing) throw new Error("The preview tab closed."); - return operation(); - }); - tab.captureLock = run.then(constVoid, constVoid); - return run; + // A capture with a deadline gives up the lock when it passes, so one that + // never gets a frame cannot stall the rest until the broker evicts the host. + const withCaptureLock = ( + tab: ServerTab, + operation: () => Promise, + deadlineMs?: number, + ): Promise => { + const remainingMs = deadlineMs === undefined ? undefined : deadlineMs - performance.now(); + if (remainingMs !== undefined && remainingMs <= 0) { + return Promise.reject(ServerBrowserPage.tooBusy()); + } + const signal = + remainingMs === undefined ? undefined : AbortSignal.timeout(Math.round(remainingMs)); + return tab.captureLock + .run(() => { + if (tab.closing) throw new Error("The preview tab closed."); + return operation(); + }, signal) + .catch((cause: unknown) => { + throw signal?.aborted && cause === signal.reason ? ServerBrowserPage.tooBusy() : cause; + }); }; const startRecording = (tab: ServerTab): Promise => { @@ -1521,27 +1538,37 @@ const make = Effect.gen(function* () { return started; }; - const withScreencastsPaused = (tab: ServerTab, capture: () => Promise): Promise => - withCaptureLock(tab, async () => { - tab.capturing += 1; - const recording = tab.recording; - try { - await Promise.all([ - ...[...tab.viewers].map((viewer) => viewer.pause()), - recording?.session.send("Page.stopScreencast").catch(constVoid), - ]); - return await capture(); - } finally { - tab.capturing -= 1; - - await Promise.all([ - ...[...tab.viewers].map((viewer) => viewer.resume()), - recording && tab.recording === recording - ? recording.session.send("Page.startScreencast", RECORDING_SCREENCAST).catch(constVoid) - : undefined, - ]); - } - }); + const withScreencastsPaused = ( + tab: ServerTab, + capture: () => Promise, + deadlineMs?: number, + ): Promise => + withCaptureLock( + tab, + async () => { + tab.capturing += 1; + const recording = tab.recording; + try { + await Promise.all([ + ...[...tab.viewers].map((viewer) => viewer.pause()), + recording?.session.send("Page.stopScreencast").catch(constVoid), + ]); + return await capture(); + } finally { + tab.capturing -= 1; + + await Promise.all([ + ...[...tab.viewers].map((viewer) => viewer.resume()), + recording && tab.recording === recording && tab.capturing === 0 + ? recording.session + .send("Page.startScreencast", RECORDING_SCREENCAST) + .catch(constVoid) + : undefined, + ]); + } + }, + deadlineMs, + ); const stopRecording = (tab: ServerTab) => withCaptureLock(tab, async () => { @@ -1895,7 +1922,7 @@ const make = Effect.gen(function* () { "A browser dialog is pending. Read preview_status.", "dialogPending", ); - return executeTabOperation(tab, request); + return executeTabOperation(tab, request, deadlineMs); }); } return tab.control.agent( @@ -1910,7 +1937,7 @@ const make = Effect.gen(function* () { tab.agentOperating = !READ_OPERATIONS.has(request.operation); updateFileChooserInterception(tab); try { - return await executeTabOperation(tab, request); + return await executeTabOperation(tab, request, deadlineMs); } finally { tab.agentOperating = false; updateFileChooserInterception(tab); @@ -1921,7 +1948,11 @@ const make = Effect.gen(function* () { ); }; - const executeTabOperation = async (tab: ServerTab, request: PreviewAutomationRequest) => { + const executeTabOperation = async ( + tab: ServerTab, + request: PreviewAutomationRequest, + deadlineMs: number, + ) => { const input = request.input; switch (request.operation) { case "navigate": { @@ -1957,11 +1988,14 @@ const make = Effect.gen(function* () { return { tabId: tab.tabId, colorScheme }; } case "snapshot": { - return withScreencastsPaused(tab, () => - ServerBrowserPage.snapshot({ - ...tab, - renderScale: tab.desktop ? "display" : RENDER_SCALE, - }), + return withScreencastsPaused( + tab, + () => + ServerBrowserPage.snapshot({ + ...tab, + renderScale: tab.desktop ? "display" : RENDER_SCALE, + }), + deadlineMs, ); } case "click": { @@ -2332,8 +2366,10 @@ const make = Effect.gen(function* () { broadcastControl(tab); pushFileChooser(tab); const pushStill = async () => { - const data = await withCaptureLock(tab, () => - ServerBrowserPage.captureViewport(session, { format: "jpeg", quality }), + const data = await withCaptureLock( + tab, + () => ServerBrowserPage.captureViewport(session, { format: "jpeg", quality }), + performance.now() + VIEWER_STILL_TIMEOUT_MS, ).catch(() => null); if (data) viewer.push({ _tag: "frame", data: Buffer.from(data, "base64"), ack: Effect.void }); }; diff --git a/apps/server/src/preview/ServerBrowserPage.ts b/apps/server/src/preview/ServerBrowserPage.ts index 9525ac7c5b579..bce93ac5f0346 100644 --- a/apps/server/src/preview/ServerBrowserPage.ts +++ b/apps/server/src/preview/ServerBrowserPage.ts @@ -232,7 +232,7 @@ const pageChanged = () => "The page changed while capturing its snapshot. Take another snapshot.", ); -const tooBusy = () => +export const tooBusy = () => new ServerBrowserOperationError( "PreviewAutomationTimeoutError", "The page was too busy to capture a snapshot in time. Wait for it to settle and take another snapshot.",