From f5cdeaa54ce5f98ec56ce9eee520d47e750cd65a Mon Sep 17 00:00:00 2001 From: Vaibhav Zope Date: Wed, 26 Aug 2026 22:43:00 +0530 Subject: [PATCH 1/2] Check a ref against the run this action is reaching, not the last one --- server/src/computer/gateway.ts | 64 +++++- server/src/computer/supervisor.ts | 33 +-- server/tests/computer-session-guard.test.ts | 218 ++++++++++++++++++++ 3 files changed, 292 insertions(+), 23 deletions(-) create mode 100644 server/tests/computer-session-guard.test.ts diff --git a/server/src/computer/gateway.ts b/server/src/computer/gateway.ts index 311c54ac..21c7d985 100644 --- a/server/src/computer/gateway.ts +++ b/server/src/computer/gateway.ts @@ -275,15 +275,23 @@ export function createComputerGateway( ); } + /** + * `address` is this action's own `/ensure`, when the caller has already made it. + * + * Locating twice for one action is two calls to the supervisor to learn the same thing, and worse, + * they can disagree: the run the action was decided against would not be the run it was sent to. + * Passing the address through means the check and the send are the same trip. + */ async function post( botId: string, path: string, payload: unknown, signal?: AbortSignal, timeoutMs?: number, + address?: string, ): Promise { return transport.post( - await locate(botId), + address ?? (await locate(botId)), botId, path, payload, @@ -292,6 +300,22 @@ export function createComputerGateway( ); } + /** + * This action's own `/ensure`, or nothing when it cannot be made. + * + * A computer that cannot be located is not a verdict this may reach on its own. The action still + * has to be decided and recorded, and the attempt failing is what writes the failure row beside the + * decision; throwing here would take the action off the trail entirely. So a failure answers + * "unknown", which leaves the generation check where it was and leaves the address to the attempt. + */ + async function locateForAction(botId: string): Promise { + try { + return await locate(botId); + } catch { + return undefined; + } + } + /* * The transport's deadline for a command, which is a backstop and not the limit. * @@ -421,13 +445,29 @@ export function createComputerGateway( /** The person's Stop, on its way to the browser. See the acting methods below. */ signal?: AbortSignal; }, - run: () => Promise, + run: (address?: string) => Promise, ): Promise { const { ref, filePath, snapshotId } = subject; // Loaded from the store, not this process's memory: the snapshot these refs belong to was very // likely taken by another replica, and resolving against a local map would find nothing there. const stored = await snapshots.load(botId); - const { session } = await sessionOf(botId); + /* + * LOCATE FIRST, THEN ASK WHICH RUN THAT WAS. The order is the check. + * + * `sessionOf` answers with what the last `/ensure` reported, and until this action has made its + * own, the last one belongs to the action before it. Asking first compared the stored snapshot + * against the previous action's run, which is the same run on every action but the first one + * after a replacement — exactly the action the check exists to catch. The click after a replaced + * container was allowed and the one after that refused, which is a guarantee arriving one action + * too late. + * + * Only for an action that cites a ref. Nothing else is resolved against a snapshot, so nothing + * else needs the run, and a scroll or a file read should not have to reach the supervisor before + * the policy has even seen it. The address that comes back is the one the attempt then uses, so + * this costs no extra call for the actions that do need it. + */ + const address = ref ? await locateForAction(botId) : undefined; + const { session } = ref ? await sessionOf(botId) : { session: undefined }; const element = resolve(stored, ref, snapshotId, session); // For a navigation the relevant page is the one being opened, not the one already loaded. Using // the stored URL would mean `page.host == "..."` could never match the destination, which is the @@ -532,7 +572,7 @@ export function createComputerGateway( `${ref} is not on the page this computer is showing, so nothing can be checked against it before acting. Take a fresh snapshot and use the refs it returns.`, ); } - result = await run(); + result = await run(address); } catch (error) { /** * A permitted action that did not happen gets its own row. @@ -782,7 +822,15 @@ export function createComputerGateway( snapshotId: input.snapshotId, ...(signal ? { signal } : {}), }, - () => post(botId, "/click", input, signal), + (address) => + post( + botId, + "/click", + input, + signal, + undefined, + address, + ), ); }, @@ -810,7 +858,8 @@ export function createComputerGateway( ...(input.submit ? { key: "Enter" } : {}), ...(signal ? { signal } : {}), }, - () => post(botId, "/type", input, signal), + (address) => + post(botId, "/type", input, signal, undefined, address), ); }, @@ -832,7 +881,8 @@ export function createComputerGateway( key: input.key, ...(signal ? { signal } : {}), }, - () => post(botId, "/key", input, signal), + (address) => + post(botId, "/key", input, signal, undefined, address), ); }, diff --git a/server/src/computer/supervisor.ts b/server/src/computer/supervisor.ts index 76ac0201..b632e9c2 100644 --- a/server/src/computer/supervisor.ts +++ b/server/src/computer/supervisor.ts @@ -16,22 +16,6 @@ import type { ComputerLocation, ComputerProvider } from "./provider"; import type { ComputerStatus } from "./schema"; -/** - * The last container start time seen for each Bot, from the `/ensure` that located it. - * - * Not a cache in front of the supervisor: `locate` still calls it every time. This only carries the - * answer the few lines to whoever needs to know which run of the computer they are talking to. - * - * PROCESS-LOCAL, AND THEREFORE NEVER THE ONLY ANSWER. On one replica the snapshot and the click that - * follows it are the same process, so this is always populated by the time anything asks. On several - * they are usually not, and a replica that has never located this Bot has nothing here. `resolve` - * reads an unknown session as "no opinion" and skips the generation check, so an empty map does not - * fail — it silently stops checking, on exactly the deployment shape the check was written for. So - * `sessionOf` falls back to asking, and this stays what it always was: a way to skip the round trip - * on the replica that just did the work. - */ -const sessions = new Map(); - type SupervisorComputerLocation = { botId: string; container?: string; @@ -67,6 +51,23 @@ export function createDockerSupervisorProvider( const hostForPort = options.hostForPort ?? ((port) => `http://localhost:${port}`); + /** + * The last container start time seen for each Bot, from the `/ensure` that located it. + * + * ONE MAP PER PROVIDER, not one per process. It was module-scope, so every provider built in a + * process shared it: two supervisors, or a test's second stack, answered each other's question + * about which run a Bot's computer is on, and the answer they gave was whichever one wrote last. + * + * Not a cache in front of the supervisor: `locate` still calls it every time, and a governed + * action locates before it asks. This only carries that answer the few lines to whoever needs to + * know which run of the computer they are talking to. + * + * Still process-local, and therefore never the only answer. A replica that has never located this + * Bot has nothing here, and an unknown run skips the generation check rather than failing it, so + * `sessionOf` falls back to asking rather than letting an empty map quietly stop checking. + */ + const sessions = new Map(); + async function call(path: string, method = "POST"): Promise { let response: Response; try { diff --git a/server/tests/computer-session-guard.test.ts b/server/tests/computer-session-guard.test.ts new file mode 100644 index 00000000..bbe9a741 --- /dev/null +++ b/server/tests/computer-session-guard.test.ts @@ -0,0 +1,218 @@ +import { describe, expect, test } from "bun:test"; +import type { AuditEventInput, AuditStore } from "../src/audit"; +import { StaleSnapshotError } from "../src/computer/client"; +import { + ActionRefusedError, + createComputerGateway, +} from "../src/computer/gateway"; +import type { ActionPolicy } from "../src/computer/policy"; +import type { SnapshotResult } from "../src/computer/schema"; +import { createInMemorySnapshotStore } from "../src/computer/snapshot-store"; +import { createDockerSupervisorProvider } from "../src/computer/supervisor"; + +/** + * A ref is only meaningful for the run of the computer that produced it. + * + * A replaced container counts its generations from one again, so a ref from the run before matches a + * stored row nothing has overwritten and resolves to an element on a page that is gone. The server + * compares the run to tell the two apart, and these are the cases where that comparison has to hold: + * the first action after a replacement, every action after that, and an action that landed on a + * replica which has never located this Bot itself. + */ + +const PERMISSIVE: ActionPolicy = { mode: "enforce", deny: [], allow: ["true"] }; +const ACTOR = { id: "dev-local-user" }; + +const FIRST_RUN = "2026-08-25T10:00:00.000Z"; +const SECOND_RUN = "2026-08-25T11:30:00.000Z"; + +const SNAPSHOT: SnapshotResult = { + snapshotId: 7, + url: "https://example.com/order", + title: "Order", + truncated: false, + elements: [ + { ref: "e1", role: "input", name: "Customer name:", type: "text" }, + { ref: "e9", role: "button", name: "Submit order" }, + ], +}; + +/** + * The supervisor, answering for whatever container exists now. + * + * Both endpoints report the same run, because a real supervisor does: `/ensure` and `/computers` + * read the same container. A stub that lists nothing would let these tests pass on an absence the + * deployment does not have, and `sessionOf` falls back to listing when it has not located a Bot + * itself. So every Bot named here is listed, whether or not this process ensured it. + */ +function fakeSupervisor(startedAt: () => string, known: string[] = []) { + const ensured = new Set(known); + const describe = (botId: string) => ({ + botId, + container: `openbot-computer-${botId}`, + status: "running", + url: "http://openbot-computer:4100", + startedAt: startedAt(), + }); + return (async (url: string) => { + const path = new URL(url).pathname; + if (path.endsWith("/ensure")) { + const asked = decodeURIComponent( + path.slice("/computers/".length, -"/ensure".length), + ); + ensured.add(asked); + return Response.json(describe(asked)); + } + return Response.json({ + computers: [...ensured].map((botId) => describe(botId)), + }); + }) as unknown as typeof fetch; +} + +function fakeComputerFetch() { + return (async (url: string) => { + const path = new URL(url).pathname; + if (path === "/snapshot") return Response.json(SNAPSHOT); + if (path === "/click") + return Response.json({ + action: "click", + url: SNAPSHOT.url, + elapsedMs: 1, + }); + return Response.json({ error: path }, { status: 404 }); + }) as unknown as typeof fetch; +} + +function fakeAudit() { + const rows: AuditEventInput[] = []; + const store: AuditStore = { insert: async (event) => void rows.push(event) }; + return { store, rows }; +} + +/** One server replica: its own provider, its own gateway, sharing only the snapshot store. */ +function stack( + startedAt: () => string, + policy: ActionPolicy = PERMISSIVE, + known: string[] = [], +) { + const snapshots = createInMemorySnapshotStore(); + const provider = createDockerSupervisorProvider({ + baseUrl: "http://supervisor:4300", + fetchImpl: fakeSupervisor(startedAt, known), + }); + const { store, rows } = fakeAudit(); + const gateway = createComputerGateway({ + provider, + fetchImpl: fakeComputerFetch(), + auditStore: store, + policy: () => policy, + snapshots, + }); + return { provider, gateway, rows, snapshots }; +} + +describe("a ref outliving the computer that produced it", () => { + test("the first action after the computer is replaced is refused, not the second", async () => { + let run = FIRST_RUN; + const { gateway, rows } = stack(() => run); + + const taken = await gateway.snapshot("bot-a"); + expect(taken.snapshotId).toBe(7); + + // The supervisor replaces the container because the image tag moved. Generations start from one + // again, nothing clears the stored snapshot, and the server is never told. + run = SECOND_RUN; + + const first = await gateway + .click("bot-a", ACTOR, { ref: "e9", snapshotId: taken.snapshotId }) + .catch((error: unknown) => error); + expect(first).toBeInstanceOf(StaleSnapshotError); + + // The audit row must not name a button on a page that is gone. + expect(JSON.stringify(rows.at(-1)?.payload.element ?? "")).not.toContain( + "Submit order", + ); + + const second = await gateway + .click("bot-a", ACTOR, { ref: "e9", snapshotId: taken.snapshotId }) + .catch((error: unknown) => error); + expect(second).toBeInstanceOf(StaleSnapshotError); + }); + + test("a replica that never located the Bot itself refuses too", async () => { + // This process has never called `/ensure` for bot-c, so it holds nothing for it. The snapshot row + // was written by the replica that took it and carries the run it belongs to. + const { gateway, snapshots } = stack(() => SECOND_RUN, PERMISSIVE, [ + "bot-c", + ]); + await snapshots.save("bot-c", { + snapshotId: 7, + url: SNAPSHOT.url, + elements: new Map(SNAPSHOT.elements.map((e) => [e.ref, e])), + session: FIRST_RUN, + }); + + const clicked = await gateway + .click("bot-c", ACTOR, { ref: "e9", snapshotId: 7 }) + .catch((error: unknown) => error); + expect(clicked).toBeInstanceOf(StaleSnapshotError); + }); + + test("a deny rule stops naming an element from a run that has ended", async () => { + // The sharper failure: with a rule keyed on the element's name, the server refused on the dead + // run's element and sent nothing, so the computer's own generation check never got a say. The + // refusal named a rule about a button the Bot never touched. + let run = FIRST_RUN; + const { gateway } = stack(() => run, { + mode: "enforce", + deny: ['contains(element.name, "Submit order")'], + allow: ["true"], + }); + + const taken = await gateway.snapshot("bot-d"); + run = SECOND_RUN; + + const clicked = await gateway + .click("bot-d", ACTOR, { ref: "e9", snapshotId: taken.snapshotId }) + .catch((error: unknown) => error); + expect(clicked).not.toBeInstanceOf(ActionRefusedError); + expect(clicked).toBeInstanceOf(StaleSnapshotError); + }); + + test("the run is read per provider, not once per process", async () => { + // The map was module-scope, so two providers in one process answered each other's question about + // which run a Bot is on, and the answer was whichever located last. + const replicaA = createDockerSupervisorProvider({ + baseUrl: "http://supervisor-a:4300", + fetchImpl: fakeSupervisor(() => FIRST_RUN, ["bot-e"]), + }); + const replicaB = createDockerSupervisorProvider({ + baseUrl: "http://supervisor-b:4300", + fetchImpl: fakeSupervisor(() => SECOND_RUN, ["bot-e"]), + }); + + await replicaA.locate("bot-e"); + expect(await replicaA.sessionOf?.("bot-e")).toBe(FIRST_RUN); + expect(await replicaB.sessionOf?.("bot-e")).toBe(SECOND_RUN); + }); + + test("CONTROL: a ref from the run that is still current still resolves", async () => { + // The same path with nothing replaced. A guard that refuses everything is not a guard, and this + // is what says the refusals above are about the run and not about the ordering itself. + const { gateway, rows } = stack(() => FIRST_RUN); + + const taken = await gateway.snapshot("bot-f"); + const clicked = await gateway.click("bot-f", ACTOR, { + ref: "e9", + snapshotId: taken.snapshotId, + }); + + expect(clicked).toMatchObject({ + action: "click", + element: { role: "button", name: "Submit order" }, + }); + expect(rows.at(-1)?.payload.element).toMatchObject({ + name: "Submit order", + }); + }); +}); From ab3f330a3eefbf67ee197851e9ed477ab15b4984 Mon Sep 17 00:00:00 2001 From: Vaibhav Zope Date: Thu, 27 Aug 2026 11:05:50 +0530 Subject: [PATCH 2/2] Prove the run check across two replicas and a real snapshot row --- ...session-guard-replicas.integration.test.ts | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 server/tests/computer-session-guard-replicas.integration.test.ts diff --git a/server/tests/computer-session-guard-replicas.integration.test.ts b/server/tests/computer-session-guard-replicas.integration.test.ts new file mode 100644 index 00000000..207dbeba --- /dev/null +++ b/server/tests/computer-session-guard-replicas.integration.test.ts @@ -0,0 +1,148 @@ +import { afterAll, beforeAll, expect, test } from "bun:test"; +import { randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import type { AuditEventInput, AuditStore } from "../src/audit"; +import { StaleSnapshotError } from "../src/computer/client"; +import { createComputerGateway } from "../src/computer/gateway"; +import type { ActionPolicy } from "../src/computer/policy"; +import type { SnapshotResult } from "../src/computer/schema"; +import { createSnapshotStore } from "../src/computer/snapshot-store"; +import { createDockerSupervisorProvider } from "../src/computer/supervisor"; +import { createDatabase } from "../src/db/client"; +import { computerSnapshot } from "../src/db/schema"; +import { TEST_POOL } from "./support/database"; + +/** + * Two replicas, one Postgres, a container replaced between them. + * + * The in-process guard tests share a Map-free provider but an in-memory snapshot store. This is the + * shape the bug is actually about: the snapshot crosses processes through Postgres, and the run it + * belongs to has to cross with it. Replica A takes the snapshot; replica B, which has never located + * this Bot, gets the click after the container was replaced. + */ + +const database = createDatabase( + process.env.DATABASE_URL ?? + "postgres://openbot:openbot@localhost:5432/openbot", + TEST_POOL, +); + +const suite = randomUUID().slice(0, 8); +const botId = `agent_replica_${suite}`; + +const FIRST_RUN = "2026-08-25T10:00:00.000Z"; +const SECOND_RUN = "2026-08-25T11:30:00.000Z"; +const PERMISSIVE: ActionPolicy = { mode: "enforce", deny: [], allow: ["true"] }; +const ACTOR = { id: "dev-local-user" }; + +const SNAPSHOT: SnapshotResult = { + snapshotId: 7, + url: "https://example.com/order", + title: "Order", + truncated: false, + elements: [{ ref: "e9", role: "button", name: "Submit order" }], +}; + +function fakeSupervisor(startedAt: () => string, known: string[] = []) { + const ensured = new Set(known); + const describe = (id: string) => ({ + botId: id, + container: `openbot-computer-${id}`, + status: "running", + url: "http://openbot-computer:4100", + startedAt: startedAt(), + }); + return (async (url: string) => { + const path = new URL(url).pathname; + if (path.endsWith("/ensure")) { + const asked = decodeURIComponent( + path.slice("/computers/".length, -"/ensure".length), + ); + ensured.add(asked); + return Response.json(describe(asked)); + } + return Response.json({ + computers: [...ensured].map((id) => describe(id)), + }); + }) as unknown as typeof fetch; +} + +function fakeComputerFetch() { + return (async (url: string) => { + const path = new URL(url).pathname; + if (path === "/snapshot") return Response.json(SNAPSHOT); + if (path === "/click") + return Response.json({ + action: "click", + url: SNAPSHOT.url, + elapsedMs: 1, + }); + return Response.json({ error: path }, { status: 404 }); + }) as unknown as typeof fetch; +} + +/** One server replica: its own provider and gateway, sharing only Postgres. */ +function replica(startedAt: () => string, known: string[] = []) { + const rows: AuditEventInput[] = []; + const store: AuditStore = { insert: async (e) => void rows.push(e) }; + const provider = createDockerSupervisorProvider({ + baseUrl: "http://supervisor:4300", + fetchImpl: fakeSupervisor(startedAt, known), + }); + const gateway = createComputerGateway({ + provider, + fetchImpl: fakeComputerFetch(), + auditStore: store, + policy: () => PERMISSIVE, + // The real store, against the real database. + snapshots: createSnapshotStore(database), + }); + return { gateway, rows }; +} + +beforeAll(async () => { + await database + .delete(computerSnapshot) + .where(eq(computerSnapshot.computerId, botId)); +}); + +afterAll(async () => { + await database + .delete(computerSnapshot) + .where(eq(computerSnapshot.computerId, botId)); +}); + +test("replica B refuses a ref that replica A took before the container was replaced", async () => { + // Replica A: takes the snapshot during the first run. Writes it to Postgres. + const a = replica(() => FIRST_RUN); + await a.gateway.snapshot(botId); + + // The row really is in Postgres, carrying the run it was taken on. + const [stored] = await database + .select({ session: computerSnapshot.session }) + .from(computerSnapshot) + .where(eq(computerSnapshot.computerId, botId)); + expect(stored?.session).toBe(FIRST_RUN); + + // Replica B: a different process. Never located this Bot, so it holds nothing for it. The + // container has since been replaced, so the supervisor now reports the second run. + const b = replica(() => SECOND_RUN, [botId]); + const clicked = await b.gateway + .click(botId, ACTOR, { ref: "e9", snapshotId: 7 }) + .catch((error: unknown) => error); + + expect(clicked).toBeInstanceOf(StaleSnapshotError); +}); + +test("CONTROL: replica B allows the same ref while the run is still current", async () => { + const a = replica(() => FIRST_RUN); + await a.gateway.snapshot(botId); + + // Same cold replica, but the container was never replaced. + const b = replica(() => FIRST_RUN, [botId]); + const result = await b.gateway.click(botId, ACTOR, { + ref: "e9", + snapshotId: 7, + }); + expect(result).toMatchObject({ action: "click" }); +});