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

package/src/fetch.ts

@@ -19,50 +19,28 @@ import { log } from "@dbx-tools/shared-core";
 import { gotScraping } from "got-scraping";
 import { assertUrlAllowed } from "./allowlist";
 import type { ResolvedWebSearchConfig } from "./config";
+import { toCallSettings, webFetchExecuteDefaults } from "./defaults";
+import { decodeHtmlEntities, htmlToText } from "./html-text";
+import { executeRead } from "./runtime";
 import type { WebFetchRequest, WebFetchResult } from "./schema";
 
 const logger = log.logger("web-search/fetch");
 
+/** The slice of a got-scraping response this module reads. */
+interface FetchedPage {
+  url: string;
+  statusCode: number;
+  contentType: string | undefined;
+  body: string;
+}
+
 /** Pull the <title> text out of an HTML document, when present. */
 function extractTitle(html: string): string | undefined {
   const match = /<title[^>]*>([\s\S]*?)<\/title>/i.exec(html);
-  const title = match?.[1] ? decodeEntities(match[1]).trim() : "";
+  const title = match?.[1] ? decodeHtmlEntities(match[1]).trim() : "";
   return title.length > 0 ? title : undefined;
 }
 
-/** Decode the handful of HTML entities that survive tag-stripping. */
-function decodeEntities(text: string): string {
-  return text
-    .replace(/&nbsp;/g, " ")
-    .replace(/&amp;/g, "&")
-    .replace(/&lt;/g, "<")
-    .replace(/&gt;/g, ">")
-    .replace(/&quot;/g, '"')
-    .replace(/&#39;/g, "'")
-    .replace(/&#(\d+);/g, (_, code) => String.fromCodePoint(Number(code)));
-}
-
-/**
- * Reduce an HTML document to readable plain text: drop `<script>` /
- * `<style>` / `<noscript>` blocks and HTML comments, turn block-level tags
- * into newlines, strip the remaining tags, decode entities, and collapse
- * runs of blank lines / trailing spaces.
- */
-export function htmlToText(html: string): string {
-  return decodeEntities(
-    html
-      .replace(/<!--[\s\S]*?-->/g, "")
-      .replace(/<(script|style|noscript)[^>]*>[\s\S]*?<\/\1>/gi, "")
-      .replace(/<\/(p|div|section|article|li|tr|h[1-6]|header|footer|br)>/gi, "\n")
-      .replace(/<br\s*\/?>/gi, "\n")
-      .replace(/<[^>]+>/g, ""),
-  )
-    .replace(/[ \t]+\n/g, "\n")
-    .replace(/\n{3,}/g, "\n\n")
-    .replace(/[ \t]{2,}/g, " ")
-    .trim();
-}
-
 /** Truncate `text` to `max` chars, reporting whether it was cut. */
 function truncate(text: string, max: number): { content: string; truncated: boolean } {
   if (text.length <= max) return { content: text, truncated: false };
@@ -73,43 +51,60 @@ function truncate(text: string, max: number): { content: string; truncated: bool
  * Fetch a single URL and return its content in the requested format.
  *
  * Throws when the URL is not permitted by the allow-list (the visible,
- * correctable failure the design calls for on the fetch path). Network /
- * HTTP errors propagate from got-scraping. The request's `maxLength` narrows
- * (never widens) the plugin's `fetchMaxLength` cap.
+ * correctable failure the design calls for on the fetch path). The request's
+ * `maxLength` narrows (never widens) the plugin's `fetchMaxLength` cap.
+ * `signal` cancels the in-flight fetch.
  */
 export async function runWebFetch(
   request: WebFetchRequest,
   config: ResolvedWebSearchConfig,
+  signal?: AbortSignal,
 ): Promise<WebFetchResult> {
   assertUrlAllowed(request.url, config.allowList);
   const cap = Math.min(request.maxLength ?? config.fetchMaxLength, config.fetchMaxLength);
 
-  const response = await gotScraping({
-    url: request.url,
-    timeout: { request: config.timeoutMs },
-    throwHttpErrors: false,
-    followRedirect: true,
-  });
+  // The response is cached before the format reduction, so the same page read
+  // as text and as html shares one network round trip.
+  const page = await executeRead(
+    "page-fetch",
+    toCallSettings(webFetchExecuteDefaults, config.timeoutMs, ["web-search", "fetch", request.url]),
+    async (executeSignal): Promise<FetchedPage> => {
+      const response = await gotScraping({
+        url: request.url,
+        // got's own timeout aborts the socket and reports which phase timed
+        // out; the interceptor timeout bounds the whole attempt around it.
+        timeout: { request: config.timeoutMs },
+        throwHttpErrors: false,
+        followRedirect: true,
+        ...(executeSignal ? { signal: executeSignal } : {}),
+      });
+      return {
+        url: response.url ?? request.url,
+        statusCode: response.statusCode,
+        contentType: response.headers["content-type"],
+        body: typeof response.body === "string" ? response.body : String(response.body ?? ""),
+      };
+    },
+    signal,
+  );
 
-  const body = typeof response.body === "string" ? response.body : String(response.body ?? "");
-  const contentType = response.headers["content-type"];
-  const isHtml = !contentType || /html|xml/i.test(contentType);
-  const rawContent = request.format === "html" || !isHtml ? body : htmlToText(body);
+  const isHtml = !page.contentType || /html|xml/i.test(page.contentType);
+  const rawContent = request.format === "html" || !isHtml ? page.body : htmlToText(page.body);
   const { content, truncated } = truncate(rawContent, cap);
-  const title = isHtml ? extractTitle(body) : undefined;
+  const title = isHtml ? extractTitle(page.body) : undefined;
 
   logger.debug("fetched", {
-    url: response.url,
-    status: response.statusCode,
-    bytes: body.length,
+    url: page.url,
+    status: page.statusCode,
+    bytes: page.body.length,
     returned: content.length,
     ...(truncated ? { truncated: true } : {}),
   });
 
   return {
-    url: response.url ?? request.url,
-    status: response.statusCode,
-    ...(contentType ? { contentType } : {}),
+    url: page.url,
+    status: page.statusCode,
+    ...(page.contentType ? { contentType: page.contentType } : {}),
     ...(title ? { title } : {}),
     content,
     truncated,

package/src/html-text.ts

@@ -0,0 +1,72 @@
+/**
+ * Turning fetched HTML into the plain text a model can read.
+ *
+ * Both scraping paths need this: `runWebFetch`'s full-page read and
+ * `runScrapeSearch`'s title/snippet extraction. The entity table is shared so
+ * the two cannot drift apart.
+ *
+ * This is deliberately regex-based rather than a DOM parse. The inputs are
+ * whole pages fetched for summarization, not documents to be queried, so the
+ * cost of a real parser buys nothing - and a malformed page still degrades to
+ * readable text instead of throwing.
+ *
+ * @module
+ */
+
+/**
+ * The named entities that survive tag-stripping in practice, plus numeric
+ * escapes. Not a complete HTML entity table - anything rarer passes through
+ * as-is, which reads acceptably in model input.
+ */
+const NAMED_ENTITIES: Record<string, string> = {
+  "&nbsp;": " ",
+  "&amp;": "&",
+  "&lt;": "<",
+  "&gt;": ">",
+  "&quot;": '"',
+  "&#39;": "'",
+  "&#x27;": "'",
+  "&apos;": "'",
+};
+
+const NAMED_ENTITY_REGEXP = new RegExp(Object.keys(NAMED_ENTITIES).join("|"), "g");
+const NUMERIC_ENTITY_REGEXP = /&#(\d+);/g;
+const TAG_REGEXP = /<[^>]+>/g;
+
+/** Decode the HTML entities that survive tag-stripping. */
+export function decodeHtmlEntities(text: string): string {
+  return text
+    .replace(NAMED_ENTITY_REGEXP, (entity) => NAMED_ENTITIES[entity] ?? entity)
+    .replace(NUMERIC_ENTITY_REGEXP, (_, code) => String.fromCodePoint(Number(code)));
+}
+
+/**
+ * Strip tags and decode entities from a short HTML fragment, collapsing all
+ * whitespace to single spaces. For inline snippets (a search result title or
+ * summary), where layout carries no meaning.
+ */
+export function htmlFragmentToText(html: string): string {
+  return decodeHtmlEntities(html.replace(TAG_REGEXP, "")).replace(/\s+/g, " ").trim();
+}
+
+/**
+ * Reduce a full HTML document to readable plain text: drop `<script>` /
+ * `<style>` / `<noscript>` blocks and comments, turn block-level tags into
+ * newlines, strip the remaining tags, decode entities, and collapse runs of
+ * blank lines / trailing spaces. Unlike {@link htmlFragmentToText}, this
+ * preserves line structure because paragraph breaks carry meaning in a page.
+ */
+export function htmlToText(html: string): string {
+  return decodeHtmlEntities(
+    html
+      .replace(/<!--[\s\S]*?-->/g, "")
+      .replace(/<(script|style|noscript)[^>]*>[\s\S]*?<\/\1>/gi, "")
+      .replace(/<\/(p|div|section|article|li|tr|h[1-6]|header|footer|br)>/gi, "\n")
+      .replace(/<br\s*\/?>/gi, "\n")
+      .replace(TAG_REGEXP, ""),
+  )
+    .replace(/[ \t]+\n/g, "\n")
+    .replace(/\n{3,}/g, "\n\n")
+    .replace(/[ \t]{2,}/g, " ")
+    .trim();
+}

package/src/plugin.ts

@@ -1,68 +1,218 @@
 /**
  * AppKit plugin (registered name: `web-search`) that owns the resolved
- * web-search runtime - the URL allow-list, result / length caps, timeout,
- * and default approval gate the {@link webSearchTool} / {@link webFetchTool}
- * read. Registering it resolves and logs the effective config (whether an
- * allow-list is active, the caps) so a misconfiguration is visible in the
- * boot logs rather than on the first search.
+ * web-search runtime - the URL policy, result / length caps, timeout, and
+ * default approval gate the {@link webSearchTool} / {@link webFetchTool}
+ * read. Registering it resolves and logs the effective config (which URL
+ * policy is in force, the caps) so a misconfiguration is visible in the boot
+ * logs rather than on the first search, and installs the plugin's
+ * `execute()` as the runtime's executor so every outbound call picks up
+ * AppKit's cache / retry / timeout / telemetry chain.
  *
- * The tools do the actual work when spread into an agent; this plugin primes
- * the shared runtime they reuse and exposes direct {@link runWebSearch} /
- * {@link runWebFetch} for non-agent callers.
+ * The plugin also implements AppKit's `ToolProvider`, so an AppKit agent gets
+ * `web_search` / `web_fetch` without going through Mastra; the Mastra tools in
+ * `tool.ts` are the other half and share the same runtime.
  *
  * @module
  */
 
-import { getExecutionContext, Plugin, toPlugin, type PluginManifest } from "@databricks/appkit";
-import { log } from "@dbx-tools/shared-core";
-import { WEB_SEARCH_CONFIG_SCHEMA, type WebSearchPluginConfig } from "./config";
+import {
+  Plugin,
+  ResourceType,
+  toPlugin,
+  type PluginManifest,
+  type ResourceRequirement,
+} from "@databricks/appkit";
+import {
+  defineTool,
+  executeFromRegistry,
+  toolsFromRegistry,
+  type AgentToolDefinition,
+  type ToolProvider,
+  type ToolRegistry,
+} from "@databricks/appkit/beta";
+import { log, string } from "@dbx-tools/shared-core";
+import {
+  MODEL_ENV,
+  SERVING_ENDPOINT_ENV,
+  WEB_SEARCH_CONFIG_SCHEMA,
+  type WebSearchPluginConfig,
+} from "./config";
 import { runWebFetch } from "./fetch";
-import { getWebSearchRuntime } from "./runtime";
+import { getWebSearchRuntime, resetWebSearchRuntime, setWebSearchExecutor } from "./runtime";
+import {
+  webFetchRequestSchema,
+  webSearchRequestSchema,
+  WEB_FETCH_TOOL_DESCRIPTION,
+  WEB_SEARCH_TOOL_DESCRIPTION,
+} from "./schema";
 import type { WebFetchRequest, WebFetchResult, WebSearchRequest, WebSearchResult } from "./schema";
-import { runWebSearch, type WebSearchContext } from "./search";
+import { resolveWebSearchContext, runWebSearch } from "./search";
+
+const logger = log.logger("web-search");
+
+/**
+ * The Model Serving endpoint the native web-search tool runs on. Declared
+ * optional because the plugin resolves a web-search-capable endpoint from the
+ * live catalogue when nothing is pinned; {@link WebSearchPlugin.getResourceRequirements}
+ * promotes it to required once a deployment names one. `CAN_QUERY` is the
+ * weakest permission that can invoke an endpoint.
+ */
+const SERVING_ENDPOINT_RESOURCE = {
+  type: ResourceType.SERVING_ENDPOINT,
+  alias: "Web Search Endpoint",
+  resourceKey: "web-search-endpoint",
+  description:
+    "Model Serving endpoint running the native web-search tool. Optional: resolved from the " +
+    "workspace catalogue (Gemini, then GPT) when no endpoint is pinned.",
+  permission: "CAN_QUERY",
+  fields: {
+    name: {
+      env: SERVING_ENDPOINT_ENV,
+      description: `Serving endpoint name for web search. ${MODEL_ENV} overrides it.`,
+      discovery: {
+        type: "cli",
+        cliCommand: "databricks serving-endpoints list --profile <PROFILE> --output json",
+        selectField: ".name",
+      },
+    },
+  },
+} satisfies Omit<ResourceRequirement, "required">;
 
 /**
  * AppKit plugin that resolves and holds the web-search runtime config used
  * by the `web_search` / `web_fetch` tools.
+ *
+ * @example
+ * ```ts
+ * import { createApp, server } from "@databricks/appkit";
+ * import { plugin as webSearchPlugin } from "@dbx-tools/appkit-web-search";
+ *
+ * await createApp({
+ *   plugins: [
+ *     server(),
+ *     webSearchPlugin.webSearch({
+ *       model: "gemini",
+ *       urlPolicy: "allowlist",
+ *       allowedUrls: ["*.databricks.com"],
+ *     }),
+ *   ],
+ * });
+ * ```
  */
-export class WebSearchPlugin extends Plugin<WebSearchPluginConfig> {
+export class WebSearchPlugin extends Plugin<WebSearchPluginConfig> implements ToolProvider {
   static manifest = {
     name: "web-search",
     displayName: "Web Search",
     description:
-      "Searches the web (via duck-duck-scrape) and fetches pages (via got-scraping), " +
-      "with an optional URL allow-list and per-tool approval gating.",
+      "Searches the web through the Databricks native web-search tool on its own " +
+      "web-capable serving endpoint, and fetches pages via got-scraping, behind an " +
+      "optional URL allow-list and per-tool approval gating.",
     stability: "beta",
     resources: {
       required: [],
-      optional: [],
+      optional: [SERVING_ENDPOINT_RESOURCE],
     },
     config: { schema: WEB_SEARCH_CONFIG_SCHEMA },
   } satisfies PluginManifest<"web-search">;
 
-  private logger = log.logger(this);
+  /**
+   * Promote the serving endpoint to a required resource once a deployment
+   * pins one, through plugin config or either environment name. Left optional
+   * otherwise, because the plugin picks a web-search-capable endpoint out of
+   * the live catalogue on its own.
+   */
+  static getResourceRequirements(config: WebSearchPluginConfig): ResourceRequirement[] {
+    const pinned =
+      string.trimToNull(config.model) ??
+      string.trimToNull(process.env[MODEL_ENV]) ??
+      string.trimToNull(process.env[SERVING_ENDPOINT_ENV]);
+    return pinned === null ? [] : [{ ...SERVING_ENDPOINT_RESOURCE, required: true }];
+  }
+
+  /**
+   * The tools this plugin offers to an AppKit agent. Both are reads and both
+   * run under the caller's identity.
+   *
+   * Neither is `autoInheritable`. `web_fetch` reaches a URL the model chose,
+   * which is reachable from inside the workspace network, so it must only
+   * appear in an agent that asked for it and accepted the URL policy that
+   * comes with it. `web_search` is safer but still spends serving tokens on
+   * every call, so it is opt-in for cost rather than for safety.
+   *
+   * `execute` re-parses its arguments with the local schema: AppKit validates
+   * against the same schema first, but re-parsing is what gives the body typed
+   * arguments instead of `unknown`.
+   */
+  private readonly tools: ToolRegistry = {
+    web_search: defineTool({
+      description: WEB_SEARCH_TOOL_DESCRIPTION,
+      schema: webSearchRequestSchema,
+      annotations: { effect: "read", requiresUserContext: true },
+      autoInheritable: false,
+      execute: async (args, signal) => this.search(webSearchRequestSchema.parse(args), signal),
+    }),
+    web_fetch: defineTool({
+      description: WEB_FETCH_TOOL_DESCRIPTION,
+      schema: webFetchRequestSchema,
+      annotations: { effect: "read", requiresUserContext: true },
+      autoInheritable: false,
+      execute: async (args, signal) => this.fetch(webFetchRequestSchema.parse(args), signal),
+    }),
+  };
 
   /**
-   * Prime the shared runtime from this plugin's config (over env) and log
-   * the effective policy so an active allow-list / caps are obvious at boot.
+   * Prime the shared runtime from this plugin's config (over env), route the
+   * tools' outbound calls through this plugin's interceptor chain, and log the
+   * effective policy so an active allow-list / caps are obvious at boot.
    */
   override async setup(): Promise<void> {
     const { config } = getWebSearchRuntime(this.config);
-    this.logger.info("ready", {
+    setWebSearchExecutor((fn, settings) => this.execute(fn, settings));
+    logger.info("ready", {
       model: config.model ?? `fallbacks:[${config.modelFallbacks.join(", ")}]`,
-      restricted: config.allowList.restricted,
+      modelSource: config.modelSource,
+      urlPolicy: config.urlPolicy,
       ...(config.allowList.restricted ? { allowedUrls: config.allowList.patterns } : {}),
       maxCitations: config.maxCitations,
       fetchMaxLength: config.fetchMaxLength,
-      approval: config.approval === false ? "none" : config.approval,
+      approval: config.approval.mode,
+      scrapeFallback: config.scrapeFallback,
     });
   }
 
-  /** Resolve the OBO client + host from the active execution context. */
-  private async searchContext(): Promise<WebSearchContext> {
-    const ctx = getExecutionContext();
-    const host = (await ctx.client.config.getHost()).toString();
-    return { client: ctx.client, host };
+  /**
+   * Drop the shared runtime so a restarted app re-resolves config and does not
+   * keep calling through a torn-down plugin's `execute()`. Bounded and
+   * idempotent: there is no connection to drain, only the memo to clear.
+   */
+  async shutdown(): Promise<void> {
+    resetWebSearchRuntime();
+  }
+
+  /**
+   * Abort in-flight work. AppKit's graceful shutdown only invokes this hook -
+   * it never calls {@link shutdown} - so the runtime memo is dropped from here
+   * to keep a restarted app from calling through a torn-down `execute()`. The
+   * teardown is synchronous and idempotent, so the un-awaited call costs
+   * nothing.
+   */
+  override abortActiveOperations(): void {
+    super.abortActiveOperations();
+    void this.shutdown();
+  }
+
+  /** AppKit `ToolProvider`: the tool definitions offered to an agent. */
+  getAgentTools(): AgentToolDefinition[] {
+    return toolsFromRegistry(this.tools);
+  }
+
+  /**
+   * AppKit `ToolProvider`: run one tool call. Arguments are validated against
+   * the tool's schema first, and a validation failure comes back as an
+   * LLM-friendly string so the model can correct itself on the next turn.
+   */
+  async executeAgentTool(name: string, args: unknown, signal?: AbortSignal): Promise<unknown> {
+    return executeFromRegistry(this.tools, name, args, signal);
   }
 
   override exports() {
@@ -72,16 +222,52 @@ export class WebSearchPlugin extends Plugin<WebSearchPluginConfig> {
        * OBO client from the active execution context and reads the shared
        * runtime config primed at setup.
        */
-      search: async (request: WebSearchRequest): Promise<WebSearchResult> =>
-        runWebSearch(request, getWebSearchRuntime().config, await this.searchContext()),
+      search: (request: WebSearchRequest, signal?: AbortSignal): Promise<WebSearchResult> =>
+        this.search(request, signal),
       /**
        * Fetch one URL directly (bypassing the agent tool). Enforces the
-       * configured allow-list. Reads the shared runtime config.
+       * configured URL policy. Reads the shared runtime config.
        */
-      fetch: (request: WebFetchRequest): Promise<WebFetchResult> =>
-        runWebFetch(request, getWebSearchRuntime().config),
+      fetch: (request: WebFetchRequest, signal?: AbortSignal): Promise<WebFetchResult> =>
+        this.fetch(request, signal),
     };
   }
+
+  private async search(request: WebSearchRequest, signal?: AbortSignal): Promise<WebSearchResult> {
+    return runWebSearch(
+      request,
+      getWebSearchRuntime().config,
+      await resolveWebSearchContext(),
+      signal,
+    );
+  }
+
+  private async fetch(request: WebFetchRequest, signal?: AbortSignal): Promise<WebFetchResult> {
+    return runWebFetch(request, getWebSearchRuntime().config, signal);
+  }
 }
 
+/**
+ * Register the web-search runtime with AppKit.
+ *
+ * @example
+ * ```ts
+ * import { createApp, server } from "@databricks/appkit";
+ * import { plugin as webSearchPlugin, tool as webTool } from "@dbx-tools/appkit-web-search";
+ * import { agents, plugin as mastraPlugin } from "@dbx-tools/appkit-mastra";
+ *
+ * const researcher = agents.createAgent({
+ *   instructions: "Research questions with web_search, then read sources with web_fetch.",
+ *   tools: () => ({ web_search: webTool.webSearchTool(), web_fetch: webTool.webFetchTool() }),
+ * });
+ *
+ * await createApp({
+ *   plugins: [
+ *     server(),
+ *     webSearchPlugin.webSearch({ model: "gemini" }),
+ *     mastraPlugin.mastra({ agents: researcher }),
+ *   ],
+ * });
+ * ```
+ */
 export const webSearch = toPlugin(WebSearchPlugin);

package/src/provider.ts

@@ -23,6 +23,9 @@
  * @module
  */
 
+import { ValidationError } from "@databricks/appkit";
+import { z } from "zod";
+
 /** A web-search-capable model provider family. */
 export type WebSearchProvider = "openai" | "gemini";
 
@@ -70,21 +73,41 @@ export function supportsWebSearch(modelId: string): boolean {
   return detectWebSearchProvider(modelId) !== null;
 }
 
+/**
+ * Runtime shape of one entry in the operator override map. The map arrives as
+ * parsed JSON (config or `WEB_SEARCH_TOOLS`), so it is validated rather than
+ * asserted.
+ */
+const providerOverrideSchema = z.object({
+  api: z.enum(["responses", "chat"]).optional(),
+  tool: z.record(z.string(), z.unknown()).optional(),
+});
+
 /**
  * Resolve the effective {@link WebSearchProviderSpec} for a provider: the
  * built-in default, with any operator override (the `webSearchTools` map,
  * keyed by provider) shallow-merged over it. An override may replace just the
- * `tool` (the common case - a new tool type) or also the `api`.
+ * `tool` (the common case - a new tool type) or also the `api`. An override
+ * that is not one of those two fields is a deployment mistake that would
+ * otherwise be silently dropped, so it throws.
  */
 export function webSearchToolSpec(
   provider: WebSearchProvider,
   overrides?: Record<string, unknown>,
 ): WebSearchProviderSpec {
   const base = WEB_SEARCH_PROVIDERS[provider];
-  const override = overrides?.[provider] as Partial<WebSearchProviderSpec> | undefined;
-  if (!override) return base;
+  const raw = overrides?.[provider];
+  if (raw === undefined) return base;
+  const parsed = providerOverrideSchema.safeParse(raw);
+  if (!parsed.success) {
+    throw ValidationError.invalidValue(
+      `webSearchTools.${provider}`,
+      raw,
+      'an object with an optional "api" ("responses" | "chat") and an optional "tool" object',
+    );
+  }
   return {
-    api: override.api ?? base.api,
-    tool: override.tool ?? base.tool,
+    api: parsed.data.api ?? base.api,
+    tool: parsed.data.tool ?? base.tool,
   };
 }