@oh-my-pi/pi-coding-agent 17.1.8 → 17.2.0

package/src/mcp/manager.ts

@@ -94,6 +94,14 @@ const STARTUP_TIMEOUT_MS = 250;
 const RECONNECT_BURST_WINDOW_MS = 30_000;
 const RECONNECT_BURST_LIMIT = 5;
 
+/**
+ * Bounded buffer for notifications received before any listener attaches.
+ * Mirrors {@link IrcBus}'s `MAILBOX_CAP` — drop-oldest on overflow. Drained
+ * into the first {@link MCPManager.addNotificationListener} subscriber, then
+ * cleared; subsequent frames deliver directly to attached listeners.
+ */
+const NOTIFICATION_BUFFER_CAP = 100;
+
 function trackPromise<T>(promise: Promise<T>): TrackedPromise<T> {
 	const tracked: TrackedPromise<T> = { promise, status: "pending" };
 	promise.then(
@@ -194,8 +202,14 @@ export class MCPManager {
 	#sources = new Map<string, SourceMeta>();
 	#authStorage: AuthStorage | null = null;
 	#authHandler?: MCPAuthHandler;
-	#onNotification?: (serverName: string, method: string, params: unknown) => void;
-	#onToolsChanged?: (tools: CustomTool<TSchema, MCPToolDetails>[]) => void;
+	#notificationListeners = new Set<(serverName: string, method: string, params: unknown) => void>();
+	/**
+	 * Notifications received before any listener attached, to be drained on
+	 * the first {@link addNotificationListener} call. Bounded by
+	 * {@link NOTIFICATION_BUFFER_CAP}, drop-oldest on overflow.
+	 */
+	#pendingNotifications: Array<{ server: string; method: string; params: unknown }> = [];
+	#onToolsChanged?: (tools: CustomTool<TSchema, MCPToolDetails>[]) => void | Promise<void>;
 	#onResourcesChanged?: (serverName: string, uri: string) => void;
 	#onPromptsChanged?: (serverName: string) => void;
 	#notificationsEnabled = false;
@@ -219,16 +233,66 @@ export class MCPManager {
 	) {}
 
 	/**
-	 * Set a callback to receive all server notifications.
+	 * Register a listener for server-initiated MCP notifications.
+	 *
+	 * The listener is called for every JSON-RPC notification received from any
+	 * connected server, AFTER the manager's own handling of known methods
+	 * (`notifications/tools/list_changed`, `notifications/resources/list_changed`,
+	 * `notifications/resources/updated`, `notifications/prompts/list_changed`).
+	 * For list-change methods the internal refresh promise is awaited before
+	 * fanout, so listeners observe up-to-date manager and tool state. Unknown
+	 * or server-custom methods are also delivered, letting consumers bridge
+	 * server-initiated events into session-level behavior (e.g. an extension
+	 * injecting a steer via `pi.sendMessage`).
+	 *
+	 * Notifications received before any listener attached are buffered
+	 * (bounded FIFO, cap {@link NOTIFICATION_BUFFER_CAP}, drop-oldest) and
+	 * drained into the first subscriber — matches {@link setOnPromptsChanged}'s
+	 * replay-on-attach and {@link IrcBus}'s mailbox semantics.
+	 *
+	 * Returns an unsubscribe function; call it to remove the listener.
+	 *
+	 * Multiple listeners are allowed; each is invoked with independent error
+	 * isolation — a listener that throws does not prevent other listeners from
+	 * firing.
 	 */
-	setOnNotification(handler: (serverName: string, method: string, params: unknown) => void): void {
-		this.#onNotification = handler;
+	addNotificationListener(listener: (serverName: string, method: string, params: unknown) => void): () => void {
+		const wasEmpty = this.#notificationListeners.size === 0;
+		this.#notificationListeners.add(listener);
+
+		// Drain startup-buffered notifications into the first attaching listener.
+		if (wasEmpty && this.#pendingNotifications.length > 0) {
+			const pending = this.#pendingNotifications.splice(0);
+			for (const frame of pending) {
+				try {
+					listener(frame.server, frame.method, frame.params);
+				} catch (error) {
+					logger.debug("MCP notification listener threw during buffered drain", {
+						path: `mcp:${frame.server}`,
+						method: frame.method,
+						error,
+					});
+				}
+			}
+		}
+
+		return () => {
+			this.#notificationListeners.delete(listener);
+		};
 	}
 
 	/**
 	 * Set a callback to fire when any server's tools change.
+	 *
+	 * May return a Promise; if so, {@link refreshServerTools} awaits it so that
+	 * downstream consumers (e.g. `mcp_notification` listeners for
+	 * `notifications/tools/list_changed`) observe not just the manager's
+	 * refreshed tool set but also any session-level rebind driven by the
+	 * handler (`session.refreshMCPTools`). Other callsites (initial connect,
+	 * disconnect, reconnect) invoke the handler synchronously — their downstream
+	 * chains don't need to serialize on the rebind.
 	 */
-	setOnToolsChanged(handler: (tools: CustomTool<TSchema, MCPToolDetails>[]) => void): void {
+	setOnToolsChanged(handler: (tools: CustomTool<TSchema, MCPToolDetails>[]) => void | Promise<void>): void {
 		this.#onToolsChanged = handler;
 	}
 
@@ -488,7 +552,7 @@ export class MCPManager {
 						this.reconnectServer(name, options);
 					const customTools = MCPTool.fromTools(connection, serverTools, reconnect);
 					this.#replaceServerTools(name, customTools);
-					this.#onToolsChanged?.(this.#tools);
+					void this.#onToolsChanged?.(this.#tools);
 					void this.toolCache?.set(name, config, serverTools);
 
 					onStatus?.({ type: "connected", serverName: name });
@@ -597,7 +661,7 @@ export class MCPManager {
 		sortMCPToolsByName(this.#tools);
 	}
 
-	#triggerNotificationRefresh(serverName: string, kind: "tools" | "resources" | "prompts"): void {
+	#triggerNotificationRefresh(serverName: string, kind: "tools" | "resources" | "prompts"): Promise<void> {
 		const refresh = (() => {
 			switch (kind) {
 				case "tools":
@@ -608,22 +672,33 @@ export class MCPManager {
 					return this.refreshServerPrompts(serverName);
 			}
 		})();
-		void refresh.catch(error => {
+		return refresh.catch(error => {
 			logger.debug("Failed MCP notification refresh", { path: `mcp:${serverName}`, kind, error });
 		});
 	}
-	#handleServerNotification(serverName: string, method: string, params: unknown): void {
+	async #handleServerNotification(serverName: string, method: string, params: unknown): Promise<void> {
 		logger.debug("MCP notification received", { path: `mcp:${serverName}`, method });
 
+		// Only trigger refresh if the connection is already stored — during the
+		// initial connect handshake, notifications may arrive before
+		// `#connections.set()` completes, and `refreshServer*` would no-op
+		// anyway. Skipping the await in that case preserves arrival order
+		// across concurrently-dispatched notifications (an awaited refresh,
+		// even a no-op, yields a microtask that lets later frames overtake).
+		const connectionKnown = this.#connections.has(serverName);
+		let refreshPromise: Promise<void> | undefined;
 		switch (method) {
 			case MCPNotificationMethods.TOOLS_LIST_CHANGED:
-				this.#triggerNotificationRefresh(serverName, "tools");
+				if (connectionKnown) refreshPromise = this.#triggerNotificationRefresh(serverName, "tools");
 				break;
 			case MCPNotificationMethods.RESOURCES_LIST_CHANGED:
-				this.#triggerNotificationRefresh(serverName, "resources");
+				if (connectionKnown) refreshPromise = this.#triggerNotificationRefresh(serverName, "resources");
 				break;
 			case MCPNotificationMethods.RESOURCES_UPDATED: {
-				const uri = (params as { uri?: string })?.uri;
+				const uri =
+					params && typeof params === "object" && "uri" in params && typeof params.uri === "string"
+						? params.uri
+						: undefined;
 				const subscribed = this.#subscribedResources.get(serverName);
 				if (uri && subscribed?.has(uri)) {
 					this.#onResourcesChanged?.(serverName, uri);
@@ -631,13 +706,40 @@ export class MCPManager {
 				break;
 			}
 			case MCPNotificationMethods.PROMPTS_LIST_CHANGED:
-				this.#triggerNotificationRefresh(serverName, "prompts");
+				if (connectionKnown) refreshPromise = this.#triggerNotificationRefresh(serverName, "prompts");
 				break;
 			default:
 				break;
 		}
 
-		this.#onNotification?.(serverName, method, params);
+		// Await internal refresh so listeners see the manager's post-refresh
+		// state (satisfies the documented "AFTER the manager's own handling"
+		// contract on `addNotificationListener` — otherwise an extension acting
+		// on `tools/list_changed` could hit stale `getTools()`).
+		if (refreshPromise) {
+			await refreshPromise;
+		}
+
+		// Buffer for late-attaching subscribers when no listener exists yet.
+		if (this.#notificationListeners.size === 0) {
+			this.#pendingNotifications.push({ server: serverName, method, params });
+			if (this.#pendingNotifications.length > NOTIFICATION_BUFFER_CAP) {
+				this.#pendingNotifications.shift();
+			}
+			return;
+		}
+
+		for (const listener of this.#notificationListeners) {
+			try {
+				listener(serverName, method, params);
+			} catch (error) {
+				logger.debug("MCP notification listener threw", {
+					path: `mcp:${serverName}`,
+					method,
+					error,
+				});
+			}
+		}
 	}
 
 	/** Handle server-to-client JSON-RPC requests (e.g. ping, roots/list). */
@@ -781,7 +883,7 @@ export class MCPManager {
 		// Remove tools from this server and notify consumers
 		const hadTools = this.#tools.some(t => t.mcpServerName === name);
 		this.#tools = this.#tools.filter(t => t.mcpServerName !== name);
-		if (hadTools) this.#onToolsChanged?.(this.#tools);
+		if (hadTools) void this.#onToolsChanged?.(this.#tools);
 
 		// Notify prompt consumers so stale commands are cleared
 		if (connection?.prompts?.length) this.#onPromptsChanged?.(name);
@@ -1022,7 +1124,7 @@ export class MCPManager {
 			const customTools = MCPTool.fromTools(connection, serverTools, reconnect);
 			void this.toolCache?.set(name, config, serverTools);
 			this.#replaceServerTools(name, customTools);
-			this.#onToolsChanged?.(this.#tools);
+			void this.#onToolsChanged?.(this.#tools);
 			void this.#loadServerResourcesAndPrompts(name, connection);
 			return connection;
 		} catch (error) {
@@ -1075,7 +1177,7 @@ export class MCPManager {
 
 		// Replace tools from this server
 		this.#replaceServerTools(name, customTools);
-		this.#onToolsChanged?.(this.#tools);
+		await this.#onToolsChanged?.(this.#tools);
 	}
 
 	/**

package/src/modes/acp/acp-agent.ts

@@ -651,12 +651,17 @@ export class AcpAgent implements Agent {
 			});
 		}
 
-		// For `thinking` the lifetime subscription pushes post-bootstrap; only
-		// push here when it's not yet installed so pre-bootstrap callers still
-		// see the change without a post-bootstrap duplicate.
-		const thinkingHandledBySubscription =
-			params.configId === THINKING_CONFIG_ID && record.lifetimeUnsubscribe !== undefined;
-		if (!thinkingHandledBySubscription) {
+		// For `model`/`thinking`, `#setModelById`/`#setThinkingLevelById` change
+		// the session model/thinking level through AgentSession, which now emits
+		// a lifetime event (`model_changed`/`thinking_level_changed`) that
+		// `#handleLifetimeEvent` turns into a push once the subscription is
+		// installed. Only push here when that subscription is not yet
+		// installed, so pre-bootstrap callers still see the change without a
+		// post-bootstrap duplicate.
+		const handledBySubscription =
+			(params.configId === THINKING_CONFIG_ID || params.configId === MODEL_CONFIG_ID) &&
+			record.lifetimeUnsubscribe !== undefined;
+		if (!handledBySubscription) {
 			await this.#pushConfigOptionUpdate(record);
 		}
 		return { configOptions: this.#buildConfigOptions(record.session) };
@@ -1150,14 +1155,15 @@ export class AcpAgent implements Agent {
 	}
 
 	async #handleLifetimeEvent(record: ManagedSessionRecord, event: AgentSessionEvent): Promise<void> {
-		if (event.type !== "thinking_level_changed") {
+		if (event.type !== "thinking_level_changed" && event.type !== "model_changed") {
 			return;
 		}
 		try {
 			await this.#pushConfigOptionUpdate(record);
 		} catch (error) {
-			logger.warn("Failed to push thinking-level config_option_update", {
+			logger.warn("Failed to push config_option_update after a lifetime event", {
 				sessionId: record.session.sessionId,
+				eventType: event.type,
 				error,
 			});
 		}
@@ -1575,7 +1581,7 @@ export class AcpAgent implements Agent {
 	#buildThinkingOptions(session: AgentSession): Array<{ value: string; name: string; description?: string }> {
 		return [
 			{ value: THINKING_OFF, name: "Off" },
-			{ value: AUTO_THINKING, name: "Auto", description: "Auto-detect per prompt (low–xhigh)" },
+			{ value: AUTO_THINKING, name: "Auto", description: "Auto-detect per prompt" },
 			...session.getAvailableThinkingLevels().map(level => ({
 				value: level,
 				name: level,

package/src/modes/components/chat-transcript-builder.ts

@@ -51,7 +51,7 @@ import {
 import { CustomMessageComponent } from "./custom-message";
 import { EvalExecutionComponent } from "./eval-execution";
 import { type LateDiagnosticsFile, LateDiagnosticsMessageComponent } from "./late-diagnostics-message";
-import { ReadToolGroupComponent, readArgsCollapseIntoGroup } from "./read-tool-group";
+import { groupedReadUsageCallIds, ReadToolGroupComponent, readArgsCollapseIntoGroup } from "./read-tool-group";
 import { SkillMessageComponent } from "./skill-message";
 import { ToolExecutionComponent } from "./tool-execution";
 import { TranscriptContainer } from "./transcript-container";
@@ -86,6 +86,7 @@ export class ChatTranscriptBuilder {
 	#pendingUsageDuration: number | undefined;
 	#pendingUsageTtft: number | undefined;
 	#pendingUsageTimestamp: number | undefined;
+	#pendingReadUsageCallIds: string[] | undefined;
 	#lastAssistantUsage: Usage | undefined;
 	#waitingPoll: ToolExecutionComponent | null = null;
 	#todoSnapshot: ToolExecutionComponent | null = null;
@@ -135,6 +136,7 @@ export class ChatTranscriptBuilder {
 		this.#pendingUsageDuration = undefined;
 		this.#pendingUsageTtft = undefined;
 		this.#pendingUsageTimestamp = undefined;
+		this.#pendingReadUsageCallIds = undefined;
 		this.#lastAssistantUsage = undefined;
 		this.#waitingPoll = null;
 		this.#todoSnapshot = null;
@@ -194,30 +196,46 @@ export class ChatTranscriptBuilder {
 		return this.#readGroup;
 	}
 
-	// The per-turn token-usage row must land below the turn's tool blocks, but
-	// normal `read` calls only materialize their group in #appendToolResult. Defer
-	// the row: stash it on the assistant message and flush once the turn's tools
-	// are placed, sealing the read run so the row sits under it.
+	// Defer per-turn metrics until the turn's tool results have materialized.
+	// Read-only invisible turns attach the metrics to their shared compact
+	// group; every other turn keeps the standalone row below its tool blocks.
 	#flushPendingUsage(): void {
 		if (!this.#pendingUsage) return;
-		this.#readGroup?.seal();
-		this.#readGroup = null;
-		this.container.addChild(
-			createUsageRowBlock(
+		const usageAttached =
+			this.#pendingReadUsageCallIds !== undefined &&
+			(this.#readGroup?.attachUsage(
+				this.#pendingReadUsageCallIds,
 				this.#pendingUsage,
 				this.#pendingUsageDuration,
 				this.#pendingUsageTtft,
 				this.#pendingUsageTimestamp,
-			),
-		);
+			) ??
+				false);
+		if (!usageAttached) {
+			this.#readGroup?.seal();
+			this.#readGroup = null;
+			this.container.addChild(
+				createUsageRowBlock(
+					this.#pendingUsage,
+					this.#pendingUsageDuration,
+					this.#pendingUsageTtft,
+					this.#pendingUsageTimestamp,
+				),
+			);
+		}
 		this.#pendingUsage = undefined;
 		this.#pendingUsageDuration = undefined;
 		this.#pendingUsageTtft = undefined;
 		this.#pendingUsageTimestamp = undefined;
+		this.#pendingReadUsageCallIds = undefined;
 	}
 
 	#appendChatMessage(message: AgentMessage): void {
 		if (message.role !== "toolResult") this.#flushPendingUsage();
+		if (message.role !== "assistant" && message.role !== "toolResult") {
+			this.#readGroup?.seal();
+			this.#readGroup = null;
+		}
 		switch (message.role) {
 			case "assistant":
 				this.#appendAssistantMessage(message);
@@ -403,6 +421,7 @@ export class ChatTranscriptBuilder {
 		this.#pendingUsageDuration = message.duration;
 		this.#pendingUsageTtft = message.ttft;
 		this.#pendingUsageTimestamp = message.timestamp;
+		this.#pendingReadUsageCallIds = this.#pendingUsage ? groupedReadUsageCallIds(message) : undefined;
 	}
 
 	#appendToolResult(message: Extract<AgentMessage, { role: "toolResult" }>): void {