@dbx-tools/appkit-web-search 0.3.29 → 0.3.31

package/src/search.ts

@@ -10,23 +10,31 @@
  *      agent's chat model (the agent may run on a model without web search).
  *      A pinned `model` (request or config) is fuzzy-matched; otherwise the
  *      configured fallback order (Gemini, then GPT) is walked to the first
- *      web-search-capable endpoint that exists in the workspace. An explicit
- *      but unsupported model is a hard error, not a silent fallback.
+ *      web-search-capable endpoint that exists. An explicit but unsupported
+ *      model is a hard error, not a silent fallback.
  *   2. POSTs to the provider's REST surface (`/serving-endpoints/responses`
  *      for OpenAI, `/serving-endpoints/chat/completions` for Gemini) with the
- *      mapped tool spec, using the OBO-scoped workspace client.
+ *      mapped tool spec, authenticated as the OBO caller.
  *   3. Returns the synthesized answer plus the cited sources, with citations
  *      silently filtered through the configured URL allow-list.
  *
  * @module
  */
 
-import { getExecutionContext } from "@databricks/appkit";
+import {
+  ConfigurationError,
+  ConnectionError,
+  ExecutionError,
+  getExecutionContext,
+  ValidationError,
+} from "@databricks/appkit";
 import { invoke, resolve, serving } from "@dbx-tools/model";
-import { error, log, string } from "@dbx-tools/shared-core";
+import { log, object, string } from "@dbx-tools/shared-core";
 import { openaiChat, openaiResponses } from "@dbx-tools/shared-model";
-import type { ResolvedWebSearchConfig } from "./config";
+import { MODEL_ENV, SERVING_ENDPOINT_ENV, type ResolvedWebSearchConfig } from "./config";
+import { toCallSettings, webSearchExecuteDefaults } from "./defaults";
 import { detectWebSearchProvider, supportsWebSearch, webSearchToolSpec } from "./provider";
+import { executeRead } from "./runtime";
 import type { WebSearchCitation, WebSearchRequest, WebSearchResult } from "./schema";
 import { runScrapeSearch } from "./scrape";
 
@@ -35,6 +43,20 @@ const logger = log.logger("web-search/search");
 const { resolveModel } = resolve;
 const { listServingEndpoints } = serving;
 
+/**
+ * How deep the grounding-metadata walk descends. Gemini nests its sources a
+ * handful of levels down and the shape varies by model version, so the walk
+ * is generic; the bound is what keeps a pathological payload from turning it
+ * into a full traversal of the response.
+ */
+const MAX_GROUNDING_WALK_DEPTH = 6;
+
+/** Lowest HTTP status treated as a server-side (retryable) serving failure. */
+const SERVER_ERROR_STATUS = 500;
+
+/** Status Model Serving uses to shed load; retryable like a 5xx. */
+const RATE_LIMITED_STATUS = 429;
+
 /** Context a search needs from the caller: the OBO client + workspace host. */
 export interface WebSearchContext {
   client: WorkspaceClientLike;
@@ -64,9 +86,13 @@ export async function resolveWebSearchContext(): Promise<WebSearchContext> {
  *
  * Returns the chosen endpoint id, or `null` when the workspace has no
  * web-search-capable model deployed (the caller then uses the scrape
- * fallback). An explicit request that resolves to an unsupported / absent
- * model throws, so a deliberate bad pin surfaces rather than silently
- * degrading.
+ * fallback). A pin the caller chose deliberately - the request's `model`, the
+ * plugin's `model`, or `WEB_SEARCH_MODEL` - throws when it resolves to an
+ * unsupported / absent endpoint, so a bad pin surfaces rather than silently
+ * degrading. A pin inherited from the shared `DATABRICKS_SERVING_ENDPOINT_NAME`
+ * binding is only a preference: that endpoint is the app's serving endpoint,
+ * not necessarily a web-search-capable one, so an unusable value falls through
+ * to the fallback order.
  */
 async function resolveWebSearchModel(
   ctx: WebSearchContext,
@@ -77,6 +103,8 @@ async function resolveWebSearchModel(
   // Only deployed, web-search-capable endpoints are candidates.
   const capable = endpoints.filter((e) => supportsWebSearch(e.name));
   const pinned = requested ?? config.model;
+  const pinIsDeliberate =
+    requested !== undefined || config.modelSource === "config" || config.modelSource === MODEL_ENV;
 
   if (pinned) {
     // Resolve the explicit ask within the capable set only.
@@ -87,19 +115,33 @@ async function resolveWebSearchModel(
     });
     // resolveModel returns the input verbatim on no match; require it to be a
     // real capable endpoint so a bad pin is a clear error, not a phantom call.
-    if (!capable.some((e) => e.name === modelId) || !supportsWebSearch(modelId)) {
-      throw new Error(
-        `web-search: requested model "${pinned}" is not a deployed web-search-capable endpoint. ` +
-          `Deployed GPT/Gemini endpoints: [${capable.map((e) => e.name).join(", ") || "none"}].`,
+    const usable = capable.some((e) => e.name === modelId) && supportsWebSearch(modelId);
+    if (usable) return modelId;
+    const deployed = capable.map((e) => e.name).join(", ") || "none";
+    if (!pinIsDeliberate) {
+      logger.info("shared-endpoint-not-web-capable", {
+        envVar: SERVING_ENDPOINT_ENV,
+        deployed,
+      });
+    } else if (requested !== undefined) {
+      throw ValidationError.invalidValue(
+        "model",
+        pinned,
+        `a deployed web-search-capable endpoint (deployed: ${deployed})`,
+      );
+    } else {
+      throw ConfigurationError.resourceNotFound(
+        "Web-search-capable serving endpoint",
+        `Deployed web-search-capable endpoints: ${deployed}. Set model or ${MODEL_ENV} to one of them.`,
       );
     }
-    return modelId;
   }
 
   if (capable.length === 0) return null;
 
-  // Nothing pinned: prefer the configured fallback order (Gemini, then GPT)
-  // when those ids are actually deployed; else take the best capable endpoint.
+  // Nothing usable pinned: prefer the configured fallback order (Gemini, then
+  // GPT) when those ids are actually deployed; else take the best capable
+  // endpoint.
   const { modelId } = resolveModel(capable, {
     fallbacks: config.modelFallbacks,
     fuzzy: config.fuzzy,
@@ -108,20 +150,57 @@ async function resolveWebSearchModel(
   return capable.some((e) => e.name === modelId) ? modelId : (capable[0]?.name ?? null);
 }
 
-/** POST a serving request through the OBO client and return the parsed JSON. */
+/**
+ * POST a serving request as the OBO caller and return the parsed JSON body.
+ *
+ * Auth headers are minted per call from the OBO client's SDK config, which
+ * refreshes the token when it is close to expiry, so the request carries the
+ * requesting user's identity.
+ *
+ * This is the one Databricks call in the repo that does not go through
+ * `apiClient.request` + `databricks.toContext` (which does forward
+ * cancellation). Retry classification here has to distinguish a load-shed or
+ * server fault from this request's own 4xx, and `fetch` exposes the HTTP status
+ * directly; the SDK raises an `ApiError` that carries the status under
+ * inconsistent keys, which is guesswork by comparison.
+ */
 async function postServing(
   ctx: WebSearchContext,
-  path: string,
+  url: string,
   body: unknown,
+  config: ResolvedWebSearchConfig,
+  cacheKey: readonly (string | number)[],
+  signal?: AbortSignal,
 ): Promise<Record<string, unknown>> {
-  const res = await ctx.client.apiClient.request({
-    path,
-    method: "POST",
-    headers: new Headers({ Accept: "application/json", "Content-Type": "application/json" }),
-    raw: false,
-    payload: body,
-  });
-  return (res ?? {}) as Record<string, unknown>;
+  const payload = await executeRead(
+    "serving-request",
+    toCallSettings(webSearchExecuteDefaults, config.timeoutMs, cacheKey),
+    async (executeSignal): Promise<unknown> => {
+      const response = await fetch(url, {
+        method: "POST",
+        headers: {
+          ...(await invoke.authHeaders(ctx.client)),
+          Accept: "application/json",
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(body),
+        ...(executeSignal ? { signal: executeSignal } : {}),
+      });
+      if (!response.ok) {
+        // Load-shedding and server faults are worth another attempt; a 4xx is
+        // this request's own problem, so it must not be retried.
+        const retryable =
+          response.status >= SERVER_ERROR_STATUS || response.status === RATE_LIMITED_STATUS;
+        const message = `web-search: Model Serving rejected the search request (HTTP ${response.status})`;
+        throw retryable
+          ? new ConnectionError(message, { context: { status: response.status } })
+          : new ExecutionError(message, { context: { status: response.status } });
+      }
+      return response.json();
+    },
+    signal,
+  );
+  return object.isRecord(payload) ? payload : {};
 }
 
 /* --------------------------- response extraction --------------------------- */
@@ -149,22 +228,22 @@ function fromChatPayload(payload: Record<string, unknown>): {
   answer: string;
   citations: WebSearchCitation[];
 } {
-  const choices = Array.isArray(payload.choices) ? (payload.choices as unknown[]) : [];
-  const message = (choices[0] as { message?: Record<string, unknown> })?.message ?? {};
+  const choices = Array.isArray(payload.choices) ? payload.choices : [];
+  const first = choices[0];
+  const message = object.isRecord(first) && object.isRecord(first.message) ? first.message : {};
   const answer = openaiChat.chatContentToText(message.content);
   const citations: WebSearchCitation[] = [];
   // Best-effort grounding extraction: walk any nested object for {uri|url,title}.
   const seen = new Set<string>();
   const visit = (v: unknown, depth: number): void => {
-    if (depth > 6 || v === null || typeof v !== "object") return;
-    const o = v as Record<string, unknown>;
-    const url = string.trimToEmpty(o.url) || string.trimToEmpty(o.uri);
+    if (depth > MAX_GROUNDING_WALK_DEPTH || !object.isRecord(v)) return;
+    const url = string.trimToEmpty(v.url) || string.trimToEmpty(v.uri);
     if (url && !seen.has(url)) {
       seen.add(url);
-      const title = string.trimToEmpty(o.title);
+      const title = string.trimToEmpty(v.title);
       citations.push({ url, ...(title ? { title } : {}) });
     }
-    for (const val of Object.values(o)) {
+    for (const val of Object.values(v)) {
       if (Array.isArray(val)) val.forEach((x) => visit(x, depth + 1));
       else if (val && typeof val === "object") visit(val, depth + 1);
     }
@@ -179,11 +258,14 @@ function fromChatPayload(payload: Record<string, unknown>): {
  * workspace has no such endpoint AND the scrape fallback is enabled, falls
  * back to a DuckDuckGo scrape so the tool still returns results instead of
  * erroring. Citations are filtered through the configured URL allow-list.
+ *
+ * `signal` cancels the whole call, including the in-flight serving request.
  */
 export async function runWebSearch(
   request: WebSearchRequest,
   config: ResolvedWebSearchConfig,
   ctx: WebSearchContext,
+  signal?: AbortSignal,
 ): Promise<WebSearchResult> {
   const modelId = await resolveWebSearchModel(ctx, config, request.model);
 
@@ -191,20 +273,21 @@ export async function runWebSearch(
     // No native web-search model deployed in this workspace.
     if (config.scrapeFallback) {
       logger.info("no-native-model:scrape-fallback", { query: request.query });
-      return runScrapeSearch(request, config);
+      return runScrapeSearch(request, config, signal);
     }
-    throw new Error(
-      "web-search: no web-search-capable model (GPT/Gemini) is deployed in this workspace, " +
-        "and the scrape fallback is disabled. Deploy a supported endpoint, set `model` / " +
-        "WEB_SEARCH_MODEL, or enable the fallback (WEB_SEARCH_SCRAPE_FALLBACK=1).",
+    throw ConfigurationError.resourceNotFound(
+      "Web-search-capable serving endpoint",
+      "No GPT/Gemini endpoint is deployed in this workspace and the scrape fallback is " +
+        `disabled. Deploy a supported endpoint, set model or ${MODEL_ENV}, or enable the ` +
+        "fallback (WEB_SEARCH_SCRAPE_FALLBACK=1).",
     );
   }
 
   const provider = detectWebSearchProvider(modelId)!; // guaranteed by resolve step
   const spec = webSearchToolSpec(provider, config.webSearchTools);
 
-  const path =
-    spec.api === "responses" ? `/${invoke.RESPONSES_PATH}` : `/${invoke.CHAT_COMPLETIONS_PATH}`;
+  const url =
+    spec.api === "responses" ? invoke.responsesUrl(ctx.host) : invoke.chatCompletionsUrl(ctx.host);
   const body =
     spec.api === "responses"
       ? {
@@ -218,13 +301,14 @@ export async function runWebSearch(
           tools: [spec.tool],
         };
 
-  let payload: Record<string, unknown>;
-  try {
-    payload = await postServing(ctx, path, body);
-  } catch (err) {
-    logger.warn("serving-error", { model: modelId, provider, error: error.errorMessage(err) });
-    throw err;
-  }
+  const payload = await postServing(
+    ctx,
+    url,
+    body,
+    config,
+    ["web-search", "serving", spec.api, modelId, request.query],
+    signal,
+  );
 
   const { answer, citations } =
     spec.api === "responses" ? fromResponsesPayload(payload) : fromChatPayload(payload);

package/src/tool.ts

@@ -15,12 +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 { 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 {
@@ -28,11 +37,25 @@ import {
   webFetchResultSchema,
   webSearchRequestSchema,
   webSearchResultSchema,
-  type WebFetchRequest,
-  type WebSearchRequest,
+  WEB_FETCH_TOOL_DESCRIPTION,
+  WEB_SEARCH_TOOL_DESCRIPTION,
 } from "./schema";
 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 {
   /** Override the tool id. */
@@ -42,12 +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;
+function effectiveGate(opts: WebSearchToolOptions): ApprovalPolicy {
+  return opts.approval === undefined
+    ? getWebSearchRuntime().config.approval
+    : toApprovalPolicy(opts.approval);
 }
 
 /**
@@ -69,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 resolveWebSearchContext());
+      const request = parseToolInput(webSearchRequestSchema, input);
+      return runWebSearch(request, config, await resolveWebSearchContext());
     },
   });
 }
@@ -106,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,7 +6,7 @@ import {
   parseAllowedUrls,
   toUrlAllowList,
 } from "../src/allowlist";
-import { approvalMatches, resolveWebSearchConfig } from "../src/config";
+import { approvalMatches, resolveWebSearchConfig, toApprovalPolicy } from "../src/config";
 import { htmlToText } from "../src/html-text";
 import {
   detectWebSearchProvider,
@@ -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", () => {