@shanepadgett/tau-agent 0.20.2 → 0.22.0

package/shared/context-pruning-state.ts

@@ -1,40 +1,38 @@
 import type { SessionEntry } from "@earendil-works/pi-coding-agent";
 
 const CONTEXT_PRUNE_TOOL = "context_prune";
-const AUTOREAD_MESSAGE_TYPE = "tau.autoread";
 
-export interface ContextPruneRefreshedFileV1 {
+export interface ContextPruneRefreshedFileV2 {
 	path: string;
 	rowId: string;
 	servedHash: string;
+	autoreadDetails: Record<string, unknown>;
 }
 
-export interface ContextPruneDeferredFileV1 {
+export interface ContextPruneDeferredFileV2 {
 	path: string;
 	reason: string;
 	relevantWhen: string;
 }
 
-export interface ContextPruneDetailsV1 {
-	v: 1;
-	status: "applied" | "skipped";
+export interface ContextPruneDetailsV2 {
+	v: 2;
 	anchorToolCallId: string;
-	newlyPrunedToolCallIds: string[];
-	newlyPrunedAutoreadRowIds: string[];
+	prunedToolCallIds: string[];
+	prunedAutoreadRowIds: string[];
 	retainedToolCallIds: string[];
 	retainedAutoreadRowIds: string[];
-	refreshedFiles: ContextPruneRefreshedFileV1[];
-	deferredFiles: ContextPruneDeferredFileV1[];
-	tokensBefore: number;
-	tokensAfter: number;
-	tokensReclaimed: number;
+	refreshedFiles: ContextPruneRefreshedFileV2[];
+	deferredFiles: ContextPruneDeferredFileV2[];
+	warnings: string[];
 }
 
 export interface ActiveContextPruningState {
 	latestAnchorToolCallId: string | undefined;
+	retainedToolCallIds: ReadonlySet<string>;
 	prunedToolCallIds: ReadonlySet<string>;
 	prunedAutoreadRowIds: ReadonlySet<string>;
-	deferredFiles: readonly ContextPruneDeferredFileV1[];
+	deferredFiles: readonly ContextPruneDeferredFileV2[];
 }
 
 const RUNTIME_ENABLED_KEY = Symbol.for("@earendil-works/tau-agent/context-pruning-enabled");
@@ -48,67 +46,35 @@ function isContextPruningEnabled(): boolean {
 	return runtimeState[RUNTIME_ENABLED_KEY] === true;
 }
 
-export function parseContextPruneDetailsV1(value: unknown): ContextPruneDetailsV1 | undefined {
-	if (
-		!hasExactKeys(value, [
-			"v",
-			"status",
-			"anchorToolCallId",
-			"newlyPrunedToolCallIds",
-			"newlyPrunedAutoreadRowIds",
-			"retainedToolCallIds",
-			"retainedAutoreadRowIds",
-			"refreshedFiles",
-			"deferredFiles",
-			"tokensBefore",
-			"tokensAfter",
-			"tokensReclaimed",
-		])
-	)
-		return undefined;
-	if (value.v !== 1 || (value.status !== "applied" && value.status !== "skipped")) return undefined;
-	if (!isNonEmptyString(value.anchorToolCallId)) return undefined;
-
-	const newlyPrunedToolCallIds = parseUniqueStrings(value.newlyPrunedToolCallIds);
-	const newlyPrunedAutoreadRowIds = parseUniqueStrings(value.newlyPrunedAutoreadRowIds);
-	const retainedToolCallIds = parseUniqueStrings(value.retainedToolCallIds);
+export function parseContextPruneDetailsV2(value: unknown): ContextPruneDetailsV2 | undefined {
+	if (!isRecord(value) || value.v !== 2 || !isNonEmptyString(value.anchorToolCallId)) return undefined;
+	const prunedToolCallIds = parseUniqueStrings(value.prunedToolCallIds);
+	const prunedAutoreadRowIds = parseUniqueStrings(value.prunedAutoreadRowIds);
+	const retainedToolCallIds = parseUniqueStrings(value.retainedToolCallIds, true);
 	const retainedAutoreadRowIds = parseUniqueStrings(value.retainedAutoreadRowIds);
 	const refreshedFiles = parseRefreshedFiles(value.refreshedFiles);
 	const deferredFiles = parseDeferredFiles(value.deferredFiles);
+	const warnings = parseStrings(value.warnings, true);
 	if (
-		!newlyPrunedToolCallIds ||
-		!newlyPrunedAutoreadRowIds ||
+		!prunedToolCallIds ||
+		!prunedAutoreadRowIds ||
 		!retainedToolCallIds ||
 		!retainedAutoreadRowIds ||
 		!refreshedFiles ||
-		!deferredFiles
-	)
-		return undefined;
-	if (
-		!isNonNegativeFinite(value.tokensBefore) ||
-		!isNonNegativeFinite(value.tokensAfter) ||
-		!isFiniteNumber(value.tokensReclaimed) ||
-		value.tokensBefore - value.tokensAfter !== value.tokensReclaimed
+		!deferredFiles ||
+		!warnings
 	)
 		return undefined;
-	if (hasOverlap(newlyPrunedToolCallIds, retainedToolCallIds)) return undefined;
-	if (hasOverlap(newlyPrunedAutoreadRowIds, retainedAutoreadRowIds)) return undefined;
-	if (value.status === "skipped" && (newlyPrunedToolCallIds.length > 0 || newlyPrunedAutoreadRowIds.length > 0))
-		return undefined;
-
 	return {
-		v: 1,
-		status: value.status,
+		v: 2,
 		anchorToolCallId: value.anchorToolCallId,
-		newlyPrunedToolCallIds,
-		newlyPrunedAutoreadRowIds,
+		prunedToolCallIds,
+		prunedAutoreadRowIds,
 		retainedToolCallIds,
 		retainedAutoreadRowIds,
 		refreshedFiles,
 		deferredFiles,
-		tokensBefore: value.tokensBefore,
-		tokensAfter: value.tokensAfter,
-		tokensReclaimed: value.tokensReclaimed,
+		warnings,
 	};
 }
 
@@ -118,233 +84,71 @@ export function replayContextPruningState(
 ): ActiveContextPruningState {
 	const prunedToolCallIds = new Set<string>();
 	const prunedAutoreadRowIds = new Set<string>();
-	const state: ActiveContextPruningState = {
-		latestAnchorToolCallId: undefined,
+	let latestAnchorToolCallId: string | undefined;
+	let retainedToolCallIds = new Set<string>();
+	let deferredFiles: readonly ContextPruneDeferredFileV2[] = [];
+	if (enabled) {
+		for (const entry of branch) {
+			if (entry.type !== "message" || entry.message.role !== "toolResult") continue;
+			const result = entry.message;
+			if (result.toolName !== CONTEXT_PRUNE_TOOL) continue;
+			const details = parseContextPruneDetailsV2(result.details);
+			if (!details || details.anchorToolCallId !== result.toolCallId) continue;
+			for (const id of details.prunedToolCallIds) prunedToolCallIds.add(id);
+			for (const id of details.prunedAutoreadRowIds) prunedAutoreadRowIds.add(id);
+			latestAnchorToolCallId = details.anchorToolCallId;
+			retainedToolCallIds = new Set(details.retainedToolCallIds);
+			deferredFiles = details.deferredFiles;
+		}
+		for (const id of retainedToolCallIds) prunedToolCallIds.delete(id);
+	}
+	return {
+		latestAnchorToolCallId,
+		retainedToolCallIds,
 		prunedToolCallIds,
 		prunedAutoreadRowIds,
-		deferredFiles: [],
+		deferredFiles,
 	};
-	if (!enabled) return state;
-
-	const toolCalls = new Map<string, { count: number; name: string; index: number }>();
-	const toolResults = new Map<string, { count: number; name: string; index: number }>();
-	const autoreads = new Map<
-		string,
-		{ count: number; path: string; servedHash: string; source: string; batchId: string; index: number }
-	>();
-	for (let index = 0; index < branch.length; index += 1) {
-		const entry = branch[index];
-		if (entry.type === "message") {
-			const message = entry.message;
-			if (message.role === "assistant") {
-				for (const block of message.content) {
-					if (block.type !== "toolCall") continue;
-					const current = toolCalls.get(block.id);
-					toolCalls.set(block.id, { count: (current?.count ?? 0) + 1, name: block.name, index });
-				}
-			} else if (message.role === "toolResult") {
-				const current = toolResults.get(message.toolCallId);
-				toolResults.set(message.toolCallId, {
-					count: (current?.count ?? 0) + 1,
-					name: message.toolName,
-					index,
-				});
-			}
-			continue;
-		}
-		if (entry.type !== "custom_message" || entry.customType !== AUTOREAD_MESSAGE_TYPE) continue;
-		const autoread = parseAutoreadReference(entry.details);
-		if (!autoread) continue;
-		const current = autoreads.get(autoread.rowId);
-		autoreads.set(autoread.rowId, {
-			count: (current?.count ?? 0) + 1,
-			path: autoread.path,
-			servedHash: autoread.servedHash,
-			source: autoread.source,
-			batchId: autoread.batchId,
-			index,
-		});
-	}
-
-	for (let resultIndex = 0; resultIndex < branch.length; resultIndex += 1) {
-		const entry = branch[resultIndex];
-		if (entry.type !== "message" || entry.message.role !== "toolResult") continue;
-		const result = entry.message;
-		if (result.toolName !== CONTEXT_PRUNE_TOOL) continue;
-		const details = parseContextPruneDetailsV1(result.details);
-		if (!details || details.status !== "applied" || details.anchorToolCallId !== result.toolCallId) continue;
-		if (!isCompleteExchange(details.anchorToolCallId, CONTEXT_PRUNE_TOOL, toolCalls, toolResults)) continue;
-		const anchorCallIndex = toolCalls.get(details.anchorToolCallId)?.index;
-		if (anchorCallIndex === undefined || anchorCallIndex >= resultIndex) continue;
-		if (details.newlyPrunedToolCallIds.includes(details.anchorToolCallId)) continue;
-		if (details.retainedToolCallIds.includes(details.anchorToolCallId)) continue;
-		if (
-			!allCompleteExchanges(details.newlyPrunedToolCallIds, toolCalls, toolResults) ||
-			!allCompleteExchanges(details.retainedToolCallIds, toolCalls, toolResults) ||
-			!allAutoreads(details.newlyPrunedAutoreadRowIds, autoreads) ||
-			!allAutoreads(details.retainedAutoreadRowIds, autoreads)
-		)
-			continue;
-		if (
-			details.newlyPrunedToolCallIds.some((id) => !exchangePrecedes(id, anchorCallIndex, toolCalls, toolResults)) ||
-			details.retainedToolCallIds.some((id) => !exchangePrecedes(id, anchorCallIndex, toolCalls, toolResults)) ||
-			details.newlyPrunedAutoreadRowIds.some((id) => {
-				const row = autoreads.get(id);
-				return row === undefined || row.index >= anchorCallIndex;
-			})
-		)
-			continue;
-		if (
-			details.newlyPrunedToolCallIds.some((id) => prunedToolCallIds.has(id)) ||
-			details.retainedToolCallIds.some((id) => prunedToolCallIds.has(id)) ||
-			details.newlyPrunedAutoreadRowIds.some((id) => prunedAutoreadRowIds.has(id)) ||
-			details.retainedAutoreadRowIds.some((id) => prunedAutoreadRowIds.has(id))
-		)
-			continue;
-		if (
-			details.refreshedFiles.some((file) => {
-				const row = autoreads.get(file.rowId);
-				return (
-					row?.count !== 1 ||
-					row.path !== file.path ||
-					row.servedHash !== file.servedHash ||
-					row.source !== "context-pruning" ||
-					row.batchId !== details.anchorToolCallId ||
-					!details.retainedAutoreadRowIds.includes(file.rowId)
-				);
-			})
-		)
-			continue;
-		const refreshedRowIds = new Set(details.refreshedFiles.map((file) => file.rowId));
-		if (
-			details.retainedAutoreadRowIds.some((id) => {
-				const row = autoreads.get(id);
-				return row === undefined || (row.index >= anchorCallIndex && !refreshedRowIds.has(id));
-			})
-		)
-			continue;
-
-		for (const id of details.newlyPrunedToolCallIds) prunedToolCallIds.add(id);
-		for (const id of details.newlyPrunedAutoreadRowIds) prunedAutoreadRowIds.add(id);
-		state.latestAnchorToolCallId = details.anchorToolCallId;
-		state.deferredFiles = details.deferredFiles;
-	}
-	return state;
 }
 
-function parseRefreshedFiles(value: unknown): ContextPruneRefreshedFileV1[] | undefined {
+function parseRefreshedFiles(value: unknown): ContextPruneRefreshedFileV2[] | undefined {
 	if (!Array.isArray(value)) return undefined;
-	const files: ContextPruneRefreshedFileV1[] = [];
-	const rowIds = new Set<string>();
+	const files: ContextPruneRefreshedFileV2[] = [];
 	for (const item of value) {
-		if (!hasExactKeys(item, ["path", "rowId", "servedHash"])) return undefined;
-		if (!isNonEmptyString(item.path) || !isNonEmptyString(item.rowId) || !isNonEmptyString(item.servedHash))
-			return undefined;
-		if (rowIds.has(item.rowId)) return undefined;
-		rowIds.add(item.rowId);
-		files.push({ path: item.path, rowId: item.rowId, servedHash: item.servedHash });
+		if (!isRecord(item) || !isNonEmptyString(item.path) || !isNonEmptyString(item.rowId)) return undefined;
+		if (!isNonEmptyString(item.servedHash) || !isRecord(item.autoreadDetails)) return undefined;
+		files.push({
+			path: item.path,
+			rowId: item.rowId,
+			servedHash: item.servedHash,
+			autoreadDetails: { ...item.autoreadDetails },
+		});
 	}
 	return files;
 }
 
-function parseDeferredFiles(value: unknown): ContextPruneDeferredFileV1[] | undefined {
+function parseDeferredFiles(value: unknown): ContextPruneDeferredFileV2[] | undefined {
 	if (!Array.isArray(value)) return undefined;
-	const files: ContextPruneDeferredFileV1[] = [];
+	const files: ContextPruneDeferredFileV2[] = [];
 	for (const item of value) {
-		if (!hasExactKeys(item, ["path", "reason", "relevantWhen"])) return undefined;
-		if (!isNonEmptyString(item.path) || !isNonEmptyString(item.reason) || !isNonEmptyString(item.relevantWhen))
-			return undefined;
+		if (!isRecord(item) || typeof item.path !== "string" || typeof item.reason !== "string") return undefined;
+		if (typeof item.relevantWhen !== "string") return undefined;
 		files.push({ path: item.path, reason: item.reason, relevantWhen: item.relevantWhen });
 	}
 	return files;
 }
 
-function parseUniqueStrings(value: unknown): string[] | undefined {
-	if (!Array.isArray(value)) return undefined;
-	const strings: string[] = [];
-	const seen = new Set<string>();
-	for (const item of value) {
-		if (!isNonEmptyString(item) || seen.has(item)) return undefined;
-		seen.add(item);
-		strings.push(item);
-	}
+function parseUniqueStrings(value: unknown, allowEmpty = false): string[] | undefined {
+	const strings = parseStrings(value, allowEmpty);
+	if (!strings || new Set(strings).size !== strings.length) return undefined;
 	return strings;
 }
 
-function parseAutoreadReference(
-	value: unknown,
-): { rowId: string; path: string; servedHash: string; source: string; batchId: string } | undefined {
-	if (!isRecord(value) || value.status !== "read") return undefined;
-	if (
-		!isNonEmptyString(value.rowId) ||
-		!isNonEmptyString(value.path) ||
-		!isNonEmptyString(value.cwd) ||
-		!isNonEmptyString(value.source) ||
-		!isNonEmptyString(value.batchId)
-	)
+function parseStrings(value: unknown, allowEmpty = false): string[] | undefined {
+	if (!Array.isArray(value) || !value.every((item) => typeof item === "string" && (allowEmpty || item.length > 0))) {
 		return undefined;
-	if (!isRecord(value.readCache) || !isNonEmptyString(value.readCache.servedHash)) return undefined;
-	return {
-		rowId: value.rowId,
-		path: value.path,
-		servedHash: value.readCache.servedHash,
-		source: value.source,
-		batchId: value.batchId,
-	};
-}
-
-function allCompleteExchanges(
-	ids: readonly string[],
-	calls: ReadonlyMap<string, { count: number; name: string; index: number }>,
-	results: ReadonlyMap<string, { count: number; name: string; index: number }>,
-): boolean {
-	return ids.every((id) => {
-		const call = calls.get(id);
-		return call !== undefined && isCompleteExchange(id, call.name, calls, results);
-	});
-}
-
-function isCompleteExchange(
-	id: string,
-	name: string,
-	calls: ReadonlyMap<string, { count: number; name: string; index: number }>,
-	results: ReadonlyMap<string, { count: number; name: string; index: number }>,
-): boolean {
-	const call = calls.get(id);
-	const result = results.get(id);
-	return (
-		call?.count === 1 &&
-		result?.count === 1 &&
-		call.name === name &&
-		result.name === name &&
-		call.index < result.index
-	);
-}
-
-function exchangePrecedes(
-	id: string,
-	anchorCallIndex: number,
-	calls: ReadonlyMap<string, { index: number }>,
-	results: ReadonlyMap<string, { index: number }>,
-): boolean {
-	const call = calls.get(id);
-	const result = results.get(id);
-	return call !== undefined && result !== undefined && call.index < result.index && result.index < anchorCallIndex;
-}
-
-function allAutoreads(ids: readonly string[], autoreads: ReadonlyMap<string, { count: number }>): boolean {
-	return ids.every((id) => autoreads.get(id)?.count === 1);
-}
-
-function hasOverlap(left: readonly string[], right: readonly string[]): boolean {
-	const rightSet = new Set(right);
-	return left.some((item) => rightSet.has(item));
-}
-
-function hasExactKeys(value: unknown, keys: readonly string[]): value is Record<string, unknown> {
-	if (!isRecord(value)) return false;
-	const actual = Object.keys(value);
-	return actual.length === keys.length && keys.every((key) => Object.hasOwn(value, key));
+	}
+	return [...value];
 }
 
 function isRecord(value: unknown): value is Record<string, unknown> {
@@ -354,11 +158,3 @@ function isRecord(value: unknown): value is Record<string, unknown> {
 function isNonEmptyString(value: unknown): value is string {
 	return typeof value === "string" && value.length > 0;
 }
-
-function isNonNegativeFinite(value: unknown): value is number {
-	return typeof value === "number" && Number.isFinite(value) && value >= 0;
-}
-
-function isFiniteNumber(value: unknown): value is number {
-	return typeof value === "number" && Number.isFinite(value);
-}

package/shared/model-fallback/index.ts

@@ -24,7 +24,7 @@ interface ModelFallbackOptions {
 
 export async function resolveCandidates(
 	ctx: Pick<ExtensionContext, "modelRegistry" | "model" | "cwd" | "isProjectTrusted">,
-	preferredModels: ReadonlyArray<{ provider: string; model: string; reasoning: ThinkingLevel }>,
+	preferredModels: ReadonlyArray<{ provider: string; model: string; reasoning: ThinkingLevel | undefined }>,
 	includeParentModel: boolean,
 ): Promise<ModelCandidate[]> {
 	const settings = await loadTauExtensionSettings(ctx, modelFallbackSettings);

package/{extensions/image-gen → src/image-generation}/client.ts

@@ -24,7 +24,6 @@ export interface EditImage {
 
 export interface GeneratedImage {
 	bytes: Buffer;
-	base64: string;
 	mimeType: EditImage["mimeType"];
 }
 
@@ -58,7 +57,7 @@ export function resolveCodexAuth(token: string): CodexAuth {
 	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();
+	if (typeof accountId !== "string" || !/^[\x21-\x7e]+$/.test(accountId.trim())) throw invalidCredential();
 	return { token, accountId: accountId.trim() };
 }
 
@@ -87,7 +86,7 @@ async function boundedError(response: HttpResponse): Promise<string> {
 		.trim();
 }
 
-function serverErrorMessage(body: string, token: string): string {
+function serverErrorMessage(body: string, secrets: readonly string[]): string {
 	let message = body;
 	try {
 		const value: unknown = JSON.parse(body);
@@ -98,7 +97,10 @@ function serverErrorMessage(body: string, token: string): string {
 	} catch {
 		// Plain-text error body.
 	}
-	return message.replaceAll(token, "[redacted]").slice(0, MAX_ERROR_MESSAGE_LENGTH).trim();
+	for (const secret of secrets) {
+		if (secret) message = message.replaceAll(secret, "[redacted]");
+	}
+	return message.slice(0, MAX_ERROR_MESSAGE_LENGTH).trim();
 }
 
 export function detectImageMimeType(bytes: Buffer): GeneratedImage["mimeType"] | undefined {
@@ -137,7 +139,7 @@ function decodeImageResponse(value: unknown, service: "OpenAI Codex" | "xAI", pn
 	if (!mimeType || (pngOnly && mimeType !== "image/png")) {
 		throw new Error(`${service} returned unsupported image data`);
 	}
-	return { bytes, base64: bytes.toString("base64"), mimeType };
+	return { bytes, mimeType };
 }
 
 async function requestOpenAIImage(
@@ -147,20 +149,28 @@ async function requestOpenAIImage(
 	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;
+	let response: HttpResponse;
+	try {
+		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;
+	} catch (error) {
+		if (signal?.aborted) throw signal.reason;
+		let message = error instanceof Error ? error.message : "OpenAI Codex image request failed";
+		for (const secret of [auth.token, auth.accountId]) message = message.replaceAll(secret, "[redacted]");
+		throw new Error(message.slice(0, MAX_ERROR_MESSAGE_LENGTH).trim() || "OpenAI Codex image request failed");
+	}
 	if (!response.ok) {
-		const message = serverErrorMessage(await boundedError(response), auth.token);
+		const message = serverErrorMessage(await boundedError(response), [auth.token, auth.accountId]);
 		throw new Error(
 			`OpenAI Codex image ${operation} failed with status ${response.status}${message ? `: ${message}` : ""}`,
 		);
@@ -232,7 +242,7 @@ async function requestXaiImage(
 			}
 			return decodeImageResponse(value, "xAI", false);
 		}
-		const message = serverErrorMessage(await boundedError(response), token);
+		const message = serverErrorMessage(await boundedError(response), [token]);
 		if (!retryable(response.status) || attempt === MAX_ATTEMPTS) {
 			throw new Error(
 				`xAI image ${operation} failed with status ${response.status}${message ? `: ${message}` : ""}`,
@@ -243,7 +253,7 @@ async function requestXaiImage(
 	throw new Error(`xAI image ${operation} failed`);
 }
 
-export function generateImage(
+export function requestGeneratedImage(
 	provider: ImageProvider,
 	prompt: string,
 	token: string,