pi-web/index.ts

// Pi-Web Extension - Registers `web_search` and `web_fetch` tools backed by
// a shared headless Firefox session. Provider config lives in
// ~/.pi/pi-web/config.json (env overrides supported).
import {
  keyHint,
  type AgentToolResult,
  type ExtensionAPI,
  type Theme,
} from "@mariozechner/pi-coding-agent";
import { Text } from "@mariozechner/pi-tui";
import { Type } from "typebox";
import { ConfigError, loadConfig, resolveSettings } from "./src/config.ts";
import { fetchPage, type FetchResult } from "./src/fetch.ts";
import { searchKagi } from "./src/providers/kagi.ts";
import { searchSearxng } from "./src/providers/searxng.ts";
import { SearchError, type SearchResult } from "./src/types.ts";

// Collapsed-view preview cap. The full list still goes to the model via `content`.
const PREVIEW_RESULTS = 3;

function formatResults(results: SearchResult[]): string {
  if (results.length === 0) return "(no results)";
  return results
    .map((r) => `## [${r.title}](${r.url})\n> ${r.description}`)
    .join("\n\n");
}

function renderSearchResult(
  result: AgentToolResult<unknown>,
  expanded: boolean,
  theme: Theme,
): Text {
  const fullText = result.content[0]?.type === "text" ? result.content[0].text : "";
  const raw = (result.details as { raw?: SearchResult[] } | undefined)?.raw;

  // Expanded or unavailable details: show the full markdown the model received.
  if (expanded || !raw || raw.length === 0) {
    return new Text(fullText, 0, 0);
  }

  let text = theme.fg("muted", `${raw.length} result(s):`);
  for (const r of raw.slice(0, PREVIEW_RESULTS)) {
    text += `\n${theme.fg("toolTitle", r.title)} ${theme.fg("dim", r.url)}`;
  }
  const more = raw.length - PREVIEW_RESULTS;
  if (more > 0) {
    text += `\n${theme.fg("dim", `... ${more} more (${keyHint("app.tools.expand", "to expand")})`)}`;
  }
  return new Text(text, 0, 0);
}

function formatFetch(r: FetchResult): string {
  const header =
    r.title && r.title.trim()
      ? `# ${r.title}\n\n<${r.finalUrl}>\n\n`
      : `<${r.finalUrl}>\n\n`;
  return `${header}${r.markdown}`;
}

async function runSearch(query: string): Promise<SearchResult[]> {
  const config = loadConfig();
  const settings = resolveSettings(config);
  switch (settings.provider) {
    case "kagi":
      return searchKagi({ query, token: settings.kagiToken ?? "" });
    case "searxng":
      return searchSearxng({ query, baseUrl: settings.searxngBaseUrl ?? "" });
  }
}

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "web_search",
    label: "Web Search",
    description:
      "Search the web. Returns a markdown list of titles, URLs, and snippets.",
    promptSnippet: "Search the web for the given query",
    parameters: Type.Object({
      query: Type.String({ description: "Search query text" }),
    }),
    renderCall(args, theme, context) {
      const query = args?.query;
      const display = query && query.length > 0
        ? query
        : theme.fg("toolOutput", "...");
      const text = new Text("", 0, 0);
      text.setText(theme.fg("toolTitle", theme.bold(`web_search "${display}"`)));
      return text;
    },
    renderResult(result, { expanded }, theme) {
      return renderSearchResult(result, expanded, theme);
    },
    async execute(_toolCallId, params) {
      try {
        const results = await runSearch(params.query);
        return {
          content: [{ type: "text", text: formatResults(results) }],
          details: { raw: results },
        };
      } catch (err) {
        if (err instanceof ConfigError || err instanceof SearchError) {
          return {
            content: [{ type: "text", text: `Search error: ${err.message}` }],
            isError: true,
            details: { raw: [] as SearchResult[] },
          };
        }
        throw err;
      }
    },
  });

  pi.registerTool({
    name: "web_fetch",
    label: "Web Fetch",
    description:
      "Fetch a URL and return the page content as readable markdown (Readability + Turndown over a headless Firefox session, so JS-rendered pages work). For raw HTML, non-text content (PDFs, images), or simple HTTP requests, use bash with curl instead.",
    promptSnippet: "Fetch a URL and convert the page to readable markdown",
    parameters: Type.Object({
      url: Type.String({ description: "Absolute URL to fetch" }),
    }),
    renderCall(args, theme, context) {
      const url = args?.url;
      const display = url && url.length > 0
        ? url
        : theme.fg("toolOutput", "...");
      const text = new Text("", 0, 0);
      text.setText(theme.fg("toolTitle", theme.bold(`web_fetch ${display}`)));
      return text;
    },
    async execute(_toolCallId, params) {
      try {
        const result = await fetchPage({ url: params.url });
        return {
          content: [{ type: "text", text: formatFetch(result) }],
          details: { raw: result },
        };
      } catch (err) {
        if (err instanceof SearchError) {
          return {
            content: [{ type: "text", text: `Fetch error: ${err.message}` }],
            isError: true,
            details: {
              raw: {
                url: params.url,
                finalUrl: params.url,
                markdown: "",
                truncated: false,
              } satisfies FetchResult,
            },
          };
        }
        throw err;
      }
    },
  });
}