@vellumai/assistant 0.12.2-staging.2 → 0.12.2-staging.4

package/src/desktop/desktop-automation-lease.ts

@@ -0,0 +1,270 @@
+import { getConfig } from "../config/loader.js";
+import type { ToolContext, ToolExecutionResult } from "../tools/types.js";
+import { getLogger } from "../util/logger.js";
+import { desktopDependencyInstaller } from "./desktop-dependencies.js";
+import {
+  type DesktopSessionManager,
+  type DesktopViewer,
+  getDesktopSessionManager,
+} from "./desktop-session-manager.js";
+import { isVirtualDesktopEnabled } from "./virtual-desktop-feature.js";
+
+const log = getLogger("desktop-automation-lease");
+const IDLE_TIMEOUT_MS = 5 * 60_000;
+const MAX_ACTIONS = 100;
+type Owner = {
+  conversationId: string;
+  actorId: string;
+  abort: AbortController;
+  holder: DesktopViewer;
+  removeAbortListener: () => void;
+  actions: number;
+  lastActivity: number;
+  desktopLost: boolean;
+};
+
+export class DesktopAutomationLease {
+  private owner: Owner | null = null;
+  private generation = 0;
+  private tail: Promise<unknown> = Promise.resolve();
+  private watchdog: ReturnType<typeof setInterval> | undefined;
+
+  constructor(
+    private readonly deps: {
+      enabled: () => boolean;
+      ready: () => boolean;
+      ensureReady: (signal: AbortSignal) => Promise<void>;
+      manager: () => DesktopSessionManager;
+    } = {
+      enabled: () => isVirtualDesktopEnabled(getConfig()),
+      ready: () => desktopDependencyInstaller.getStatus().state === "ready",
+      ensureReady: (signal) => desktopDependencyInstaller.ensureReady(signal),
+      manager: getDesktopSessionManager,
+    },
+  ) {}
+
+  private exclusive<T>(run: () => Promise<T>): Promise<T> {
+    const next = this.tail.then(run, run);
+    this.tail = next.catch(() => {});
+    return next;
+  }
+
+  private assertAvailable(): void {
+    if (!this.deps.enabled()) {
+      throw new Error(
+        "Virtual desktop browser automation is available only on enabled platform-hosted assistants",
+      );
+    }
+    if (!this.deps.ready()) {
+      throw new Error(
+        "Open the Virtual desktop panel and wait for automatic installation to finish before using desktop browser automation",
+      );
+    }
+  }
+
+  private async release(): Promise<void> {
+    const owner = this.owner;
+    if (!owner) {
+      return;
+    }
+    owner.abort.abort();
+    owner.removeAbortListener();
+    clearInterval(this.watchdog);
+    this.watchdog = undefined;
+    if (!owner.desktopLost) {
+      await this.deps.manager().browser.release();
+    }
+    this.deps.manager().releaseAutomationSlot(owner.holder);
+    this.owner = null;
+  }
+
+  private cancel(owner: Owner): void {
+    if (this.owner !== owner) {
+      return;
+    }
+    this.generation += 1;
+    owner.abort.abort();
+    void this.exclusive(() => this.release()).catch((err) =>
+      log.warn({ err }, "Desktop browser session cleanup failed"),
+    );
+  }
+
+  private bindCancellation(owner: Owner, signal?: AbortSignal): void {
+    owner.removeAbortListener();
+    const cancel = () => this.cancel(owner);
+    signal?.addEventListener("abort", cancel, { once: true });
+    owner.removeAbortListener = () =>
+      signal?.removeEventListener("abort", cancel);
+    if (signal?.aborted) {
+      cancel();
+    }
+  }
+
+  private async claim(context: ToolContext): Promise<Owner> {
+    const holder: DesktopViewer = {
+      onDesktopLost: () => {
+        owner.desktopLost = true;
+        this.cancel(owner);
+      },
+    };
+    const owner: Owner = {
+      conversationId: context.conversationId,
+      actorId: context.sourceActorPrincipalId!,
+      abort: new AbortController(),
+      holder,
+      actions: 0,
+      lastActivity: Date.now(),
+      desktopLost: false,
+      removeAbortListener: () => {},
+    };
+    const slot = this.deps.manager().acquireAutomationSlot(holder);
+    if (!slot.ok) {
+      throw new Error("The desktop is busy or shutting down");
+    }
+    this.owner = owner;
+    this.bindCancellation(owner, context.signal);
+    const cancel = () => this.cancel(owner);
+    this.watchdog = setInterval(() => {
+      try {
+        if (
+          Date.now() - owner.lastActivity > IDLE_TIMEOUT_MS ||
+          !this.deps.enabled() ||
+          !this.deps.ready()
+        ) {
+          cancel();
+        }
+      } catch (err) {
+        log.warn({ err }, "Desktop browser session availability check failed");
+        cancel();
+      }
+    }, 1_000);
+    this.watchdog.unref?.();
+    try {
+      await this.deps.manager().ensureDesktopRunning();
+      owner.abort.signal.throwIfAborted();
+      this.assertAvailable();
+      await this.deps.manager().browser.release();
+      return owner;
+    } catch (err) {
+      await this.release().catch((cleanupError) =>
+        log.warn(
+          { err: cleanupError },
+          "Desktop browser session cleanup failed",
+        ),
+      );
+      throw err;
+    }
+  }
+
+  runBrowser(
+    context: ToolContext,
+    operation: (signal: AbortSignal) => Promise<ToolExecutionResult>,
+    done = false,
+  ): Promise<ToolExecutionResult> {
+    const generation = this.generation;
+    return this.exclusive(async () => {
+      if (
+        context.trustClass !== "guardian" ||
+        !context.sourceActorPrincipalId ||
+        !context.conversationId
+      ) {
+        throw new Error(
+          "Desktop browser session requires an identified guardian conversation",
+        );
+      }
+      if (
+        this.owner &&
+        (this.owner.conversationId !== context.conversationId ||
+          this.owner.actorId !== context.sourceActorPrincipalId)
+      ) {
+        throw new Error("Another conversation is controlling the desktop");
+      }
+      if (done) {
+        await this.release();
+        return { content: "Desktop browser session released.", isError: false };
+      }
+      context.signal?.throwIfAborted();
+      try {
+        if (
+          !this.owner &&
+          generation === this.generation &&
+          this.deps.enabled() &&
+          !this.deps.ready()
+        ) {
+          const abort = new AbortController();
+          const signal = AbortSignal.any([
+            abort.signal,
+            ...(context.signal ? [context.signal] : []),
+          ]);
+          const timeout = setTimeout(
+            () =>
+              abort.abort(
+                new Error(
+                  "Virtual desktop setup is taking too long. Check installation progress in the Virtual desktop panel. No browser action was performed.",
+                ),
+              ),
+            8 * 60_000,
+          );
+          const watchdog = setInterval(() => {
+            if (!this.deps.enabled()) {
+              abort.abort(
+                new Error("Virtual desktop was disabled during setup"),
+              );
+            }
+          }, 1_000);
+          try {
+            await this.deps.ensureReady(signal);
+            signal.throwIfAborted();
+          } finally {
+            clearTimeout(timeout);
+            clearInterval(watchdog);
+          }
+        }
+        this.assertAvailable();
+      } catch (error) {
+        await this.release();
+        throw error;
+      }
+      if (generation !== this.generation) {
+        return {
+          content:
+            "Desktop browser session was interrupted. Take a fresh snapshot before continuing.",
+          isError: true,
+        };
+      }
+      if (this.owner) {
+        this.bindCancellation(this.owner, context.signal);
+      }
+      const owner = this.owner ?? (await this.claim(context));
+      const signal = context.signal
+        ? AbortSignal.any([context.signal, owner.abort.signal])
+        : owner.abort.signal;
+      owner.lastActivity = Date.now();
+      signal.throwIfAborted();
+      try {
+        this.assertAvailable();
+        if (++owner.actions > MAX_ACTIONS) {
+          throw new Error(
+            "Desktop action limit reached. Finish this session before continuing.",
+          );
+        }
+        const result = await operation(signal);
+        signal.throwIfAborted();
+        if (result.isError) {
+          await this.release();
+        }
+        return result;
+      } catch (err) {
+        await this.release().catch((cleanupError) =>
+          log.warn(
+            { err: cleanupError },
+            "Desktop browser session cleanup failed",
+          ),
+        );
+        throw err;
+      }
+    });
+  }
+}
+
+export const desktopAutomationLease = new DesktopAutomationLease();

package/src/desktop/desktop-browser-client.test.ts

@@ -0,0 +1,413 @@
+import { afterEach, expect, test } from "bun:test";
+
+import { waitFor } from "../__tests__/helpers/wait-for.js";
+import { executeBrowserOperation } from "../browser/operations.js";
+import { browserManager } from "../tools/browser/browser-manager.js";
+import {
+  type CdpTransportEvent,
+  type CdpWsTransport,
+  CdpWsTransportError,
+} from "../tools/browser/cdp-client/cdp-inspect/ws-transport.js";
+import { DesktopBrowserClient } from "./desktop-browser-client.js";
+
+function fixture() {
+  const calls: {
+    connection: number;
+    method: string;
+    params: Record<string, unknown>;
+    session?: string;
+  }[] = [];
+  const targets = [
+    {
+      targetId: "page-1",
+      type: "page",
+      url: "https://example.com",
+      title: "Example",
+    },
+  ];
+  const listeners: ((event: CdpTransportEvent) => void)[] = [];
+  let connections = 0;
+  let disconnect = () => {};
+  let reject:
+    | ((
+        method: string,
+        params: Record<string, unknown>,
+      ) => boolean | Promise<boolean>)
+    | undefined;
+  const browser = new DesktopBrowserClient(async () => {
+    const connection = ++connections;
+    let closed = false;
+    disconnect = () => {
+      closed = true;
+    };
+    return {
+      get closed() {
+        return closed;
+      },
+      async send<T>(
+        method: string,
+        params: Record<string, unknown> = {},
+        options?: { sessionId?: string; signal?: AbortSignal },
+      ): Promise<T> {
+        options?.signal?.throwIfAborted();
+        if (closed) {
+          throw new CdpWsTransportError("closed");
+        }
+        calls.push({ connection, method, params, session: options?.sessionId });
+        if (await reject?.(method, params)) {
+          throw new CdpWsTransportError("closed");
+        }
+        let result: unknown = {};
+        if (method === "Target.getTargets") {
+          result = { targetInfos: targets };
+        }
+        if (method === "Target.attachToTarget") {
+          result = { sessionId: `session-${connection}-${params.targetId}` };
+        }
+        if (method === "Target.createTarget") {
+          const targetId = `page-${targets.length + 1}`;
+          targets.push({
+            targetId,
+            type: "page",
+            url: "about:blank",
+            title: "New",
+          });
+          result = { targetId };
+        }
+        if (method === "Runtime.evaluate") {
+          result = { result: { value: "Example" } };
+        }
+        if (method === "DOM.getBoxModel") {
+          result = { model: { content: [10, 10, 110, 10, 110, 50, 10, 50] } };
+        }
+        if (method === "Accessibility.getFullAXTree") {
+          result = {
+            nodes: [
+              {
+                nodeId: "1",
+                backendDOMNodeId: 42,
+                role: { value: "button" },
+                name: { value: "Save" },
+                properties: [],
+              },
+            ],
+          };
+        }
+        return result as T;
+      },
+      addEventListener(listener) {
+        listeners.push(listener);
+        return () => {};
+      },
+      dispose() {
+        closed = true;
+      },
+    } satisfies CdpWsTransport;
+  });
+  return {
+    browser,
+    calls,
+    targets,
+    disconnect: () => disconnect(),
+    emit: (method: string, params = {}, sessionId?: string) =>
+      listeners.forEach((listener) => listener({ method, params, sessionId })),
+    fail: (predicate?: typeof reject) => {
+      reject = predicate;
+    },
+  };
+}
+const cleanups: (() => void)[] = [];
+afterEach(() => {
+  for (const cleanup of cleanups.splice(0)) {
+    cleanup();
+  }
+});
+
+async function session() {
+  const f = fixture();
+  cleanups.push(() => f.browser.dispose());
+  const abort = new AbortController();
+  const cdp = await f.browser.client("conv-123", abort.signal);
+  return { ...f, cdp, abort };
+}
+
+test("shared snapshot and click use desktop element IDs without personal browser state", async () => {
+  const f = await session();
+  browserManager.storeSnapshotBackendNodeMap(
+    "conv-123",
+    new Map([["e1", 999]]),
+  );
+  cleanups.push(() => browserManager.clearSnapshotBackendNodeMap("conv-123"));
+  const context = {
+    workingDir: "/tmp",
+    trustClass: "guardian" as const,
+    conversationId: "conv-123",
+    cdpClient: f.cdp,
+    signal: f.abort.signal,
+  };
+  const snapshot = await executeBrowserOperation("snapshot", {}, context);
+  expect(snapshot.isError).toBe(false);
+  expect(snapshot.content).toContain("Save");
+  const element = /\b(e\d+)\b/.exec(snapshot.content)?.[1];
+  expect(element).toBeDefined();
+  const result = await executeBrowserOperation(
+    "click",
+    { element_id: element },
+    context,
+  );
+  expect(result.isError).toBe(false);
+  expect(
+    f.calls.find((call) => call.method === "DOM.getBoxModel")?.params
+      .backendNodeId,
+  ).toBe(42);
+  expect(browserManager.resolveSnapshotBackendNodeId("conv-123", "e1")).toBe(
+    999,
+  );
+  const pressIndex = f.calls.findIndex(
+    (call) => call.params.type === "mousePressed",
+  );
+  expect(pressIndex).toBeGreaterThan(0);
+  expect(f.calls[pressIndex - 1]?.params.expression).toContain(
+    "data-vellum-desktop-cursor",
+  );
+  expect(
+    f.calls.filter((call) => call.params.type === "mousePressed"),
+  ).toHaveLength(1);
+  f.emit("DOM.documentUpdated");
+  const stale = await executeBrowserOperation(
+    "click",
+    { element_id: element },
+    context,
+  );
+  expect(stale.isError).toBe(true);
+  expect(
+    f.calls.filter((call) => call.params.type === "mousePressed"),
+  ).toHaveLength(1);
+});
+
+test("release invalidates borrowed clients and releases uncertain input through a fresh connection", async () => {
+  const f = await session();
+  f.fail(
+    (method, params) =>
+      method === "Input.dispatchKeyEvent" && params.type === "keyDown",
+  );
+  await expect(
+    f.cdp.send("Input.dispatchKeyEvent", {
+      type: "keyDown",
+      key: "Shift",
+      code: "ShiftLeft",
+    }),
+  ).rejects.toThrow("do not retry");
+  f.abort.abort();
+  f.fail();
+  await f.browser.release();
+  const releases = f.calls.filter((call) => call.params.type === "keyUp");
+  expect(releases).toHaveLength(1);
+  expect(releases[0]?.connection).toBe(2);
+  expect(releases[0]?.session).toBe("session-2-page-1");
+  expect(f.calls.filter((call) => call.params.type === "keyDown")).toHaveLength(
+    1,
+  );
+  await expect(f.cdp.listTabs()).rejects.toBeDefined();
+  await expect(
+    f.cdp.send("Input.insertText", { text: "late" }),
+  ).rejects.toBeDefined();
+});
+
+test("failed cleanup retains held input for a later attempt", async () => {
+  const f = await session();
+  await f.cdp.send("Input.dispatchMouseEvent", {
+    type: "mousePressed",
+    button: "left",
+    x: 10,
+    y: 20,
+  });
+  f.fail((_method, params) => params.type === "mouseReleased");
+  await expect(f.browser.release()).rejects.toBeDefined();
+  f.fail();
+  await f.browser.release();
+  expect(
+    f.calls
+      .filter((call) => call.params.type === "mouseReleased")
+      .map((call) => call.connection),
+  ).toEqual([2, 3]);
+});
+
+test("detachment preserves uncertain key state until target cleanup", async () => {
+  const f = await session();
+  await f.cdp.send("Input.dispatchKeyEvent", {
+    type: "rawKeyDown",
+    key: "Alt",
+    code: "AltLeft",
+  });
+  f.emit("Target.detachedFromTarget", { sessionId: "session-1-page-1" });
+  await f.browser.release();
+  expect(f.calls.filter((call) => call.params.type === "keyUp")).toHaveLength(
+    1,
+  );
+});
+
+test("tab changes clear element references and aliases are not reused across leases", async () => {
+  const f = await session();
+  const first = (await f.cdp.listTabs())[0]!.tabId!;
+  browserManager.storeSnapshotBackendNodeMap(
+    f.cdp.conversationId,
+    new Map([["e1", 42]]),
+  );
+  const created = await f.cdp.send<{ tabId: number }>("Vellum.createTab");
+  expect(
+    browserManager.resolveSnapshotBackendNodeId(f.cdp.conversationId, "e1"),
+  ).toBeNull();
+  await f.cdp.selectTab(first);
+  expect((await f.cdp.listTabs()).find((tab) => tab.active)?.tabId).toBe(first);
+  f.targets.splice(0, 1);
+  f.emit("Target.detachedFromTarget", { sessionId: "session-1-page-1" });
+  await expect(
+    f.cdp.send("Input.insertText", { text: "wrong page" }),
+  ).rejects.toThrow("tab closed");
+  expect(f.calls.some((call) => call.params.text === "wrong page")).toBe(false);
+  await f.browser.release();
+  const next = await f.browser.client("conv-456", new AbortController().signal);
+  expect(
+    (await next.listTabs()).every((tab) => tab.tabId! > created.tabId),
+  ).toBe(true);
+});
+
+test("navigation during cursor animation prevents input on the replacement page", async () => {
+  const f = await session();
+  f.fail((method, params) => {
+    if (
+      method === "Runtime.evaluate" &&
+      String(params.expression).includes("data-vellum-desktop-cursor")
+    ) {
+      f.emit("Page.frameNavigated");
+    }
+    return false;
+  });
+  await expect(
+    f.cdp.send("Input.dispatchMouseEvent", {
+      type: "mousePressed",
+      x: 50,
+      y: 60,
+      button: "left",
+    }),
+  ).rejects.toThrow("page changed");
+  expect(f.calls.some((call) => call.params.type === "mousePressed")).toBe(
+    false,
+  );
+});
+
+test("navigation restores the cursor in its tab without replaying mouse input", async () => {
+  const f = await session();
+  await f.cdp.send("Input.dispatchMouseEvent", {
+    type: "mouseMoved",
+    x: 50,
+    y: 60,
+  });
+  f.calls.length = 0;
+  f.emit("Page.domContentEventFired", {}, "session-1-page-2");
+  f.emit("Page.domContentEventFired", {}, "session-1-page-1");
+  await Bun.sleep(0);
+  expect(f.calls).toHaveLength(1);
+  expect(f.calls[0]).toMatchObject({
+    method: "Runtime.evaluate",
+    session: "session-1-page-1",
+  });
+  expect(f.calls[0]?.params.expression).toContain("translate(50px, 60px)");
+});
+
+test("keyboard input restores an existing cursor without synthesizing mouse input", async () => {
+  const f = await session();
+  await f.cdp.send("Input.dispatchMouseEvent", {
+    type: "mouseMoved",
+    x: 50,
+    y: 60,
+  });
+  f.calls.length = 0;
+  await f.cdp.send("Input.dispatchKeyEvent", { type: "keyUp", key: "Enter" });
+  expect(f.calls.map((call) => call.method)).toEqual([
+    "Input.dispatchKeyEvent",
+    "Runtime.evaluate",
+  ]);
+  expect(f.calls[1]?.params.expression).toContain("translate(50px, 60px)");
+});
+
+test("release suppresses pending cursor restoration and clears it before detaching", async () => {
+  const f = await session();
+  await f.cdp.send("Input.dispatchMouseEvent", {
+    type: "mouseMoved",
+    x: 50,
+    y: 60,
+  });
+  f.calls.length = 0;
+  f.emit("Page.domContentEventFired", {}, "session-1-page-1");
+  await f.browser.release();
+  f.emit("Page.domContentEventFired", {}, "session-1-page-1");
+  await Bun.sleep(0);
+  const expressions = f.calls.filter(
+    (call) => call.method === "Runtime.evaluate",
+  );
+  expect(expressions).toHaveLength(1);
+  expect(expressions[0]?.params.expression).toContain("?.remove()");
+});
+
+test("release waits for an in-flight cursor restoration before removing it", async () => {
+  const f = await session();
+  await f.cdp.send("Input.dispatchMouseEvent", {
+    type: "mouseMoved",
+    x: 50,
+    y: 60,
+  });
+  let finish!: (reject: boolean) => void;
+  const restoring = new Promise<boolean>((resolve) => {
+    finish = resolve;
+  });
+  let started = false;
+  f.fail((method, params) => {
+    if (
+      method === "Runtime.evaluate" &&
+      String(params.expression).includes("translate(")
+    ) {
+      started = true;
+      return restoring;
+    }
+    return false;
+  });
+  f.emit("Page.domContentEventFired", {}, "session-1-page-1");
+  await waitFor(() => started);
+  const released = f.browser.release();
+  try {
+    await Bun.sleep(0);
+    expect(f.calls.some((call) => call.connection === 2)).toBe(false);
+  } finally {
+    finish(false);
+    await released;
+  }
+  expect(f.calls.at(-1)?.params.expression).toContain("?.remove()");
+});
+
+test("closed connections reconnect after cleanup without retrying input or reusing old clients", async () => {
+  const f = await session();
+  await f.cdp.send("Input.dispatchKeyEvent", {
+    type: "keyDown",
+    key: "Shift",
+    code: "ShiftLeft",
+  });
+  let previous = f.cdp;
+  for (let round = 0; round < 2; round++) {
+    f.disconnect();
+    const current = await f.browser.client("conv-123", f.abort.signal);
+    await expect(previous.listTabs()).rejects.toThrow("session expired");
+    await current.send("Runtime.evaluate", { expression: "document.title" });
+    expect(f.calls.at(-1)?.connection).toBe(3 + round * 2);
+    previous = current;
+  }
+  expect(f.calls.filter((call) => call.params.type === "keyDown")).toHaveLength(
+    1,
+  );
+  expect(
+    f.calls
+      .filter((call) => call.params.type === "keyUp")
+      .map((call) => call.connection),
+  ).toEqual([2]);
+});