tau-coding-agent 0.1.6 → 0.2.0

package/extensions/usage/openrouter.ts

@@ -95,8 +95,16 @@ async function fetchUsage(ctx: ExtensionContext, signal?: AbortSignal): Promise<
     throw keyResult.reason;
   }
 
-  const totalCredits = readNumber(creditsPayload.data?.total_credits) ?? 0;
-  const totalUsage = readNumber(creditsPayload.data?.total_usage) ?? 0;
+  const totalCredits = readNumber(creditsPayload?.data?.total_credits);
+  const totalUsage = readNumber(creditsPayload?.data?.total_usage);
+  if (
+    totalCredits === undefined ||
+    totalCredits < 0 ||
+    totalUsage === undefined ||
+    totalUsage < 0
+  ) {
+    throw new Error("OpenRouter returned invalid credit totals.");
+  }
   const balance = Math.max(0, totalCredits - totalUsage);
 
   const keyLimit = readNumber(keyPayload?.data?.limit);

package/extensions/websearch/README.md

@@ -1,36 +1,14 @@
 # websearch
 
-Web search via Gemini, OpenAI, or Claude, leveraging Pi or browser sessions.
+Search the web with Gemini, OpenAI, or Claude using Pi credentials or existing browser sessions. The `websearch` tool tries the configured routes and falls back when a search fails.
 
-Current routes:
+## Configuration
 
-- `pi:openai-codex`
-- `pi:anthropic`
-- `pi:gemini`
-- `firefox:gemini`
-- `firefox:openai-codex`
-- `chromium:gemini`
-- `chromium:openai-codex`
-
-## Config
-
-Optional config file:
-
-- `~/.pi/agent/websearch.json`
-
-Example:
+Optionally create `websearch.json` in your Pi agent directory (normally `~/.pi/agent`):
 
 ```json
 {
-  "routes": [
-    "pi:openai-codex",
-    "pi:anthropic",
-    "pi:gemini",
-    "firefox:gemini",
-    "firefox:openai-codex",
-    "chromium:gemini",
-    "chromium:openai-codex"
-  ],
+  "routes": ["pi:openai-codex", "pi:anthropic", "pi:gemini", "firefox:gemini", "chromium:gemini"],
   "profiles": {
     "firefox": "default-release",
     "chromium": "Default"
@@ -38,31 +16,18 @@ Example:
 }
 ```
 
-Defaults:
+The routes shown are the default order. Set `routes` to choose which routes to use and their order.
 
-- default routes match the example above
-- if `routes` is set, only those routes are tried, in that order
-- Pi routes try the current Pi model first within the matching provider family, then fall back to another available model from that family
+Pi routes use Pi's authentication. Browser routes use existing Google sessions for Gemini. Browser profiles are discovered automatically unless pinned with `profiles`.
 
-`profiles` only applies to browser-backed routes. It pins which local Firefox or Chromium profile to read cookies from.
+Browser access supports Firefox and Chromium-family browsers on macOS and Linux. macOS may prompt for Keychain access to decrypt cookies.
 
 ## Usage
 
-In Pi, the extension exposes:
-
-- tool: `websearch`
-
-Examples:
-
-- `websearch({ query: "latest Todoist release notes" })`
-- `websearch({ query: "how does Pi compaction work" })`
+```js
+websearch({ query: "how does Pi compaction work" });
+```
 
-## Notes
+## Output
 
-- Pi-backed API routes are preferred, with browser fallbacks after that.
-- Browser routes use existing local browser sessions.
-- Browser profiles are auto-discovered unless pinned in config.
-- Supported Chromium-family browsers are Chromium, Chrome, Brave, and Edge.
-- Browser discovery currently supports macOS and Linux, not Windows.
-- On macOS, Chromium-based browsers may prompt for Keychain access so cookies can be decrypted.
-- If one route fails, the extension falls through to the next one.
+Output is limited to 2,000 lines or 50 KB. Longer results include a path to the full text, including sources, in a temporary file.

package/extensions/websearch/config.ts

@@ -11,9 +11,7 @@ const DEFAULT_ROUTES: WebsearchRouteId[] = [
   "pi:anthropic",
   "pi:gemini",
   "firefox:gemini",
-  "firefox:openai-codex",
   "chromium:gemini",
-  "chromium:openai-codex",
 ];
 
 function sanitizeRoutes(value: unknown): WebsearchRouteId[] {
@@ -26,6 +24,11 @@ function sanitizeRoutes(value: unknown): WebsearchRouteId[] {
   const result: WebsearchRouteId[] = [];
 
   for (const item of value) {
+    if (item === "firefox:openai-codex" || item === "chromium:openai-codex") {
+      throw new Error(
+        `Invalid websearch route: ${item} (removed). Use pi:openai-codex and sign in with Pi /login.`,
+      );
+    }
     if (typeof item !== "string" || !valid.has(item as WebsearchRouteId)) {
       throw new Error(`Invalid websearch route: ${String(item)}`);
     }

package/extensions/websearch/index.ts

@@ -1,6 +1,9 @@
 import type { Api, Model } from "@earendil-works/pi-ai";
 import {
+  DEFAULT_MAX_BYTES,
+  DEFAULT_MAX_LINES,
   defineTool,
+  formatSize,
   getMarkdownTheme,
   keyHint,
   type ExtensionAPI,
@@ -12,15 +15,15 @@ import { Type } from "typebox";
 import { createBrowserSession, discoverProfiles } from "./browser/discovery.js";
 import { loadConfig } from "./config.js";
 import { renderSearchResultMarkdown } from "./normalize.js";
+import { limitOutput } from "./output.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 { getPiModelCandidates, selectNextPiModel } from "./providers/pi-model.shared.js";
+import { isModelUnavailableError } from "./providers/shared.js";
 import type {
-  BrowserProfile,
   WebsearchAuthSource,
   WebsearchBackendId,
   WebsearchBrowserFamily,
@@ -29,10 +32,13 @@ import type {
   WebsearchRouteId,
 } from "./types.js";
 
-const PI_ROUTE_HANDLERS: Record<`pi:${WebsearchBackendId}`, PiRouteHandler> = {
+type PiRouteId = `pi:${WebsearchBackendId}`;
+
+const PI_ROUTE_HANDLERS: Record<PiRouteId, PiRouteHandler> = {
   "pi:openai-codex": {
     predicate: isPiOpenAICodexModel,
     search: searchWithPiOpenAICodex,
+    fallbackModels: ["gpt-5.6-luna", "gpt-5.5"],
   },
   "pi:anthropic": {
     predicate: isPiAnthropicModel,
@@ -51,12 +57,17 @@ interface SearchSummary {
   authSource: WebsearchAuthSource;
   browserName?: string;
   profile?: string;
-  accountLabel?: string;
   sources: number;
 }
 
+interface PiRoutePlan {
+  remaining: Model<Api>[];
+  attemptsLeft: number;
+}
+
 interface PiRouteHandler {
   predicate: (model: Model<Api>) => boolean;
+  fallbackModels?: readonly string[];
   search: (
     selection: PiModelSelection,
     query: string,
@@ -64,6 +75,64 @@ interface PiRouteHandler {
   ) => Promise<WebsearchResult>;
 }
 
+async function runSearch(
+  ctx: ExtensionContext,
+  query: string,
+  signal?: AbortSignal,
+): Promise<SearchSummary> {
+  const config = loadConfig();
+  const plans = new Map<PiRouteId, PiRoutePlan>();
+  let pending = config.routes;
+  let lastError: string | null = null;
+
+  while (pending.length > 0) {
+    const retries: PiRouteId[] = [];
+    for (const route of pending) {
+      throwIfAborted(signal);
+      let selection: PiModelSelection | null = null;
+      let plan: PiRoutePlan | undefined;
+      try {
+        if (isPiRoute(route)) {
+          const handler = PI_ROUTE_HANDLERS[route];
+          plan = plans.get(route);
+          if (!plan) {
+            const remaining = getPiModelCandidates(ctx, handler.predicate, handler.fallbackModels);
+            plan = { remaining, attemptsLeft: handler.fallbackModels ? remaining.length : 2 };
+            plans.set(route, plan);
+          }
+          selection = await selectNextPiModel(ctx, plan.remaining, signal);
+          if (!selection) continue;
+          throwIfAborted(signal);
+          plan.attemptsLeft--;
+          return toSearchSummary(route, await handler.search(selection, query, signal));
+        }
+        const result = await searchBrowserRoute(route, config, query, signal);
+        if (result) return result;
+      } catch (error) {
+        if (isAbortError(error, signal)) throw error;
+        lastError = error instanceof Error ? error.message : String(error);
+        if (
+          isPiRoute(route) &&
+          selection &&
+          plan &&
+          plan.attemptsLeft > 0 &&
+          plan.remaining.length > 0 &&
+          isModelUnavailableError(error, selection.model)
+        ) {
+          retries.push(route);
+        }
+      }
+    }
+    pending = retries;
+  }
+
+  throw new Error(lastError ?? "No configured websearch route is available.");
+}
+
+function isPiRoute(route: WebsearchRouteId): route is PiRouteId {
+  return route.startsWith("pi:");
+}
+
 function toSearchSummary(route: WebsearchRouteId, result: WebsearchResult): SearchSummary {
   return {
     result: renderSearchResultMarkdown(result.answer, result.sources),
@@ -72,53 +141,10 @@ function toSearchSummary(route: WebsearchRouteId, result: WebsearchResult): Sear
     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");
 }
@@ -128,32 +154,17 @@ function throwIfAborted(signal?: AbortSignal): void {
   throw signal.reason ?? new Error("The operation was aborted.");
 }
 
-async function getBrowserProfiles(
-  cache: Partial<Record<WebsearchBrowserFamily, BrowserProfile[]>>,
-  family: WebsearchBrowserFamily,
-  config: WebsearchConfig,
-): Promise<BrowserProfile[]> {
-  const cached = cache[family];
-  if (cached) return cached;
-
-  const profiles = await 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 = await getBrowserProfiles(profileCache, profileFamily, config);
+  const profiles = await discoverProfiles([profileFamily], config.profiles);
   let lastError: string | null = null;
 
   if (configuredProfileName && profiles.length === 0) {
@@ -164,7 +175,7 @@ async function searchBrowserRoute(
     throwIfAborted(signal);
 
     try {
-      const session = await createBrowserSession(profile, browserProvider.domains);
+      const session = await createBrowserSession(profile, browserGemini.domains);
       if (!session) {
         if (configuredProfileName) {
           throw new Error(
@@ -173,7 +184,7 @@ async function searchBrowserRoute(
         }
         continue;
       }
-      return toSearchSummary(route, await browserProvider.search(session, query, signal));
+      return toSearchSummary(route, await browserGemini.search(session, query, signal));
     } catch (error) {
       if (isAbortError(error, signal)) throw error;
       lastError = error instanceof Error ? error.message : String(error);
@@ -184,43 +195,6 @@ async function searchBrowserRoute(
   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[] {
@@ -261,11 +235,11 @@ export default function (pi: ExtensionAPI) {
     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",
+      description: `Search the web for sourced answers. Output is limited to ${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)} (whichever is hit first), including the truncation notice. Full truncated output is saved to a temporary file.`,
+      promptSnippet: "Discover web sources for current or external information",
       promptGuidelines: [
         "Use websearch for recent facts, live service behavior, or external documentation that is not already present in the repo.",
+        "When asked to read a supplied URL, fetch that page directly. Use websearch to discover sources or independently verify information.",
         "Do not use websearch when repository files or supplied context already answer the question.",
       ],
       parameters: Type.Object({
@@ -293,16 +267,27 @@ export default function (pi: ExtensionAPI) {
         return container;
       },
       async execute(_toolCallId, params, signal, _onUpdate, ctx) {
-        const result = await runSearch(ctx, params.query, signal);
+        let result: SearchSummary;
+        try {
+          result = await runSearch(ctx, params.query, signal);
+        } catch (error) {
+          if (isAbortError(error, signal)) throw error;
+          const output = await limitOutput(
+            error instanceof Error ? error.message : String(error),
+            signal,
+          );
+          throw new Error(output.text);
+        }
+        const output = await limitOutput(result.result, signal);
         return {
-          content: [{ type: "text", text: result.result }],
+          content: [{ type: "text", text: output.text }],
           details: {
+            ...output.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/output.ts

@@ -0,0 +1,51 @@
+import { mkdtemp, rm, writeFile } from "node:fs/promises";
+import os from "node:os";
+import path from "node:path";
+import {
+  DEFAULT_MAX_BYTES,
+  DEFAULT_MAX_LINES,
+  formatSize,
+  truncateHead,
+  withFileMutationQueue,
+  type TruncationResult,
+} from "@earendil-works/pi-coding-agent";
+
+interface OutputDetails {
+  truncation?: Omit<TruncationResult, "content">;
+  fullOutputPath?: string;
+}
+
+/** Limit the complete model-facing text; retain full output only in a private, published file. */
+export async function limitOutput(
+  output: string,
+  signal?: AbortSignal,
+): Promise<{ text: string; details: OutputDetails }> {
+  signal?.throwIfAborted();
+  if (!truncateHead(output).truncated) return { text: output, details: {} };
+
+  const directory = await mkdtemp(path.join(os.tmpdir(), "pi-websearch-"));
+  try {
+    const fullOutputPath = path.join(directory, "output.txt");
+    const notice = `\n\n[Output truncated (${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)} limit). Full output saved to: ${fullOutputPath}]`;
+    const maxBytes = DEFAULT_MAX_BYTES - Buffer.byteLength(notice, "utf8");
+    const maxLines = DEFAULT_MAX_LINES - (notice.split("\n").length - 1);
+    if (maxBytes <= 0 || maxLines <= 0) {
+      throw new Error("Websearch output path is too long for the truncation notice.");
+    }
+    const { content, ...truncation } = truncateHead(output, { maxBytes, maxLines });
+    await withFileMutationQueue(fullOutputPath, async () => {
+      await writeFile(fullOutputPath, output, {
+        encoding: "utf8",
+        mode: 0o600,
+        flag: "wx",
+        signal,
+      });
+    });
+    signal?.throwIfAborted();
+    // Published files remain readable after completion/shutdown, like native tool output files.
+    return { text: content + notice, details: { truncation, fullOutputPath } };
+  } catch (error) {
+    await rm(directory, { recursive: true, force: true });
+    throw error;
+  }
+}

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

@@ -4,7 +4,7 @@ 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";
+import { applyResolvedHeaders, fetchText, getResolvedHeader, withTimeout } from "./shared.js";
 
 export async function searchWithPiAnthropic(
   selection: PiModelSelection,
@@ -55,35 +55,33 @@ export async function searchWithPiAnthropic(
 
 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}` }),
+    return applyResolvedHeaders(
+      {
+        ...(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",
+      },
+      selection.headers,
+    );
+  }
+
+  return applyResolvedHeaders(
+    {
+      ...(credential ? { "x-api-key": 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",
+      "anthropic-beta": "web-search-2025-03-05",
       "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",
-  };
+    },
+    selection.headers,
+  );
 }
 
 function resolveAnthropicMessagesUrl(baseUrl?: string): string {
@@ -104,32 +102,17 @@ function getAnthropicCredential(selection: PiModelSelection): string | undefined
   return (
     selection.apiKey ??
     getBearerToken(selection.headers) ??
-    getHeaderValue(selection.headers, "x-api-key")
+    getResolvedHeader(selection.headers, "x-api-key")
   );
 }
 
-function getBearerToken(headers?: Record<string, string>): string | undefined {
-  const authorization = getHeaderValue(headers, "authorization");
+function getBearerToken(headers?: Record<string, string | null>): string | undefined {
+  const authorization = getResolvedHeader(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"));
 }