Harden preview ownership and option-based secret handling (#3172)

Author: juliusmarminge

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

@@ -107,7 +107,15 @@ it.effect("registers annotated tools and preserves authenticated request context
     Effect.gen(function* () {
       const server = yield* McpServer.McpServer;
       const broker = yield* PreviewAutomationBroker.PreviewAutomationBroker;
-      const requests = yield* broker.connect("mcp-test-client");
+      const requests = yield* broker.connect({
+        clientId: "mcp-test-client",
+        environmentId,
+        threadId,
+        tabId,
+        visible: true,
+        supportsAutomation: true,
+        focusedAt: "2026-06-11T00:00:00.000Z",
+      });
       yield* Stream.runForEach(requests, (request) =>
         broker.respond({
           requestId: request.requestId,

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

@@ -2,10 +2,13 @@ import { expect, it } from "@effect/vitest";
 import {
   EnvironmentId,
   PreviewAutomationNoFocusedOwnerError,
+  PreviewAutomationUnavailableError,
   ProviderInstanceId,
   ThreadId,
+  type PreviewAutomationOwner,
 } from "@t3tools/contracts";
 import * as Effect from "effect/Effect";
+import * as Fiber from "effect/Fiber";
 import * as Stream from "effect/Stream";
 
 import * as PreviewAutomationBroker from "./PreviewAutomationBroker.ts";
@@ -20,11 +23,22 @@ const scope = {
   expiresAt: 2,
 };
 
-it.effect("routes a request to the focused owner and correlates its response", () =>
+const makeOwner = (overrides: Partial<PreviewAutomationOwner> = {}): PreviewAutomationOwner => ({
+  clientId: "client-1",
+  environmentId: scope.environmentId,
+  threadId: scope.threadId,
+  tabId: null,
+  visible: false,
+  supportsAutomation: true,
+  focusedAt: "2026-06-11T00:00:00.000Z",
+  ...overrides,
+});
+
+it.effect("atomically registers a connected owner and correlates its response", () =>
   Effect.scoped(
     Effect.gen(function* () {
       const broker = yield* PreviewAutomationBroker.__testing.make;
-      const requests = yield* broker.connect("client-1");
+      const requests = yield* broker.connect(makeOwner());
       yield* Stream.runForEach(requests, (request) =>
         broker.respond({
           requestId: request.requestId,
@@ -33,15 +47,6 @@ it.effect("routes a request to the focused owner and correlates its response", (
         }),
       ).pipe(Effect.forkScoped);
       yield* Effect.yieldNow;
-      yield* broker.reportOwner({
-        clientId: "client-1",
-        environmentId: scope.environmentId,
-        threadId: scope.threadId,
-        tabId: null,
-        visible: false,
-        supportsAutomation: true,
-        focusedAt: "2026-06-11T00:00:00.000Z",
-      });
 
       const result = yield* broker.invoke<{ available: boolean }>({
         scope,
@@ -68,22 +73,119 @@ it.effect("routes interactive commands to a hidden durable browser host", () =>
   Effect.scoped(
     Effect.gen(function* () {
       const broker = yield* PreviewAutomationBroker.__testing.make;
-      const requests = yield* broker.connect("client-hidden");
+      const requests = yield* broker.connect(
+        makeOwner({ clientId: "client-hidden", tabId: "tab-hidden" }),
+      );
+      yield* Stream.runForEach(requests, (request) =>
+        broker.respond({ requestId: request.requestId, ok: true }),
+      ).pipe(Effect.forkScoped);
+      yield* Effect.yieldNow;
+
+      yield* broker.invoke<void>({ scope, operation: "click", input: { x: 10, y: 10 } });
+    }),
+  ),
+);
+
+it.effect("lets the browser host resolve an active tab that has not been reported yet", () =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const broker = yield* PreviewAutomationBroker.__testing.make;
+      const requests = yield* broker.connect(makeOwner({ tabId: null }));
+      let routedTabId: string | undefined;
+      yield* Stream.runForEach(requests, (request) => {
+        routedTabId = request.tabId;
+        return broker.respond({ requestId: request.requestId, ok: true });
+      }).pipe(Effect.forkScoped);
+      yield* Effect.yieldNow;
+
+      yield* broker.invoke<void>({ scope, operation: "click", input: { x: 10, y: 10 } });
+
+      expect(routedTabId).toBeUndefined();
+    }),
+  ),
+);
+
+it.effect("preserves current owner metadata when its request stream reconnects", () =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const broker = yield* PreviewAutomationBroker.__testing.make;
+      const firstRequests = yield* broker.connect(makeOwner());
+      yield* Stream.runDrain(firstRequests).pipe(Effect.forkScoped);
+      yield* broker.reportOwner(makeOwner({ tabId: "tab-current", visible: true }));
+
+      const reconnectedRequests = yield* broker.connect(makeOwner());
+      let routedTabId: string | undefined;
+      yield* Stream.runForEach(reconnectedRequests, (request) => {
+        routedTabId = request.tabId;
+        return broker.respond({ requestId: request.requestId, ok: true });
+      }).pipe(Effect.forkScoped);
+      yield* Effect.yieldNow;
+
+      yield* broker.invoke<void>({ scope, operation: "click", input: { x: 10, y: 10 } });
+
+      expect(routedTabId).toBe("tab-current");
+    }),
+  ),
+);
+
+it.effect("ignores stale owner cleanup after the client moves to another thread", () =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const broker = yield* PreviewAutomationBroker.__testing.make;
+      const requests = yield* broker.connect(makeOwner());
       yield* Stream.runForEach(requests, (request) =>
         broker.respond({ requestId: request.requestId, ok: true }),
       ).pipe(Effect.forkScoped);
       yield* Effect.yieldNow;
-      yield* broker.reportOwner({
-        clientId: "client-hidden",
+
+      yield* broker.clearOwner({
+        clientId: "client-1",
         environmentId: scope.environmentId,
-        threadId: scope.threadId,
-        tabId: "tab-hidden",
-        visible: false,
-        supportsAutomation: true,
-        focusedAt: "2026-06-11T00:00:00.000Z",
+        threadId: ThreadId.make("thread-stale"),
       });
 
-      yield* broker.invoke<void>({ scope, operation: "click", input: { x: 10, y: 10 } });
+      yield* broker.invoke<void>({ scope, operation: "status", input: {} });
+    }),
+  ),
+);
+
+it.effect("fails requests assigned to a browser stream when that stream reconnects", () =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const broker = yield* PreviewAutomationBroker.__testing.make;
+      const _requests = yield* broker.connect(makeOwner());
+      const pending = yield* broker
+        .invoke<void>({ scope, operation: "status", input: {} })
+        .pipe(Effect.flip, Effect.forkScoped);
+      yield* Effect.yieldNow;
+
+      const _replacementRequests = yield* broker.connect(makeOwner());
+
+      const error = yield* Fiber.join(pending);
+      expect(error).toBeInstanceOf(PreviewAutomationUnavailableError);
+    }),
+  ),
+);
+
+it.effect("falls back to an older connected owner when a newer report is not connected", () =>
+  Effect.scoped(
+    Effect.gen(function* () {
+      const broker = yield* PreviewAutomationBroker.__testing.make;
+      const requests = yield* broker.connect(makeOwner({ clientId: "client-connected" }));
+      yield* Stream.runForEach(requests, (request) =>
+        broker.respond({ requestId: request.requestId, ok: true, result: "connected" }),
+      ).pipe(Effect.forkScoped);
+      yield* Effect.yieldNow;
+      yield* broker.reportOwner(
+        makeOwner({
+          clientId: "client-report-only",
+          focusedAt: "2026-06-11T00:00:01.000Z",
+        }),
+      );
+
+      const result = yield* broker.invoke<string>({ scope, operation: "status", input: {} });
+
+      expect(result).toBe("connected");
     }),
   ),
 );

apps/server/src/mcp/PreviewAutomationBroker.ts

@@ -11,6 +11,7 @@ import {
   type PreviewAutomationError,
   type PreviewAutomationOperation,
   type PreviewAutomationOwner,
+  type PreviewAutomationOwnerIdentity,
   type PreviewAutomationRequest,
   type PreviewAutomationResponse,
   type PreviewTabId,
@@ -35,11 +36,13 @@ export interface PreviewAutomationInvokeInput {
 }
 
 export interface PreviewAutomationBrokerShape {
-  readonly connect: (clientId: string) => Effect.Effect<Stream.Stream<PreviewAutomationRequest>>;
+  readonly connect: (
+    owner: PreviewAutomationOwner,
+  ) => Effect.Effect<Stream.Stream<PreviewAutomationRequest>>;
   readonly reportOwner: (
     owner: PreviewAutomationOwner,
   ) => Effect.Effect<void, PreviewAutomationError>;
-  readonly clearOwner: (clientId: string) => Effect.Effect<void>;
+  readonly clearOwner: (owner: PreviewAutomationOwnerIdentity) => Effect.Effect<void>;
   readonly respond: (
     response: PreviewAutomationResponse,
   ) => Effect.Effect<void, PreviewAutomationError>;
@@ -63,7 +66,7 @@ interface ClientConnection {
 }
 
 interface PendingRequest {
-  readonly clientId: string;
+  readonly queue: ClientConnection["queue"];
   readonly deferred: Deferred.Deferred<unknown, PreviewAutomationError>;
 }
 
@@ -133,17 +136,16 @@ const make = Effect.gen(function* PreviewAutomationBrokerMake() {
     queue: ClientConnection["queue"],
   ) {
     const toFail = yield* SynchronizedRef.modify(state, (current) => {
-      if (current.clients.get(clientId)?.queue !== queue) {
-        return [[] as ReadonlyArray<PendingRequest>, current] as const;
-      }
       const clients = new Map(current.clients);
       const owners = new Map(current.owners);
       const pending = new Map(current.pending);
       const disconnected: PendingRequest[] = [];
-      clients.delete(clientId);
-      owners.delete(clientId);
+      if (current.clients.get(clientId)?.queue === queue) {
+        clients.delete(clientId);
+        owners.delete(clientId);
+      }
       for (const [requestId, entry] of pending) {
-        if (entry.clientId === clientId) {
+        if (entry.queue === queue) {
           pending.delete(requestId);
           disconnected.push(entry);
         }
@@ -166,12 +168,22 @@ const make = Effect.gen(function* PreviewAutomationBrokerMake() {
 
   const connect: PreviewAutomationBrokerShape["connect"] = Effect.fn(
     "PreviewAutomationBroker.connect",
-  )(function* (clientId) {
+  )(function* (owner) {
+    const clientId = owner.clientId;
     const queue = yield* Queue.unbounded<import("@t3tools/contracts").PreviewAutomationRequest>();
     const previous = yield* SynchronizedRef.modify(state, (current) => {
       const clients = new Map(current.clients);
+      const owners = new Map(current.owners);
+      const existingOwner = current.owners.get(clientId);
       clients.set(clientId, { clientId, queue });
-      return [current.clients.get(clientId), { ...current, clients }] as const;
+      owners.set(
+        clientId,
+        existingOwner?.environmentId === owner.environmentId &&
+          existingOwner.threadId === owner.threadId
+          ? { ...existingOwner, supportsAutomation: owner.supportsAutomation }
+          : owner,
+      );
+      return [current.clients.get(clientId), { ...current, clients, owners }] as const;
     });
     if (previous) yield* disconnect(clientId, previous.queue);
     return Stream.fromQueue(queue).pipe(Stream.ensuring(disconnect(clientId, queue)));
@@ -189,10 +201,18 @@ const make = Effect.gen(function* PreviewAutomationBrokerMake() {
 
   const clearOwner: PreviewAutomationBrokerShape["clearOwner"] = Effect.fn(
     "PreviewAutomationBroker.clearOwner",
-  )(function* (clientId) {
+  )(function* (owner) {
     yield* SynchronizedRef.update(state, (current) => {
+      const currentOwner = current.owners.get(owner.clientId);
+      if (
+        !currentOwner ||
+        currentOwner.environmentId !== owner.environmentId ||
+        currentOwner.threadId !== owner.threadId
+      ) {
+        return current;
+      }
       const owners = new Map(current.owners);
-      owners.delete(clientId);
+      owners.delete(owner.clientId);
       return { ...current, owners };
     });
   });
@@ -234,8 +254,13 @@ const make = Effect.gen(function* PreviewAutomationBrokerMake() {
           owner.supportsAutomation,
       )
       .sort((left, right) => right.focusedAt.localeCompare(left.focusedAt));
-    const owner = candidates[0];
+    const owner = candidates.find((candidate) => current.clients.has(candidate.clientId));
     if (!owner) {
+      if (candidates.length > 0) {
+        return yield* new PreviewAutomationUnavailableError({
+          message: "The browser host is not connected.",
+        });
+      }
       return yield* new PreviewAutomationNoFocusedOwnerError({
         message: "No desktop browser host is available for this thread.",
       });
@@ -246,22 +271,12 @@ const make = Effect.gen(function* PreviewAutomationBrokerMake() {
         message: "The browser host is not connected.",
       });
     }
-    if (
-      input.operation !== "open" &&
-      input.operation !== "status" &&
-      !owner.tabId &&
-      !input.tabId
-    ) {
-      return yield* new PreviewAutomationTabNotFoundError({
-        message: "The browser host does not have an active tab.",
-      });
-    }
     const timeoutMs = input.timeoutMs ?? 15_000;
     const deferred = yield* Deferred.make<unknown, PreviewAutomationError>();
     const requestId = yield* SynchronizedRef.modify(state, (next) => {
       const requestId = `preview-${next.requestSequence}`;
       const pending = new Map(next.pending);
-      pending.set(requestId, { clientId: owner.clientId, deferred });
+      pending.set(requestId, { queue: connection.queue, deferred });
       return [requestId, { ...next, pending, requestSequence: next.requestSequence + 1 }] as const;
     });
     const removePending = SynchronizedRef.update(state, (next) => {

apps/server/src/ws.ts

@@ -1489,7 +1489,7 @@ const makeWsRpcLayer = (currentSession: AuthenticatedSession) =>
         [WS_METHODS.previewAutomationConnect]: (input) =>
           observeRpcStreamEffect(
             WS_METHODS.previewAutomationConnect,
-            previewAutomationBroker.connect(input.clientId),
+            previewAutomationBroker.connect(input),
             { "rpc.aggregate": "preview-automation" },
           ),
         [WS_METHODS.previewAutomationRespond]: (input) =>
@@ -1507,7 +1507,7 @@ const makeWsRpcLayer = (currentSession: AuthenticatedSession) =>
         [WS_METHODS.previewAutomationClearOwner]: (input) =>
           observeRpcEffect(
             WS_METHODS.previewAutomationClearOwner,
-            previewAutomationBroker.clearOwner(input.clientId),
+            previewAutomationBroker.clearOwner(input),
             { "rpc.aggregate": "preview-automation" },
           ),
         [WS_METHODS.subscribePreviewEvents]: (_input) =>

apps/web/src/components/preview/PreviewAutomationOwner.test.ts

@@ -3,11 +3,11 @@ import { describe, expect, it } from "vite-plus/test";
 import { observeAutomationOwnerConnectedGeneration } from "./PreviewAutomationOwner";
 
 describe("observeAutomationOwnerConnectedGeneration", () => {
-  it("re-reports ownership only after a later transport generation connects", () => {
+  it("reports ownership when the initial transport generation connects", () => {
     const initial = observeAutomationOwnerConnectedGeneration(null, 1);
     expect(initial).toEqual({
       nextGeneration: 1,
-      shouldReport: false,
+      shouldReport: true,
     });
 
     const disconnected = observeAutomationOwnerConnectedGeneration(initial.nextGeneration, null);