@vellumai/assistant 0.8.5 → 0.8.6

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

@@ -0,0 +1,282 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+interface MockPlatformClient {
+  platformAssistantId: string;
+  fetch: ReturnType<typeof mock>;
+}
+
+let mockClient: MockPlatformClient | null = null;
+
+mock.module("../../../platform/client.js", () => ({
+  VellumPlatformClient: {
+    create: async () => mockClient,
+  },
+}));
+
+import { callManagedSearchProxy } from "../managed-search-proxy.js";
+
+describe("callManagedSearchProxy", () => {
+  beforeEach(() => {
+    mockClient = {
+      platformAssistantId: "asst-123",
+      fetch: mock(async () => {
+        return new Response(
+          JSON.stringify({
+            status: 200,
+            headers: { "content-type": "application/json" },
+            body: { web: { results: [] } },
+          }),
+          { status: 200, headers: { "content-type": "application/json" } },
+        );
+      }),
+    };
+  });
+
+  afterEach(() => {
+    mockClient = null;
+  });
+
+  test("posts to the managed search proxy endpoint for the assistant and provider", async () => {
+    await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(mockClient!.fetch).toHaveBeenCalledTimes(1);
+    expect(mockClient!.fetch.mock.calls[0][0]).toBe(
+      "/v1/assistants/asst-123/managed-search-proxy/brave/",
+    );
+  });
+
+  test("encodes URL path segments", async () => {
+    mockClient!.platformAssistantId = "asst/123";
+
+    await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(mockClient!.fetch.mock.calls[0][0]).toBe(
+      "/v1/assistants/asst%2F123/managed-search-proxy/brave/",
+    );
+  });
+
+  test("sends the platform contract request envelope", async () => {
+    await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+      query: {
+        q: "kimi k2 fireworks web search",
+        count: "10",
+        offset: "0",
+        freshness: "pw",
+      },
+      headers: {
+        Accept: "application/json",
+      },
+      body: null,
+    });
+
+    const init = mockClient!.fetch.mock.calls[0][1] as RequestInit;
+    expect(init.method).toBe("POST");
+    expect(new Headers(init.headers).get("Content-Type")).toBe(
+      "application/json",
+    );
+    expect(JSON.parse(init.body as string)).toEqual({
+      request: {
+        method: "GET",
+        path: "/res/v1/web/search",
+        query: {
+          q: "kimi k2 fireworks web search",
+          count: "10",
+          offset: "0",
+          freshness: "pw",
+        },
+        headers: {
+          Accept: "application/json",
+        },
+        body: null,
+      },
+    });
+  });
+
+  test("fills absent query, headers, and body with contract defaults", async () => {
+    await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    const init = mockClient!.fetch.mock.calls[0][1] as RequestInit;
+    expect(JSON.parse(init.body as string)).toEqual({
+      request: {
+        method: "GET",
+        path: "/res/v1/web/search",
+        query: {},
+        headers: {},
+        body: null,
+      },
+    });
+  });
+
+  test("passes the abort signal through to platform fetch", async () => {
+    const controller = new AbortController();
+
+    await callManagedSearchProxy(
+      "brave",
+      {
+        method: "GET",
+        path: "/res/v1/web/search",
+      },
+      controller.signal,
+    );
+
+    const init = mockClient!.fetch.mock.calls[0][1] as RequestInit;
+    expect(init.signal).toBe(controller.signal);
+  });
+
+  test("parses the response envelope and preserves provider JSON body", async () => {
+    const providerBody = {
+      web: {
+        results: [{ title: "Result", url: "https://example.com" }],
+      },
+    };
+    mockClient!.fetch = mock(async () => {
+      return new Response(
+        JSON.stringify({
+          status: 200,
+          headers: {
+            "content-type": "application/json",
+            "x-request-id": "req-123",
+          },
+          body: providerBody,
+        }),
+        { status: 200, headers: { "content-type": "application/json" } },
+      );
+    });
+
+    const result = await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(result).toEqual({
+      ok: true,
+      status: 200,
+      headers: {
+        "content-type": "application/json",
+        "x-request-id": "req-123",
+      },
+      body: providerBody,
+    });
+  });
+
+  test("returns a typed platform error for non-2xx platform responses", async () => {
+    mockClient!.fetch = mock(async () => {
+      return new Response(
+        JSON.stringify({ detail: "unsupported managed search provider" }),
+        {
+          status: 400,
+          headers: {
+            "content-type": "application/json",
+            "x-platform-request-id": "platform-req-123",
+          },
+        },
+      );
+    });
+
+    const result = await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(result).toEqual({
+      ok: false,
+      kind: "platform-error",
+      status: 400,
+      headers: {
+        "content-type": "application/json",
+        "x-platform-request-id": "platform-req-123",
+      },
+      body: { detail: "unsupported managed search provider" },
+      message:
+        "Managed search proxy returned status 400: unsupported managed search provider",
+    });
+  });
+
+  test("preserves insufficient-balance platform errors", async () => {
+    mockClient!.fetch = mock(async () => {
+      return new Response(JSON.stringify({ detail: "Insufficient balance" }), {
+        status: 402,
+        headers: { "content-type": "application/json" },
+      });
+    });
+
+    const result = await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(result).toEqual({
+      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",
+    });
+  });
+
+  test("returns unavailable when platform context is missing", async () => {
+    mockClient = null;
+
+    const result = await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(result).toEqual({
+      ok: false,
+      kind: "unavailable",
+      message: "Managed search proxy is unavailable in this environment.",
+    });
+  });
+
+  test("returns unavailable when platform assistant ID is missing", async () => {
+    mockClient!.platformAssistantId = "";
+
+    const result = await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(result).toEqual({
+      ok: false,
+      kind: "unavailable",
+      message:
+        "Managed search proxy is unavailable: platform assistant ID is missing.",
+    });
+  });
+
+  test("returns an invalid-response result for malformed platform envelopes", async () => {
+    mockClient!.fetch = mock(async () => {
+      return new Response(JSON.stringify({ ok: true }), {
+        status: 200,
+        headers: { "content-type": "application/json" },
+      });
+    });
+
+    const result = await callManagedSearchProxy("brave", {
+      method: "GET",
+      path: "/res/v1/web/search",
+    });
+
+    expect(result).toEqual({
+      ok: false,
+      kind: "invalid-response",
+      body: { ok: true },
+      message: "Managed search proxy returned an invalid response envelope.",
+    });
+  });
+});

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 () => {