@vellumai/assistant 0.6.2 → 0.6.3

package/src/__tests__/headless-browser-read-tools.test.ts

@@ -9,6 +9,55 @@ mock.module("../util/logger.js", () => ({
     }),
 }));
 
+/**
+ * Fake CDP session used by the screenshot/extract/wait_for tests in
+ * this file. The tests configure `sendHandler` before invoking a
+ * tool; each `session.send(method, params)` call is recorded in
+ * `sendCalls` and routed to the handler. The handler returns either
+ * a CDP response object (e.g. `{ result: { value: ... } }` for
+ * `Runtime.evaluate`, or `{ data }` for `Page.captureScreenshot`) or
+ * an `Error` to simulate a CDP failure.
+ *
+ * The fake session is exposed via `mockPage.context().newCDPSession(page)`
+ * which is what the real `LocalCdpClient` calls internally. Going
+ * through the real `LocalCdpClient` (instead of mocking the factory
+ * or the cdp-client submodules) avoids polluting the global module
+ * cache that the `factory.test.ts` and LocalCdpClient/
+ * ExtensionCdpClient unit tests rely on.
+ */
+interface SendCall {
+  method: string;
+  params: Record<string, unknown> | undefined;
+}
+
+let sendCalls: SendCall[];
+let sendHandler: (
+  method: string,
+  params: Record<string, unknown> | undefined,
+) => unknown;
+
+function resetCdpMock() {
+  sendCalls = [];
+  sendHandler = () => ({});
+}
+
+const fakeCdpSession = {
+  send: async (method: string, params?: Record<string, unknown>) => {
+    sendCalls.push({ method, params });
+    const value = sendHandler(method, params);
+    if (value instanceof Error) throw value;
+    return value;
+  },
+  // Provided so `LocalCdpClient.dispose()` can call `session.detach()`
+  // without throwing. The tests don't assert on detach calls.
+  detach: async () => {},
+};
+
+// The mock page is served by `browserManager.getOrCreateSessionPage`
+// and is consumed indirectly: LocalCdpClient calls
+// `page.context().newCDPSession(page)` to obtain a CDP session and
+// then dispatches raw CDP methods against it. The page's
+// `context().newCDPSession` is wired to return `fakeCdpSession` above.
 let mockPage: {
   click: ReturnType<typeof mock>;
   fill: ReturnType<typeof mock>;
@@ -22,25 +71,17 @@ let mockPage: {
   waitForSelector: ReturnType<typeof mock>;
   waitForFunction: ReturnType<typeof mock>;
   keyboard: { press: ReturnType<typeof mock> };
+  context: () => {
+    newCDPSession: (page: unknown) => Promise<typeof fakeCdpSession>;
+  };
 };
 
-let snapshotMaps: Map<string, Map<string, string>>;
-
 mock.module("../tools/browser/browser-manager.js", () => {
-  snapshotMaps = new Map();
   return {
     browserManager: {
       getOrCreateSessionPage: async () => mockPage,
       closeSessionPage: async () => {},
       closeAllPages: async () => {},
-      storeSnapshotMap: (conversationId: string, map: Map<string, string>) => {
-        snapshotMaps.set(conversationId, map);
-      },
-      resolveSnapshotSelector: (conversationId: string, elementId: string) => {
-        const map = snapshotMaps.get(conversationId);
-        if (!map) return null;
-        return map.get(elementId) ?? null;
-      },
     },
   };
 });
@@ -55,8 +96,9 @@ mock.module("../tools/network/url-safety.js", () => ({
 
 import {
   executeBrowserExtract,
-  executeBrowserPressKey,
+  executeBrowserScreenshot,
   executeBrowserWaitFor,
+  EXTRACT_LINKS_EXPRESSION,
 } from "../tools/browser/browser-execution.js";
 import type { ToolContext } from "../tools/types.js";
 
@@ -83,72 +125,76 @@ function resetMockPage() {
     waitForSelector: mock(async () => null),
     waitForFunction: mock(async () => null),
     keyboard: { press: mock(async () => {}) },
+    // `LocalCdpClient.ensureSession()` calls `page.context().newCDPSession(
+    // page)` to create a Playwright CDPSession. For these tests we
+    // return the in-file `fakeCdpSession` which records every send()
+    // into `sendCalls` and lets each test set `sendHandler` to shape
+    // the responses.
+    context: () => ({
+      newCDPSession: async (_page: unknown) => fakeCdpSession,
+    }),
   };
 }
 
-// ── browser_press_key ────────────────────────────────────────────────
+// executeBrowserPressKey tests live in
+// `headless-browser-interactions.test.ts` alongside the other
+// CDP-migrated interaction tools.
 
-describe("executeBrowserPressKey", () => {
+// ── browser_screenshot ───────────────────────────────────────────────
+
+describe("executeBrowserScreenshot", () => {
   beforeEach(() => {
     resetMockPage();
-    snapshotMaps.clear();
+    resetCdpMock();
   });
 
-  test("presses key on focused element (no target)", async () => {
-    const result = await executeBrowserPressKey({ key: "Enter" }, ctx);
-    expect(result.isError).toBe(false);
-    expect(result.content).toContain('Pressed "Enter"');
-    expect(mockPage.keyboard.press).toHaveBeenCalledWith("Enter");
-  });
+  test("captures viewport screenshot via Page.captureScreenshot", async () => {
+    // Base64 for "abc" = "YWJj"
+    sendHandler = () => ({ data: "YWJj" });
+
+    const result = await executeBrowserScreenshot({}, ctx);
 
-  test("presses key on targeted element via element_id", async () => {
-    snapshotMaps.set(
-      "test-conversation",
-      new Map([["e1", '[data-vellum-eid="e1"]']]),
-    );
-    const result = await executeBrowserPressKey(
-      { key: "Tab", element_id: "e1" },
-      ctx,
-    );
     expect(result.isError).toBe(false);
-    expect(result.content).toContain('Pressed "Tab" on element');
-    expect(mockPage.press).toHaveBeenCalledWith(
-      '[data-vellum-eid="e1"]',
-      "Tab",
-    );
+    expect(result.content).toContain("Screenshot captured");
+    expect(result.content).toContain("viewport");
+    expect(result.contentBlocks).toHaveLength(1);
+    const block = result.contentBlocks![0]!;
+
+    const source = (block as any).source;
+    expect(source.media_type).toBe("image/jpeg");
+    expect(source.data).toBe("YWJj");
+    // Page.captureScreenshot was called with jpeg quality 80
+    expect(sendCalls).toHaveLength(1);
+    expect(sendCalls[0]!.method).toBe("Page.captureScreenshot");
+    expect(sendCalls[0]!.params).toEqual({
+      format: "jpeg",
+      quality: 80,
+      captureBeyondViewport: false,
+    });
   });
 
-  test("presses key on targeted element via selector", async () => {
-    const result = await executeBrowserPressKey(
-      { key: "Escape", selector: "#modal" },
-      ctx,
-    );
+  test("captures full-page screenshot when full_page is true", async () => {
+    sendHandler = () => ({ data: "YWJj" });
+
+    const result = await executeBrowserScreenshot({ full_page: true }, ctx);
+
     expect(result.isError).toBe(false);
-    expect(mockPage.press).toHaveBeenCalledWith("#modal", "Escape");
+    expect(result.content).toContain("full page");
+    expect(sendCalls[0]!.params).toEqual({
+      format: "jpeg",
+      quality: 80,
+      captureBeyondViewport: true,
+    });
   });
 
-  test("errors when key is missing", async () => {
-    const result = await executeBrowserPressKey({}, ctx);
-    expect(result.isError).toBe(true);
-    expect(result.content).toContain("key is required");
-  });
+  test("surfaces CDP failure as an error result", async () => {
+    sendHandler = () => new Error("CDP crashed");
 
-  test("errors when element_id not found", async () => {
-    const result = await executeBrowserPressKey(
-      { key: "Enter", element_id: "e99" },
-      ctx,
-    );
-    expect(result.isError).toBe(true);
-    expect(result.content).toContain('element_id "e99" not found');
-  });
+    const result = await executeBrowserScreenshot({}, ctx);
 
-  test("handles press error", async () => {
-    mockPage.keyboard.press = mock(async () => {
-      throw new Error("key not recognized");
-    });
-    const result = await executeBrowserPressKey({ key: "InvalidKey" }, ctx);
     expect(result.isError).toBe(true);
-    expect(result.content).toContain("Press key failed");
+    expect(result.content).toContain("Screenshot failed");
+    expect(result.content).toContain("CDP crashed");
   });
 });
 
@@ -157,28 +203,67 @@ describe("executeBrowserPressKey", () => {
 describe("executeBrowserWaitFor", () => {
   beforeEach(() => {
     resetMockPage();
+    resetCdpMock();
   });
 
-  test("waits for selector", async () => {
+  test("waits for selector (CDP polling)", async () => {
+    // waitForSelector polls via Runtime.evaluate until the element
+    // exists, then calls querySelectorBackendNodeId (which triggers
+    // DOM.getDocument / DOM.querySelector / DOM.describeNode).
+    sendHandler = (method, _params) => {
+      if (method === "Runtime.evaluate") {
+        return { result: { value: true } };
+      }
+      if (method === "DOM.getDocument") return { root: { nodeId: 1 } };
+      if (method === "DOM.querySelector") return { nodeId: 42 };
+      if (method === "DOM.describeNode") {
+        return { node: { backendNodeId: 100 } };
+      }
+      return {};
+    };
+
     const result = await executeBrowserWaitFor({ selector: "#loaded" }, ctx);
+
     expect(result.isError).toBe(false);
     expect(result.content).toContain('Element matching "#loaded" appeared');
-    expect(mockPage.waitForSelector).toHaveBeenCalledWith("#loaded", {
-      timeout: 30_000,
-    });
+    // First call is Runtime.evaluate checking existence
+    const evaluateCall = sendCalls.find((c) => c.method === "Runtime.evaluate");
+    expect(evaluateCall).toBeDefined();
+    expect((evaluateCall!.params as { expression: string }).expression).toBe(
+      'document.querySelector("#loaded") !== null',
+    );
   });
 
-  test("waits for text", async () => {
+  test("waits for text (CDP polling)", async () => {
+    sendHandler = (method) => {
+      if (method === "Runtime.evaluate") {
+        return { result: { value: true } };
+      }
+      return {};
+    };
+
     const result = await executeBrowserWaitFor({ text: "Success" }, ctx);
+
     expect(result.isError).toBe(false);
     expect(result.content).toContain('Text "Success" appeared');
-    expect(mockPage.waitForFunction).toHaveBeenCalled();
+    const evaluateCall = sendCalls.find((c) => c.method === "Runtime.evaluate");
+    expect(evaluateCall).toBeDefined();
+    expect(
+      (evaluateCall!.params as { expression: string }).expression,
+    ).toContain('"Success"');
+    expect(
+      (evaluateCall!.params as { expression: string }).expression,
+    ).toContain(".includes(");
   });
 
-  test("waits for duration", async () => {
+  test("waits for duration (no CDP client acquired)", async () => {
     const result = await executeBrowserWaitFor({ duration: 10 }, ctx);
+
     expect(result.isError).toBe(false);
     expect(result.content).toContain("Waited 10ms");
+    // Duration mode is transport-agnostic and must not allocate a
+    // CdpClient — so neither send nor dispose should have fired.
+    expect(sendCalls).toHaveLength(0);
   });
 
   test("errors when no mode specified", async () => {
@@ -187,6 +272,8 @@ describe("executeBrowserWaitFor", () => {
     expect(result.content).toContain(
       "Exactly one of selector, text, or duration",
     );
+    // Validation rejects before any CDP work is attempted.
+    expect(sendCalls).toHaveLength(0);
   });
 
   test("errors when multiple modes specified", async () => {
@@ -196,17 +283,7 @@ describe("executeBrowserWaitFor", () => {
     );
     expect(result.isError).toBe(true);
     expect(result.content).toContain("exactly one");
-  });
-
-  test("respects custom timeout", async () => {
-    const result = await executeBrowserWaitFor(
-      { selector: "#el", timeout: 5000 },
-      ctx,
-    );
-    expect(result.isError).toBe(false);
-    expect(mockPage.waitForSelector).toHaveBeenCalledWith("#el", {
-      timeout: 5000,
-    });
+    expect(sendCalls).toHaveLength(0);
   });
 
   test("caps duration at MAX_WAIT_MS", async () => {
@@ -217,14 +294,16 @@ describe("executeBrowserWaitFor", () => {
     expect(result.content).toContain("Waited 50ms");
   });
 
-  test("handles wait error (timeout)", async () => {
-    mockPage.waitForSelector = mock(async () => {
-      throw new Error("Timeout 30000ms exceeded");
-    });
-    const result = await executeBrowserWaitFor({ selector: "#missing" }, ctx);
+  test("surfaces CDP transport failure as a wait error", async () => {
+    sendHandler = () => new Error("CDP transport failed");
+
+    const result = await executeBrowserWaitFor(
+      { selector: "#missing", timeout: 100 },
+      ctx,
+    );
+
     expect(result.isError).toBe(true);
     expect(result.content).toContain("Wait failed");
-    expect(result.content).toContain("Timeout");
   });
 });
 
@@ -233,60 +312,144 @@ describe("executeBrowserWaitFor", () => {
 describe("executeBrowserExtract", () => {
   beforeEach(() => {
     resetMockPage();
+    resetCdpMock();
   });
 
-  test("extracts page text content", async () => {
-    mockPage.evaluate = mock(async () => "Hello World");
+  test("extracts page text content via CDP", async () => {
+    sendHandler = (method, params) => {
+      if (method !== "Runtime.evaluate") return {};
+      const expression = (params as { expression: string }).expression;
+      if (expression === "document.location.href") {
+        return { result: { value: "https://example.com/" } };
+      }
+      if (expression === "document.title") {
+        return { result: { value: "Test Page" } };
+      }
+      if (expression === "document.body?.innerText ?? ''") {
+        return { result: { value: "Hello World" } };
+      }
+      return { result: { value: null } };
+    };
+
     const result = await executeBrowserExtract({}, ctx);
+
     expect(result.isError).toBe(false);
     expect(result.content).toContain("URL: https://example.com/");
     expect(result.content).toContain("Title: Test Page");
     expect(result.content).toContain("Hello World");
+    // URL, title, body innerText = 3 Runtime.evaluate calls
+    const evaluateCalls = sendCalls.filter(
+      (c) => c.method === "Runtime.evaluate",
+    );
+    expect(evaluateCalls).toHaveLength(3);
   });
 
   test("shows (empty page) for empty content", async () => {
-    mockPage.evaluate = mock(async () => "");
+    sendHandler = (method, params) => {
+      if (method !== "Runtime.evaluate") return {};
+      const expression = (params as { expression: string }).expression;
+      if (expression === "document.location.href") {
+        return { result: { value: "https://example.com/" } };
+      }
+      if (expression === "document.title") {
+        return { result: { value: "Test Page" } };
+      }
+      return { result: { value: "" } };
+    };
+
     const result = await executeBrowserExtract({}, ctx);
     expect(result.content).toContain("(empty page)");
   });
 
   test("truncates long content", async () => {
     const longText = "x".repeat(60_000);
-    mockPage.evaluate = mock(async () => longText);
+    sendHandler = (method, params) => {
+      if (method !== "Runtime.evaluate") return {};
+      const expression = (params as { expression: string }).expression;
+      if (expression === "document.location.href") {
+        return { result: { value: "https://example.com/" } };
+      }
+      if (expression === "document.title") {
+        return { result: { value: "Test Page" } };
+      }
+      return { result: { value: longText } };
+    };
+
     const result = await executeBrowserExtract({}, ctx);
     expect(result.content).toContain("... (truncated)");
     // Content should be capped
     expect(result.content.length).toBeLessThan(60_000);
   });
 
-  test("includes links when requested", async () => {
-    let callCount = 0;
-    mockPage.evaluate = mock(async () => {
-      callCount++;
-      if (callCount === 1) return "Page text";
-      return [
-        { text: "About", href: "https://example.com/about" },
-        { text: "Contact", href: "https://example.com/contact" },
-      ];
-    });
+  test("includes links when requested using EXTRACT_LINKS_EXPRESSION", async () => {
+    sendHandler = (method, params) => {
+      if (method !== "Runtime.evaluate") return {};
+      const expression = (params as { expression: string }).expression;
+      if (expression === "document.location.href") {
+        return { result: { value: "https://example.com/" } };
+      }
+      if (expression === "document.title") {
+        return { result: { value: "Test Page" } };
+      }
+      if (expression === "document.body?.innerText ?? ''") {
+        return { result: { value: "Page text" } };
+      }
+      if (expression === EXTRACT_LINKS_EXPRESSION) {
+        return {
+          result: {
+            value: [
+              { text: "About", href: "https://example.com/about" },
+              { text: "Contact", href: "https://example.com/contact" },
+            ],
+          },
+        };
+      }
+      return { result: { value: null } };
+    };
 
     const result = await executeBrowserExtract({ include_links: true }, ctx);
     expect(result.isError).toBe(false);
     expect(result.content).toContain("Links:");
     expect(result.content).toContain("[About](https://example.com/about)");
     expect(result.content).toContain("[Contact](https://example.com/contact)");
+    // EXTRACT_LINKS_EXPRESSION was actually used (assert the expression appears in a call)
+    const linksCall = sendCalls.find(
+      (c) =>
+        c.method === "Runtime.evaluate" &&
+        (c.params as { expression: string }).expression ===
+          EXTRACT_LINKS_EXPRESSION,
+    );
+    expect(linksCall).toBeDefined();
   });
 
   test("does not include links by default", async () => {
-    mockPage.evaluate = mock(async () => "Page text");
+    sendHandler = (method, params) => {
+      if (method !== "Runtime.evaluate") return {};
+      const expression = (params as { expression: string }).expression;
+      if (expression === "document.location.href") {
+        return { result: { value: "https://example.com/" } };
+      }
+      if (expression === "document.title") {
+        return { result: { value: "Test Page" } };
+      }
+      return { result: { value: "Page text" } };
+    };
+
     const result = await executeBrowserExtract({}, ctx);
     expect(result.content).not.toContain("Links:");
+    // EXTRACT_LINKS_EXPRESSION should NOT have been evaluated
+    const linksCall = sendCalls.find(
+      (c) =>
+        c.method === "Runtime.evaluate" &&
+        (c.params as { expression: string }).expression ===
+          EXTRACT_LINKS_EXPRESSION,
+    );
+    expect(linksCall).toBeUndefined();
   });
 
-  test("handles extract error", async () => {
-    mockPage.evaluate = mock(async () => {
-      throw new Error("page crashed");
-    });
+  test("surfaces CDP failure as an extract error", async () => {
+    sendHandler = () => new Error("page crashed");
+
     const result = await executeBrowserExtract({}, ctx);
     expect(result.isError).toBe(true);
     expect(result.content).toContain("Extract failed");