tau-all-agent 0.1.0

package/extensions/websearch/index.ts

@@ -0,0 +1,312 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+import {
+  defineTool,
+  getMarkdownTheme,
+  keyHint,
+  type ExtensionAPI,
+  type ExtensionContext,
+} from "@mariozechner/pi-coding-agent";
+import { Container, Markdown, Spacer, Text } from "@mariozechner/pi-tui";
+import { Type } from "typebox";
+
+import { createBrowserSession, discoverProfiles } from "./browser/discovery.js";
+import { loadConfig } from "./config.js";
+import { renderSearchResultMarkdown } from "./normalize.js";
+import { isPiAnthropicModel, searchWithPiAnthropic } from "./providers/anthropic.pi.js";
+import { browserGemini } from "./providers/gemini.browser.js";
+import { isPiGeminiModel, searchWithPiGemini } from "./providers/gemini.pi.js";
+import { browserOpenAICodex } from "./providers/openai-codex.browser.js";
+import { isPiOpenAICodexModel, searchWithPiOpenAICodex } from "./providers/openai-codex.pi.js";
+import type { PiModelSelection } from "./providers/pi-model.shared.js";
+import { selectCurrentPiModel, selectFallbackPiModel } from "./providers/pi-model.shared.js";
+import type {
+  BrowserProfile,
+  WebsearchAuthSource,
+  WebsearchBackendId,
+  WebsearchBrowserFamily,
+  WebsearchConfig,
+  WebsearchResult,
+  WebsearchRouteId,
+} from "./types.js";
+
+const PI_ROUTE_HANDLERS: Record<`pi:${WebsearchBackendId}`, PiRouteHandler> = {
+  "pi:openai-codex": {
+    predicate: isPiOpenAICodexModel,
+    search: searchWithPiOpenAICodex,
+  },
+  "pi:anthropic": {
+    predicate: isPiAnthropicModel,
+    search: searchWithPiAnthropic,
+  },
+  "pi:gemini": {
+    predicate: isPiGeminiModel,
+    search: searchWithPiGemini,
+  },
+};
+
+interface SearchSummary {
+  result: string;
+  route: WebsearchRouteId;
+  backend: WebsearchBackendId;
+  authSource: WebsearchAuthSource;
+  browserName?: string;
+  profile?: string;
+  accountLabel?: string;
+  sources: number;
+}
+
+interface PiRouteHandler {
+  predicate: (model: Model<Api>) => boolean;
+  search: (
+    selection: PiModelSelection,
+    query: string,
+    signal?: AbortSignal,
+  ) => Promise<WebsearchResult>;
+}
+
+function toSearchSummary(route: WebsearchRouteId, result: WebsearchResult): SearchSummary {
+  return {
+    result: renderSearchResultMarkdown(result.answer, result.sources),
+    route,
+    backend: result.backend,
+    authSource: result.authSource,
+    browserName: result.browserName,
+    profile: result.profile,
+    accountLabel: result.accountLabel,
+    sources: result.sources.length,
+  };
+}
+
+function currentPiRoute(ctx: ExtensionContext): `pi:${WebsearchBackendId}` | null {
+  if (!ctx.model) return null;
+
+  for (const [route, handler] of Object.entries(PI_ROUTE_HANDLERS) as Array<
+    [`pi:${WebsearchBackendId}`, PiRouteHandler]
+  >) {
+    if (handler.predicate(ctx.model)) return route;
+  }
+
+  return null;
+}
+
+async function executePiRoute(
+  route: `pi:${WebsearchBackendId}`,
+  ctx: ExtensionContext,
+  query: string,
+  mode: "current" | "fallback",
+  signal?: AbortSignal,
+): Promise<SearchSummary | null> {
+  const handler = PI_ROUTE_HANDLERS[route];
+  const selection =
+    mode === "current"
+      ? await selectCurrentPiModel(ctx, handler.predicate)
+      : await selectFallbackPiModel(ctx, handler.predicate);
+
+  return selection ? toSearchSummary(route, await handler.search(selection, query, signal)) : null;
+}
+
+async function searchPiRoute(
+  route: `pi:${WebsearchBackendId}`,
+  ctx: ExtensionContext,
+  query: string,
+  signal?: AbortSignal,
+): Promise<SearchSummary | null> {
+  if (route === currentPiRoute(ctx)) {
+    const current = await executePiRoute(route, ctx, query, "current", signal);
+    if (current) return current;
+  }
+
+  return executePiRoute(route, ctx, query, "fallback", signal);
+}
+
+function isAbortError(error: unknown, signal?: AbortSignal): boolean {
+  return signal?.aborted === true || (error instanceof Error && error.name === "AbortError");
+}
+
+function throwIfAborted(signal?: AbortSignal): void {
+  if (!signal?.aborted) return;
+  throw signal.reason ?? new Error("The operation was aborted.");
+}
+
+function getBrowserProfiles(
+  cache: Partial<Record<WebsearchBrowserFamily, BrowserProfile[]>>,
+  family: WebsearchBrowserFamily,
+  config: WebsearchConfig,
+): BrowserProfile[] {
+  const cached = cache[family];
+  if (cached) return cached;
+
+  const profiles = discoverProfiles([family], config.profiles);
+  cache[family] = profiles;
+  return profiles;
+}
+
+async function searchBrowserRoute(
+  route: WebsearchRouteId,
+  config: WebsearchConfig,
+  query: string,
+  profileCache: Partial<Record<WebsearchBrowserFamily, BrowserProfile[]>>,
+  signal?: AbortSignal,
+): Promise<SearchSummary | null> {
+  const profileFamily: WebsearchBrowserFamily = route.startsWith("firefox:")
+    ? "firefox"
+    : "chromium";
+  const browserProvider = route.endsWith(":gemini") ? browserGemini : browserOpenAICodex;
+  const configuredProfileName = config.profiles[profileFamily];
+  const profiles = getBrowserProfiles(profileCache, profileFamily, config);
+  let lastError: string | null = null;
+
+  if (configuredProfileName && profiles.length === 0) {
+    throw new Error(`Configured ${profileFamily} profile was not found: ${configuredProfileName}`);
+  }
+
+  for (const profile of profiles) {
+    throwIfAborted(signal);
+
+    try {
+      const session = createBrowserSession(profile, browserProvider.domains);
+      if (!session) {
+        if (configuredProfileName) {
+          throw new Error(
+            `Configured ${profileFamily} profile has no usable session: ${configuredProfileName}`,
+          );
+        }
+        continue;
+      }
+      return toSearchSummary(route, await browserProvider.search(session, query, signal));
+    } catch (error) {
+      if (isAbortError(error, signal)) throw error;
+      lastError = error instanceof Error ? error.message : String(error);
+    }
+  }
+
+  if (lastError) throw new Error(lastError);
+  return null;
+}
+
+async function searchRoute(
+  route: WebsearchRouteId,
+  ctx: ExtensionContext,
+  config: WebsearchConfig,
+  query: string,
+  profileCache: Partial<Record<WebsearchBrowserFamily, BrowserProfile[]>>,
+  signal?: AbortSignal,
+): Promise<SearchSummary | null> {
+  return route.startsWith("pi:")
+    ? await searchPiRoute(route as `pi:${WebsearchBackendId}`, ctx, query, signal)
+    : await searchBrowserRoute(route, config, query, profileCache, signal);
+}
+
+async function runSearch(
+  ctx: ExtensionContext,
+  query: string,
+  signal?: AbortSignal,
+): Promise<SearchSummary> {
+  const config = loadConfig();
+  const profileCache: Partial<Record<WebsearchBrowserFamily, BrowserProfile[]>> = {};
+  let lastError: string | null = null;
+
+  for (const route of config.routes) {
+    throwIfAborted(signal);
+
+    try {
+      const result = await searchRoute(route, ctx, config, query, profileCache, signal);
+      if (result) return result;
+    } catch (error) {
+      if (isAbortError(error, signal)) throw error;
+      lastError = error instanceof Error ? error.message : String(error);
+    }
+  }
+
+  throw new Error(lastError ?? "No configured websearch route is available.");
+}
+
+const COLLAPSED_RESULT_LINES = 10;
+
+function trimTrailingEmptyLines(lines: string[]): string[] {
+  let end = lines.length;
+  while (end > 0 && lines[end - 1] === "") {
+    end--;
+  }
+  return lines.slice(0, end);
+}
+
+function formatWebsearchCall(query: string, theme: any): string {
+  let text = theme.fg("toolTitle", theme.bold("websearch"));
+  if (query) {
+    text += ` ${theme.fg("accent", query)}`;
+  }
+  return text;
+}
+
+function formatCollapsedWebsearchResult(resultText: string, theme: any): string {
+  const lines = trimTrailingEmptyLines(resultText.split("\n"));
+  if (lines.length === 0) {
+    return `\n${theme.fg("muted", "(no output)")}`;
+  }
+
+  const displayLines = lines.slice(0, COLLAPSED_RESULT_LINES);
+  const remaining = lines.length - displayLines.length;
+
+  let text = `\n${displayLines.map((line) => theme.fg("toolOutput", line)).join("\n")}`;
+  if (remaining > 0) {
+    text += `${theme.fg("muted", `\n... (${remaining} more lines,`)} ${keyHint("app.tools.expand", "to expand")})`;
+  }
+
+  return text;
+}
+
+export default function (pi: ExtensionAPI) {
+  pi.registerTool(
+    defineTool({
+      name: "websearch",
+      label: "Websearch",
+      description: "Web search via Gemini, OpenAI, or Claude, leveraging Pi or browser sessions.",
+      promptSnippet:
+        "Search the web for current or external information unavailable in local files",
+      promptGuidelines: [
+        "Use this for recent facts, live service behavior, or external documentation that is not already present in the repo.",
+        "Do not use websearch when repository files or supplied context already answer the question.",
+      ],
+      parameters: Type.Object({
+        query: Type.String({ description: "What to search for" }),
+      }),
+      renderShell: "self",
+      renderCall(args, theme) {
+        return new Text(formatWebsearchCall(args.query, theme), 0, 0);
+      },
+      renderResult(result, { expanded }, theme) {
+        const content = result.content.find((item) => item.type === "text");
+        const text = content?.type === "text" ? content.text : "";
+
+        if (!expanded) {
+          return new Text(formatCollapsedWebsearchResult(text, theme), 0, 0);
+        }
+
+        if (!text.trim()) {
+          return new Text(`\n${theme.fg("muted", "(no output)")}`, 0, 0);
+        }
+
+        const container = new Container();
+        container.addChild(new Spacer(1));
+        container.addChild(new Markdown(text.trim(), 0, 0, getMarkdownTheme()));
+        return container;
+      },
+      async execute(_toolCallId, params, signal, _onUpdate, ctx) {
+        const result = await runSearch(ctx, params.query, signal);
+        return {
+          content: [{ type: "text", text: result.result }],
+          details: {
+            route: result.route,
+            backend: result.backend,
+            authSource: result.authSource,
+            browserName: result.browserName,
+            profile: result.profile,
+            accountLabel: result.accountLabel,
+            sources: result.sources,
+          },
+        };
+      },
+    }),
+  );
+}

package/extensions/websearch/normalize.ts

@@ -0,0 +1,155 @@
+import type { BrowserCookie, WebsearchSource } from "./types.js";
+
+export const DEFAULT_BROWSER_USER_AGENT =
+  "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:148.0) Gecko/20100101 Firefox/148.0";
+
+export function buildCookieHeader(url: string, cookies: BrowserCookie[]): string {
+  return cookiesForUrl(url, cookies)
+    .map((cookie) => `${cookie.name}=${cookie.value}`)
+    .join("; ");
+}
+
+export function hasCookie(url: string, cookies: BrowserCookie[], name: string): boolean {
+  return cookiesForUrl(url, cookies).some((cookie) => cookie.name === name);
+}
+
+function cookiesForUrl(url: string, cookies: BrowserCookie[]): BrowserCookie[] {
+  const target = new URL(url);
+  const requestPath = target.pathname || "/";
+
+  return [...cookies]
+    .filter((cookie) => {
+      return (
+        cookie.value.length > 0 &&
+        matchesCookieDomain(target.hostname, cookie.domain) &&
+        matchesCookiePath(requestPath, cookie.path) &&
+        (!cookie.secure || target.protocol === "https:")
+      );
+    })
+    .sort((left, right) => right.path.length - left.path.length);
+}
+
+function matchesCookieDomain(hostname: string, domain: string): boolean {
+  const normalizedHost = hostname.toLowerCase();
+  const normalizedDomain = domain.toLowerCase();
+
+  if (normalizedDomain.startsWith(".")) {
+    const suffix = normalizedDomain.slice(1);
+    return normalizedHost === suffix || normalizedHost.endsWith(`.${suffix}`);
+  }
+
+  return normalizedHost === normalizedDomain;
+}
+
+function matchesCookiePath(requestPath: string, cookiePath: string): boolean {
+  const normalizedCookiePath = cookiePath || "/";
+  if (requestPath === normalizedCookiePath) return true;
+  if (!requestPath.startsWith(normalizedCookiePath)) return false;
+  if (normalizedCookiePath.endsWith("/")) return true;
+  return requestPath.charAt(normalizedCookiePath.length) === "/";
+}
+
+export function browserHeaders(options: {
+  cookieHeader?: string;
+  origin: string;
+  referer: string;
+  accept?: string;
+  contentType?: string;
+  authorization?: string;
+  extra?: Record<string, string | undefined>;
+}): Record<string, string> {
+  const headers: Record<string, string> = {
+    "User-Agent": DEFAULT_BROWSER_USER_AGENT,
+    Accept: options.accept ?? "application/json,text/plain,*/*",
+    "Accept-Language": "en-US,en;q=0.5",
+    Referer: options.referer,
+    Origin: options.origin,
+    "Sec-Fetch-Dest": "empty",
+    "Sec-Fetch-Mode": "cors",
+    "Sec-Fetch-Site": "same-origin",
+    Pragma: "no-cache",
+    "Cache-Control": "no-cache",
+  };
+
+  if (options.cookieHeader) headers.Cookie = options.cookieHeader;
+  if (options.contentType) headers["Content-Type"] = options.contentType;
+  if (options.authorization) headers.Authorization = options.authorization;
+
+  if (options.extra) {
+    for (const [key, value] of Object.entries(options.extra)) {
+      if (value) headers[key] = value;
+    }
+  }
+
+  return headers;
+}
+
+export function normalizeSourceTitle(url: string, title?: string | null): string {
+  const trimmed = title?.trim();
+  if (trimmed) return trimmed;
+
+  try {
+    const parsed = new URL(url);
+    return parsed.hostname.replace(/^www\./, "");
+  } catch {
+    return url;
+  }
+}
+
+export function extractMarkdownSources(text: string): WebsearchSource[] {
+  const sources: WebsearchSource[] = [];
+  const seen = new Set<string>();
+
+  const markdownLinkPattern = /\[([^\]]+)\]\((https?:\/\/[^)]+)\)/g;
+  for (const match of text.matchAll(markdownLinkPattern)) {
+    const url = match[2];
+    if (!url || seen.has(url)) continue;
+    seen.add(url);
+    sources.push({ title: normalizeSourceTitle(url, match[1]), url });
+  }
+
+  const urlPattern = /https?:\/\/[^\s)\]}>,]+/g;
+  for (const match of text.matchAll(urlPattern)) {
+    const url = match[0].replace(/[.,;:]+$/, "");
+    if (!url || seen.has(url)) continue;
+    seen.add(url);
+    sources.push({ title: normalizeSourceTitle(url), url });
+  }
+
+  return sources;
+}
+
+export function dedupeSources(sources: WebsearchSource[]): WebsearchSource[] {
+  const deduped = new Map<string, WebsearchSource>();
+
+  for (const source of sources) {
+    if (!source.url) continue;
+    if (!deduped.has(source.url)) {
+      deduped.set(source.url, {
+        title: normalizeSourceTitle(source.url, source.title),
+        url: source.url,
+        snippet: source.snippet,
+      });
+    }
+  }
+
+  return [...deduped.values()];
+}
+
+export function renderSearchResultMarkdown(answer: string, sources: WebsearchSource[]): string {
+  const trimmedAnswer = answer.trim();
+  if (sources.length === 0) return trimmedAnswer;
+
+  const embeddedSources = new Set(
+    extractMarkdownSources(trimmedAnswer).map((source) => source.url),
+  );
+  const extraSources = sources.filter((source) => !embeddedSources.has(source.url));
+  if (extraSources.length === 0) return trimmedAnswer;
+
+  const lines = [trimmedAnswer, "", "Sources:"];
+  for (const source of extraSources) {
+    lines.push(`- [${normalizeSourceTitle(source.url, source.title)}](${source.url})`);
+  }
+
+  return lines.join("\n").trim();
+}

package/extensions/websearch/providers/anthropic.pi.ts

@@ -0,0 +1,139 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+
+import { dedupeSources, extractMarkdownSources } from "../normalize.js";
+import type { WebsearchResult } from "../types.js";
+import type { PiModelSelection } from "./pi-model.shared.js";
+import { buildWebsearchPrompt, WEBSEARCH_SYSTEM_PROMPT } from "./search-prompt.shared.js";
+import { fetchText, withTimeout } from "./shared.js";
+
+export async function searchWithPiAnthropic(
+  selection: PiModelSelection,
+  query: string,
+  signal?: AbortSignal,
+): Promise<WebsearchResult> {
+  const payload = await fetchText(resolveAnthropicMessagesUrl(selection.model.baseUrl), {
+    method: "POST",
+    headers: buildAnthropicHeaders(selection),
+    body: JSON.stringify({
+      model: selection.model.id,
+      max_tokens: 1800,
+      system: buildAnthropicSystem(getAnthropicCredential(selection)),
+      tools: [{ type: "web_search_20250305", name: "web_search", max_uses: 5 }],
+      messages: [{ role: "user", content: buildWebsearchPrompt(query) }],
+    }),
+    signal: withTimeout(signal, 120_000),
+  });
+
+  let parsed: Record<string, unknown>;
+  try {
+    parsed = JSON.parse(payload) as Record<string, unknown>;
+  } catch {
+    throw new Error("Anthropic returned non-JSON response.");
+  }
+
+  const answer = (Array.isArray(parsed.content) ? parsed.content : [])
+    .filter(
+      (item) =>
+        item && typeof item === "object" && (item as Record<string, unknown>).type === "text",
+    )
+    .map((item) => (item as Record<string, unknown>).text)
+    .filter((text): text is string => typeof text === "string")
+    .join("\n\n")
+    .trim();
+
+  if (!answer) {
+    throw new Error("Anthropic returned no text content.");
+  }
+
+  return {
+    backend: "anthropic",
+    authSource: "pi",
+    answer,
+    sources: dedupeSources(extractMarkdownSources(answer)),
+  };
+}
+
+function buildAnthropicHeaders(selection: PiModelSelection): Record<string, string> {
+  const credential = getAnthropicCredential(selection);
+  const headers = { ...selection.headers };
+
+  if (isAnthropicOAuthToken(credential)) {
+    return {
+      ...headers,
+      ...(hasHeader(headers, "authorization") || !credential
+        ? {}
+        : { authorization: `Bearer ${credential}` }),
+      "anthropic-version": "2023-06-01",
+      "anthropic-beta":
+        "claude-code-20250219,oauth-2025-04-20,fine-grained-tool-streaming-2025-05-14,web-search-2025-03-05",
+      "anthropic-dangerous-direct-browser-access": "true",
+      "content-type": "application/json",
+      accept: "application/json",
+      "x-app": "cli",
+      "user-agent": "claude-cli/2.1.62",
+    };
+  }
+
+  return {
+    ...headers,
+    ...(hasHeader(headers, "authorization") || hasHeader(headers, "x-api-key") || !credential
+      ? {}
+      : { "x-api-key": credential }),
+    "anthropic-version": "2023-06-01",
+    "anthropic-beta": "web-search-2025-03-05",
+    "content-type": "application/json",
+    accept: "application/json",
+  };
+}
+
+function resolveAnthropicMessagesUrl(baseUrl?: string): string {
+  const normalized = String(baseUrl || "https://api.anthropic.com").replace(/\/+$/, "");
+  return normalized.endsWith("/v1/messages") ? normalized : `${normalized}/v1/messages`;
+}
+
+function buildAnthropicSystem(apiKey?: string): string | Array<{ type: "text"; text: string }> {
+  if (!isAnthropicOAuthToken(apiKey)) return WEBSEARCH_SYSTEM_PROMPT;
+
+  return [
+    { type: "text", text: "You are Claude Code, Anthropic's official CLI for Claude." },
+    { type: "text", text: WEBSEARCH_SYSTEM_PROMPT },
+  ];
+}
+
+function getAnthropicCredential(selection: PiModelSelection): string | undefined {
+  return (
+    selection.apiKey ??
+    getBearerToken(selection.headers) ??
+    getHeaderValue(selection.headers, "x-api-key")
+  );
+}
+
+function getBearerToken(headers?: Record<string, string>): string | undefined {
+  const authorization = getHeaderValue(headers, "authorization");
+  if (!authorization) return undefined;
+  const match = authorization.match(/^Bearer\s+(.+)$/i);
+  return match?.[1];
+}
+
+function hasHeader(headers: Record<string, string>, name: string): boolean {
+  return Object.keys(headers).some((key) => key.toLowerCase() === name.toLowerCase());
+}
+
+function getHeaderValue(
+  headers: Record<string, string> | undefined,
+  name: string,
+): string | undefined {
+  if (!headers) return undefined;
+  const key = Object.keys(headers).find(
+    (headerName) => headerName.toLowerCase() === name.toLowerCase(),
+  );
+  return key ? headers[key] : undefined;
+}
+
+function isAnthropicOAuthToken(apiKey?: string): boolean {
+  return Boolean(apiKey?.includes("sk-ant-oat"));
+}
+
+export function isPiAnthropicModel(model: Model<Api>): boolean {
+  return model.api === "anthropic-messages" && model.provider === "anthropic";
+}