pi-web-search 1.3.0 → 1.3.1

package/README.md

@@ -12,6 +12,7 @@ Search the web using your currently selected model. Automatically picks the righ
 |---|---|
 | Google Gemini | Grounding with Google Search |
 | OpenAI | Responses API web search |
+| OpenAI Codex | Codex Responses API web search |
 | Anthropic | Messages API web search |
 
 Supports passing up to 20 additional URLs to analyze alongside the query.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-web-search",
-  "version": "1.3.0",
+  "version": "1.3.1",
   "description": "Provider-native web search for pi across Google Gemini, OpenAI, and Anthropic, plus Gemini URL Context",
   "type": "module",
   "main": "src/index.ts",
@@ -37,8 +37,8 @@
     "test:real:url-context": "node tests/real-url-context.mjs"
   },
   "devDependencies": {
-    "@earendil-works/pi-ai": "^0.75.4",
-    "@earendil-works/pi-coding-agent": "^0.75.4",
+    "@earendil-works/pi-ai": "^0.80.3",
+    "@earendil-works/pi-coding-agent": "^0.80.3",
     "@types/node": "^25.9.1",
     "typebox": "^1.1.38",
     "typescript": "^6.0.3"
@@ -49,8 +49,8 @@
     ]
   },
   "peerDependencies": {
-    "@earendil-works/pi-ai": "*",
-    "@earendil-works/pi-coding-agent": "*",
+    "@earendil-works/pi-ai": ">=0.80.3",
+    "@earendil-works/pi-coding-agent": ">=0.80.3",
     "typebox": "*"
   }
 }

package/src/api.ts

@@ -1,12 +1,13 @@
 import type { ExtensionContext, AgentToolUpdateCallback } from "@earendil-works/pi-coding-agent";
-import { type Model } from "@earendil-works/pi-ai";
+import type { Api, Model } from "@earendil-works/pi-ai";
+import { getEnvApiKey } from "@earendil-works/pi-ai/compat";
 import { TextEncoder, TextDecoder } from "util";
 
 // --- Provider Configuration ---
 
 type ProviderKind = "google" | "openai" | "anthropic" | "unsupported";
 
-type GoogleRequestBuilder = (model: Model<any>, body: any) => { url: string; headers: Record<string, string>; body: any };
+type GoogleRequestBuilder = (model: Model<Api>, body: any) => { url: string; headers: Record<string, string>; body: any };
 
 type ProviderConfig = {
     kind: ProviderKind;
@@ -31,14 +32,14 @@ const GOOGLE_PROVIDERS: Record<string, ProviderConfig> = {
     }
 };
 
-export function getProviderKind(model: Model<any>): ProviderKind {
+export function getProviderKind(model: Model<Api>): ProviderKind {
     if (GOOGLE_PROVIDERS[model.provider] || GOOGLE_PROVIDERS[model.api]) return "google";
-    if (model.api === "openai-responses") return "openai";
+    if (model.api === "openai-responses" || model.api === "openai-codex-responses") return "openai";
     if (model.api === "anthropic-messages") return "anthropic";
     return "unsupported";
 }
 
-export function getConfig(model: Model<any>): ProviderConfig {
+export function getConfig(model: Model<Api>): ProviderConfig {
     const googleConfig = GOOGLE_PROVIDERS[model.provider] || GOOGLE_PROVIDERS[model.api];
     if (googleConfig) return googleConfig;
     const kind = getProviderKind(model);
@@ -51,28 +52,34 @@ type ResolvedAuth =
     | { ok: true; apiKey?: string; headers?: Record<string, string>; }
     | { ok: false; error: string; };
 
+function getEnvAuth(model: Model<Api>): Extract<ResolvedAuth, { ok: true }> | undefined {
+    const apiKey = getEnvApiKey(model.provider);
+    return apiKey ? { ok: true, apiKey } : undefined;
+}
+
 /**
  * Get API key and headers for a model.
- * Compatible with both new pi versions (getApiKeyAndHeaders) and old versions (getApiKey).
  */
-async function getAuth(ctx: ExtensionContext, model: Model<any>): Promise<ResolvedAuth> {
-    const registry = ctx.modelRegistry as any;
-    
-    // Try new API first (pi >= 0.63.0)
-    if (typeof registry.getApiKeyAndHeaders === 'function') {
-        return await registry.getApiKeyAndHeaders(model);
-    }
-    
-    // Fallback to old API (pi < 0.63.0)
-    if (typeof registry.getApiKey === 'function') {
-        const apiKey = await registry.getApiKey(model);
-        if (apiKey === undefined || apiKey === null) {
-            return { ok: false, error: "No API key configured for model" };
-        }
-        return { ok: true, apiKey };
-    }
-    
-    return { ok: false, error: "Model registry does not support API key retrieval" };
+async function getAuth(ctx: ExtensionContext, model: Model<Api>): Promise<ResolvedAuth> {
+    const resolved = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+    if (!resolved.ok) return resolved;
+
+    // pi-coding-agent 0.80.1+ returns { ok: true } from getApiKeyAndHeaders()
+    // when auth only comes from provider env vars such as ANTHROPIC_API_KEY.
+    // The main agent still works because pi-ai's streamSimple() performs its own
+    // getEnvApiKey() fallback, but this extension calls fetch() directly, so it
+    // must mirror that fallback while preserving explicit model/auth headers.
+    const envAuth = !resolved.apiKey && !hasAuthHeader(resolved.headers) ? getEnvAuth(model) : undefined;
+    return envAuth ? { ...resolved, apiKey: envAuth.apiKey } : resolved;
+}
+
+function hasAuthHeader(headers?: Record<string, string>): boolean {
+    if (!headers) return false;
+    return Object.entries(headers).some(([name, value]) => {
+        if (!value) return false;
+        const normalized = name.toLowerCase();
+        return normalized === "authorization" || normalized === "x-api-key" || normalized === "x-goog-api-key";
+    });
 }
 
 // --- Streaming API Call ---
@@ -126,7 +133,7 @@ type SseEvent = {
 async function readSseEvents(
     response: Response,
     signal: AbortSignal | undefined,
-    onEvent: (event: SseEvent) => void | Promise<void>
+    onEvent: (event: SseEvent) => boolean | void | Promise<boolean | void>
 ): Promise<void> {
     if (!response.body) {
         throw new Error("No response body");
@@ -137,59 +144,78 @@ async function readSseEvents(
     let buffer = "";
     let currentEventData = "";
     let currentEventName = "";
+    let stopRequested = false;
+    let reachedEof = false;
 
-    const flushEvent = async () => {
-        if (!currentEventData) return;
+    const flushEvent = async (): Promise<boolean> => {
+        if (!currentEventData) return false;
         const raw = currentEventData.trim();
         currentEventData = "";
         const eventName = currentEventName;
         currentEventName = "";
-        if (!raw || raw === "[DONE]") return;
+        if (!raw || raw === "[DONE]") return false;
 
         let data: any;
         try {
             data = JSON.parse(raw);
         } catch {
-            return;
+            return false;
         }
-        await onEvent({ event: eventName, data });
+        return await onEvent({ event: eventName, data }) === true;
     };
 
-    while (true) {
-        if (signal?.aborted) {
-            throw new Error("Request was aborted");
-        }
+    try {
+        readLoop: while (true) {
+            if (signal?.aborted) {
+                throw new Error("Request was aborted");
+            }
 
-        const { done, value } = await reader.read();
-        if (done) break;
+            const { done, value } = await reader.read();
+            if (done) {
+                reachedEof = true;
+                break;
+            }
 
-        buffer += decoder.decode(value, { stream: true });
-        const lines = buffer.split("\n");
-        buffer = lines.pop() || "";
+            buffer += decoder.decode(value, { stream: true });
+            const lines = buffer.split("\n");
+            buffer = lines.pop() || "";
+
+            for (const line of lines) {
+                if (line === "" || line === "\r") {
+                    if (await flushEvent()) {
+                        stopRequested = true;
+                        break readLoop;
+                    }
+                    continue;
+                }
 
-        for (const line of lines) {
-            if (line === "" || line === "\r") {
-                await flushEvent();
-                continue;
+                if (line.startsWith("data:")) {
+                    const data = line.slice(5).trim();
+                    currentEventData = currentEventData ? currentEventData + "\n" + data : data;
+                } else if (line.startsWith("event:")) {
+                    currentEventName = line.slice(6).trim();
+                }
             }
+        }
 
+        if (!stopRequested && buffer.trim()) {
+            const line = buffer.trim();
             if (line.startsWith("data:")) {
                 const data = line.slice(5).trim();
                 currentEventData = currentEventData ? currentEventData + "\n" + data : data;
-            } else if (line.startsWith("event:")) {
-                currentEventName = line.slice(6).trim();
             }
         }
-    }
-
-    if (buffer.trim()) {
-        const line = buffer.trim();
-        if (line.startsWith("data:")) {
-            const data = line.slice(5).trim();
-            currentEventData = currentEventData ? currentEventData + "\n" + data : data;
+        if (!stopRequested) stopRequested = await flushEvent();
+    } finally {
+        if (!reachedEof) {
+            try {
+                await reader.cancel();
+            } catch {
+                // Ignore cancellation failures while cleaning up an interrupted stream.
+            }
         }
+        reader.releaseLock();
     }
-    await flushEvent();
 }
 
 function extractPromptFromGeminiBody(body: any): string {
@@ -210,6 +236,34 @@ function trimTrailingSlash(value: string): string {
     return value.replace(/\/+$/, "");
 }
 
+function isOpenAICodexModel(model: Model<Api>): boolean {
+    return model.api === "openai-codex-responses";
+}
+
+function resolveOpenAIResponsesUrl(model: Model<Api>): string {
+    const base = trimTrailingSlash(model.baseUrl);
+    if (!isOpenAICodexModel(model)) return `${base}/responses`;
+    if (base.endsWith("/codex/responses")) return base;
+    if (base.endsWith("/codex")) return `${base}/responses`;
+    return `${base}/codex/responses`;
+}
+
+function extractOpenAICodexAccountId(token: string): string {
+    try {
+        const parts = token.split(".");
+        if (parts.length !== 3) throw new Error("Invalid token");
+        const base64 = parts[1].replace(/-/g, "+").replace(/_/g, "/");
+        const padded = base64.padEnd(Math.ceil(base64.length / 4) * 4, "=");
+        const bytes = Uint8Array.from(atob(padded), (char) => char.charCodeAt(0));
+        const payload = JSON.parse(new TextDecoder().decode(bytes));
+        const accountId = payload?.["https://api.openai.com/auth"]?.chatgpt_account_id;
+        if (typeof accountId !== "string" || !accountId) throw new Error("Missing account ID");
+        return accountId;
+    } catch {
+        throw new Error("Failed to extract ChatGPT account ID from openai-codex credentials");
+    }
+}
+
 function resolveAnthropicMessagesUrl(baseUrl: string): string {
     const base = trimTrailingSlash(baseUrl);
     return base.endsWith("/v1") ? `${base}/messages` : `${base}/v1/messages`;
@@ -459,7 +513,7 @@ function extractGoogleSearchDetails(groundingMetadata: any): { searchQueries: st
 
 async function callGoogleStream(
     ctx: ExtensionContext,
-    model: Model<any>,
+    model: Model<Api>,
     body: any,
     onUpdate?: AgentToolUpdateCallback,
     signal?: AbortSignal
@@ -551,7 +605,7 @@ async function callGoogleStream(
 
 async function callOpenAIStream(
     ctx: ExtensionContext,
-    model: Model<any>,
+    model: Model<Api>,
     prompt: string,
     onUpdate?: AgentToolUpdateCallback,
     signal?: AbortSignal
@@ -561,31 +615,55 @@ async function callOpenAIStream(
         throw new Error(auth.error || "Failed to get API key and headers");
     }
 
-    const headers: Record<string, string> = {
-        "Content-Type": "application/json",
-        "Accept": "text/event-stream",
-        ...(model.headers || {}),
-        ...(auth.headers || {}),
-    };
-    if (auth.apiKey && !headers.Authorization && !headers.authorization) {
-        headers.Authorization = `Bearer ${auth.apiKey}`;
+    const headers = new Headers();
+    for (const [name, value] of Object.entries(model.headers || {})) headers.set(name, value);
+    for (const [name, value] of Object.entries(auth.headers || {})) headers.set(name, value);
+    if (!headers.has("Content-Type")) headers.set("Content-Type", "application/json");
+    if (!headers.has("Accept")) headers.set("Accept", "text/event-stream");
+    if (auth.apiKey && !headers.has("Authorization")) headers.set("Authorization", `Bearer ${auth.apiKey}`);
+
+    const isCodex = isOpenAICodexModel(model);
+    if (isCodex) {
+        const authorization = headers.get("Authorization");
+        const hasBearerAuth = typeof authorization === "string" && /^Bearer\s+\S+/i.test(authorization);
+        if (!auth.apiKey && !hasBearerAuth) {
+            throw new Error("No OAuth credential configured for openai-codex model");
+        }
+        if (!headers.has("chatgpt-account-id")) {
+            if (!auth.apiKey) {
+                throw new Error("No ChatGPT account ID configured for openai-codex model");
+            }
+            headers.set("chatgpt-account-id", extractOpenAICodexAccountId(auth.apiKey));
+        }
+        if (!headers.has("originator")) headers.set("originator", "codex_cli_rs");
     }
+    const requestHeaders = Object.fromEntries(headers.entries());
 
     const requestBody: any = {
         model: model.id,
-        input: prompt,
+        input: isCodex
+            ? [{ role: "user", content: [{ type: "input_text", text: prompt }] }]
+            : prompt,
         tools: [{ type: "web_search" }],
-        include: ["web_search_call.action.sources", "web_search_call.results"],
+        include: isCodex
+            ? ["web_search_call.action.sources"]
+            : ["web_search_call.action.sources", "web_search_call.results"],
         stream: true,
         store: false,
     };
     if (model.reasoning) {
         requestBody.reasoning = { effort: "none" };
     }
+    if (isCodex) {
+        requestBody.instructions = "Answer the user's request using web search when needed.";
+        requestBody.text = { verbosity: "low" };
+        requestBody.tool_choice = "required";
+        requestBody.parallel_tool_calls = true;
+    }
 
-    const response = await fetch(`${trimTrailingSlash(model.baseUrl)}/responses`, {
+    const response = await fetch(resolveOpenAIResponsesUrl(model), {
         method: "POST",
-        headers,
+        headers: requestHeaders,
         body: JSON.stringify(requestBody),
         signal
     });
@@ -649,7 +727,14 @@ async function callOpenAIStream(
                 raw: action,
             });
         }
-        nativeSearchCalls.push(call);
+        const existingCall = call.id ? nativeSearchCalls.find((existing) => existing.id === call.id) : undefined;
+        if (existingCall) {
+            Object.assign(existingCall, Object.fromEntries(
+                Object.entries(call).filter(([, value]) => value !== undefined)
+            ));
+        } else {
+            nativeSearchCalls.push(call);
+        }
     };
 
     const collectFromResponse = (response: any) => {
@@ -664,7 +749,10 @@ async function callOpenAIStream(
     };
 
     await readSseEvents(response, signal, ({ data: event }) => {
-        if (event.type === "response.output_text.delta") {
+        if (event.type === "error" || event.type === "response.failed") {
+            const message = event.message || event.error?.message || event.response?.error?.message;
+            throw new Error(message || JSON.stringify(event.error || event.response?.error || event));
+        } else if (event.type === "response.output_text.delta") {
             accumulatedText += event.delta || "";
             onUpdate?.({
                 content: [{ type: "text", text: accumulatedText }],
@@ -674,8 +762,12 @@ async function callOpenAIStream(
             collectAnnotation(event.annotation);
         } else if (event.type === "response.output_item.added" || event.type === "response.output_item.done") {
             collectWebSearchCall(event.item);
-        } else if (event.type === "response.completed") {
+        } else if (event.type === "response.incomplete" || event.response?.status === "incomplete") {
             collectFromResponse(event.response);
+            if (isCodex) return true;
+        } else if (event.type === "response.completed" || event.type === "response.done") {
+            collectFromResponse(event.response);
+            if (isCodex) return true;
         } else if (event.type === "response.web_search_call.in_progress" || event.type === "response.web_search_call.searching" || event.type === "response.web_search_call.completed") {
             pushNativeSearchEvent(nativeSearchEvents, event.type);
             const call = nativeSearchCalls.find((item) => item.id === event.item_id);
@@ -687,11 +779,6 @@ async function callOpenAIStream(
                     details: { streaming: true, searching: true }
                 });
             }
-        } else if (event.type === "response.failed") {
-            const error = event.response?.error;
-            throw new Error(error?.message || JSON.stringify(error || event.response || event));
-        } else if (event.type === "error") {
-            throw new Error(event.message || JSON.stringify(event));
         }
     });
 
@@ -725,7 +812,7 @@ async function callOpenAIStream(
 
 async function callAnthropicStream(
     ctx: ExtensionContext,
-    model: Model<any>,
+    model: Model<Api>,
     prompt: string,
     onUpdate?: AgentToolUpdateCallback,
     signal?: AbortSignal
@@ -891,7 +978,7 @@ async function callAnthropicStream(
 
 export async function callApiStream(
     ctx: ExtensionContext,
-    model: Model<any>,
+    model: Model<Api>,
     body: any,
     onUpdate?: AgentToolUpdateCallback,
     signal?: AbortSignal

package/src/index.ts

@@ -1,5 +1,5 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { type Model } from "@earendil-works/pi-ai";
+import type { Api, Model } from "@earendil-works/pi-ai";
 import { getProviderKind } from "./api.ts";
 import { webSearch, WebSearchSchema } from "./web_search.ts";
 import { urlContext, UrlContextSchema } from "./url_context.ts";
@@ -7,7 +7,7 @@ import { urlContext, UrlContextSchema } from "./url_context.ts";
 const WEB_SEARCH_TOOL = "web_search";
 const URL_CONTEXT_TOOL = "url_context";
 
-function supportsUrlContext(model: Model<any> | undefined) {
+function supportsUrlContext(model: Model<Api> | undefined) {
     return !!model && getProviderKind(model) === "google";
 }
 
@@ -24,7 +24,7 @@ export function createModelScopedToolManager(pi: Pick<ExtensionAPI, "getActiveTo
     let lastAppliedActiveTools: Set<string> | undefined;
     let suppressedTools = new Set<string>();
 
-    const sync = (model: Model<any> | undefined) => {
+    const sync = (model: Model<Api> | undefined) => {
         const currentActiveTools = new Set(pi.getActiveTools());
 
         if (!preferredActiveTools) {

package/src/utils.ts

@@ -1,5 +1,5 @@
 import type { ExtensionContext, AgentToolResult } from "@earendil-works/pi-coding-agent";
-import { type Model } from "@earendil-works/pi-ai";
+import type { Api, Model } from "@earendil-works/pi-ai";
 import { truncateHead, DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES } from "@earendil-works/pi-coding-agent";
 import { readFileSync } from "node:fs";
 import { homedir } from "node:os";
@@ -18,7 +18,7 @@ export function formatResult(text: string, details: any): AgentToolResult<any> {
 
 // --- Model Selection ---
 
-const SUPPORTED_PROVIDERS = ["google-generative-ai", "openai-responses", "anthropic-messages"];
+const SUPPORTED_PROVIDERS = ["google-generative-ai", "openai-responses", "openai-codex-responses", "anthropic-messages"];
 const WEB_SEARCH_CONFIG_PATH = join(homedir(), ".pi", "agent", "web-search.json");
 
 type WebSearchModelConfig =
@@ -26,12 +26,12 @@ type WebSearchModelConfig =
     | { status: "configured"; path: string; provider: string; modelId: string; }
     | { status: "invalid"; path: string; error: string; };
 
-function isSupportedSearchModel(model: Model<any> | undefined): model is Model<any> {
+function isSupportedSearchModel(model: Model<Api> | undefined): model is Model<Api> {
     if (!model) return false;
     return getProviderKind(model) !== "unsupported";
 }
 
-function describeModel(model: Model<any>): string {
+function describeModel(model: Model<Api>): string {
     return `${model.id} (${model.provider}/${model.api})`;
 }
 
@@ -82,13 +82,13 @@ function getAvailableSupportedModels(ctx: ExtensionContext): string[] {
     }
 }
 
-export async function getModel(ctx: ExtensionContext): Promise<Model<any> | undefined> {
+export async function getModel(ctx: ExtensionContext): Promise<Model<Api> | undefined> {
     // Only use the currently selected model. Do not silently fall back to another
     // configured model, because that can surprise users with unexpected API costs.
     return isSupportedSearchModel(ctx.model) ? ctx.model : undefined;
 }
 
-export async function getWebSearchModel(ctx: ExtensionContext): Promise<Model<any> | undefined> {
+export async function getWebSearchModel(ctx: ExtensionContext): Promise<Model<Api> | undefined> {
     const config = readWebSearchModelConfig();
     if (config.status === "invalid") return undefined;