diff --git a/apps/server/src/mcp/PreviewAutomationBroker.test.ts b/apps/server/src/mcp/PreviewAutomationBroker.test.ts index 0371c97c5fd99e23b767d35f8256a12594b51bf7..ef552f9b04d06e51dc55d9faf51681b461e99c3a 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.test.ts +++ b/apps/server/src/mcp/PreviewAutomationBroker.test.ts @@ -10,18 +10,22 @@ import { PreviewTabId, ProviderInstanceId, ThreadId, + WS_METHODS, + WsRpcGroup, type PreviewAutomationHost, type PreviewAutomationRequest, type PreviewAutomationStreamEvent, } from "@t3tools/contracts"; import * as Effect from "effect/Effect"; import * as Exit from "effect/Exit"; -import * as Cause from "effect/Cause"; import * as Deferred from "effect/Deferred"; import * as Fiber from "effect/Fiber"; import * as Result from "effect/Result"; +import * as Scheduler from "effect/Scheduler"; import * as Stream from "effect/Stream"; import * as TestClock from "effect/testing/TestClock"; +import * as RpcGroup from "effect/unstable/rpc/RpcGroup"; +import * as RpcTest from "effect/unstable/rpc/RpcTest"; import * as PreviewAutomationBroker from "./PreviewAutomationBroker.ts"; @@ -1118,7 +1122,19 @@ it.effect("evicts an unanswered host and lets later calls use a healthy runtime" const otherReceived = yield* Deferred.make(); const otherCompleted = yield* Deferred.make(); const oldTab = PreviewTabId.make("tab-on-frozen-host"); - const events = yield* broker.connect(makeHost()); + const group = RpcGroup.make( + ...Array.from(WsRpcGroup.requests.values()).filter( + (rpc) => rpc._tag === WS_METHODS.previewAutomationConnect, + ), + ); + const client = yield* RpcTest.makeClient(group).pipe( + Effect.provide( + group.toLayer({ + [WS_METHODS.previewAutomationConnect]: (host) => Stream.unwrap(broker.connect(host)), + }), + ), + ); + const events = client[WS_METHODS.previewAutomationConnect](makeHost()); const consumer = yield* Stream.runForEach(events, (event) => { if (event.type === "connected") return Deferred.succeed(connected, event.connectionId); const request = { ...event.request, connectionId: event.connectionId }; @@ -1183,10 +1199,7 @@ it.effect("evicts an unanswered host and lets later calls use a healthy runtime" _tag: "PreviewAutomationClientDisconnectedError", }); const consumerExit = yield* Fiber.await(consumer); - expect(Exit.isFailure(consumerExit)).toBe(true); - if (Exit.isFailure(consumerExit)) { - expect(Cause.hasInterruptsOnly(consumerExit.cause)).toBe(true); - } + expect(Exit.isSuccess(consumerExit)).toBe(true); // Late traffic from the evicted connection cannot restore its assignment. yield* broker.respond({ @@ -1209,6 +1222,115 @@ it.effect("evicts an unanswered host and lets later calls use a healthy runtime" ), ); +it.effect("discards buffered actions before completing an evicted host stream", () => + Effect.scoped( + Effect.gen(function* () { + const broker = yield* makeBroker; + const connected = yield* Deferred.make(); + const received = yield* Deferred.make(); + const actionRouted = yield* Deferred.make(); + const releaseConsumer = yield* Deferred.make(); + const operations: string[] = []; + const consumer = yield* Stream.runForEach(yield* broker.connect(makeHost()), (event) => { + if (event.type === "connected") return Deferred.succeed(connected, undefined); + operations.push(event.request.operation); + return Deferred.succeed(received, undefined).pipe( + Effect.andThen(Deferred.await(releaseConsumer)), + ); + }).pipe(Effect.forkScoped); + yield* Deferred.await(connected); + const timedOut = yield* broker + .invoke({ scope, operation: "snapshot", input: {}, timeoutMs: 1_000 }) + .pipe(Effect.flip, Effect.forkScoped); + yield* Deferred.await(received); + const buffered = yield* broker + .invoke({ + scope, + operation: "click", + input: {}, + timeoutMs: 10_000, + onTargetTab: () => { + Deferred.doneUnsafe(actionRouted, Effect.void); + }, + }) + .pipe(Effect.flip, Effect.forkScoped); + yield* Deferred.await(actionRouted); + yield* TestClock.adjust(1_000); + expect(yield* Fiber.join(timedOut)).toMatchObject({ _tag: "PreviewAutomationTimeoutError" }); + expect(yield* Fiber.join(buffered)).toMatchObject({ + _tag: "PreviewAutomationClientDisconnectedError", + }); + yield* Deferred.succeed(releaseConsumer, undefined); + expect(Exit.isSuccess(yield* Fiber.await(consumer))).toBe(true); + expect(operations).toEqual(["snapshot"]); + }), + ), +); + +it.effect("rejects a routed action when its generation is evicted before delivery", () => + Effect.scoped( + Effect.gen(function* () { + const broker = yield* makeBroker; + const connected = yield* Deferred.make(); + const received = yield* Deferred.make(); + const actionRouted = yield* Deferred.make(); + const operations: string[] = []; + const consumer = yield* Stream.runForEach(yield* broker.connect(makeHost()), (event) => { + if (event.type === "connected") return Deferred.succeed(connected, undefined); + operations.push(event.request.operation); + return Deferred.succeed(received, undefined); + }).pipe(Effect.forkScoped); + yield* Deferred.await(connected); + const timedOut = yield* broker + .invoke({ scope, operation: "snapshot", input: {}, timeoutMs: 1_000 }) + .pipe(Effect.flip, Effect.forkScoped); + yield* Deferred.await(received); + + // Suspend only this invocation in the gap between route selection and delivery. + const tasks: Array<() => void> = []; + let paused = false; + const dispatcher: Scheduler.SchedulerDispatcher = { + scheduleTask: (task) => tasks.push(task), + flush: () => { + let task: (() => void) | undefined; + while ((task = tasks.shift()) !== undefined) task(); + }, + }; + const scheduler: Scheduler.Scheduler = { + executionMode: "async", + makeDispatcher: () => dispatcher, + shouldYield: () => paused, + }; + const action = yield* broker + .invoke({ + scope, + operation: "click", + input: {}, + onTargetTab: () => { + paused = true; + Deferred.doneUnsafe(actionRouted, Effect.void); + }, + }) + .pipe( + Effect.flip, + Effect.provideService(Scheduler.Scheduler, scheduler), + Effect.forkScoped, + ); + yield* Deferred.await(actionRouted); + yield* TestClock.adjust(1_000); + expect(yield* Fiber.join(timedOut)).toMatchObject({ _tag: "PreviewAutomationTimeoutError" }); + expect(Exit.isSuccess(yield* Fiber.await(consumer))).toBe(true); + + paused = false; + dispatcher.flush(); + expect(yield* Fiber.join(action)).toMatchObject({ + _tag: "PreviewAutomationClientDisconnectedError", + }); + expect(operations).toEqual(["snapshot"]); + }), + ), +); + it.effect("keeps a host that responds with an operation timeout", () => Effect.scoped( Effect.gen(function* () { diff --git a/apps/server/src/mcp/PreviewAutomationBroker.ts b/apps/server/src/mcp/PreviewAutomationBroker.ts index 418835ae81d926b2b48af967872c59479c0aee21..527f71f5915d63c1d094e82b0b5b81a121dd2d43 100644 --- a/apps/server/src/mcp/PreviewAutomationBroker.ts +++ b/apps/server/src/mcp/PreviewAutomationBroker.ts @@ -26,6 +26,7 @@ import { type PreviewAutomationStreamEvent, } from "@t3tools/contracts"; import * as Context from "effect/Context"; +import type * as Cause from "effect/Cause"; import * as Crypto from "effect/Crypto"; import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; @@ -73,7 +74,7 @@ interface ClientConnection { readonly supportedOperations: ReadonlySet; readonly focused: boolean; readonly focusOrder: number; - readonly queue: Queue.Queue; + readonly queue: Queue.Queue; } interface PendingRequest { @@ -326,32 +327,47 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() { const closeConnection = Effect.fn("PreviewAutomationBroker.closeConnection")(function* ( queue: ClientConnection["queue"], disconnected: ReadonlyArray, + completeStream = false, ) { + if (completeStream) { + // Discard this generation's commands and complete the RPC stream so a + // responsive desktop can re-register after a timeout eviction. + yield* Queue.clear(queue); + yield* Queue.end(queue); + } else { + // Replaced registrations must not reconnect and displace their successor. + yield* Queue.shutdown(queue); + } yield* Effect.forEach( disconnected, ({ deferred, context }) => Deferred.fail(deferred, new PreviewAutomationClientDisconnectedError(context)), { discard: true }, ); - yield* Queue.shutdown(queue); }); const disconnect = Effect.fn("PreviewAutomationBroker.disconnect")(function* ( clientId: string, queue: ClientConnection["queue"], + completeStream = false, ) { - const disconnected = yield* SynchronizedRef.modify(state, (current) => { + yield* SynchronizedRef.modifyEffect(state, (current) => { + // Retired generations were already closed by their replacement or eviction. + if (current.clients.get(clientId)?.queue !== queue) { + return Effect.succeed([undefined, current] as const); + } const removed = removeConnectionFromState(current, clientId, queue); - return [removed.disconnected, removed.state] as const; + return closeConnection(queue, removed.disconnected, completeStream).pipe( + Effect.as([undefined, removed.state] as const), + ); }); - yield* closeConnection(queue, disconnected); }); const acquireConnection = Effect.fn("PreviewAutomationBroker.acquireConnection")(function* ( host: PreviewAutomationHost, ) { const clientId = host.clientId; - const queue = yield* Queue.unbounded(); + const queue = yield* Queue.unbounded(); const connectionId = yield* crypto.randomUUIDv4.pipe(Effect.orDie); yield* Queue.offer(queue, { type: "connected", connectionId }); const connection: ClientConnection = { @@ -553,18 +569,28 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() { return { ...next, pending }; }); const awaitResponse = Effect.fn("PreviewAutomationBroker.awaitResponse")(function* () { - const offered = yield* Queue.offer(connection.queue, { - type: "request", - connectionId: connection.connectionId, - request: { - requestId, - threadId: input.scope.threadId, - tabId: requestContext.tabId, - tabIdExplicit: input.tabId !== undefined, - operation: input.operation, - input: input.input, - timeoutMs, - }, + const offered = yield* SynchronizedRef.modifyEffect(state, (current) => { + // A route can outlive its generation while another request evicts it. + // Serialize the live-generation check and offer with queue closure. + if ( + current.clients.get(connection.clientId)?.queue !== connection.queue || + !current.pending.has(requestId) + ) { + return Effect.succeed([false, current] as const); + } + return Queue.offer(connection.queue, { + type: "request", + connectionId: connection.connectionId, + request: { + requestId, + threadId: input.scope.threadId, + tabId: requestContext.tabId, + tabIdExplicit: input.tabId !== undefined, + operation: input.operation, + input: input.input, + timeoutMs, + }, + }).pipe(Effect.map((offered) => [offered, current] as const)); }); if (!offered) { const completion = yield* Deferred.poll(deferred); @@ -579,7 +605,7 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() { Effect.gen(function* () { // An unanswered request invalidates this connection. Do not replay // actions: the client may have applied them before becoming unreachable. - yield* disconnect(connection.clientId, connection.queue); + yield* disconnect(connection.clientId, connection.queue, true); return yield* new PreviewAutomationTimeoutError(requestContext); }), onSome: (value) => Effect.succeed(value as A), diff --git a/packages/client-runtime/src/rpc/client.test.ts b/packages/client-runtime/src/rpc/client.test.ts index 9e4e8a600d559ddad92dad8ae6a33f2b8146a42e..032603f28a96b67b377b9d1887c65dba10a096bb 100644 --- a/packages/client-runtime/src/rpc/client.test.ts +++ b/packages/client-runtime/src/rpc/client.test.ts @@ -1,6 +1,9 @@ import { DEFAULT_SERVER_SETTINGS, EnvironmentId, + PreviewTabId, + ThreadId, + type PreviewAutomationStreamEvent, type RelayClientInstallProgressEvent, type ServerConfigStreamEvent, type ServerLifecycleStreamEvent, @@ -88,6 +91,156 @@ const makeHarness = Effect.fn("TestEnvironmentRpc.makeHarness")(function* () { }); describe("environment RPC", () => { + it.effect("registers a fresh preview host after completion without replaying requests", () => + Effect.gen(function* () { + const firstCompleted = yield* Deferred.make(); + const reconnected = yield* Deferred.make(); + const requests: string[] = []; + const connections: string[] = []; + let attempts = 0; + const client = { + [WS_METHODS.previewAutomationConnect]: () => + Stream.suspend(() => { + attempts += 1; + const connected: PreviewAutomationStreamEvent = { + type: "connected", + connectionId: `connection-${attempts}`, + }; + return attempts === 1 + ? Stream.make(connected, { + type: "request", + connectionId: connected.connectionId, + request: { + requestId: "timed-out-action", + operation: "click", + threadId: ThreadId.make("thread-1"), + tabId: PreviewTabId.make("tab-1"), + input: {}, + timeoutMs: 1_000, + }, + } satisfies PreviewAutomationStreamEvent).pipe( + Stream.ensuring(Deferred.succeed(firstCompleted, undefined)), + ) + : Stream.succeed(connected).pipe(Stream.concat(Stream.never)); + }), + } as unknown as WsRpcProtocolClient; + const { activeSession, supervisor } = yield* makeHarness(); + yield* SubscriptionRef.set(activeSession, Option.some(session(client))); + const consumer = yield* subscribe(WS_METHODS.previewAutomationConnect, { + clientId: "preview-host", + environmentId: TARGET.environmentId, + }).pipe( + Stream.runForEach((event) => { + if (event.type === "request") { + requests.push(event.request.requestId); + return Effect.void; + } + connections.push(event.connectionId); + return connections.length === 2 ? Deferred.succeed(reconnected, undefined) : Effect.void; + }), + Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor), + Effect.forkChild, + ); + yield* Deferred.await(firstCompleted); + yield* TestClock.adjust(999); + expect(attempts).toBe(1); + yield* TestClock.adjust(1); + expect(attempts).toBe(2); + yield* Deferred.await(reconnected); + expect(connections).toEqual(["connection-1", "connection-2"]); + expect(requests).toEqual(["timed-out-action"]); + yield* Fiber.interrupt(consumer); + }), + ); + + it.effect("does not re-register an unmounted preview host during the recovery delay", () => + Effect.gen(function* () { + const completed = yield* Deferred.make(); + let attempts = 0; + const client = { + [WS_METHODS.previewAutomationConnect]: () => + Stream.suspend(() => { + attempts += 1; + return Stream.empty.pipe(Stream.ensuring(Deferred.succeed(completed, undefined))); + }), + } as unknown as WsRpcProtocolClient; + const { activeSession, supervisor } = yield* makeHarness(); + yield* SubscriptionRef.set(activeSession, Option.some(session(client))); + const consumer = yield* subscribe(WS_METHODS.previewAutomationConnect, { + clientId: "preview-host", + environmentId: TARGET.environmentId, + }).pipe( + Stream.runDrain, + Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor), + Effect.forkChild, + ); + yield* Deferred.await(completed); + yield* Fiber.interrupt(consumer); + yield* TestClock.adjust(10_000); + expect(attempts).toBe(1); + }), + ); + + it.effect.each(["completion", "transport failure"] as const)( + "keeps preview recovery tied to the active session after %s", + (reason) => + Effect.gen(function* () { + const completed = yield* Deferred.make(); + const nextConnected = yield* Deferred.make(); + let oldAttempts = 0; + let nextAttempts = 0; + const firstClient = { + [WS_METHODS.previewAutomationConnect]: () => + Stream.suspend(() => { + oldAttempts += 1; + return ( + reason === "completion" + ? Stream.empty + : Stream.fail( + new RpcClientError.RpcClientError({ + reason: new RpcClientError.RpcClientDefect({ + message: "socket closed", + cause: new Error("socket closed"), + }), + }), + ) + ).pipe(Stream.ensuring(Deferred.succeed(completed, undefined))); + }), + } as unknown as WsRpcProtocolClient; + const nextClient = { + [WS_METHODS.previewAutomationConnect]: () => + Stream.suspend(() => { + nextAttempts += 1; + return Stream.fromEffect(Deferred.succeed(nextConnected, undefined)).pipe( + Stream.drain, + Stream.concat(Stream.never), + ); + }), + } as unknown as WsRpcProtocolClient; + const { activeSession, supervisor } = yield* makeHarness(); + yield* SubscriptionRef.set(activeSession, Option.some(session(firstClient))); + const consumer = yield* subscribe(WS_METHODS.previewAutomationConnect, { + clientId: "preview-host", + environmentId: TARGET.environmentId, + }).pipe( + Stream.runDrain, + Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor), + Effect.forkChild, + ); + yield* Deferred.await(completed); + if (reason === "transport failure") { + yield* TestClock.adjust(10_000); + expect(oldAttempts).toBe(1); + } + yield* SubscriptionRef.set(activeSession, Option.some(session(nextClient))); + yield* Deferred.await(nextConnected); + yield* TestClock.adjust(10_000); + expect(oldAttempts).toBe(1); + expect(nextAttempts).toBe(1); + yield* Fiber.interrupt(consumer); + }), + ); + it.effect("reuses the session config stream instead of opening a duplicate subscription", () => Effect.gen(function* () { const event: ServerConfigStreamEvent = { diff --git a/packages/client-runtime/src/rpc/client.ts b/packages/client-runtime/src/rpc/client.ts index e80a2f0f4b12f37c85d34aaf659bf2968f129235..4f725378ffbf3ed6fef078fdcc1c7322ea319e33 100644 --- a/packages/client-runtime/src/rpc/client.ts +++ b/packages/client-runtime/src/rpc/client.ts @@ -5,6 +5,7 @@ import type * as Duration from "effect/Duration"; import * as Effect from "effect/Effect"; import * as Option from "effect/Option"; import * as Schema from "effect/Schema"; +import * as Schedule from "effect/Schedule"; import * as Stream from "effect/Stream"; import * as SubscriptionRef from "effect/SubscriptionRef"; import { RpcClientError } from "effect/unstable/rpc"; @@ -233,9 +234,15 @@ function subscribeDynamicMapped( method: tag, input, }); - return mapStream(session, method(input)).pipe( - Stream.ensuring(completeObservation), - ); + const stream = mapStream(session, method(input)); + // An evicted preview host completes its registration stream. + // Re-register only after completion; failures still follow the + // session recovery policy and browser actions are never replayed. + return ( + tag === WS_METHODS.previewAutomationConnect + ? stream.pipe(Stream.repeat(Schedule.spaced("1 second"))) + : stream + ).pipe(Stream.ensuring(completeObservation)); }), ).pipe( Stream.tapCause((cause) =>