@vellumai/assistant 0.8.5 → 0.8.7

package/src/tools/network/__tests__/web-search.test.ts

@@ -2,9 +2,16 @@ import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
 
 // Mutable mock state - set per test
 let mockWebSearchProvider: string | undefined = "perplexity";
+let mockWebSearchMode: string | undefined = "your-own";
 let mockBraveSecureKey: string | undefined;
 let mockPerplexitySecureKey: string | undefined;
 let mockTavilySecureKey: string | undefined;
+let mockManagedSearchProxyResult: any;
+let mockManagedSearchProxyCalls: Array<{
+  provider: string;
+  request: Record<string, unknown>;
+  signal?: AbortSignal;
+}> = [];
 
 // Capture the registered tool
 let capturedTool: any = null;
@@ -18,7 +25,10 @@ mock.module("../../registry.js", () => ({
 mock.module("../../../config/loader.js", () => ({
   getConfig: () => ({
     services: {
-      "web-search": { provider: mockWebSearchProvider },
+      "web-search": {
+        mode: mockWebSearchMode,
+        provider: mockWebSearchProvider,
+      },
     },
   }),
 }));
@@ -43,6 +53,17 @@ mock.module("../../../permissions/types.js", () => ({
   RiskLevel: { Low: "low", Medium: "medium", High: "high" },
 }));
 
+mock.module("../managed-search-proxy.js", () => ({
+  callManagedSearchProxy: async (
+    provider: string,
+    request: Record<string, unknown>,
+    signal?: AbortSignal,
+  ) => {
+    mockManagedSearchProxyCalls.push({ provider, request, signal });
+    return mockManagedSearchProxyResult;
+  },
+}));
+
 // Force the module to load (triggers registerTool)
 await import("../web-search.js");
 
@@ -52,17 +73,25 @@ describe("web_search tool", () => {
   beforeEach(() => {
     originalFetch = globalThis.fetch;
     mockWebSearchProvider = "perplexity";
+    mockWebSearchMode = "your-own";
     mockBraveSecureKey = undefined;
     mockPerplexitySecureKey = undefined;
     mockTavilySecureKey = undefined;
+    mockManagedSearchProxyCalls = [];
+    mockManagedSearchProxyResult = {
+      ok: true,
+      status: 200,
+      headers: { "content-type": "application/json" },
+      body: { web: { results: [] } },
+    };
   });
 
   afterEach(() => {
     globalThis.fetch = originalFetch;
   });
 
-  function execute(input: Record<string, unknown>) {
-    return capturedTool.execute(input, {} as any);
+  function execute(input: Record<string, unknown>, context: any = {}) {
+    return capturedTool.execute(input, context);
   }
 
   // ---- Input validation ---------------------------------------------------
@@ -350,6 +379,185 @@ describe("web_search tool", () => {
     expect(callCount).toBe(4);
   });
 
+  // ---- Managed Brave provider -------------------------------------------
+
+  test("managed mode uses Brave proxy without BYOK provider keys", async () => {
+    mockWebSearchMode = "managed";
+    mockWebSearchProvider = "inference-provider-native";
+    mockManagedSearchProxyResult = {
+      ok: true,
+      status: 200,
+      headers: { "content-type": "application/json" },
+      body: {
+        web: {
+          results: [
+            {
+              title: "Managed Result",
+              url: "https://example.com/managed",
+              description: "Managed Brave result",
+            },
+          ],
+        },
+      },
+    };
+
+    const result = await execute({
+      query: "managed query",
+      count: 5,
+      offset: 2,
+      freshness: "pw",
+    });
+
+    expect(result.isError).toBe(false);
+    expect(result.content).toContain("Managed Result");
+    expect(mockManagedSearchProxyCalls).toHaveLength(1);
+    expect(mockManagedSearchProxyCalls[0]).toMatchObject({
+      provider: "brave",
+      request: {
+        method: "GET",
+        path: "/res/v1/web/search",
+        query: {
+          q: "managed query",
+          count: "5",
+          offset: "2",
+          freshness: "pw",
+        },
+        headers: {
+          Accept: "application/json",
+        },
+        body: null,
+      },
+    });
+  });
+
+  test("managed Brave formats results like direct Brave", async () => {
+    const braveBody = {
+      web: {
+        results: [
+          {
+            title: "Same Result",
+            url: "https://example.com/same",
+            description: "Same description",
+            age: "1 day ago",
+            extra_snippets: ["Same snippet"],
+          },
+        ],
+      },
+    };
+
+    mockWebSearchMode = "your-own";
+    mockWebSearchProvider = "brave";
+    mockBraveSecureKey = "brave-key";
+    globalThis.fetch = (async () => {
+      return new Response(JSON.stringify(braveBody), {
+        status: 200,
+        headers: { "content-type": "application/json" },
+      });
+    }) as any;
+
+    const directResult = await execute({ query: "same query" });
+
+    mockWebSearchMode = "managed";
+    mockBraveSecureKey = undefined;
+    mockManagedSearchProxyResult = {
+      ok: true,
+      status: 200,
+      headers: { "content-type": "application/json" },
+      body: braveBody,
+    };
+
+    const managedResult = await execute({ query: "same query" });
+
+    expect(managedResult.isError).toBe(false);
+    expect(managedResult.content).toBe(directResult.content);
+    expect(managedResult.activityMetadata.webSearch.results).toEqual(
+      directResult.activityMetadata.webSearch.results,
+    );
+  });
+
+  test("managed mode passes abort signal to the platform proxy", async () => {
+    mockWebSearchMode = "managed";
+    const controller = new AbortController();
+
+    await execute({ query: "abortable query" }, { signal: controller.signal });
+
+    expect(mockManagedSearchProxyCalls).toHaveLength(1);
+    expect(mockManagedSearchProxyCalls[0].signal).toBe(controller.signal);
+  });
+
+  test("managed mode maps insufficient balance to a managed usage error", async () => {
+    mockWebSearchMode = "managed";
+    mockManagedSearchProxyResult = {
+      ok: false,
+      kind: "platform-error",
+      status: 402,
+      headers: { "content-type": "application/json" },
+      body: { detail: "Insufficient balance" },
+      message: "Managed search proxy returned status 402: Insufficient balance",
+    };
+
+    const result = await execute({ query: "billing query" });
+
+    expect(result.isError).toBe(true);
+    expect(result.content).toContain("Managed web search");
+    expect(result.content).toContain("account balance");
+    expect(result.content).toContain("Your Own mode");
+  });
+
+  test("managed mode maps proxied provider errors to a tool error", async () => {
+    mockWebSearchMode = "managed";
+    mockManagedSearchProxyResult = {
+      ok: true,
+      status: 400,
+      headers: { "content-type": "application/json" },
+      body: { error: "Bad upstream request" },
+    };
+
+    const result = await execute({ query: "provider error query" });
+
+    expect(result.isError).toBe(true);
+    expect(result.content).toContain(
+      "Managed Brave Search provider returned status 400",
+    );
+  });
+
+  test("managed mode returns a clear error when platform context is unavailable", async () => {
+    mockWebSearchMode = "managed";
+    mockManagedSearchProxyResult = {
+      ok: false,
+      kind: "unavailable",
+      message: "Managed search proxy is unavailable in this environment.",
+    };
+
+    const result = await execute({ query: "local managed query" });
+
+    expect(result.isError).toBe(true);
+    expect(result.content).toContain("Managed search proxy is unavailable");
+    expect(result.content).toContain("Log in to Vellum");
+  });
+
+  test("your-own mode keeps direct Brave BYOK behavior unchanged", async () => {
+    mockWebSearchMode = "your-own";
+    mockWebSearchProvider = "brave";
+    mockBraveSecureKey = "brave-direct-key";
+    let capturedHeaders: Headers | undefined;
+    globalThis.fetch = (async (_url: string, init?: RequestInit) => {
+      capturedHeaders = new Headers(init?.headers);
+      return new Response(JSON.stringify({ web: { results: [] } }), {
+        status: 200,
+        headers: { "content-type": "application/json" },
+      });
+    }) as any;
+
+    const result = await execute({ query: "direct query" });
+
+    expect(result.isError).toBe(false);
+    expect(mockManagedSearchProxyCalls).toHaveLength(0);
+    expect(capturedHeaders!.get("X-Subscription-Token")).toBe(
+      "brave-direct-key",
+    );
+  });
+
   // ---- Tavily provider ----------------------------------------------------
 
   test("executes Tavily search successfully", async () => {

package/src/tools/network/managed-search-proxy.ts

@@ -0,0 +1,183 @@
+import { VellumPlatformClient } from "../../platform/client.js";
+
+export type ManagedSearchProxyProvider = "brave";
+
+export interface ManagedSearchProxyRequest {
+  method: string;
+  path: string;
+  query?: Record<string, string>;
+  headers?: Record<string, string>;
+  body?: unknown;
+}
+
+export interface ManagedSearchProxyResponse {
+  status: number;
+  headers: Record<string, string>;
+  body: unknown;
+}
+
+export type ManagedSearchProxyResult =
+  | ({ ok: true } & ManagedSearchProxyResponse)
+  | {
+      ok: false;
+      kind: "unavailable";
+      message: string;
+    }
+  | {
+      ok: false;
+      kind: "platform-error";
+      status: number;
+      headers: Record<string, string>;
+      body: unknown;
+      message: string;
+    }
+  | {
+      ok: false;
+      kind: "invalid-response";
+      body: unknown;
+      message: string;
+    };
+
+interface ManagedSearchProxyEnvelope {
+  request: {
+    method: string;
+    path: string;
+    query: Record<string, string>;
+    headers: Record<string, string>;
+    body: unknown;
+  };
+}
+
+export async function callManagedSearchProxy(
+  provider: ManagedSearchProxyProvider,
+  request: ManagedSearchProxyRequest,
+  signal?: AbortSignal,
+): Promise<ManagedSearchProxyResult> {
+  const client = await VellumPlatformClient.create();
+  if (!client) {
+    return {
+      ok: false,
+      kind: "unavailable",
+      message: "Managed search proxy is unavailable in this environment.",
+    };
+  }
+
+  if (!client.platformAssistantId) {
+    return {
+      ok: false,
+      kind: "unavailable",
+      message:
+        "Managed search proxy is unavailable: platform assistant ID is missing.",
+    };
+  }
+
+  const path = `/v1/assistants/${encodeURIComponent(
+    client.platformAssistantId,
+  )}/managed-search-proxy/${encodeURIComponent(provider)}/`;
+
+  const envelope: ManagedSearchProxyEnvelope = {
+    request: {
+      method: request.method,
+      path: request.path,
+      query: request.query ?? {},
+      headers: request.headers ?? {},
+      body: request.body ?? null,
+    },
+  };
+
+  const response = await client.fetch(path, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify(envelope),
+    signal,
+  });
+
+  if (!response.ok) {
+    const body = await readResponseBody(response);
+    return {
+      ok: false,
+      kind: "platform-error",
+      status: response.status,
+      headers: responseHeadersToRecord(response.headers),
+      body,
+      message: platformErrorMessage(response.status, body),
+    };
+  }
+
+  const body = await readResponseBody(response);
+  if (!isManagedSearchProxyResponse(body)) {
+    return {
+      ok: false,
+      kind: "invalid-response",
+      body,
+      message: "Managed search proxy returned an invalid response envelope.",
+    };
+  }
+
+  return {
+    ok: true,
+    status: body.status,
+    headers: body.headers,
+    body: body.body,
+  };
+}
+
+function isManagedSearchProxyResponse(
+  value: unknown,
+): value is ManagedSearchProxyResponse {
+  if (!isRecord(value)) return false;
+  if (typeof value.status !== "number") return false;
+  if (!isStringRecord(value.headers)) return false;
+  return "body" in value;
+}
+
+async function readResponseBody(response: Response): Promise<unknown> {
+  const text = await response.text().catch(() => "");
+  if (!text) return null;
+
+  const contentType = response.headers.get("content-type") ?? "";
+  if (contentType.includes("application/json")) {
+    try {
+      return JSON.parse(text);
+    } catch {
+      return text;
+    }
+  }
+
+  return text;
+}
+
+function responseHeadersToRecord(headers: Headers): Record<string, string> {
+  return Object.fromEntries(headers.entries());
+}
+
+function platformErrorMessage(status: number, body: unknown): string {
+  const detail = extractErrorDetail(body);
+  if (detail) {
+    return `Managed search proxy returned status ${status}: ${detail}`;
+  }
+  return `Managed search proxy returned status ${status}.`;
+}
+
+function extractErrorDetail(body: unknown): string | undefined {
+  if (typeof body === "string") return body || undefined;
+  if (!isRecord(body)) return undefined;
+
+  for (const key of ["detail", "error", "message"]) {
+    const value = body[key];
+    if (typeof value === "string" && value) return value;
+  }
+
+  return undefined;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function isStringRecord(value: unknown): value is Record<string, string> {
+  if (!isRecord(value)) return false;
+  return Object.values(value).every((entry) => typeof entry === "string");
+}

package/src/tools/network/web-fetch.ts

@@ -12,7 +12,11 @@ import { faviconUrlForDomain } from "../../util/favicon.js";
 import { getLogger } from "../../util/logger.js";
 import { safeStringSlice } from "../../util/unicode.js";
 import { registerTool } from "../registry.js";
-import type { Tool, ToolContext, ToolExecutionResult } from "../types.js";
+import type {
+  ToolContext,
+  ToolDefinition,
+  ToolExecutionResult,
+} from "../types.js";
 import { extractDomain } from "./domain-normalize.js";
 import {
   buildHostHeader,
@@ -984,56 +988,55 @@ export async function executeWebFetch(
   }
 }
 
-class WebFetchTool implements Tool {
-  name = "web_fetch";
-  description =
-    "Fetch a webpage and return LLM-friendly extracted text with metadata. Use this after web_search when you need to read a specific result. To find pages on a site without guessing slugs, fetch /sitemap.xml first — it has ground-truth paths and works even when pages are JS-rendered.";
-  category = "network";
-  executionTarget = "sandbox" as const;
-  defaultRiskLevel = RiskLevel.Low;
-
-  input_schema = {
-        type: "object",
-        properties: {
-          url: {
-            type: "string",
-            description:
-              "The target webpage URL. If scheme is missing, https:// is assumed.",
-          },
-          max_chars: {
-            type: "number",
-            description: `Maximum characters of content to return (1-${MAX_MAX_CHARS}, default ${DEFAULT_MAX_CHARS})`,
-          },
-          start_index: {
-            type: "number",
-            description:
-              "Character index to start returning content from (default 0). Useful for paging large pages.",
-          },
-          timeout_seconds: {
-            type: "number",
-            description: `Request timeout in seconds (1-${MAX_TIMEOUT_SECONDS}, default ${DEFAULT_TIMEOUT_SECONDS})`,
-          },
-          raw: {
-            type: "boolean",
-            description:
-              "If true, return normalized raw response text instead of extracted plain text for HTML pages.",
-          },
-          allow_private_network: {
-            type: "boolean",
-            description:
-              "If true, allows requests to localhost/private-network hosts. Disabled by default for SSRF safety.",
-          },
-        },
-        required: ["url"],
-      };
+export const webFetchTool = {
+  name: "web_fetch",
+  description:
+    "Fetch a webpage and return LLM-friendly extracted text with metadata. Use this after web_search when you need to read a specific result. To find pages on a site without guessing slugs, fetch /sitemap.xml first — it has ground-truth paths and works even when pages are JS-rendered.",
+  category: "network",
+  executionTarget: "sandbox",
+  defaultRiskLevel: RiskLevel.Low,
+
+  input_schema: {
+    type: "object",
+    properties: {
+      url: {
+        type: "string",
+        description:
+          "The target webpage URL. If scheme is missing, https:// is assumed.",
+      },
+      max_chars: {
+        type: "number",
+        description: `Maximum characters of content to return (1-${MAX_MAX_CHARS}, default ${DEFAULT_MAX_CHARS})`,
+      },
+      start_index: {
+        type: "number",
+        description:
+          "Character index to start returning content from (default 0). Useful for paging large pages.",
+      },
+      timeout_seconds: {
+        type: "number",
+        description: `Request timeout in seconds (1-${MAX_TIMEOUT_SECONDS}, default ${DEFAULT_TIMEOUT_SECONDS})`,
+      },
+      raw: {
+        type: "boolean",
+        description:
+          "If true, return normalized raw response text instead of extracted plain text for HTML pages.",
+      },
+      allow_private_network: {
+        type: "boolean",
+        description:
+          "If true, allows requests to localhost/private-network hosts. Disabled by default for SSRF safety.",
+      },
+    },
+    required: ["url"],
+  },
 
   async execute(
     input: Record<string, unknown>,
     context: ToolContext,
   ): Promise<ToolExecutionResult> {
     return executeWebFetch(input, { signal: context.signal });
-  }
-}
+  },
+} satisfies ToolDefinition;
 
-export const webFetchTool = new WebFetchTool();
 registerTool(webFetchTool);