diff --git a/apps/server/src/mcp/PreviewAutomationBroker.test.ts b/apps/server/src/mcp/PreviewAutomationBroker.test.ts index 0f86dfafa1a3c..95171a5a15818 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.test.ts +++ b/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({ 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; diff --git a/apps/server/src/mcp/PreviewAutomationBroker.ts b/apps/server/src/mcp/PreviewAutomationBroker.ts index cc3230e5b68b5..87f9eff4676ee 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.ts +++ b/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; 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; @@ -130,14 +148,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; @@ -145,7 +173,7 @@ const removeConnectionFromState = ( 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({ 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) { 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.ts b/apps/server/src/preview/ServerBrowser.ts index c87e829130e25..1fa499aa98688 100644 --- a/apps/server/src/preview/ServerBrowser.ts +++ b/apps/server/src/preview/ServerBrowser.ts @@ -92,6 +92,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"; @@ -395,7 +396,7 @@ interface ServerTab { } | null; recordingStart: Promise | null; - captureLock: Promise; + captureLock: CaptureLock; capturing: number; } @@ -1084,7 +1085,7 @@ const make = Effect.gen(function* () { navigationGeneration: 0, stoppedGeneration: null, abortedNavigation: null, - captureLock: Promise.resolve(), + captureLock: new CaptureLock(), capturing: 0, }; if (!desktop) { @@ -2064,13 +2065,29 @@ const make = Effect.gen(function* () { assertNavigationSucceeded(tab); }; - 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 timeoutMs = + deadlineMs === undefined ? undefined : Math.round(deadlineMs - performance.now()); + if (timeoutMs !== undefined && timeoutMs <= 0) { + return Promise.reject(ServerBrowserPage.snapshotTimedOut(0)); + } + const signal = timeoutMs === undefined ? undefined : AbortSignal.timeout(timeoutMs); + return tab.captureLock + .run(() => { + if (tab.closing) throw new Error("The preview tab closed."); + return operation(); + }, signal) + .catch((cause: unknown) => { + throw timeoutMs !== undefined && signal?.aborted && cause === signal.reason + ? ServerBrowserPage.snapshotTimedOut(timeoutMs) + : cause; + }); }; const startRecording = (tab: ServerTab): Promise => { @@ -2148,27 +2165,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 () => { @@ -2690,12 +2717,16 @@ const make = Effect.gen(function* () { return { tabId: tab.tabId, colorScheme }; } case "snapshot": { - const snapshot = await withScreencastsPaused(tab, () => - ServerBrowserPage.snapshot({ - ...tab, - renderScale: tab.desktop ? "display" : RENDER_SCALE, - fullScreenshot: (input as PreviewAutomationSnapshotInput | undefined)?.fullScreenshot, - }), + const snapshot = await withScreencastsPaused( + tab, + () => + ServerBrowserPage.snapshot({ + ...tab, + deadlineMs, + renderScale: tab.desktop ? "display" : RENDER_SCALE, + fullScreenshot: (input as PreviewAutomationSnapshotInput | undefined)?.fullScreenshot, + }), + deadlineMs, ); const loadError = loadErrorOf(tab); return loadError ? { ...snapshot, url: loadError.url, loadError } : snapshot; diff --git a/apps/server/src/preview/ServerBrowserPage.test.ts b/apps/server/src/preview/ServerBrowserPage.test.ts index 36b133cf185d8..a025d6c03dcd0 100644 --- a/apps/server/src/preview/ServerBrowserPage.test.ts +++ b/apps/server/src/preview/ServerBrowserPage.test.ts @@ -632,6 +632,48 @@ describe("server browser snapshot screenshot", () => { }).toEqual({ size: [1280, 800], colors: ["blue", "red", "red"] }); }); + it("a late metrics reply after the request deadline preserves a newer snapshot's refs", async () => { + const page = await context.newPage(); + await page.setContent(""); + const cdp = await context.newCDPSession(page); + const started = Promise.withResolvers(); + const metrics = Promise.withResolvers<{ + ratio: number; + zoom: number; + width: number; + height: number; + }>(); + const send = cdp.send.bind(cdp); + using _send = vi.spyOn(cdp, "send").mockImplementation(async (method, params) => { + if ( + method === "Runtime.evaluate" && + params && + "expression" in params && + params.expression.includes("devicePixelRatio") + ) { + started.resolve(); + return { result: { value: await metrics.promise } }; + } + return send(method, params); + }); + const input = { page, cdp, consoleEntries: [], networkEntries: [], actionTimeline: [] }; + using now = vi.spyOn(performance, "now").mockReturnValue(0); + const late = ServerBrowserPage.snapshot({ + ...input, + renderScale: "display", + deadlineMs: 10_000, + }); + const rejected = expect(late).rejects.toThrow(); + await started.promise; + const current = await ServerBrowserPage.snapshot({ ...input, renderScale: 2 }); + now.mockReturnValue(10_001); + metrics.resolve({ ratio: 2, zoom: 1, width: 1280, height: 800 }); + await rejected; + const ref = /\[ref=([^\]]+)\]/.exec(String(current.accessibilityTree))![1]!; + await ServerBrowserPage.click(page, { locator: `aria-ref=${ref}` }); + expect(await page.locator("button").innerText()).toBe("clicked"); + }); + it("captures a plain page at the shrunk size with its content in place", async () => { const { screenshot, fullScreenshot, pixels } = await screenshotOf(halves("")); diff --git a/apps/server/src/preview/ServerBrowserPage.ts b/apps/server/src/preview/ServerBrowserPage.ts index bcfbc296a2cb0..8c492f3c11016 100644 --- a/apps/server/src/preview/ServerBrowserPage.ts +++ b/apps/server/src/preview/ServerBrowserPage.ts @@ -289,12 +289,12 @@ interface SnapshotStages { readonly timedOut: Set; } -const snapshotTimedOut = (stages?: SnapshotStages) => +export const snapshotTimedOut = (timeoutMs: number, stages?: SnapshotStages) => new ServerBrowserOperationError({ tag: "PreviewAutomationTimeoutError", - message: `Snapshot capture timed out after ${SNAPSHOT_TIMEOUT_MS}ms. Try the snapshot again.`, + message: `Snapshot capture timed out after ${timeoutMs}ms. Try the snapshot again.`, detail: { - timeoutMs: SNAPSHOT_TIMEOUT_MS, + timeoutMs, pendingStages: stages ? [...new Set([...stages.timedOut, ...stages.running.keys()])] : [], }, }); @@ -603,6 +603,7 @@ const shrinkPng = async (cdp: CDPSession, png: string, scale: number) => { }; type SnapshotInput = { + readonly deadlineMs?: number; readonly page: Page; readonly cdp: CDPSession; /** Headless pages render at the scale the server set; a desktop page at its display's. */ @@ -633,6 +634,7 @@ const renderingOf = async (input: SnapshotInput) => { const capture = async ( input: SnapshotInput, deadline: number, + timeoutMs: number, stages: SnapshotStages, ): Promise => { const { running } = stages; @@ -653,7 +655,7 @@ const capture = async ( }; const { renderScale, viewport } = await track("metrics", () => renderingOf(input)); // A late metrics reply must not begin a capture after viewers have resumed. - if (Date.now() >= deadline) throw snapshotTimedOut(); + if (performance.now() >= deadline) throw snapshotTimedOut(timeoutMs); const scale = Math.min(1, MAX_SCREENSHOT_WIDTH / (viewport.width * renderScale)); const state = refsFor(input.page); invalidateRefs(input.page); @@ -689,7 +691,7 @@ const capture = async ( input.page.ariaSnapshot({ mode: "ai", boxes: true, - timeout: Math.max(1, deadline - Date.now() - 2_000), + timeout: Math.max(1, deadline - performance.now() - 2_000), }), ); const after = await track("frames", () => read(before)); @@ -744,26 +746,28 @@ const capture = async ( const snapshotTurns = new WeakMap(); export const snapshot = async (input: SnapshotInput): Promise => { - const deadline = Date.now() + SNAPSHOT_TIMEOUT_MS; + const deadline = Math.min(performance.now() + SNAPSHOT_TIMEOUT_MS, input.deadlineMs ?? Infinity); + const timeoutMs = Math.round(deadline - performance.now()); + if (timeoutMs <= 0) throw snapshotTimedOut(SNAPSHOT_TIMEOUT_MS); const stages: SnapshotStages = { running: new Map(), timedOut: new Set() }; let timer: ReturnType | undefined; const expired = new Promise((_, reject) => { - timer = setTimeout(() => reject(snapshotTimedOut(stages)), SNAPSHOT_TIMEOUT_MS); + timer = setTimeout(() => reject(snapshotTimedOut(timeoutMs, stages)), timeoutMs); }); const turn = (snapshotTurns.get(input.page) ?? 0) + 1; snapshotTurns.set(input.page, turn); // Retry only before the deadline and the next snapshot, whose refs a later retry would revoke. - const captured = capture(input, deadline, stages).catch((cause: unknown) => - Date.now() < deadline && + const captured = capture(input, deadline, timeoutMs, stages).catch((cause: unknown) => + performance.now() < deadline && snapshotTurns.get(input.page) === turn && !(cause instanceof Error && cause.message === VALUE_TOO_LONG) - ? capture(input, deadline, stages) + ? capture(input, deadline, timeoutMs, stages) : Promise.reject(cause), ); try { return await Promise.race([captured, expired]); } catch (cause) { - throw isTimeout(cause) ? snapshotTimedOut(stages) : cause; + throw isTimeout(cause) ? snapshotTimedOut(timeoutMs, stages) : cause; } finally { clearTimeout(timer); }