tau-coding-agent 0.1.0

package/extensions/websearch/providers/gemini.browser.ts

@@ -0,0 +1,185 @@
+import type { BrowserCookie, BrowserSession, WebsearchResult } from "../types.js";
+import {
+  browserHeaders,
+  buildCookieHeader,
+  dedupeSources,
+  extractMarkdownSources,
+  hasCookie,
+} from "../normalize.js";
+import { buildWebsearchPrompt } from "./search-prompt.shared.js";
+import { fetchText, withTimeout } from "./shared.js";
+
+const GEMINI_APP_URL = "https://gemini.google.com/app";
+const GEMINI_STREAM_GENERATE_URL =
+  "https://gemini.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate";
+const GOOGLE_LIST_ACCOUNTS_URL =
+  "https://accounts.google.com/ListAccounts?gpsia=1&source=ChromiumBrowser&laf=b64bin&json=standard";
+const REQUIRED_COOKIE_NAMES = ["__Secure-1PSID", "__Secure-1PSIDTS"];
+
+export const browserGemini = {
+  backend: "gemini" as const,
+  domains: ["google.com"],
+  async search(
+    session: BrowserSession,
+    query: string,
+    signal?: AbortSignal,
+  ): Promise<WebsearchResult> {
+    for (const cookieName of REQUIRED_COOKIE_NAMES) {
+      if (!hasCookie(GEMINI_APP_URL, session.cookies, cookieName)) {
+        throw new Error(`Missing required Google cookie: ${cookieName}`);
+      }
+    }
+
+    const accessToken = await fetchAccessToken(session.cookies, signal);
+    let accountLabel: string | null | undefined;
+    void getActiveGoogleEmail(session.cookies, signal)
+      .then((value) => {
+        accountLabel = value;
+      })
+      .catch(() => {
+        accountLabel = null;
+      });
+
+    const answer = await queryGemini(query, session.cookies, accessToken, signal);
+
+    return {
+      backend: "gemini",
+      authSource: session.profile.family,
+      browserName: session.profile.browserName,
+      profile: session.profile.profileName,
+      accountLabel: accountLabel ?? undefined,
+      answer,
+      sources: dedupeSources(extractMarkdownSources(answer)),
+    };
+  },
+};
+
+async function queryGemini(
+  query: string,
+  cookies: BrowserCookie[],
+  accessToken: string,
+  signal?: AbortSignal,
+): Promise<string> {
+  const body = new URLSearchParams();
+  body.set("at", accessToken);
+  body.set(
+    "f.req",
+    JSON.stringify([null, JSON.stringify([[buildWebsearchPrompt(query)], null, null])]),
+  );
+
+  const rawText = await fetchText(GEMINI_STREAM_GENERATE_URL, {
+    method: "POST",
+    headers: {
+      ...browserHeaders({
+        cookieHeader: buildCookieHeader(GEMINI_STREAM_GENERATE_URL, cookies),
+        origin: "https://gemini.google.com",
+        referer: "https://gemini.google.com/",
+        contentType: "application/x-www-form-urlencoded;charset=utf-8",
+      }),
+      host: "gemini.google.com",
+      "x-same-domain": "1",
+    },
+    body: body.toString(),
+    signal: withTimeout(signal, 120_000),
+  });
+
+  const parsed = parseGeminiResponse(rawText);
+  if (!parsed.trim()) {
+    throw new Error("Gemini Web returned an empty response.");
+  }
+
+  return parsed.trim();
+}
+
+async function fetchAccessToken(cookies: BrowserCookie[], signal?: AbortSignal): Promise<string> {
+  const html = await fetchText(GEMINI_APP_URL, {
+    headers: browserHeaders({
+      cookieHeader: buildCookieHeader(GEMINI_APP_URL, cookies),
+      origin: "https://gemini.google.com",
+      referer: "https://gemini.google.com/",
+      accept: "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
+    }),
+    signal: withTimeout(signal, 20_000),
+  });
+
+  for (const key of ["SNlM0e", "thykhd"]) {
+    const match = html.match(new RegExp(`"${key}":"([^\\"]*)"`));
+    if (match?.[1]) return match[1];
+  }
+
+  throw new Error("Could not authenticate with Gemini Web.");
+}
+
+async function getActiveGoogleEmail(
+  cookies: BrowserCookie[],
+  signal?: AbortSignal,
+): Promise<string | null> {
+  try {
+    const response = await fetchText(GOOGLE_LIST_ACCOUNTS_URL, {
+      headers: browserHeaders({
+        cookieHeader: buildCookieHeader(GOOGLE_LIST_ACCOUNTS_URL, cookies),
+        origin: "https://accounts.google.com",
+        referer: "https://accounts.google.com/",
+      }),
+      signal: withTimeout(signal, 10_000),
+    });
+
+    return findFirstEmail(response);
+  } catch {
+    return null;
+  }
+}
+
+function parseGeminiResponse(rawText: string): string {
+  const start = rawText.indexOf("[");
+  const end = rawText.lastIndexOf("]");
+  if (start === -1 || end === -1 || end <= start) {
+    throw new Error("Gemini Web returned an invalid payload.");
+  }
+
+  const responseJson = JSON.parse(rawText.slice(start, end + 1));
+  const parts = Array.isArray(responseJson) ? responseJson : [];
+
+  for (const part of parts) {
+    const payload = getNestedValue(part, [2]);
+    if (typeof payload !== "string") continue;
+
+    try {
+      const parsed = JSON.parse(payload);
+      const candidateList = getNestedValue(parsed, [4]);
+      const firstCandidate = Array.isArray(candidateList) ? candidateList[0] : undefined;
+      const text = getNestedValue(firstCandidate, [1, 0]);
+      if (typeof text === "string" && text.trim().length > 0) {
+        return text;
+      }
+
+      const alternateText = getNestedValue(firstCandidate, [22, 0]);
+      if (typeof alternateText === "string" && alternateText.trim().length > 0) {
+        return alternateText;
+      }
+    } catch {
+      // Ignore non-message chunks.
+    }
+  }
+
+  throw new Error("Gemini Web returned no assistant text.");
+}
+
+function getNestedValue(value: unknown, path: number[]): unknown {
+  let current: unknown = value;
+  for (const index of path) {
+    if (!Array.isArray(current)) return undefined;
+    current = current[index];
+  }
+  return current;
+}
+
+function findFirstEmail(value: string): string | null {
+  const normalized = value
+    .replace(/\\u0040/gi, "@")
+    .replace(/\\x40/gi, "@")
+    .replace(/&#64;/gi, "@")
+    .replace(/&commat;/gi, "@");
+  const match = normalized.match(/\b[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}\b/i);
+  return match?.[0] ?? null;
+}

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

@@ -0,0 +1,108 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+
+import { dedupeSources, extractMarkdownSources, normalizeSourceTitle } from "../normalize.js";
+import type { WebsearchResult, WebsearchSource } from "../types.js";
+import type { PiModelSelection } from "./pi-model.shared.js";
+import { buildWebsearchPrompt } from "./search-prompt.shared.js";
+import { fetchJson, withTimeout } from "./shared.js";
+
+export async function searchWithPiGemini(
+  selection: PiModelSelection,
+  query: string,
+  signal?: AbortSignal,
+): Promise<WebsearchResult> {
+  const interaction = await fetchJson<Record<string, unknown>>(
+    resolveGeminiInteractionsUrl(selection.model.baseUrl),
+    {
+      method: "POST",
+      headers: buildGeminiHeaders(selection),
+      body: JSON.stringify({
+        model: selection.model.id,
+        input: buildWebsearchPrompt(query),
+        tools: [{ googleSearch: {} }],
+      }),
+      signal: withTimeout(signal, 120_000),
+    },
+  );
+
+  const textOutputs = (Array.isArray(interaction.outputs) ? interaction.outputs : []).filter(
+    (output) => {
+      return Boolean(
+        output &&
+        typeof output === "object" &&
+        (output as { type?: unknown }).type === "text" &&
+        typeof (output as { text?: unknown }).text === "string",
+      );
+    },
+  ) as Array<{ text: string; annotations?: Array<{ source?: string }> }>;
+
+  const answer = textOutputs
+    .map((output) => output.text)
+    .join("\n\n")
+    .trim();
+
+  if (!answer) {
+    throw new Error("Gemini API returned no text content.");
+  }
+
+  const annotationSources = textOutputs.flatMap((output) =>
+    extractAnnotationSources(output as { annotations?: Array<{ source?: string }> }),
+  );
+  return {
+    backend: "gemini",
+    authSource: "pi",
+    answer,
+    sources: dedupeSources([...annotationSources, ...extractMarkdownSources(answer)]),
+  };
+}
+
+function resolveGeminiInteractionsUrl(baseUrl?: string): string {
+  const normalized = String(baseUrl || "https://generativelanguage.googleapis.com/v1beta").replace(
+    /\/+$/,
+    "",
+  );
+  return normalized.endsWith("/interactions") ? normalized : `${normalized}/interactions`;
+}
+
+function buildGeminiHeaders(selection: PiModelSelection): Record<string, string> {
+  const headers = { ...selection.headers };
+
+  return {
+    ...headers,
+    ...(hasHeader(headers, "x-goog-api-key") || !selection.apiKey
+      ? {}
+      : { "x-goog-api-key": selection.apiKey }),
+    "content-type": "application/json",
+    accept: "application/json",
+  };
+}
+
+function hasHeader(headers: Record<string, string>, name: string): boolean {
+  return Object.keys(headers).some((key) => key.toLowerCase() === name.toLowerCase());
+}
+
+function extractAnnotationSources(output: {
+  annotations?: Array<{ source?: string; url?: string; title?: string }>;
+}): WebsearchSource[] {
+  const sources: WebsearchSource[] = [];
+
+  for (const annotation of output.annotations ?? []) {
+    const url =
+      typeof annotation?.url === "string"
+        ? annotation.url
+        : typeof annotation?.source === "string"
+          ? annotation.source
+          : null;
+    if (!url || !/^https?:\/\//i.test(url)) continue;
+    sources.push({
+      title: normalizeSourceTitle(url, annotation.title),
+      url,
+    });
+  }
+
+  return sources;
+}
+
+export function isPiGeminiModel(model: Model<Api>): boolean {
+  return model.api === "google-generative-ai" && model.provider === "google";
+}

package/extensions/websearch/providers/openai-codex.browser.ts

@@ -0,0 +1,77 @@
+import type { BrowserSession, WebsearchResult } from "../types.js";
+import { browserHeaders, buildCookieHeader } from "../normalize.js";
+import { fetchJson, withTimeout } from "./shared.js";
+import { decodeJwtAccountId, runOpenAICodexSearch } from "./openai-codex.shared.js";
+
+const CHATGPT_BASE_URL = "https://chatgpt.com";
+const SESSION_URL = `${CHATGPT_BASE_URL}/api/auth/session`;
+const PREFERRED_MODELS = ["gpt-5.4", "gpt-5.3-codex-spark", "gpt-5.1-codex-mini"];
+
+interface ChatGptSessionPayload {
+  accessToken?: string;
+  user?: {
+    email?: string;
+    name?: string;
+  };
+}
+
+export const browserOpenAICodex = {
+  backend: "openai-codex" as const,
+  domains: ["chatgpt.com"],
+  async search(
+    session: BrowserSession,
+    query: string,
+    signal?: AbortSignal,
+  ): Promise<WebsearchResult> {
+    const cookieHeader = buildCookieHeader(SESSION_URL, session.cookies);
+    if (!cookieHeader) {
+      throw new Error("No ChatGPT session cookies found in browser profile.");
+    }
+
+    const chatSession = await fetchJson<ChatGptSessionPayload>(SESSION_URL, {
+      headers: browserHeaders({
+        cookieHeader,
+        origin: CHATGPT_BASE_URL,
+        referer: `${CHATGPT_BASE_URL}/`,
+      }),
+      signal: withTimeout(signal, 20_000),
+    });
+
+    if (!chatSession.accessToken) {
+      throw new Error("No ChatGPT access token found in session.");
+    }
+
+    const result = await searchWithPreferredModel(chatSession.accessToken, query, signal);
+
+    return {
+      backend: "openai-codex",
+      authSource: session.profile.family,
+      browserName: session.profile.browserName,
+      profile: session.profile.profileName,
+      accountLabel: chatSession.user?.email ?? chatSession.user?.name,
+      answer: result.answer,
+      sources: result.sources,
+    };
+  },
+};
+
+async function searchWithPreferredModel(accessToken: string, query: string, signal?: AbortSignal) {
+  let lastError: string | null = null;
+
+  for (const model of PREFERRED_MODELS) {
+    try {
+      return await runOpenAICodexSearch({
+        apiKey: accessToken,
+        accountId: decodeJwtAccountId(accessToken),
+        model,
+        query,
+        baseUrl: `${CHATGPT_BASE_URL}/backend-api`,
+        signal,
+      });
+    } catch (error) {
+      lastError = error instanceof Error ? error.message : String(error);
+    }
+  }
+
+  throw new Error(lastError ?? "ChatGPT Codex search failed.");
+}

package/extensions/websearch/providers/openai-codex.pi.ts

@@ -0,0 +1,49 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+
+import type { WebsearchResult } from "../types.js";
+import { decodeJwtAccountId, runOpenAICodexSearch } from "./openai-codex.shared.js";
+import type { PiModelSelection } from "./pi-model.shared.js";
+
+export async function searchWithPiOpenAICodex(
+  selection: PiModelSelection,
+  query: string,
+  signal?: AbortSignal,
+): Promise<WebsearchResult> {
+  const apiKey = resolveApiKey(selection);
+  if (!apiKey) {
+    throw new Error("OpenAI Codex auth is not configured.");
+  }
+
+  const result = await runOpenAICodexSearch({
+    apiKey,
+    accountId: decodeJwtAccountId(apiKey),
+    model: selection.model.id,
+    query,
+    baseUrl: selection.model.baseUrl,
+    headers: selection.headers,
+    signal,
+  });
+
+  return {
+    backend: "openai-codex",
+    authSource: "pi",
+    answer: result.answer,
+    sources: result.sources,
+  };
+}
+
+function resolveApiKey(selection: PiModelSelection): string | undefined {
+  if (selection.apiKey) return selection.apiKey;
+
+  const authorization = Object.entries(selection.headers ?? {}).find(
+    ([name]) => name.toLowerCase() === "authorization",
+  )?.[1];
+  if (!authorization) return undefined;
+
+  const match = authorization.match(/^Bearer\s+(.+)$/i);
+  return match?.[1];
+}
+
+export function isPiOpenAICodexModel(model: Model<Api>): boolean {
+  return model.api === "openai-codex-responses" && model.provider === "openai-codex";
+}

package/extensions/websearch/providers/openai-codex.shared.ts

@@ -0,0 +1,123 @@
+import type { WebsearchSource } from "../types.js";
+import { dedupeSources, extractMarkdownSources } from "../normalize.js";
+import { buildWebsearchPrompt, WEBSEARCH_SYSTEM_PROMPT } from "./search-prompt.shared.js";
+import { readEventStream, withTimeout } from "./shared.js";
+
+export function decodeJwtAccountId(jwt: string | undefined): string | undefined {
+  if (!jwt || typeof jwt !== "string") return undefined;
+
+  try {
+    const parts = jwt.split(".");
+    if (parts.length !== 3) return undefined;
+    const payload = JSON.parse(Buffer.from(parts[1], "base64url").toString("utf8")) as Record<
+      string,
+      unknown
+    >;
+    const auth = payload["https://api.openai.com/auth"];
+    return auth && typeof auth === "object"
+      ? ((auth as Record<string, unknown>).chatgpt_account_id as string | undefined)
+      : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+export function resolveCodexUrl(baseUrl = "https://chatgpt.com/backend-api"): string {
+  const normalized = String(baseUrl || "https://chatgpt.com/backend-api").replace(/\/+$/, "");
+  if (normalized.endsWith("/codex/responses")) return normalized;
+  if (normalized.endsWith("/codex")) return `${normalized}/responses`;
+  return `${normalized}/codex/responses`;
+}
+
+export async function runOpenAICodexSearch(options: {
+  apiKey: string;
+  model: string;
+  query: string;
+  baseUrl?: string;
+  accountId?: string;
+  headers?: Record<string, string>;
+  signal?: AbortSignal;
+}): Promise<{ answer: string; sources: WebsearchSource[] }> {
+  const response = await fetch(resolveCodexUrl(options.baseUrl), {
+    method: "POST",
+    headers: {
+      ...options.headers,
+      authorization: `Bearer ${options.apiKey}`,
+      ...(options.accountId ? { "chatgpt-account-id": options.accountId } : {}),
+      "content-type": "application/json",
+      accept: "text/event-stream",
+      "OpenAI-Beta": "responses=experimental",
+      originator: "pi-websearch",
+    },
+    body: JSON.stringify({
+      model: options.model,
+      store: false,
+      stream: true,
+      instructions: WEBSEARCH_SYSTEM_PROMPT,
+      input: [{ role: "user", content: buildWebsearchPrompt(options.query) }],
+      tools: [{ type: "web_search" }],
+      tool_choice: "auto",
+    }),
+    signal: withTimeout(options.signal, 120_000),
+  });
+
+  let answer = "";
+  let fallbackAnswer = "";
+
+  await readEventStream(response, ({ data }) => {
+    if (!data.trim()) return;
+
+    try {
+      const event = JSON.parse(data) as Record<string, unknown>;
+
+      if (event.type === "response.output_text.delta" && typeof event.delta === "string") {
+        answer += event.delta;
+      }
+
+      if (event.type === "response.output_item.done") {
+        const item = event.item as Record<string, unknown> | undefined;
+        const content = Array.isArray(item?.content) ? item.content : [];
+        const fullText = content
+          .filter(
+            (part) =>
+              part &&
+              typeof part === "object" &&
+              (part as Record<string, unknown>).type === "output_text",
+          )
+          .map((part) => (part as Record<string, unknown>).text)
+          .filter((text): text is string => typeof text === "string")
+          .join("\n");
+        if (fullText) fallbackAnswer = fullText;
+      }
+
+      if (event.type === "response.failed" || event.type === "error") {
+        const failedResponse = event.response;
+        const eventMessage = typeof event.message === "string" ? event.message : undefined;
+        if (failedResponse && typeof failedResponse === "object") {
+          const error = (failedResponse as Record<string, unknown>).error;
+          if (
+            error &&
+            typeof error === "object" &&
+            typeof (error as Record<string, unknown>).message === "string"
+          ) {
+            throw new Error((error as Record<string, unknown>).message as string);
+          }
+        }
+        throw new Error(eventMessage ?? "OpenAI Codex search failed.");
+      }
+    } catch (error) {
+      if (error instanceof SyntaxError) return;
+      throw error;
+    }
+  });
+
+  const finalAnswer = (answer || fallbackAnswer).trim();
+  if (!finalAnswer) {
+    throw new Error("OpenAI Codex returned an empty response.");
+  }
+
+  return {
+    answer: finalAnswer,
+    sources: dedupeSources(extractMarkdownSources(finalAnswer)),
+  };
+}

package/extensions/websearch/providers/pi-model.shared.ts

@@ -0,0 +1,86 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
+
+export interface PiModelSelection {
+  model: Model<Api>;
+  apiKey?: string;
+  headers?: Record<string, string>;
+}
+
+export async function selectCurrentPiModel(
+  ctx: ExtensionContext,
+  predicate: (model: Model<Api>) => boolean,
+): Promise<PiModelSelection | null> {
+  if (!ctx.model || !predicate(ctx.model)) return null;
+
+  return resolvePiModelSelection(ctx.model, ctx);
+}
+
+export async function selectFallbackPiModel(
+  ctx: ExtensionContext,
+  predicate: (model: Model<Api>) => boolean,
+): Promise<PiModelSelection | null> {
+  const currentKey =
+    ctx.model && predicate(ctx.model) ? `${ctx.model.provider}:${ctx.model.id}` : null;
+  const candidates = rankPiModels(
+    ctx.modelRegistry.getAvailable().filter((model) => {
+      if (!predicate(model)) return false;
+      return `${model.provider}:${model.id}` !== currentKey;
+    }),
+  );
+
+  for (const model of candidates) {
+    const selection = await resolvePiModelSelection(model, ctx);
+    if (selection) return selection;
+  }
+
+  return null;
+}
+
+async function resolvePiModelSelection(
+  model: Model<Api>,
+  ctx: Pick<ExtensionContext, "modelRegistry">,
+): Promise<PiModelSelection | null> {
+  const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+  return auth.ok ? { model, apiKey: auth.apiKey, headers: auth.headers } : null;
+}
+
+function rankPiModels(models: Model<Api>[]): Model<Api>[] {
+  const seen = new Set<string>();
+
+  return [...models]
+    .filter((model) => {
+      const key = `${model.provider}:${model.id}`;
+      if (seen.has(key)) return false;
+      seen.add(key);
+      return true;
+    })
+    .sort(comparePiModels);
+}
+
+function comparePiModels(left: Model<Api>, right: Model<Api>): number {
+  return (
+    compareBoolean(Boolean(right.reasoning), Boolean(left.reasoning)) ||
+    compareNumber(right.contextWindow, left.contextWindow) ||
+    compareNumber(right.maxTokens, left.maxTokens) ||
+    compareNumber(modelCostScore(right), modelCostScore(left)) ||
+    right.id.localeCompare(left.id)
+  );
+}
+
+function modelCostScore(model: Model<Api>): number {
+  return (
+    (model.cost?.input ?? 0) +
+    (model.cost?.output ?? 0) +
+    (model.cost?.cacheRead ?? 0) +
+    (model.cost?.cacheWrite ?? 0)
+  );
+}
+
+function compareBoolean(left: boolean, right: boolean): number {
+  return Number(left) - Number(right);
+}
+
+function compareNumber(left: number, right: number): number {
+  return left - right;
+}

package/extensions/websearch/providers/search-prompt.shared.ts

@@ -0,0 +1,17 @@
+export const WEBSEARCH_SYSTEM_PROMPT =
+  "You are a web research assistant. Always use web search, produce practical summaries, and include full canonical URLs (no shortened links).";
+
+export function buildWebsearchPrompt(query: string): string {
+  return `Search the internet for: ${query}
+
+Requirements:
+- Always use web search.
+- Use current web information.
+- Prefer primary or official sources.
+- Open and read relevant source pages before answering.
+- Base the answer on the content of the pages you read, not on search-result snippets alone.
+- Be concise.
+- Call out important source disagreements.
+- Do not mention internal tools or implementation details.
+- End with a Sources section that includes full canonical URLs.`;
+}