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

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

@@ -0,0 +1,208 @@
+import { afterEach, expect, mock, test } from "bun:test";
+
+import { DesktopAutomationLease } from "./desktop-automation-lease.js";
+import { DesktopDependencyInstaller } from "./desktop-dependencies.js";
+import type { DesktopSessionManager } from "./desktop-session-manager.js";
+
+const context = {
+  conversationId: "conv-123",
+  sourceActorPrincipalId: "user-123",
+  trustClass: "guardian" as const,
+  workingDir: "/tmp",
+};
+const cleanups: DesktopAutomationLease[] = [];
+afterEach(async () => {
+  await Promise.all(
+    cleanups
+      .splice(0)
+      .map((lease) => lease.runBrowser(context, operation, true)),
+  );
+});
+function fixture() {
+  let enabled = true;
+  let ready = true;
+  const setup = Promise.withResolvers<void>();
+  const install = mock(() => setup.promise);
+  const installer = new DesktopDependencyInstaller({
+    supported: () => true,
+    ready: () => ready,
+    install,
+    notify: async () => {},
+  });
+  const ensureReady = mock((signal: AbortSignal) =>
+    installer.ensureReady(signal),
+  );
+  const started = mock(async () => {});
+  const released = mock(() => {});
+  const releaseBrowser = mock(async () => {});
+  const lease = new DesktopAutomationLease({
+    enabled: () => enabled,
+    ready: () => ready,
+    ensureReady,
+    manager: () =>
+      ({
+        browser: { release: releaseBrowser },
+        acquireAutomationSlot: () => ({ ok: true }),
+        releaseAutomationSlot: released,
+        ensureDesktopRunning: started,
+      }) as unknown as DesktopSessionManager,
+  });
+  cleanups.push(lease);
+  return {
+    lease,
+    ensureReady,
+    install,
+    failSetup: () => setup.reject(new Error("download failed")),
+    completeSetup: () => {
+      ready = true;
+      setup.resolve();
+    },
+    started,
+    released,
+    releaseBrowser,
+    disable: () => {
+      enabled = false;
+    },
+    uninstall: () => {
+      ready = false;
+    },
+  };
+}
+const operation = async () => ({ content: "ok", isError: false });
+
+test("browser session reuses one automation slot and isolates conversations", async () => {
+  const f = fixture();
+  await f.lease.runBrowser(context, operation);
+  await f.lease.runBrowser(context, operation);
+  expect(f.started).toHaveBeenCalledTimes(1);
+  await expect(
+    f.lease.runBrowser({ ...context, conversationId: "conv-456" }, operation),
+  ).rejects.toThrow("Another conversation");
+  expect(f.released).not.toHaveBeenCalled();
+  await f.lease.runBrowser(context, operation, true);
+  expect(f.released).toHaveBeenCalledTimes(1);
+});
+
+for (const lose of ["disable", "uninstall"] as const) {
+  test(`${lose} blocks startup, reuse and startup races`, async () => {
+    const unavailable = fixture();
+    unavailable[lose]();
+    if (lose === "uninstall") {
+      unavailable.ensureReady.mockRejectedValueOnce(new Error("setup failed"));
+    }
+    await expect(
+      unavailable.lease.runBrowser(context, operation),
+    ).rejects.toThrow();
+    expect(unavailable.started).not.toHaveBeenCalled();
+    const race = fixture();
+    race.started.mockImplementationOnce(async () => race[lose]());
+    const callback = mock(operation);
+    await expect(race.lease.runBrowser(context, callback)).rejects.toThrow();
+    expect(callback).not.toHaveBeenCalled();
+    expect(race.released).toHaveBeenCalledTimes(1);
+    const reused = fixture();
+    await reused.lease.runBrowser(context, operation);
+    reused[lose]();
+    await expect(reused.lease.runBrowser(context, callback)).rejects.toThrow();
+    expect(callback).not.toHaveBeenCalled();
+    expect(reused.released).toHaveBeenCalledTimes(1);
+  });
+}
+
+test("cancellation stops running and queued browser commands without blocking a new session", async () => {
+  const f = fixture();
+  const started = Promise.withResolvers<void>();
+  const abort = new AbortController();
+  const running = f.lease
+    .runBrowser({ ...context, signal: abort.signal }, async (signal) => {
+      started.resolve();
+      await new Promise<void>((_resolve, reject) => {
+        signal.addEventListener("abort", () => reject(new Error("aborted")), {
+          once: true,
+        });
+      });
+      return operation();
+    })
+    .catch((error: Error) => error);
+  await started.promise;
+  const callback = mock(operation);
+  const queued = f.lease.runBrowser(context, callback);
+  abort.abort();
+  expect(await running).toBeInstanceOf(Error);
+  expect((await queued).isError).toBe(true);
+  expect(callback).not.toHaveBeenCalled();
+  await f.lease.runBrowser(context, callback);
+  expect(callback).toHaveBeenCalledTimes(1);
+});
+
+test("failed browser cleanup retains ownership until release succeeds", async () => {
+  const f = fixture();
+  await f.lease.runBrowser(context, operation);
+  f.releaseBrowser.mockRejectedValueOnce(new Error("Chrome busy"));
+  await expect(f.lease.runBrowser(context, operation, true)).rejects.toThrow(
+    "Chrome busy",
+  );
+  expect(f.released).not.toHaveBeenCalled();
+  await f.lease.runBrowser(context, operation, true);
+  expect(f.released).toHaveBeenCalledTimes(1);
+});
+
+test("first browser call waits for one shared install then executes exactly once", async () => {
+  const f = fixture();
+  f.uninstall();
+  const callback = mock(operation);
+  const first = f.lease.runBrowser(context, callback);
+  await Bun.sleep(0);
+  expect(f.install).toHaveBeenCalledTimes(1);
+  expect(f.started).not.toHaveBeenCalled();
+  expect(callback).not.toHaveBeenCalled();
+  f.completeSetup();
+  expect((await first).isError).toBe(false);
+  expect(callback).toHaveBeenCalledTimes(1);
+  await f.lease.runBrowser(context, callback);
+  expect(f.install).toHaveBeenCalledTimes(1);
+  expect(f.started).toHaveBeenCalledTimes(1);
+});
+
+for (const interrupt of ["cancel", "disable", "failure"] as const) {
+  test(`${interrupt} during setup prevents a delayed browser action`, async () => {
+    const f = fixture();
+    f.uninstall();
+    const abort = new AbortController();
+    const callback = mock(operation);
+    const result = f.lease
+      .runBrowser({ ...context, signal: abort.signal }, callback)
+      .catch((error: unknown) => error);
+    await Bun.sleep(0);
+    if (interrupt === "cancel") {
+      abort.abort();
+    } else if (interrupt === "disable") {
+      f.disable();
+    } else {
+      f.failSetup();
+    }
+    expect(await result).toBeInstanceOf(Error);
+    f.completeSetup();
+    await Bun.sleep(0);
+    expect(f.started).not.toHaveBeenCalled();
+    expect(callback).not.toHaveBeenCalled();
+  });
+}
+
+test("disabled, unidentified, cancelled and released browser calls cannot install", async () => {
+  const f = fixture();
+  f.uninstall();
+  const cancelled = new AbortController();
+  cancelled.abort();
+  for (const caller of [
+    { ...context, trustClass: "unknown" as const },
+    { ...context, sourceActorPrincipalId: undefined },
+    { ...context, signal: cancelled.signal },
+  ]) {
+    await expect(f.lease.runBrowser(caller, operation)).rejects.toThrow();
+  }
+  await f.lease.runBrowser(context, operation, true);
+  f.disable();
+  await expect(f.lease.runBrowser(context, operation)).rejects.toThrow();
+  expect(f.ensureReady).not.toHaveBeenCalled();
+});

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();