@shanepadgett/tau-agent 0.18.1 → 0.20.0

package/extensions/cache-diagnostics/README.md

@@ -0,0 +1,11 @@
+# Cache Diagnostics
+
+Cache Diagnostics records compact fingerprints of Tau's final provider requests. It helps separate local prompt changes from unexplained provider cache misses without storing prompts, source code, or credentials.
+
+The extension runs automatically. Logs are written under:
+
+```text
+~/.pi/agent/cache-diagnostics/
+```
+
+Run `/cache-debug` after suspicious misses. Tau writes a bounded report under `~/.pi/agent/cache-diagnostics/reports/` for later investigation. Logs and reports older than 30 days are removed when a session starts.

package/extensions/cache-diagnostics/index.ts

@@ -0,0 +1,655 @@
+import { createHash, randomUUID } from "node:crypto";
+import { appendFile, type FileHandle, mkdir, open, readdir, rename, stat, unlink, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+import { getAgentDir, type BuildSystemPromptOptions, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+const RETENTION_MS = 30 * 24 * 60 * 60 * 1000;
+const CACHE_NOISE_FLOOR = 1_024;
+const MAX_MEMORY_RECORDS = 500;
+const MAX_REPORT_INCIDENTS = 20;
+const RECENT_REPORT_REQUESTS = 30;
+const MAX_LOG_TAIL_BYTES = 10 * 1_024 * 1_024;
+
+interface HashFingerprint {
+	hash: string;
+	bytes: number;
+}
+
+interface ItemFingerprint extends HashFingerprint {
+	type: string | null;
+	role: string | null;
+	name: string | null;
+}
+
+interface ToolFingerprint extends HashFingerprint {
+	name: string | null;
+}
+
+interface FieldFingerprint extends HashFingerprint {
+	name: string;
+}
+
+interface PromptState {
+	hash: string;
+	systemPromptHash: string;
+	selectedTools: string[];
+	toolSnippetsHash: string | null;
+	promptGuidelinesHash: string | null;
+	appendSystemPromptHash: string | null;
+	customPromptHash: string | null;
+	contextFiles: Array<{ path: string; hash: string }>;
+	skills: Array<{ name: string | null; path: string | null }>;
+}
+
+interface PayloadFingerprint {
+	model: string | null;
+	promptCacheKeyHash: string | null;
+	instructionsHash: string | null;
+	toolsHash: string | null;
+	nonSequenceHash: string;
+	sequenceField: "input" | "messages" | "none";
+	fields: FieldFingerprint[];
+	tools: ToolFingerprint[];
+	items: ItemFingerprint[];
+}
+
+interface RequestRecord {
+	kind: "request";
+	version: 2;
+	timestamp: string;
+	id: string;
+	turnIndex: number | null;
+	previousRequestId: string | null;
+	model: string | null;
+	promptCacheKeyHash: string | null;
+	instructionsHash: string | null;
+	toolsHash: string | null;
+	nonSequenceHash: string;
+	sequenceField: "input" | "messages" | "none";
+	fields: FieldFingerprint[];
+	tools: ToolFingerprint[];
+	items: ItemFingerprint[];
+	previousItems: number;
+	commonPrefixItems: number;
+	stableEnvelope: boolean;
+	previousExactPrefix: boolean;
+	previousPromptTokens: number;
+	changes: {
+		envelopeFields: string[];
+		toolsAdded: string[];
+		toolsRemoved: string[];
+		toolsChanged: string[];
+		firstChangedItem: number | null;
+		promptStateChanged: boolean;
+	};
+	promptState: PromptState | null;
+}
+
+interface ResponseRecord {
+	kind: "response";
+	version: 2;
+	timestamp: string;
+	id: string;
+	status: number;
+	headers: Record<string, string>;
+}
+
+interface ResultRecord {
+	kind: "result";
+	version: 2;
+	timestamp: string;
+	id: string;
+	provider: string;
+	model: string;
+	stopReason: string;
+	usage: {
+		input: number;
+		cacheRead: number;
+		cacheWrite: number;
+		promptTokens: number;
+	};
+	previousExactPrefix: boolean;
+	missedTokens: number;
+	cacheMiss: boolean;
+	baselinePromoted: boolean;
+}
+
+interface MarkerRecord {
+	kind: "marker";
+	version: 2;
+	timestamp: string;
+	name: string;
+	details: Record<string, string | number | boolean | null>;
+}
+
+interface Attempt {
+	id: string;
+	fingerprint: PayloadFingerprint;
+	request: RequestRecord;
+	responseStatus: number | undefined;
+}
+
+export default function cacheDiagnosticsExtension(pi: ExtensionAPI): void {
+	const directory = join(getAgentDir(), "cache-diagnostics");
+	const reportsDirectory = join(directory, "reports");
+	const runtimeId = randomUUID();
+	let logFile = "";
+	let cwd = "";
+	let sessionId = "";
+	let sessionFile: string | null = null;
+	let requestSequence = 0;
+	let turnIndex: number | null = null;
+	let latestPromptState: PromptState | null = null;
+	let previousPromptState: PromptState | null = null;
+	let previousFingerprint: PayloadFingerprint | undefined;
+	let previousRequestId: string | null = null;
+	let previousPromptTokens = 0;
+	let cacheActivitySeen = false;
+	let awaitingResponses: Attempt[] = [];
+	let turnAttempts: Attempt[] = [];
+	let requests: RequestRecord[] = [];
+	let responses: ResponseRecord[] = [];
+	let results: ResultRecord[] = [];
+	let markers: MarkerRecord[] = [];
+
+	const appendRecord = async (record: object): Promise<void> => {
+		const path = logFile;
+		if (!path) return;
+		await appendFile(path, `${JSON.stringify(record)}\n`, "utf8");
+	};
+	const addMarker = async (
+		name: string,
+		details: Record<string, string | number | boolean | null> = {},
+	): Promise<void> => {
+		const marker: MarkerRecord = {
+			kind: "marker",
+			version: 2,
+			timestamp: new Date().toISOString(),
+			name,
+			details,
+		};
+		markers = keepRecent([...markers, marker]);
+		await appendRecord(marker);
+	};
+	const resetComparison = () => {
+		previousFingerprint = undefined;
+		previousRequestId = null;
+		previousPromptState = null;
+		previousPromptTokens = 0;
+		cacheActivitySeen = false;
+	};
+
+	pi.registerCommand("cache-debug", {
+		description: "Write a bounded prompt-cache diagnostic report for this session",
+		async handler(_args, ctx) {
+			await mkdir(reportsDirectory, { recursive: true });
+			const incidentResults = results.filter((result) => result.cacheMiss).slice(-MAX_REPORT_INCIDENTS);
+			const selectedIds = new Set<string>();
+			for (const result of incidentResults) {
+				selectedIds.add(result.id);
+				const request = requests.find((candidate) => candidate.id === result.id);
+				if (request?.previousRequestId) selectedIds.add(request.previousRequestId);
+			}
+			for (const request of requests.slice(-RECENT_REPORT_REQUESTS)) selectedIds.add(request.id);
+			const selectedRequests = requests.filter((request) => selectedIds.has(request.id));
+			const selectedResponses = responses.filter((response) => selectedIds.has(response.id));
+			const selectedResults = results.filter((result) => selectedIds.has(result.id));
+			const createdAt = new Date().toISOString();
+			const safeSessionId = (sessionId || "ephemeral").replaceAll(/[^a-zA-Z0-9_-]/g, "_");
+			const reportFile = join(reportsDirectory, `${safeSessionId}-${createdAt.replaceAll(/[:.]/g, "-")}.json`);
+			const temporaryFile = `${reportFile}.${randomUUID()}.tmp`;
+			const report = {
+				kind: "tau-cache-debug",
+				version: 1,
+				createdAt,
+				cwd: cwd || ctx.cwd,
+				runtimeId,
+				sessionId: sessionId || null,
+				sessionFile,
+				sourceLog: logFile || null,
+				summary: {
+					requestsObserved: requests.length,
+					resultsObserved: results.length,
+					cacheMissesObserved: results.filter((result) => result.cacheMiss).length,
+					requestsIncluded: selectedRequests.length,
+					models: [...new Set(results.map((result) => `${result.provider}/${result.model}`))],
+				},
+				selection: {
+					maxIncidents: MAX_REPORT_INCIDENTS,
+					recentRequests: RECENT_REPORT_REQUESTS,
+				},
+				markers,
+				requests: selectedRequests,
+				responses: selectedResponses,
+				results: selectedResults,
+			};
+			try {
+				await writeFile(temporaryFile, `${JSON.stringify(report, null, "\t")}\n`, { encoding: "utf8", flag: "wx" });
+				await rename(temporaryFile, reportFile);
+			} catch (error) {
+				await unlink(temporaryFile).catch(() => undefined);
+				throw error;
+			}
+			ctx.ui.notify(`Cache debug written: ${reportFile}`, "info");
+		},
+	});
+
+	pi.on("session_start", async (_event, ctx) => {
+		cwd = ctx.cwd;
+		sessionId = ctx.sessionManager.getSessionId();
+		sessionFile = ctx.sessionManager.getSessionFile() ?? null;
+		requestSequence = 0;
+		turnIndex = null;
+		latestPromptState = null;
+		awaitingResponses = [];
+		turnAttempts = [];
+		resetComparison();
+		await mkdir(reportsDirectory, { recursive: true });
+		logFile = join(directory, `${sessionId.replaceAll(/[^a-zA-Z0-9_-]/g, "_")}.jsonl`);
+		const cutoff = Date.now() - RETENTION_MS;
+		for (const parent of [directory, reportsDirectory]) {
+			for (const entry of await readdir(parent, { withFileTypes: true })) {
+				if (
+					!entry.isFile() ||
+					(!entry.name.endsWith(".jsonl") && !entry.name.endsWith(".json") && !entry.name.endsWith(".tmp"))
+				)
+					continue;
+				const path = join(parent, entry.name);
+				if ((await stat(path)).mtimeMs < cutoff) await unlink(path);
+			}
+		}
+		const persisted = await readRecentLogRecords(logFile);
+		requests = keepRecent(
+			persisted
+				.filter((record) => record.kind === "request" && record.version === 2)
+				.map((record) => record as unknown as RequestRecord),
+		);
+		responses = keepRecent(
+			persisted
+				.filter((record) => record.kind === "response" && record.version === 2)
+				.map((record) => record as unknown as ResponseRecord),
+		);
+		results = keepRecent(
+			persisted
+				.filter((record) => record.kind === "result" && record.version === 2)
+				.map((record) => record as unknown as ResultRecord),
+		);
+		markers = keepRecent(
+			persisted
+				.filter((record) => record.kind === "marker" && record.version === 2)
+				.map((record) => record as unknown as MarkerRecord),
+		);
+		let latestModelSelectIndex = -1;
+		for (let index = persisted.length - 1; index >= 0; index -= 1) {
+			const record = persisted[index];
+			if (record?.kind !== "marker" || record.version !== 2 || record.name !== "model-select") continue;
+			latestModelSelectIndex = index;
+			break;
+		}
+		const comparisonRecords = persisted.slice(latestModelSelectIndex + 1);
+		const comparisonResults = comparisonRecords
+			.filter((record) => record.kind === "result" && record.version === 2)
+			.map((record) => record as unknown as ResultRecord);
+		const latestPromotedResult = [...comparisonResults].reverse().find((result) => result.baselinePromoted);
+		const latestRequest = latestPromotedResult
+			? requests.find((request) => request.id === latestPromotedResult.id)
+			: undefined;
+		if (latestPromotedResult && latestRequest) {
+			previousFingerprint = payloadFingerprintFromRequest(latestRequest);
+			previousRequestId = latestRequest.id;
+			previousPromptState = latestRequest.promptState;
+			previousPromptTokens = latestPromotedResult.usage.promptTokens;
+			cacheActivitySeen = comparisonResults.some(
+				(result) => result.baselinePromoted && result.usage.cacheRead + result.usage.cacheWrite > 0,
+			);
+		}
+		await appendRecord({
+			kind: "runtime",
+			version: 2,
+			timestamp: new Date().toISOString(),
+			runtimeId,
+			cwd,
+			sessionId,
+			sessionFile,
+		});
+	});
+
+	pi.on("before_agent_start", (event) => {
+		latestPromptState = fingerprintPromptState(event.systemPrompt, event.systemPromptOptions);
+	});
+
+	pi.on("turn_start", (event) => {
+		turnIndex = event.turnIndex;
+		turnAttempts = [];
+	});
+
+	pi.on("before_provider_request", async (event) => {
+		const unresolvedIds = new Set(
+			turnAttempts.filter((attempt) => attempt.responseStatus === undefined).map((attempt) => attempt.id),
+		);
+		if (unresolvedIds.size > 0) {
+			awaitingResponses = awaitingResponses.filter((attempt) => !unresolvedIds.has(attempt.id));
+		}
+		requestSequence += 1;
+		const id = `${runtimeId}:${requestSequence}`;
+		const fingerprint = fingerprintPayload(event.payload);
+		let commonPrefixItems = 0;
+		if (previousFingerprint) {
+			const limit = Math.min(previousFingerprint.items.length, fingerprint.items.length);
+			while (
+				commonPrefixItems < limit &&
+				previousFingerprint.items[commonPrefixItems]?.hash === fingerprint.items[commonPrefixItems]?.hash
+			) {
+				commonPrefixItems += 1;
+			}
+		}
+		const stableEnvelope = previousFingerprint?.nonSequenceHash === fingerprint.nonSequenceHash;
+		const previousExactPrefix =
+			previousFingerprint !== undefined &&
+			stableEnvelope &&
+			previousFingerprint.sequenceField === fingerprint.sequenceField &&
+			commonPrefixItems === previousFingerprint.items.length;
+		const request: RequestRecord = {
+			kind: "request",
+			version: 2,
+			timestamp: new Date().toISOString(),
+			id,
+			turnIndex,
+			previousRequestId,
+			model: fingerprint.model,
+			promptCacheKeyHash: fingerprint.promptCacheKeyHash,
+			instructionsHash: fingerprint.instructionsHash,
+			toolsHash: fingerprint.toolsHash,
+			nonSequenceHash: fingerprint.nonSequenceHash,
+			sequenceField: fingerprint.sequenceField,
+			fields: fingerprint.fields,
+			tools: fingerprint.tools,
+			items: fingerprint.items,
+			previousItems: previousFingerprint?.items.length ?? 0,
+			commonPrefixItems,
+			stableEnvelope,
+			previousExactPrefix,
+			previousPromptTokens,
+			changes: compareFingerprints(previousFingerprint, fingerprint, previousPromptState, latestPromptState),
+			promptState: latestPromptState,
+		};
+		const attempt: Attempt = { id, fingerprint, request, responseStatus: undefined };
+		awaitingResponses.push(attempt);
+		turnAttempts.push(attempt);
+		requests = keepRecent([...requests, request]);
+		await appendRecord(request);
+	});
+
+	pi.on("after_provider_response", async (event) => {
+		const attempt = awaitingResponses.shift();
+		if (!attempt) return;
+		attempt.responseStatus = event.status;
+		const selectedHeaders = Object.fromEntries(
+			Object.entries(event.headers).filter(([name]) => {
+				const normalized = name.toLowerCase();
+				return normalized.includes("request-id") || normalized === "cf-ray";
+			}),
+		);
+		const response: ResponseRecord = {
+			kind: "response",
+			version: 2,
+			timestamp: new Date().toISOString(),
+			id: attempt.id,
+			status: event.status,
+			headers: selectedHeaders,
+		};
+		responses = keepRecent([...responses, response]);
+		await appendRecord(response);
+	});
+
+	pi.on("message_end", async (event) => {
+		if (event.message.role !== "assistant") return;
+		const attempt = turnAttempts.at(-1);
+		const completedAttemptIds = new Set(turnAttempts.map((candidate) => candidate.id));
+		awaitingResponses = awaitingResponses.filter((candidate) => !completedAttemptIds.has(candidate.id));
+		turnAttempts = [];
+		if (!attempt) return;
+		const usage = event.message.usage;
+		const promptTokens = usage.input + usage.cacheRead + usage.cacheWrite;
+		const baselinePromoted =
+			event.message.stopReason !== "error" &&
+			event.message.stopReason !== "aborted" &&
+			(attempt.responseStatus === undefined || (attempt.responseStatus >= 200 && attempt.responseStatus < 300));
+		const missedTokens =
+			baselinePromoted && attempt.request.previousExactPrefix
+				? Math.max(0, Math.min(attempt.request.previousPromptTokens, promptTokens) - usage.cacheRead)
+				: 0;
+		const cacheMiss = baselinePromoted && cacheActivitySeen && missedTokens > CACHE_NOISE_FLOOR;
+		const result: ResultRecord = {
+			kind: "result",
+			version: 2,
+			timestamp: new Date().toISOString(),
+			id: attempt.id,
+			provider: event.message.provider,
+			model: event.message.model,
+			stopReason: event.message.stopReason,
+			usage: {
+				input: usage.input,
+				cacheRead: usage.cacheRead,
+				cacheWrite: usage.cacheWrite,
+				promptTokens,
+			},
+			previousExactPrefix: attempt.request.previousExactPrefix,
+			missedTokens,
+			cacheMiss,
+			baselinePromoted,
+		};
+		results = keepRecent([...results, result]);
+		if (baselinePromoted) {
+			previousFingerprint = attempt.fingerprint;
+			previousRequestId = attempt.id;
+			previousPromptState = attempt.request.promptState;
+			previousPromptTokens = promptTokens;
+			cacheActivitySeen ||= usage.cacheRead + usage.cacheWrite > 0;
+		}
+		await appendRecord(result);
+	});
+
+	pi.on("model_select", async (event) => {
+		resetComparison();
+		await addMarker("model-select", {
+			provider: event.model.provider,
+			model: event.model.id,
+			previousProvider: event.previousModel?.provider ?? null,
+			previousModel: event.previousModel?.id ?? null,
+			source: event.source,
+		});
+	});
+	pi.on("thinking_level_select", (event) =>
+		addMarker("thinking-level-select", { level: event.level, previousLevel: event.previousLevel }),
+	);
+	pi.on("session_compact", (event) =>
+		addMarker("session-compact", {
+			reason: event.reason,
+			fromExtension: event.fromExtension,
+			willRetry: event.willRetry,
+		}),
+	);
+	pi.on("session_tree", () => addMarker("session-tree"));
+	pi.on("tool_execution_end", (event) => {
+		if (event.toolName !== "context_prune" && event.toolName !== "load_tools") return;
+		return addMarker("cache-affecting-tool", { tool: event.toolName, isError: event.isError });
+	});
+}
+
+async function readRecentLogRecords(path: string): Promise<Array<Record<string, unknown>>> {
+	let handle: FileHandle | undefined;
+	try {
+		handle = await open(path, "r");
+		const size = (await handle.stat()).size;
+		const start = Math.max(0, size - MAX_LOG_TAIL_BYTES);
+		const buffer = Buffer.alloc(size - start);
+		await handle.read(buffer, 0, buffer.length, start);
+		let text = buffer.toString("utf8");
+		if (start > 0) {
+			const firstNewline = text.indexOf("\n");
+			text = firstNewline < 0 ? "" : text.slice(firstNewline + 1);
+		}
+		const records: Array<Record<string, unknown>> = [];
+		for (const line of text.split("\n")) {
+			if (!line) continue;
+			try {
+				const value: unknown = JSON.parse(line);
+				if (isRecord(value)) records.push(value);
+			} catch {
+				// Ignore a partial final record left by an interrupted append.
+			}
+		}
+		return records;
+	} catch (error) {
+		if (isRecord(error) && error.code === "ENOENT") return [];
+		throw error;
+	} finally {
+		await handle?.close();
+	}
+}
+
+function payloadFingerprintFromRequest(request: RequestRecord): PayloadFingerprint {
+	return {
+		model: request.model,
+		promptCacheKeyHash: request.promptCacheKeyHash,
+		instructionsHash: request.instructionsHash,
+		toolsHash: request.toolsHash,
+		nonSequenceHash: request.nonSequenceHash,
+		sequenceField: request.sequenceField,
+		fields: request.fields,
+		tools: request.tools,
+		items: request.items,
+	};
+}
+
+function fingerprintPromptState(systemPrompt: string, options: BuildSystemPromptOptions): PromptState {
+	const state = {
+		systemPromptHash: hashJson(systemPrompt).hash,
+		selectedTools: options.selectedTools ?? [],
+		toolSnippetsHash: valueHash(options.toolSnippets),
+		promptGuidelinesHash: valueHash(options.promptGuidelines),
+		appendSystemPromptHash: valueHash(options.appendSystemPrompt),
+		customPromptHash: valueHash(options.customPrompt),
+		contextFiles: (options.contextFiles ?? []).map((file) => ({
+			path: file.path,
+			hash: hashJson(file.content).hash,
+		})),
+		skills: (options.skills ?? []).map((skill) => {
+			const record = skill as unknown as Record<string, unknown>;
+			return {
+				name: typeof record.name === "string" ? record.name : null,
+				path: typeof record.path === "string" ? record.path : null,
+			};
+		}),
+	};
+	return { hash: hashJson(state).hash, ...state };
+}
+
+function fingerprintPayload(payload: unknown): PayloadFingerprint {
+	const record = isRecord(payload) ? payload : {};
+	const sequenceField = Array.isArray(record.input) ? "input" : Array.isArray(record.messages) ? "messages" : "none";
+	const sequence = sequenceField === "none" ? [] : (record[sequenceField] as unknown[]);
+	const nonSequence = Object.fromEntries(Object.entries(record).filter(([key]) => key !== sequenceField));
+	const tools = Array.isArray(record.tools)
+		? record.tools.map((tool) => {
+				const fingerprint = hashJson(tool);
+				return { ...fingerprint, name: toolName(tool) };
+			})
+		: [];
+	return {
+		model: typeof record.model === "string" ? record.model : null,
+		promptCacheKeyHash: valueHash(record.prompt_cache_key),
+		instructionsHash: valueHash(record.instructions ?? record.system),
+		toolsHash: valueHash(record.tools),
+		nonSequenceHash: hashJson(nonSequence).hash,
+		sequenceField,
+		fields: Object.entries(nonSequence)
+			.map(([name, value]) => ({ name, ...hashJson(value) }))
+			.sort((left, right) => left.name.localeCompare(right.name)),
+		tools,
+		items: sequence.map((item) => {
+			const fingerprint = hashJson(item);
+			const itemRecord = isRecord(item) ? item : {};
+			return {
+				...fingerprint,
+				type: typeof itemRecord.type === "string" ? itemRecord.type : null,
+				role: typeof itemRecord.role === "string" ? itemRecord.role : null,
+				name: typeof itemRecord.name === "string" ? itemRecord.name : null,
+			};
+		}),
+	};
+}
+
+function compareFingerprints(
+	previous: PayloadFingerprint | undefined,
+	current: PayloadFingerprint,
+	previousPromptState: PromptState | null,
+	currentPromptState: PromptState | null,
+): RequestRecord["changes"] {
+	const previousFields = new Map(previous?.fields.map((field) => [field.name, field.hash]) ?? []);
+	const currentFields = new Map(current.fields.map((field) => [field.name, field.hash]));
+	const envelopeFields = [...new Set([...previousFields.keys(), ...currentFields.keys()])]
+		.filter((name) => previousFields.get(name) !== currentFields.get(name))
+		.sort();
+	const previousTools = namedToolHashes(previous?.tools ?? []);
+	const currentTools = namedToolHashes(current.tools);
+	const toolsAdded = [...currentTools.keys()].filter((name) => !previousTools.has(name));
+	const toolsRemoved = [...previousTools.keys()].filter((name) => !currentTools.has(name));
+	const toolsChanged = [...currentTools.keys()].filter(
+		(name) => previousTools.has(name) && previousTools.get(name) !== currentTools.get(name),
+	);
+	let firstChangedItem: number | null = null;
+	if (previous) {
+		const limit = Math.max(previous.items.length, current.items.length);
+		for (let index = 0; index < limit; index += 1) {
+			if (previous.items[index]?.hash === current.items[index]?.hash) continue;
+			firstChangedItem = index;
+			break;
+		}
+	}
+	return {
+		envelopeFields,
+		toolsAdded,
+		toolsRemoved,
+		toolsChanged,
+		firstChangedItem,
+		promptStateChanged:
+			previousPromptState !== null &&
+			currentPromptState !== null &&
+			previousPromptState.hash !== currentPromptState.hash,
+	};
+}
+
+function namedToolHashes(tools: ToolFingerprint[]): Map<string, string> {
+	const named = new Map<string, string>();
+	for (const [index, tool] of tools.entries()) named.set(tool.name ?? `<unnamed:${index}>`, tool.hash);
+	return named;
+}
+
+function toolName(value: unknown): string | null {
+	if (!isRecord(value)) return null;
+	if (typeof value.name === "string") return value.name;
+	if (isRecord(value.function) && typeof value.function.name === "string") return value.function.name;
+	return null;
+}
+
+function valueHash(value: unknown): string | null {
+	return value === undefined ? null : hashJson(value).hash;
+}
+
+function hashJson(value: unknown): HashFingerprint {
+	const serialized = JSON.stringify(value) ?? "undefined";
+	return {
+		hash: createHash("sha256").update(serialized).digest("hex"),
+		bytes: Buffer.byteLength(serialized, "utf8"),
+	};
+}
+
+function keepRecent<T>(items: T[]): T[] {
+	return items.length <= MAX_MEMORY_RECORDS ? items : items.slice(-MAX_MEMORY_RECORDS);
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null;
+}

package/extensions/context-pruning/README.md

@@ -1,22 +1,135 @@
 # Context Pruning
 
-Context Pruning removes stale tool evidence from the context sent to the model while leaving the session transcript intact. The agent can create a prune anchor after broad exploration has converged or enough irrelevant evidence has accumulated.
+Context Pruning removes old tool calls and their results from the messages sent to the model. It does not delete them from the saved conversation. Branch switching, transcript review, and normal Pi compaction still see the original records.
 
-Before calling `context_prune`, the agent records durable conclusions and its next action in visible prose. It can retain exact tool exchanges, retain current knowledge of selected files, and defer files until a stated condition applies. The tool prepares any required file snapshots before applying the anchor, skips attempts with low estimated savings, and continues the current agent run automatically.
+The feature can also remove old full-file autoread messages and old private reasoning. User messages, visible assistant prose, branch summaries, compaction summaries, shell records, and unrelated custom messages remain.
 
-After tool-using turns, Tau occasionally adds a compact context marker. At lower usage it is informational. Under pressure it asks the agent to finish a coherent step and prune only when the savings are worthwhile. The full steering text stays hidden, and the agent is told not to discuss context management.
+## The important distinction
 
-Run `/prune` with no arguments to ask the agent to create an anchor and continue its unfinished work. The same validation and minimum-savings rules still apply, so a low-value attempt is skipped.
+Context usage and pruning are separate decisions.
 
-Prune anchors, automatic marker boundaries, and warning-colored Tau tool rows follow the active session branch. Switching branches restores that branch's state. Pi's manual, threshold, and overflow compaction continue to work normally.
+- The context percentage decides when Tau gives the agent a pruning hint.
+- A `context_prune` call decides what old evidence the agent wants to remove or retain.
+- Tau then checks whether that exact removal is safe and whether it saves enough tokens.
 
-Manual UI check: Tau-owned tool rows and autoreads should turn warning-colored when pruned. Native Pi rows and fallback rows that don't use Tau's row-state renderer won't change color; pruning and read-cache behavior still apply to them.
+Being at 70% context does not force Tau to apply a prune. With the default settings, 70% is high enough for a pressure hint, but every prune still has to pass the checks below.
+
+## What happens when the agent calls `context_prune`
+
+The agent supplies three required lists:
+
+- `keepFiles`: files whose complete current contents must remain available.
+- `keepToolCalls`: exact earlier tool calls and results that must remain available.
+- `deferFiles`: files that are currently irrelevant but may matter under a stated condition.
+
+Unless retained by one of the first two lists, every complete earlier tool exchange and every autoread currently eligible for pruning is selected for removal. The current `context_prune` call is the anchor: later messages and tool results are unaffected by that anchor.
+
+The agent is instructed to put durable conclusions, conditional information, and its next action in visible prose immediately before the call. That prose survives pruning. Tau enforces that `context_prune` is the only tool call in its assistant message, but it does not judge whether the preceding prose is complete or useful.
+
+Tau prepares the result without changing the active context. It estimates the token count before and after the proposed removal, including any replacement file contents and deferred-file note. It applies the prune only when the estimated saving is at least `minimumReclaimTokens`.
+
+After a successful call, Tau removes the selected evidence from future model input. The saved conversation remains unchanged. The tool result records exactly which tool exchanges and autoreads were removed, retained, or refreshed.
+
+## Why a prune is skipped
+
+A skipped prune removes nothing and publishes no deferred-file note or replacement file contents. Tau skips when any of these checks fails:
+
+1. Context Pruning is disabled.
+2. Tau does not have a usable copy of the exact messages most recently sent to the model.
+3. The `context_prune` call is missing from the active branch.
+4. Another tool call appears beside `context_prune` in the same assistant message.
+5. The current model input contains duplicate, mismatched, or orphaned tool records that Tau cannot safely remove. An unmatched call from an explicitly aborted assistant response is discarded and does not block pruning.
+6. A retained tool-call ID is duplicated, missing, or is not a complete exchange in the current model input.
+7. Two selected file paths resolve to the same file, including aliases and symbolic links across `keepFiles` and `deferFiles`.
+8. A retained file does not have an earlier complete-file read, has malformed read evidence, cannot currently be read as UTF-8, or exceeds the 1 MiB complete-file limit.
+9. The estimated saving is below `minimumReclaimTokens`.
+
+Cancellation, a session lifecycle change during preparation, and failures while publishing an already prepared prune are reported as tool failures rather than skipped prunes.
+
+### “The latest provider-context projection is unavailable”
+
+This error uses an internal term. In plain language, Tau did not have a usable copy of the exact message list the model had just received.
+
+Tau normally saves that list while Pi prepares a model turn. Tau clears it on session start, branch changes, compaction, and shutdown. It also refuses to save a list with invalid tool-call/result pairing. If the list is missing or belongs to an earlier session state, `context_prune` stops before checking the agent’s selections or estimating savings.
+
+Pi can save an aborted assistant response that requested a tool but never received a result. Tau discards that abandoned call when preparing later model input. Other unmatched calls and results still fail validation because Tau cannot prove whether their evidence is incomplete.
+
+For example, a row that says `context_prune 29 selections` followed by this error means:
+
+- the agent supplied 29 items across the three selection lists;
+- Tau did not inspect or apply those selections;
+- no old tool calls were removed;
+- the context percentage did not cause the rejection;
+- the current implementation does not reconstruct the missing list during that call.
+
+The response tells the agent not to retry immediately because repeated calls in the same state would usually fail for the same reason. A later model turn should normally give Tau a new list. Repeated occurrences indicate a feature defect or a session history that Tau cannot safely process.
+
+## Retaining files
+
+`keepFiles` preserves complete current file knowledge, not necessarily the original read row.
+
+Each selected file must have been read completely earlier in the current model input. A partial read is insufficient. Tau then reads the file from disk again and chooses the cheaper valid representation:
+
+- Keep an existing complete snapshot when it still matches the file.
+- Keep a baseline plus later read diffs when that chain reconstructs the current file and costs no more than a fresh snapshot.
+- Add one fresh complete snapshot when the earlier evidence is stale, broken, or more expensive.
+
+The entire prune is skipped if any selected file fails. Every retained file must currently exist, be valid UTF-8, and fit within the 1 MiB complete-file snapshot limit. Read the complete file before selecting it for retention.
+
+Paths are resolved from the session working directory. Paths outside it are stored as absolute paths. Existing symbolic links and other aliases are resolved before duplicate checks.
+
+The required `relevance` text explains the agent’s selection. It does not loosen any validation rule.
+
+## Retaining tool exchanges
+
+`keepToolCalls` retains both sides of an earlier exchange: the assistant’s tool call and the matching tool result. The ID must occur exactly once, the tool names must match, and the result must follow the call.
+
+Retention is exact. Selecting one call does not retain nearby calls from the same assistant message. Tau can remove one parallel tool exchange while keeping another, then removes the empty assistant message if no content remains.
+
+The required `relevance` text explains why the exchange matters. Tau does not use that text to infer additional exchanges to retain.
+
+## Deferring files
+
+`deferFiles` does not read or preserve file contents. It adds a hidden advisory note telling the model why each file was deferred and when to reconsider it. The latest successful anchor replaces the previous deferred-file list.
+
+A deferred path may be missing. It still participates in canonical path and duplicate checks. Its `reason` and `relevantWhen` text are passed to the model as written.
+
+## Automatic hints
+
+Tau checks context usage only after a turn that produced at least one tool result. It emits at most one marker for the strongest newly crossed boundary.
+
+With the defaults:
+
+- `nudgeEveryPercent: 20` allows markers after enough context growth to cross 20% intervals.
+- `pressurePercent: 50` makes a marker a pressure hint only when raw usage is greater than 50%.
+
+Below the pressure threshold, the hidden instruction says no prune is required unless broad exploration has converged or substantial evidence is already irrelevant. Above it, the instruction asks the agent to finish the current coherent step and prune when the projected saving is substantial. A hint never calls the tool or bypasses its checks.
+
+After a successful prune, the first later tool-using turn records a new usage baseline and emits no automatic marker. Tau waits for context usage to grow by another `nudgeEveryPercent` from that baseline. The baseline and crossed boundaries are stored with the active branch so compaction and branch navigation do not immediately repeat hints.
+
+The visible marker is compact. The full instruction stays hidden, and the agent is told not to discuss internal context management.
+
+## Manual requests
+
+Run `/prune` with no arguments to ask the agent to create an anchor and continue unfinished work. The command starts an agent turn immediately. It does not force a successful prune or bypass file, tool-pairing, current-message, lifecycle, or minimum-savings checks.
+
+Extra arguments produce `Usage: /prune`. When Context Pruning is disabled, the command reports that state and does not start a prune turn.
+
+Asking the agent to prune in ordinary chat has the same execution boundaries once it calls `context_prune`.
+
+## Branches, compaction, and display
+
+Applied anchors belong to the active session branch. Switching branches rebuilds the removed-evidence set, deferred-file list, automatic-hint baseline, and warning-colored rows from that branch’s valid tool results. A skipped or malformed prune result does not become an anchor.
+
+Context Pruning does not cancel or replace Pi’s manual, threshold, or overflow compaction. Compaction clears Tau’s saved model-input list; Tau must receive another model-input event before a later prune can use it.
+
+Tau-owned tool rows and autoreads turn warning-colored when their evidence has been pruned. Native Pi rows and fallback rows that do not use Tau’s row-state renderer do not change color. Their evidence is still removed from model input.
 
 ## Settings
 
 Settings live under `extensions.contextPruning` in Tau settings.
 
-- `enabled`: enables the tool, `/prune`, automatic markers, branch replay, context projection, and pruned-row display. Defaults to `true`.
-- `nudgeEveryPercent`: context-growth interval between automatic markers. Defaults to `20`.
-- `pressurePercent`: usage above which a marker suggests pruning after the current coherent step. Defaults to `50`.
-- `minimumReclaimTokens`: minimum estimated saving required to apply a prune. Defaults to `8000`.
+- `enabled`: enables the tool, `/prune`, automatic markers, branch replay, context filtering, and pruned-row display. Defaults to `true`.
+- `nudgeEveryPercent`: integer context-growth interval between automatic hints. Allowed range: `1` through `100`. Defaults to `20`.
+- `pressurePercent`: integer usage percentage above which hints use pressure wording. Allowed range: `1` through `99`. Defaults to `50`.
+- `minimumReclaimTokens`: positive integer minimum estimated saving required to apply a prune. Defaults to `8000`.

package/extensions/context-pruning/projection.ts

@@ -10,7 +10,8 @@ export function projectContext(
 	messages: readonly ContextMessage[],
 	state: ActiveContextPruningState,
 ): ContextMessage[] {
-	const inputPairs = indexToolPairs(messages);
+	const abandonedToolCallIds = new Set<string>();
+	const inputPairs = indexToolPairs(messages, abandonedToolCallIds);
 	const anchorBoundary = visibleAnchorBoundary(messages, state.latestAnchorToolCallId, inputPairs);
 	const prunedToolCallIds = state.latestAnchorToolCallId === undefined ? new Set<string>() : state.prunedToolCallIds;
 	const prunedAutoreadRowIds =
@@ -43,7 +44,7 @@ export function projectContext(
 		const content = message.content.filter((block) => {
 			const remove =
 				(removeThinking && block.type === "thinking") ||
-				(block.type === "toolCall" && prunedToolCallIds.has(block.id));
+				(block.type === "toolCall" && (prunedToolCallIds.has(block.id) || abandonedToolCallIds.has(block.id)));
 			if (remove) changed = true;
 			return !remove;
 		});
@@ -77,8 +78,11 @@ function visibleAnchorBoundary(
 	return block?.type === "toolCall" && block.name === "context_prune" ? pair.resultIndex : undefined;
 }
 
-function indexToolPairs(messages: readonly ContextMessage[]): Map<string, { callIndex: number; resultIndex: number }> {
-	const calls = new Map<string, number>();
+function indexToolPairs(
+	messages: readonly ContextMessage[],
+	abandonedToolCallIds?: Set<string>,
+): Map<string, { callIndex: number; resultIndex: number }> {
+	const calls = new Map<string, { index: number; aborted: boolean }>();
 	const results = new Map<string, number>();
 	for (let index = 0; index < messages.length; index += 1) {
 		const message = messages[index];
@@ -86,7 +90,7 @@ function indexToolPairs(messages: readonly ContextMessage[]): Map<string, { call
 			for (const block of message.content) {
 				if (block.type !== "toolCall") continue;
 				if (calls.has(block.id)) throw new Error(`Duplicate tool call in projected context: ${block.id}`);
-				calls.set(block.id, index);
+				calls.set(block.id, { index, aborted: message.stopReason === "aborted" });
 			}
 		} else if (message.role === "toolResult") {
 			if (results.has(message.toolCallId))
@@ -96,10 +100,16 @@ function indexToolPairs(messages: readonly ContextMessage[]): Map<string, { call
 	}
 
 	const pairs = new Map<string, { callIndex: number; resultIndex: number }>();
-	for (const [id, callIndex] of calls) {
+	for (const [id, call] of calls) {
 		const resultIndex = results.get(id);
-		if (resultIndex === undefined) throw new Error(`Orphaned tool call in projected context: ${id}`);
-		pairs.set(id, { callIndex, resultIndex });
+		if (resultIndex === undefined) {
+			if (call.aborted && abandonedToolCallIds) {
+				abandonedToolCallIds.add(id);
+				continue;
+			}
+			throw new Error(`Orphaned tool call in projected context: ${id}`);
+		}
+		pairs.set(id, { callIndex: call.index, resultIndex });
 	}
 	for (const id of results.keys()) {
 		if (!calls.has(id)) throw new Error(`Orphaned tool result in projected context: ${id}`);

package/extensions/image-gen/README.md

@@ -1,11 +1,13 @@
 # Image Generation
 
-`image_gen` generates raster images and edits up to three local raster images with Grok Imagine. It uses `grok-imagine-image-quality` and saves results under `~/.local/share/tau-agent/images/` by default. Pass an explicit path with the expected image extension when the image should be saved in the current repository or another chosen location.
+`image_gen` generates raster images and edits up to three local raster images with OpenAI GPT Image or xAI Grok Imagine. It uses `gpt-image-2` for OpenAI and `grok-imagine-image-quality` for xAI.
 
-Run `/login xai` and choose either a subscription or API-key login before invoking the tool.
+By default, the tool follows the parent model: GPT models prefer OpenAI and Grok models prefer xAI. If that provider has no configured authentication, it tries the other provider. The model can pass `provider: "openai"` or `provider: "xai"` to override automatic selection.
+
+Run `/login openai-codex` for GPT Image or `/login xai` for Grok Imagine before invoking the tool. OpenAI image generation uses the existing Codex subscription login; xAI supports its configured login methods.
 
 Run `/reload` after installing or changing the extension.
 
-The model invokes `image_gen` with a prompt. For edits, it also supplies one to three local PNG, JPEG, or WebP paths. Successful images up to 12 MiB are returned inline for inspection; larger results remain available at the saved path.
+Results are saved under `~/.local/share/tau-agent/images/` by default. Pass an explicit path with the expected image extension when the image should be saved in the current repository or another chosen location. For edits, supply one to three local PNG, JPEG, or WebP paths. Successful images up to 12 MiB are returned inline for inspection; larger results remain available at the saved path.
 
-xAI controls model availability and subscription entitlements. A successful login does not guarantee that every account can use Grok Imagine.
+The OpenAI path calls a private Codex backend, which may change without notice. xAI controls Grok Imagine availability and subscription entitlements. A successful login does not guarantee image access.

package/extensions/image-gen/client.ts

@@ -1,4 +1,10 @@
-import { XAI_API_BASE_URL, XAI_IMAGE_MODEL } from "./constants.ts";
+import {
+	OPENAI_IMAGE_API_BASE_URL,
+	OPENAI_IMAGE_MODEL,
+	XAI_API_BASE_URL,
+	XAI_IMAGE_MODEL,
+	type ImageProvider,
+} from "./constants.ts";
 
 const MAX_ERROR_BODY_BYTES = 8192;
 const MAX_ERROR_MESSAGE_LENGTH = 2000;
@@ -6,6 +12,11 @@ const REQUEST_TIMEOUT_MS = 60_000;
 const MAX_ATTEMPTS = 3;
 const RETRY_BASE_DELAY_MS = 500;
 
+export interface CodexAuth {
+	token: string;
+	accountId: string;
+}
+
 export interface EditImage {
 	mimeType: "image/png" | "image/jpeg" | "image/webp";
 	data: string;
@@ -28,6 +39,29 @@ function isRecord(value: unknown): value is Record<string, unknown> {
 	return typeof value === "object" && value !== null && !Array.isArray(value);
 }
 
+export function resolveCodexAuth(token: string): CodexAuth {
+	const invalidCredential = () =>
+		new Error("The OpenAI Codex credential does not contain a usable ChatGPT account ID. Run /login again.");
+	const segments = token.split(".");
+	if (segments.length !== 3 || !segments[1]) throw invalidCredential();
+	if (!/^[A-Za-z0-9_-]+$/.test(segments[1]) || segments[1].length % 4 === 1) throw invalidCredential();
+
+	let payload: unknown;
+	try {
+		const decoded = Buffer.from(segments[1], "base64url");
+		if (decoded.toString("base64url") !== segments[1]) throw invalidCredential();
+		payload = JSON.parse(decoded.toString("utf8"));
+	} catch {
+		throw invalidCredential();
+	}
+	if (!isRecord(payload)) throw invalidCredential();
+	const authClaim = payload["https://api.openai.com/auth"];
+	if (!isRecord(authClaim)) throw invalidCredential();
+	const accountId = authClaim.chatgpt_account_id;
+	if (typeof accountId !== "string" || !accountId.trim()) throw invalidCredential();
+	return { token, accountId: accountId.trim() };
+}
+
 async function boundedError(response: HttpResponse): Promise<string> {
 	if (!response.body) return "";
 	const reader = response.body.getReader();
@@ -85,25 +119,61 @@ export function detectImageMimeType(bytes: Buffer): GeneratedImage["mimeType"] |
 	return undefined;
 }
 
-function decodeImageResponse(value: unknown): GeneratedImage {
+function decodeImageResponse(value: unknown, service: "OpenAI Codex" | "xAI", pngOnly: boolean): GeneratedImage {
 	if (!isRecord(value) || !Array.isArray(value.data) || !isRecord(value.data[0])) {
-		throw new Error("xAI returned an invalid image response");
+		throw new Error(`${service} returned an invalid image response`);
 	}
 	const encoded = value.data[0].b64_json;
-	if (typeof encoded !== "string" || !encoded.trim()) throw new Error("xAI returned an invalid image response");
+	if (typeof encoded !== "string" || !encoded.trim()) throw new Error(`${service} returned an invalid image response`);
 	const base64 = encoded.trim();
 	if (!/^[A-Za-z0-9+/]*={0,2}$/.test(base64) || base64.length % 4 === 1) {
-		throw new Error("xAI returned invalid base64 image data");
+		throw new Error(`${service} returned invalid base64 image data`);
 	}
 	const bytes = Buffer.from(base64, "base64");
 	if (bytes.toString("base64").replace(/=+$/, "") !== base64.replace(/=+$/, "")) {
-		throw new Error("xAI returned invalid base64 image data");
+		throw new Error(`${service} returned invalid base64 image data`);
 	}
 	const mimeType = detectImageMimeType(bytes);
-	if (!mimeType) throw new Error("xAI returned unsupported image data");
+	if (!mimeType || (pngOnly && mimeType !== "image/png")) {
+		throw new Error(`${service} returned unsupported image data`);
+	}
 	return { bytes, base64: bytes.toString("base64"), mimeType };
 }
 
+async function requestOpenAIImage(
+	operation: "generation" | "edit",
+	body: Record<string, unknown>,
+	auth: CodexAuth,
+	signal?: AbortSignal,
+): Promise<GeneratedImage> {
+	const route = operation === "generation" ? "generations" : "edits";
+	const response = (await fetch(`${OPENAI_IMAGE_API_BASE_URL}/${route}`, {
+		method: "POST",
+		headers: {
+			Accept: "application/json",
+			Authorization: `Bearer ${auth.token}`,
+			"chatgpt-account-id": auth.accountId,
+			"Content-Type": "application/json",
+			originator: "pi",
+		},
+		body: JSON.stringify(body),
+		signal,
+	})) as HttpResponse;
+	if (!response.ok) {
+		const message = serverErrorMessage(await boundedError(response), auth.token);
+		throw new Error(
+			`OpenAI Codex image ${operation} failed with status ${response.status}${message ? `: ${message}` : ""}`,
+		);
+	}
+	let value: unknown;
+	try {
+		value = await response.json();
+	} catch {
+		throw new Error("OpenAI Codex returned a non-JSON image response");
+	}
+	return decodeImageResponse(value, "OpenAI Codex", true);
+}
+
 function retryable(status: number): boolean {
 	return status === 408 || status === 409 || status === 425 || status === 429 || status >= 500;
 }
@@ -124,7 +194,7 @@ async function delay(milliseconds: number, signal?: AbortSignal): Promise<void>
 	});
 }
 
-async function requestImage(
+async function requestXaiImage(
 	operation: "generation" | "edit",
 	body: Record<string, unknown>,
 	token: string,
@@ -160,7 +230,7 @@ async function requestImage(
 			} catch {
 				throw new Error("xAI returned a non-JSON image response");
 			}
-			return decodeImageResponse(value);
+			return decodeImageResponse(value, "xAI", false);
 		}
 		const message = serverErrorMessage(await boundedError(response), token);
 		if (!retryable(response.status) || attempt === MAX_ATTEMPTS) {
@@ -173,8 +243,21 @@ async function requestImage(
 	throw new Error(`xAI image ${operation} failed`);
 }
 
-export function generateImage(prompt: string, token: string, signal?: AbortSignal): Promise<GeneratedImage> {
-	return requestImage(
+export function generateImage(
+	provider: ImageProvider,
+	prompt: string,
+	token: string,
+	signal?: AbortSignal,
+): Promise<GeneratedImage> {
+	if (provider === "openai") {
+		return requestOpenAIImage(
+			"generation",
+			{ prompt, model: OPENAI_IMAGE_MODEL, background: "auto", quality: "auto", size: "auto" },
+			resolveCodexAuth(token),
+			signal,
+		);
+	}
+	return requestXaiImage(
 		"generation",
 		{ model: XAI_IMAGE_MODEL, prompt, n: 1, resolution: "1k", response_format: "b64_json" },
 		token,
@@ -183,13 +266,29 @@ export function generateImage(prompt: string, token: string, signal?: AbortSigna
 }
 
 export function editImage(
+	provider: ImageProvider,
 	prompt: string,
 	images: readonly EditImage[],
 	token: string,
 	signal?: AbortSignal,
 ): Promise<GeneratedImage> {
+	if (provider === "openai") {
+		return requestOpenAIImage(
+			"edit",
+			{
+				images: images.map((image) => ({ image_url: `data:${image.mimeType};base64,${image.data}` })),
+				prompt,
+				model: OPENAI_IMAGE_MODEL,
+				background: "auto",
+				quality: "auto",
+				size: "auto",
+			},
+			resolveCodexAuth(token),
+			signal,
+		);
+	}
 	const references = images.map((image) => ({ url: `data:${image.mimeType};base64,${image.data}` }));
-	return requestImage(
+	return requestXaiImage(
 		"edit",
 		{
 			model: XAI_IMAGE_MODEL,

package/extensions/image-gen/constants.ts

@@ -1,3 +1,8 @@
+export const OPENAI_PROVIDER = "openai-codex";
+export const OPENAI_IMAGE_MODEL = "gpt-image-2";
+export const OPENAI_IMAGE_API_BASE_URL = "https://chatgpt.com/backend-api/codex/images";
 export const XAI_PROVIDER = "xai";
 export const XAI_IMAGE_MODEL = "grok-imagine-image-quality";
 export const XAI_API_BASE_URL = "https://api.x.ai/v1";
+
+export type ImageProvider = "openai" | "xai";

package/extensions/image-gen/index.ts

@@ -5,7 +5,13 @@ import { homedir } from "node:os";
 import { basename, dirname, extname, isAbsolute, join, resolve } from "node:path";
 import { type Static, Type } from "typebox";
 import { detectImageMimeType, editImage, generateImage, type EditImage, type GeneratedImage } from "./client.ts";
-import { XAI_IMAGE_MODEL, XAI_PROVIDER } from "./constants.ts";
+import {
+	OPENAI_IMAGE_MODEL,
+	OPENAI_PROVIDER,
+	XAI_IMAGE_MODEL,
+	XAI_PROVIDER,
+	type ImageProvider,
+} from "./constants.ts";
 
 const MAX_INPUT_BYTES = 50 * 1024 * 1024;
 const MAX_INLINE_BYTES = 12 * 1024 * 1024;
@@ -13,6 +19,12 @@ const MAX_INLINE_BYTES = 12 * 1024 * 1024;
 const imageGenSchema = Type.Object(
 	{
 		prompt: Type.String({ minLength: 1 }),
+		provider: Type.Optional(
+			Type.Union([Type.Literal("openai"), Type.Literal("xai")], {
+				description:
+					"Image provider override. Omit to follow the parent model, preferring OpenAI for GPT and xAI for Grok.",
+			}),
+		),
 		path: Type.Optional(
 			Type.String({ description: "Explicit image destination path; defaults to Tau's external image store" }),
 		),
@@ -25,7 +37,8 @@ type ImageGenParams = Static<typeof imageGenSchema>;
 
 interface ImageGenDetails {
 	path: string;
-	model: typeof XAI_IMAGE_MODEL;
+	provider: ImageProvider;
+	model: typeof OPENAI_IMAGE_MODEL | typeof XAI_IMAGE_MODEL;
 	operation: "generate" | "edit";
 }
 
@@ -41,7 +54,7 @@ export default function imageGenExtension(pi: ExtensionAPI): void {
 			name: "image_gen",
 			label: "Image Generation",
 			description:
-				"Generate a requested raster image or AI-edit existing images with configured xAI authentication. Omit referenced_image_paths to generate; pass one to three local paths to edit or compose. Omit path to use Tau's external image store; pass path only when the user explicitly requests a repository file or other destination. Returns the image for inspection.",
+				"Generate a requested raster image or AI-edit existing images with OpenAI GPT Image or xAI Grok Imagine. Omit provider to follow the parent model; set it to openai or xai to override. Omit referenced_image_paths to generate; pass one to three local paths to edit or compose. Omit path to use Tau's external image store; pass path only when the user explicitly requests a repository file or other destination. Returns the image for inspection.",
 			parameters: imageGenSchema,
 			async execute(_toolCallId, params: ImageGenParams, signal, onUpdate, ctx) {
 				signal?.throwIfAborted();
@@ -61,10 +74,38 @@ export default function imageGenExtension(pi: ExtensionAPI): void {
 					throw new Error("Image path must end in .jpg, .jpeg, .png, or .webp");
 				}
 
-				const token = await ctx.modelRegistry.getApiKeyForProvider(XAI_PROVIDER);
-				if (!token) {
-					throw new Error("xAI authentication is unavailable. Run /login xai and choose a login method.");
+				const parentUsesXai =
+					ctx.model?.provider.toLowerCase() === XAI_PROVIDER || ctx.model?.id.toLowerCase().includes("grok") === true;
+				const preferredProvider: ImageProvider = params.provider ?? (parentUsesXai ? "xai" : "openai");
+				const providers: readonly ImageProvider[] = params.provider
+					? [params.provider]
+					: preferredProvider === "xai"
+						? ["xai", "openai"]
+						: ["openai", "xai"];
+				let provider: ImageProvider | undefined;
+				let token: string | undefined;
+				for (const candidate of providers) {
+					const candidateToken = await ctx.modelRegistry.getApiKeyForProvider(
+						candidate === "openai" ? OPENAI_PROVIDER : XAI_PROVIDER,
+					);
+					if (candidateToken) {
+						provider = candidate;
+						token = candidateToken;
+						break;
+					}
+				}
+				if (!provider || !token) {
+					if (params.provider === "openai") {
+						throw new Error("OpenAI Codex authentication is unavailable. Run /login openai-codex.");
+					}
+					if (params.provider === "xai") {
+						throw new Error("xAI authentication is unavailable. Run /login xai and choose a login method.");
+					}
+					throw new Error(
+						"Image generation authentication is unavailable. Run /login for OpenAI Codex or xAI.",
+					);
 				}
+				const model = provider === "openai" ? OPENAI_IMAGE_MODEL : XAI_IMAGE_MODEL;
 				const images: EditImage[] = [];
 				for (const path of params.referenced_image_paths ?? []) {
 					const rawPath = path.startsWith("@") ? path.slice(1) : path;
@@ -88,16 +129,16 @@ export default function imageGenExtension(pi: ExtensionAPI): void {
 							type: "text",
 							text:
 								operation === "generate"
-									? `Generating image with ${XAI_IMAGE_MODEL}...`
-									: `Editing image with ${XAI_IMAGE_MODEL}...`,
+									? `Generating image with ${model}...`
+									: `Editing image with ${model}...`,
 						},
 					],
 					details: undefined,
 				});
 				const generated =
 					operation === "generate"
-						? await generateImage(prompt, token, signal)
-						: await editImage(prompt, images, token, signal);
+						? await generateImage(provider, prompt, token, signal)
+						: await editImage(provider, prompt, images, token, signal);
 				signal?.throwIfAborted();
 				const generatedExtension = outputExtension(generated);
 				if (requestedAbsolutePath) {
@@ -106,7 +147,9 @@ export default function imageGenExtension(pi: ExtensionAPI): void {
 						requestedExtension === generatedExtension ||
 						(generatedExtension === ".jpg" && requestedExtension === ".jpeg");
 					if (!matches)
-						throw new Error(`xAI returned ${generated.mimeType}; destination must end in ${generatedExtension}`);
+						throw new Error(
+							`${provider === "openai" ? "OpenAI Codex" : "xAI"} returned ${generated.mimeType}; destination must end in ${generatedExtension}`,
+						);
 				}
 				const absolutePath =
 					requestedAbsolutePath ??
@@ -129,13 +172,13 @@ export default function imageGenExtension(pi: ExtensionAPI): void {
 				});
 
 				const verb = operation === "generate" ? "Generated" : "Edited";
-				const details: ImageGenDetails = { path: absolutePath, model: XAI_IMAGE_MODEL, operation };
+				const details: ImageGenDetails = { path: absolutePath, provider, model, operation };
 				if (generated.bytes.length > MAX_INLINE_BYTES) {
 					return {
 						content: [
 							{
 								type: "text",
-								text: `${verb} image saved to ${absolutePath}. The PNG exceeds the 12 MiB attachment limit, so it was not added to model context.`,
+								text: `${verb} image saved to ${absolutePath}. The image exceeds the 12 MiB attachment limit, so it was not added to model context.`,
 							},
 						],
 						details,

package/extensions/tau-help/help.md

@@ -18,6 +18,10 @@ Names sessions from their first request so saved sessions remain findable.
 
 Adds `/branch` to create and switch Git branches from the TUI.
 
+## cache-diagnostics
+
+Records private prompt-cache fingerprints without storing prompt content. Run `/cache-debug` after suspicious cache misses to write a bounded investigation report under `~/.pi/agent/cache-diagnostics/reports/`.
+
 ## clear-screen
 
 Adds `/clear-screen` to clear terminal output without changing the session.
@@ -48,7 +52,7 @@ Adds `/ideas` to log rough ideas or open the ideas browser.
 
 ## image-gen
 
-Gives the agent a Grok Imagine generation and editing tool. Run `/login xai` and choose a subscription or API-key login before use. Generated images are saved for inspection.
+Gives the agent an OpenAI GPT Image and xAI Grok Imagine generation and editing tool. It follows the parent model by default and can override the provider per request. Run `/login openai-codex` or `/login xai` before use. Generated images are saved for inspection.
 
 ## manage-sessions
 

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@shanepadgett/tau-agent",
-	"version": "0.18.1",
+	"version": "0.20.0",
 	"description": "Tau is a custom agentic harness built with pi extensions",
 	"type": "module",
 	"license": "MIT",
@@ -28,7 +28,7 @@
 		"README.md"
 	],
 	"dependencies": {
-		"@shanepadgett/tau-tui": "0.18.1",
+		"@shanepadgett/tau-tui": "0.20.0",
 		"@toon-format/toon": "2.3.0",
 		"smol-toml": "1.7.0"
 	},