@dbx-tools/appkit-web-search 0.3.28 → 0.3.30

package/src/tool.ts

@@ -15,13 +15,21 @@
  * before the call) treats a pattern gate as "always gate". `approval` falls
  * back to the plugin's `approval` config when a tool omits its own.
  *
+ * The same two tools are exposed to AppKit's own agents through the plugin's
+ * `ToolProvider` (see `plugin.ts`); this module is the Mastra half.
+ *
  * @module
  */
 
-import { getExecutionContext } from "@databricks/appkit";
-import { string } from "@dbx-tools/shared-core";
+import { ValidationError } from "@databricks/appkit";
 import { createTool } from "@mastra/core/tools";
-import { approvalMatches, type ApprovalGate } from "./config";
+import type { z } from "zod";
+import {
+  approvalMatches,
+  toApprovalPolicy,
+  type ApprovalGate,
+  type ApprovalPolicy,
+} from "./config";
 import { runWebFetch } from "./fetch";
 import { getWebSearchRuntime } from "./runtime";
 import {
@@ -29,10 +37,24 @@ import {
   webFetchResultSchema,
   webSearchRequestSchema,
   webSearchResultSchema,
-  type WebFetchRequest,
-  type WebSearchRequest,
+  WEB_FETCH_TOOL_DESCRIPTION,
+  WEB_SEARCH_TOOL_DESCRIPTION,
 } from "./schema";
-import { runWebSearch, type WebSearchContext } from "./search";
+import { resolveWebSearchContext, runWebSearch } from "./search";
+
+/**
+ * Validate a tool call's arguments at the runtime boundary. Mastra checks the
+ * input schema before dispatch, but the argument still arrives typed as
+ * `unknown` and the model is the one filling it in. The rejected value is not
+ * echoed back.
+ */
+function parseToolInput<S extends z.ZodType>(schema: S, input: unknown): z.infer<S> {
+  const parsed = schema.safeParse(input);
+  if (!parsed.success) {
+    throw ValidationError.invalidValue("input", input, "arguments matching the tool's schema");
+  }
+  return parsed.data;
+}
 
 /** Options shared by both web tools. */
 export interface WebSearchToolOptions {
@@ -43,24 +65,14 @@ export interface WebSearchToolOptions {
    * gates every call; a URL-pattern (or list) gates only matching calls;
    * omit / `false` for no approval. See {@link ApprovalGate}.
    */
-  approval?: ApprovalGate;
+  approval?: ApprovalGate | ApprovalPolicy;
 }
 
 /** Resolve the effective gate: explicit tool option, else the plugin default. */
-function effectiveGate(opts: WebSearchToolOptions): ApprovalGate {
-  return opts.approval ?? getWebSearchRuntime().config.approval;
-}
-
-/**
- * Resolve the OBO workspace client + host from the active AppKit execution
- * context. Runs inside `agent.stream`'s `asUser(req)` scope, so the search
- * hits the serving endpoint as the requesting user; outside a user context it
- * falls back to the service principal.
- */
-async function webSearchContext(): Promise<WebSearchContext> {
-  const ctx = getExecutionContext();
-  const host = (await ctx.client.config.getHost()).toString();
-  return { client: ctx.client, host };
+function effectiveGate(opts: WebSearchToolOptions): ApprovalPolicy {
+  return opts.approval === undefined
+    ? getWebSearchRuntime().config.approval
+    : toApprovalPolicy(opts.approval);
 }
 
 /**
@@ -82,24 +94,16 @@ export function webSearchTool(opts: WebSearchToolOptions = {}) {
   const gate = effectiveGate(opts);
   return createTool({
     id: opts.id ?? "web_search",
-    description: string.toDescription(`
-      Search the web for current information and get an answer synthesized from
-      live results, with the sources it used. Pass a natural-language query;
-      the search runs inside a web-search-capable model (chosen independently
-      of your own model). Optionally pass a model name to use a specific
-      web-search model. Use it whenever a question needs up-to-date or external
-      information you don't already have.
-    `),
+    description: WEB_SEARCH_TOOL_DESCRIPTION,
     inputSchema: webSearchRequestSchema,
     outputSchema: webSearchResultSchema,
     // A search's result URLs aren't known before the call, so a pattern gate
-    // is treated as "always gate"; boolean gates pass through.
-    ...(gate === false || gate === undefined
-      ? {}
-      : { requireApproval: () => (typeof gate === "boolean" ? gate : true) }),
+    // is treated as "always gate".
+    ...(gate.mode === "none" ? {} : { requireApproval: () => true }),
     execute: async (input) => {
       const { config } = getWebSearchRuntime();
-      return runWebSearch(input as WebSearchRequest, config, await webSearchContext());
+      const request = parseToolInput(webSearchRequestSchema, input);
+      return runWebSearch(request, config, await resolveWebSearchContext());
     },
   });
 }
@@ -119,26 +123,21 @@ export function webFetchTool(opts: WebSearchToolOptions = {}) {
   const gate = effectiveGate(opts);
   return createTool({
     id: opts.id ?? "web_fetch",
-    description: string.toDescription(`
-      Fetch a single web page and return its readable contents. Pass an
-      absolute URL (including https://); set format to "html" for raw markup
-      instead of extracted text. Use it to read a page returned by web_search
-      or provided by the user. Content is length-capped; fetching a URL outside
-      the configured allow-list is refused.
-    `),
+    description: WEB_FETCH_TOOL_DESCRIPTION,
     inputSchema: webFetchRequestSchema,
     outputSchema: webFetchResultSchema,
     // A fetch knows its single target URL, so a pattern gate is evaluated
-    // precisely against it; boolean gates pass through.
-    ...(gate === false || gate === undefined
+    // precisely against it.
+    ...(gate.mode === "none"
       ? {}
       : {
           requireApproval: (input: unknown) =>
-            approvalMatches(gate, [(input as WebFetchRequest).url]),
+            gate.mode === "always" ||
+            approvalMatches(gate, [parseToolInput(webFetchRequestSchema, input).url]),
         }),
     execute: async (input) => {
       const { config } = getWebSearchRuntime();
-      return runWebFetch(input as WebFetchRequest, config);
+      return runWebFetch(parseToolInput(webFetchRequestSchema, input), config);
     },
   });
 }

package/test/allowlist.test.ts

@@ -6,8 +6,8 @@ import {
   parseAllowedUrls,
   toUrlAllowList,
 } from "../src/allowlist";
-import { approvalMatches, resolveWebSearchConfig } from "../src/config";
-import { htmlToText } from "../src/fetch";
+import { approvalMatches, resolveWebSearchConfig, toApprovalPolicy } from "../src/config";
+import { htmlToText } from "../src/html-text";
 import {
   detectWebSearchProvider,
   supportsWebSearch,
@@ -54,7 +54,7 @@ describe("web-search allow-list", () => {
 
   it("assertUrlAllowed throws for a disallowed URL, passes an allowed one", () => {
     const list = toUrlAllowList(parseAllowedUrls(["databricks.com"]));
-    assert.throws(() => assertUrlAllowed("https://evil.example.com/", list), /not permitted/);
+    assert.throws(() => assertUrlAllowed("https://evil.example.com/", list), /allow-list/);
     assert.doesNotThrow(() => assertUrlAllowed("https://databricks.com/", list));
   });
 });
@@ -97,8 +97,9 @@ describe("web-search provider detection", () => {
 describe("web-search config", () => {
   it("defaults model fallbacks to Gemini-then-GPT and approval to none", () => {
     const c = resolveWebSearchConfig();
-    assert.equal(c.approval, false);
+    assert.deepEqual(c.approval, { mode: "none" });
     assert.equal(c.model, undefined);
+    assert.equal(c.modelSource, "none");
     assert.match(c.modelFallbacks[0]!, /gemini/);
     assert.ok(c.modelFallbacks.some((m) => m.includes("gpt")));
   });
@@ -112,6 +113,71 @@ describe("web-search config", () => {
     assert.equal(resolveWebSearchConfig().scrapeFallback, true);
     assert.equal(resolveWebSearchConfig({ scrapeFallback: false }).scrapeFallback, false);
   });
+
+  it("records where a pinned endpoint came from", () => {
+    assert.equal(resolveWebSearchConfig({ model: "gemini" }).modelSource, "config");
+    process.env.DATABRICKS_SERVING_ENDPOINT_NAME = "databricks-claude-sonnet-4-6";
+    try {
+      const c = resolveWebSearchConfig();
+      assert.equal(c.model, "databricks-claude-sonnet-4-6");
+      assert.equal(c.modelSource, "DATABRICKS_SERVING_ENDPOINT_NAME");
+      // The dedicated override wins over the shared resource binding.
+      assert.equal(resolveWebSearchConfig({ model: "gemini" }).modelSource, "config");
+    } finally {
+      delete process.env.DATABRICKS_SERVING_ENDPOINT_NAME;
+    }
+  });
+
+  it("names the URL policy, defaulting from whether an allow-list was given", () => {
+    assert.equal(resolveWebSearchConfig().urlPolicy, "unrestricted");
+    assert.equal(resolveWebSearchConfig().allowList.restricted, false);
+    const restricted = resolveWebSearchConfig({ allowedUrls: ["databricks.com"] });
+    assert.equal(restricted.urlPolicy, "allowlist");
+    assert.equal(restricted.allowList.restricted, true);
+    assert.equal(
+      resolveWebSearchConfig({ urlPolicy: "allowlist", allowedUrls: ["databricks.com"] }).urlPolicy,
+      "allowlist",
+    );
+  });
+
+  it("fails loudly on a URL policy that contradicts its allow-list", () => {
+    assert.throws(() => resolveWebSearchConfig({ urlPolicy: "allowlist" }), /allowedUrls/);
+    assert.throws(
+      () => resolveWebSearchConfig({ urlPolicy: "unrestricted", allowedUrls: ["databricks.com"] }),
+      /allowedUrls/,
+    );
+  });
+
+  it("fails loudly on a WEB_SEARCH_TOOLS value that is not a JSON object", () => {
+    process.env.WEB_SEARCH_TOOLS = "not-json";
+    try {
+      assert.throws(() => resolveWebSearchConfig(), /WEB_SEARCH_TOOLS/);
+    } finally {
+      delete process.env.WEB_SEARCH_TOOLS;
+    }
+  });
+
+  it("rejects a webSearchTools override that is not a tool spec", () => {
+    assert.throws(() => webSearchToolSpec("gemini", { gemini: "nope" }), /webSearchTools.gemini/);
+  });
+});
+
+describe("web-search approval policy", () => {
+  it("normalizes every accepted spelling", () => {
+    assert.deepEqual(toApprovalPolicy(undefined), { mode: "none" });
+    assert.deepEqual(toApprovalPolicy(false), { mode: "none" });
+    assert.deepEqual(toApprovalPolicy(true), { mode: "always" });
+    assert.deepEqual(toApprovalPolicy("*.internal.example.com"), {
+      mode: "urls",
+      patterns: ["*.internal.example.com"],
+    });
+    assert.deepEqual(toApprovalPolicy({ mode: "always" }), { mode: "always" });
+  });
+
+  it("collapses an empty pattern list to no approval", () => {
+    assert.deepEqual(toApprovalPolicy("  "), { mode: "none" });
+    assert.deepEqual(toApprovalPolicy({ mode: "urls", patterns: [] }), { mode: "none" });
+  });
 });
 
 describe("web-search approval gate", () => {