diff --git a/docs/developers/qwen-serve-protocol.md b/docs/developers/qwen-serve-protocol.md index f6b97771d91..7198a41eff3 100644 --- a/docs/developers/qwen-serve-protocol.md +++ b/docs/developers/qwen-serve-protocol.md @@ -101,6 +101,7 @@ registry. Clients **must** gate UI off `features`, not off `mode` (per design 'workspace_providers', 'workspace_env', 'workspace_preflight', 'session_context', 'session_supported_commands', 'session_close', 'session_metadata', 'mcp_guardrails', + 'mcp_guardrail_events', 'workspace_file_read', 'workspace_file_bytes', 'workspace_file_write', 'session_approval_mode_control', 'workspace_tool_toggle', 'workspace_init', 'workspace_mcp_restart'] @@ -141,6 +142,13 @@ routes and require a configured bearer token even on loopback. `mcp_guardrails` is **not** in this conditional table — it's an always-on tag, advertised whenever the binary supports the new `/workspace/mcp` budget fields, regardless of whether the operator configured a budget. Operators who haven't set `--mcp-client-budget` still get the new fields (with `budgetMode: 'off'`, `budgets: []`). +`mcp_guardrail_events` (issue [#4175](https://github.com/QwenLM/qwen-code/issues/4175) PR 14b) advertises the typed SSE push events that surface MCP budget state crossings without a poll loop. Two frame types arrive on `GET /session/:id/events`: + +- `mcp_budget_warning` — fires once on the upward 75% crossing of `reservedSlots.size / clientBudget`. Re-arms only after the ratio drops below 37.5% (`MCP_BUDGET_REARM_FRACTION`). Mirrors PR 10's `slow_client_warning` hysteresis, but at the manager level rather than the per-subscriber backlog level. Payload: `{ liveCount, reservedCount, budget, thresholdRatio: 0.75, mode: 'warn' | 'enforce' }`. Fires under both `warn` and `enforce` modes; never under `off`. +- `mcp_child_refused_batch` — fires at end of each `discoverAllMcpTools*` pass when one or more servers were refused, AND as a length-1 batch on the `readResource` lazy-spawn refusal path. Payload: `{ refusedServers: [{ name, transport, reason: 'budget_exhausted' }, ...], budget, liveCount, reservedCount, mode: 'enforce' }`. `mode` is the literal `'enforce'` because `warn` mode never refuses. + +Both events live in the per-session SSE replay ring (they carry an `id`) so a client reconnecting with `Last-Event-ID` resumes through them; the snapshot at `GET /workspace/mcp` is still the source-of-truth for state-after-extended-disconnect. Always-on once advertised — there is no conditional toggle. SDK reducer state (`DaemonSessionViewState`) exposes `mcpBudgetWarningCount`, `lastMcpBudgetWarning`, `mcpChildRefusedBatchCount`, `lastMcpChildRefusedBatch` for adapters that want simple lag-style UI. + ## Routes ### `GET /health` @@ -332,13 +340,16 @@ vars only; proxy URLs are stripped of credentials and reduced to Budget enforcement in PR 14 v1 is **per-session, not per-workspace**. Although Mode B daemons are `1 daemon = 1 workspace × N sessions` post-#4113 at the process level, the `McpClientManager` is constructed inside each ACP session's `Config` via `acpAgent.newSessionConfig`, so N sessions each enforce their own copy of the cap. The snapshot represents the bootstrap session's view. Wave 5 PR 23 introduces a workspace-scoped shared MCP pool that graduates this to true per-workspace enforcement. -**Detecting budget pressure in v1 (no push events yet).** PR 14 v1 is snapshot-only — typed SSE push events (`mcp_budget_warning` + `mcp_child_refused_batch`) ship in PR 14b. Until then, operator dashboards poll `GET /workspace/mcp` and inspect the per-session budget cell (`budgets[0]`): +**Detecting budget pressure.** Two surfaces, both populated post-PR-14b: + +- **Push events** (advertised via `mcp_guardrail_events`): subscribe to `GET /session/:id/events` and narrow `mcp_budget_warning` / `mcp_child_refused_batch` frames through `KnownDaemonEvent`. The state machine fires once per upward 75% crossing (re-armed below 37.5%); refusals are coalesced once per discovery pass under `enforce` mode. +- **Snapshot poll** (advertised via `mcp_guardrails`): `GET /workspace/mcp` and inspect the per-session budget cell (`budgets[0]`): - `budgets[0].status === 'warning'` ⇔ `liveCount >= 0.75 * clientBudget` (matches the hysteresis threshold PR 14b's push event will use). - `budgets[0].status === 'error'` ⇔ `refusedCount > 0` (one or more servers refused this discovery pass). - `budgets[0].status === 'ok'` ⇔ below the 75% threshold AND no refusals. -Recommended poll cadence: aligned with whatever already polls `/workspace/mcp`; the snapshot is cheap and the budget cell carries no extra discovery cost. +Recommended poll cadence: aligned with whatever already polls `/workspace/mcp`; the snapshot is cheap and the budget cell carries no extra discovery cost. SDK clients that subscribe to push events still benefit from the snapshot for state-after-extended-disconnect (the SSE replay ring depth is finite — `--event-ring-size`, default 8000 — so a client offline longer than the ring's coverage falls back to snapshot resync). ### `GET /workspace/skills` diff --git a/docs/users/qwen-serve.md b/docs/users/qwen-serve.md index 2a4ec170c14..638e971e111 100644 --- a/docs/users/qwen-serve.md +++ b/docs/users/qwen-serve.md @@ -209,6 +209,8 @@ The token comparison is constant-time (SHA-256 + `crypto.timingSafeEqual`); 401 > ``` > > This is **not** the same as claude-code's `MCP_SERVER_CONNECTION_BATCH_SIZE` (which gates startup concurrency); they're orthogonal. PR 23 will add a real shared MCP pool (a `scope: 'workspace'` cell in `budgets[]` alongside the per-session cell); PR 14 v1 is the in-process counter + soft enforcement on the existing per-session manager. +> +> **Push events (issue [#4175](https://github.com/QwenLM/qwen-code/issues/4175) PR 14b).** SDK clients subscribed to `GET /session/:id/events` receive typed frames when budget thresholds cross — `mcp_budget_warning` (synthetic, fires once per upward 75% crossing with hysteresis re-arm at 37.5%, advertised via `mcp_guardrail_events`) and `mcp_child_refused_batch` (coalesced once per discovery pass under `enforce` mode; length-1 from `readResource` lazy-spawn refusal). The snapshot at `GET /workspace/mcp` is still the source-of-truth for state-after-reconnect; events are change-edges. Useful when dashboarding in real-time without polling. ## Default deployment threat model diff --git a/integration-tests/cli/qwen-serve-baseline.test.ts b/integration-tests/cli/qwen-serve-baseline.test.ts index b6cfba1027f..97ca6b48d92 100644 --- a/integration-tests/cli/qwen-serve-baseline.test.ts +++ b/integration-tests/cli/qwen-serve-baseline.test.ts @@ -462,6 +462,78 @@ async function measureRssAtSessionCount(sessionCount: number): Promise<{ fs.rmSync(ws, { recursive: true, force: true }); } }, 120_000); + + // PR 14b cross-check: validate the daemon's in-process MCP + // accounting against external `pgrep -P` measurement. The + // snapshot at `GET /workspace/mcp` exposes `clientCount` + // (live CONNECTED clients, `getMcpClientAccounting().total`) + // — that's the field SDK consumers and dashboards actually + // see, and it's the same source the push-event channel + // (`mcp_budget_warning` / `mcp_child_refused_batch`) reads. + // If `clientCount` diverges from what `pgrep -P` observes for + // the daemon's MCP grandchildren, the events lie. + // + // (Codex round 3 doc fix — codex/copilot finding: this test + // was named "in-process subprocessCount matches external pgrep" + // but actually asserts on `clientCount`. The snapshot's + // `clientCount` and the manager-internal `subprocessCount` + // (`stdio + websocket`) match for stdio-only fixtures, so the + // assertion is numerically correct — but the test name now + // matches the field it actually validates.) + // + // Skip-gated like the parent describe (POSIX, non-sandbox); + // idle MCP fixtures are stdio-only so `clientCount` should + // equal `mcpGrandchildren.length` exactly (no amplification + // slack required). + it('clientCount matches external pgrep observation', async () => { + const ws = makeTempWorkspace('mcp-counter'); + let daemon: SpawnedDaemon | undefined; + try { + writeWorkspaceSettings(ws, { + mcpServers: { + idle1: { command: 'node', args: [IDLE_MCP_PATH] }, + idle2: { command: 'node', args: [IDLE_MCP_PATH] }, + }, + }); + daemon = await spawnDaemon({ workspaceCwd: ws }); + await daemon.client.createOrAttachSession({ workspaceCwd: ws }); + + // Wait for MCP grandchildren to be observable via pgrep, + // then read both numbers atomically (pgrep first to lock + // the comparison floor, snapshot second so the daemon + // can't sneak in a new connect between the two reads). + const observed = await waitForMcpGrandchildren( + daemon.daemon.pid!, + MCP_SERVERS_CONFIGURED, + ); + const snapshot = await daemon.client.workspaceMcp(); + + // PR 14b invariant: stdio-only fixtures → + // `clientCount === mcpGrandchildren.length`. The PR 14 + // amplification slack + // (`MCP_SERVERS_CONFIGURED * mcpAmplificationFactor`) is + // for connect-storm transient overhead, not steady-state + // counter drift. At idle the daemon's accounting MUST + // match `pgrep -P` exactly (no zombies, no orphans). + // + // `clientCount` is the snapshot's authoritative live + // count; validating it against pgrep closes the loop on + // PR 14b's event-source assumption (the push events read + // the same accounting). + expect(snapshot.clientCount).toBe(MCP_SERVERS_CONFIGURED); + expect(observed.mcpGrandchildren.length).toBe(MCP_SERVERS_CONFIGURED); + // Defense-in-depth: even if a future race lets the OS + // observe a process the daemon already considers + // disconnected, `clientCount` must NEVER exceed the + // observed pgrep count. Equality at idle, `<=` always. + expect(snapshot.clientCount).toBeLessThanOrEqual( + observed.mcpGrandchildren.length, + ); + } finally { + if (daemon) await daemon.dispose(); + fs.rmSync(ws, { recursive: true, force: true }); + } + }, 120_000); }); describe('SSE backpressure (unit)', () => { diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index 1580731d567..59596c3d530 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -2095,6 +2095,155 @@ describe('QwenAgent MCP SSE/HTTP support', () => { mockConnectionState.resolve(); await agentPromise; }); + + // PR 14b: budget-event push channel. After codex review fix #2, the + // callback is wired via `Config.setMcpBudgetEventCallback` BEFORE + // `config.initialize()`, so MCP discovery (which can fire events + // synchronously in legacy blocking mode and races with background + // discovery in progressive mode) sees the callback wired from the + // first pass. The Config-level shim stashes the callback and applies + // it inside `createToolRegistry` to the freshly-constructed manager. + it('newSession wires Config.setMcpBudgetEventCallback BEFORE initialize() (codex fix #2)', async () => { + const sessionId = 'session-budget-events'; + const innerConfig = await setupSessionMocks(sessionId); + // Stub `setMcpBudgetEventCallback` on the inner Config. The + // production path delegates the manager apply to Config; the test + // captures the callback at the Config boundary and verifies the + // ordering vs `initialize()`. + let capturedCallback: + | ((event: Record) => void) + | undefined; + const callOrder: string[] = []; + (innerConfig as unknown as Record)[ + 'setMcpBudgetEventCallback' + ] = vi.fn((cb: (event: Record) => void) => { + callOrder.push('setMcpBudgetEventCallback'); + capturedCallback = cb; + }); + // Wrap `initialize` to record its position in `callOrder`. The + // critical invariant codex review fix #2 enforces: setter runs + // BEFORE initialize. + const originalInitialize = innerConfig.initialize; + innerConfig.initialize = vi.fn().mockImplementation(async () => { + callOrder.push('initialize'); + return originalInitialize(); + }); + + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + + // Spy connection: only `extNotification` is exercised here, but + // the AgentSideConnection contract is wide. Stubbing only what the + // PR 14b code path touches keeps the test focused. + const extNotification = vi.fn().mockResolvedValue(undefined); + const fakeConn = { + get closed() { + return mockConnectionState.promise; + }, + extNotification, + }; + const agent = capturedAgentFactory!( + fakeConn as unknown as AgentSideConnectionLike, + ) as AgentLike; + + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + + // Strict ordering invariant — codex review fix #2. + expect(callOrder).toEqual(['setMcpBudgetEventCallback', 'initialize']); + expect(typeof capturedCallback).toBe('function'); + + // Fire a synthetic budget_warning through the captured callback — + // the wired extNotification must receive the same shape with + // `sessionId` inserted and `v: 1` envelope. + const warningEvent = { + kind: 'budget_warning' as const, + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75 as const, + mode: 'warn' as const, + }; + capturedCallback!(warningEvent); + + expect(extNotification).toHaveBeenCalledTimes(1); + expect(extNotification).toHaveBeenCalledWith( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId, + ...warningEvent, + }, + ); + + // Fire a refused_batch through the same callback — same routing, + // discriminated union shape preserved verbatim. + const refusedEvent = { + kind: 'refused_batch' as const, + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce' as const, + }; + capturedCallback!(refusedEvent); + + expect(extNotification).toHaveBeenCalledTimes(2); + expect(extNotification).toHaveBeenLastCalledWith( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId, + ...refusedEvent, + }, + ); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('newSession is a no-op for budget wiring when setMcpBudgetEventCallback is absent (defensive)', async () => { + // Codex review fix #2: the wiring path now goes through + // `Config.setMcpBudgetEventCallback`, not the manager directly. + // Older / stubbed `Config` shapes may omit it; the `typeof check` + // in newSessionConfig keeps the absence silent. + const innerConfig = await setupSessionMocks('session-no-cb-setter'); + // `setupSessionMocks`/`makeInnerConfig` returns a Config without + // `setMcpBudgetEventCallback` defined — that's the defensive case. + + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + + const extNotification = vi.fn().mockResolvedValue(undefined); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + extNotification, + } as unknown as AgentSideConnectionLike) as AgentLike; + + await agent.newSession({ cwd: '/tmp', mcpServers: [] }); + + // No setter on Config → no wiring → no extNotification fires. + expect( + (innerConfig as unknown as Record)[ + 'setMcpBudgetEventCallback' + ], + ).toBeUndefined(); + expect(extNotification).not.toHaveBeenCalled(); + + mockConnectionState.resolve(); + await agentPromise; + }); }); // Regression coverage for the MR-review finding that ACP renameSession diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index 565ab2dc587..517259a075e 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -1812,6 +1812,66 @@ class QwenAgent implements Agent { projectHooks: this.settings.getProjectHooks(), }, ); + // PR 14b fix #2 (codex review round 1): register the MCP guardrail + // budget-event callback BEFORE `config.initialize()`. Pre-fix the + // registration ran AFTER initialize, which (a) missed end-of-pass + // events under `QWEN_CODE_LEGACY_MCP_BLOCKING=1` (synchronous + // discovery completes inside initialize, before our setter runs) + // and (b) raced against background-discovery completion under the + // default progressive mode. `Config.setMcpBudgetEventCallback` + // stashes the callback and `createToolRegistry` applies it to the + // manager BEFORE `discoverAllTools` / `startMcpDiscoveryInBackground` + // fires, closing both windows. + // + // sessionId source: `config.getSessionId()` reads the Config's own + // session id (auto-assigned via `randomUUID()` in the Config + // constructor when no override is passed — see `config.ts:849`), + // so the value is available immediately after `loadCliConfig` + // returns. The closure pins it for the manager's whole lifetime. + // + // Defensive `typeof` checks tolerate stub Configs / ToolRegistries + // in older tests (older fixtures may omit `setMcpBudgetEventCallback` + // or `getSessionId`). + const wiredSessionId = + typeof config.getSessionId === 'function' + ? config.getSessionId() + : undefined; + if ( + typeof config.setMcpBudgetEventCallback === 'function' && + wiredSessionId !== undefined + ) { + const sid = wiredSessionId; + config.setMcpBudgetEventCallback((event) => { + // Fire-and-forget: `extNotification` returns Promise but + // the manager's call site doesn't await. `.catch` suppresses + // unhandled rejections — a mid-flight ACP disconnect would + // otherwise crash the child. Snapshot still carries the state + // for clients that reconnect. + // + // PR 14b fix (codex round 3 — DeepSeek): pre-fix the catch + // handler was `() => {}`, silently dropping every error + // including "real" ones (serialization bugs, protocol + // violations) — operators had no debug trail. Now logs at + // `debug` level: ACP channel closure during shutdown is the + // expected case and would spam at higher levels, but `debug` + // is opt-in so when an oncall engineer DOES turn it on for + // an MCP guardrail incident, they see exactly which event + // dropped and why. + void this.connection + .extNotification('qwen/notify/session/mcp-budget-event', { + v: 1, + sessionId: sid, + ...event, + }) + .catch((err: unknown) => { + debugLogger.debug( + `MCP budget extNotification dropped ` + + `(session=${sid}, kind=${event.kind}): ` + + `${err instanceof Error ? err.message : String(err)}`, + ); + }); + }); + } await config.initialize(); // Same reasoning as the top-level runAcpAgent path: ACP feeds session // messages to the model immediately, so we cannot return a Config whose diff --git a/packages/cli/src/serve/capabilities.ts b/packages/cli/src/serve/capabilities.ts index 1d17adf01b4..1509dc6b520 100644 --- a/packages/cli/src/serve/capabilities.ts +++ b/packages/cli/src/serve/capabilities.ts @@ -96,6 +96,18 @@ export const SERVE_CAPABILITY_REGISTRY = { // `require_auth` is the only conditional tag, kept last for // visibility in `Object.keys(SERVE_CAPABILITY_REGISTRY)`. mcp_guardrails: { since: 'v1', modes: ['warn', 'enforce'] }, + // Issue #4175 PR 14b. Daemon emits typed push events for MCP budget + // state crossings: `mcp_budget_warning` (synthetic, fires once per + // upward 75% crossing with hysteresis re-arm at 37.5%) and + // `mcp_child_refused_batch` (coalesced, one per discovery pass / + // length-1 per readResource refusal, only in `enforce` mode). SDK + // reducer narrows both via `KnownDaemonEvent` (`DaemonSessionViewState` + // exposes `mcpBudgetWarningCount`, `lastMcpBudgetWarning`, + // `mcpChildRefusedBatchCount`, `lastMcpChildRefusedBatch`). Always-on once + // PR 14b lands; orthogonal to `mcp_guardrails` (the snapshot + // surface). Listed alongside `mcp_guardrails` to keep the MCP-related + // tags grouped. + mcp_guardrail_events: { since: 'v1' }, // Issue #4175 PR 19. Daemon supports the read-only workspace file // surface: `GET /file`, `GET /list`, `GET /glob`, `GET /stat`. The // four routes are gated as a single feature because they share the diff --git a/packages/cli/src/serve/httpAcpBridge.test.ts b/packages/cli/src/serve/httpAcpBridge.test.ts index d1f46da1995..7083738ab2e 100644 --- a/packages/cli/src/serve/httpAcpBridge.test.ts +++ b/packages/cli/src/serve/httpAcpBridge.test.ts @@ -4787,6 +4787,510 @@ describe('createHttpAcpBridge', () => { }); }); + // PR 14b: ext-notification handler for child→bridge MCP budget events. + // Translates `qwen/notify/session/mcp-budget-event` into session-scoped + // SSE frames (`mcp_budget_warning` / `mcp_child_refused_batch`). + describe('extNotification — MCP budget events (PR 14b)', () => { + it('publishes mcp_budget_warning when the child fires the warning event', async () => { + let capturedConn: AgentSideConnection | undefined; + const factory: ChannelFactory = async () => { + const { clientStream, agentStream } = createInMemoryChannel(); + const fakeAgent = new FakeAgent(); + capturedConn = new AgentSideConnection(() => fakeAgent, agentStream); + return { + stream: clientStream, + exited: new Promise< + | { exitCode: number | null; signalCode: NodeJS.Signals | null } + | undefined + >(() => {}), + kill: async () => {}, + killSync: () => {}, + }; + }; + const bridge = makeBridge({ channelFactory: factory }); + const session = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + + const abort = new AbortController(); + const iter = bridge.subscribeEvents(session.sessionId, { + signal: abort.signal, + }); + + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId: session.sessionId, + kind: 'budget_warning', + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + ); + + const collected: Array<{ id?: number; type: string; data: unknown }> = []; + for await (const e of iter) { + collected.push({ id: e.id, type: e.type, data: e.data }); + if (collected.length === 1) break; + } + expect(collected[0]?.type).toBe('mcp_budget_warning'); + // PR 14b drops the routing fields (`v`, `sessionId`, `kind`) + // from `data` since the SSE envelope already encodes them. + expect(collected[0]?.data).toEqual({ + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }); + expect(collected[0]?.id).toBe(1); + + abort.abort(); + await bridge.shutdown(); + }); + + it('publishes mcp_child_refused_batch when the child fires the refused-batch event', async () => { + let capturedConn: AgentSideConnection | undefined; + const factory: ChannelFactory = async () => { + const { clientStream, agentStream } = createInMemoryChannel(); + const fakeAgent = new FakeAgent(); + capturedConn = new AgentSideConnection(() => fakeAgent, agentStream); + return { + stream: clientStream, + exited: new Promise< + | { exitCode: number | null; signalCode: NodeJS.Signals | null } + | undefined + >(() => {}), + kill: async () => {}, + killSync: () => {}, + }; + }; + const bridge = makeBridge({ channelFactory: factory }); + const session = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + + const abort = new AbortController(); + const iter = bridge.subscribeEvents(session.sessionId, { + signal: abort.signal, + }); + + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId: session.sessionId, + kind: 'refused_batch', + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }, + ); + + const collected: Array<{ type: string; data: unknown }> = []; + for await (const e of iter) { + collected.push({ type: e.type, data: e.data }); + if (collected.length === 1) break; + } + expect(collected[0]?.type).toBe('mcp_child_refused_batch'); + expect(collected[0]?.data).toEqual({ + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }); + + abort.abort(); + await bridge.shutdown(); + }); + + it('drops unknown extNotification methods, kinds, and missing sessionIds silently', async () => { + let capturedConn: AgentSideConnection | undefined; + const factory: ChannelFactory = async () => { + const { clientStream, agentStream } = createInMemoryChannel(); + const fakeAgent = new FakeAgent(); + capturedConn = new AgentSideConnection(() => fakeAgent, agentStream); + return { + stream: clientStream, + exited: new Promise< + | { exitCode: number | null; signalCode: NodeJS.Signals | null } + | undefined + >(() => {}), + kill: async () => {}, + killSync: () => {}, + }; + }; + const bridge = makeBridge({ channelFactory: factory }); + const session = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + + const abort = new AbortController(); + const iter = bridge.subscribeEvents(session.sessionId, { + signal: abort.signal, + }); + + // Unknown method — drop. + void capturedConn!.extNotification('qwen/notify/session/unknown-event', { + sessionId: session.sessionId, + kind: 'budget_warning', + }); + // Missing sessionId — drop. + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { kind: 'budget_warning' }, + ); + // Unknown kind — drop. + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { sessionId: session.sessionId, kind: 'mystery_kind' }, + ); + // Resolvable sessionId but session id doesn't exist — drop. + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { + sessionId: 'nonexistent', + kind: 'budget_warning', + liveCount: 1, + reservedCount: 1, + budget: 1, + thresholdRatio: 0.75, + mode: 'warn', + }, + ); + // Real event — must arrive AFTER all drops above. + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId: session.sessionId, + kind: 'budget_warning', + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + ); + + const collected: Array<{ type: string }> = []; + for await (const e of iter) { + collected.push({ type: e.type }); + if (collected.length === 1) break; + } + // Exactly one event got through. Codex review fix #1 changed + // the "unknown sessionId" path from drop to buffer — the + // `nonexistent` frame above is now sitting in the early-event + // buffer (it never registers, so it'll TTL out). All other + // drops (unknown method, missing sessionId, unknown kind) + // remain hard-drops. + expect(collected).toEqual([{ type: 'mcp_budget_warning' }]); + + abort.abort(); + await bridge.shutdown(); + }); + + it('buffers events for a not-yet-registered sessionId, drains them on registration (codex fix #1)', async () => { + // Codex review round 1, finding #1: budget events fired during + // a session's startup window (between `connection.newSession` + // dispatching and `byId.set`) reach `BridgeClient.extNotification` + // with a valid sessionId but no matching entry. Pre-fix those + // were dropped silently; post-fix they're buffered and replayed + // via `drainEarlyEvents` so SSE subscribers see them as the + // FIRST frames of the new session. + // + // This test exercises the buffer + drain mechanism directly, + // pre-buffering for a sessionId that doesn't yet exist, then + // creating that session via newSessionImpl-controlled id and + // verifying the drain replayed the frame onto the new EventBus. + // (Forcing the actual production race window is timing-flaky; + // the mechanism is the invariant we care about.) + let capturedConn: AgentSideConnection | undefined; + // Use sessionScope: 'thread' + a deterministic id-prefix so + // `spawnOrAttach` returns an id we can pre-target. + const factory: ChannelFactory = async () => { + const { clientStream, agentStream } = createInMemoryChannel(); + const fakeAgent = new FakeAgent({ sessionIdPrefix: 'pre-buffer' }); + capturedConn = new AgentSideConnection(() => fakeAgent, agentStream); + return { + stream: clientStream, + exited: new Promise< + | { exitCode: number | null; signalCode: NodeJS.Signals | null } + | undefined + >(() => {}), + kill: async () => {}, + killSync: () => {}, + }; + }; + const bridge = makeBridge({ + channelFactory: factory, + sessionScope: 'thread', + }); + + // Boot ANY session first to get the channel + BridgeClient + // alive (factory + AgentSideConnection are constructed lazily + // on first spawn). After this, subsequent spawns share the + // channel and BridgeClient. + const seed = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + // Pre-buffer for the NEXT thread-scope session id. FakeAgent + // names them `:#`; the seed was call 1 + // (suffix ''), the next will be call 2 (suffix '#2'). + const futureSessionId = `pre-buffer:${WS_A}#2`; + expect(seed.sessionId).not.toBe(futureSessionId); + + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId: futureSessionId, + kind: 'budget_warning', + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + ); + + // Give the bridge's reader loop a tick to dispatch the + // notification onto BridgeClient.extNotification — it goes + // through `bufferEarlyEvent` because `futureSessionId` isn't + // in `byId` yet. + await new Promise((r) => setTimeout(r, 50)); + + // Now create the future session. `createSessionEntry`'s new + // `drainEarlyEvents` call replays the buffered frame onto the + // freshly-constructed EventBus. + const target = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + expect(target.sessionId).toBe(futureSessionId); + + // Subscribe with `lastEventId: 0` so the replay-ring drain + // path runs (live-only subscriptions skip the ring per + // `eventBus.ts` semantics). Production SSE clients reconnecting + // with `Last-Event-ID: 0` get this same behavior. + const abort = new AbortController(); + const iter = bridge.subscribeEvents(target.sessionId, { + signal: abort.signal, + lastEventId: 0, + }); + const collected: Array<{ id?: number; type: string }> = []; + for await (const e of iter) { + collected.push({ id: e.id, type: e.type }); + if (collected.length === 1) break; + } + expect(collected[0]?.type).toBe('mcp_budget_warning'); + // Drained frame went through `events.publish`, so it gets an + // `id` — PR 14b events are session-scoped + replayable. + expect(collected[0]?.id).toBe(1); + + abort.abort(); + await bridge.shutdown(); + }); + + it('tombstones closed sessionIds so late notifications cannot leak into a future load of the same id (codex round 5 fix)', async () => { + // Codex round 5 finding: pre-fix, after a session was killed + // / closed, a late `extNotification` from its dying child for + // the same id would land in `earlyEvents`. If the SAME + // sessionId came back via `session/load`/`session/resume` + // within the 60s TTL, `drainEarlyEvents` would replay stale + // prior-session telemetry onto the NEW subscriber. + // + // Fix: every `byId.delete(sid)` site now calls + // `BridgeClient.markSessionClosed(sid)`, which tombstones the + // id (rejecting future `bufferEarlyEvent` calls for it) and + // purges any frames already buffered for it. + let capturedConn: AgentSideConnection | undefined; + const factory: ChannelFactory = async () => { + const { clientStream, agentStream } = createInMemoryChannel(); + const fakeAgent = new FakeAgent({ + loadSessionImpl: () => ({ configOptions: [] }), + }); + capturedConn = new AgentSideConnection(() => fakeAgent, agentStream); + return { + stream: clientStream, + exited: new Promise< + | { exitCode: number | null; signalCode: NodeJS.Signals | null } + | undefined + >(() => {}), + kill: async () => {}, + killSync: () => {}, + }; + }; + const bridge = makeBridge({ channelFactory: factory }); + + // 1) Spawn session A — id = SESS_A. + const sess = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + const sessionId = sess.sessionId; + expect(sessionId).toBe(SESS_A); + + // 2) Close session A — calls byId.delete + markSessionClosed. + await bridge.closeSession(sessionId); + + // 3) Simulate a LATE notification from the (now-defunct) + // child for the closed sessionId. Pre-fix this would land in + // `earlyEvents`. Post-fix the tombstone rejects it. + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId, + kind: 'budget_warning', + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + ); + // Give the bridge's read loop time to dispatch the notification. + await new Promise((r) => setTimeout(r, 50)); + + // 4) Re-load the SAME persisted sessionId via session/load. + // createSessionEntry runs drainEarlyEvents — pre-fix the stale + // frame would be replayed onto the new session's bus. + const loaded = await bridge.loadSession({ + sessionId, + workspaceCwd: WS_A, + }); + expect(loaded.sessionId).toBe(sessionId); + + // 5) Subscribe with lastEventId: 0 to drain the replay ring. + // Post-fix, no `mcp_budget_warning` should be in the ring + // (the late notification was dropped at buffer time, not + // drained on registration). + const abort = new AbortController(); + const iter = bridge.subscribeEvents(loaded.sessionId, { + signal: abort.signal, + lastEventId: 0, + }); + const collected: Array<{ type: string }> = []; + const drainPromise = (async () => { + for await (const e of iter) { + collected.push({ type: e.type }); + } + })(); + // Give the iterator a tick to pull replay frames. + await new Promise((r) => setTimeout(r, 50)); + abort.abort(); + await drainPromise; + + // No mcp_budget_warning leaked through. + expect(collected.filter((e) => e.type === 'mcp_budget_warning')).toEqual( + [], + ); + + await bridge.shutdown(); + }); + + it('purges buffered guardrail events when restore fails so retry-success does not replay stale frames (codex round 7 fix)', async () => { + // Codex round 7 finding: round-6 added `markRestoreInFlight` + // so `bufferEarlyEvent` accepts frames for tombstoned ids + // during a restore. If the restore FAILS, pre-fix + // `clearRestoreInFlight` only released the allow-list and + // left buffered frames in `earlyEvents[id]`. A subsequent + // successful retry (`session/load` of the same id within + // 60s) would `drainEarlyEvents` those stale frames into the + // new session. + // + // Fix: failure path now calls `markSessionClosed` which both + // re-tombstones the id AND purges `earlyEvents[id]`. + let capturedConn: AgentSideConnection | undefined; + let loadAttempt = 0; + const factory: ChannelFactory = async () => { + const { clientStream, agentStream } = createInMemoryChannel(); + // First load attempt fails; second attempt succeeds. The + // child's notification fires DURING the failing first + // attempt — pre-fix it would survive the failure. + const fakeAgent = new FakeAgent({ + loadSessionImpl: async (req, agent) => { + loadAttempt += 1; + if (loadAttempt === 1) { + // Buffer a guardrail event for this restore window + // BEFORE failing, simulating the round-6-allow-list + // behavior. + void agent; + void capturedConn!.extNotification( + 'qwen/notify/session/mcp-budget-event', + { + v: 1, + sessionId: req.sessionId, + kind: 'budget_warning', + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + ); + // Tiny yield so the bridge dispatches the notification + // before we throw. + await new Promise((r) => setTimeout(r, 5)); + throw new Error('simulated transient load failure'); + } + return { configOptions: [] }; + }, + }); + capturedConn = new AgentSideConnection(() => fakeAgent, agentStream); + return { + stream: clientStream, + exited: new Promise< + | { exitCode: number | null; signalCode: NodeJS.Signals | null } + | undefined + >(() => {}), + kill: async () => {}, + killSync: () => {}, + }; + }; + const bridge = makeBridge({ channelFactory: factory }); + + // Pre-tombstone: spawn + close session with the id we'll later load. + const sess = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + const sessionId = sess.sessionId; + await bridge.closeSession(sessionId); + + // First load — fails after the child queues a guardrail event. + // ACP wraps the agent throw as a JSON-RPC "Internal error"; + // the original message lives in `data.details` but the assertion + // only needs to verify the load rejected. + await expect( + bridge.loadSession({ sessionId, workspaceCwd: WS_A }), + ).rejects.toThrow(); + + // Retry — succeeds. Pre-fix this would replay the queued + // guardrail event onto the new session's bus. + const loaded = await bridge.loadSession({ + sessionId, + workspaceCwd: WS_A, + }); + expect(loaded.sessionId).toBe(sessionId); + + // Verify no stale guardrail event leaked. + const abort = new AbortController(); + const iter = bridge.subscribeEvents(loaded.sessionId, { + signal: abort.signal, + lastEventId: 0, + }); + const collected: Array<{ type: string }> = []; + const drainPromise = (async () => { + for await (const e of iter) { + collected.push({ type: e.type }); + } + })(); + await new Promise((r) => setTimeout(r, 50)); + abort.abort(); + await drainPromise; + expect(collected.filter((e) => e.type === 'mcp_budget_warning')).toEqual( + [], + ); + + await bridge.shutdown(); + }); + }); + describe('maxSessions cap (chiga0 Rec 3)', () => { it('refuses NEW spawns past the cap with SessionLimitExceededError', async () => { let n = 0; diff --git a/packages/cli/src/serve/httpAcpBridge.ts b/packages/cli/src/serve/httpAcpBridge.ts index 61c067dd1e0..0d613aff175 100644 --- a/packages/cli/src/serve/httpAcpBridge.ts +++ b/packages/cli/src/serve/httpAcpBridge.ts @@ -1204,6 +1204,28 @@ export class InvalidPermissionOptionError extends Error { const MAX_DISPLAY_NAME_LENGTH = 256; +/** + * PR 14b fix #1 (codex review round 1): bounded buffering for ACP + * `extNotification` frames that arrive on `BridgeClient` before the + * matching session has been registered in `byId`. The bridge populates + * `byId` only AFTER `connection.newSession` returns, but the child's + * MCP discovery runs INSIDE `newSession` and may fire budget events + * synchronously before the response makes it back. Without buffering, + * those frames hit `resolveEntry → undefined` and are silently dropped + * — the very first replay-ring slot for the new session is missing + * the events that fired during its creation. + * + * The triple bound (max sessions × max events per session × TTL) + * caps worst-case heap retention even if a malicious / buggy child + * spammed `extNotification` for sessionIds that never register: + * 64 × 32 × ~200B ≈ 400 KB total. TTL is generous (60s — far longer + * than realistic session creation latency of seconds) so brief + * scheduling pauses don't cause real warnings to be evicted. + */ +const MAX_EARLY_EVENT_SESSIONS = 64; +const MAX_EARLY_EVENTS_PER_SESSION = 32; +const EARLY_EVENT_TTL_MS = 60_000; + function hasControlCharacter(value: string): boolean { for (let i = 0; i < value.length; i += 1) { const code = value.charCodeAt(i); @@ -1472,6 +1494,301 @@ class BridgeClient implements Client { }); } + /** + * PR 14b fix #1 (codex review round 1): bounded early-event buffer. + * Frames are keyed by sessionId; each entry tracks its `expiresAt` + * for lazy TTL-based eviction in `bufferEarlyEvent`. Drained by + * `drainEarlyEvents` whenever the bridge registers a session with + * a matching id. See MAX_EARLY_EVENT_* constants for capacity + * bounds. + */ + private readonly earlyEvents = new Map< + string, + { + frames: Array>; + expiresAt: number; + } + >(); + + /** + * PR 14b fix (codex review round 5): tombstone for closed/killed + * session ids. Pre-fix, `extNotification` buffered events for any + * unknown sessionId — including ids of just-closed sessions whose + * dying child fired one last `extNotification` between + * `byId.delete(sid)` and the channel actually exiting. If the SAME + * id was later re-registered via `session/load` or `session/resume` + * within the buffer's 60s TTL, `drainEarlyEvents` would replay + * stale prior-session telemetry (false budget warnings, refused + * server names from the OLD session) onto the NEW subscriber. + * + * Tombstone semantics: + * - Marked when the bridge removes a sessionId from `byId` (kill + * path, channel.exited handler, closeSession). + * - Concurrently purges any in-flight `earlyEvents[id]` so a + * buffered-but-undelivered frame can't leak either. + * - `bufferEarlyEvent` rejects tombstoned ids (the dying child's + * late notification just gets dropped). + * - `drainEarlyEvents` clears the tombstone — a fresh + * `createSessionEntry` for the same id is the legitimate + * "load/resume of a persisted session id" case, and at that + * point any stale event has already been rejected at buffer time. + * - TTL = `EARLY_EVENT_TTL_MS` (60s) — same as the early-event + * buffer, so by the time a tombstone expires there can be no + * stale frame for that id anywhere in the system. + */ + private readonly tombstonedSessionIds = new Map(); + + /** + * PR 14b fix (codex review round 6): allow-list of sessionIds that + * are currently being restored via `session/load` / + * `session/resume`. Bypasses the tombstone check in + * `bufferEarlyEvent` so restore-time guardrail events for a + * previously-closed id flow through to the future + * `createSessionEntry → drainEarlyEvents` call. + * + * Pre-fix the round-5 tombstone protected against post-mortem + * stale events from dying children (correct), but it ALSO + * rejected legitimate restore-time events for the same id + * because `markSessionClosed` (60s TTL) is set BEFORE a future + * `load` can clear the tombstone via `drainEarlyEvents` (which + * only runs AFTER `createSessionEntry`, which only runs AFTER the + * ACP `loadSession`/`unstable_resumeSession` returns). The + * restored child's MCP discovery firing during that ACP call + * window had its budget events silently dropped. + * + * Bridge factory enters the set before awaiting the ACP restore + * call and exits the set on settle (success or failure). Multi- + * waiter coalescing on the same id is naturally handled — the + * Set is idempotent on add and the cleanup is paired with the + * IIFE that does the ACP call (only one such IIFE per id at a + * time). + */ + private readonly inFlightRestoreIds = new Set(); + + /** + * PR 14b: handle child→bridge ACP `extNotification` calls. Only one + * method is recognized today — `qwen/notify/session/mcp-budget-event` + * — translating the McpClientManager's budget-event payload into a + * session-scoped SSE frame. Unknown methods, unknown event kinds, + * and missing sessionIds are dropped silently for forward-compat + * (a future child can add new notification methods without breaking + * this handler; an older daemon can ignore them cleanly). + * + * Codex review fix #1: when the sessionId IS present but the + * `byId`-resolvable entry is not yet registered (the child fired + * the event during its own `newSession` handler, before + * `connection.newSession` returned to `doSpawn`), buffer the frame + * and replay it on `drainEarlyEvents`. + */ + async extNotification( + method: string, + params: Record, + ): Promise { + if (method !== 'qwen/notify/session/mcp-budget-event') return; + const sessionId = params['sessionId']; + if (typeof sessionId !== 'string') return; + const kind = params['kind']; + const type = + kind === 'budget_warning' + ? 'mcp_budget_warning' + : kind === 'refused_batch' + ? 'mcp_child_refused_batch' + : undefined; + if (!type) return; + // Strip the routing fields (`v`, `sessionId`, `kind`) from the + // outbound `data` payload — the SSE frame already carries `v` at + // the envelope level (`EVENT_SCHEMA_VERSION`) and the session id + // is implicit from the endpoint, so duplicating them in `data` + // would be noise. `kind` is encoded as the frame `type`. + const { v: _v, sessionId: _sid, kind: _kind, ...rest } = params; + void _v; + void _sid; + void _kind; + const entry = this.resolveEntry(sessionId); + const frame: Omit = { + type, + data: rest, + ...(entry?.activePromptOriginatorClientId + ? { originatorClientId: entry.activePromptOriginatorClientId } + : {}), + }; + if (entry) { + entry.events.publish(frame); + return; + } + // No entry yet — buffer for `drainEarlyEvents`. The bridge calls + // `drainEarlyEvents` immediately after `byId.set(sessionId, entry)` + // in `createSessionEntry`; if the session never registers (spawn + // failure), the entry is GC'd by TTL after EARLY_EVENT_TTL_MS. + this.bufferEarlyEvent(sessionId, frame); + } + + /** + * PR 14b fix #1: enqueue `frame` for `sessionId`. Lazy TTL sweep + * runs first so caller doesn't pay for stale entries before + * deciding whether the session-cap is reached. New sessionIds + * past `MAX_EARLY_EVENT_SESSIONS` are dropped (defense against a + * malicious / buggy child fanning out fake sessionIds); same- + * sessionId frames past `MAX_EARLY_EVENTS_PER_SESSION` are dropped + * to bound per-session memory. + */ + private bufferEarlyEvent( + sessionId: string, + frame: Omit, + ): void { + const now = Date.now(); + // PR 14b fix (codex round 5): drop frames for ids the bridge has + // already marked closed/killed. Sweep + check before any other + // work so a malicious / buggy child can't keep appending + // post-mortem frames against an old id. Live ids that re-register + // (load/resume) clear their tombstone in `drainEarlyEvents`. + // + // Round 6 amendment: skip the tombstone check for ids currently + // being restored. Pre-amendment a `close → load same id` sequence + // within 60s lost any restore-time guardrail events because the + // tombstone outlived `bufferEarlyEvent` but `drainEarlyEvents` + // (which clears it) only runs after the ACP restore returns. + this.sweepExpiredTombstones(now); + if ( + this.tombstonedSessionIds.has(sessionId) && + !this.inFlightRestoreIds.has(sessionId) + ) { + writeStderrLine( + `qwen serve: dropping mcp guardrail extNotification ` + + `for tombstoned session ${JSON.stringify(sessionId)} ` + + `(post-close stale event)`, + ); + return; + } + this.sweepExpiredEarlyEvents(now); + let buf = this.earlyEvents.get(sessionId); + if (!buf) { + if (this.earlyEvents.size >= MAX_EARLY_EVENT_SESSIONS) { + // PR 14b fix (codex round 6): observability. Other drop + // sites in this PR all log; the silent return here was the + // outlier. Stays at stderr (visible without debug=true) + // because hitting this cap means the daemon is under + // notification pressure from 64+ concurrent sessions — + // worth surfacing. + writeStderrLine( + `qwen serve: dropping mcp guardrail extNotification — ` + + `early-event buffer at MAX_EARLY_EVENT_SESSIONS ` + + `(${MAX_EARLY_EVENT_SESSIONS}); possible session-id fanout abuse`, + ); + return; + } + buf = { frames: [], expiresAt: now + EARLY_EVENT_TTL_MS }; + this.earlyEvents.set(sessionId, buf); + } + if (buf.frames.length >= MAX_EARLY_EVENTS_PER_SESSION) { + writeStderrLine( + `qwen serve: dropping mcp guardrail extNotification ` + + `for session ${JSON.stringify(sessionId)} — per-session ` + + `cap (${MAX_EARLY_EVENTS_PER_SESSION}) reached`, + ); + return; + } + buf.frames.push(frame); + } + + private sweepExpiredEarlyEvents(now: number): void { + for (const [sid, buf] of this.earlyEvents) { + if (buf.expiresAt <= now) this.earlyEvents.delete(sid); + } + } + + private sweepExpiredTombstones(now: number): void { + for (const [sid, expiresAt] of this.tombstonedSessionIds) { + if (expiresAt <= now) this.tombstonedSessionIds.delete(sid); + } + } + + /** + * PR 14b fix (codex round 5): mark a sessionId as closed so a late + * `extNotification` from the dying child can't leak into the + * early-event buffer. Bridge factory calls this from every + * `byId.delete(sid)` site (kill path, channel.exited handler, + * closeSession). Idempotent on already-tombstoned ids — refreshes + * the TTL so a recently-killed id stays dead long enough for any + * in-flight stale frames to expire. + */ + markSessionClosed(sessionId: string): void { + const now = Date.now(); + // PR 14b fix (codex round 7): bound `tombstonedSessionIds` under + // session churn. Pre-fix `sweepExpiredTombstones` was only called + // inside `bufferEarlyEvent`; on a daemon that closes/kills many + // sessions but rarely receives extNotifications (the common + // production pattern when MCP guardrail mode is `off`), the map + // grew monotonically and the documented 60s TTL didn't bound + // memory. Sweeping at every close is O(map size) but cheap (one + // integer compare per entry); under any realistic workload the + // map stays small. + this.sweepExpiredTombstones(now); + this.tombstonedSessionIds.set(sessionId, now + EARLY_EVENT_TTL_MS); + // Purge any frames already buffered for this id — they're now + // stale by definition (their session is dead). + this.earlyEvents.delete(sessionId); + } + + /** + * PR 14b fix (codex round 6): mark a sessionId as currently being + * restored via `session/load` / `session/resume`. While in this set, + * `bufferEarlyEvent` accepts frames for the id even if it's + * tombstoned — so restore-time guardrail events from the freshly- + * restored child reach `drainEarlyEvents` instead of being rejected + * by the close-window tombstone. + * + * Bridge factory calls this BEFORE awaiting the ACP restore call. + * `clearRestoreInFlight` is paired in the matching `finally` so a + * failed restore doesn't leave a dangling allow-list entry. + * Idempotent — safe to call repeatedly during coalesced restores. + */ + markRestoreInFlight(sessionId: string): void { + this.inFlightRestoreIds.add(sessionId); + } + + /** + * PR 14b fix (codex round 6): companion to `markRestoreInFlight`. + * Bridge factory calls this when the restore IIFE settles — + * after `createSessionEntry` runs (success) or after the ACP + * restore call fails (error). After the entry is registered, + * `bufferEarlyEvent` is no longer reached for this id (notifications + * route through `entry.events.publish`), so the allow-list entry + * has no further effect — but cleared anyway to prevent the Set + * from growing forever under high restore churn. + */ + clearRestoreInFlight(sessionId: string): void { + this.inFlightRestoreIds.delete(sessionId); + } + + /** + * PR 14b fix #1: drain any frames buffered for `sessionId` onto + * `entry.events`. Bridge calls this immediately after + * `byId.set(sessionId, entry)` in `createSessionEntry`. The frames + * were captured before the entry existed (e.g. MCP discovery during + * the child's `newSession` handler), so draining them now lands + * them in the replay ring as the FIRST events of this session — + * SDK consumers reconnecting with `Last-Event-ID: 0` see them on + * their initial subscription. + * + * Public so the bridge factory can call it directly. Idempotent on + * unknown sessionIds. + */ + drainEarlyEvents(sessionId: string, entry: SessionEntry): void { + // PR 14b fix (codex round 5): a fresh registration clears any + // tombstone for this id — this is the legitimate + // "load/resume of a persisted session id" case. Any stale + // pre-tombstone frame was already rejected by `bufferEarlyEvent` + // above; clearing the tombstone now means subsequent + // notifications for this re-attached session (which is now in + // `byId`) flow through the normal `entry.events.publish` path. + this.tombstonedSessionIds.delete(sessionId); + const buf = this.earlyEvents.get(sessionId); + if (!buf) return; + for (const frame of buf.frames) entry.events.publish(frame); + this.earlyEvents.delete(sessionId); + } + async writeTextFile( params: WriteTextFileRequest, ): Promise { @@ -2219,6 +2536,11 @@ export function createHttpAcpBridge(opts: BridgeOptions): HttpAcpBridge { /* bus already closed */ } byId.delete(sid); + // PR 14b fix (codex round 5): tombstone the id so any + // late `extNotification` from the dying child can't leak + // into the early-event buffer for a future load/resume of + // the same persisted session id. + info.client.markSessionClosed(sid); if (defaultEntry === sessEntry) defaultEntry = undefined; sessEntry.events.close(); } @@ -2656,6 +2978,11 @@ export function createHttpAcpBridge(opts: BridgeOptions): HttpAcpBridge { }; ci.sessionIds.add(entry.sessionId); byId.set(entry.sessionId, entry); + // PR 14b fix #1 (codex review round 1): drain any guardrail + // events that fired during this session's `newSession` handler + // (before this entry registered) onto the freshly-created + // EventBus. Idempotent on unknown sessionIds. + ci.client.drainEarlyEvents(entry.sessionId, entry); return entry; }; @@ -2789,6 +3116,15 @@ export function createHttpAcpBridge(opts: BridgeOptions): HttpAcpBridge { pendingRestoreEvents.set(req.sessionId, restoreEvents); ci = await ensureChannel(); ci.pendingRestoreIds.add(req.sessionId); + // PR 14b fix (codex round 6): mark this id as in-flight restore + // BEFORE the ACP `loadSession`/`unstable_resumeSession` call. + // Restore-time guardrail events arriving on the bridge during + // that ACP call hit `bufferEarlyEvent` BEFORE the + // post-restore `createSessionEntry → drainEarlyEvents` clears + // the (close-window) tombstone, so without this allow-list the + // tombstone would silently drop them. Cleared in the matching + // `finally` below regardless of success / failure. + ci.client.markRestoreInFlight(req.sessionId); // Restore is a low-frequency one-shot path, so we register a // fresh `channel.exited` listener per call instead of going // through `getTransportClosedReject` (which exists to keep @@ -2922,9 +3258,27 @@ export function createHttpAcpBridge(opts: BridgeOptions): HttpAcpBridge { }; })().finally(() => { ci?.pendingRestoreIds.delete(req.sessionId); + // PR 14b fix (codex round 6): pair with `markRestoreInFlight`. + // Once the IIFE settles, either `createSessionEntry` ran + // (`drainEarlyEvents` already cleared the tombstone) or the + // restore failed (handled below). + ci?.client.clearRestoreInFlight(req.sessionId); pendingRestoreEvents.delete(req.sessionId); if (!registeredEntry) { restoreEvents.close(); + // PR 14b fix (codex round 7): on restore failure, purge any + // guardrail events that the child buffered during this + // restore window AND re-tombstone the id. Pre-fix the + // round-6 allow-list (`markRestoreInFlight`) let + // `bufferEarlyEvent` accept frames during the ACP call; + // failure here only cleared the allow-list entry, leaving + // queued frames in `earlyEvents`. A subsequent successful + // `session/load`/`session/resume` for the same id within + // 60s would then `drainEarlyEvents` those stale frames into + // the new session — exactly the leak round 5's tombstone + // was meant to prevent. `markSessionClosed` already does + // both: refresh tombstone + delete `earlyEvents[id]`. + ci?.client.markSessionClosed(req.sessionId); } }); @@ -3385,6 +3739,11 @@ export function createHttpAcpBridge(opts: BridgeOptions): HttpAcpBridge { resolvePending(id, { outcome: { outcome: 'cancelled' } }); } byId.delete(sessionId); + // PR 14b fix (codex round 5): tombstone the closed sessionId + // so any late `extNotification` from the (now-defunct) child + // can't seed the early-event buffer and leak into a future + // load/resume of the same persisted id. + ci?.client.markSessionClosed(sessionId); try { entry.events.publish({ type: 'session_closed', @@ -4143,6 +4502,12 @@ export function createHttpAcpBridge(opts: BridgeOptions): HttpAcpBridge { if (ci && ci.channel === entry.channel) { ci.sessionIds.delete(sessionId); } + // PR 14b fix (codex round 5): tombstone the killed sessionId + // so any in-flight `extNotification` from the (about-to-be- + // killed) child can't seed the early-event buffer for a + // subsequent load/resume of the same persisted id. See the + // matching guard in BridgeClient.bufferEarlyEvent. + ci?.client.markSessionClosed(sessionId); // Resolve any still-pending permission as cancelled (matches the // shutdown path) so callers awaiting requestPermission unwind. for (const id of Array.from(entry.pendingPermissionIds)) { diff --git a/packages/cli/src/serve/server.test.ts b/packages/cli/src/serve/server.test.ts index 080ede586ae..f0469504053 100644 --- a/packages/cli/src/serve/server.test.ts +++ b/packages/cli/src/serve/server.test.ts @@ -113,6 +113,10 @@ const EXPECTED_STAGE1_FEATURES = [ // `budgets[]` on `/workspace/mcp`, `disabledReason: 'budget'` on // refused per-server cells). 'mcp_guardrails', + // Issue #4175 PR 14b. Always-on. Daemon emits typed push events for + // MCP budget state crossings (`mcp_budget_warning` with hysteresis, + // `mcp_child_refused_batch` coalesced per pass). + 'mcp_guardrail_events', // Issue #4175 PR 19. Always-on. Daemon exposes the read-only file // surface: `GET /file`, `GET /list`, `GET /glob`, `GET /stat`. 'workspace_file_read', @@ -832,6 +836,18 @@ describe('createServeApp', () => { }); }); + it('registers mcp_guardrail_events as a baseline tag (#4175 PR 14b)', () => { + // PR 14b's push events are unconditional once advertised — there's + // no operator toggle. So no `modes`, no entry in + // `CONDITIONAL_SERVE_FEATURES`. SDK consumers feature-detect via + // `caps.features.includes('mcp_guardrail_events')` before + // narrowing `mcp_budget_warning` / `mcp_child_refused_batch` + // frames through `KnownDaemonEvent`. + expect(SERVE_CAPABILITY_REGISTRY['mcp_guardrail_events']).toEqual({ + since: 'v1', + }); + }); + it('returns protocol version metadata with a fresh supported array', () => { const versions = getServeProtocolVersions(); expect(versions).toEqual({ current: 'v1', supported: ['v1'] }); diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 172b3983d19..99ea9eebff5 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -91,6 +91,34 @@ vi.mock('../tools/tool-registry', () => { ToolRegistryMock.prototype.getAllToolNames = vi.fn(() => []); ToolRegistryMock.prototype.getTool = vi.fn(); ToolRegistryMock.prototype.getFunctionDeclarations = vi.fn(() => []); + // PR 14b fix (codex round 4): per-instance manager stub so the + // `setMcpBudgetEventCallback → createToolRegistry → manager.setOnBudgetEvent` + // integration test can observe each instance's callback wiring. + // The mock constructor stamps a fresh `__mcpManagerMock` onto each + // ToolRegistry instance so tests can inspect it via + // `(registry as unknown as { __mcpManagerMock }).__mcpManagerMock` + // (escape hatch — production code reads it via `getMcpClientManager`). + ToolRegistryMock.mockImplementation(function (this: { + __mcpManagerMock: { + setOnBudgetEvent: Mock; + discoverAllMcpToolsIncremental: Mock; + }; + }) { + this.__mcpManagerMock = { + setOnBudgetEvent: vi.fn(), + // Stubbed so `Config.startMcpDiscoveryInBackground` (kicked off + // at the tail of `initialize`) doesn't crash on missing method. + // Test cares only about the `setOnBudgetEvent` wiring; discovery + // itself is a no-op here. + discoverAllMcpToolsIncremental: vi.fn().mockResolvedValue(undefined), + }; + return this; + }); + ToolRegistryMock.prototype.getMcpClientManager = function (this: { + __mcpManagerMock: { setOnBudgetEvent: Mock }; + }) { + return this.__mcpManagerMock; + }; return { ToolRegistry: ToolRegistryMock }; }); @@ -2079,6 +2107,116 @@ describe('Server Config (config.ts)', () => { ); }); }); + + // PR 14b fix (codex round 4 — wenshao gpt-5.5 review): the + // `Config.setMcpBudgetEventCallback → pendingMcpBudgetCallback → + // createToolRegistry → registry.getMcpClientManager().setOnBudgetEvent` + // boundary previously had NO test. The acpAgent test stubs the + // setter (proves QwenAgent calls it pre-`initialize`); the manager + // tests bypass Config by passing `onBudgetEvent` directly to + // `McpClientManager`. Neither covers the actual stash + apply path + // inside Config — and that path is the safety net that prevents + // startup-window MCP guardrail events from being dropped under + // legacy blocking discovery + closes the progressive-mode race + // window. These two tests exercise both call orderings (pre-init + // and late-call). + describe('setMcpBudgetEventCallback handoff to McpClientManager', () => { + it('applies pending callback when registry is created during initialize()', async () => { + const config = new Config(baseParams); + const cb = vi.fn(); + // Setter called BEFORE initialize — value stashed on + // `pendingMcpBudgetCallback` and applied inside + // `createToolRegistry` after the manager is constructed but + // BEFORE `discoverAllTools` / background discovery fires. + config.setMcpBudgetEventCallback(cb); + await config.initialize(); + + const registry = config.getToolRegistry() as unknown as { + __mcpManagerMock: { setOnBudgetEvent: Mock }; + }; + expect(registry.__mcpManagerMock.setOnBudgetEvent).toHaveBeenCalledWith( + cb, + ); + // Exactly once — the apply path fires only once per + // `createToolRegistry` invocation. + expect( + registry.__mcpManagerMock.setOnBudgetEvent.mock.calls, + ).toHaveLength(1); + }); + + it('applies callback directly to existing manager when called after initialize()', async () => { + const config = new Config(baseParams); + // Initialize WITHOUT a pending callback first — the + // createToolRegistry apply branch is a no-op. + await config.initialize(); + const registry = config.getToolRegistry() as unknown as { + __mcpManagerMock: { setOnBudgetEvent: Mock }; + }; + // Sanity: no apply happened during init since callback was + // never registered. + expect(registry.__mcpManagerMock.setOnBudgetEvent).not.toHaveBeenCalled(); + + // Late-call path: setter dispatches DIRECTLY to the existing + // manager via the `if (this.toolRegistry)` branch in + // `setMcpBudgetEventCallback`. This is the path tests/adapters + // use when they discover the manager only after Config is up. + const cb = vi.fn(); + config.setMcpBudgetEventCallback(cb); + expect(registry.__mcpManagerMock.setOnBudgetEvent).toHaveBeenCalledWith( + cb, + ); + + // Calling with `undefined` clears the registration on the + // manager (parity with the constructor-time `off`-mode strip + // in McpClientManager). + config.setMcpBudgetEventCallback(undefined); + expect( + registry.__mcpManagerMock.setOnBudgetEvent, + ).toHaveBeenLastCalledWith(undefined); + }); + + it('does NOT stash the callback when called after initialize() (codex round 7 fix — subagent isolation)', async () => { + // Codex round 7 finding: pre-fix, the late-call path assigned + // to `pendingMcpBudgetCallback` BEFORE applying directly to + // the existing manager. A subsequent `createToolRegistry` + // (e.g. subagent override via `createApprovalModeOverride` / + // `buildSubagentContextOverride`) would inherit the stash and + // wire the parent session's ACP push callback into the + // subagent's fresh manager, routing subagent telemetry + // through the wrong session. + // + // Fix: late-call path applies directly + sets + // `pendingMcpBudgetCallback = undefined`. Pre-init path still + // stashes (the only way to reach a manager that doesn't + // exist yet — round 1 fix #2 contract). + const config = new Config(baseParams); + await config.initialize(); + const registry = config.getToolRegistry() as unknown as { + __mcpManagerMock: { setOnBudgetEvent: Mock }; + }; + + // Late-call: apply. + const cb = vi.fn(); + config.setMcpBudgetEventCallback(cb); + expect(registry.__mcpManagerMock.setOnBudgetEvent).toHaveBeenCalledWith( + cb, + ); + + // Now rebuild a registry as if for a subagent override. With + // the round-7 fix, the new manager should NOT receive the + // parent session's callback — pre-fix this would re-apply + // `cb` to the new manager. + const subagentRegistry = (await config.createToolRegistry(undefined, { + skipDiscovery: true, + forSubAgent: true, + })) as unknown as { + __mcpManagerMock: { setOnBudgetEvent: Mock }; + }; + expect( + subagentRegistry.__mcpManagerMock.setOnBudgetEvent, + ).not.toHaveBeenCalled(); + }); + }); }); describe('setApprovalMode with folder trust', () => { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 4d074c83220..89c6859ee79 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -59,6 +59,7 @@ import { setGeminiMdFilename } from '../memory/const.js'; import { canUseRipgrep } from '../utils/ripgrepUtils.js'; import { recordStartupEvent } from '../utils/startupEventSink.js'; import { ToolRegistry, type ToolFactory } from '../tools/tool-registry.js'; +import type { McpBudgetEvent } from '../tools/mcp-client-manager.js'; import { ToolNames } from '../tools/tool-names.js'; import type { LspClient, LspStatusSnapshot } from '../lsp/types.js'; @@ -712,6 +713,17 @@ export class Config { private sessionData?: ResumedSessionData; private debugLogger: DebugLogger; private toolRegistry!: ToolRegistry; + /** + * PR 14b fix #2 (codex review round 1): callback stashed BEFORE + * `initialize()` runs and applied as soon as `toolRegistry` is up, + * so the manager's `setOnBudgetEvent` is wired before + * `startMcpDiscoveryInBackground` (or legacy blocking discovery) + * fires the first pass. Pre-fix the acpAgent registered after + * `initialize()` returned, missing the first pass entirely under + * `QWEN_CODE_LEGACY_MCP_BLOCKING=1` and racing against background + * discovery completion under the default mode. + */ + private pendingMcpBudgetCallback?: (event: McpBudgetEvent) => void; private promptRegistry!: PromptRegistry; private subagentManager!: SubagentManager; private readonly backgroundTaskRegistry = new BackgroundTaskRegistry(); @@ -3675,6 +3687,34 @@ export class Config { return new MonitorTool(this); }); + // PR 14b fix #2 (codex review round 1): apply any pending MCP + // budget-event callback BEFORE `discoverAllTools` (legacy blocking + // mode runs MCP discovery synchronously in there) and BEFORE the + // post-`createToolRegistry` `startMcpDiscoveryInBackground` (default + // mode). Either way the manager has its callback wired at the + // moment the first discovery pass fires, so end-of-pass events + // for that pass are routed through the SDK push channel. + if (this.pendingMcpBudgetCallback) { + const mgr = registry.getMcpClientManager(); + if (mgr && typeof mgr.setOnBudgetEvent === 'function') { + mgr.setOnBudgetEvent(this.pendingMcpBudgetCallback); + } + // PR 14b fix (codex round 6): clear after consumption so a + // subsequent `createToolRegistry` call (e.g. subagent override + // via `createApprovalModeOverride` / + // `buildSubagentContextOverride`) doesn't re-apply the parent + // session's callback to a fresh manager. Subagent contexts run + // their own MCP clients but should NOT push budget events + // through the parent's ACP session — that would route subagent + // telemetry to the wrong subscriber. + // + // Late-call setter (`setMcpBudgetEventCallback` after + // `initialize()`) is unaffected: it dispatches directly to the + // existing manager via the `if (this.toolRegistry)` branch, + // not through `pendingMcpBudgetCallback`. + this.pendingMcpBudgetCallback = undefined; + } + if (!options?.skipDiscovery) { await registry.discoverAllTools(); } @@ -3683,4 +3723,45 @@ export class Config { ); return registry; } + + /** + * PR 14b fix #2 (codex review round 1): register the MCP guardrail + * push-event callback. Acceptable to call at any point in the + * Config lifecycle — before, during, or after `initialize()`. + * + * Two paths: + * - **Pre-init** (no `toolRegistry` yet): stash on + * `pendingMcpBudgetCallback`. `createToolRegistry` will apply it + * to the freshly-constructed manager and clear the stash (round + * 6 fix). The stash is the ONLY way to reach a manager that + * doesn't exist yet. + * - **Late** (`toolRegistry` already exists): dispatch directly to + * the existing manager. **DO NOT** also stash — that's the + * round-7 fix. Pre-fix, both paths assigned to + * `pendingMcpBudgetCallback` regardless, so a subsequent + * `createToolRegistry` (subagent override via + * `createApprovalModeOverride` / + * `buildSubagentContextOverride`) would re-apply the parent + * session's callback to the subagent's fresh manager — routing + * subagent telemetry through the wrong ACP session. + * + * `cb: undefined` clears the registration. `off`-mode managers + * silently drop the callback (their state machine never runs). + */ + setMcpBudgetEventCallback( + cb: ((event: McpBudgetEvent) => void) | undefined, + ): void { + if (this.toolRegistry) { + // Late-call path: apply directly. Do NOT stash — see comment + // above for the subagent isolation rationale. + const mgr = this.toolRegistry.getMcpClientManager?.(); + if (mgr && typeof mgr.setOnBudgetEvent === 'function') { + mgr.setOnBudgetEvent(cb); + } + this.pendingMcpBudgetCallback = undefined; + return; + } + // Pre-init path: stash for `createToolRegistry` to consume. + this.pendingMcpBudgetCallback = cb; + } } diff --git a/packages/core/src/tools/mcp-client-manager.test.ts b/packages/core/src/tools/mcp-client-manager.test.ts index 9c728f91f57..ea4a8f3e2ca 100644 --- a/packages/core/src/tools/mcp-client-manager.test.ts +++ b/packages/core/src/tools/mcp-client-manager.test.ts @@ -2028,3 +2028,558 @@ describe('McpClientManager — PR 14 guardrails', () => { expect(manager.getMcpClientAccounting().reservedSlots).toEqual(['a']); }); }); + +// Issue #4175 PR 14b: push events + hysteresis state machine. Kept in +// its own describe so a future revert of PR 14b drops a single +// contiguous block. Mirrors PR 14's testing style (mock `McpClient`, +// fluent `configWithServers` helper). Imports are dynamic to keep +// the spy on `McpClient` cleanly bound per test (vi.mocked module +// already mocked at file top). +describe('McpClientManager — PR 14b push events + hysteresis', () => { + afterEach(() => { + vi.restoreAllMocks(); + delete process.env['QWEN_SERVE_MCP_CLIENT_BUDGET']; + delete process.env['QWEN_SERVE_MCP_BUDGET_MODE']; + }); + + function makeConnectedMcpClientMock() { + const state = { status: undefined as unknown }; + return { + connect: vi.fn().mockImplementation(async () => { + const { MCPServerStatus } = await import('./mcp-client.js'); + state.status = MCPServerStatus.CONNECTED; + }), + discover: vi.fn().mockResolvedValue(undefined), + disconnect: vi.fn().mockResolvedValue(undefined), + getStatus: vi.fn(() => state.status), + readResource: vi.fn().mockResolvedValue({ contents: [] }), + }; + } + + function configWithServers( + servers: Record, + overrides: Partial = {}, + ): Config { + return { + isTrustedFolder: () => true, + getMcpServers: () => servers, + getMcpServerCommand: () => undefined, + getPromptRegistry: () => ({}) as PromptRegistry, + getWorkspaceContext: () => ({}) as WorkspaceContext, + getDebugMode: () => false, + isMcpServerDisabled: () => false, + ...overrides, + } as unknown as Config; + } + + it('exports MCP_BUDGET_REARM_FRACTION = 0.375', async () => { + const { MCP_BUDGET_REARM_FRACTION } = await import( + './mcp-client-manager.js' + ); + expect(MCP_BUDGET_REARM_FRACTION).toBe(0.375); + }); + + it('budget_warning fires once on first 75% upward crossing', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + // 4-server config, budget 4, ratio after pass = 4/4 = 1.0 ≥ 0.75 + // → exactly one warning fires. + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + c: { command: 'node' }, + d: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 4, + budgetMode: 'warn', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + const warnings = events.filter( + (e) => (e as { kind: string }).kind === 'budget_warning', + ); + expect(warnings).toHaveLength(1); + // PR 14b fix #4 (codex review round 1): hysteresis fires inline on + // the upward crossing, so the payload reflects the moment ratio + // first hits 0.75 — `reservedCount: 3` (3 of 4 reserved). Pre-fix + // the test saw the post-stabilization `reservedCount: 4` because + // the standalone end-of-pass `evaluateBudgetState` ran after every + // reservation completed. + expect(warnings[0]).toMatchObject({ + kind: 'budget_warning', + reservedCount: 3, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }); + }); + + it('budget_warning does NOT fire when ratio stays below 75%', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + // 2 of 4 → 0.5 < 0.75 → no fire. + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 4, + budgetMode: 'warn', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toEqual([]); + }); + + it('budget_warning hysteresis re-arms only after dropping below 37.5%', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + // Budget 4. Pass 1: 4/4 = 1.0 fires. Pass 2 after disconnecting + // 2 (-> 2/4=0.5, above 37.5%) does NOT re-arm. Pass 3 after + // disconnecting one more (-> 1/4=0.25 below 37.5%) re-arms. + // Re-arming alone doesn't fire — the next upward crossing fires. + let servers: Record = { + a: { command: 'node' }, + b: { command: 'node' }, + c: { command: 'node' }, + d: { command: 'node' }, + }; + const cfgGetter = () => servers; + const config = configWithServers({}, { + getMcpServers: cfgGetter, + } as Partial); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 4, + budgetMode: 'warn', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(1); + + // Drop to 50% via disconnect: 2/4 = 0.5 — above 37.5%, NO re-arm + // (warning stays disabled). + await manager.disconnectServer('c'); + await manager.disconnectServer('d'); + // Force a state evaluation: a successful per-server reconnect path + // is the cleanest in-band trigger; emulate one by re-discovering + // 'a'. (`evaluateBudgetState` is private — we exercise it via the + // public path instead.) + await manager.discoverMcpToolsForServer('a', config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(1); // still 1 — not re-fired + + // Drop to 25% via disconnect — below 37.5% — should re-arm but + // not fire yet (re-arming alone doesn't trigger an event). + await manager.disconnectServer('b'); + await manager.discoverMcpToolsForServer('a', config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(1); + + // Now refill back to 4/4 — re-armed state plus upward crossing + // fires the second warning. + servers = { + a: { command: 'node' }, + b: { command: 'node' }, + c: { command: 'node' }, + d: { command: 'node' }, + }; + await manager.discoverAllMcpToolsIncremental(config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(2); + }); + + it('off mode never fires budget_warning', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { budgetMode: 'off', onBudgetEvent: (e) => events.push(e) }, + ); + await manager.discoverAllMcpTools(config); + expect(events).toEqual([]); + }); + + it('refused_batch coalesces multi-refusal into one event per pass', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + // budget 1, 3 servers → a connects, b+c refused. + const config = configWithServers({ + a: { command: 'node' }, + b: { httpUrl: 'http://b' }, + c: { url: 'http://c' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 1, + budgetMode: 'enforce', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + const batches = events.filter( + (e) => (e as { kind: string }).kind === 'refused_batch', + ); + expect(batches).toHaveLength(1); + expect(batches[0]).toMatchObject({ + kind: 'refused_batch', + budget: 1, + mode: 'enforce', + refusedServers: [ + { name: 'b', transport: 'http', reason: 'budget_exhausted' }, + { name: 'c', transport: 'sse', reason: 'budget_exhausted' }, + ], + }); + }); + + it('refused_batch does NOT fire when no servers are refused', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 5, + budgetMode: 'enforce', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'refused_batch'), + ).toEqual([]); + }); + + it('readResource refusal emits a length-1 refused_batch then throws', async () => { + const { BudgetExhaustedError } = await import('./mcp-client-manager.js'); + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 1, + budgetMode: 'enforce', + onBudgetEvent: (e) => events.push(e), + }, + ); + // First pass fills the budget with `a`. `b` is refused — that's + // the bulk refusal (length-1 batch). + await manager.discoverAllMcpTools(config); + // Clear bulk events so the assertion below tracks only the + // readResource path. + events.length = 0; + // Now lazy-spawn against b — slot full, throws + emits a + // length-1 batch. + await expect(manager.readResource('b', 'mcp://b/resource')).rejects.toThrow( + BudgetExhaustedError, + ); + const batches = events.filter( + (e) => (e as { kind: string }).kind === 'refused_batch', + ); + expect(batches).toHaveLength(1); + expect(batches[0]).toMatchObject({ + kind: 'refused_batch', + mode: 'enforce', + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + ], + }); + }); + + it('off-mode constructor strips onBudgetEvent (defense in depth)', async () => { + // Off-mode never runs the state machine; the constructor stashes + // `undefined` for `onBudgetEvent` so even a stray internal call + // can't fire. Verified externally by observing that no events + // arrive. + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { budgetMode: 'off', onBudgetEvent: (e) => events.push(e) }, + ); + await manager.discoverAllMcpTools(config); + // Force discovery refusal would be impossible in off mode (no + // budget). Disconnect-then-rediscover also no-ops the state + // machine. End-to-end no events. + await manager.disconnectServer('a'); + await manager.discoverMcpToolsForServer('a', config); + expect(events).toEqual([]); + }); + + it('refused_batch transports preserve the per-server family at refusal time', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + // Mixed transports refused; budget 1 admits the first only. + const config = configWithServers({ + a: { command: 'node' }, // stdio (admitted) + b: { httpUrl: 'http://b' }, // http (refused) + c: { url: 'http://c' }, // sse (refused) + d: { tcp: 'ws://d' }, // websocket (refused) + e: { type: 'sdk', command: 'sdk' }, // sdk (refused) + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 1, + budgetMode: 'enforce', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + const batches = events.filter( + (e) => (e as { kind: string }).kind === 'refused_batch', + ) as Array<{ refusedServers: Array<{ name: string; transport: string }> }>; + expect(batches).toHaveLength(1); + expect( + batches[0].refusedServers.map((r) => `${r.name}:${r.transport}`), + ).toEqual(['b:http', 'c:sse', 'd:websocket', 'e:sdk']); + }); + + it('warn mode never emits refused_batch (only enforce refuses)', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + c: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 1, + budgetMode: 'warn', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + // warn mode: no refusals, but the warning may fire (3/1 ratio crosses 0.75). + expect( + events.filter((e) => (e as { kind: string }).kind === 'refused_batch'), + ).toEqual([]); + }); + + it('stop() re-arms the warning state machine for the next session', async () => { + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + c: { command: 'node' }, + d: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 4, + budgetMode: 'warn', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + // First crossing fired one warning. + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(1); + // stop() resets state. Next discovery pass that crosses 75% + // fires anew. discoverAllMcpTools internally calls stop() at + // the top, so calling it again is sufficient. + await manager.discoverAllMcpTools(config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(2); + }); + + it('discoverAllMcpToolsIncremental coalesces multi-server refusals into ONE batch (codex review fix #3)', async () => { + // Codex review round 1, finding #3: pre-fix, when + // `discoverAllMcpToolsIncremental` walked N new servers and the + // budget was full, each per-server refusal called + // `emitRefusedBatchIfAny` inline → N length-1 batch events + // instead of 1 length-N batch. This test pins the documented + // "one batch per pass" contract via the `bulkPassDepth` guard. + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + // Budget 1, 4 servers — 1 admitted, 3 refused. Pre-fix this + // produced 3 length-1 batches via `discoverMcpToolsForServer` → + // `discoverMcpToolsForServerInternal`. Post-fix: 1 length-3 batch. + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + c: { command: 'node' }, + d: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 1, + budgetMode: 'enforce', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpToolsIncremental(config); + const batches = events.filter( + (e) => (e as { kind: string }).kind === 'refused_batch', + ) as Array<{ refusedServers: Array<{ name: string }> }>; + // Strict invariant: ONE batch event, not N. + expect(batches).toHaveLength(1); + expect(batches[0].refusedServers.map((r) => r.name)).toEqual([ + 'b', + 'c', + 'd', + ]); + }); + + it('disconnectServer drives the hysteresis re-arm path (codex review fix #4)', async () => { + // Codex review round 1, finding #4: pre-fix `disconnectServer` / + // `removeServer` deleted from `reservedSlots` without invoking + // `evaluateBudgetState`, so `warnArmed` stayed `false` after a + // 75% fire even though the ratio dropped below 37.5%. This test + // exercises the operator-driven release path: 4/4 → fire #1 → + // disconnect 3 servers (1/4, below re-arm) → reconnect 3 → 4/4 + // → fire #2. Pre-fix: only one fire. Post-fix: two fires. + vi.mocked(McpClient).mockImplementation( + () => makeConnectedMcpClientMock() as unknown as McpClient, + ); + const events: unknown[] = []; + const config = configWithServers({ + a: { command: 'node' }, + b: { command: 'node' }, + c: { command: 'node' }, + d: { command: 'node' }, + }); + const manager = new McpClientManager( + config, + {} as ToolRegistry, + undefined, + undefined, + undefined, + { + clientBudget: 4, + budgetMode: 'warn', + onBudgetEvent: (e) => events.push(e), + }, + ); + await manager.discoverAllMcpTools(config); + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(1); + // Drop to 1/4 via operator disconnects — each release crosses + // through 0.75 → 0.5 → 0.25, the last one crossing 37.5% inline + // re-arms `warnArmed` via `releaseSlotName`'s evaluate. + await manager.disconnectServer('b'); + await manager.disconnectServer('c'); + await manager.disconnectServer('d'); + // Reconnect via direct discoverMcpToolsForServer (bypasses + // discoverAllMcpTools' bulk-pass reset, exercises the re-armed + // state through inline `tryReserveSlot` evaluate calls). + await manager.discoverMcpToolsForServer('b', config); + await manager.discoverMcpToolsForServer('c', config); + // 3/4 = 0.75 — fire #2. + expect( + events.filter((e) => (e as { kind: string }).kind === 'budget_warning'), + ).toHaveLength(2); + }); +}); diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts index 527d6a2d0ec..936eeeed97c 100644 --- a/packages/core/src/tools/mcp-client-manager.ts +++ b/packages/core/src/tools/mcp-client-manager.ts @@ -46,13 +46,12 @@ const DEFAULT_HEALTH_CONFIG: MCPHealthMonitorConfig = { }; /** - * Single-threshold warning fraction for the snapshot-based budget cell - * (PR 14 v1). When `liveCount >= MCP_BUDGET_WARN_FRACTION * budget` the - * `budgets[0].status` flips to `'warning'`. Exported and consumed by - * (a) `acpAgent.buildBudgetCells` (snapshot status) and (b) - * `commands/serve.ts` (stderr boot breadcrumb) — pre-extract these - * shared so PR 14b can swap to a dual-threshold hysteresis pair - * (`armed` boolean per opencode `cli/heap.ts`) by editing one file. + * Upper threshold of the dual-threshold hysteresis used by both the + * snapshot-based budget cell (PR 14 v1) and the push-event state + * machine (PR 14b). When `reservedSlots.size / clientBudget` crosses + * this fraction upward, a `budget_warning` event fires and the + * armed-state flips to "fired"; the next fire requires the ratio to + * drop below `MCP_BUDGET_REARM_FRACTION` first. * * Picked 0.75 to mirror PR 10's `slow_client_warning` * (`eventBus.ts:WARN_THRESHOLD_RATIO`) — same rationale: "warning" @@ -60,6 +59,15 @@ const DEFAULT_HEALTH_CONFIG: MCPHealthMonitorConfig = { */ export const MCP_BUDGET_WARN_FRACTION = 0.75 as const; +/** + * Lower threshold for the hysteresis state machine (PR 14b). After a + * warning fires, the ratio must drop below this fraction before the + * state machine re-arms — so a server that flaps just above 0.75 + * doesn't produce a flood of identical warnings. Mirrors PR 10's + * `eventBus.ts:WARN_RESET_RATIO` (0.375 = half of the warn fraction). + */ +export const MCP_BUDGET_REARM_FRACTION = 0.375 as const; + /** * Budget enforcement mode for MCP client guardrails (issue #4175 PR 14). * @@ -89,8 +97,62 @@ export interface McpBudgetConfig { clientBudget?: number; /** Behavior at and above the cap. `off` when `clientBudget` is undefined. */ budgetMode: McpBudgetMode; + /** + * PR 14b: optional callback invoked by the manager when a budget + * threshold is crossed (`'budget_warning'`) or one or more servers + * are refused during a discovery pass (`'refused_batch'`). The + * manager stays decoupled from ACP wire types — the callback is + * provided by `acpAgent.newSessionConfig` and translates each event + * into a `connection.extNotification(...)` call carrying the + * sessionId. Absent in `off` mode (state machine is dormant). + */ + onBudgetEvent?: (event: McpBudgetEvent) => void; } +/** + * One refused-server entry in a `'refused_batch'` event payload (PR 14b). + * `transport` is the family resolved at refusal time via `mcpTransportOf`; + * `reason` is `'budget_exhausted'` until additional refusal causes are + * defined (Wave 5+). + */ +export interface McpRefusedServer { + name: string; + transport: McpTransportKind; + reason: 'budget_exhausted'; +} + +/** + * Discriminated union of guardrail events emitted to `onBudgetEvent`. + * + * - `budget_warning` fires on the upward crossing of + * `reservedSlots.size / clientBudget >= MCP_BUDGET_WARN_FRACTION`, + * then re-arms only after the ratio drops below + * `MCP_BUDGET_REARM_FRACTION`. Carries both `liveCount` (CONNECTED + * clients) and `reservedCount` (configured-set, including in-flight + * reservations) so SDK consumers can render either lens. + * - `refused_batch` fires once per `discoverAllMcpTools*` pass when + * `lastRefusedServerNames.length > 0`, OR as a length-1 batch on the + * `readResource` lazy-spawn refusal path. `mode` is the literal + * `'enforce'` because `warn` mode never refuses. + */ +export type McpBudgetEvent = + | { + kind: 'budget_warning'; + liveCount: number; + reservedCount: number; + budget: number; + thresholdRatio: typeof MCP_BUDGET_WARN_FRACTION; + mode: 'warn' | 'enforce'; + } + | { + kind: 'refused_batch'; + refusedServers: McpRefusedServer[]; + budget: number; + liveCount: number; + reservedCount: number; + mode: 'enforce'; + }; + /** Transport family per `MCPServerConfig`. `unknown` covers misconfigured entries. */ export type McpTransportKind = | 'stdio' @@ -321,6 +383,80 @@ export class McpClientManager { * refusals to operators. */ private lastRefusedServerNames: string[] = []; + /** + * PR 14b: transport family (`stdio`/`http`/...) resolved for each + * entry in `lastRefusedServerNames`, captured at refusal time. The + * `'refused_batch'` event payload includes the per-server transport + * so dashboards can break down "which kind of servers got refused" + * without re-walking config. + * + * Lifetime mirrors `lastRefusedServerNames`: reset at the start of + * each `discoverAllMcpTools*` pass + on `stop()` + on + * `dropRefusalEntry` (operator removed/disconnected the server). + * NOT cleared on `emitRefusedBatchIfAny` — the snapshot-visible + * refusal state survives between passes per the PR 14 contract, + * so a snapshot taken between passes still reports the last + * refusal set with correct transport metadata. The push-event + * idempotency invariant is held by the separate + * `pendingRefusalNames` queue, not by clearing this map. + */ + private lastRefusedTransports = new Map(); + /** + * PR 14b: queue of refusal names NOT YET emitted as a push event. + * `lastRefusedServerNames` is the snapshot-visible state and MUST + * survive between passes (PR 14 contract). The push-event path + * needs separate accounting so a length-1 batch fired by a single- + * server / readResource refusal doesn't get re-emitted by the + * bulk-pass end-of-pass call. `refuseAndLog` adds to both; + * `emitRefusedBatchIfAny` drains and clears this set without + * touching `lastRefusedServerNames`. Empty whenever there are no + * unsent refusals, regardless of pass. + */ + private pendingRefusalNames = new Set(); + /** + * PR 14b: hysteresis state for `'budget_warning'` events. `true` + * means "next 75% upward crossing fires"; `false` means "warning + * already fired, waiting for ratio to drop below 37.5% to re-arm". + * Stays `true` permanently in `off` mode (the state machine + * short-circuits before touching it). Initial value `true` so the + * first crossing during a session always fires. + */ + private warnArmed = true; + /** + * PR 14b fix #3 (codex review round 1): re-entrant counter that + * tracks whether a bulk discovery pass is currently in flight. + * Incremented on entry to `discoverAllMcpTools` / + * `discoverAllMcpToolsIncremental`; decremented in the matching + * `finally`. While > 0, `emitRefusedBatchIfAny` short-circuits so + * per-server refusals queue up; the bulk pass's own end-of-pass + * call (which runs AFTER `bulkPassDepth--`) drains the queue once + * as a coalesced batch — preserving the documented "one batch per + * pass" contract regardless of which inner code path enqueued the + * refusals (`discoverMcpToolsForServerInternal` from incremental, + * inline `refuseAndLog` from legacy bulk). + * + * Counter rather than boolean to defend against re-entry (a future + * code path that nests bulk passes — e.g. a discovery hook that + * itself triggers reload — wouldn't accidentally clear the flag + * mid-outer-pass). + */ + private bulkPassDepth = 0; + /** + * PR 14b: optional callback set at construction time OR via + * `setOnBudgetEvent` after construction. When non-`null` and + * `budgetMode !== 'off'`, the manager fires it on every threshold + * crossing or non-empty refusal batch. Decouples core from ACP + * wire types; `acpAgent.newSessionConfig` provides the adapter + * that translates events into `connection.extNotification`. + * + * The setter exists because the production construction path + * (`ToolRegistry` constructor → `loadCliConfig`) doesn't expose a + * hook to thread the callback through. acpAgent registers the + * callback after `loadCliConfig` returns but BEFORE + * `config.initialize()` fires the first discovery — so no events + * are missed. + */ + private onBudgetEvent?: (event: McpBudgetEvent) => void; constructor( config: Config, @@ -372,6 +508,12 @@ export class McpClientManager { } this.clientBudget = resolved.clientBudget; this.budgetMode = resolvedMode; + // PR 14b: capture the optional event callback only when enforcement + // is actually live. In `off` mode the state machine never runs, so + // a stray callback would never fire — stash `undefined` to make + // that invariant visible at the field level. + this.onBudgetEvent = + resolvedMode === 'off' ? undefined : resolved.onBudgetEvent; } /** @@ -399,9 +541,32 @@ export class McpClientManager { } // `warn` mode (and `enforce` under cap) — track in the configured set. this.reservedSlots.add(serverName); + // PR 14b fix #4 (codex review round 1): drive the hysteresis state + // machine on every upward slot mutation so a 75% crossing during + // bulk discovery fires inline, not at end-of-pass. Pre-fix the + // bulk path's terminal evaluate saw the post-stabilization ratio + // and missed transient crossings. + this.evaluateBudgetState(); return 'reserved'; } + /** + * PR 14b fix #4 (codex review round 1): single release path for + * `reservedSlots`. Delete + re-evaluate hysteresis on every + * downward mutation so re-arming through the 37.5% boundary + * happens whether the release came from operator + * `disconnectServer`, config-driven `removeServer`, discovery + * timeout cleanup, or a connect-failure catch block. + * + * Returns `true` when the name was actually held (parity with + * `Set.delete`'s return); idempotent on already-released names. + */ + private releaseSlotName(name: string): boolean { + const had = this.reservedSlots.delete(name); + if (had) this.evaluateBudgetState(); + return had; + } + /** * Snapshot the manager's MCP accounting for the daemon's read-only * `GET /workspace/mcp` route. Cheap to call — iterates `this.clients` @@ -449,6 +614,25 @@ export class McpClientManager { return this.clientBudget; } + /** + * PR 14b: register (or replace) the budget-event callback. Production + * code path: acpAgent constructs Config (which constructs the + * manager via env-var defaults) then calls this BEFORE + * `config.initialize()` so the callback is wired before the first + * discovery pass fires. + * + * No-op in `off` mode — the state machine never runs, so a callback + * here would never fire. Tests can pass a callback at construction + * via `budgetConfig.onBudgetEvent` instead, which avoids this + * setter path. + */ + setOnBudgetEvent( + callback: ((event: McpBudgetEvent) => void) | undefined, + ): void { + if (this.budgetMode === 'off') return; + this.onBudgetEvent = callback; + } + /** * Whether a discovery / reconnect for `serverName` is currently in * flight (started but not yet resolved). Used by the daemon's @@ -477,6 +661,16 @@ export class McpClientManager { if (idx >= 0) { this.lastRefusedServerNames.splice(idx, 1); } + // PR 14b: keep the transport map aligned with the names list so a + // late-cleared refusal (e.g. operator removed the server) doesn't + // leave stale transport metadata that would surface in a future + // batch event if the same name later got refused again. + this.lastRefusedTransports.delete(serverName); + // PR 14b: drop the name from the unsent-refusals queue too. If it + // was queued but not yet emitted, the operator action that + // cleared it (disconnect, server removed) makes the queued + // event stale; if it was already emitted, this is a no-op. + this.pendingRefusalNames.delete(serverName); } /** @@ -494,10 +688,27 @@ export class McpClientManager { * grown. The stderr line still fires so the operator sees the * refusal at every reproduction. */ - private refuseAndLog(serverName: string): void { + private refuseAndLog( + serverName: string, + serverConfig: MCPServerConfig | undefined, + ): void { if (!this.lastRefusedServerNames.includes(serverName)) { this.lastRefusedServerNames.push(serverName); } + // PR 14b: record the transport family at refusal time so the + // `refused_batch` event payload can break it down. Latest-write + // wins: a duplicate refusal in the same pass updates the entry + // instead of growing the names list (mirrors the `.includes` + // guard above). + this.lastRefusedTransports.set( + serverName, + serverConfig ? mcpTransportOf(serverConfig) : 'unknown', + ); + // PR 14b: queue the name for the next push-event emit. Set + // semantics make repeated `refuseAndLog` for the same name in + // one pass collapse into one queued entry (matches the + // `lastRefusedServerNames.includes` guard above). + this.pendingRefusalNames.add(serverName); process.stderr.write( `qwen serve: MCP server '${serverName}' refused (budget exhausted, ` + `budget=${this.clientBudget}, mode=enforce)\n`, @@ -530,6 +741,191 @@ export class McpClientManager { }); } + /** + * PR 14b: hysteresis state machine for `'budget_warning'` events. + * Called at end of each discovery pass and in the `readResource` + * lazy-spawn path after a successful slot reservation. + * + * Invariants: + * - In `off` mode or with no budget configured: hard no-op. + * `warnArmed` stays at its initial `true`, never read or + * mutated. The constructor's `onBudgetEvent` capture is + * `undefined` in `off` mode, so an accidental call wouldn't + * fire anyway — defense in depth. + * - Trigger is `reservedSlots.size / clientBudget`, NOT + * `liveCount / clientBudget`. Reservations include in-flight + * connects and survive transient `disconnectServer` calls, + * making the trigger stable against connect/disconnect + * chatter. Payload exposes BOTH so SDK consumers can pick. + * - One fire per upward 75% crossing; no fire while the ratio + * stays at or above 0.75; re-arms only on dropping below + * 0.375. Mirrors `slow_client_warning`'s hysteresis exactly. + */ + private evaluateBudgetState(): void { + if (this.budgetMode === 'off' || this.clientBudget === undefined) return; + const ratio = this.reservedSlots.size / this.clientBudget; + if (this.warnArmed && ratio >= MCP_BUDGET_WARN_FRACTION) { + this.warnArmed = false; + // PR 14b fix #1 (codex round 3): visibility for oncall — + // pre-fix `evaluateBudgetState` had ZERO log output, so + // operators couldn't distinguish "events emitted but + // dropped downstream" from "events never emitted." Mirrors + // the stderr breadcrumb in `refuseAndLog` for the refusal + // side; warning side now has its own debug trail. + debugLogger.info( + `MCP budget warning fired (ratio=${ratio.toFixed(2)}, ` + + `reservedCount=${this.reservedSlots.size}, ` + + `budget=${this.clientBudget}, mode=${this.budgetMode})`, + ); + this.emitBudgetEvent({ + kind: 'budget_warning', + liveCount: this.getMcpClientAccounting().total, + reservedCount: this.reservedSlots.size, + budget: this.clientBudget, + thresholdRatio: MCP_BUDGET_WARN_FRACTION, + mode: this.budgetMode, + }); + } else if (!this.warnArmed && ratio < MCP_BUDGET_REARM_FRACTION) { + this.warnArmed = true; + // PR 14b fix #1 (codex round 3): re-arm transitions are silent + // by design (no SDK event), but operators dashboarding budget + // pressure benefit from knowing the manager has re-armed — + // the next 75% crossing will fire a fresh warning. + debugLogger.info( + `MCP budget warning re-armed (ratio=${ratio.toFixed(2)}, ` + + `budget=${this.clientBudget}; next 75% crossing will fire)`, + ); + } + } + + /** + * PR 14b: coalesce per-pass refusals into a single `'refused_batch'` + * event. Called at end of `discoverAllMcpTools` and + * `discoverAllMcpToolsIncremental`, plus the `readResource` lazy- + * spawn refusal path (where it emits a length-1 batch for shape + * consistency). + * + * Idempotent on empty queue: when `pendingRefusalNames.size === 0` + * the call short-circuits without firing or clearing. + * + * What gets cleared on a successful emit: + * - `pendingRefusalNames` — drained, so a follow-up + * `emitRefusedBatchIfAny` in the same pass is a no-op. + * + * What does NOT get cleared on emit (codex round 3 doc fix): + * - `lastRefusedServerNames` — snapshot-visible, must survive + * between passes so `GET /workspace/mcp` reports the last + * refusal set even after the push event fired. + * - `lastRefusedTransports` — sidecar of the names list, same + * lifetime: reset at start of each pass / `stop()` / + * `dropRefusalEntry`, NOT on emit. + * + * `mode: 'enforce'` is a literal: `warn` mode never refuses, so the + * code path that calls `refuseAndLog` (the only writer of + * `lastRefusedServerNames`) is reachable only under `enforce`. + */ + private emitRefusedBatchIfAny(): void { + // PR 14b fix #3 (codex review round 1): suppress inline emit while + // a bulk pass is active. The bulk pass's terminal emit (after + // `bulkPassDepth--` in its `finally`) will drain the queue once. + // This preserves the documented "one batch per `discoverAllMcpTools*` + // pass" contract — pre-fix, every per-server refusal inside an + // incremental pass produced its own length-1 batch, breaking the + // contract for the most common refusal scenario. + if (this.bulkPassDepth > 0) return; + if (this.pendingRefusalNames.size === 0) return; + if (this.clientBudget === undefined || this.budgetMode !== 'enforce') { + // Defensive: refusals queued without `enforce` + budget means + // some upstream path mis-reserved. Drain the queue so it + // doesn't loop into the next pass; skip the emit (we can't + // build a truthful payload without a real budget value). + // + // PR 14b fix (codex round 6): pre-fix this branch was silent. + // The two writers of `pendingRefusalNames` (`refuseAndLog`) + // are gated on `enforce` mode, so reaching this point means + // an invariant violation. Surface the regression at debug + // level so a future bug can be diagnosed by flipping debug + // on, not by reverse-engineering missing telemetry. + debugLogger.warn( + `MCP guardrail: dropped ${this.pendingRefusalNames.size} ` + + `pending refusal(s) — invariant violation ` + + `(budget=${this.clientBudget}, mode=${this.budgetMode}). ` + + `This branch should be unreachable; investigate the ` + + `refuseAndLog call sites.`, + ); + this.pendingRefusalNames.clear(); + return; + } + // PR 14b: emit names in `lastRefusedServerNames` insertion order, + // restricted to the not-yet-emitted set. Insertion order matches + // config-declaration order (the loop in `discoverAllMcpTools*` + // uses `Object.entries`), giving SDK consumers a deterministic + // ordering across reconnects. + const namesInOrder = this.lastRefusedServerNames.filter((n) => + this.pendingRefusalNames.has(n), + ); + if (namesInOrder.length === 0) { + // The pending set is non-empty but none of the names appear in + // `lastRefusedServerNames` — shouldn't happen given `refuseAndLog` + // adds to both. Drain defensively to avoid a stuck queue. + // + // PR 14b fix (codex round 6): same rationale as the + // budget/mode invariant branch above — surface unreachable + // states so future regressions are diagnosable. + debugLogger.warn( + `MCP guardrail: dropped ${this.pendingRefusalNames.size} ` + + `pending refusal(s) — names absent from ` + + `lastRefusedServerNames (the two writers in refuseAndLog ` + + `are paired; reaching this branch indicates a sync gap).`, + ); + this.pendingRefusalNames.clear(); + return; + } + const refusedServers: McpRefusedServer[] = namesInOrder.map((name) => ({ + name, + transport: this.lastRefusedTransports.get(name) ?? 'unknown', + reason: 'budget_exhausted' as const, + })); + this.emitBudgetEvent({ + kind: 'refused_batch', + refusedServers, + budget: this.clientBudget, + liveCount: this.getMcpClientAccounting().total, + reservedCount: this.reservedSlots.size, + mode: 'enforce', + }); + this.pendingRefusalNames.clear(); + } + + /** + * PR 14b fix (codex round 3): single boundary for `onBudgetEvent` + * invocation. The manager's state machine and refused-batch + * coalescer both call this — the production ACP adapter wraps its + * extNotification in `void ... .catch()` so async failures don't + * leak, but the callback ITSELF could throw synchronously (a future + * test fixture, a buggy adapter, an unexpected serialization + * crash). Without this guard, the throw would propagate into MCP + * discovery / `readResource` / `disconnectServer` paths and abort + * unrelated work — budget push events are best-effort telemetry, + * NEVER critical-path. + * + * Logs at `debug` level so production daemons stay quiet on the + * happy path; oncall flips debug on when investigating an MCP + * guardrail incident and sees both delivery successes (via + * `evaluateBudgetState`'s info logs) and failures. + */ + private emitBudgetEvent(event: McpBudgetEvent): void { + if (!this.onBudgetEvent) return; + try { + this.onBudgetEvent(event); + } catch (err) { + debugLogger.debug( + `MCP budget event callback threw (kind=${event.kind}): ` + + `${err instanceof Error ? err.message : String(err)}`, + ); + } + } + /** * Initiates the tool discovery process for all configured MCP servers. * It connects to each server, discovers its available tools, and registers @@ -546,110 +942,136 @@ export class McpClientManager { this.cliConfig.getMcpServerCommand(), ); - this.discoveryState = MCPDiscoveryState.IN_PROGRESS; - // Reset per-pass refusal log so a snapshot taken after this pass - // reflects THIS pass's refusals, not a stale one. Reservations - // (this.reservedSlots) persist across passes — they're keyed by - // server name, which is the operator's intent unit. - this.lastRefusedServerNames = []; + // PR 14b fix #3 (codex review round 1): mark the bulk pass active + // so per-server `emitRefusedBatchIfAny` calls (which the inner + // `discoverMcpToolsForServer` path makes when it refuses a slot) + // queue the names instead of firing length-1 batches inline. The + // matching `bulkPassDepth--` + terminal `emitRefusedBatchIfAny` + // run after `Promise.all` resolves, draining the queue once as + // a coalesced length-N batch. + this.bulkPassDepth++; + try { + this.discoveryState = MCPDiscoveryState.IN_PROGRESS; + // Reset per-pass refusal log so a snapshot taken after this pass + // reflects THIS pass's refusals, not a stale one. Reservations + // (this.reservedSlots) persist across passes — they're keyed by + // server name, which is the operator's intent unit. + this.lastRefusedServerNames = []; + // PR 14b: keep the transport sidecar aligned with the names list, + // and drain any unsent refusal queue from a prior pass so it + // can't bleed into this pass's batch. + this.lastRefusedTransports.clear(); + this.pendingRefusalNames.clear(); - this.eventEmitter?.emit('mcp-client-update', this.clients); - const discoveryPromises = Object.entries(servers).map( - async ([name, config]) => { - // Skip disabled servers - if (cliConfig.isMcpServerDisabled(name)) { - debugLogger.debug(`Skipping disabled MCP server: ${name}`); - return; - } + this.eventEmitter?.emit('mcp-client-update', this.clients); + const discoveryPromises = Object.entries(servers).map( + async ([name, config]) => { + // Skip disabled servers + if (cliConfig.isMcpServerDisabled(name)) { + debugLogger.debug(`Skipping disabled MCP server: ${name}`); + return; + } - // Budget gate (PR 14): synchronous slot reservation BEFORE the - // `await client.connect()` below. Refusal only happens under - // `enforce` mode; `warn` mode reserves regardless so accounting - // reflects the configured set. `off` is a no-op. - const reservation = this.tryReserveSlot(name); - if (reservation === 'refused') { - this.refuseAndLog(name); - return; - } + // Budget gate (PR 14): synchronous slot reservation BEFORE the + // `await client.connect()` below. Refusal only happens under + // `enforce` mode; `warn` mode reserves regardless so accounting + // reflects the configured set. `off` is a no-op. + const reservation = this.tryReserveSlot(name); + if (reservation === 'refused') { + this.refuseAndLog(name, config); + return; + } - // For SDK MCP servers, pass the sendSdkMcpMessage callback - const sdkCallback = isSdkMcpServerConfig(config) - ? this.sendSdkMcpMessage - : undefined; - - const client = new McpClient( - name, - config, - this.toolRegistry, - this.cliConfig.getPromptRegistry(), - this.cliConfig.getWorkspaceContext(), - this.cliConfig.getDebugMode(), - sdkCallback, - ); - this.clients.set(name, client); + // For SDK MCP servers, pass the sendSdkMcpMessage callback + const sdkCallback = isSdkMcpServerConfig(config) + ? this.sendSdkMcpMessage + : undefined; + + const client = new McpClient( + name, + config, + this.toolRegistry, + this.cliConfig.getPromptRegistry(), + this.cliConfig.getWorkspaceContext(), + this.cliConfig.getDebugMode(), + sdkCallback, + ); + this.clients.set(name, client); - this.eventEmitter?.emit('mcp-client-update', this.clients); - try { - await client.connect(); - await client.discover(cliConfig); this.eventEmitter?.emit('mcp-client-update', this.clients); - } catch (error) { - // PR 14 fix (review #4247 wenshao C2): zombie slot leak. - // `tryReserveSlot(name)` reserved a slot above. If `connect()` - // throws, the slot would stay reserved forever and the client - // entry would stay in `this.clients` in a never-CONNECTED - // state, blocking other servers in `enforce` mode until a - // full discovery restart. Release both so the budget cap - // reflects actual usable capacity. - // - // Slot bookkeeping in this bulk path is partially redundant - // with `await this.stop()` at the top of - // `discoverAllMcpTools` (line ~320) — the next bulk run - // wipes `reservedSlots` regardless. But the SAME catch - // ALSO needs to handle the transport (see below): the - // client object held by `clients.delete(name)` only had - // its tracking reference removed, not its underlying - // transport closed. Leaving the orphan transport alive - // would leak the stdio child / WebSocket / HTTP socket - // for the rest of the process — `stop()` can't clean it - // because we just removed it from the map. - // - // The per-server reconnect path - // (`discoverMcpToolsForServerInternal`) keeps the slot - // when `weReservedSlot === false` so health-monitor retry - // doesn't have to compete for capacity — different - // lifecycle, different contract. Bulk path always releases - // because every server is "fresh" here (preceded by - // stop()). - // - // PR 14 fix (review #4247 wenshao R8 #1 line 532): also - // call `await client.disconnect()` BEFORE dropping the - // reference. R7 #3 fixed the analogous leak in the - // per-server path; this is the bulk-path mirror. Errors - // intentionally swallowed (we're already in a discovery- - // failure catch; double-throwing would lose the original - // error context). try { - await client.disconnect(); - } catch { - // best-effort transport cleanup + await client.connect(); + await client.discover(cliConfig); + this.eventEmitter?.emit('mcp-client-update', this.clients); + } catch (error) { + // PR 14 fix (review #4247 wenshao C2): zombie slot leak. + // `tryReserveSlot(name)` reserved a slot above. If `connect()` + // throws, the slot would stay reserved forever and the client + // entry would stay in `this.clients` in a never-CONNECTED + // state, blocking other servers in `enforce` mode until a + // full discovery restart. Release both so the budget cap + // reflects actual usable capacity. + // + // Slot bookkeeping in this bulk path is partially redundant + // with `await this.stop()` at the top of + // `discoverAllMcpTools` (line ~320) — the next bulk run + // wipes `reservedSlots` regardless. But the SAME catch + // ALSO needs to handle the transport (see below): the + // client object held by `clients.delete(name)` only had + // its tracking reference removed, not its underlying + // transport closed. Leaving the orphan transport alive + // would leak the stdio child / WebSocket / HTTP socket + // for the rest of the process — `stop()` can't clean it + // because we just removed it from the map. + // + // The per-server reconnect path + // (`discoverMcpToolsForServerInternal`) keeps the slot + // when `weReservedSlot === false` so health-monitor retry + // doesn't have to compete for capacity — different + // lifecycle, different contract. Bulk path always releases + // because every server is "fresh" here (preceded by + // stop()). + // + // PR 14 fix (review #4247 wenshao R8 #1 line 532): also + // call `await client.disconnect()` BEFORE dropping the + // reference. R7 #3 fixed the analogous leak in the + // per-server path; this is the bulk-path mirror. Errors + // intentionally swallowed (we're already in a discovery- + // failure catch; double-throwing would lose the original + // error context). + try { + await client.disconnect(); + } catch { + // best-effort transport cleanup + } + this.releaseSlotName(name); + this.clients.delete(name); + this.eventEmitter?.emit('mcp-client-update', this.clients); + // Log the error but don't let a single failed server stop the others + debugLogger.error( + `Error during discovery for server '${name}': ${getErrorMessage( + error, + )}`, + ); } - this.reservedSlots.delete(name); - this.clients.delete(name); - this.eventEmitter?.emit('mcp-client-update', this.clients); - // Log the error but don't let a single failed server stop the others - debugLogger.error( - `Error during discovery for server '${name}': ${getErrorMessage( - error, - )}`, - ); - } - }, - ); + }, + ); - await Promise.all(discoveryPromises); - this.discoveryState = MCPDiscoveryState.COMPLETED; - this.emitBudgetTelemetry(Object.keys(servers).length); + await Promise.all(discoveryPromises); + this.discoveryState = MCPDiscoveryState.COMPLETED; + this.emitBudgetTelemetry(Object.keys(servers).length); + } finally { + // PR 14b fix #3: drop the bulk-pass marker BEFORE the terminal + // emit so `emitRefusedBatchIfAny` actually fires (its early- + // return guard reads `bulkPassDepth`). The warning event fires + // inline from `tryReserveSlot` / `releaseSlotName` whenever a + // slot mutation crosses the 75% threshold — codex review fix + // #4 — so no terminal `evaluateBudgetState` is needed here. + // Refused batch is the only deferred emit (coalesced over the + // whole pass — fix #3 makes this a strict invariant). + this.bulkPassDepth--; + this.emitRefusedBatchIfAny(); + } } /** @@ -726,7 +1148,13 @@ export class McpClientManager { // failure. const reservation = this.tryReserveSlot(serverName); if (reservation === 'refused') { - this.refuseAndLog(serverName); + this.refuseAndLog(serverName, serverConfig); + // PR 14b: single-server refusal (e.g. health-monitor retry into + // a full budget, `/mcp reconnect `) emits a length-1 + // batch for shape consistency with the bulk-pass refusal. + // Operators / dashboards see one event shape regardless of + // entrypoint. + this.emitRefusedBatchIfAny(); return; } // PR 14 fix (review #4247 wenshao R3-R4): track whether THIS call @@ -808,6 +1236,11 @@ export class McpClientManager { // snapshots immediately reflect reality. Mirrors the same // pattern in `readResource`'s late-reserve branch. this.dropRefusalEntry(serverName); + // PR 14b fix #4: hysteresis is driven inline by + // `tryReserveSlot` (upward) and `releaseSlotName` (downward). + // The standalone `evaluateBudgetState` that used to live here + // is now redundant — the reservation that opened this branch + // already fired the warning if it crossed 75%. } catch (error) { // PR 14 fix (review #4247 wenshao R3 line 546): two-mode // cleanup for connect failure, matching the `readResource` @@ -847,7 +1280,7 @@ export class McpClientManager { } catch { // best-effort transport cleanup } - this.reservedSlots.delete(serverName); + this.releaseSlotName(serverName); this.clients.delete(serverName); } // Log the error but don't throw: callers expect best-effort discovery. @@ -899,6 +1332,13 @@ export class McpClientManager { this.reservedSlots.clear(); this.freshReservations.clear(); this.lastRefusedServerNames = []; + // PR 14b: post-`stop` the manager is fresh — clear refusal + // transport sidecar, drain the unsent-refusal queue, and re-arm + // the warning state machine so the next discovery pass that + // crosses 75% fires anew. + this.lastRefusedTransports.clear(); + this.pendingRefusalNames.clear(); + this.warnArmed = true; } /** @@ -934,7 +1374,7 @@ export class McpClientManager { // internal reconnect path (`discoverMcpToolsForServerInternal`) // calls `existingClient.disconnect()` directly, NOT this public // method, so reconnect still doesn't release the slot. - this.reservedSlots.delete(serverName); + this.releaseSlotName(serverName); this.dropRefusalEntry(serverName); } @@ -1106,164 +1546,190 @@ export class McpClientManager { this.cliConfig.getMcpServerCommand(), ); - this.discoveryState = MCPDiscoveryState.IN_PROGRESS; - // Reset per-pass refusal log; see the sibling reset in - // `discoverAllMcpTools` for rationale. - this.lastRefusedServerNames = []; - recordStartupEvent('mcp_discovery_start', { - serverCount: Object.keys(servers).length, - incremental: true, - }); - // Mirrors `discoverAllMcpTools`: announce IN_PROGRESS so UI subscribers - // (MCP status pill, AppContainer batch-flush effect) know discovery - // started, even when no servers need updates this pass. - this.eventEmitter?.emit('mcp-client-update', this.clients); - - // Tracks the first successful server discover so we can emit the - // `mcp_first_tool_registered` event exactly once. "First successful - // discover" rather than a tool-count delta — simpler and aligns with the - // user-perceived metric ("first MCP server is ready"). - let firstToolEventFired = false; - - // Find servers that are new or have changed configuration - const serversToUpdate: string[] = []; - const currentServerNames = new Set(this.clients.keys()); - const newServerNames = new Set(Object.keys(servers)); - - // PR 14 fix (review #4247): process removals BEFORE the new-server - // reservation pass so freed slots are visible to `tryReserveSlot`. - // Scenario: budget=2, currently `{a, b}` reserved, new config - // `{a, c}`. Pre-fix order refused `c` because `b`'s slot was only - // freed after the new-server loop. Now `b` is removed first → - // reservedSlots={a} → `c` reservation succeeds. Disabled-mid-session - // removals stay inline (below) because they also release slots - // via `removeServer`'s `reservedSlots.delete` — same call, just - // reached from a different branch. - for (const name of currentServerNames) { - if (!newServerNames.has(name)) { - // Server was removed from configuration - await this.removeServer(name); - } - } + // PR 14b fix #3 (codex review round 1): suppress per-server + // length-1 batches inside this incremental pass — the + // `discoverMcpToolsForServerInternal` calls below would otherwise + // emit one batch per refused server, breaking the documented + // "one batch per pass" contract. The terminal + // `emitRefusedBatchIfAny` (after `bulkPassDepth--`) drains the + // queue once. + this.bulkPassDepth++; + try { + this.discoveryState = MCPDiscoveryState.IN_PROGRESS; + // Reset per-pass refusal log; see the sibling reset in + // `discoverAllMcpTools` for rationale. + this.lastRefusedServerNames = []; + // PR 14b: keep the transport sidecar aligned with the names list, + // and drain any unsent refusal queue from a prior pass. + this.lastRefusedTransports.clear(); + this.pendingRefusalNames.clear(); + recordStartupEvent('mcp_discovery_start', { + serverCount: Object.keys(servers).length, + incremental: true, + }); + // Mirrors `discoverAllMcpTools`: announce IN_PROGRESS so UI subscribers + // (MCP status pill, AppContainer batch-flush effect) know discovery + // started, even when no servers need updates this pass. + this.eventEmitter?.emit('mcp-client-update', this.clients); - // Check for new servers or configuration changes - for (const [name] of Object.entries(servers)) { - // Mirror `discoverAllMcpTools` (line ~102): users who explicitly - // disabled a server via `mcpServers..disabled: true` must not - // see it reconnected by the incremental path. Without this, the - // PR-A background path silently re-registers tools the user has - // told us to ignore. - if (cliConfig.isMcpServerDisabled(name)) { - debugLogger.debug(`Skipping disabled MCP server: ${name}`); - // If the server was previously enabled and got connected, we now - // need to tear it down — otherwise its client, registered tools - // and health checks linger after an enabled→disabled mid-session - // transition (e.g. via `/mcp disable `). `removeServer` - // disconnects, drops the client entry, removes tools from the - // registry, stops the health check, and removes the global - // status so the Footer pill stops counting it. - if (this.clients.has(name)) { + // Tracks the first successful server discover so we can emit the + // `mcp_first_tool_registered` event exactly once. "First successful + // discover" rather than a tool-count delta — simpler and aligns with the + // user-perceived metric ("first MCP server is ready"). + let firstToolEventFired = false; + + // Find servers that are new or have changed configuration + const serversToUpdate: string[] = []; + const currentServerNames = new Set(this.clients.keys()); + const newServerNames = new Set(Object.keys(servers)); + + // PR 14 fix (review #4247): process removals BEFORE the new-server + // reservation pass so freed slots are visible to `tryReserveSlot`. + // Scenario: budget=2, currently `{a, b}` reserved, new config + // `{a, c}`. Pre-fix order refused `c` because `b`'s slot was only + // freed after the new-server loop. Now `b` is removed first → + // reservedSlots={a} → `c` reservation succeeds. Disabled-mid-session + // removals stay inline (below) because they also release slots + // via `removeServer`'s `reservedSlots.delete` — same call, just + // reached from a different branch. + for (const name of currentServerNames) { + if (!newServerNames.has(name)) { + // Server was removed from configuration await this.removeServer(name); } - continue; } - const existingClient = this.clients.get(name); - if (!existingClient) { - // PR 14 fix (review #4247 wenshao R6 line 956): pre-reservation - // here was a TOCTOU race. The inner - // `discoverMcpToolsForServerInternal` ALSO does `tryReserveSlot` - // (added in R1 fix #1). With BOTH sites reserving, the - // reservation lifecycle didn't align with the timeout - // cleanup site — `runWithDiscoveryTimeout`'s timeout handler - // could release the slot mid-flight while the inner - // `connect()` later resolves successfully, leaving a - // CONNECTED client with NO reservation. Next pass admits - // another new server because `reservedSlots.size < budget`, - // and `enforce` mode silently exceeds the cap. - // - // Fix: delete the pre-reservation. `discoverMcpToolsForServerInternal` - // owns the reservation lifecycle end-to-end (reserve → - // try-catch around connect → release on weReservedSlot - // failure path → cleared by timeout handler if it fires). - // Refusal still happens — just inside the inner call. The - // operator-visible behavior is identical; only the race is - // closed. - serversToUpdate.push(name); - } else if (existingClient.getStatus() === MCPServerStatus.DISCONNECTED) { - // Disconnected server, try to reconnect - serversToUpdate.push(name); + + // Check for new servers or configuration changes + for (const [name] of Object.entries(servers)) { + // Mirror `discoverAllMcpTools` (line ~102): users who explicitly + // disabled a server via `mcpServers..disabled: true` must not + // see it reconnected by the incremental path. Without this, the + // PR-A background path silently re-registers tools the user has + // told us to ignore. + if (cliConfig.isMcpServerDisabled(name)) { + debugLogger.debug(`Skipping disabled MCP server: ${name}`); + // If the server was previously enabled and got connected, we now + // need to tear it down — otherwise its client, registered tools + // and health checks linger after an enabled→disabled mid-session + // transition (e.g. via `/mcp disable `). `removeServer` + // disconnects, drops the client entry, removes tools from the + // registry, stops the health check, and removes the global + // status so the Footer pill stops counting it. + if (this.clients.has(name)) { + await this.removeServer(name); + } + continue; + } + const existingClient = this.clients.get(name); + if (!existingClient) { + // PR 14 fix (review #4247 wenshao R6 line 956): pre-reservation + // here was a TOCTOU race. The inner + // `discoverMcpToolsForServerInternal` ALSO does `tryReserveSlot` + // (added in R1 fix #1). With BOTH sites reserving, the + // reservation lifecycle didn't align with the timeout + // cleanup site — `runWithDiscoveryTimeout`'s timeout handler + // could release the slot mid-flight while the inner + // `connect()` later resolves successfully, leaving a + // CONNECTED client with NO reservation. Next pass admits + // another new server because `reservedSlots.size < budget`, + // and `enforce` mode silently exceeds the cap. + // + // Fix: delete the pre-reservation. `discoverMcpToolsForServerInternal` + // owns the reservation lifecycle end-to-end (reserve → + // try-catch around connect → release on weReservedSlot + // failure path → cleared by timeout handler if it fires). + // Refusal still happens — just inside the inner call. The + // operator-visible behavior is identical; only the race is + // closed. + serversToUpdate.push(name); + } else if ( + existingClient.getStatus() === MCPServerStatus.DISCONNECTED + ) { + // Disconnected server, try to reconnect + serversToUpdate.push(name); + } + // Note: Configuration change detection would require comparing + // the old and new config, which is not implemented here } - // Note: Configuration change detection would require comparing - // the old and new config, which is not implemented here - } - // Update only the servers that need it. Each per-server discover is - // wrapped in a discovery-only timeout (stdio default 30s, remote 5s, - // per-server override via `discoveryTimeoutMs`). Tool-call timeout is - // intentionally left alone — a long-running tool invocation is not a - // startup pathology. - const discoveryPromises = serversToUpdate.map(async (name) => { - const serverConfig = servers[name]; - try { - await this.runWithDiscoveryTimeout(name, serverConfig, () => - this.discoverMcpToolsForServer(name, cliConfig), - ); - // `discoverMcpToolsForServerInternal` swallows connect/discover - // errors (best-effort discovery semantics — see its catch block), - // so the try here resolves even for failed servers. Only the - // timeout path reaches the catch below. Consult the actual - // server status to decide which outcome to record, otherwise - // every auth failure / crash / "no tools found" looks like - // `ready` in the startup profile. - const client = this.clients.get(name); - const actuallyReady = - !!client && getMCPServerStatus(name) === MCPServerStatus.CONNECTED; - if (actuallyReady) { - if (!firstToolEventFired) { - firstToolEventFired = true; - recordStartupEvent('mcp_first_tool_registered', { - serverName: name, + // Update only the servers that need it. Each per-server discover is + // wrapped in a discovery-only timeout (stdio default 30s, remote 5s, + // per-server override via `discoveryTimeoutMs`). Tool-call timeout is + // intentionally left alone — a long-running tool invocation is not a + // startup pathology. + const discoveryPromises = serversToUpdate.map(async (name) => { + const serverConfig = servers[name]; + try { + await this.runWithDiscoveryTimeout(name, serverConfig, () => + this.discoverMcpToolsForServer(name, cliConfig), + ); + // `discoverMcpToolsForServerInternal` swallows connect/discover + // errors (best-effort discovery semantics — see its catch block), + // so the try here resolves even for failed servers. Only the + // timeout path reaches the catch below. Consult the actual + // server status to decide which outcome to record, otherwise + // every auth failure / crash / "no tools found" looks like + // `ready` in the startup profile. + const client = this.clients.get(name); + const actuallyReady = + !!client && getMCPServerStatus(name) === MCPServerStatus.CONNECTED; + if (actuallyReady) { + if (!firstToolEventFired) { + firstToolEventFired = true; + recordStartupEvent('mcp_first_tool_registered', { + serverName: name, + }); + } + recordStartupEvent(`mcp_server_ready:${name}`, { + outcome: 'ready', + }); + } else { + recordStartupEvent(`mcp_server_ready:${name}`, { + outcome: 'failed', + reason: 'connect or discover error', }); } - recordStartupEvent(`mcp_server_ready:${name}`, { outcome: 'ready' }); - } else { + } catch (error) { + // Defensive cleanup: the dedup Map entry is normally removed by + // `discoverMcpToolsForServer`'s `finally`, but `runWithDiscoveryTimeout` + // can reject before that finally runs (the timeout also disconnects + // the client to abort the underlying handshake). Without this + // explicit delete, a brief window exists where a subsequent + // `discoverMcpToolsForServer(name)` call would short-circuit on + // a now-doomed promise. + this.serverDiscoveryPromises.delete(name); recordStartupEvent(`mcp_server_ready:${name}`, { outcome: 'failed', - reason: 'connect or discover error', + reason: getErrorMessage(error), }); + debugLogger.error( + `Error during incremental discovery for server '${name}': ${getErrorMessage(error)}`, + ); } - } catch (error) { - // Defensive cleanup: the dedup Map entry is normally removed by - // `discoverMcpToolsForServer`'s `finally`, but `runWithDiscoveryTimeout` - // can reject before that finally runs (the timeout also disconnects - // the client to abort the underlying handshake). Without this - // explicit delete, a brief window exists where a subsequent - // `discoverMcpToolsForServer(name)` call would short-circuit on - // a now-doomed promise. - this.serverDiscoveryPromises.delete(name); - recordStartupEvent(`mcp_server_ready:${name}`, { - outcome: 'failed', - reason: getErrorMessage(error), - }); - debugLogger.error( - `Error during incremental discovery for server '${name}': ${getErrorMessage(error)}`, - ); - } - }); + }); - await Promise.all(discoveryPromises); + await Promise.all(discoveryPromises); - // Start health checks for all connected servers - if (this.healthConfig.autoReconnect) { - this.startAllHealthChecks(); - } + // Start health checks for all connected servers + if (this.healthConfig.autoReconnect) { + this.startAllHealthChecks(); + } - this.discoveryState = MCPDiscoveryState.COMPLETED; - recordStartupEvent('mcp_all_servers_settled', { - serverCount: Object.keys(servers).length, - incremental: true, - }); - this.emitBudgetTelemetry(Object.keys(servers).length); + this.discoveryState = MCPDiscoveryState.COMPLETED; + recordStartupEvent('mcp_all_servers_settled', { + serverCount: Object.keys(servers).length, + incremental: true, + }); + this.emitBudgetTelemetry(Object.keys(servers).length); + } finally { + // PR 14b fix #3: drop the bulk marker BEFORE the terminal + // emit so `emitRefusedBatchIfAny` actually fires the coalesced + // batch. Warning fires inline from `tryReserveSlot` / + // `releaseSlotName` (fix #4) — no terminal + // `evaluateBudgetState` here. + this.bulkPassDepth--; + this.emitRefusedBatchIfAny(); + } // Trailing `mcp-client-update` AFTER flipping discoveryState to // COMPLETED. Without this the per-server updates above all fire while // the state is still IN_PROGRESS, so the AppContainer batch-flush @@ -1349,7 +1815,7 @@ export class McpClientManager { // R8 #4 caught the asymmetry with the connect-failure // path's `weReservedSlot` guard. Now they match. if (this.freshReservations.has(serverName)) { - this.reservedSlots.delete(serverName); + this.releaseSlotName(serverName); this.freshReservations.delete(serverName); } // And drop any stale refusal entry — operator intent shifts @@ -1442,7 +1908,7 @@ export class McpClientManager { // the budget slot too — operator intent is "this server should not // be running", so it must not block a different server from taking // its place on the next discovery pass. - this.reservedSlots.delete(serverName); + this.releaseSlotName(serverName); // PR 14 fix (review #4247): also drop the entry from the per-pass // refusal log so a snapshot taken between discoveries doesn't // stale-tag the (now-disabled or now-removed) server as @@ -1526,7 +1992,12 @@ export class McpClientManager { // throw so operators get the same stderr trail as bulk // discovery refusals — the throw alone doesn't surface to // stderr (caller decides what to do with the typed error). - this.refuseAndLog(serverName); + this.refuseAndLog(serverName, serverConfig); + // PR 14b: lazy-spawn refusal emits a length-1 batch BEFORE + // throwing so SDK consumers see the structured event whether + // or not they catch the typed error. Order matches the + // discovery paths: emit, then throw / return. + this.emitRefusedBatchIfAny(); throw new BudgetExhaustedError( serverName, this.clientBudget as number, @@ -1551,6 +2022,9 @@ export class McpClientManager { // the next snapshot reflects the late-reservation success. if (weReservedSlot) { this.dropRefusalEntry(serverName); + // PR 14b fix #4 (codex review round 1): no inline evaluate + // needed — `tryReserveSlot` already fired the warning if + // the upward crossing happened during reservation. } const sdkCallback = isSdkMcpServerConfig(serverConfig) @@ -1666,7 +2140,7 @@ export class McpClientManager { } catch { // best-effort transport cleanup } - this.reservedSlots.delete(serverName); + this.releaseSlotName(serverName); this.clients.delete(serverName); this.eventEmitter?.emit('mcp-client-update', this.clients); } diff --git a/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts b/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts index 07303a6877c..cae2ca111bb 100644 --- a/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts +++ b/packages/sdk-typescript/src/daemon/DaemonSessionClient.ts @@ -81,14 +81,34 @@ export class DaemonSessionClient { clientId?: string, ): Promise { const session = await client.createOrAttachSession(req, clientId); - // `modelServiceId` switch failures are reported on SSE, not the - // create/attach HTTP response. Seed the first subscription from the - // daemon replay ring so create-then-subscribe clients observe attach-time - // `model_switch_failed` / `model_switched` events. The daemon treats - // Last-Event-ID: 0 as "replay from the beginning of the bounded ring"; - // if older events have already been evicted, clients receive the retained - // suffix and continue live from there. - const lastEventId = req.modelServiceId ? 0 : undefined; + // Seed the first subscription from the daemon replay ring whenever + // events can fire during the session-creation window — otherwise + // they land in the per-session ring before the consumer's first + // `events()` call and never reach the live stream. + // + // Two such windows exist today: + // - **Newly-created sessions** (`session.attached === false`): the + // child's `newSession` handler runs MCP discovery synchronously + // in legacy blocking mode and as background work in progressive + // mode. PR 14b's `mcp_budget_warning` / `mcp_child_refused_batch` + // push events fire during this window and are buffered on + // `BridgeClient.earlyEvents` until `byId.set` runs, then drained + // into the per-session bus before `spawnOrAttach` returns. The + // guardrail events advertised via `mcp_guardrail_events` are + // useless without this seed because they predate any live + // subscription. + // - **Pre-PR 14b carve-out**: `modelServiceId` switch failures are + // reported on SSE, not the create/attach HTTP response. The + // original carve-out covered just this case; the unified rule + // below subsumes it (newly-created sessions always seed) while + // preserving the semantics for re-attached sessions where the + // caller may have an existing event cursor it doesn't want to + // reset. + // + // The daemon treats Last-Event-ID: 0 as "replay from the beginning + // of the bounded ring"; if older events have already been evicted, + // clients receive the retained suffix and continue live from there. + const lastEventId = !session.attached || req.modelServiceId ? 0 : undefined; return new DaemonSessionClient({ client, session, lastEventId }); } diff --git a/packages/sdk-typescript/src/daemon/events.ts b/packages/sdk-typescript/src/daemon/events.ts index 27a52ed01dc..f24511ffa6d 100644 --- a/packages/sdk-typescript/src/daemon/events.ts +++ b/packages/sdk-typescript/src/daemon/events.ts @@ -4,7 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { DaemonEvent, PermissionOutcome } from './types.js'; +import type { + DaemonEvent, + DaemonMcpTransport, + PermissionOutcome, +} from './types.js'; const DAEMON_KNOWN_EVENT_TYPE_VALUES = [ 'session_update', @@ -19,6 +23,13 @@ const DAEMON_KNOWN_EVENT_TYPE_VALUES = [ 'client_evicted', 'slow_client_warning', 'stream_error', + // PR 14b — MCP guardrail push events. See `mcp_guardrail_events` + // capability tag. Both fire on the per-session SSE bus; consumers + // should pre-flight `caps.features.includes('mcp_guardrail_events')` + // before relying on these for non-snapshot UX (the `GET /workspace/mcp` + // snapshot still encodes the same state). + 'mcp_budget_warning', + 'mcp_child_refused_batch', // Issue #4175 PR 16: workspace-level mutation signals fanned out // through every active session's bus. Non-terminal — informational // for adapters that want to render "memory just changed" / "agent X @@ -146,6 +157,59 @@ export interface DaemonStreamErrorData { [key: string]: unknown; } +/** + * PR 14b: payload for the `mcp_budget_warning` SSE frame. Fired on the + * upward 75% crossing of `reservedSlots.size / clientBudget`. Re-arms + * only after the ratio drops below 37.5% — so a budget that flaps just + * above the threshold doesn't produce a flood of identical warnings. + * + * `liveCount` (CONNECTED clients) and `reservedCount` (configured set, + * including in-flight reservations) are exposed separately so SDK + * consumers can render either lens. The snapshot (`GET /workspace/mcp`) + * is the source of truth for state-after-reconnect; this event is the + * change-edge. + * + * `mode` is `'warn' | 'enforce'` because the warning fires in either + * mode (only `'off'` skips the state machine entirely). + */ +export interface DaemonMcpBudgetWarningData { + liveCount: number; + reservedCount: number; + budget: number; + thresholdRatio: 0.75; + mode: 'warn' | 'enforce'; + [key: string]: unknown; +} + +/** + * PR 14b: per-server entry inside a `mcp_child_refused_batch` payload. + * `transport` is the family resolved at refusal time via the daemon's + * `mcpTransportOf` helper; future refusal causes (Wave 5+) would + * extend `reason` beyond `'budget_exhausted'`. + */ +export interface DaemonMcpRefusedServer { + name: string; + transport: DaemonMcpTransport; + reason: 'budget_exhausted'; + [key: string]: unknown; +} + +/** + * PR 14b: payload for the `mcp_child_refused_batch` SSE frame. Fires + * once per `discoverAllMcpTools*` pass when at least one server was + * refused, OR as a length-1 batch on the `readResource` lazy-spawn + * refusal path. `mode` is the literal `'enforce'` because `warn` mode + * never refuses (so this event never fires under `warn`). + */ +export interface DaemonMcpChildRefusedBatchData { + refusedServers: DaemonMcpRefusedServer[]; + budget: number; + liveCount: number; + reservedCount: number; + mode: 'enforce'; + [key: string]: unknown; +} + /** * Issue #4175 PR 16: a `POST /workspace/memory` write completed * successfully. `scope` records which file was touched (workspace QWEN.md @@ -379,6 +443,14 @@ export type DaemonStreamErrorEvent = DaemonEventEnvelope< 'stream_error', DaemonStreamErrorData >; +export type DaemonMcpBudgetWarningEvent = DaemonEventEnvelope< + 'mcp_budget_warning', + DaemonMcpBudgetWarningData +>; +export type DaemonMcpChildRefusedBatchEvent = DaemonEventEnvelope< + 'mcp_child_refused_batch', + DaemonMcpChildRefusedBatchData +>; export type DaemonMemoryChangedEvent = DaemonEventEnvelope< 'memory_changed', DaemonMemoryChangedData @@ -459,6 +531,17 @@ export type DaemonStreamLifecycleEvent = | DaemonSlowClientWarningEvent | DaemonStreamErrorEvent; +/** + * PR 14b: MCP guardrail push events. Grouped as their own union member + * (rather than folded into `DaemonStreamLifecycleEvent`) because they + * report McpClientManager state, not the SSE subscriber's queue health + * or the daemon's stream lifecycle. Adapters that only care about + * "is the stream alive" can ignore this whole branch. + */ +export type DaemonMcpGuardrailEvent = + | DaemonMcpBudgetWarningEvent + | DaemonMcpChildRefusedBatchEvent; + /** * Issue #4175 PR 16: workspace-level mutation signals fanned out * through every active session's bus. Non-terminal; clients use them @@ -472,6 +555,7 @@ export type KnownDaemonEvent = | DaemonSessionEvent | DaemonControlEvent | DaemonStreamLifecycleEvent + | DaemonMcpGuardrailEvent | DaemonWorkspaceMutationEvent | DaemonAuthEvent; @@ -509,6 +593,25 @@ export interface DaemonSessionViewState { */ slowClientWarningCount: number; lastSlowClientWarning?: DaemonSlowClientWarningData; + /** + * PR 14b: count of `mcp_budget_warning` frames this stream has + * observed. Non-terminal — warning fires on the upward 75% crossing + * and re-arms below 37.5%, so a flapping budget produces at most + * one warning per crossing episode. Adapters tap this counter to + * surface MCP-pressure UI; the snapshot at `GET /workspace/mcp` + * still carries the authoritative state-after-reconnect. + */ + mcpBudgetWarningCount: number; + lastMcpBudgetWarning?: DaemonMcpBudgetWarningData; + /** + * PR 14b: count of `mcp_child_refused_batch` frames this stream has + * observed. Each frame is a single batch (per discovery pass, or + * length-1 from `readResource`'s lazy-spawn refusal); the count + * reflects batches not refused-server entries. Mirrors the + * snapshot's `disabledReason: 'budget'` per-server tag. + */ + mcpChildRefusedBatchCount: number; + lastMcpChildRefusedBatch?: DaemonMcpChildRefusedBatchData; /** * Issue #4175 PR 16: most recent workspace mutation observed on this * stream (memory or agent change). Non-terminal — adapters render a @@ -581,6 +684,10 @@ export function createDaemonSessionViewState( seed.lastUnmatchedPermissionResolutionId, slowClientWarningCount: seed.slowClientWarningCount ?? 0, lastSlowClientWarning: seed.lastSlowClientWarning, + mcpBudgetWarningCount: seed.mcpBudgetWarningCount ?? 0, + lastMcpBudgetWarning: seed.lastMcpBudgetWarning, + mcpChildRefusedBatchCount: seed.mcpChildRefusedBatchCount ?? 0, + lastMcpChildRefusedBatch: seed.lastMcpChildRefusedBatch, lastWorkspaceMutation: seed.lastWorkspaceMutation, lastWorkspaceMutationType: seed.lastWorkspaceMutationType, approvalMode: seed.approvalMode, @@ -663,6 +770,14 @@ export function asKnownDaemonEvent( return isStreamErrorData(event.data) ? (event as DaemonStreamErrorEvent) : undefined; + case 'mcp_budget_warning': + return isMcpBudgetWarningData(event.data) + ? (event as DaemonMcpBudgetWarningEvent) + : undefined; + case 'mcp_child_refused_batch': + return isMcpChildRefusedBatchData(event.data) + ? (event as DaemonMcpChildRefusedBatchEvent) + : undefined; case 'memory_changed': return isMemoryChangedData(event.data) ? (event as DaemonMemoryChangedEvent) @@ -847,6 +962,24 @@ export function reduceDaemonSessionEvent( streamError: event.data, pendingPermissions: {}, }; + case 'mcp_budget_warning': + // Non-terminal: budget pressure is a status signal, not a stream + // close. Count + capture latest so adapters can render + // "MCP pressure" UI; `alive` and `pendingPermissions` unchanged. + return { + ...base, + mcpBudgetWarningCount: base.mcpBudgetWarningCount + 1, + lastMcpBudgetWarning: event.data, + }; + case 'mcp_child_refused_batch': + // Non-terminal: refusals are operator-actionable signals (raise + // budget / drop servers), not stream lifecycle events. The + // session keeps running with a smaller MCP fleet. + return { + ...base, + mcpChildRefusedBatchCount: base.mcpChildRefusedBatchCount + 1, + lastMcpChildRefusedBatch: event.data, + }; case 'memory_changed': // Non-terminal: adapters render a "memory just changed" hint and // re-fetch `GET /workspace/memory` to get the canonical state. We @@ -1299,6 +1432,71 @@ function isStreamErrorData(value: unknown): value is DaemonStreamErrorData { return isRecord(value) && isNonEmptyString(value['error']); } +function isMcpBudgetWarningData( + value: unknown, +): value is DaemonMcpBudgetWarningData { + // PR 14b fix (codex round 6): `thresholdRatio` is validated as a + // finite number, NOT pinned to the literal `0.75`. The SDK's + // role here is wire-shape validation; threshold semantics are + // owned by the daemon's `MCP_BUDGET_WARN_FRACTION` constant + // (`packages/core/src/tools/mcp-client-manager.ts`) and documented + // in `qwen-serve-protocol.md`. Pinning the literal in the SDK + // would mean a daemon-side change to e.g. 0.80 silently routes + // every warning through `unrecognizedKnownEventCount` — a + // cross-package coordination hazard with no operator-visible + // failure mode. The `DaemonMcpBudgetWarningData.thresholdRatio` + // type still narrows to `0.75` for current daemons; future + // multi-threshold support (e.g. 0.5 critical) would extend the + // type AND the wire shape via a `severity` discriminator field. + return ( + isRecord(value) && + isFiniteNumber(value['liveCount']) && + isFiniteNumber(value['reservedCount']) && + isFiniteNumber(value['budget']) && + isFiniteNumber(value['thresholdRatio']) && + (value['mode'] === 'warn' || value['mode'] === 'enforce') + ); +} + +function isMcpRefusedServerEntry( + value: unknown, +): value is DaemonMcpRefusedServer { + if (!isRecord(value)) return false; + if (!isNonEmptyString(value['name'])) return false; + if (value['reason'] !== 'budget_exhausted') return false; + // Transport family must be one of the known kinds. Reject silently + // for forward-compat: a daemon emitting an unknown transport is + // likely speaking a newer wire than this SDK release. + const transport = value['transport']; + return ( + transport === 'stdio' || + transport === 'sse' || + transport === 'http' || + transport === 'websocket' || + transport === 'sdk' || + transport === 'unknown' + ); +} + +function isMcpChildRefusedBatchData( + value: unknown, +): value is DaemonMcpChildRefusedBatchData { + return ( + isRecord(value) && + Array.isArray(value['refusedServers']) && + value['refusedServers'].every(isMcpRefusedServerEntry) && + isFiniteNumber(value['budget']) && + isFiniteNumber(value['liveCount']) && + isFiniteNumber(value['reservedCount']) && + // `mode` is a literal `'enforce'` — `warn` mode never refuses, so + // `'warn'`-tagged refusal payloads are protocol garbage. Reject + // them so the reducer sees the raw event under the + // `unrecognizedKnownEventCount` branch instead of silently + // accepting a malformed shape. + value['mode'] === 'enforce' + ); +} + function isMemoryChangedData(value: unknown): value is DaemonMemoryChangedData { if (!isRecord(value)) return false; const scope = value['scope']; diff --git a/packages/sdk-typescript/src/daemon/index.ts b/packages/sdk-typescript/src/daemon/index.ts index a034676a18e..8800eda53f8 100644 --- a/packages/sdk-typescript/src/daemon/index.ts +++ b/packages/sdk-typescript/src/daemon/index.ts @@ -61,6 +61,14 @@ export type { DaemonWorkspaceInitializedEvent, DaemonEventEnvelope, DaemonKnownEventType, + // PR 14b — MCP guardrail push-event types. See `mcp_guardrail_events` + // capability tag and the `DaemonMcpGuardrailEvent` union below. + DaemonMcpBudgetWarningData, + DaemonMcpBudgetWarningEvent, + DaemonMcpChildRefusedBatchData, + DaemonMcpChildRefusedBatchEvent, + DaemonMcpGuardrailEvent, + DaemonMcpRefusedServer, DaemonMemoryChangedData, DaemonMemoryChangedEvent, DaemonModelSwitchedData, diff --git a/packages/sdk-typescript/src/index.ts b/packages/sdk-typescript/src/index.ts index 5967b78f163..c9e8281012b 100644 --- a/packages/sdk-typescript/src/index.ts +++ b/packages/sdk-typescript/src/index.ts @@ -50,6 +50,13 @@ export { type DaemonEvent, type DaemonEventEnvelope, type DaemonKnownEventType, + // PR 14b — MCP guardrail push-event types. + type DaemonMcpBudgetWarningData, + type DaemonMcpBudgetWarningEvent, + type DaemonMcpChildRefusedBatchData, + type DaemonMcpChildRefusedBatchEvent, + type DaemonMcpGuardrailEvent, + type DaemonMcpRefusedServer, type DaemonMcpDiscoveryState, type DaemonMcpServerRuntimeStatus, type DaemonMcpTransport, diff --git a/packages/sdk-typescript/test/unit/DaemonSessionClient.test.ts b/packages/sdk-typescript/test/unit/DaemonSessionClient.test.ts index 1146e3e19a9..1943d62eb61 100644 --- a/packages/sdk-typescript/test/unit/DaemonSessionClient.test.ts +++ b/packages/sdk-typescript/test/unit/DaemonSessionClient.test.ts @@ -254,6 +254,51 @@ describe('DaemonSessionClient', () => { expect(calls[1]?.headers['last-event-id']).toBe('0'); }); + it('replays from id 0 on freshly-created sessions so startup-window guardrail events are observable (codex review fix #1)', async () => { + // Codex review round 2, finding #1: PR 14b's + // `mcp_budget_warning` / `mcp_child_refused_batch` events fire + // during the child's `newSession` handler and are buffered on + // `BridgeClient.earlyEvents` until `byId.set(sessionId, entry)` + // runs. The bridge drains them onto the per-session bus before + // `spawnOrAttach` returns, so they live in the replay ring with + // ids — but the SDK's old default of `lastEventId: undefined` + // started subscriptions live, so consumers never observed them. + // + // Fix: when `session.attached === false` (newly-created), seed + // `Last-Event-ID: 0` to replay the startup-window events. The + // existing `modelServiceId` carve-out still triggers seed for + // re-attached sessions where attach-time switch events need to + // replay. + const { fetch, calls } = recordingFetch((req) => { + if (req.url.endsWith('/session')) { + return jsonResponse(200, { + sessionId: 's-1', + workspaceCwd: '/work/a', + attached: false, + }); + } + if (req.url.endsWith('/session/s-1/events')) { + return sseResponse(''); + } + return jsonResponse(500, { error: `unexpected ${req.url}` }); + }); + const client = new DaemonClient({ baseUrl: 'http://daemon', fetch }); + + const session = await DaemonSessionClient.createOrAttach(client, { + workspaceCwd: '/work/a', + // No `modelServiceId` — the only signal that triggered seed + // pre-fix. With the fix, `attached: false` alone is enough. + }); + + for await (const _event of session.events()) { + /* empty */ + } + + expect(session.attached).toBe(false); + expect(calls[1]?.url).toBe('http://daemon/session/s-1/events'); + expect(calls[1]?.headers['last-event-id']).toBe('0'); + }); + it('starts live when createOrAttach has no model service replay need', async () => { const { fetch, calls } = recordingFetch((req) => { if (req.url.endsWith('/session')) { diff --git a/packages/sdk-typescript/test/unit/daemonEvents.test.ts b/packages/sdk-typescript/test/unit/daemonEvents.test.ts index c5149347a55..779a9f2a672 100644 --- a/packages/sdk-typescript/test/unit/daemonEvents.test.ts +++ b/packages/sdk-typescript/test/unit/daemonEvents.test.ts @@ -661,12 +661,28 @@ describe('daemon event schema', () => { }); it('recognizes slow_client_warning frames as known events', () => { + // PR 14b fix (codex round 8 — sibling consistency): `satisfies + // DaemonEvent` keeps `v: 1` / `type: 'slow_client_warning'` + // narrow rather than widening to `number` / `string`. The same + // pattern was applied to PR 14b's own fixtures in round 3 + // (`mcp_budget_warning` + `mcp_child_refused_batch`); this is the + // closest sibling fixture in the same describe block, so + // matching it here keeps the sdk-test typing style coherent. + // + // Note: a tsconfig audit found ~17 OTHER fixtures in this file + // with the same widening shape (PR 4 / PR 10 / PR 11 era). They + // remain unfixed because (a) they're outside PR 14b's scope, and + // (b) the sdk package's `tsconfig.json` excludes the test + // directory from `tsc --noEmit`, so none of them block CI today. + // A future PR that opts tests into the typecheck scope can fix + // all of them at once. Round 3 only signed up for PR 14b's own + // fixtures. const warning = { // No `id` on synthetic frames (matches the daemon's emit shape). v: 1, type: 'slow_client_warning', data: { queueSize: 192, maxQueued: 256, lastEventId: 42 }, - }; + } satisfies DaemonEvent; const known = asKnownDaemonEvent(warning); expect(known?.type).toBe('slow_client_warning'); @@ -748,6 +764,336 @@ describe('daemon event schema', () => { expect(state.lastEventId).toBe(1); }); + // PR 14b: MCP guardrail push events. Mirrors the slow_client_warning + // test patterns (predicate validation + reducer state) — the two + // event types are siblings on the per-session SSE bus and use the + // same KnownDaemonEvent narrowing. + it('recognizes mcp_budget_warning frames as known events', () => { + // PR 14b fix (codex round 3): `satisfies DaemonEvent` keeps the + // discriminator literals (`v: 1`, `type: 'mcp_budget_warning'`) + // narrow without widening to `number`/`string`. Required so the + // fixture passes through `asKnownDaemonEvent`'s `event.type` + // switch under strict typecheck. The sdk package's tsconfig + // currently scopes `tsc --noEmit` to `src/**/*.ts` only — tests + // aren't gated yet — but the fixture stays type-safe for when + // they are. + const warning = { + id: 7, + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + } satisfies DaemonEvent; + const known = asKnownDaemonEvent(warning); + expect(known?.type).toBe('mcp_budget_warning'); + + // Schema: required numeric fields, exact-literal `thresholdRatio`, + // and `mode` constrained to `'warn' | 'enforce'`. Bad shapes are + // rejected so the reducer routes them through the + // `unrecognizedKnownEventCount` branch. + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_budget_warning', + data: { + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + }), + ).toBeUndefined(); + // PR 14b fix (codex round 6): `thresholdRatio` is validated as a + // finite number rather than the literal 0.75 — the SDK's role is + // wire-shape validation, not threshold-value enforcement. Pinning + // the literal would mean a daemon-side bump to e.g. 0.80 silently + // routes every warning through `unrecognizedKnownEventCount` (a + // cross-package coordination hazard). Forward-compat for a future + // 0.5 critical threshold falls out for free; the daemon constant + // and protocol docs are the source of truth for threshold values. + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.5, // forward-compat threshold value + mode: 'warn', + }, + }), + ).toBeDefined(); + // Non-finite values (NaN / Infinity) are still rejected — the + // predicate uses `isFiniteNumber`, not bare `typeof === 'number'`. + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: Number.NaN, + mode: 'warn', + }, + }), + ).toBeUndefined(); + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'off', // off-mode never fires the warning — bad payload. + }, + }), + ).toBeUndefined(); + }); + + it('reduces mcp_budget_warning into the view state without ending the stream', () => { + const state = reduceDaemonSessionEvents([ + { + id: 1, + v: 1, + type: 'session_update', + data: { sessionId: 's-1', phase: 'prompting' }, + }, + { + id: 2, + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 3, + reservedCount: 3, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + }, + { + id: 3, + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'enforce', + }, + }, + ]); + + expect(state.mcpBudgetWarningCount).toBe(2); + expect(state.lastMcpBudgetWarning).toEqual({ + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'enforce', + }); + // Non-terminal — stream stays alive. + expect(state.alive).toBe(true); + expect(state.terminalEvent).toBeUndefined(); + expect(state.lastEventId).toBe(3); + }); + + it('recognizes mcp_child_refused_batch frames as known events', () => { + // PR 14b fix (codex round 3): `satisfies DaemonEvent` preserves + // the literal discriminator (`v: 1`, `type: + // 'mcp_child_refused_batch'`) — see sibling fixture above for + // the full rationale. + const batch = { + id: 9, + v: 1, + type: 'mcp_child_refused_batch', + data: { + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + { name: 'c', transport: 'http', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }, + } satisfies DaemonEvent; + const known = asKnownDaemonEvent(batch); + expect(known?.type).toBe('mcp_child_refused_batch'); + + // `mode: 'warn'` must be rejected — warn mode never refuses, so a + // refused-batch tagged with warn is protocol garbage. The + // reducer's safety net (`unrecognizedKnownEventCount`) catches it + // instead of letting the `last*` field hold a malformed shape. + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_child_refused_batch', + data: { + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'warn', + }, + }), + ).toBeUndefined(); + + // Unknown transport family rejected (forward-compat: a future + // daemon emitting a new transport speaks a newer wire than this + // SDK release). + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_child_refused_batch', + data: { + refusedServers: [ + { name: 'b', transport: 'quic', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }, + }), + ).toBeUndefined(); + + // Bad reason rejected — only `'budget_exhausted'` is valid in + // PR 14b. Future causes extend the literal set. + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_child_refused_batch', + data: { + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'something_else' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }, + }), + ).toBeUndefined(); + + // Empty `refusedServers` is structurally valid (the daemon would + // never emit an empty batch — `emitRefusedBatchIfAny` is gated on + // `lastRefusedServerNames.length > 0` — but the SDK predicate + // doesn't enforce that invariant; it's a daemon-side correctness + // property, not a wire-format requirement). Verify the predicate + // accepts it so a future daemon contract change doesn't break + // adapters. + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_child_refused_batch', + data: { + refusedServers: [], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }, + }), + ).toBeDefined(); + }); + + it('reduces mcp_child_refused_batch into the view state without ending the stream', () => { + const state = reduceDaemonSessionEvents([ + { + id: 1, + v: 1, + type: 'session_update', + data: { sessionId: 's-1', phase: 'prompting' }, + }, + { + id: 2, + v: 1, + type: 'mcp_child_refused_batch', + data: { + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }, + }, + // Length-1 batch from `readResource` lazy-spawn refusal + // arrives next. + { + id: 3, + v: 1, + type: 'mcp_child_refused_batch', + data: { + refusedServers: [ + { name: 'c', transport: 'http', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }, + }, + ]); + + expect(state.mcpChildRefusedBatchCount).toBe(2); + expect(state.lastMcpChildRefusedBatch).toEqual({ + refusedServers: [ + { name: 'c', transport: 'http', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'enforce', + }); + expect(state.alive).toBe(true); + expect(state.terminalEvent).toBeUndefined(); + expect(state.lastEventId).toBe(3); + }); + + it('rejected MCP guardrail payloads route through unrecognizedKnownEventCount', () => { + // The reducer's safety net for "type matches a known type but + // schema fails": increments `unrecognizedKnownEventCount` and + // captures the raw event in `lastUnrecognizedKnownEvent`. Mirrors + // the slow_client_warning sibling pattern. + const state = reduceDaemonSessionEvent(reduceDaemonSessionEvents([]), { + id: 1, + v: 1, + type: 'mcp_child_refused_batch', + data: { + // `mode: 'warn'` is invalid (warn never refuses) — predicate + // rejects, reducer routes through the unrecognized branch. + refusedServers: [ + { name: 'b', transport: 'stdio', reason: 'budget_exhausted' }, + ], + budget: 1, + liveCount: 1, + reservedCount: 1, + mode: 'warn', + }, + }); + expect(state.unrecognizedKnownEventCount).toBe(1); + expect(state.lastUnrecognizedKnownEvent?.type).toBe( + 'mcp_child_refused_batch', + ); + // Refused-batch counter NOT incremented — the malformed payload + // didn't reach the typed reducer arm. + expect(state.mcpChildRefusedBatchCount).toBe(0); + expect(state.lastMcpChildRefusedBatch).toBeUndefined(); + }); it('narrows memory_changed events and rejects malformed payloads', () => { const valid: DaemonEvent = { id: 7,