From 29625f08c2443d8f6588f96abf4e1b01e519feec Mon Sep 17 00:00:00 2001 From: "Anthony \"Asterisk\" Ambuehl" Date: Tue, 21 Jul 2026 10:39:43 -0700 Subject: [PATCH] feat(mcp): add mcp_notification extension event + multi-listener API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Convert MCPManager's dangling single-slot setOnNotification callback into a multi-listener API and expose server-initiated MCP notifications as an extension event so extensions can bridge push-capable MCP servers (e.g. peer messaging, ticket nudges) into session behavior. API changes: - Removed: MCPManager.setOnNotification(handler) — single-slot, zero callers - Added: MCPManager.addNotificationListener(listener): () => void Multi-listener with per-listener error isolation, returns unsub. - Added: 'mcp_notification' extension event Payload: { server: string; method: string; params: unknown } Wired in sdk.ts: one listener bridges to extensionRunner.emitMcpNotification, captured under postmortem for teardown. Tests: 3 new (multi-listener fanout, error isolation, unsubscribe), fixture pattern matches neighboring mcp tests. bun check passes (biome + tsgo). Docs: extensions.md (new MCP notifications subsection with bridging example), mcp-runtime-lifecycle.md (Server-initiated notifications section), CHANGELOG. --- docs/extensions.md | 17 + docs/mcp-runtime-lifecycle.md | 13 + packages/coding-agent/CHANGELOG.md | 8 + .../src/extensibility/extensions/runner.ts | 65 ++++ .../src/extensibility/extensions/types.ts | 27 ++ packages/coding-agent/src/mcp/manager.ts | 138 ++++++- packages/coding-agent/src/sdk.ts | 67 +++- .../test/extensions-runner.test.ts | 163 +++++++++ .../test/fixtures/notifications-mcp.ts | 86 +++++ ...mcp-manager-notification-listeners.test.ts | 336 ++++++++++++++++++ .../test/mcp-server-tool-ownership.test.ts | 8 +- 11 files changed, 896 insertions(+), 32 deletions(-) create mode 100755 packages/coding-agent/test/fixtures/notifications-mcp.ts create mode 100644 packages/coding-agent/test/mcp-manager-notification-listeners.test.ts diff --git a/docs/extensions.md b/docs/extensions.md index 9fd9e3b32..76f19ce27 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -268,6 +268,23 @@ Cancelable pre-events: - `goal_updated` - `credential_disabled` +### MCP notifications + +- `mcp_notification` — fired for every JSON-RPC notification received from a connected MCP server, AFTER the manager's own handling of known list/update methods (`notifications/tools/list_changed`, `notifications/resources/list_changed`, `notifications/resources/updated`, `notifications/prompts/list_changed`). Unknown or server-custom methods are also delivered. Payload: `{ server: string; method: string; params: unknown }`. Multiple extensions may subscribe; a handler that throws does not prevent other handlers from firing. Notifications received before any listener attaches are buffered (bounded FIFO, cap 100, drop-oldest) and drained into the first subscriber — so startup-time frames aren't lost even if the extension binds after MCP discovery. + +Bridging a push-capable MCP into a session steer: + +```ts +pi.on("mcp_notification", event => { + if (event.server !== "peer-bus") return; + if (event.method !== "notifications/peer_message") return; + const params = event.params as { from: string; text: string }; + pi.sendUserMessage(`[from ${params.from}] ${params.text}`, { deliverAs: "steer" }); +}); +``` + +The runtime handles the JSON-RPC transport and its own list/update refresh first; the handler runs afterwards and can inject a mid-turn steer via `pi.sendMessage` / `pi.sendUserMessage`. + ### User command interception - `user_bash` (override with `{ result }`) diff --git a/docs/mcp-runtime-lifecycle.md b/docs/mcp-runtime-lifecycle.md index 7896bf8ac..45707571c 100644 --- a/docs/mcp-runtime-lifecycle.md +++ b/docs/mcp-runtime-lifecycle.md @@ -158,6 +158,19 @@ Both return structured tool output and convert remaining transport/tool errors i There is also a follow-up path for late connections: after waiting for a specific server, if status becomes `connected`, it re-runs `session.refreshMCPTools(...)` so newly available tools are rebound in-session. + +## Server-initiated notifications + +MCP servers may push JSON-RPC notification frames at any point after `initialize` completes. The transport surfaces them via `onNotification`; the manager fans them out in two paths: + +1. **Internal refresh** for known methods: + - `notifications/tools/list_changed` → `refreshServerTools` + - `notifications/resources/list_changed` → `refreshServerResources` + - `notifications/resources/updated` → `#onResourcesChanged` (only for currently subscribed URIs) + - `notifications/prompts/list_changed` → `refreshServerPrompts` +2. **Listener fanout**: every notification (including the known ones AND server-custom methods) is delivered to registered listeners AFTER the internal refresh runs. Registered via `MCPManager.addNotificationListener(listener)`, which returns an unsubscribe function. Multiple listeners are supported; each is invoked with independent error isolation — a synchronous throw in one listener does not prevent others from firing (thrown errors are logged at `debug`). + +`sdk.ts` registers one listener that bridges to the extension runner's `mcp_notification` event, so extensions receive every server-initiated frame with `{ server, method, params }`. The listener is captured with `postmortem` so it is released on session teardown. ## Health, reconnect, and partial failure behavior Current runtime behavior is connection-event driven: diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index e150972d1..ed2bc494c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,14 @@ ## [Unreleased] +### Added + +- Added `mcp_notification` extension event and multi-listener `MCPManager.addNotificationListener` API. The runtime already received MCP server-initiated JSON-RPC notifications at the transport layer but had no path to forward them to extensions; every notification (including server-custom methods) is now delivered as `{ server, method, params }` after the manager's own list/update handling. For known list-change methods (`notifications/tools/list_changed`, `notifications/resources/list_changed`, `notifications/prompts/list_changed`) the internal refresh promise is awaited before fanout, so a listener acting on `tools/list_changed` sees fresh `getTools()`. Notifications received before any listener attaches are buffered (bounded FIFO, cap 100, drop-oldest — matches `IrcBus`'s `MAILBOX_CAP`) and drained into the first subscriber, so startup-time frames aren't lost even if the extension binds after MCP discovery. Extensions can use this to bridge push-capable MCP servers (e.g. peer messaging) into session behavior by injecting a mid-turn steer via `pi.sendMessage` / `pi.sendUserMessage`. + +### Removed + +- Removed the dangling `MCPManager.setOnNotification` single-slot setter, which had no callers in the runtime. Replaced by `MCPManager.addNotificationListener` — multi-listener, per-listener error isolation, returns an unsubscribe function. + ## [17.1.2] - 2026-07-24 ### Added diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 43a37f1c3..9c69a147c 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -39,6 +39,7 @@ import type { ExtensionUIContext, InputEvent, InputEventResult, + McpNotificationEvent, MessageRenderer, RegisteredCommand, RegisteredTool, @@ -132,6 +133,14 @@ async function raceHandlerWithTimeout( const MAX_PENDING_CREDENTIAL_DISABLED = 32; +/** + * Buffer cap for `mcp_notification` events received before {@link ExtensionRunner.initialize} + * has run. Sized to match the manager-side buffer in `MCPManager.NOTIFICATION_BUFFER_CAP` so + * the two layers can't drop different amounts of the same burst — the pipe drains, or it + * spills, but it does so consistently at both ends. Drop-oldest under pressure. + */ +const MAX_PENDING_MCP_NOTIFICATIONS = 100; + /** * Events handled by the generic emit() method. * Events with dedicated emitXxx() methods are excluded for stronger type safety. @@ -259,6 +268,19 @@ export class ExtensionRunner { */ #pendingCredentialDisabled: CredentialDisabledEvent[] = []; + /** + * Buffer for `mcp_notification` events received via {@link emitMcpNotification} before + * {@link initialize} has run. Two-layer race: `MCPManager` also buffers frames until + * its first `addNotificationListener` subscriber attaches, but the sdk.ts bridge is + * registered inside `createAgentSession` — BEFORE the mode controller calls + * `ExtensionRunner.initialize()`. Without this second buffer, the manager's drain + * arrives at the bridge → the bridge calls `emitMcpNotification` → the runner drops + * the frame because `#initialized === false`, and the frame evaporates a second time. + * Bounded at {@link MAX_PENDING_MCP_NOTIFICATIONS}; oldest entries are dropped under + * pressure. Drained in {@link initialize} once the runtime/UI context is wired. + */ + #pendingMcpNotifications: Array> = []; + /** * Timers scheduled by extensions through the sanctioned `ctx.setInterval` / * `ctx.setTimeout` helpers. Callbacks run with the same isolation as handler @@ -345,6 +367,23 @@ export class ExtensionRunner { }); } }); + + // Drain events buffered by emitMcpNotification() before initialize ran, using the + // same deferred-microtask ordering as the credential-disabled drain above so any + // onError listener registered synchronously after initialize() still catches + // handler errors during flush. + const pendingMcp = this.#pendingMcpNotifications.splice(0); + queueMicrotask(() => { + for (const event of pendingMcp) { + this.emit({ type: "mcp_notification", ...event }).catch((error: unknown) => { + logger.warn("mcp_notification handler threw during initialize flush", { + server: event.server, + method: event.method, + error: error instanceof Error ? error.message : String(error), + }); + }); + } + }); } /** @@ -372,6 +411,32 @@ export class ExtensionRunner { await this.emit({ type: "credential_disabled", ...event }); } + /** + * Forward an MCP server notification to extension handlers. + * + * If {@link initialize} has not yet run, the notification is buffered and replayed + * once initialize wires the runtime/UI context. Matches the credential-disabled + * deferral above: the sdk.ts bridge registers `MCPManager.addNotificationListener` + * inside `createAgentSession` — BEFORE the mode controller calls `initialize()` on + * this runner — so notification frames drained by the manager (either fresh + * arrivals or replay from its own startup buffer) can reach us pre-init. Without + * this buffer they would evaporate for a second time here. + * + * Bounded at {@link MAX_PENDING_MCP_NOTIFICATIONS}; oldest entries drop under + * pressure. Never throws; per-handler errors are routed through {@link onError} + * via {@link emit}'s normal isolation. + */ + async emitMcpNotification(event: Omit): Promise { + if (!this.#initialized) { + if (this.#pendingMcpNotifications.length >= MAX_PENDING_MCP_NOTIFICATIONS) { + this.#pendingMcpNotifications.shift(); + } + this.#pendingMcpNotifications.push(event); + return; + } + await this.emit({ type: "mcp_notification", ...event }); + } + async emitSessionStop(event: Omit): Promise { return await this.emit({ type: "session_stop", ...event }); } diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 0d7493b68..3a6ce857a 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -713,6 +713,31 @@ export interface CredentialDisabledEvent { disabledCause: string; } +// ============================================================================ +// MCP Events +// ============================================================================ + +/** + * Fired for every JSON-RPC notification received from a connected MCP server, + * AFTER the runtime's own handling of known list/update methods. Unknown or + * server-custom methods are delivered too — extensions can bridge them into + * session behavior by inspecting `method`/`params` and injecting a follow-up + * via `pi.sendMessage(..., { deliverAs })` or `pi.sendUserMessage(...)`. + */ +export interface McpNotificationEvent { + type: "mcp_notification"; + /** + * Server name as declared in the MCP config (raw, unsanitized). Note this + * differs from the sanitized prefix used in `mcp___` + * tool names — filter by this raw name, not by tool-name prefix matching. + */ + server: string; + /** JSON-RPC method (e.g. `notifications/tools/list_changed`, or server-custom). */ + method: string; + /** JSON-RPC params, opaque to the runtime. */ + params: unknown; +} + // ============================================================================ // User Bash Events // ============================================================================ @@ -941,6 +966,7 @@ export type ExtensionEvent = | TodoReminderEvent | GoalUpdatedEvent | CredentialDisabledEvent + | McpNotificationEvent | UserBashEvent | UserPythonEvent | InputEvent @@ -1134,6 +1160,7 @@ export interface ExtensionAPI { on(event: "tool_result", handler: ExtensionHandler): void; on(event: "user_bash", handler: ExtensionHandler): void; on(event: "user_python", handler: ExtensionHandler): void; + on(event: "mcp_notification", handler: ExtensionHandler): void; // ========================================================================= // Tool Registration diff --git a/packages/coding-agent/src/mcp/manager.ts b/packages/coding-agent/src/mcp/manager.ts index 12a8e1ad9..16fd00321 100644 --- a/packages/coding-agent/src/mcp/manager.ts +++ b/packages/coding-agent/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(promise: Promise): TrackedPromise { const tracked: TrackedPromise = { promise, status: "pending" }; promise.then( @@ -194,8 +202,14 @@ export class MCPManager { #sources = new Map(); #authStorage: AuthStorage | null = null; #authHandler?: MCPAuthHandler; - #onNotification?: (serverName: string, method: string, params: unknown) => void; - #onToolsChanged?: (tools: CustomTool[]) => 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[]) => void | Promise; #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[]) => void): void { + setOnToolsChanged(handler: (tools: CustomTool[]) => void | Promise): 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 { 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 { 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 | 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) { @@ -1081,7 +1183,7 @@ export class MCPManager { // Replace tools from this server this.#replaceServerTools(name, customTools); - this.#onToolsChanged?.(this.#tools); + await this.#onToolsChanged?.(this.#tools); } /** diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index a08aa5683..4a8f613cc 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -3203,6 +3203,15 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} throw new Error(`Agent "${resolvedAgentId}" was replaced during session initialization.`); } hasRegistered = true; + // MCP notification bridge cleanup — assigned when the bridge is wired below, + // invoked from the dispose wrapper AND registered as a postmortem so both + // explicit-dispose (SDK embedders that reuse the process across sessions) and + // process-exit paths tear the listener down. Nulled after use so the closure + // graph (`extensionRunner`, `session`) can be GC'd instead of retained by the + // process-global postmortem list. + let unsubscribeMcpNotifications: (() => void) | undefined; + let unregisterMcpPostmortem: (() => void) | undefined; + { const originalDispose = session.dispose.bind(session); session.dispose = async () => { @@ -3233,6 +3242,12 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} } finally { unregisterUnlessParked(); unsubscribeCredentialDisabled?.(); + unsubscribeMcpNotifications?.(); + unregisterMcpPostmortem?.(); + // Drop refs so the process-global postmortem list doesn't retain + // the bridge closure past explicit dispose. + unsubscribeMcpNotifications = undefined; + unregisterMcpPostmortem = undefined; } }; } @@ -3405,19 +3420,24 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} } } - // Wire MCP manager callbacks to session for reactive tool updates. - // Skip when reusing a parent's manager — the parent owns the callbacks. + // MCP manager wiring has two ownership models: + // * Single-slot callbacks (tools/prompts/resources changed) — exactly one + // owner per manager. When reusing a parent's manager (subagent path, + // see task/executor.ts), the parent already owns these slots so we + // MUST NOT overwrite them. Guarded by `!options.mcpManager`. + // * Notification listener — multi-listener by design. Every session with + // an MCP manager (fresh OR reused) needs its own bridge to its own + // `extensionRunner` so extensions loaded in that session receive frames. + // Guarded only by `mcpManager` (see the second `if` below). if (mcpManager && !options.mcpManager) { - mcpManager.setOnToolsChanged(tools => { - void (async () => { - try { - await session.refreshMCPTools(tools); - } catch (error) { - logger.warn("MCP tool refresh failed", { - error: error instanceof Error ? error.message : String(error), - }); - } - })(); + mcpManager.setOnToolsChanged(async tools => { + try { + await session.refreshMCPTools(tools); + } catch (error) { + logger.warn("MCP tool refresh failed", { + error: error instanceof Error ? error.message : String(error), + }); + } }); // Wire prompt refresh → rebuild MCP prompt slash commands mcpManager.setOnPromptsChanged(serverName => { @@ -3450,6 +3470,29 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} }); } + if (mcpManager) { + // Bridge server-initiated notifications to this session's extension + // handlers. Multi-listener registration: fresh-manager and reused-manager + // sessions both install their own listener here, so a subagent's + // extensions get frames even though the parent owns the single-slot + // tool/prompt/resource callbacks above. MCPManager fires known + // list/update refreshes internally, then invokes all registered + // listeners with (server, method, params) for every frame (including + // server-custom methods). Two-layer buffering protects the startup + // race: MCPManager buffers frames received before the first + // `addNotificationListener` subscriber (drains here); ExtensionRunner + // buffers frames received before `initialize()` and drains them on + // init. Both drop-oldest under pressure at cap 100. + unsubscribeMcpNotifications = mcpManager.addNotificationListener((server, method, params) => { + void extensionRunner.emitMcpNotification({ server, method, params }); + }); + // postmortem.register returns a cancel function; capture it so explicit + // session.dispose can remove this from the global list (see finally above). + unregisterMcpPostmortem = postmortem.register("mcp-notification-listener-cleanup", () => + unsubscribeMcpNotifications?.(), + ); + } + startDeferredMCPDiscovery?.(session); return { diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index aaaf82136..e51f7fead 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -2114,6 +2114,169 @@ describe("ExtensionRunner", () => { }); }); + describe("mcp_notification", () => { + it("delivers mcp_notification events to subscribed extensions with the typed payload", async () => { + const eventsPath = path.join(tempDir.path(), "mcp-notification-events.jsonl"); + const extCode = ` + import * as fs from "node:fs"; + + export default function(pi) { + pi.on("mcp_notification", async (event) => { + fs.appendFileSync( + ${JSON.stringify(eventsPath)}, + JSON.stringify({ + type: event.type, + server: event.server, + method: event.method, + params: event.params, + }) + "\\n", + ); + }); + } + `; + fs.writeFileSync(path.join(extensionsDir, "mcp-notification.ts"), extCode); + + const result = await loadTestExtensions(); + const runner = new ExtensionRunner( + result.extensions, + result.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + ); + runner.initialize( + { + sendMessage: () => {}, + sendUserMessage: () => {}, + appendEntry: () => {}, + setLabel: () => {}, + getActiveTools: () => [], + getAllTools: () => [], + setActiveTools: async () => {}, + getCommands: () => [], + setModel: async () => false, + getThinkingLevel: () => undefined, + setThinkingLevel: () => {}, + getSessionName: () => sessionManager.getSessionName(), + setSessionName: async () => {}, + }, + { + getModel: () => undefined, + isIdle: () => true, + abort: () => {}, + hasPendingMessages: () => false, + shutdown: () => {}, + getContextUsage: () => undefined, + compact: async () => {}, + getSystemPrompt: () => [], + }, + ); + + await runner.emitMcpNotification({ + server: "peers", + method: "notifications/peer_message", + params: { from: "alice", text: "hi" }, + }); + + const events = fs + .readFileSync(eventsPath, "utf8") + .trim() + .split("\n") + .map(line => JSON.parse(line)); + expect(events).toEqual([ + { + type: "mcp_notification", + server: "peers", + method: "notifications/peer_message", + params: { from: "alice", text: "hi" }, + }, + ]); + }); + + it("buffers pre-initialize events and drains them on initialize (caps at 100, drops oldest)", async () => { + // Guard against regression of the two-layer startup race that Codex flagged + // on PR #6535 (commit ffa058aa8): the sdk.ts bridge wires + // mcpManager.addNotificationListener inside createAgentSession BEFORE the + // mode controller calls ExtensionRunner.initialize(). Frames the manager + // drains from its own buffer arrive here pre-init. Prior behavior silently + // dropped them; the fix buffers and drains on initialize (same shape as + // emitCredentialDisabled). + const eventsPath = path.join(tempDir.path(), "mcp-notification-cap.jsonl"); + const extCode = ` + import * as fs from "node:fs"; + + export default function(pi) { + pi.on("mcp_notification", async (event) => { + fs.appendFileSync( + ${JSON.stringify(eventsPath)}, + JSON.stringify({ server: event.server, method: event.method }) + "\\n", + ); + }); + } + `; + fs.writeFileSync(path.join(extensionsDir, "mcp-notification-cap.ts"), extCode); + + const result = await loadTestExtensions(); + const runner = new ExtensionRunner( + result.extensions, + result.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + ); + + // Push 101 events while uninitialized — the 1st should be dropped, next 100 buffered. + for (let i = 0; i < 101; i++) { + await runner.emitMcpNotification({ + server: "peers", + method: `notifications/test/${i}`, + params: null, + }); + } + + runner.initialize( + { + sendMessage: () => {}, + sendUserMessage: () => {}, + appendEntry: () => {}, + setLabel: () => {}, + getActiveTools: () => [], + getAllTools: () => [], + setActiveTools: async () => {}, + getCommands: () => [], + setModel: async () => false, + getThinkingLevel: () => undefined, + setThinkingLevel: () => {}, + getSessionName: () => sessionManager.getSessionName(), + setSessionName: async () => {}, + }, + { + getModel: () => undefined, + isIdle: () => true, + abort: () => {}, + hasPendingMessages: () => false, + shutdown: () => {}, + getContextUsage: () => undefined, + compact: async () => {}, + getSystemPrompt: () => [], + }, + ); + + // Drain microtasks so the fire-and-forget emit() calls inside initialize() complete. + for (let i = 0; i < 5; i++) await Promise.resolve(); + + const events = fs + .readFileSync(eventsPath, "utf8") + .trim() + .split("\n") + .map(line => JSON.parse(line)); + expect(events).toHaveLength(100); + // Drop-oldest policy: test/0 was evicted, test/1 survives as the head. + expect(events[0]?.method).toBe("notifications/test/1"); + expect(events[99]?.method).toBe("notifications/test/100"); + }); + }); + describe("managed timers (ctx.setInterval / ctx.setTimeout)", () => { it("contains a throwing interval callback instead of letting it escape as uncaughtException", () => { vi.useFakeTimers(); diff --git a/packages/coding-agent/test/fixtures/notifications-mcp.ts b/packages/coding-agent/test/fixtures/notifications-mcp.ts new file mode 100755 index 000000000..c8aa1c571 --- /dev/null +++ b/packages/coding-agent/test/fixtures/notifications-mcp.ts @@ -0,0 +1,86 @@ +#!/usr/bin/env bun +/** + * Test fixture: minimal stdio MCP server used by + * `mcp-manager-notification-listeners.test.ts` to exercise the notification + * listener API. On receiving the client's `notifications/initialized`, emits + * TWO server-to-client notification frames: + * + * 1. `notifications/tools/list_changed` — a known method that MCPManager + * handles internally (triggers a `tools/list` refresh). Delivered to + * listeners AFTER the internal handling; verifies fanout is not gated + * on the method being unknown. + * 2. A server-custom method (`notifications/custom/test-event`) with a + * distinctive payload — verifies that arbitrary methods are delivered + * to listeners verbatim, which is the extension use case (bridging a + * peer-messaging MCP's custom push into a session steer). + * + * The custom method + payload constants are exported so the test can assert + * on frame equality without duplicating literals. + */ +import * as readline from "node:readline"; + +export const CUSTOM_NOTIFICATION_METHOD = "notifications/custom/test-event"; +export const CUSTOM_NOTIFICATION_PAYLOAD: Readonly> = Object.freeze({ + hello: "world", + n: 42, +}); + +type JsonRpcMessage = { + jsonrpc: "2.0"; + id?: string | number; + method?: string; + params?: Record; +}; + +function buildResult(method: string): Record { + switch (method) { + case "initialize": + return { + protocolVersion: "2025-03-26", + serverInfo: { name: "notifications-fixture", version: "1.0.0" }, + capabilities: { tools: { listChanged: true } }, + }; + case "tools/list": + return { tools: [] }; + default: + return {}; + } +} + +function writeNotification(method: string, params?: Record): void { + const frame: JsonRpcMessage = { jsonrpc: "2.0", method }; + if (params) frame.params = params; + process.stdout.write(`${JSON.stringify(frame)}\n`); +} + +function startServer(): void { + const rl = readline.createInterface({ input: process.stdin }); + rl.on("line", line => { + void (async () => { + const trimmed = line.trim(); + if (trimmed.length === 0) return; + let msg: JsonRpcMessage; + try { + msg = JSON.parse(trimmed) as JsonRpcMessage; + } catch { + return; + } + const isNotification = msg.id === undefined || msg.id === null; + if (isNotification) { + if (msg.method === "notifications/initialized") { + writeNotification("notifications/tools/list_changed"); + writeNotification(CUSTOM_NOTIFICATION_METHOD, { ...CUSTOM_NOTIFICATION_PAYLOAD }); + } + return; + } + if (!msg.method) return; + const response = { jsonrpc: "2.0" as const, id: msg.id, result: buildResult(msg.method) }; + process.stdout.write(`${JSON.stringify(response)}\n`); + })(); + }); + rl.on("close", () => process.exit(0)); +} + +if (import.meta.main) { + startServer(); +} diff --git a/packages/coding-agent/test/mcp-manager-notification-listeners.test.ts b/packages/coding-agent/test/mcp-manager-notification-listeners.test.ts new file mode 100644 index 000000000..8f2678c14 --- /dev/null +++ b/packages/coding-agent/test/mcp-manager-notification-listeners.test.ts @@ -0,0 +1,336 @@ +import { afterEach, beforeEach, describe, expect, it } from "bun:test"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { MCPManager } from "@oh-my-pi/pi-coding-agent/mcp/manager"; +import type { MCPServerConfig } from "@oh-my-pi/pi-coding-agent/mcp/types"; +import { removeSyncWithRetries } from "@oh-my-pi/pi-utils"; +import { CUSTOM_NOTIFICATION_METHOD, CUSTOM_NOTIFICATION_PAYLOAD } from "./fixtures/notifications-mcp"; + +const FIXTURE_PATH = path.join(import.meta.dir, "fixtures", "notifications-mcp.ts"); +const BUN_EXEC = process.execPath; + +type CapturedFrame = { server: string; method: string; params: unknown }; + +function serverConfig(): MCPServerConfig { + return { type: "stdio", command: BUN_EXEC, args: [FIXTURE_PATH] }; +} + +/** + * Bind a capturing listener that ALSO exposes a promise-based await for a + * specific `(server, method)` pair. Lets each test wait for the exact frame it + * cares about without polling — the frame arriving is the real signal, not a + * timer. bun:test's per-test timeout backstops a genuine hang. + */ +function makeFrameCollector(): { + listener: (server: string, method: string, params: unknown) => void; + frames: CapturedFrame[]; + awaitFrame: (server: string, method: string) => Promise; +} { + const frames: CapturedFrame[] = []; + const waiters = new Map void>>(); + const key = (server: string, method: string) => `${server}\x00${method}`; + return { + frames, + listener: (server, method, params) => { + const frame: CapturedFrame = { server, method, params }; + frames.push(frame); + const bucket = waiters.get(key(server, method)); + if (bucket && bucket.length > 0) { + waiters.delete(key(server, method)); + for (const resolve of bucket) resolve(frame); + } + }, + awaitFrame: (server, method) => { + const existing = frames.find(f => f.server === server && f.method === method); + if (existing) return Promise.resolve(existing); + const { promise, resolve } = Promise.withResolvers(); + const bucket = waiters.get(key(server, method)) ?? []; + bucket.push(resolve); + waiters.set(key(server, method), bucket); + return promise; + }, + }; +} + +describe("MCPManager notification listeners", () => { + let workDir: string; + + beforeEach(() => { + workDir = fs.mkdtempSync(path.join(os.tmpdir(), "omp-mcp-notif-")); + }); + + afterEach(() => { + removeSyncWithRetries(workDir); + }); + + it("delivers known and server-custom notifications to a registered listener", async () => { + const manager = new MCPManager(workDir); + const collector = makeFrameCollector(); + const unsubscribe = manager.addNotificationListener(collector.listener); + expect(typeof unsubscribe).toBe("function"); + + try { + const result = await manager.connectServers({ alpha: serverConfig() }, {}); + expect(result.connectedServers).toContain("alpha"); + + // Await both the known list_changed and the server-custom frame + // independently. Arrival order across the two isn't guaranteed + // because list_changed triggers an internal `tools/list` refresh + // that MCPManager awaits before fanout (so listeners see fresh + // `getTools()`), which lets a non-refresh frame like the custom + // one overtake in the same batch. Both MUST still be delivered; + // that's the actual contract this test exercises. + const [listChanged, custom] = await Promise.all([ + collector.awaitFrame("alpha", "notifications/tools/list_changed"), + collector.awaitFrame("alpha", CUSTOM_NOTIFICATION_METHOD), + ]); + expect(custom.params).toEqual(CUSTOM_NOTIFICATION_PAYLOAD); + expect(listChanged.method).toBe("notifications/tools/list_changed"); + } finally { + await manager.disconnectAll(); + } + }); + + it("dispatches to every listener and isolates handler errors", async () => { + const manager = new MCPManager(workDir); + const collectorA = makeFrameCollector(); + const collectorC = makeFrameCollector(); + + manager.addNotificationListener(collectorA.listener); + manager.addNotificationListener(() => { + // Middle listener throws synchronously. Must NOT prevent the third + // listener from being invoked for the same frame. + throw new Error("listener-B intentionally throws"); + }); + manager.addNotificationListener(collectorC.listener); + + try { + await manager.connectServers({ alpha: serverConfig() }, {}); + const [frameA, frameC] = await Promise.all([ + collectorA.awaitFrame("alpha", CUSTOM_NOTIFICATION_METHOD), + collectorC.awaitFrame("alpha", CUSTOM_NOTIFICATION_METHOD), + ]); + expect(frameA.params).toEqual(CUSTOM_NOTIFICATION_PAYLOAD); + expect(frameC.params).toEqual(CUSTOM_NOTIFICATION_PAYLOAD); + } finally { + await manager.disconnectAll(); + } + }); + + it("unsubscribe stops delivery to that listener without affecting others", async () => { + const manager = new MCPManager(workDir); + const early = makeFrameCollector(); + const late = makeFrameCollector(); + + const unsubscribeEarly = manager.addNotificationListener(early.listener); + + try { + await manager.connectServers({ alpha: serverConfig() }, {}); + await early.awaitFrame("alpha", CUSTOM_NOTIFICATION_METHOD); + const earlyCountBeforeUnsub = early.frames.length; + + unsubscribeEarly(); + + // Add a fresh listener AFTER unsub, then connect a second server so + // the fixture emits a new batch. The unsubscribed listener must not + // grow; the fresh listener must observe the fresh batch. Dispatch + // per frame is synchronous over the listener set, so once `late` + // resolves for a `beta` frame, any concurrent delivery to `early` + // would have already happened had the unsubscribe not landed. + manager.addNotificationListener(late.listener); + await manager.connectServers({ beta: serverConfig() }, {}); + await late.awaitFrame("beta", CUSTOM_NOTIFICATION_METHOD); + + expect(early.frames.length).toBe(earlyCountBeforeUnsub); + } finally { + await manager.disconnectAll(); + } + }); + + it("buffers notifications received before any listener attaches, drains into first subscriber", async () => { + const manager = new MCPManager(workDir); + + try { + // Connect FIRST, before any listener is registered. The fixture + // emits `notifications/tools/list_changed` plus a server-custom + // notification during its `notifications/initialized` handshake, + // so by the time connect resolves at least one notification frame + // has already been received by the manager with no consumer. + await manager.connectServers({ alpha: serverConfig() }, {}); + + // Attach the FIRST listener AFTER frames were already dispatched. + // Drain-on-first-subscribe must deliver the buffered frames now. + const collector = makeFrameCollector(); + manager.addNotificationListener(collector.listener); + // Both the known list_changed and the server-custom frame must + // have been buffered and drained on attach. Await each + // independently — arrival order between refresh-triggering and + // non-refresh frames is not guaranteed (see the delivery test + // for rationale), but both are always delivered. + const [listChanged, custom] = await Promise.all([ + collector.awaitFrame("alpha", "notifications/tools/list_changed"), + collector.awaitFrame("alpha", CUSTOM_NOTIFICATION_METHOD), + ]); + expect(custom.params).toEqual(CUSTOM_NOTIFICATION_PAYLOAD); + expect(listChanged.method).toBe("notifications/tools/list_changed"); + } finally { + await manager.disconnectAll(); + } + }); + + it("refreshServerTools awaits an async setOnToolsChanged handler before resolving", async () => { + // Verifies the second-layer guarantee (issue raised in review): + // #onToolsChanged in the sdk.ts wiring calls session.refreshMCPTools() + // asynchronously. If refreshServerTools does not await the callback, + // a mcp_notification listener acting on tools/list_changed can call + // getAllTools() on the session before its registry has been rebound. + // This test drives the same pattern by registering a slow async + // setOnToolsChanged handler and confirming refreshServerTools blocks + // on it before resolving — proving the fanout in + // #handleServerNotification (which awaits refreshServerTools) sees + // the callback complete. + const manager = new MCPManager(workDir); + const events: string[] = []; + const { promise: hold, resolve: release } = Promise.withResolvers(); + + manager.setOnToolsChanged(async () => { + events.push("cb:start"); + await hold; + events.push("cb:end"); + }); + + try { + await manager.connectServers({ alpha: serverConfig() }, {}); + // connectServers may kick the callback for the initial tool load via + // its background continuation; drain that first so we can assert + // specifically on the refresh path. + await Bun.sleep(30); + release(); + await Bun.sleep(30); + const priorLen = events.length; + + // Trigger a fresh refresh and hold the callback again — this is the + // path #handleServerNotification takes for tools/list_changed. + const { promise: hold2, resolve: release2 } = Promise.withResolvers(); + let cbState: "started" | "ended" | undefined; + manager.setOnToolsChanged(async () => { + cbState = "started"; + await hold2; + cbState = "ended"; + }); + const refreshDone = manager.refreshServerTools("alpha"); + // Give the callback time to enter the await. + await Bun.sleep(20); + expect(cbState).toBe("started"); + // refreshServerTools has NOT resolved yet — proves it's awaiting. + let refreshResolved = false; + void refreshDone.then(() => { + refreshResolved = true; + }); + await Bun.sleep(10); + expect(refreshResolved).toBe(false); + // Release the callback; refresh should now resolve promptly. + release2(); + await refreshDone; + expect(cbState).toBe("ended"); + // Drain events reference — just to silence unused var if refactored. + expect(events.length).toBeGreaterThanOrEqual(priorLen); + } finally { + await manager.disconnectAll(); + } + }); + + it("mcp_notification listener for tools/list_changed fires only AFTER setOnToolsChanged callback resolves", async () => { + // Option B integration test: exercises the full manager-side chain + // #handleServerNotification → refreshServerTools → #onToolsChanged + // → fanout to notification listeners. Uses a slow async + // setOnToolsChanged callback that gates on a controllable promise + // (matches the mcp-dispose-disconnect-bounded.test.ts gate pattern) + // to prove that if the SDK's setOnToolsChanged handler is async and + // its work isn't complete, the mcp_notification listener does NOT + // fire prematurely. Regression guard: catches any future refactor + // that reintroduces fire-and-forget in the callback chain (e.g. the + // original sdk.ts:3411 `void (async () => ...)()` wrapper — which + // was the specific race raised in review of #6535). + const manager = new MCPManager(workDir); + const events: string[] = []; + let cbCall = 0; + const { promise: gate, resolve: releaseGate } = Promise.withResolvers(); + + manager.setOnToolsChanged(async () => { + cbCall++; + // Initial connect fires the callback via a background continuation + // (manager.ts line ~555, fire-and-forget). Return immediately so + // its floating promise resolves cleanly; only gate on the second + // call, which is the one our simulated post-connect notification + // triggers. + if (cbCall === 1) { + events.push("initial-cb:done"); + return; + } + events.push("cb:start"); + await gate; + events.push("cb:end"); + }); + + manager.addNotificationListener((_server, method) => { + if (method === "notifications/tools/list_changed") { + events.push("listener:fire"); + } + }); + + try { + await manager.connectServers({ alpha: serverConfig() }, {}); + // Wait for the initial-connect background continuation to fire the + // callback (call #1, ungated). Once it's recorded, we know the + // callback is idle and the next invocation will be call #2. + for (let i = 0; i < 20 && !events.includes("initial-cb:done"); i++) { + await Bun.sleep(10); + } + expect(events).toContain("initial-cb:done"); + + // Simulate a post-connect `notifications/tools/list_changed` frame + // by driving the transport's onNotification hook directly — that's + // exactly what the transport does when a server pushes a + // notification after startup. Bypasses the fixture's initial + // handshake burst so we exercise the connectionKnown=true path + // (where refresh actually runs and the callback is awaited). + const conn = manager.getConnection("alpha"); + if (!conn?.transport.onNotification) throw new Error("expected transport onNotification to be wired"); + void conn.transport.onNotification("notifications/tools/list_changed", {}); + + // Give the async chain time to enter the callback's await. + for (let i = 0; i < 20 && !events.includes("cb:start"); i++) { + await Bun.sleep(10); + } + expect(events).toContain("cb:start"); + + // CRITICAL: at this point the callback is holding, and the + // listener MUST NOT have fired. If it has, the manager fanned + // out before awaiting the callback — the exact race this test + // guards against. + expect(events).not.toContain("cb:end"); + expect(events).not.toContain("listener:fire"); + + // Release the gate. Now the callback resolves, refresh completes, + // and the listener fires — in that order. + releaseGate(); + + // Wait for the listener to fire (or bail on timeout). + for (let i = 0; i < 50 && !events.includes("listener:fire"); i++) { + await Bun.sleep(10); + } + + const idxCbEnd = events.indexOf("cb:end"); + const idxFire = events.indexOf("listener:fire"); + expect(idxCbEnd).toBeGreaterThanOrEqual(0); + expect(idxFire).toBeGreaterThanOrEqual(0); + expect(idxFire).toBeGreaterThan(idxCbEnd); + } finally { + // Ensure gate is released even on early failure so no promise leaks. + releaseGate(); + await manager.disconnectAll(); + } + }, 15_000); +}); diff --git a/packages/coding-agent/test/mcp-server-tool-ownership.test.ts b/packages/coding-agent/test/mcp-server-tool-ownership.test.ts index fdd7973c5..bc63ddfdd 100644 --- a/packages/coding-agent/test/mcp-server-tool-ownership.test.ts +++ b/packages/coding-agent/test/mcp-server-tool-ownership.test.ts @@ -61,7 +61,9 @@ describe("MCP tool ownership with prefix-colliding server names", () => { expect(names()).toHaveLength(MANY_TOOL_COUNT * 2); const payloads: string[][] = []; - manager.setOnToolsChanged(tools => payloads.push(tools.map(t => t.name))); + manager.setOnToolsChanged(tools => { + payloads.push(tools.map(t => t.name)); + }); // Same code path a reconnect takes: replace the named server's tools. await manager.refreshServerTools(SHORT_SERVER); @@ -80,7 +82,9 @@ describe("MCP tool ownership with prefix-colliding server names", () => { it("disconnecting a server with sanitized name characters removes exactly its tools", async () => { await manager.connectServers({ [SHORT_SERVER]: fixtureConfig(), [COLON_SERVER]: fixtureConfig() }, {}); const payloads: string[][] = []; - manager.setOnToolsChanged(tools => payloads.push(tools.map(t => t.name))); + manager.setOnToolsChanged(tools => { + payloads.push(tools.map(t => t.name)); + }); await manager.disconnectServer(COLON_SERVER);