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:
@@ -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 }`)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -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
@@ -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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user