@oh-my-pi/pi-coding-agent 17.0.4 → 17.0.6

package/src/vibe/runtime.ts

@@ -33,6 +33,7 @@ import { type AgentDefinition, type AgentProgress, oneLineLabel, type SingleResu
 import type { ToolSession } from "../tools";
 import { formatDuration } from "../tools/render-utils";
 import { ToolError } from "../tools/tool-errors";
+import { calculateTokensPerSecond } from "../utils/token-rate";
 
 /** The two worker CLI flavors the director drives. */
 export type VibeCli = "fast" | "good";
@@ -207,6 +208,26 @@ export class VibeSessionRegistry {
 		VibeSessionRegistry.#global = undefined;
 	}
 
+	/**
+	 * Insert a bare worker record without the spawn/job machinery. Test-only —
+	 * lets {@link aggregateVibeWorkerTokensPerSecond} be exercised against a
+	 * fake roster + AgentRegistry session without driving a real turn.
+	 */
+	registerRecordForTests(record: { id: string; cli?: VibeCli; ownerId: string; state?: VibeSessionState }): void {
+		this.#records.set(record.id, {
+			id: record.id,
+			cli: record.cli ?? "fast",
+			ownerId: record.ownerId,
+			agent: getBundledAgent("sonic")!,
+			state: record.state ?? "running",
+			createdAt: Date.now(),
+			lastActivityAt: Date.now(),
+			queue: [],
+			turnCount: 0,
+			killed: false,
+		});
+	}
+
 	readonly #records = new Map<string, VibeRecord>();
 
 	#manager(session: ToolSession): AsyncJobManager {
@@ -708,3 +729,32 @@ export class VibeSessionRegistry {
 		return text;
 	}
 }
+
+/**
+ * Aggregate tok/s across every live vibe worker session owned by `ownerId`.
+ * Returns null when no workers are streaming (so callers can fall back to
+ * their own rate unchanged). The director is often idle while workers stream,
+ * so without this aggregation the status-line tok/s badge would show a stale
+ * value while parallel work is actively generating tokens.
+ *
+ * Reads each worker's last assistant message via {@link calculateTokensPerSecond}
+ * — the same leaf calculator the main status line uses — so worker rates are
+ * computed identically to the main session's rate.
+ */
+export function aggregateVibeWorkerTokensPerSecond(ownerId: string): number | null {
+	const ids = VibeSessionRegistry.global().listIds(ownerId);
+	if (ids.length === 0) return null;
+	let total = 0;
+	let any = false;
+	const registry = AgentRegistry.global();
+	for (const id of ids) {
+		const workerSession = registry.get(id)?.session;
+		if (!workerSession?.isStreaming) continue;
+		const rate = calculateTokensPerSecond(workerSession.state.messages, true);
+		if (rate !== null) {
+			total += rate;
+			any = true;
+		}
+	}
+	return any ? total : null;
+}

package/src/web/search/index.ts

@@ -134,10 +134,7 @@ async function executeSearch(
 	const explicitProvider = params.provider;
 	let candidates: SearchProviderCandidate[];
 	if (explicitProvider && explicitProvider !== "auto") {
-		const provider = await getSearchProvider(explicitProvider);
-		candidates = (await provider.isExplicitlyAvailable(authStorage))
-			? [{ id: explicitProvider, explicit: true }]
-			: resolveProviderCandidates("auto");
+		candidates = [{ id: explicitProvider, explicit: true }];
 	} else if (explicitProvider === "auto") {
 		// Explicit `--provider auto` bypasses the configured preferred provider
 		// for this invocation; exclusions still apply.
@@ -175,7 +172,13 @@ async function executeSearch(
 			const available = candidate.explicit
 				? await provider.isExplicitlyAvailable(authStorage)
 				: await provider.isAvailable(authStorage);
-			if (!available) continue;
+			if (!available && !candidate.explicit) continue;
+			if (!available && candidate.explicit) {
+				throw new SearchProviderError(
+					provider.id,
+					`${provider.label} web search is unavailable. Configure its credentials or select the automatic provider chain.`,
+				);
+			}
 			availableProviderCount++;
 			lastProvider = provider;
 
@@ -213,6 +216,7 @@ async function executeSearch(
 			// summary error), masking the cancellation.
 			throwIfAborted(signal);
 			failures.push({ provider: provider ?? providerMeta, error });
+			if (candidate.explicit) break;
 		}
 	}
 

package/src/web/search/providers/codex.ts

@@ -1,28 +1,40 @@
 /**
  * OpenAI Codex Web Search Provider
  *
- * Uses Codex's built-in web_search tool via the Responses API.
- * Auth is resolved through `AuthStorage.getOAuthAccess("openai-codex")` so the
- * broker is the sole refresh authority — this module never opens a sibling
- * SQLite store, never POSTs the broker sentinel to an OpenAI token endpoint.
+ * Uses the configured Codex Responses transport for proxy/API-key setups and
+ * the official ChatGPT backend for OAuth logins.
  */
 import * as os from "node:os";
-import { type AuthStorage, type FetchImpl, type Model, type OAuthAccess, withOAuthAccess } from "@oh-my-pi/pi-ai";
-import { decodeJwt } from "@oh-my-pi/pi-ai/oauth/openai-codex";
+import {
+	type AuthStorage,
+	type FetchImpl,
+	type Model,
+	type OAuthAccess,
+	withAuth,
+	withOAuthAccess,
+} from "@oh-my-pi/pi-ai";
 import { applyCodexResponsesLiteShape } from "@oh-my-pi/pi-ai/providers/openai-codex/request-transformer";
-import { createOpenAICodexCompatibilityMetadata } from "@oh-my-pi/pi-ai/providers/openai-codex-responses";
+import {
+	createOpenAICodexCompatibilityMetadata,
+	resolveCodexResponsesUrl,
+} from "@oh-my-pi/pi-ai/providers/openai-codex-responses";
 import { getBundledModels } from "@oh-my-pi/pi-catalog/models";
-import { CODEX_CLIENT_VERSION, OPENAI_HEADER_VALUES, OPENAI_HEADERS } from "@oh-my-pi/pi-catalog/wire/codex";
+import {
+	CODEX_BASE_URL,
+	CODEX_CLIENT_VERSION,
+	getCodexAccountId,
+	OPENAI_HEADER_VALUES,
+	OPENAI_HEADERS,
+} from "@oh-my-pi/pi-catalog/wire/codex";
 import { $env, readSseJson } from "@oh-my-pi/pi-utils";
 import packageJson from "../../../../package.json" with { type: "json" };
+import type { ModelRegistry } from "../../../config/model-registry";
 import type { SearchResponse, SearchSource } from "../../../web/search/types";
 import { SearchProviderError } from "../../../web/search/types";
 import type { SearchParams } from "./base";
 import { SearchProvider } from "./base";
 import { classifyProviderHttpError, withHardTimeout } from "./utils";
 
-const CODEX_BASE_URL = "https://chatgpt.com/backend-api";
-const CODEX_RESPONSES_PATH = "/codex/responses";
 const FALLBACK_MODEL = "gpt-5.5";
 const DEFAULT_MODEL_PREFERENCES = [
 	"gpt-5.6-luna",
@@ -37,7 +49,6 @@ const DEFAULT_MODEL_PREFERENCES = [
 	"gpt-5.1-codex",
 	"gpt-5-codex-mini",
 ];
-const JWT_CLAIM_PATH = "https://api.openai.com/auth";
 const DEFAULT_INSTRUCTIONS =
 	"You are a helpful assistant with web search capabilities. Search the web to answer the user's question accurately and cite your sources.";
 
@@ -48,6 +59,21 @@ interface CodexModelCandidate {
 	catalogModel?: CodexSearchModel;
 }
 
+interface CodexSearchTransport {
+	baseUrl: string;
+	url: string;
+	headers: Record<string, string>;
+	customEndpoint: boolean;
+}
+
+interface CodexSearchResult {
+	answer: string;
+	sources: SearchSource[];
+	model: string;
+	requestId: string;
+	usage?: { inputTokens: number; outputTokens: number; totalTokens: number };
+}
+
 function getBundledCodexModels(): CodexSearchModel[] {
 	const models: CodexSearchModel[] = [];
 	for (const model of getBundledModels("openai-codex")) {
@@ -295,18 +321,6 @@ function extractTextSources(text: string): SearchSource[] {
 	return sources;
 }
 
-/**
- * Extracts account ID from a Codex access token.
- * @param accessToken - JWT access token
- * @returns Account ID string, or null if not found
- */
-function getAccountIdFromJwt(accessToken: string): string | null {
-	const payload = decodeJwt(accessToken);
-	const auth = payload?.[JWT_CLAIM_PATH] as { chatgpt_account_id?: string } | undefined;
-	const accountId = auth?.chatgpt_account_id;
-	return typeof accountId === "string" && accountId.length > 0 ? accountId : null;
-}
-
 /**
  * Resolve a Codex bearer + accountId through {@link AuthStorage} — the single
  * refresh authority. Returns `null` when no OAuth credential is configured,
@@ -320,25 +334,55 @@ async function findCodexAuth(
 ): Promise<{ access: OAuthAccess; accountId: string } | null> {
 	const access = await authStorage.getOAuthAccess("openai-codex", sessionId, { signal });
 	if (!access) return null;
-	const accountId = access.accountId ?? getAccountIdFromJwt(access.accessToken);
+	const accountId = access.accountId ?? getCodexAccountId(access.accessToken);
 	if (!accountId) return null;
 	return { access, accountId };
 }
 
+function resolveCodexSearchTransport(modelRegistry: ModelRegistry | undefined, modelId: string): CodexSearchTransport {
+	const registryModel = modelRegistry?.find("openai-codex", modelId);
+	const bundledModel = getBundledCodexModels().find(model => model.id === modelId);
+	const providerBaseUrl = modelRegistry?.getProviderBaseUrl("openai-codex");
+	let baseUrl = providerBaseUrl ?? registryModel?.baseUrl ?? CODEX_BASE_URL;
+	if (registryModel?.baseUrl && registryModel.baseUrl !== (bundledModel?.baseUrl ?? CODEX_BASE_URL)) {
+		baseUrl = registryModel.baseUrl;
+	}
+
+	const url = resolveCodexResponsesUrl(baseUrl);
+	return {
+		baseUrl,
+		url,
+		headers: {
+			...(modelRegistry?.getProviderHeaders("openai-codex") ?? {}),
+			...(registryModel?.headers ?? {}),
+		},
+		customEndpoint: url !== resolveCodexResponsesUrl(CODEX_BASE_URL),
+	};
+}
+
 /**
  * Builds HTTP headers for Codex API requests.
  */
-function buildCodexHeaders(accessToken: string, accountId: string): Record<string, string> {
-	return {
-		Authorization: `Bearer ${accessToken}`,
-		[OPENAI_HEADERS.ACCOUNT_ID]: accountId,
-		[OPENAI_HEADERS.BETA]: OPENAI_HEADER_VALUES.BETA_RESPONSES,
-		[OPENAI_HEADERS.ORIGINATOR]: OPENAI_HEADER_VALUES.ORIGINATOR_CODEX,
-		[OPENAI_HEADERS.VERSION]: CODEX_CLIENT_VERSION,
-		"User-Agent": `pi/${packageJson.version} (${os.platform()} ${os.release()}; ${os.arch()})`,
-		Accept: "text/event-stream",
-		"Content-Type": "application/json",
-	};
+function buildCodexHeaders(
+	accessToken: string,
+	accountId: string | undefined,
+	configuredHeaders: Record<string, string>,
+): Headers {
+	const headers = new Headers(configuredHeaders);
+	headers.delete("x-api-key");
+	headers.set("Authorization", `Bearer ${accessToken}`);
+	if (accountId) {
+		headers.set(OPENAI_HEADERS.ACCOUNT_ID, accountId);
+	} else {
+		headers.delete(OPENAI_HEADERS.ACCOUNT_ID);
+	}
+	headers.set(OPENAI_HEADERS.BETA, OPENAI_HEADER_VALUES.BETA_RESPONSES);
+	headers.set(OPENAI_HEADERS.ORIGINATOR, OPENAI_HEADER_VALUES.ORIGINATOR_CODEX);
+	headers.set(OPENAI_HEADERS.VERSION, CODEX_CLIENT_VERSION);
+	headers.set("User-Agent", `pi/${packageJson.version} (${os.platform()} ${os.release()}; ${os.arch()})`);
+	headers.set("Accept", "text/event-stream");
+	headers.set("Content-Type", "application/json");
+	return headers;
 }
 
 /**
@@ -348,7 +392,7 @@ function buildCodexHeaders(accessToken: string, accountId: string): Record<strin
  * overrides from the default ChatGPT-account model-selection path.
  */
 async function callCodexSearch(
-	auth: { accessToken: string; accountId: string },
+	auth: { accessToken: string; accountId?: string },
 	query: string,
 	options: {
 		signal?: AbortSignal;
@@ -357,16 +401,10 @@ async function callCodexSearch(
 		model: CodexModelCandidate;
 		sessionId?: string;
 		fetch?: FetchImpl;
+		transport: CodexSearchTransport;
 	},
-): Promise<{
-	answer: string;
-	sources: SearchSource[];
-	model: string;
-	requestId: string;
-	usage?: { inputTokens: number; outputTokens: number; totalTokens: number };
-}> {
-	const url = `${CODEX_BASE_URL}${CODEX_RESPONSES_PATH}`;
-	const headers = buildCodexHeaders(auth.accessToken, auth.accountId);
+): Promise<CodexSearchResult> {
+	const headers = buildCodexHeaders(auth.accessToken, auth.accountId, options.transport.headers);
 
 	const requestedModel = options.model.modelId;
 	const usesResponsesLite = options.model.catalogModel?.useResponsesLite === true;
@@ -397,15 +435,18 @@ async function callCodexSearch(
 			requestKind: "turn",
 			startNewTurn: true,
 		});
-		Object.assign(headers, metadata.headers);
-		headers[OPENAI_HEADERS.RESPONSES_LITE] = "true";
+		for (const name in metadata.headers) {
+			const value = metadata.headers[name];
+			if (value !== undefined) headers.set(name, value);
+		}
+		headers.set(OPENAI_HEADERS.RESPONSES_LITE, "true");
 		body.client_metadata = metadata.clientMetadata;
 		body.reasoning = { context: "all_turns" };
 		applyCodexResponsesLiteShape(body);
 	}
 
 	const fetchImpl = options.fetch ?? fetch;
-	const response = await fetchImpl(url, {
+	const response = await fetchImpl(options.transport.url, {
 		method: "POST",
 		headers,
 		body: JSON.stringify(body),
@@ -528,6 +569,39 @@ async function callCodexSearch(
 	};
 }
 
+async function runCodexSearchCandidates(options: {
+	auth: { accessToken: string; accountId?: string };
+	params: SearchParams;
+	modelCandidates: CodexModelCandidate[];
+	modelWasConfigured: boolean;
+	transport: CodexSearchTransport;
+}): Promise<CodexSearchResult> {
+	let lastError: unknown;
+	for (let index = 0; index < options.modelCandidates.length; index += 1) {
+		const candidate = options.modelCandidates[index];
+		if (!candidate) continue;
+
+		try {
+			return await callCodexSearch(options.auth, options.params.query, {
+				signal: options.params.signal,
+				systemPrompt: options.params.systemPrompt,
+				searchContextSize: "high",
+				model: candidate,
+				sessionId: options.params.sessionId,
+				fetch: options.params.fetch,
+				transport: options.transport,
+			});
+		} catch (error) {
+			lastError = error;
+			const isLastCandidate = index === options.modelCandidates.length - 1;
+			if (options.modelWasConfigured || isLastCandidate || !shouldRetryWithNextDefaultModel(error)) {
+				throw error;
+			}
+		}
+	}
+	throw lastError ?? new Error("Codex search failed without returning a result");
+}
+
 /**
  * Executes a web search using OpenAI Codex's built-in web search tool.
  *
@@ -541,55 +615,81 @@ async function callCodexSearch(
  *   rejects.
  */
 export async function searchCodex(params: SearchParams): Promise<SearchResponse> {
-	const seed = await findCodexAuth(params.authStorage, params.sessionId, params.signal);
-	if (!seed) {
-		throw new Error(
-			"No Codex OAuth credentials found. Login with 'omp /login openai-codex' to enable Codex web search.",
-		);
-	}
-
 	const configuredModel = getConfiguredModel();
 	const modelCandidates = configuredModel ? [configuredModel] : getDefaultModelCandidates();
+	const firstCandidate = modelCandidates[0];
+	if (!firstCandidate) {
+		throw new SearchProviderError("codex", "No Codex web search model is configured.");
+	}
+	const transport = resolveCodexSearchTransport(params.modelRegistry, firstCandidate.modelId);
+
+	let result: CodexSearchResult;
+	if (transport.customEndpoint) {
+		// ModelRegistry resolves command-backed provider keys before consulting
+		// its AuthStorage, so a lower-priority OAuth origin is irrelevant when
+		// that command source is configured.
+		const credentialSource = params.modelRegistry?.authStorage ?? params.authStorage;
+		const credentialOrigin = credentialSource.getCredentialOrigin("openai-codex");
+		const hasCommandBackedKey = params.modelRegistry?.hasCommandBackedApiKey("openai-codex") === true;
+		if (!hasCommandBackedKey && (credentialOrigin?.kind === "oauth" || credentialOrigin?.kind === "env")) {
+			throw new SearchProviderError(
+				"codex",
+				`Refusing to send official Codex OAuth credentials to custom endpoint ${transport.baseUrl}. Configure an API key for provider "openai-codex".`,
+			);
+		}
 
-	const result = await withOAuthAccess(
-		params.authStorage,
-		"openai-codex",
-		async access => {
-			// Derive ALL auth material from the access this attempt received —
-			// a refreshed/rotated credential carries a different bearer and
-			// ChatGPT account id than the seed.
-			const accountId = access.accountId ?? getAccountIdFromJwt(access.accessToken);
-			if (!accountId) {
-				throw new Error("Codex OAuth credential is missing a ChatGPT account id");
-			}
-			const auth = { accessToken: access.accessToken, accountId };
-
-			let lastError: unknown;
-			for (let index = 0; index < modelCandidates.length; index += 1) {
-				const candidate = modelCandidates[index];
-				if (!candidate) continue;
-
-				try {
-					return await callCodexSearch(auth, params.query, {
-						signal: params.signal,
-						systemPrompt: params.systemPrompt,
-						searchContextSize: "high",
-						model: candidate,
-						sessionId: params.sessionId,
-						fetch: params.fetch,
-					});
-				} catch (error) {
-					lastError = error;
-					const isLastCandidate = index === modelCandidates.length - 1;
-					if (configuredModel || isLastCandidate || !shouldRetryWithNextDefaultModel(error)) {
-						throw error;
-					}
+		const resolverOptions = {
+			sessionId: params.sessionId,
+			baseUrl: transport.baseUrl,
+			modelId: firstCandidate.modelId,
+		};
+		const keyOrResolver = params.modelRegistry
+			? params.modelRegistry.resolver("openai-codex", resolverOptions)
+			: params.authStorage.resolver("openai-codex", resolverOptions);
+		result = await withAuth(
+			keyOrResolver,
+			accessToken =>
+				runCodexSearchCandidates({
+					auth: { accessToken },
+					params,
+					modelCandidates,
+					modelWasConfigured: configuredModel !== undefined,
+					transport,
+				}),
+			{
+				signal: params.signal,
+				missingKeyMessage: 'Codex credentials not found. Configure an API key for provider "openai-codex".',
+			},
+		);
+	} else {
+		const seed = await findCodexAuth(params.authStorage, params.sessionId, params.signal);
+		if (!seed) {
+			throw new Error(
+				"No Codex OAuth credentials found. Login with 'omp /login openai-codex' to enable Codex web search.",
+			);
+		}
+
+		result = await withOAuthAccess(
+			params.authStorage,
+			"openai-codex",
+			access => {
+				// A refreshed/rotated credential can carry a different bearer and
+				// ChatGPT account id than the seed used to select the first attempt.
+				const accountId = access.accountId ?? getCodexAccountId(access.accessToken);
+				if (!accountId) {
+					throw new Error("Codex OAuth credential is missing a ChatGPT account id");
 				}
-			}
-			throw lastError ?? new Error("Codex search failed without returning a result");
-		},
-		{ sessionId: params.sessionId, signal: params.signal, seed: seed.access },
-	);
+				return runCodexSearchCandidates({
+					auth: { accessToken: access.accessToken, accountId },
+					params,
+					modelCandidates,
+					modelWasConfigured: configuredModel !== undefined,
+					transport,
+				});
+			},
+			{ sessionId: params.sessionId, signal: params.signal, seed: seed.access },
+		);
+	}
 
 	let sources = result.sources;
 
@@ -615,14 +715,10 @@ export async function searchCodex(params: SearchParams): Promise<SearchResponse>
 }
 
 /**
- * Checks if Codex web search is available.
+ * Checks whether Codex web search has an API key or OAuth credential.
  */
 export async function hasCodexSearch(authStorage: AuthStorage): Promise<boolean> {
-	// `isAvailable` runs before every request — keep the probe cheap.
-	// `hasOAuth(...)` is a synchronous in-memory check that returns true as soon
-	// as a Codex OAuth credential is loaded, without driving the refresh
-	// pipeline. The actual refresh happens lazily in `searchCodex`.
-	return authStorage.hasOAuth("openai-codex");
+	return authStorage.hasAuth("openai-codex");
 }
 
 /** Search provider for OpenAI Codex web search. */