diff --git a/apps/server/src/mcp/McpHttpServer.test.ts b/apps/server/src/mcp/McpHttpServer.test.ts index 15518e47a4974..0a73db5e925b4 100644 --- a/apps/server/src/mcp/McpHttpServer.test.ts +++ b/apps/server/src/mcp/McpHttpServer.test.ts @@ -4,12 +4,14 @@ import * as NodeServices from "@effect/platform-node/NodeServices"; import { EnvironmentId, PreviewTabId, ProviderInstanceId, ThreadId } 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"; import * as Layer from "effect/Layer"; import * as Option from "effect/Option"; import * as Path from "effect/Path"; import * as Schema from "effect/Schema"; import * as Stream from "effect/Stream"; +import * as TestClock from "effect/testing/TestClock"; import { McpProtocol, McpSchema, McpServer } from "effect/unstable/ai"; import { HttpBody, HttpClient, HttpRouter, HttpServerResponse } from "effect/unstable/http"; @@ -766,3 +768,76 @@ it.effect("registers annotated tools and preserves authenticated request context }), ).pipe(Effect.provide(TestLayer)), ); + +it.effect("keeps the host after an honest preview_wait_for miss at the caller deadline", () => + Effect.scoped( + Effect.gen(function* () { + const server = yield* McpServer.McpServer; + const broker = yield* PreviewAutomationBroker.PreviewAutomationBroker; + const clientId = "mcp-wait-for-client"; + const connected = yield* Deferred.make(); + const waitForReceived = yield* Deferred.make(); + const routedOperations: Array = []; + const statusResult = { + available: true, + visible: true, + tabId, + url: "http://example.test/", + title: "Example", + loading: false, + }; + const events = yield* broker.connect({ clientId, environmentId }); + yield* Stream.runForEach(events, (event) => { + if (event.type === "connected") return Deferred.succeed(connected, undefined); + routedOperations.push(event.request.operation); + const target = { + clientId, + connectionId: event.connectionId, + requestId: event.request.requestId, + }; + if (event.request.operation !== "waitFor") { + return broker.respond({ ...target, ok: true, result: statusResult }); + } + // The desktop polls until the caller's own deadline and sleeps one more + // tick before it reports the miss, so its reply lands after timeoutMs. + const input = event.request.input as { readonly timeoutMs?: number }; + return Deferred.succeed(waitForReceived, undefined).pipe( + Effect.andThen(Effect.sleep((input.timeoutMs ?? 15_000) + 100)), + Effect.andThen( + broker.respond({ + ...target, + ok: false, + error: { _tag: "PreviewAutomationTimeoutError", message: "Wait timed out" }, + }), + ), + ); + }).pipe(Effect.forkScoped); + yield* Deferred.await(connected); + + const waited = yield* server + .callTool({ name: "preview_wait_for", arguments: { text: "Missing", timeoutMs: 5_000 } }) + .pipe( + Effect.provideService(McpInvocationContext.McpInvocationContext, invocation), + Effect.provideService(McpSchema.McpServerClient, client), + Effect.forkScoped, + ); + yield* Deferred.await(waitForReceived); + yield* TestClock.adjust(5_100); + const missed = yield* Fiber.join(waited); + expect(missed.isError).toBe(true); + expect(missed.content).toEqual([ + { type: "text", text: "Preview automation waitFor timed out after 5000ms." }, + ]); + + const status = yield* server + .callTool({ name: "preview_status", arguments: {} }) + .pipe( + Effect.provideService(McpInvocationContext.McpInvocationContext, invocation), + Effect.provideService(McpSchema.McpServerClient, client), + ); + expect(status.isError).toBe(false); + expect(status.structuredContent).toMatchObject({ available: true, tabId }); + expect(routedOperations).toEqual(["waitFor", "status"]); + }), + ).pipe(Effect.provide(TestLayer)), +); diff --git a/apps/server/src/mcp/PreviewAutomationBroker.test.ts b/apps/server/src/mcp/PreviewAutomationBroker.test.ts index ef552f9b04d06..66dfad64024bf 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.test.ts +++ b/apps/server/src/mcp/PreviewAutomationBroker.test.ts @@ -1359,3 +1359,58 @@ it.effect("keeps a host that responds with an operation timeout", () => }), ), ); + +it.effect("keeps a host whose operation timeout reply lands inside the reply grace", () => + Effect.scoped( + Effect.gen(function* () { + const broker = yield* makeBroker; + const connected = yield* Deferred.make(); + const waitForReceived = yield* Deferred.make(); + const routed: string[] = []; + const events = yield* broker.connect(makeHost()); + yield* Stream.runForEach(events, (event) => { + if (event.type === "connected") return Deferred.succeed(connected, undefined); + routed.push(event.request.operation); + const target = { + clientId: "client-1", + connectionId: event.connectionId, + requestId: event.request.requestId, + }; + if (event.request.operation !== "waitFor") { + return broker.respond({ ...target, ok: true, result: "responsive" }); + } + // The desktop polls until the caller deadline, then reports the miss. + return Deferred.succeed(waitForReceived, undefined).pipe( + Effect.andThen(Effect.sleep(event.request.timeoutMs + 100)), + Effect.andThen( + broker.respond({ + ...target, + ok: false, + error: { _tag: "PreviewAutomationTimeoutError", message: "Selector timed out" }, + }), + ), + ); + }).pipe(Effect.forkScoped); + yield* Deferred.await(connected); + + const missed = yield* broker + .invoke({ + scope, + operation: "waitFor", + input: {}, + timeoutMs: 5_000, + replyGraceMs: 1_500, + }) + .pipe(Effect.flip, Effect.forkScoped); + yield* Deferred.await(waitForReceived); + yield* TestClock.adjust(5_100); + expect(yield* Fiber.join(missed)).toMatchObject({ + _tag: "PreviewAutomationTimeoutError", + timeoutMs: 5_000, + remoteTag: "PreviewAutomationTimeoutError", + }); + expect(yield* broker.invoke({ scope, operation: "status", input: {} })).toBe("responsive"); + expect(routed).toEqual(["waitFor", "status"]); + }), + ), +); diff --git a/apps/server/src/mcp/PreviewAutomationBroker.ts b/apps/server/src/mcp/PreviewAutomationBroker.ts index 527f71f5915d6..f4c8339d6a287 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.ts +++ b/apps/server/src/mcp/PreviewAutomationBroker.ts @@ -45,6 +45,11 @@ export interface PreviewAutomationInvokeInput { readonly input: unknown; readonly tabId?: PreviewTabId; readonly timeoutMs?: number; + /** + * Extra time to wait for a host that answers at its own `timeoutMs`, such as + * a `waitFor` miss. Keeps an honest late reply from reading as a dead host. + */ + readonly replyGraceMs?: number; /** Background metadata reads must not change the agent's current tab. */ readonly updateCurrentTab?: boolean; /** Capture the routed tab before another request changes the current assignment. */ @@ -471,6 +476,7 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() { input: Parameters[0], ): Effect.fn.Return { const timeoutMs = input.timeoutMs ?? 15_000; + const replyGraceMs = input.replyGraceMs ?? 0; const deferred = yield* Deferred.make(); const route = yield* SynchronizedRef.modify(state, (current) => { const assignments = new Map( @@ -599,7 +605,9 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() { } return yield* new PreviewAutomationRequestQueueClosedError(requestContext); } - const result = yield* Deferred.await(deferred).pipe(Effect.timeoutOption(timeoutMs)); + const result = yield* Deferred.await(deferred).pipe( + Effect.timeoutOption(timeoutMs + replyGraceMs), + ); return yield* Option.match(result, { onNone: () => Effect.gen(function* () { diff --git a/apps/server/src/mcp/toolkits/preview/handlers.ts b/apps/server/src/mcp/toolkits/preview/handlers.ts index caa4cbd157cf9..56962641a86b4 100644 --- a/apps/server/src/mcp/toolkits/preview/handlers.ts +++ b/apps/server/src/mcp/toolkits/preview/handlers.ts @@ -50,11 +50,19 @@ export function normalizePreviewOpenInput( }; } +/** + * Desktop `waitFor` polls for the caller's full `timeoutMs` and only then + * reports the miss, so its reply lands after the broker deadline when both + * clocks match. The grace lets that reply arrive before the host is evicted. + */ +const PREVIEW_WAIT_FOR_REPLY_GRACE_MS = 1_500; + const invoke = Effect.fn("PreviewToolkit.invoke")(function* ( operation: PreviewAutomationOperation, input: unknown, timeoutMs?: number, tabId?: PreviewTabId, + replyGraceMs?: number, ): Effect.fn.Return< { result: A; toolIcon?: ToolActivityIcon }, import("@t3tools/contracts").PreviewAutomationError, @@ -72,6 +80,7 @@ const invoke = Effect.fn("PreviewToolkit.invoke")(function* ( input, ...(timeoutMs === undefined ? {} : { timeoutMs }), ...(tabId === undefined ? {} : { tabId }), + ...(replyGraceMs === undefined ? {} : { replyGraceMs }), }); if (["status", "open", "navigate", "snapshot"].includes(operation)) return { result }; const statusTabId = @@ -103,9 +112,10 @@ const invokeTargeted = ( readonly [key: string]: unknown; }, timeoutMs?: number, + replyGraceMs?: number, ) => { const { tabId, ...operationInput } = input; - return invoke(operation, operationInput, timeoutMs, tabId).pipe( + return invoke(operation, operationInput, timeoutMs, tabId, replyGraceMs).pipe( Effect.map(({ result, toolIcon }) => ({ ...result, ...(toolIcon ? { toolIcon } : {}), @@ -212,7 +222,8 @@ const handlers = { ...(toolIcon ? { toolIcon } : {}), })), ), - preview_wait_for: (input) => invokeTargeted("waitFor", input, input.timeoutMs), + preview_wait_for: (input) => + invokeTargeted("waitFor", input, input.timeoutMs, PREVIEW_WAIT_FOR_REPLY_GRACE_MS), preview_recording_start: (input) => invokeTargeted("recordingStart", input ?? {}), preview_recording_stop: (input) =>