feat(mcp): add mcp_notification extension event + multi-listener API

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.
This commit is contained in:
Anthony "Asterisk" Ambuehl
2026-07-21 10:39:43 -07:00
parent a38cd95d7d
commit 29625f08c2
11 changed files with 896 additions and 32 deletions
+17
View File
@@ -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 }`)
+13
View File
@@ -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:
+8
View File
@@ -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
@@ -39,6 +39,7 @@ import type {
ExtensionUIContext,
InputEvent,
InputEventResult,
McpNotificationEvent,
MessageRenderer,
RegisteredCommand,
RegisteredTool,
@@ -132,6 +133,14 @@ async function raceHandlerWithTimeout<T>(
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<Omit<McpNotificationEvent, "type">> = [];
/**
* 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<McpNotificationEvent, "type">): Promise<void> {
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<SessionStopEvent, "type">): Promise<SessionStopEventResult | undefined> {
return await this.emit({ type: "session_stop", ...event });
}
@@ -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__<sanitized_server>_<tool>`
* 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<ToolResultEvent, ToolResultEventResult>): void;
on(event: "user_bash", handler: ExtensionHandler<UserBashEvent, UserBashEventResult>): void;
on(event: "user_python", handler: ExtensionHandler<UserPythonEvent, UserPythonEventResult>): void;
on(event: "mcp_notification", handler: ExtensionHandler<McpNotificationEvent>): void;
// =========================================================================
// Tool Registration
+120 -18
View File
@@ -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) {
@@ -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);
}
/**
+55 -12
View File
@@ -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 {
@@ -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();
+86
View File
@@ -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<Record<string, unknown>> = Object.freeze({
hello: "world",
n: 42,
});
type JsonRpcMessage = {
jsonrpc: "2.0";
id?: string | number;
method?: string;
params?: Record<string, unknown>;
};
function buildResult(method: string): Record<string, unknown> {
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<string, unknown>): 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();
}
@@ -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<CapturedFrame>;
} {
const frames: CapturedFrame[] = [];
const waiters = new Map<string, Array<(frame: CapturedFrame) => 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<CapturedFrame>();
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<void>();
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<void>();
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<void>();
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);
});
@@ -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);