Add configurable notification buffering

Author: ethan-huo

README.md

@@ -215,6 +215,10 @@ tool and then continues the original command.
 Tool calls use a 5 minute request timeout by default. Set
 `MCPX_TOOL_CALL_TIMEOUT_MS` when a remote MCP server legitimately needs longer.
 
+mcpx buffers MCP subscription notifications and returns them with the next daemon
+tool-call result by default. Set `MCPX_NOTIFICATION_MODE=discard` to ignore those
+notifications for a command; accepted values are `buffer` and `discard`.
+
 ## Output
 
 mcpx optimizes output for humans and agents by default:

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "mcpx",
-	"version": "0.9.11",
+	"version": "0.9.12",
 	"license": "MIT",
 	"bin": {
 		"mcpx": "./src/main.ts"

src/daemon-client.ts

@@ -8,6 +8,7 @@ import {
 	DAEMON_PROTOCOL_VERSION,
 	buildServerKey,
 	helloMessage,
+	notificationModeFromEnv,
 	type ClientMessage,
 	type DaemonMessage,
 	type DaemonStatus,
@@ -68,6 +69,7 @@ export async function callToolViaDaemon(
 		server,
 		toolName,
 		input,
+		notificationMode: notificationModeFromEnv(),
 	}
 	if (context.headers) message.headers = context.headers
 	const response = await requestDaemonMessage(

src/daemon-protocol.ts

@@ -8,6 +8,9 @@ import { MCPX_VERSION } from './version'
 export const DAEMON_PROTOCOL_VERSION = 2
 export const DAEMON_ENV = 'MCPX_DAEMON_SERVER'
 export const DISABLE_DAEMON_ENV = 'MCPX_DISABLE_DAEMON'
+export const NOTIFICATION_MODE_ENV = 'MCPX_NOTIFICATION_MODE'
+
+export type NotificationMode = 'buffer' | 'discard'
 
 export type McpNotification =
 	| {
@@ -63,7 +66,7 @@ export type ClientMessage =
 			headers?: Record<string, string>
 			toolName: string
 			input: Record<string, unknown>
-			notificationMode?: 'buffer' | 'discard'
+			notificationMode?: NotificationMode
 	  }
 	| { op: 'status' }
 	| { op: 'stop' }
@@ -89,6 +92,15 @@ export function shouldUseDaemon(): boolean {
 	)
 }
 
+export function notificationModeFromEnv(): NotificationMode {
+	const raw = process.env[NOTIFICATION_MODE_ENV]
+	if (!raw || raw === 'buffer') return 'buffer'
+	if (raw === 'discard') return 'discard'
+	throw new Error(
+		`Invalid ${NOTIFICATION_MODE_ENV} value "${raw}". Expected "buffer" or "discard".`,
+	)
+}
+
 export function buildServerKey(server: ServerConfig): string {
 	const payload =
 		server.transport === 'stdio'

src/daemon-server.ts

@@ -197,13 +197,13 @@ async function callTool(
 		message.headers,
 		async (session) => {
 			const buffer = createNotificationBuffer()
-			session.currentBuffer =
-				message.notificationMode === 'discard' ? undefined : buffer
+			const bufferNotifications = message.notificationMode !== 'discard'
+			session.currentBuffer = bufferNotifications ? buffer : undefined
 			try {
 				const result = await callToolWithRetainedSessionFallback(
 					session,
 					message,
-					buffer,
+					bufferNotifications ? buffer : undefined,
 				)
 				const notifications = await flushNotifications(buffer)
 				const toolsChanged =
@@ -228,7 +228,7 @@ async function callTool(
 async function callToolWithRetainedSessionFallback(
 	session: ManagedSession,
 	message: Extract<ClientMessage, { op: 'call' }>,
-	buffer: NotificationBuffer,
+	buffer: NotificationBuffer | undefined,
 ): Promise<unknown> {
 	try {
 		return await callToolOnConnectedSession(session, message, buffer)
@@ -249,24 +249,25 @@ async function callToolWithRetainedSessionFallback(
 async function callToolOnConnectedSession(
 	session: ManagedSession,
 	message: Extract<ClientMessage, { op: 'call' }>,
-	buffer: NotificationBuffer,
+	buffer: NotificationBuffer | undefined,
 ): Promise<unknown> {
 	const connection = await ensureConnected(session)
+	const options = toolCallRequestOptions()
+	if (buffer) {
+		options.onprogress = (progress) => {
+			buffer.add({
+				method: 'notifications/progress',
+				params: { progressToken: message.callId, ...progress },
+			})
+		}
+	}
 	return connection.client.callTool(
 		{
 			name: message.toolName,
 			arguments: message.input,
 		},
 		undefined,
-		{
-			onprogress: (progress) => {
-				buffer.add({
-					method: 'notifications/progress',
-					params: { progressToken: message.callId, ...progress },
-				})
-			},
-			...toolCallRequestOptions(),
-		},
+		options,
 	)
 }
 

tests/daemon-protocol.test.ts

@@ -2,7 +2,11 @@ import { describe, expect, it } from 'bun:test'
 
 import type { HttpServerConfig } from '../src/types'
 
-import { buildServerKey } from '../src/daemon-protocol'
+import {
+	NOTIFICATION_MODE_ENV,
+	buildServerKey,
+	notificationModeFromEnv,
+} from '../src/daemon-protocol'
 
 describe('daemon protocol', () => {
 	it('keeps HTTP server keys stable across resolved token values', () => {
@@ -75,4 +79,40 @@ describe('daemon protocol', () => {
 
 		expect(buildServerKey(second)).not.toBe(buildServerKey(first))
 	})
+
+	it('defaults notification buffering unless explicitly discarded', () => {
+		const previous = process.env[NOTIFICATION_MODE_ENV]
+		try {
+			delete process.env[NOTIFICATION_MODE_ENV]
+			expect(notificationModeFromEnv()).toBe('buffer')
+
+			process.env[NOTIFICATION_MODE_ENV] = 'buffer'
+			expect(notificationModeFromEnv()).toBe('buffer')
+
+			process.env[NOTIFICATION_MODE_ENV] = 'discard'
+			expect(notificationModeFromEnv()).toBe('discard')
+		} finally {
+			if (previous === undefined) {
+				delete process.env[NOTIFICATION_MODE_ENV]
+			} else {
+				process.env[NOTIFICATION_MODE_ENV] = previous
+			}
+		}
+	})
+
+	it('rejects invalid notification mode env values', () => {
+		const previous = process.env[NOTIFICATION_MODE_ENV]
+		try {
+			process.env[NOTIFICATION_MODE_ENV] = 'off'
+			expect(() => notificationModeFromEnv()).toThrow(
+				'Invalid MCPX_NOTIFICATION_MODE value "off". Expected "buffer" or "discard".',
+			)
+		} finally {
+			if (previous === undefined) {
+				delete process.env[NOTIFICATION_MODE_ENV]
+			} else {
+				process.env[NOTIFICATION_MODE_ENV] = previous
+			}
+		}
+	})
 })

tests/notification-dogfood.test.ts

@@ -112,6 +112,16 @@ describe('notification fixture dogfood', () => {
 			},
 		])
 
+		const discardedRaw = JSON.parse(
+			(
+				await runMcpx(
+					['notification-fixture', 'progress-stream', '--count', '4', '--raw'],
+					{ MCPX_NOTIFICATION_MODE: 'discard' },
+				)
+			).stdout,
+		)
+		expect(discardedRaw).toEqual({ tool: 'progress-stream', count: 4 })
+
 		const toolsChanged = JSON.parse(
 			(await runMcpx(['notification-fixture', 'notify-tools-changed', '--raw']))
 				.stdout,
@@ -180,12 +190,14 @@ describe('notification fixture dogfood', () => {
 
 async function runMcpx(
 	args: string[],
+	env: Record<string, string> = {},
 ): Promise<{ stdout: string; stderr: string }> {
 	const proc = Bun.spawn([process.execPath, mainPath, ...args], {
 		env: {
 			...process.env,
 			HOME: home,
 			MCPX_HOME: home,
+			...env,
 		},
 		stdin: 'ignore',
 		stdout: 'pipe',