@aaroncql/pim-agent 0.0.1

package/src/extensions/web-fetch/schema.ts

@@ -0,0 +1,23 @@
+import { StringEnum } from "@earendil-works/pi-ai";
+import { type Static, Type } from "typebox";
+
+export const WEB_FETCH_INLINE_BYTES = 32 * 1024;
+
+export const FETCH_FORMATS = ["markdown", "html"] as const;
+export type WebFetchFormat = (typeof FETCH_FORMATS)[number];
+export type WebFetchResolvedFormat = WebFetchFormat;
+
+export const webFetchSchema = Type.Object({
+  url: Type.String({
+    minLength: 1,
+    description: "Must be a public http(s) URL.",
+  }),
+  format: Type.Optional(
+    StringEnum(FETCH_FORMATS, {
+      description:
+        "`markdown`: used by default. `html`: use only when raw source is required.",
+    })
+  ),
+});
+
+export type WebFetchInput = Static<typeof webFetchSchema>;

package/src/extensions/web-search/ExaMcpClient.test.ts

@@ -0,0 +1,143 @@
+import { expect, test } from "bun:test";
+import { ExaMcpClient } from "./ExaMcpClient";
+
+type MockFetch = (
+  input: Parameters<typeof fetch>[0],
+  init?: Parameters<typeof fetch>[1]
+) => ReturnType<typeof fetch>;
+
+const captureMethod = async (
+  input: Parameters<typeof fetch>[0],
+  init: Parameters<typeof fetch>[1] | undefined
+): Promise<string> => {
+  if (input instanceof Request) {
+    const body = (await input.clone().json()) as { method?: string };
+    return body.method ?? "";
+  }
+  const body = JSON.parse(String(init?.body)) as { method?: string };
+  return body.method ?? "";
+};
+
+const handshakeOr = (toolCallResponse: () => Response): MockFetch => {
+  return async (input, init) => {
+    const method = await captureMethod(input, init);
+
+    if (method === "initialize") {
+      return Response.json(
+        { jsonrpc: "2.0", id: 1, result: {} },
+        { headers: { "mcp-session-id": "session" } }
+      );
+    }
+
+    if (method === "notifications/initialized") {
+      return new Response(null, { status: 202 });
+    }
+
+    return toolCallResponse();
+  };
+};
+
+test("parses Exa JSON results", async () => {
+  const client = new ExaMcpClient({
+    fetch: handshakeOr(() =>
+      Response.json({
+        jsonrpc: "2.0",
+        id: 2,
+        result: {
+          content: [
+            {
+              type: "text",
+              text: JSON.stringify({
+                results: [
+                  {
+                    title: "Pim docs",
+                    url: "https://example.test/pim",
+                    snippet: "A concise result.",
+                  },
+                ],
+              }),
+            },
+          ],
+        },
+      })
+    ),
+  });
+
+  await expect(
+    client.search({ query: "pim agent", numResults: 3 })
+  ).resolves.toEqual([
+    {
+      title: "Pim docs",
+      url: "https://example.test/pim",
+      snippet: "A concise result.",
+    },
+  ]);
+});
+
+test("parses Exa plain-text result blocks", async () => {
+  const client = new ExaMcpClient({
+    fetch: handshakeOr(() =>
+      Response.json({
+        jsonrpc: "2.0",
+        id: 2,
+        result: {
+          content: [
+            {
+              type: "text",
+              text: [
+                "Title: First text result",
+                "URL: https://example.test/first",
+                "Published: N/A",
+                "Author: N/A",
+                "Highlights:",
+                "First highlighted sentence.",
+                "[...]",
+                "Second highlighted sentence.",
+                "",
+                "---",
+                "",
+                "Title: Second text result",
+                "URL: https://example.test/second",
+                "Highlights:",
+                "Another result.",
+              ].join("\n"),
+            },
+          ],
+        },
+      })
+    ),
+  });
+
+  await expect(
+    client.search({ query: "text result", numResults: 2 })
+  ).resolves.toEqual([
+    {
+      title: "First text result",
+      url: "https://example.test/first",
+      snippet: "First highlighted sentence. Second highlighted sentence.",
+    },
+    {
+      title: "Second text result",
+      url: "https://example.test/second",
+      snippet: "Another result.",
+    },
+  ]);
+});
+
+test("throws clean errors for malformed tool envelopes", async () => {
+  const client = new ExaMcpClient({
+    fetch: handshakeOr(() =>
+      Response.json({
+        jsonrpc: "2.0",
+        id: 2,
+        result: {
+          content: [{ type: "image", url: "https://example.test/image.png" }],
+        },
+      })
+    ),
+  });
+
+  await expect(client.search({ query: "pim", numResults: 1 })).rejects.toThrow(
+    "Exa returned malformed tool content."
+  );
+});

package/src/extensions/web-search/ExaMcpClient.ts

@@ -0,0 +1,258 @@
+import { McpClient, type McpFetch } from "../../shared/McpClient";
+
+type ExaMcpClientOptions = {
+  readonly endpoint?: string;
+  readonly apiKey?: string;
+  readonly fetch?: McpFetch;
+};
+
+type ExaSearchInput = {
+  readonly query: string;
+  readonly numResults: number;
+  readonly signal?: AbortSignal;
+};
+
+export type ExaSearchResult = {
+  readonly title: string;
+  readonly url: string;
+  readonly snippet: string;
+};
+
+class ExaSearchError extends Error {
+  public constructor(message: string) {
+    super(message);
+    this.name = "ExaSearchError";
+  }
+}
+
+export class ExaMcpClient {
+  private static readonly defaultEndpoint = "https://mcp.exa.ai/mcp";
+  private static readonly toolName = "web_search_exa";
+
+  private readonly client: McpClient;
+
+  public constructor(options: ExaMcpClientOptions = {}) {
+    this.client = new McpClient({
+      endpoint: options.endpoint ?? ExaMcpClient.defaultEndpoint,
+      ...(options.apiKey === undefined || options.apiKey.length === 0
+        ? {}
+        : { headers: { "x-api-key": options.apiKey } }),
+      ...(options.fetch === undefined ? {} : { fetch: options.fetch }),
+    });
+  }
+
+  public async search(
+    input: ExaSearchInput
+  ): Promise<readonly ExaSearchResult[]> {
+    const result = await this.client.callTool({
+      name: ExaMcpClient.toolName,
+      arguments: {
+        query: input.query,
+        numResults: input.numResults,
+      },
+      ...(input.signal === undefined ? {} : { signal: input.signal }),
+    });
+
+    return extractResults(result);
+  }
+}
+
+function extractResults(result: unknown): readonly ExaSearchResult[] {
+  const record = asRecord(result);
+  const content = record?.["content"];
+
+  if (!Array.isArray(content) || content.length === 0) {
+    throw new ExaSearchError("Exa returned malformed tool content.");
+  }
+
+  const textBlocks = content.map(readTextBlock);
+  const plainTextResults = extractPlainTextResults(textBlocks);
+
+  if (plainTextResults !== undefined) {
+    return plainTextResults;
+  }
+
+  const resultObjects = findFirstObjectArray(
+    textBlocks
+      .map((block) => tryParseJson(block))
+      .filter((value) => value !== undefined)
+  );
+
+  if (resultObjects === undefined) {
+    throw new ExaSearchError("Exa returned malformed search results.");
+  }
+
+  return resultObjects.map(projectResult);
+}
+
+function readTextBlock(block: unknown): string {
+  const record = asRecord(block);
+
+  if (record?.["type"] !== "text" || typeof record["text"] !== "string") {
+    throw new ExaSearchError("Exa returned malformed tool content.");
+  }
+
+  return record["text"];
+}
+
+function extractPlainTextResults(
+  textBlocks: readonly string[]
+): readonly ExaSearchResult[] | undefined {
+  for (const textBlock of textBlocks) {
+    const blocks = textBlock
+      .split(/\n---\n/u)
+      .map((block) => block.trim())
+      .filter((block) => block.length > 0);
+    const results = blocks
+      .map((block) => parsePlainTextResult(block))
+      .filter((result) => result !== undefined);
+
+    if (results.length > 0) {
+      return results;
+    }
+  }
+
+  return undefined;
+}
+
+function parsePlainTextResult(block: string): ExaSearchResult | undefined {
+  const lines = block
+    .split(/\r?\n/u)
+    .map((line) => line.trim())
+    .filter((line) => line.length > 0);
+  const title = readLabeledLine(lines, "Title");
+  const url = readLabeledLine(lines, "URL");
+
+  if (title === undefined || url === undefined) {
+    return undefined;
+  }
+
+  return {
+    title,
+    url,
+    snippet: readPlainTextSnippet(lines),
+  };
+}
+
+function readLabeledLine(
+  lines: readonly string[],
+  label: string
+): string | undefined {
+  const prefix = `${label}:`;
+  const line = lines.find((candidate) =>
+    candidate.toLowerCase().startsWith(prefix.toLowerCase())
+  );
+
+  return line?.slice(prefix.length).trim();
+}
+
+function readPlainTextSnippet(lines: readonly string[]): string {
+  const highlightsIndex = lines.findIndex((line) =>
+    line.toLowerCase().startsWith("highlights:")
+  );
+  const snippetLines =
+    highlightsIndex === -1 ? lines : lines.slice(highlightsIndex + 1);
+  const skipPrefixes = ["title:", "url:", "published:", "author:"];
+  const snippet = snippetLines
+    .filter((line) => {
+      if (line.startsWith("[...]")) {
+        return false;
+      }
+      const lower = line.toLowerCase();
+      return !skipPrefixes.some((prefix) => lower.startsWith(prefix));
+    })
+    .join(" ")
+    .replace(/\s+/gu, " ")
+    .trim();
+
+  return snippet.length > 500 ? `${snippet.slice(0, 500)}...` : snippet;
+}
+
+function findFirstObjectArray(
+  values: readonly unknown[]
+): readonly Readonly<Record<string, unknown>>[] | undefined {
+  for (const value of values) {
+    if (
+      Array.isArray(value) &&
+      value.every((item) => asRecord(item) !== undefined)
+    ) {
+      return value as readonly Readonly<Record<string, unknown>>[];
+    }
+
+    const record = asRecord(value);
+
+    if (record === undefined) {
+      continue;
+    }
+
+    for (const nestedValue of Object.values(record)) {
+      if (
+        Array.isArray(nestedValue) &&
+        nestedValue.every((item) => asRecord(item) !== undefined)
+      ) {
+        return nestedValue as readonly Readonly<Record<string, unknown>>[];
+      }
+    }
+  }
+
+  return undefined;
+}
+
+function projectResult(
+  result: Readonly<Record<string, unknown>>
+): ExaSearchResult {
+  return {
+    title: readResultString(result, "title"),
+    url: readResultString(result, "url"),
+    snippet: readSnippet(result),
+  };
+}
+
+function readSnippet(result: Readonly<Record<string, unknown>>): string {
+  return (
+    readOptionalResultString(result, "snippet") ??
+    readOptionalResultString(result, "text") ??
+    readOptionalResultString(result, "summary") ??
+    ""
+  );
+}
+
+function readResultString(
+  result: Readonly<Record<string, unknown>>,
+  name: string
+): string {
+  const value = readOptionalResultString(result, name);
+
+  if (value === undefined) {
+    throw new ExaSearchError(`Exa returned a result without ${name}.`);
+  }
+
+  return value;
+}
+
+function readOptionalResultString(
+  result: Readonly<Record<string, unknown>>,
+  name: string
+): string | undefined {
+  const value = result[name];
+
+  return typeof value === "string" ? value : undefined;
+}
+
+function tryParseJson(text: string): unknown | undefined {
+  try {
+    return JSON.parse(text);
+  } catch {
+    return undefined;
+  }
+}
+
+function asRecord(
+  value: unknown
+): Readonly<Record<string, unknown>> | undefined {
+  if (typeof value !== "object" || value === null || Array.isArray(value)) {
+    return undefined;
+  }
+
+  return value as Readonly<Record<string, unknown>>;
+}

package/src/extensions/web-search/index.ts

@@ -0,0 +1,118 @@
+import type { ExtensionAPI, Theme } from "@earendil-works/pi-coding-agent";
+import {
+  Renderer,
+  type StatefulToolCallTitleContext,
+  type StatefulToolCallTitleState,
+} from "../../shared/Renderer";
+import { PimSettings } from "../../shared/PimSettings";
+import { Tools } from "../../shared/Tools";
+import { ExaMcpClient } from "./ExaMcpClient";
+import { formatTitle } from "./render";
+import { type WebSearchInput, webSearchSchema } from "./schema";
+import { clampNumResults, formatResults } from "./search";
+
+const PREVIEW_LINES = 6;
+
+type WebSearchCallState = StatefulToolCallTitleState & {
+  resultCount?: number;
+};
+
+type WebSearchRenderContext = StatefulToolCallTitleContext & {
+  readonly args?: WebSearchInput;
+};
+
+async function createClient(): Promise<ExaMcpClient> {
+  const apiKey = await PimSettings.getExaApiKey();
+  return new ExaMcpClient(apiKey ? { apiKey } : {});
+}
+
+function renderTitle(
+  input: Partial<WebSearchInput>,
+  theme: Theme,
+  context: WebSearchRenderContext
+) {
+  const state = context.state as WebSearchCallState;
+  const count = state.resultCount ?? clampNumResults(input.numResults);
+  return Renderer.renderStatefulToolCallTitle({
+    label: "Web Search",
+    title: formatTitle(input.query, count),
+    theme,
+    context,
+  });
+}
+
+export default function (pi: ExtensionAPI): void {
+  let clientPromise: Promise<ExaMcpClient> | undefined;
+  const getClient = () => (clientPromise ??= createClient());
+
+  Tools.register(pi, {
+    name: "web_search",
+    label: "web_search",
+    description:
+      "Search the web. Returns ranked results with title, URL, and a short snippet.",
+    parameters: webSearchSchema,
+    renderShell: "self",
+    executionMode: "parallel",
+    async execute(_id, params, signal) {
+      const { query, numResults } = params as WebSearchInput;
+
+      if (signal?.aborted) {
+        throw new Error("Web search aborted before execution.");
+      }
+
+      const trimmed = query.trim();
+      if (trimmed.length === 0) {
+        throw new Error(
+          "Web search query is empty. Provide a non-empty query."
+        );
+      }
+
+      const clamped = clampNumResults(numResults);
+      const client = await getClient();
+      const results = await client.search({
+        query: trimmed,
+        numResults: clamped,
+        ...(signal === undefined ? {} : { signal }),
+      });
+
+      if (results.length === 0) {
+        throw new Error(
+          `No web results for "${trimmed}". Try broader keywords or different phrasing.`
+        );
+      }
+
+      return {
+        content: [{ type: "text", text: formatResults(results) }],
+        details: {
+          query: trimmed,
+          numResults: clamped,
+          count: results.length,
+        },
+      };
+    },
+    renderCall(args, theme, context) {
+      return renderTitle(
+        (args ?? {}) as Partial<WebSearchInput>,
+        theme,
+        context
+      );
+    },
+    renderResult(result, options, theme, context) {
+      const state = context.state as WebSearchCallState;
+      const details = result.details as { readonly count?: number } | undefined;
+
+      if (details?.count !== undefined) {
+        state.resultCount = details.count;
+        renderTitle(context.args ?? {}, theme, context);
+      }
+
+      return Renderer.renderBorderedResult({
+        result,
+        options,
+        theme,
+        context,
+        previewLines: PREVIEW_LINES,
+      });
+    },
+  });
+}

package/src/extensions/web-search/render.test.ts

@@ -0,0 +1,21 @@
+import { describe, expect, test } from "bun:test";
+import { DEFAULT_NUM_RESULTS } from "./schema";
+import { formatTitle } from "./render";
+
+describe("formatTitle", () => {
+  test("always includes the default count in parentheses", () => {
+    expect(formatTitle("bun release notes", undefined)).toBe(
+      `bun release notes (${DEFAULT_NUM_RESULTS})`
+    );
+  });
+
+  test("includes explicit counts in parentheses", () => {
+    expect(formatTitle("pi agent", 3)).toBe("pi agent (3)");
+  });
+
+  test("uses a placeholder while keeping the count visible", () => {
+    expect(formatTitle(undefined, undefined)).toBe(
+      `... (${DEFAULT_NUM_RESULTS})`
+    );
+  });
+});

package/src/extensions/web-search/render.ts

@@ -0,0 +1,9 @@
+import { DEFAULT_NUM_RESULTS } from "./schema";
+
+export function formatTitle(
+  query: string | undefined,
+  n: number | undefined = DEFAULT_NUM_RESULTS
+): string {
+  const q = query ?? "...";
+  return `${q} (${n})`;
+}

package/src/extensions/web-search/schema.ts

@@ -0,0 +1,21 @@
+import { type Static, Type } from "typebox";
+
+export const MIN_NUM_RESULTS = 1;
+export const MAX_NUM_RESULTS = 10;
+export const DEFAULT_NUM_RESULTS = 5;
+
+export const webSearchSchema = Type.Object({
+  query: Type.String({
+    minLength: 1,
+    description: "Search query. Prefer natural language over keywords.",
+  }),
+  numResults: Type.Optional(
+    Type.Integer({
+      minimum: MIN_NUM_RESULTS,
+      maximum: MAX_NUM_RESULTS,
+      description: `Number of results, ${MIN_NUM_RESULTS}-${MAX_NUM_RESULTS}. Defaults to ${DEFAULT_NUM_RESULTS}.`,
+    })
+  ),
+});
+
+export type WebSearchInput = Static<typeof webSearchSchema>;

package/src/extensions/web-search/search.test.ts

@@ -0,0 +1,53 @@
+import { describe, expect, test } from "bun:test";
+import { clampNumResults, formatResults } from "./search";
+
+describe("clampNumResults", () => {
+  test("defaults when undefined", () => {
+    expect(clampNumResults(undefined)).toBe(5);
+  });
+
+  test("clamps above maximum", () => {
+    expect(clampNumResults(25)).toBe(10);
+  });
+
+  test("clamps below minimum", () => {
+    expect(clampNumResults(0)).toBe(1);
+  });
+
+  test("passes through valid values", () => {
+    expect(clampNumResults(3)).toBe(3);
+  });
+});
+
+describe("formatResults", () => {
+  test("renders results as deterministic plain text", () => {
+    expect(
+      formatResults([
+        {
+          title: "First",
+          url: "https://example.test/first",
+          snippet: "First snippet.",
+        },
+        {
+          title: "Second",
+          url: "https://example.test/second",
+          snippet: "Second snippet.",
+        },
+      ])
+    ).toBe(
+      [
+        "title: First",
+        "url: https://example.test/first",
+        "snippet: First snippet.",
+        "",
+        "title: Second",
+        "url: https://example.test/second",
+        "snippet: Second snippet.",
+      ].join("\n")
+    );
+  });
+
+  test("returns empty string for empty input", () => {
+    expect(formatResults([])).toBe("");
+  });
+});

package/src/extensions/web-search/search.ts

@@ -0,0 +1,23 @@
+import type { ExaSearchResult } from "./ExaMcpClient";
+import {
+  DEFAULT_NUM_RESULTS,
+  MAX_NUM_RESULTS,
+  MIN_NUM_RESULTS,
+} from "./schema";
+
+export function clampNumResults(value: number | undefined): number {
+  const requested = value ?? DEFAULT_NUM_RESULTS;
+  return Math.min(MAX_NUM_RESULTS, Math.max(MIN_NUM_RESULTS, requested));
+}
+
+export function formatResults(results: readonly ExaSearchResult[]): string {
+  return results
+    .map((result) =>
+      [
+        `title: ${result.title}`,
+        `url: ${result.url}`,
+        `snippet: ${result.snippet}`,
+      ].join("\n")
+    )
+    .join("\n\n");
+}

package/src/extensions/working-indicator/index.test.ts

@@ -0,0 +1,21 @@
+import { describe, expect, test } from "bun:test";
+import { formatElapsed } from "./index";
+
+describe("working indicator formatting", () => {
+  test("formats sub-minute elapsed time", () => {
+    expect(formatElapsed(0)).toBe("0s");
+    expect(formatElapsed(999)).toBe("0s");
+    expect(formatElapsed(32_000)).toBe("32s");
+  });
+
+  test("formats minute elapsed time", () => {
+    expect(formatElapsed(60_000)).toBe("1m 0s");
+    expect(formatElapsed(92_000)).toBe("1m 32s");
+    expect(formatElapsed(3_599_000)).toBe("59m 59s");
+  });
+
+  test("formats hour elapsed time", () => {
+    expect(formatElapsed(3_600_000)).toBe("1h 0m 0s");
+    expect(formatElapsed(3_692_000)).toBe("1h 1m 32s");
+  });
+});