@oh-my-pi/pi-coding-agent 14.2.1 → 14.4.0

package/src/sdk.ts

@@ -128,7 +128,7 @@ import {
 	warmupLspServers,
 } from "./tools";
 import { ToolContextStore } from "./tools/context";
-import { getGeminiImageTools } from "./tools/gemini-image";
+import { getImageGenTools } from "./tools/image-gen";
 import { wrapToolWithMetaNotice } from "./tools/output-meta";
 import { queueResolveHandler } from "./tools/resolve";
 import { EventBus } from "./utils/event-bus";
@@ -194,6 +194,8 @@ export interface CreateAgentSessionOptions {
 
 	/** Enable MCP server discovery from .mcp.json files. Default: true */
 	enableMCP?: boolean;
+	/** Existing MCP manager to reuse (skips discovery, propagates to toolSession). */
+	mcpManager?: MCPManager;
 
 	/** Enable LSP integration (tool, formatting, diagnostics, warmup). Default: true */
 	enableLsp?: boolean;
@@ -207,8 +209,8 @@ export interface CreateAgentSessionOptions {
 
 	/** Output schema for structured completion (subagents) */
 	outputSchema?: unknown;
-	/** Whether to include the submit_result tool by default */
-	requireSubmitResultTool?: boolean;
+	/** Whether to include the yield tool by default */
+	requireYieldTool?: boolean;
 	/** Task recursion depth (for subagent sessions). Default: 0 */
 	taskDepth?: number;
 	/** Parent task ID prefix for nested artifact naming (e.g., "6-Extensions") */
@@ -671,7 +673,12 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 	}
 
 	const imageProvider = settings.get("providers.image");
-	if (imageProvider === "auto" || imageProvider === "gemini" || imageProvider === "openrouter") {
+	if (
+		imageProvider === "auto" ||
+		imageProvider === "openai" ||
+		imageProvider === "gemini" ||
+		imageProvider === "openrouter"
+	) {
 		setPreferredImageProvider(imageProvider);
 	}
 
@@ -914,7 +921,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 			skills,
 			eventBus,
 			outputSchema: options.outputSchema,
-			requireSubmitResultTool: options.requireSubmitResultTool,
+			requireYieldTool: options.requireYieldTool,
 			taskDepth: options.taskDepth ?? 0,
 			getSessionFile: () => sessionManager.getSessionFile() ?? null,
 			getPythonKernelOwnerId: () => pythonKernelOwnerId,
@@ -1005,10 +1012,10 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 		const builtinTools = await logger.time("createAllTools", createTools, toolSession, options.toolNames);
 
 		// Discover MCP tools from .mcp.json files
-		let mcpManager: MCPManager | undefined;
+		let mcpManager: MCPManager | undefined = options.mcpManager;
 		const enableMCP = options.enableMCP ?? true;
 		const customTools: CustomTool[] = [];
-		if (enableMCP) {
+		if (enableMCP && !mcpManager) {
 			const mcpResult = await logger.time("discoverAndLoadMCPTools", discoverAndLoadMCPTools, cwd, {
 				onConnecting: serverNames => {
 					if (options.hasUI && serverNames.length > 0) {
@@ -1024,7 +1031,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 				authStorage,
 			});
 			mcpManager = mcpResult.manager;
-			toolSession.mcpManager = mcpManager;
 
 			if (settings.get("mcp.notifications")) {
 				mcpManager.setNotificationsEnabled(true);
@@ -1044,11 +1050,12 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 				customTools.push(...mcpResult.tools.map(loaded => loaded.tool));
 			}
 		}
+		toolSession.mcpManager = mcpManager;
 
-		// Add Gemini image tools if GEMINI_API_KEY (or GOOGLE_API_KEY) is available
-		const geminiImageTools = await logger.time("getGeminiImageTools", getGeminiImageTools);
-		if (geminiImageTools.length > 0) {
-			customTools.push(...(geminiImageTools as unknown as CustomTool[]));
+		// Add image tools when the active model or configured image providers can generate images.
+		const imageGenTools = await logger.time("getImageGenTools", () => getImageGenTools(modelRegistry, model));
+		if (imageGenTools.length > 0) {
+			customTools.push(...(imageGenTools as unknown as CustomTool[]));
 		}
 
 		// Add web search tools
@@ -1663,8 +1670,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 			}),
 		);
 
-		// Wire MCP manager callbacks to session for reactive tool updates
-		if (mcpManager) {
+		// Wire MCP manager callbacks to session for reactive tool updates.
+		// Skip when reusing a parent's manager — the parent owns the callbacks.
+		if (mcpManager && !options.mcpManager) {
 			mcpManager.setOnToolsChanged(tools => {
 				void session.refreshMCPTools(tools);
 			});

package/src/session/agent-session.ts

@@ -525,6 +525,7 @@ export class AgentSession {
 	#obfuscator: SecretObfuscator | undefined;
 	#checkpointState: CheckpointState | undefined = undefined;
 	#pendingRewindReport: string | undefined = undefined;
+	#lastSuccessfulYieldToolCallId: string | undefined = undefined;
 	#promptGeneration = 0;
 	#providerSessionState = new Map<string, ProviderSessionState>();
 
@@ -789,6 +790,9 @@ export class AgentSession {
 				this.#toolChoiceQueue.resolve();
 			}
 		}
+		if (event.type === "tool_execution_end" && event.toolName === "yield" && !event.isError) {
+			this.#lastSuccessfulYieldToolCallId = event.toolCallId;
+		}
 		if (event.type === "turn_end" && this.#pendingRewindReport) {
 			const report = this.#pendingRewindReport;
 			this.#pendingRewindReport = undefined;
@@ -1026,7 +1030,10 @@ export class AgentSession {
 				.find((message): message is AssistantMessage => message.role === "assistant");
 			const msg = this.#lastAssistantMessage ?? fallbackAssistant;
 			this.#lastAssistantMessage = undefined;
-			if (!msg) return;
+			if (!msg) {
+				this.#lastSuccessfulYieldToolCallId = undefined;
+				return;
+			}
 
 			// Invalidate GitHub Copilot credentials on auth failure so stale tokens
 			// aren't reused on the next request
@@ -1040,9 +1047,16 @@ export class AgentSession {
 
 			if (this.#skipPostTurnMaintenanceAssistantTimestamp === msg.timestamp) {
 				this.#skipPostTurnMaintenanceAssistantTimestamp = undefined;
+				this.#lastSuccessfulYieldToolCallId = undefined;
 				return;
 			}
 
+			if (this.#assistantEndedWithSuccessfulYield(msg)) {
+				this.#lastSuccessfulYieldToolCallId = undefined;
+				return;
+			}
+			this.#lastSuccessfulYieldToolCallId = undefined;
+
 			// Check for retryable errors first (overloaded, rate limit, server errors)
 			if (this.#isRetryableError(msg)) {
 				const didRetry = await this.#handleRetryableError(msg);
@@ -1179,6 +1193,29 @@ export class AgentSession {
 		);
 	}
 
+	#scheduleAutoContinuePrompt(generation: number): void {
+		const continuePrompt = async () => {
+			await this.#promptWithMessage(
+				{
+					role: "developer",
+					content: [{ type: "text", text: "Continue if you have next steps." }],
+					attribution: "agent",
+					timestamp: Date.now(),
+				},
+				"Continue if you have next steps.",
+				{ skipPostPromptRecoveryWait: true },
+			);
+		};
+		this.#schedulePostPromptTask(
+			async signal => {
+				await Promise.resolve();
+				if (signal.aborted) return;
+				await continuePrompt();
+			},
+			{ generation },
+		);
+	}
+
 	#cancelPostPromptTasks(): void {
 		this.#postPromptTasksAbortController.abort();
 		this.#postPromptTasksAbortController = new AbortController();
@@ -3204,7 +3241,6 @@ export class AgentSession {
 				id: task.id,
 				content: task.content,
 				status: task.status,
-				notes: task.notes,
 			})),
 		}));
 	}
@@ -4207,6 +4243,16 @@ export class AgentSession {
 			}
 		}
 	}
+	#assistantEndedWithSuccessfulYield(assistantMessage: AssistantMessage): boolean {
+		const toolCallId = this.#lastSuccessfulYieldToolCallId;
+		if (!toolCallId) return false;
+		const lastToolCall = assistantMessage.content
+			.slice()
+			.reverse()
+			.find((content): content is ToolCall => content.type === "toolCall");
+		return lastToolCall?.name === "yield" && lastToolCall.id === toolCallId;
+	}
+
 	#enforceRewindBeforeYield(): boolean {
 		if (!this.#checkpointState || this.#pendingRewindReport) {
 			return false;
@@ -4813,6 +4859,9 @@ export class AgentSession {
 						aborted: false,
 						willRetry: false,
 					});
+					if (!autoCompactionSignal.aborted && reason !== "idle" && compactionSettings.autoContinue !== false) {
+						this.#scheduleAutoContinuePrompt(generation);
+					}
 					return;
 				}
 			}
@@ -5064,26 +5113,7 @@ export class AgentSession {
 			await this.#emitSessionEvent({ type: "auto_compaction_end", action, result, aborted: false, willRetry });
 
 			if (!willRetry && reason !== "idle" && compactionSettings.autoContinue !== false) {
-				const continuePrompt = async () => {
-					await this.#promptWithMessage(
-						{
-							role: "developer",
-							content: [{ type: "text", text: "Continue if you have next steps." }],
-							attribution: "agent",
-							timestamp: Date.now(),
-						},
-						"Continue if you have next steps.",
-						{ skipPostPromptRecoveryWait: true },
-					);
-				};
-				this.#schedulePostPromptTask(
-					async signal => {
-						await Promise.resolve();
-						if (signal.aborted) return;
-						await continuePrompt();
-					},
-					{ generation },
-				);
+				this.#scheduleAutoContinuePrompt(generation);
 			}
 
 			if (willRetry) {
@@ -5604,6 +5634,14 @@ export class AgentSession {
 	// Bash Execution
 	// =========================================================================
 
+	async #saveBashOriginalArtifact(originalText: string): Promise<string | undefined> {
+		try {
+			return await this.sessionManager.saveArtifact(originalText, "bash-original");
+		} catch {
+			return undefined;
+		}
+	}
+
 	/**
 	 * Execute a bash command.
 	 * Adds result to agent context and session.
@@ -5640,6 +5678,7 @@ export class AgentSession {
 				signal: this.#bashAbortController.signal,
 				sessionKey: this.sessionId,
 				timeout: clampTimeout("bash") * 1000,
+				onMinimizedSave: originalText => this.#saveBashOriginalArtifact(originalText),
 			});
 
 			this.recordBashResult(command, result, options);

package/src/session/session-manager.ts

@@ -1264,74 +1264,245 @@ function extractTextFromContent(content: Message["content"]): string {
 		.join(" ");
 }
 
-async function collectSessionsFromFiles(files: string[], storage: SessionStorage): Promise<SessionInfo[]> {
-	const sessions: SessionInfo[] = [];
+const SESSION_LIST_PREFIX_BYTES = 1024;
+const SESSION_LIST_PARALLEL_THRESHOLD = 64;
+const SESSION_LIST_MAX_WORKERS = 16;
+const sessionListPrefixDecoder = new TextDecoder("utf-8", { fatal: false });
 
-	// Collect session info for all files in parallel
-	await Promise.all(
-		files.map(async file => {
-			try {
-				const content = await storage.readText(file);
-				const entries = parseJsonlLenient<Record<string, unknown>>(content);
-				if (entries.length === 0) return;
-
-				// Check first entry for valid session header
-				type SessionHeaderShape = {
-					type: string;
-					id: string;
-					cwd?: string;
-					title?: string;
-					titleSource?: "auto" | "user";
-					timestamp: string;
-				};
-				const header = entries[0] as SessionHeaderShape;
-				if (header.type !== "session" || !header.id) return;
+async function readSessionListPrefix(file: string, storage: SessionStorage, buffer: Buffer): Promise<string> {
+	if (!(storage instanceof FileSessionStorage)) {
+		return storage.readTextPrefix(file, buffer.byteLength);
+	}
+
+	const handle = await fs.promises.open(file, "r");
+	try {
+		const { bytesRead } = await handle.read(buffer, 0, buffer.byteLength, 0);
+		return sessionListPrefixDecoder.decode(buffer.subarray(0, bytesRead));
+	} finally {
+		await handle.close();
+	}
+}
+
+function decodeJsonStringFragment(value: string): string {
+	const safeValue = value.endsWith("\\") ? value.slice(0, -1) : value;
+	try {
+		return JSON.parse(`"${safeValue}"`) as string;
+	} catch {
+		return safeValue
+			.replace(/\\n/g, "\n")
+			.replace(/\\r/g, "\r")
+			.replace(/\\t/g, "\t")
+			.replace(/\\"/g, '"')
+			.replace(/\\\\/g, "\\");
+	}
+}
 
-				let messageCount = 0;
-				let firstMessage = "";
-				const allMessages: string[] = [];
-				let shortSummary: string | undefined;
+function extractStringProperty(source: string, name: string, startIndex = 0): string | undefined {
+	const propertyIndex = source.indexOf(`"${name}"`, startIndex);
+	if (propertyIndex === -1) return undefined;
 
-				for (let i = 1; i < entries.length; i++) {
-					const entry = entries[i] as { type?: string; message?: Message; shortSummary?: string };
+	const colonIndex = source.indexOf(":", propertyIndex + name.length + 2);
+	if (colonIndex === -1) return undefined;
 
-					if (entry.type === "compaction" && typeof entry.shortSummary === "string") {
-						shortSummary = entry.shortSummary;
-					}
+	let valueIndex = colonIndex + 1;
+	while (valueIndex < source.length) {
+		const char = source.charCodeAt(valueIndex);
+		if (char !== 32 && char !== 9 && char !== 10 && char !== 13) break;
+		valueIndex++;
+	}
+	if (source.charCodeAt(valueIndex) !== 34) return undefined;
+
+	const valueStart = valueIndex + 1;
+	let escaped = false;
+	for (let i = valueStart; i < source.length; i++) {
+		const char = source.charCodeAt(i);
+		if (escaped) {
+			escaped = false;
+			continue;
+		}
+		if (char === 92) {
+			escaped = true;
+			continue;
+		}
+		if (char === 34) {
+			return decodeJsonStringFragment(source.slice(valueStart, i));
+		}
+	}
+
+	return decodeJsonStringFragment(source.slice(valueStart));
+}
+
+function countMessageMarkers(content: string): number {
+	let count = 0;
+	let index = 0;
+	while (index < content.length) {
+		const typeIndex = content.indexOf('"type"', index);
+		if (typeIndex === -1) break;
+		const colonIndex = content.indexOf(":", typeIndex + 6);
+		if (colonIndex === -1) break;
+		const type = extractStringProperty(content, "type", typeIndex);
+		if (type === "message") count++;
+		index = colonIndex + 1;
+	}
+	return count;
+}
+
+function extractFirstUserMessageFromPrefix(content: string): string | undefined {
+	const roleIndex = content.indexOf('"role"');
+	if (roleIndex === -1) return undefined;
+
+	let index = roleIndex;
+	while (index !== -1) {
+		const role = extractStringProperty(content, "role", index);
+		if (role === "user") {
+			return extractStringProperty(content, "content", index) ?? extractStringProperty(content, "text", index);
+		}
+		index = content.indexOf('"role"', index + 6);
+	}
+
+	return undefined;
+}
+
+interface SessionListHeader {
+	type: "session";
+	id: string;
+	cwd?: string;
+	title?: string;
+	parentSession?: string;
+	timestamp?: string;
+}
 
-					if (entry.type === "message" && entry.message) {
-						messageCount++;
+function parseSessionListHeader(
+	content: string,
+	entries: Array<Record<string, unknown>>,
+): SessionListHeader | undefined {
+	const parsedHeader = entries[0];
+	if (parsedHeader?.type === "session" && typeof parsedHeader.id === "string") {
+		return {
+			type: "session",
+			id: parsedHeader.id,
+			cwd: typeof parsedHeader.cwd === "string" ? parsedHeader.cwd : undefined,
+			title: typeof parsedHeader.title === "string" ? parsedHeader.title : undefined,
+			parentSession: typeof parsedHeader.parentSession === "string" ? parsedHeader.parentSession : undefined,
+			timestamp: typeof parsedHeader.timestamp === "string" ? parsedHeader.timestamp : undefined,
+		};
+	}
+
+	const firstLineEnd = content.indexOf("\n");
+	const firstLine = firstLineEnd === -1 ? content : content.slice(0, firstLineEnd);
+	if (extractStringProperty(firstLine, "type") !== "session") return undefined;
+
+	const id = extractStringProperty(firstLine, "id");
+	if (!id) return undefined;
+
+	return {
+		type: "session",
+		id,
+		cwd: extractStringProperty(firstLine, "cwd"),
+		title: extractStringProperty(firstLine, "title"),
+		parentSession: extractStringProperty(firstLine, "parentSession"),
+		timestamp: extractStringProperty(firstLine, "timestamp"),
+	};
+}
+
+function getSessionListWorkerCount(fileCount: number): number {
+	if (fileCount <= SESSION_LIST_PARALLEL_THRESHOLD) return 1;
+	return Math.min(
+		SESSION_LIST_MAX_WORKERS,
+		os.availableParallelism(),
+		Math.ceil(fileCount / SESSION_LIST_PARALLEL_THRESHOLD),
+	);
+}
+
+async function collectSessionFromFile(
+	file: string,
+	storage: SessionStorage,
+	buffer: Buffer,
+): Promise<SessionInfo | undefined> {
+	try {
+		const content = await readSessionListPrefix(file, storage, buffer);
+		const entries = parseJsonlLenient<Record<string, unknown>>(content);
+		const header = parseSessionListHeader(content, entries);
+		if (!header) return undefined;
+
+		let parsedMessageCount = 0;
+		let firstMessage = "";
+		const allMessages: string[] = [];
+		let shortSummary: string | undefined;
+
+		for (let i = 1; i < entries.length; i++) {
+			const entry = entries[i] as { type?: string; message?: Message; shortSummary?: string };
+
+			if (entry.type === "compaction" && typeof entry.shortSummary === "string") {
+				shortSummary = entry.shortSummary;
+			}
 
-						if (entry.message.role === "user" || entry.message.role === "assistant") {
-							const textContent = extractTextFromContent(entry.message.content);
+			if (entry.type === "message" && entry.message) {
+				parsedMessageCount++;
 
-							if (textContent) {
-								allMessages.push(textContent);
+				if (entry.message.role === "user" || entry.message.role === "assistant") {
+					const textContent = extractTextFromContent(entry.message.content);
 
-								if (!firstMessage && entry.message.role === "user") {
-									firstMessage = textContent;
-								}
-							}
+					if (textContent) {
+						allMessages.push(textContent);
+
+						if (!firstMessage && entry.message.role === "user") {
+							firstMessage = textContent;
 						}
 					}
 				}
+			}
+		}
 
-				const stats = storage.statSync(file);
-				sessions.push({
-					path: file,
-					id: header.id,
-					cwd: typeof header.cwd === "string" ? header.cwd : "",
-					title: header.title ?? shortSummary,
-					parentSessionPath: (header as SessionHeader).parentSession,
-					created: new Date(header.timestamp),
-					modified: stats.mtime,
-					messageCount,
-					firstMessage: firstMessage || "(no messages)",
-					allMessagesText: allMessages.join(" "),
-				});
-			} catch {}
-		}),
-	);
+		firstMessage ||= extractFirstUserMessageFromPrefix(content) ?? "";
+		const messageCount = Math.max(parsedMessageCount, countMessageMarkers(content));
+		const stats = storage.statSync(file);
+		return {
+			path: file,
+			id: header.id,
+			cwd: header.cwd ?? "",
+			title: header.title ?? shortSummary,
+			parentSessionPath: header.parentSession,
+			created: new Date(header.timestamp ?? ""),
+			modified: stats.mtime,
+			messageCount,
+			firstMessage: firstMessage || "(no messages)",
+			allMessagesText: allMessages.length > 0 ? allMessages.join(" ") : firstMessage,
+		};
+	} catch {
+		return undefined;
+	}
+}
+
+async function collectSessionsFromFileStride(
+	files: string[],
+	storage: SessionStorage,
+	startIndex: number,
+	stride: number,
+): Promise<SessionInfo[]> {
+	const sessions: SessionInfo[] = [];
+	const buffer = Buffer.allocUnsafe(SESSION_LIST_PREFIX_BYTES);
+
+	for (let i = startIndex; i < files.length; i += stride) {
+		const session = await collectSessionFromFile(files[i], storage, buffer);
+		if (session) sessions.push(session);
+	}
+
+	return sessions;
+}
+
+async function collectSessionsFromFiles(files: string[], storage: SessionStorage): Promise<SessionInfo[]> {
+	const workerCount = getSessionListWorkerCount(files.length);
+	const sessions =
+		workerCount === 1
+			? await collectSessionsFromFileStride(files, storage, 0, 1)
+			: (
+					await Promise.all(
+						Array.from({ length: workerCount }, (_, workerIndex) =>
+							collectSessionsFromFileStride(files, storage, workerIndex, workerCount),
+						),
+					)
+				).flat();
 
 	sessions.sort((a, b) => b.modified.getTime() - a.modified.getTime());
 	return sessions;
@@ -2771,7 +2942,7 @@ export class SessionManager {
 	static async listAll(storage: SessionStorage = new FileSessionStorage()): Promise<SessionInfo[]> {
 		const sessionsRoot = path.join(getDefaultAgentDir(), "sessions");
 		try {
-			const files = Array.from(new Bun.Glob("**/*.jsonl").scanSync(sessionsRoot)).map(name =>
+			const files = await Array.fromAsync(new Bun.Glob("*/*.jsonl").scan(sessionsRoot), name =>
 				path.join(sessionsRoot, name),
 			);
 			return await collectSessionsFromFiles(files, storage);

package/src/session/streaming-output.ts

@@ -680,6 +680,17 @@ export class OutputSink {
 		});
 	}
 
+	/**
+	 * Replace the in-memory buffer with the given text while preserving the
+	 * streaming counters (totalLines/totalBytes reflect the raw chunks that
+	 * already reached the sink). Used when an upstream minimizer rewrites the
+	 * captured output after the raw bytes have already been streamed.
+	 */
+	replace(text: string): void {
+		this.#buffer = text;
+		this.#bufferBytes = Buffer.byteLength(text, "utf-8");
+	}
+
 	async dump(notice?: string): Promise<OutputSummary> {
 		const noticeLine = notice ? `[${notice}]\n` : "";
 		const outputLines = this.#buffer.length > 0 ? countNewlines(this.#buffer) + 1 : 0;

package/src/system-prompt.ts

@@ -275,13 +275,18 @@ async function getCachedGpu(): Promise<string | undefined> {
 }
 async function getEnvironmentInfo(): Promise<Array<{ label: string; value: string }>> {
 	const gpu = await getCachedGpu();
-	const cpus = os.cpus();
+	let cpuModel: string | undefined;
+	try {
+		cpuModel = os.cpus()[0]?.model;
+	} catch {
+		cpuModel = undefined;
+	}
 	const entries: Array<{ label: string; value: string | undefined }> = [
 		{ label: "OS", value: `${os.platform()} ${os.release()}` },
 		{ label: "Distro", value: os.type() },
 		{ label: "Kernel", value: os.version() },
 		{ label: "Arch", value: os.arch() },
-		{ label: "CPU", value: `${cpus[0]?.model}` },
+		{ label: "CPU", value: cpuModel },
 		{ label: "GPU", value: gpu },
 		{ label: "Terminal", value: getTerminalName() },
 	];