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

package/src/plugin.ts

@@ -1,31 +1,105 @@
 /**
  * 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 { 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 { 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",
@@ -36,29 +110,111 @@ export class WebSearchPlugin extends Plugin<WebSearchPluginConfig> {
     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,
     });
   }
 
+  /**
+   * 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() {
     return {
       /**
@@ -66,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 resolveWebSearchContext()),
+      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,
   };
 }

package/src/runtime.ts

@@ -5,23 +5,64 @@
  * plugin at setup) primes it from the plugin's config; later callers (the
  * tools' `execute`) reuse it.
  *
+ * The runtime also carries the {@link WebSearchExecutor} every outbound call
+ * runs through. The plugin installs its own `execute()` there at setup, which
+ * is how the tools - plain functions with no plugin instance in scope - still
+ * get AppKit's cache / retry / timeout / telemetry chain. Without a
+ * registered plugin (a direct call from a script or a test) the calls still
+ * run, just without interceptors.
+ *
  * Unlike the email runtime there is no connection to pool - the backend is
- * stateless HTTP per call - so the runtime holds only the resolved config.
+ * stateless HTTP per call - so the runtime holds only the resolved config and
+ * that executor.
  *
  * @module
  */
 
+import { AppKitError, ExecutionError, type ExecutionResult } from "@databricks/appkit";
+import { async, error, log } from "@dbx-tools/shared-core";
 import {
   resolveWebSearchConfig,
   type ResolvedWebSearchConfig,
   type WebSearchPluginConfig,
 } from "./config";
+import type { WebSearchExecutionSettings } from "./defaults";
+
+const logger = log.logger("web-search/runtime");
+
+/**
+ * Runs one outbound call through AppKit's interceptor chain. Matches
+ * `Plugin.execute()`, which never throws: a failure comes back as
+ * `{ ok: false }`.
+ */
+export type WebSearchExecutor = <T>(
+  fn: (signal?: AbortSignal) => Promise<T>,
+  settings: WebSearchExecutionSettings,
+) => Promise<ExecutionResult<T>>;
 
-/** The shared resolved config. */
+/** The shared resolved config plus the executor outbound calls run through. */
 export interface WebSearchRuntime {
   config: ResolvedWebSearchConfig;
+  execute: WebSearchExecutor;
 }
 
+/**
+ * Executor used until (or unless) the plugin installs its own: run the call
+ * directly, mapping a throw onto the same {@link ExecutionResult} shape so
+ * call sites branch on `ok` either way.
+ */
+const directExecute: WebSearchExecutor = async (fn) => {
+  try {
+    return { ok: true, data: await fn() };
+  } catch (err) {
+    return {
+      ok: false,
+      status: err instanceof AppKitError ? err.statusCode : 500,
+      message: error.errorMessage(err),
+    };
+  }
+};
+
 let runtime: WebSearchRuntime | undefined;
 
 /**
@@ -33,12 +74,55 @@ let runtime: WebSearchRuntime | undefined;
  */
 export function getWebSearchRuntime(overrides?: WebSearchPluginConfig): WebSearchRuntime {
   if (!runtime) {
-    runtime = { config: resolveWebSearchConfig(overrides) };
+    runtime = { config: resolveWebSearchConfig(overrides), execute: directExecute };
   }
   return runtime;
 }
 
+/**
+ * Install the executor outbound calls run through. The plugin calls this at
+ * setup with its own `execute()`; a second call replaces the previous one, so
+ * a re-registered plugin does not leave the tools bound to a dead instance.
+ */
+export function setWebSearchExecutor(execute: WebSearchExecutor): void {
+  getWebSearchRuntime().execute = execute;
+}
+
 /** Drop the memoized runtime so the next {@link getWebSearchRuntime} rebuilds it. */
 export function resetWebSearchRuntime(): void {
   runtime = undefined;
 }
+
+/**
+ * Run one idempotent read through the shared executor and unwrap it.
+ *
+ * `execute()` never throws, so a failed call arrives as `{ ok: false }` with
+ * a status the interceptors already sanitized; it is logged here and re-raised
+ * as a stable {@link ExecutionError} so an upstream message never becomes the
+ * caller's error text. `signal` is the caller's own cancellation (an agent
+ * run, a request teardown); it is merged with the signal the timeout
+ * interceptor supplies so either one unwinds the I/O.
+ */
+export async function executeRead<T>(
+  operation: string,
+  settings: WebSearchExecutionSettings,
+  fn: (signal?: AbortSignal) => Promise<T>,
+  signal?: AbortSignal,
+): Promise<T> {
+  const { execute } = getWebSearchRuntime();
+  const result = await execute(
+    (executeSignal) => fn(async.combineAbortSignals(executeSignal, signal)),
+    settings,
+  );
+  if (result.ok) return result.data;
+  // A caller that cancelled is not a failure worth reporting as one.
+  if (signal?.aborted) throw ExecutionError.canceled();
+  logger.warn("execution-failed", {
+    operation,
+    status: result.status,
+    error: result.message,
+  });
+  throw new ExecutionError(`web-search: ${operation} failed`, {
+    context: { operation, status: result.status },
+  });
+}

package/src/schema.ts

@@ -20,6 +20,28 @@
 import { string } from "@dbx-tools/shared-core";
 import { z } from "zod";
 
+/**
+ * Description the model reads for `web_search`. Shared by the Mastra tool and
+ * the AppKit tool provider so both hosts describe the tool identically.
+ */
+export const WEB_SEARCH_TOOL_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 the model reads for `web_fetch`. Shared like {@link WEB_SEARCH_TOOL_DESCRIPTION}. */
+export const WEB_FETCH_TOOL_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.
+`);
+
 /** Schema for the `web_search` tool input. */
 export const webSearchRequestSchema = z.object({
   query: z

package/src/scrape.ts

@@ -21,7 +21,9 @@
 import { log } from "@dbx-tools/shared-core";
 import { gotScraping } from "got-scraping";
 import type { ResolvedWebSearchConfig } from "./config";
+import { scrapeSearchExecuteDefaults, toCallSettings } from "./defaults";
 import { htmlFragmentToText } from "./html-text";
+import { executeRead } from "./runtime";
 import type { WebSearchCitation, WebSearchRequest, WebSearchResult } from "./schema";
 
 const logger = log.logger("web-search/scrape");
@@ -71,21 +73,40 @@ function parseDdgHtml(html: string): WebSearchCitation[] {
  * Run a scraping search over DuckDuckGo. Returns the same
  * {@link WebSearchResult} shape as the native path, with `model` set to
  * `"scrape:duckduckgo"` so callers can tell how the result was produced.
- * Citations are filtered through the configured URL allow-list.
+ * Citations are filtered through the configured URL allow-list. `signal`
+ * cancels the in-flight request.
  */
 export async function runScrapeSearch(
   request: WebSearchRequest,
   config: ResolvedWebSearchConfig,
+  signal?: AbortSignal,
 ): Promise<WebSearchResult> {
-  const response = await gotScraping({
-    url: `${DDG_HTML_URL}?q=${encodeURIComponent(request.query)}`,
-    method: "GET",
-    timeout: { request: config.timeoutMs },
-    throwHttpErrors: false,
-    followRedirect: true,
-  });
-  const body = typeof response.body === "string" ? response.body : String(response.body ?? "");
-  const all = parseDdgHtml(body);
+  const page = await executeRead(
+    "scrape-search",
+    toCallSettings(scrapeSearchExecuteDefaults, config.timeoutMs, [
+      "web-search",
+      "scrape",
+      request.query,
+    ]),
+    async (executeSignal): Promise<{ body: string; statusCode: number }> => {
+      const response = await gotScraping({
+        url: `${DDG_HTML_URL}?q=${encodeURIComponent(request.query)}`,
+        method: "GET",
+        // 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 {
+        body: typeof response.body === "string" ? response.body : String(response.body ?? ""),
+        statusCode: response.statusCode,
+      };
+    },
+    signal,
+  );
+  const all = parseDdgHtml(page.body);
   const permitted = all.filter((c) => config.allowList.allows(c.url));
   const citations = permitted.slice(0, config.maxCitations);
   const answer =
@@ -99,7 +120,7 @@ export async function runScrapeSearch(
     query: request.query,
     found: all.length,
     returned: citations.length,
-    status: response.statusCode,
+    status: page.statusCode,
   });
   return { query: request.query, answer, citations, model: "scrape:duckduckgo" };
 }