@oh-my-pi/pi-coding-agent 13.14.0 → 13.15.2

package/src/session/agent-session.ts

@@ -46,6 +46,7 @@ import {
 	calculateRateLimitBackoffMs,
 	getSupportedEfforts,
 	isContextOverflow,
+	isUsageLimitError,
 	modelsAreEqual,
 	parseRateLimitReason,
 } from "@oh-my-pi/pi-ai";
@@ -224,6 +225,8 @@ export interface AgentSessionConfig {
 	mcpDiscoveryEnabled?: boolean;
 	/** MCP tool names to activate for the current session when discovery mode is enabled. */
 	initialSelectedMCPToolNames?: string[];
+	/** Whether constructor-provided MCP defaults should be persisted immediately. */
+	persistInitialMCPToolSelection?: boolean;
 	/** MCP server names whose tools should seed discovery-mode sessions whenever those servers are connected. */
 	defaultSelectedMCPServerNames?: string[];
 	/** MCP tool names that should seed brand-new sessions created from this AgentSession. */
@@ -363,6 +366,7 @@ export class AgentSession {
 	#followUpMessages: string[] = [];
 	/** Messages queued to be included with the next user prompt as context ("asides"). */
 	#pendingNextTurnMessages: CustomMessage[] = [];
+	#scheduledHiddenNextTurnGeneration: number | undefined = undefined;
 	#planModeState: PlanModeState | undefined;
 	#planReferenceSent = false;
 	#planReferencePath = "local://PLAN.md";
@@ -483,8 +487,11 @@ export class AgentSession {
 		this.#pruneSelectedMCPToolNames();
 		const persistedSelectedMCPToolNames = this.sessionManager.buildSessionContext().selectedMCPToolNames;
 		const currentSelectedMCPToolNames = this.getSelectedMCPToolNames();
+		const persistInitialMCPToolSelection =
+			config.persistInitialMCPToolSelection ?? this.sessionManager.getBranch().length === 0;
 		if (
 			this.#mcpDiscoveryEnabled &&
+			persistInitialMCPToolSelection &&
 			!this.#selectedMCPToolNamesMatch(persistedSelectedMCPToolNames, currentSelectedMCPToolNames)
 		) {
 			this.sessionManager.appendMCPToolSelection(currentSelectedMCPToolNames);
@@ -781,7 +788,6 @@ export class AgentSession {
 						attempt: this.#retryAttempt,
 					});
 					this.#retryAttempt = 0;
-					this.#resolveRetry();
 				}
 			}
 
@@ -857,6 +863,7 @@ export class AgentSession {
 				const didRetry = await this.#handleRetryableError(msg);
 				if (didRetry) return; // Retry was initiated, don't proceed to compaction
 			}
+			this.#resolveRetry();
 
 			if (msg.stopReason === "aborted" && this.#checkpointState) {
 				this.#checkpointState = undefined;
@@ -2566,6 +2573,74 @@ export class AgentSession {
 		});
 	}
 
+	#queueHiddenNextTurnMessage(message: CustomMessage, triggerTurn: boolean): void {
+		this.#pendingNextTurnMessages.push(message);
+		if (!triggerTurn) return;
+		const generation = this.#promptGeneration;
+		if (this.#scheduledHiddenNextTurnGeneration === generation) {
+			return;
+		}
+		this.#scheduledHiddenNextTurnGeneration = generation;
+		this.#schedulePostPromptTask(
+			async () => {
+				if (this.#scheduledHiddenNextTurnGeneration === generation) {
+					this.#scheduledHiddenNextTurnGeneration = undefined;
+				}
+				if (this.#pendingNextTurnMessages.length === 0) {
+					return;
+				}
+				try {
+					await this.#promptQueuedHiddenNextTurnMessages();
+				} catch {
+					// Leave the hidden next-turn messages queued for the next explicit prompt.
+				}
+			},
+			{
+				generation,
+				onSkip: () => {
+					if (this.#scheduledHiddenNextTurnGeneration === generation) {
+						this.#scheduledHiddenNextTurnGeneration = undefined;
+					}
+				},
+			},
+		);
+	}
+
+	async #promptQueuedHiddenNextTurnMessages(): Promise<void> {
+		if (this.#pendingNextTurnMessages.length === 0) {
+			return;
+		}
+
+		const queuedMessages = [...this.#pendingNextTurnMessages];
+		this.#pendingNextTurnMessages = [];
+		const message = queuedMessages[queuedMessages.length - 1];
+		if (!message) {
+			return;
+		}
+
+		const prependMessages = queuedMessages.slice(0, -1);
+		const textContent = this.#getCustomMessageTextContent(message);
+		try {
+			await this.#promptWithMessage(message, textContent, {
+				prependMessages,
+				skipPostPromptRecoveryWait: true,
+			});
+		} catch (error) {
+			this.#pendingNextTurnMessages = [...queuedMessages, ...this.#pendingNextTurnMessages];
+			throw error;
+		}
+	}
+
+	#getCustomMessageTextContent(message: Pick<CustomMessage, "content">): string {
+		if (typeof message.content === "string") {
+			return message.content;
+		}
+		return message.content
+			.filter((content): content is TextContent => content.type === "text")
+			.map(content => content.text)
+			.join("");
+	}
+
 	/**
 	 * Throw an error if the text is an extension command.
 	 */
@@ -2606,7 +2681,7 @@ export class AgentSession {
 		};
 		if (this.isStreaming) {
 			if (options?.deliverAs === "nextTurn") {
-				this.#pendingNextTurnMessages.push(appMessage);
+				this.#queueHiddenNextTurnMessage(appMessage, options?.triggerTurn ?? false);
 				return;
 			}
 
@@ -2618,6 +2693,22 @@ export class AgentSession {
 			return;
 		}
 
+		if (options?.deliverAs === "nextTurn") {
+			if (options?.triggerTurn) {
+				await this.agent.prompt(appMessage);
+				return;
+			}
+			this.agent.appendMessage(appMessage);
+			this.sessionManager.appendCustomMessageEntry(
+				message.customType,
+				message.content,
+				message.display,
+				message.details,
+				message.attribution ?? "agent",
+			);
+			return;
+		}
+
 		if (options?.triggerTurn) {
 			await this.agent.prompt(appMessage);
 			return;
@@ -2685,9 +2776,9 @@ export class AgentSession {
 		return { steering, followUp };
 	}
 
-	/** Number of pending messages (includes both steering and follow-up) */
+	/** Number of pending messages (includes steering, follow-up, and next-turn messages) */
 	get queuedMessageCount(): number {
-		return this.#steeringMessages.length + this.#followUpMessages.length;
+		return this.#steeringMessages.length + this.#followUpMessages.length + this.#pendingNextTurnMessages.length;
 	}
 
 	/** Get pending messages (read-only) */
@@ -2829,6 +2920,7 @@ export class AgentSession {
 	async abort(): Promise<void> {
 		this.abortRetry();
 		this.#promptGeneration++;
+		this.#scheduledHiddenNextTurnGeneration = undefined;
 		this.#resolveTtsrResume();
 		this.#cancelPostPromptTasks();
 		this.agent.abort();
@@ -2878,6 +2970,7 @@ export class AgentSession {
 		this.#steeringMessages = [];
 		this.#followUpMessages = [];
 		this.#pendingNextTurnMessages = [];
+		this.#scheduledHiddenNextTurnGeneration = undefined;
 
 		this.sessionManager.appendThinkingLevelChange(this.thinkingLevel);
 		this.sessionManager.appendServiceTierChange(this.serviceTier ?? null);
@@ -3611,6 +3704,7 @@ export class AgentSession {
 			this.#steeringMessages = [];
 			this.#followUpMessages = [];
 			this.#pendingNextTurnMessages = [];
+			this.#scheduledHiddenNextTurnGeneration = undefined;
 			this.#todoReminderCount = 0;
 
 			// Inject the handoff document as a custom message
@@ -4278,7 +4372,9 @@ export class AgentSession {
 							const shouldRetry =
 								retrySettings.enabled &&
 								attempt < retrySettings.maxRetries &&
-								(retryAfterMs !== undefined || this.#isRetryableErrorMessage(message));
+								(retryAfterMs !== undefined ||
+									this.#isTransientErrorMessage(message) ||
+									isUsageLimitError(message));
 							if (!shouldRetry) {
 								lastError = error;
 								break;
@@ -4476,8 +4572,9 @@ export class AgentSession {
 	// =========================================================================
 
 	/**
-	 * Check if an error is retryable (overloaded, rate limit, server errors).
+	 * Check if an error is retryable (transient errors or usage limits).
 	 * Context overflow errors are NOT retryable (handled by compaction instead).
+	 * Usage-limit errors are retryable because the retry handler performs credential switching.
 	 */
 	#isRetryableError(message: AssistantMessage): boolean {
 		if (message.stopReason !== "error" || !message.errorMessage) return false;
@@ -4487,20 +4584,17 @@ export class AgentSession {
 		if (isContextOverflow(message, contextWindow)) return false;
 
 		const err = message.errorMessage;
-		return this.#isRetryableErrorMessage(err);
+		return this.#isTransientErrorMessage(err) || isUsageLimitError(err);
 	}
 
-	#isRetryableErrorMessage(errorMessage: string): boolean {
-		// Match: overloaded_error, rate limit, usage limit, 429, 500, 502, 503, 504, service unavailable, connection error, fetch failed, retry delay exceeded, stream stall
-		return /overloaded|rate.?limit|usage.?limit|too many requests|429|500|502|503|504|service.?unavailable|server error|internal error|connection.?error|unable to connect|fetch failed|retry delay|stream stall/i.test(
+	#isTransientErrorMessage(errorMessage: string): boolean {
+		// Match: overloaded_error, provider returned error, rate limit, 429, 500, 502, 503, 504,
+		// service unavailable, network/connection errors, fetch failed, terminated, retry delay exceeded
+		return /overloaded|provider.?returned.?error|rate.?limit|too many requests|429|500|502|503|504|service.?unavailable|server.?error|internal.?error|network.?error|connection.?error|connection.?refused|other side closed|fetch failed|upstream.?connect|reset before headers|socket hang up|timed? out|timeout|terminated|retry delay|stream stall/i.test(
 			errorMessage,
 		);
 	}
 
-	#isUsageLimitErrorMessage(errorMessage: string): boolean {
-		return /usage.?limit|usage_limit_reached|limit_reached|quota.?exceeded|resource.?exhausted/i.test(errorMessage);
-	}
-
 	#parseRetryAfterMsFromError(errorMessage: string): number | undefined {
 		const now = Date.now();
 		const retryAfterMsMatch = /retry-after-ms\s*[:=]\s*(\d+)/i.exec(errorMessage);
@@ -4582,7 +4676,7 @@ export class AgentSession {
 		const errorMessage = message.errorMessage || "Unknown error";
 		let delayMs = retrySettings.baseDelayMs * 2 ** (this.#retryAttempt - 1);
 
-		if (this.model && this.#isUsageLimitErrorMessage(errorMessage)) {
+		if (this.model && isUsageLimitError(errorMessage)) {
 			const retryAfterMs =
 				this.#parseRetryAfterMsFromError(errorMessage) ??
 				calculateRateLimitBackoffMs(parseRateLimitReason(errorMessage));
@@ -4960,6 +5054,7 @@ export class AgentSession {
 		this.#steeringMessages = [];
 		this.#followUpMessages = [];
 		this.#pendingNextTurnMessages = [];
+		this.#scheduledHiddenNextTurnGeneration = undefined;
 
 		// Flush pending writes before switching
 		await this.sessionManager.flush();
@@ -5003,21 +5098,18 @@ export class AgentSession {
 		const hasThinkingEntry = this.sessionManager.getBranch().some(entry => entry.type === "thinking_level_change");
 		const hasServiceTierEntry = this.sessionManager.getBranch().some(entry => entry.type === "service_tier_change");
 		const defaultThinkingLevel = this.settings.get("defaultThinkingLevel");
-
-		if (hasThinkingEntry) {
-			this.setThinkingLevel(sessionContext.thinkingLevel as ThinkingLevel | undefined);
-		} else {
-			const effectiveDefaultThinkingLevel = resolveThinkingLevelForModel(this.model, defaultThinkingLevel);
-			this.#thinkingLevel = effectiveDefaultThinkingLevel;
-			this.agent.setThinkingLevel(toReasoningEffort(effectiveDefaultThinkingLevel));
-			this.sessionManager.appendThinkingLevelChange(effectiveDefaultThinkingLevel);
-		}
-
-		if (hasServiceTierEntry) {
-			this.agent.serviceTier = sessionContext.serviceTier;
-		} else {
-			this.sessionManager.appendServiceTierChange(this.serviceTier ?? null);
-		}
+		const configuredServiceTier = this.settings.get("serviceTier");
+		const nextThinkingLevel = resolveThinkingLevelForModel(
+			this.model,
+			hasThinkingEntry ? (sessionContext.thinkingLevel as ThinkingLevel | undefined) : defaultThinkingLevel,
+		);
+		this.#thinkingLevel = nextThinkingLevel;
+		this.agent.setThinkingLevel(toReasoningEffort(nextThinkingLevel));
+		this.agent.serviceTier = hasServiceTierEntry
+			? sessionContext.serviceTier
+			: configuredServiceTier === "none"
+				? undefined
+				: configuredServiceTier;
 
 		this.#reconnectToAgent();
 		return true;
@@ -5059,6 +5151,7 @@ export class AgentSession {
 
 		// Clear pending messages (bound to old session state)
 		this.#pendingNextTurnMessages = [];
+		this.#scheduledHiddenNextTurnGeneration = undefined;
 
 		// Flush pending writes before branching
 		await this.sessionManager.flush();

package/src/session/session-manager.ts

@@ -1302,21 +1302,19 @@ async function collectSessionsFromFiles(files: string[], storage: SessionStorage
 					}
 				}
 
-				if (messageCount) {
-					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(" "),
-					});
-				}
+				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 {}
 		}),
 	);
@@ -1381,6 +1379,7 @@ export class SessionManager {
 	#sessionName: string | undefined;
 	#sessionFile: string | undefined;
 	#flushed: boolean = false;
+	#needsFullRewriteOnNextPersist: boolean = false;
 	#fileEntries: FileEntry[] = [];
 	#byId: Map<string, SessionEntry> = new Map();
 	#labelsById: Map<string, string> = new Map();
@@ -1443,9 +1442,7 @@ export class SessionManager {
 			this.#sessionId = header?.id ?? Snowflake.next();
 			this.#sessionName = header?.title;
 
-			if (migrateToCurrentVersion(this.#fileEntries)) {
-				await this.#rewriteFile();
-			}
+			this.#needsFullRewriteOnNextPersist = migrateToCurrentVersion(this.#fileEntries);
 
 			await resolveBlobRefsInEntries(this.#fileEntries, this.#blobStore);
 
@@ -1632,6 +1629,7 @@ export class SessionManager {
 		this.#labelsById.clear();
 		this.#leafId = null;
 		this.#flushed = false;
+		this.#needsFullRewriteOnNextPersist = false;
 		this.#usageStatistics = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, premiumRequests: 0, cost: 0 };
 
 		if (this.persist) {
@@ -1774,6 +1772,7 @@ export class SessionManager {
 				this.#fileEntries.map(entry => prepareEntryForPersistence(entry, this.#blobStore)),
 			);
 			await this.#writeEntriesAtomically(entries);
+			this.#needsFullRewriteOnNextPersist = false;
 			this.#flushed = true;
 		});
 	}
@@ -1782,6 +1781,16 @@ export class SessionManager {
 		return this.persist;
 	}
 
+	/**
+	 * Force-persist all current entries to disk, even when no assistant message exists yet.
+	 * Used by ACP mode where session/new must create a discoverable session immediately.
+	 */
+	async ensureOnDisk(): Promise<void> {
+		if (!this.persist || !this.#sessionFile) return;
+		if (this.#flushed && !this.#needsFullRewriteOnNextPersist) return;
+		await this.#rewriteFile();
+	}
+
 	/** Flush pending writes to disk. Call before switching sessions or on shutdown. */
 	async flush(): Promise<void> {
 		await this.#queuePersistTask(async () => {
@@ -1911,23 +1920,15 @@ export class SessionManager {
 
 		const hasAssistant = this.#fileEntries.some(e => e.type === "message" && e.message.role === "assistant");
 		if (!hasAssistant) {
-			// Mark as not flushed so when assistant arrives, all entries get written
+			// Mark as not flushed so when assistant arrives, all entries get written.
 			this.#flushed = false;
 			return;
 		}
 
-		if (!this.#flushed) {
-			this.#flushed = true;
-			void this.#queuePersistTask(async () => {
-				const writer = this.#ensurePersistWriter();
-				if (!writer) return;
-				const entries = await Promise.all(
-					this.#fileEntries.map(e => prepareEntryForPersistence(e, this.#blobStore)),
-				);
-				for (const persistedEntry of entries) {
-					await writer.write(persistedEntry);
-				}
-			});
+		if (this.#needsFullRewriteOnNextPersist || !this.#flushed) {
+			// Full flush: rewrite the entire file atomically to avoid
+			// duplicating entries if the file already exists (e.g. from ensureOnDisk).
+			void this.#rewriteFile();
 		} else {
 			void this.#queuePersistTask(async () => {
 				const writer = this.#ensurePersistWriter();

package/src/session/streaming-output.ts

@@ -32,6 +32,8 @@ export interface OutputSinkOptions {
 	artifactId?: string;
 	spillThreshold?: number;
 	onChunk?: (chunk: string) => void;
+	/** Minimum ms between onChunk calls. 0 = every chunk (default). */
+	chunkThrottleMs?: number;
 }
 
 export interface TruncationResult {
@@ -521,6 +523,7 @@ export class OutputSink {
 	#totalBytes = 0;
 	#sawData = false;
 	#truncated = false;
+	#lastChunkTime = 0;
 
 	#file?: {
 		path: string;
@@ -528,22 +531,46 @@ export class OutputSink {
 		sink: Bun.FileSink;
 	};
 
+	// Queue of chunks waiting for the file sink to be created.
+	#pendingFileWrites?: string[];
+	#fileReady = false;
+
 	readonly #artifactPath?: string;
 	readonly #artifactId?: string;
 	readonly #spillThreshold: number;
 	readonly #onChunk?: (chunk: string) => void;
+	readonly #chunkThrottleMs: number;
 
 	constructor(options?: OutputSinkOptions) {
-		const { artifactPath, artifactId, spillThreshold = DEFAULT_MAX_BYTES, onChunk } = options ?? {};
+		const {
+			artifactPath,
+			artifactId,
+			spillThreshold = DEFAULT_MAX_BYTES,
+			onChunk,
+			chunkThrottleMs = 0,
+		} = options ?? {};
 		this.#artifactPath = artifactPath;
 		this.#artifactId = artifactId;
 		this.#spillThreshold = spillThreshold;
 		this.#onChunk = onChunk;
+		this.#chunkThrottleMs = chunkThrottleMs;
 	}
 
-	async push(chunk: string): Promise<void> {
+	/**
+	 * Push a chunk of output. The buffer management and onChunk callback run
+	 * synchronously. File sink writes are deferred and serialized internally.
+	 */
+	push(chunk: string): void {
 		chunk = sanitizeWithOptionalSixelPassthrough(chunk, sanitizeText);
-		this.#onChunk?.(chunk);
+
+		// Throttled onChunk: only call the callback when enough time has passed.
+		if (this.#onChunk) {
+			const now = Date.now();
+			if (now - this.#lastChunkTime >= this.#chunkThrottleMs) {
+				this.#lastChunkTime = now;
+				this.#onChunk(chunk);
+			}
+		}
 
 		const dataBytes = Buffer.byteLength(chunk, "utf-8");
 		this.#totalBytes += dataBytes;
@@ -556,10 +583,9 @@ export class OutputSink {
 		const threshold = this.#spillThreshold;
 		const willOverflow = this.#bufferBytes + dataBytes > threshold;
 
-		// Write to file if already spilling or about to overflow
-		if (this.#file != null || willOverflow) {
-			const sink = await this.#ensureFileSink();
-			await sink?.write(chunk);
+		// Write to artifact file if configured and past the threshold
+		if (this.#artifactPath && (this.#file != null || willOverflow)) {
+			this.#writeToFile(chunk);
 		}
 
 		if (!willOverflow) {
@@ -589,14 +615,64 @@ export class OutputSink {
 		if (this.#file) this.#truncated = true;
 	}
 
+	/**
+	 * Write a chunk to the artifact file. Handles the async file sink creation
+	 * by queuing writes until the sink is ready, then draining synchronously.
+	 */
+	#writeToFile(chunk: string): void {
+		if (this.#fileReady && this.#file) {
+			// Fast path: file sink exists, write synchronously
+			this.#file.sink.write(chunk);
+			return;
+		}
+		// File sink not yet created — queue this chunk and kick off creation
+		if (!this.#pendingFileWrites) {
+			this.#pendingFileWrites = [chunk];
+			void this.#createFileSink();
+		} else {
+			this.#pendingFileWrites.push(chunk);
+		}
+	}
+
+	async #createFileSink(): Promise<void> {
+		if (!this.#artifactPath || this.#fileReady) return;
+		try {
+			const sink = Bun.file(this.#artifactPath).writer();
+			this.#file = { path: this.#artifactPath, artifactId: this.#artifactId, sink };
+
+			// Flush existing buffer to file BEFORE it gets trimmed further.
+			if (this.#buffer.length > 0) {
+				sink.write(this.#buffer);
+			}
+
+			// Drain any chunks that arrived while the sink was being created
+			if (this.#pendingFileWrites) {
+				for (const pending of this.#pendingFileWrites) {
+					sink.write(pending);
+				}
+				this.#pendingFileWrites = undefined;
+			}
+
+			this.#fileReady = true;
+		} catch {
+			try {
+				await this.#file?.sink?.end();
+			} catch {
+				/* ignore */
+			}
+			this.#file = undefined;
+			this.#pendingFileWrites = undefined;
+		}
+	}
+
 	createInput(): WritableStream<Uint8Array | string> {
 		const dec = new TextDecoder("utf-8", { ignoreBOM: true });
-		const finalize = async () => {
-			await this.push(dec.decode());
+		const finalize = () => {
+			this.push(dec.decode());
 		};
 		return new WritableStream({
-			write: async chunk => {
-				await this.push(typeof chunk === "string" ? chunk : dec.decode(chunk, { stream: true }));
+			write: chunk => {
+				this.push(typeof chunk === "string" ? chunk : dec.decode(chunk, { stream: true }));
 			},
 			close: finalize,
 			abort: finalize,
@@ -620,32 +696,6 @@ export class OutputSink {
 			artifactId: this.#file?.artifactId,
 		};
 	}
-
-	// -- private ---------------------------------------------------------------
-
-	async #ensureFileSink(): Promise<Bun.FileSink | null> {
-		if (!this.#artifactPath) return null;
-		if (this.#file) return this.#file.sink;
-
-		try {
-			const sink = Bun.file(this.#artifactPath).writer();
-			this.#file = { path: this.#artifactPath, artifactId: this.#artifactId, sink };
-
-			// Flush existing buffer to file BEFORE it gets trimmed further.
-			if (this.#buffer.length > 0) {
-				await sink.write(this.#buffer);
-			}
-			return sink;
-		} catch {
-			try {
-				await this.#file?.sink?.end();
-			} catch {
-				/* ignore */
-			}
-			this.#file = undefined;
-			return null;
-		}
-	}
 }
 
 // =============================================================================