From 22ac97bcc2404ac8cef583aab7710f26b67b9ebc Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Mon, 18 May 2026 15:13:55 +0800 Subject: [PATCH 1/9] feat(serve): MCP guardrail push events + hysteresis (#4175 Wave 3 PR 14b) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds typed SSE push events for the MCP budget surface introduced by PR 14 (#4247). Operators dashboarding `/workspace/mcp` snapshot already see budget pressure via `budgets[0].status` + per-server `disabledReason: 'budget'`; PR 14b layers a real-time push channel on top so SDK clients can react without polling. Two new event types 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. Fires under both `warn` and `enforce` modes. - `mcp_child_refused_batch` — fires at end of each `discoverAllMcpTools*` pass when ≥1 server was refused, AND as a length-1 batch on the `readResource` lazy-spawn refusal path. `mode` is the literal `'enforce'` since `warn` mode never refuses. Wired end-to-end: - Core `McpClientManager` gains `evaluateBudgetState()` + `emitRefusedBatchIfAny()`, a `pendingRefusalNames` queue distinct from the snapshot-visible `lastRefusedServerNames` (PR 14 contract: refusals survive between passes for the snapshot — the emit log is independent), and a `setOnBudgetEvent(cb)` setter so acpAgent can register the callback after `loadCliConfig` returns but before discovery fires. - ACP child→bridge transport: `connection.extNotification` with method `qwen/notify/session/mcp-budget-event` (mirrors the `authenticate/update` precedent at `acpAgent.ts:355`). Payload carries `v: 1` envelope + sessionId for routing. - BridgeClient gains `extNotification` handler; resolves session via `byId.get(sessionId)` and republishes as a session-scoped SSE frame. Unknown methods, unknown kinds, and missing sessionIds drop silently for forward-compat. - Capability tag `mcp_guardrail_events` (always-on, distinct from PR 14's `mcp_guardrails` snapshot tag). - SDK typed-event registry: `KnownDaemonEvent` gains `DaemonMcpGuardrailEvent` union member; reducer adds 4 fields (`mcpBudgetWarningCount`, `lastMcpBudgetWarning`, `mcpRefusedBatchCount`, `lastMcpRefusedBatch`) on `DaemonSessionViewState`. Predicate guards reject `mode: 'warn'` on refused-batch (literal-`'enforce'` invariant) + unknown transport families (forward-compat: a daemon emitting a new transport speaks newer wire than this SDK). Tests: - 11 new manager tests (state machine arms/fires/rearms, refused batch coalescing, transport preservation, off-mode no-op, `readResource` length-1 batch, stop-resets-state). - 2 new acpAgent tests (extNotification dispatch with sessionId closure; defensive no-op when `getMcpClientManager` absent). - 3 new bridge tests (warning + refused-batch publish; drop unknown methods / kinds / sessionIds). - 5 new SDK tests (predicate accepts good shapes, rejects `thresholdRatio !== 0.75`, rejects `mode: 'warn'` on refused-batch, rejects unknown transport; reducer counters increment + last* capture; safety-net routes malformed payloads through `unrecognizedKnownEventCount`). - 1 new integration test: `pgrep -P` × `subprocessCount` cross-check validates the in-process counter (PR 14b's event source) against external observation. Skip-gated like PR 1 (POSIX, non-sandbox). - 1 new capability shape test for `mcp_guardrail_events`. - 1 added EXPECTED_STAGE1_FEATURES entry. Backward compat: every new field optional; reducer state initialized to zero defaults; old daemons advertise `mcp_guardrails` only and SDK consumers fall back to snapshot polling per the existing PR 14 contract. `EVENT_SCHEMA_VERSION` stays at `1` — push events are an additive, narrowable extension, not a wire bump. Verified: 865/865 tests pass across 30 files (manager, serve, acp-integration, sdk-typescript); typecheck clean across 4 workspaces; lint clean on 11 touched files. Refs #4175. --- docs/developers/qwen-serve-protocol.md | 17 +- docs/users/qwen-serve.md | 2 + .../cli/qwen-serve-baseline.test.ts | 61 +++ .../cli/src/acp-integration/acpAgent.test.ts | 144 ++++++ packages/cli/src/acp-integration/acpAgent.ts | 66 +++ packages/cli/src/serve/capabilities.ts | 12 + packages/cli/src/serve/httpAcpBridge.test.ts | 204 ++++++++ packages/cli/src/serve/httpAcpBridge.ts | 45 ++ packages/cli/src/serve/server.test.ts | 16 + .../core/src/tools/mcp-client-manager.test.ts | 453 ++++++++++++++++++ packages/core/src/tools/mcp-client-manager.ts | 332 ++++++++++++- packages/sdk-typescript/src/daemon/events.ts | 189 +++++++- .../test/unit/daemonEvents.test.ts | 299 ++++++++++++ 13 files changed, 1824 insertions(+), 16 deletions(-) diff --git a/docs/developers/qwen-serve-protocol.md b/docs/developers/qwen-serve-protocol.md index 9e85688ba58..0924830db66 100644 --- a/docs/developers/qwen-serve-protocol.md +++ b/docs/developers/qwen-serve-protocol.md @@ -99,7 +99,8 @@ registry. Clients **must** gate UI off `features`, not off `mode` (per design 'session_set_model', 'client_identity', 'client_heartbeat', 'session_permission_vote', 'permission_vote', 'workspace_mcp', 'workspace_skills', 'workspace_providers', 'session_context', 'session_supported_commands', - 'session_close', 'session_metadata', 'mcp_guardrails'] + 'session_close', 'session_metadata', 'mcp_guardrails', + 'mcp_guardrail_events'] ``` `session_scope_override` is the negotiation handle for the per-request `sessionScope` field on `POST /session` (see below). Older daemons silently ignore the field, so SDK clients should pre-flight `caps.features` for this tag before sending it. @@ -126,6 +127,13 @@ registry. Clients **must** gate UI off `features`, not off `mode` (per design `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`, `mcpRefusedBatchCount`, `lastMcpRefusedBatch` for adapters that want simple lag-style UI. + ## Routes ### `GET /health` @@ -317,13 +325,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 31b3a609b98..c7759e2594f 100644 --- a/docs/users/qwen-serve.md +++ b/docs/users/qwen-serve.md @@ -194,6 +194,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..a0704eda19c 100644 --- a/integration-tests/cli/qwen-serve-baseline.test.ts +++ b/integration-tests/cli/qwen-serve-baseline.test.ts @@ -462,6 +462,67 @@ async function measureRssAtSessionCount(sessionCount: number): Promise<{ fs.rmSync(ws, { recursive: true, force: true }); } }, 120_000); + + // PR 14b cross-check: validate the in-process counter against + // external `pgrep -P` measurement. The push-event channel + // (`mcp_budget_warning` / `mcp_child_refused_batch`) reads + // `getMcpClientAccounting().total` for `liveCount` and the + // `subprocessCount` field is the same arithmetic + // (`stdio + websocket`). If the daemon's in-process counter + // diverges from what `pgrep -P` actually observes, the events + // would lie. Skip-gated like the parent describe (POSIX, non- + // sandbox); idle MCP fixtures are stdio-only so + // `subprocessCount` should equal `mcpGrandchildren.length` + // exactly (no amplification slack required). + it('in-process subprocessCount 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 → + // `subprocessCount === 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` (CONNECTED clients) is allowed to be + // lower than `subprocessCount` if the OS still sees a + // process the daemon already considers disconnected + // (rare race, kept as `<=`); but for fresh-spawn idle + // fixtures we expect equality. + expect(snapshot.clientCount).toBe(MCP_SERVERS_CONFIGURED); + expect(observed.mcpGrandchildren.length).toBe(MCP_SERVERS_CONFIGURED); + // `clientCount` is the snapshot's authoritative live count. + // Validating it against pgrep closes the loop on PR 14b's + // event-source assumption. + 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..530975d2fb8 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -2095,6 +2095,150 @@ describe('QwenAgent MCP SSE/HTTP support', () => { mockConnectionState.resolve(); await agentPromise; }); + + // PR 14b: budget-event push channel. The manager's `setOnBudgetEvent` + // callback is wired by `newSessionConfig` AFTER `loadCliConfig` returns + // so the very first discovery pass's events are still in the + // registration window. The callback translates each event into a + // `connection.extNotification(qwen/notify/session/mcp-budget-event, ...)` + // payload that downstream BridgeClient turns into an SSE frame. + it('newSession wires McpClientManager.setOnBudgetEvent → extNotification with sessionId', async () => { + const sessionId = 'session-budget-events'; + const innerConfig = await setupSessionMocks(sessionId); + // Stub `getToolRegistry().getMcpClientManager()` and capture the + // callback registered via `setOnBudgetEvent`. The fake manager + // doesn't care about the callback shape — the test invokes it + // synchronously with a hand-built event after registration. + let capturedCallback: + | ((event: Record) => void) + | undefined; + const fakeManager = { + setOnBudgetEvent: vi.fn( + (cb: (event: Record) => void) => { + capturedCallback = cb; + }, + ), + }; + (innerConfig as unknown as Record)['getToolRegistry'] = vi + .fn() + .mockReturnValue({ + getAllTools: () => [], + getMcpClientManager: () => fakeManager, + }); + + 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: [] }); + + // Manager's setOnBudgetEvent must have been called once with a + // function (the closure capturing the sessionId). + expect(fakeManager.setOnBudgetEvent).toHaveBeenCalledTimes(1); + 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 getMcpClientManager is absent (defensive)', async () => { + // Older / stubbed `ToolRegistry` shapes may not expose + // `getMcpClientManager` (the bulk of older test fixtures stub + // ToolRegistry as `{ getAllTools: () => [] }`). The PR 14b code + // path uses optional chaining so the absence is silent — no + // throw, no extNotification call. + const innerConfig = await setupSessionMocks('session-no-mgr'); + (innerConfig as unknown as Record)['getToolRegistry'] = vi + .fn() + .mockReturnValue({ getAllTools: () => [] }); + + 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 manager → no wiring → no extNotification fires. + 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 4b70ed88546..16a8a1bf392 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -13,6 +13,7 @@ import { QwenOAuth2Event, qwenOAuth2Events, MCP_BUDGET_WARN_FRACTION, + type McpBudgetEvent, MCPServerConfig, SessionService, SESSION_TITLE_MAX_LENGTH, @@ -1658,6 +1659,71 @@ class QwenAgent implements Agent { }, ); await config.initialize(); + // PR 14b: wire the manager's budget-event callback to push the + // events out as ACP `extNotification` frames. Done AFTER + // `initialize()` (which constructs the tool registry / manager) + // and BEFORE `waitForMcpReady()` (which awaits the still-in- + // flight first discovery pass) so the registration completes + // before discovery finishes and emits its end-of-pass events. + // + // Best-effort: the manager is `off` mode (and `setOnBudgetEvent` + // is a no-op) when no budget is configured — production cost is + // a single property read. `getMcpClientManager()` may be absent + // on stubbed test ToolRegistries, hence the optional chain. + // + // sessionId capture: this Config is per-session (see + // `newSessionConfig` JSDoc + #4175 PR 14 R4 scope correction), + // so the sessionId stays stable for the manager's lifetime. + // Source: `config.getSessionId()` — Core's Config auto-assigns a + // randomUUID at construction when no `sessionId` is passed in + // (`config.ts:849`), so the value is always present after + // `loadCliConfig` returns. Defensive optional chain keeps this + // safe if a stub Config in tests omits the method. + // Best-effort lookup: older / stubbed test ToolRegistry shapes may + // omit `getMcpClientManager`, so a `typeof` check is required in + // addition to the optional chain (the optional chain only protects + // against the registry itself being missing). + const toolRegistry = config.getToolRegistry?.() as + | { getMcpClientManager?: () => unknown } + | undefined; + const budgetManager = + typeof toolRegistry?.getMcpClientManager === 'function' + ? (toolRegistry.getMcpClientManager() as + | { + setOnBudgetEvent?: ( + cb: (event: McpBudgetEvent) => void, + ) => void; + } + | undefined) + : undefined; + const wiredSessionId = + typeof config.getSessionId === 'function' + ? config.getSessionId() + : undefined; + if ( + budgetManager && + typeof budgetManager.setOnBudgetEvent === 'function' && + wiredSessionId !== undefined + ) { + const sid = wiredSessionId; + budgetManager.setOnBudgetEvent!((event) => { + // Fire-and-forget: extNotification returns Promise but + // the manager's call site doesn't await it. The .catch + // suppresses unhandled rejections — a mid-flight ACP + // disconnect would otherwise crash the child via uncaught + // rejection. Snapshot still carries the state for clients + // that reconnect. + void this.connection + .extNotification('qwen/notify/session/mcp-budget-event', { + v: 1, + sessionId: sid, + ...event, + }) + .catch(() => { + // ACP channel closed or peer disconnected — drop event. + }); + }); + } // Same reasoning as the top-level runAcpAgent path: ACP feeds session // messages to the model immediately, so we cannot return a Config whose // MCP discovery is still in flight. diff --git a/packages/cli/src/serve/capabilities.ts b/packages/cli/src/serve/capabilities.ts index 5948a71a186..b900ddeddd1 100644 --- a/packages/cli/src/serve/capabilities.ts +++ b/packages/cli/src/serve/capabilities.ts @@ -85,6 +85,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`, + // `mcpRefusedBatchCount`, `lastMcpRefusedBatch`). 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 15. Daemon was booted with `--require-auth` (or // `requireAuth: true`), so even loopback callers must carry a bearer // token. Advertised CONDITIONALLY — only when the flag is on — so diff --git a/packages/cli/src/serve/httpAcpBridge.test.ts b/packages/cli/src/serve/httpAcpBridge.test.ts index 80a191f23ff..8c47a44c72c 100644 --- a/packages/cli/src/serve/httpAcpBridge.test.ts +++ b/packages/cli/src/serve/httpAcpBridge.test.ts @@ -4369,6 +4369,210 @@ 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 — the real one. The 4 dropped + // notifications produced no SSE frames. + expect(collected).toEqual([{ type: 'mcp_budget_warning' }]); + + abort.abort(); + 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 9190e8875ae..d2cf0120c2d 100644 --- a/packages/cli/src/serve/httpAcpBridge.ts +++ b/packages/cli/src/serve/httpAcpBridge.ts @@ -1226,6 +1226,51 @@ class BridgeClient implements Client { }); } + /** + * 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-or-unresolvable 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). + */ + 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 entry = this.resolveEntry(sessionId); + if (!entry) 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; + entry.events.publish({ + type, + data: rest, + ...(entry.activePromptOriginatorClientId + ? { originatorClientId: entry.activePromptOriginatorClientId } + : {}), + }); + } + async writeTextFile( params: WriteTextFileRequest, ): Promise { diff --git a/packages/cli/src/serve/server.test.ts b/packages/cli/src/serve/server.test.ts index 4f82c4ea69f..f7a3ea3b327 100644 --- a/packages/cli/src/serve/server.test.ts +++ b/packages/cli/src/serve/server.test.ts @@ -107,6 +107,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', ] as const; // Issue #4175 PR 15. `require_auth` is registered but conditionally @@ -666,6 +670,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/tools/mcp-client-manager.test.ts b/packages/core/src/tools/mcp-client-manager.test.ts index 9c728f91f57..a7bac31f214 100644 --- a/packages/core/src/tools/mcp-client-manager.test.ts +++ b/packages/core/src/tools/mcp-client-manager.test.ts @@ -2028,3 +2028,456 @@ 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); + expect(warnings[0]).toMatchObject({ + kind: 'budget_warning', + reservedCount: 4, + 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); + }); +}); diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts index b6aefa4fb26..a669cb6cd76 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,52 @@ 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. Same lifetime as + * `lastRefusedServerNames` — reset per pass, cleared on emit. + */ + 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: 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 +480,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; } /** @@ -449,6 +563,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; + } + /** * PR 14 fix (review #4247 wenshao R7 line 464): drop a server's * entry from the per-pass refusal log, if present. The @@ -463,6 +596,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); } /** @@ -480,10 +623,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`, @@ -516,6 +676,101 @@ 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; + this.onBudgetEvent?.({ + 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: 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 no-refusals: an empty `lastRefusedServerNames` + * short-circuits without firing or clearing. Clears both the names + * list and the transport map after firing so the next pass starts + * fresh. + * + * `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 { + 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). + 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. + this.pendingRefusalNames.clear(); + return; + } + const refusedServers: McpRefusedServer[] = namesInOrder.map((name) => ({ + name, + transport: this.lastRefusedTransports.get(name) ?? 'unknown', + reason: 'budget_exhausted' as const, + })); + this.onBudgetEvent?.({ + kind: 'refused_batch', + refusedServers, + budget: this.clientBudget, + liveCount: this.getMcpClientAccounting().total, + reservedCount: this.reservedSlots.size, + mode: 'enforce', + }); + this.pendingRefusalNames.clear(); + } + /** * Initiates the tool discovery process for all configured MCP servers. * It connects to each server, discovers its available tools, and registers @@ -538,6 +793,11 @@ export class McpClientManager { // (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( @@ -554,7 +814,7 @@ export class McpClientManager { // reflects the configured set. `off` is a no-op. const reservation = this.tryReserveSlot(name); if (reservation === 'refused') { - this.refuseAndLog(name); + this.refuseAndLog(name, config); return; } @@ -636,6 +896,14 @@ export class McpClientManager { await Promise.all(discoveryPromises); this.discoveryState = MCPDiscoveryState.COMPLETED; this.emitBudgetTelemetry(Object.keys(servers).length); + // PR 14b: end-of-pass push events. Order is intentional — + // `refused_batch` first so SDK consumers see refusals before + // any warning the same pass might have crossed (the warning + // can fire on a high reservedCount even with all-success + // connects). Both calls are no-ops when there's nothing to emit + // / `off` mode / no callback registered. + this.emitRefusedBatchIfAny(); + this.evaluateBudgetState(); } /** @@ -712,7 +980,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 @@ -794,6 +1068,11 @@ export class McpClientManager { // snapshots immediately reflect reality. Mirrors the same // pattern in `readResource`'s late-reserve branch. this.dropRefusalEntry(serverName); + // PR 14b: a successful per-server (re)discover may push the + // ratio past 75% — e.g. operator did `/mcp reconnect` on the + // last server filling the budget. The bulk-pass + // `evaluateBudgetState` won't run here, so fire it inline. + this.evaluateBudgetState(); } catch (error) { // PR 14 fix (review #4247 wenshao R3 line 546): two-mode // cleanup for connect failure, matching the `readResource` @@ -885,6 +1164,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; } /** @@ -1096,6 +1382,10 @@ export class McpClientManager { // 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, @@ -1250,6 +1540,16 @@ export class McpClientManager { incremental: true, }); this.emitBudgetTelemetry(Object.keys(servers).length); + // PR 14b: end-of-pass push events. Mirrors `discoverAllMcpTools` + // — refused_batch first, then warning. The single-server + // `discoverMcpToolsForServerInternal` calls inside this pass + // already evaluate state on success/refusal individually, but the + // bulk-shape contract is "one batch per pass" so we coalesce + // anything missed by the single-server hooks (e.g. multiple + // discoveries refused but the bulk pass also produced its own + // refused entries via `removeServer` hitting cap edge cases). + this.emitRefusedBatchIfAny(); + this.evaluateBudgetState(); // 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 @@ -1512,7 +1812,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, @@ -1537,6 +1842,11 @@ export class McpClientManager { // the next snapshot reflects the late-reservation success. if (weReservedSlot) { this.dropRefusalEntry(serverName); + // PR 14b: a successful late reservation (lazy-spawn against a + // freshly-freed slot) may push the ratio past 75%. Run the + // state machine inline since this path bypasses the bulk + // `evaluateBudgetState` at end-of-pass. + this.evaluateBudgetState(); } const sdkCallback = isSdkMcpServerConfig(serverConfig) diff --git a/packages/sdk-typescript/src/daemon/events.ts b/packages/sdk-typescript/src/daemon/events.ts index b81bf064d46..0b6c78c5721 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', ] as const; const DAEMON_KNOWN_EVENT_TYPES: ReadonlySet = new Set( @@ -125,6 +136,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; +} + export type DaemonSessionUpdateEvent = DaemonEventEnvelope< 'session_update', DaemonSessionUpdateData @@ -173,6 +237,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 DaemonSessionEvent = | DaemonSessionUpdateEvent @@ -192,10 +264,22 @@ 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; + export type KnownDaemonEvent = | DaemonSessionEvent | DaemonControlEvent - | DaemonStreamLifecycleEvent; + | DaemonStreamLifecycleEvent + | DaemonMcpGuardrailEvent; export interface DaemonSessionViewState { lastEventId?: number; @@ -231,6 +315,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. + */ + mcpRefusedBatchCount: number; + lastMcpRefusedBatch?: DaemonMcpChildRefusedBatchData; } export function createDaemonSessionViewState( @@ -257,6 +360,10 @@ export function createDaemonSessionViewState( seed.lastUnmatchedPermissionResolutionId, slowClientWarningCount: seed.slowClientWarningCount ?? 0, lastSlowClientWarning: seed.lastSlowClientWarning, + mcpBudgetWarningCount: seed.mcpBudgetWarningCount ?? 0, + lastMcpBudgetWarning: seed.lastMcpBudgetWarning, + mcpRefusedBatchCount: seed.mcpRefusedBatchCount ?? 0, + lastMcpRefusedBatch: seed.lastMcpRefusedBatch, }; } @@ -326,6 +433,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; default: return undefined; } @@ -462,6 +577,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, + mcpRefusedBatchCount: base.mcpRefusedBatchCount + 1, + lastMcpRefusedBatch: event.data, + }; default: { const _exhaustive: never = event; return _exhaustive; @@ -619,6 +752,58 @@ function isStreamErrorData(value: unknown): value is DaemonStreamErrorData { return isRecord(value) && isNonEmptyString(value['error']); } +function isMcpBudgetWarningData( + value: unknown, +): value is DaemonMcpBudgetWarningData { + return ( + isRecord(value) && + isFiniteNumber(value['liveCount']) && + isFiniteNumber(value['reservedCount']) && + isFiniteNumber(value['budget']) && + value['thresholdRatio'] === 0.75 && + (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 isPermissionOption(value: unknown): value is DaemonPermissionOption { return isRecord(value) && isNonEmptyString(value['optionId']); } diff --git a/packages/sdk-typescript/test/unit/daemonEvents.test.ts b/packages/sdk-typescript/test/unit/daemonEvents.test.ts index e58aa6ff745..0a0d6ec736b 100644 --- a/packages/sdk-typescript/test/unit/daemonEvents.test.ts +++ b/packages/sdk-typescript/test/unit/daemonEvents.test.ts @@ -744,4 +744,303 @@ describe('daemon event schema', () => { // id observed (the original session_update at id=1). 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', () => { + const warning = { + id: 7, + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 4, + reservedCount: 4, + budget: 4, + thresholdRatio: 0.75, + mode: 'warn', + }, + }; + 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(); + expect( + asKnownDaemonEvent({ + v: 1, + type: 'mcp_budget_warning', + data: { + liveCount: 4, + reservedCount: 4, + budget: 4, + // `thresholdRatio` must be the literal 0.75 (the only ratio + // PR 14b emits). A future PR adding 0.50 / 0.95 thresholds + // would extend both the daemon emit + this predicate. + thresholdRatio: 0.5, + 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', () => { + 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', + }, + }; + 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.mcpRefusedBatchCount).toBe(2); + expect(state.lastMcpRefusedBatch).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.mcpRefusedBatchCount).toBe(0); + expect(state.lastMcpRefusedBatch).toBeUndefined(); + }); }); From 195d9000b06028fc52a3cfe0f02e283ede44fb77 Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Mon, 18 May 2026 16:58:30 +0800 Subject: [PATCH 2/9] fixup(serve): address PR 14b review (codex P2 findings 1-4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex review round 1 of #4271 flagged four P2 issues; this fixup commit addresses all of them with bounded changes: #1 Bridge early-event buffer (httpAcpBridge.ts + `BridgeClient.earlyEvents` Map / `bufferEarlyEvent` / `sweepExpiredEarlyEvents` / `drainEarlyEvents`): Budget events fired during the child's `newSession` handler reach `BridgeClient.extNotification` before `byId.set` populates the session entry — pre-fix those frames hit `resolveEntry → undefined` and got dropped, so SSE subscribers never saw the events that crossed during session creation. Now we buffer per-sessionId with triple bounds (64 sessions × 32 frames × 60s TTL ≈ 400 KB worst case) and drain on `createSessionEntry`'s `byId.set`. Lazy TTL sweep + bounded capacity defends against malicious / buggy children spamming bogus sessionIds. #2 Pre-init callback registration (config.ts new `setMcpBudgetEventCallback` shim + acpAgent moves registration): Pre-fix acpAgent registered `setOnBudgetEvent` AFTER `config.initialize()` returned. In `QWEN_CODE_LEGACY_MCP_BLOCKING=1` mode discovery completes synchronously inside `initialize`, so end-of-pass events for the first pass were lost 100%; in default progressive mode there was a real race window. The Config-level shim stashes the callback and applies it inside `createToolRegistry` right after manager construction, so the manager has its callback wired BEFORE either `discoverAllTools` (legacy) or `startMcpDiscoveryInBackground` (default) fires. #3 Bulk-pass refused-batch coalescing (mcp-client-manager.ts new `bulkPassDepth` counter + `emitRefusedBatchIfAny` early-return guard): Pre-fix `discoverAllMcpToolsIncremental` walked N new servers and the per-server `discoverMcpToolsForServerInternal` refusal path emitted N length-1 batches inline, breaking the documented "one batch per `discoverAllMcpTools*` pass" contract. Now the bulk entry/exit increments/decrements `bulkPassDepth` in try/finally; while > 0, `emitRefusedBatchIfAny` no-ops. The bulk-pass `finally` decrements to 0 BEFORE its own emit, so the terminal call drains the queue once as a coalesced length-N batch. #4 Hysteresis re-arm on slot release (mcp-client-manager.ts new `releaseSlotName` helper + inline `evaluateBudgetState` in `tryReserveSlot`): Pre-fix the state machine only ran on bulk-pass end-of-pass + per-server success, so slot-release paths (`disconnectServer`, `removeServer`, `runWithDiscoveryTimeout`, connect-failure catch blocks) deleted from `reservedSlots` without touching `warnArmed`. Operator scenario "4/4 fire → drop to 1/4 → up to 4/4" never fired the second warning. Now every slot mutation drives the state machine via a single helper pair — `tryReserveSlot` calls evaluate after `add`, `releaseSlotName` wraps `delete + evaluate`. All 6 `reservedSlots.delete` sites migrated; redundant standalone `evaluateBudgetState` calls dropped from end-of-pass + per-server success + readResource late-reserve paths (the helper covers them). Tests: - `mcp-client-manager.test.ts`: +2 new tests (`discoverAllMcpToolsIncremental coalesces multi-server refusals into ONE batch`, `disconnectServer drives the hysteresis re-arm path`); 1 existing test's `reservedCount` assertion updated from 4→3 to reflect inline-at-crossing semantics. - `acpAgent.test.ts`: pre-init wiring test rewritten to assert the strict ordering invariant `setMcpBudgetEventCallback` → `initialize` via `callOrder` array; defensive test now exercises the new `Config.setMcpBudgetEventCallback` absence path. - `httpAcpBridge.test.ts`: +1 new test (`buffers events for a not-yet-registered sessionId, drains them on registration`) exercising the buffer + drain mechanism via thread-scope pre-buffering for a future session id. Verified: 868/868 tests pass across 30 files (manager 71 incl. 14 PR 14b + 2 fixup; serve 137 incl. 4 PR 14b + 1 fixup; acpAgent 55 incl. 2 PR 14b updated; sdk 27 PR 14b unchanged); typecheck clean across 4 workspaces; lint clean on 7 touched files. --- .../cli/src/acp-integration/acpAgent.test.ts | 81 +-- packages/cli/src/acp-integration/acpAgent.ts | 72 +- packages/cli/src/serve/httpAcpBridge.test.ts | 105 ++- packages/cli/src/serve/httpAcpBridge.ts | 125 +++- packages/core/src/config/config.ts | 50 ++ .../core/src/tools/mcp-client-manager.test.ts | 104 ++- packages/core/src/tools/mcp-client-manager.ts | 635 ++++++++++-------- 7 files changed, 795 insertions(+), 377 deletions(-) diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index 530975d2fb8..59596c3d530 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -2096,35 +2096,38 @@ describe('QwenAgent MCP SSE/HTTP support', () => { await agentPromise; }); - // PR 14b: budget-event push channel. The manager's `setOnBudgetEvent` - // callback is wired by `newSessionConfig` AFTER `loadCliConfig` returns - // so the very first discovery pass's events are still in the - // registration window. The callback translates each event into a - // `connection.extNotification(qwen/notify/session/mcp-budget-event, ...)` - // payload that downstream BridgeClient turns into an SSE frame. - it('newSession wires McpClientManager.setOnBudgetEvent → extNotification with sessionId', async () => { + // 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 `getToolRegistry().getMcpClientManager()` and capture the - // callback registered via `setOnBudgetEvent`. The fake manager - // doesn't care about the callback shape — the test invokes it - // synchronously with a hand-built event after registration. + // 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 fakeManager = { - setOnBudgetEvent: vi.fn( - (cb: (event: Record) => void) => { - capturedCallback = cb; - }, - ), - }; - (innerConfig as unknown as Record)['getToolRegistry'] = vi - .fn() - .mockReturnValue({ - getAllTools: () => [], - getMcpClientManager: () => fakeManager, - }); + 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, @@ -2149,9 +2152,8 @@ describe('QwenAgent MCP SSE/HTTP support', () => { await agent.newSession({ cwd: '/tmp', mcpServers: [] }); - // Manager's setOnBudgetEvent must have been called once with a - // function (the closure capturing the sessionId). - expect(fakeManager.setOnBudgetEvent).toHaveBeenCalledTimes(1); + // 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 — @@ -2205,16 +2207,14 @@ describe('QwenAgent MCP SSE/HTTP support', () => { await agentPromise; }); - it('newSession is a no-op for budget wiring when getMcpClientManager is absent (defensive)', async () => { - // Older / stubbed `ToolRegistry` shapes may not expose - // `getMcpClientManager` (the bulk of older test fixtures stub - // ToolRegistry as `{ getAllTools: () => [] }`). The PR 14b code - // path uses optional chaining so the absence is silent — no - // throw, no extNotification call. - const innerConfig = await setupSessionMocks('session-no-mgr'); - (innerConfig as unknown as Record)['getToolRegistry'] = vi - .fn() - .mockReturnValue({ getAllTools: () => [] }); + 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, @@ -2233,7 +2233,12 @@ describe('QwenAgent MCP SSE/HTTP support', () => { await agent.newSession({ cwd: '/tmp', mcpServers: [] }); - // No manager → no wiring → no extNotification fires. + // No setter on Config → no wiring → no extNotification fires. + expect( + (innerConfig as unknown as Record)[ + 'setMcpBudgetEventCallback' + ], + ).toBeUndefined(); expect(extNotification).not.toHaveBeenCalled(); mockConnectionState.resolve(); diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index 16a8a1bf392..ad096a94bc9 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -13,7 +13,6 @@ import { QwenOAuth2Event, qwenOAuth2Events, MCP_BUDGET_WARN_FRACTION, - type McpBudgetEvent, MCPServerConfig, SessionService, SESSION_TITLE_MAX_LENGTH, @@ -1658,61 +1657,41 @@ class QwenAgent implements Agent { projectHooks: this.settings.getProjectHooks(), }, ); - await config.initialize(); - // PR 14b: wire the manager's budget-event callback to push the - // events out as ACP `extNotification` frames. Done AFTER - // `initialize()` (which constructs the tool registry / manager) - // and BEFORE `waitForMcpReady()` (which awaits the still-in- - // flight first discovery pass) so the registration completes - // before discovery finishes and emits its end-of-pass events. + // 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. // - // Best-effort: the manager is `off` mode (and `setOnBudgetEvent` - // is a no-op) when no budget is configured — production cost is - // a single property read. `getMcpClientManager()` may be absent - // on stubbed test ToolRegistries, hence the optional chain. + // 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. // - // sessionId capture: this Config is per-session (see - // `newSessionConfig` JSDoc + #4175 PR 14 R4 scope correction), - // so the sessionId stays stable for the manager's lifetime. - // Source: `config.getSessionId()` — Core's Config auto-assigns a - // randomUUID at construction when no `sessionId` is passed in - // (`config.ts:849`), so the value is always present after - // `loadCliConfig` returns. Defensive optional chain keeps this - // safe if a stub Config in tests omits the method. - // Best-effort lookup: older / stubbed test ToolRegistry shapes may - // omit `getMcpClientManager`, so a `typeof` check is required in - // addition to the optional chain (the optional chain only protects - // against the registry itself being missing). - const toolRegistry = config.getToolRegistry?.() as - | { getMcpClientManager?: () => unknown } - | undefined; - const budgetManager = - typeof toolRegistry?.getMcpClientManager === 'function' - ? (toolRegistry.getMcpClientManager() as - | { - setOnBudgetEvent?: ( - cb: (event: McpBudgetEvent) => void, - ) => void; - } - | undefined) - : undefined; + // 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 ( - budgetManager && - typeof budgetManager.setOnBudgetEvent === 'function' && + typeof config.setMcpBudgetEventCallback === 'function' && wiredSessionId !== undefined ) { const sid = wiredSessionId; - budgetManager.setOnBudgetEvent!((event) => { - // Fire-and-forget: extNotification returns Promise but - // the manager's call site doesn't await it. The .catch - // suppresses unhandled rejections — a mid-flight ACP - // disconnect would otherwise crash the child via uncaught - // rejection. Snapshot still carries the state for clients - // that reconnect. + 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. void this.connection .extNotification('qwen/notify/session/mcp-budget-event', { v: 1, @@ -1724,6 +1703,7 @@ class QwenAgent implements Agent { }); }); } + 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 // MCP discovery is still in flight. diff --git a/packages/cli/src/serve/httpAcpBridge.test.ts b/packages/cli/src/serve/httpAcpBridge.test.ts index 8c47a44c72c..969bc538044 100644 --- a/packages/cli/src/serve/httpAcpBridge.test.ts +++ b/packages/cli/src/serve/httpAcpBridge.test.ts @@ -4564,13 +4564,114 @@ describe('createHttpAcpBridge', () => { collected.push({ type: e.type }); if (collected.length === 1) break; } - // Exactly one event got through — the real one. The 4 dropped - // notifications produced no SSE frames. + // 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(); + }); }); describe('maxSessions cap (chiga0 Rec 3)', () => { diff --git a/packages/cli/src/serve/httpAcpBridge.ts b/packages/cli/src/serve/httpAcpBridge.ts index d2cf0120c2d..3713648ef41 100644 --- a/packages/cli/src/serve/httpAcpBridge.ts +++ b/packages/cli/src/serve/httpAcpBridge.ts @@ -1019,6 +1019,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); @@ -1226,15 +1248,36 @@ 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: 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-or-unresolvable 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). + * 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, @@ -1243,8 +1286,6 @@ class BridgeClient implements Client { if (method !== 'qwen/notify/session/mcp-budget-event') return; const sessionId = params['sessionId']; if (typeof sessionId !== 'string') return; - const entry = this.resolveEntry(sessionId); - if (!entry) return; const kind = params['kind']; const type = kind === 'budget_warning' @@ -1262,13 +1303,74 @@ class BridgeClient implements Client { void _v; void _sid; void _kind; - entry.events.publish({ + const entry = this.resolveEntry(sessionId); + const frame: Omit = { type, data: rest, - ...(entry.activePromptOriginatorClientId + ...(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(); + this.sweepExpiredEarlyEvents(now); + let buf = this.earlyEvents.get(sessionId); + if (!buf) { + if (this.earlyEvents.size >= MAX_EARLY_EVENT_SESSIONS) return; + buf = { frames: [], expiresAt: now + EARLY_EVENT_TTL_MS }; + this.earlyEvents.set(sessionId, buf); + } + if (buf.frames.length >= MAX_EARLY_EVENTS_PER_SESSION) 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); + } + } + + /** + * 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 { + 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( @@ -2416,6 +2518,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; }; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index a8e53969242..f5f1cd5cbe2 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'; @@ -683,6 +684,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(); @@ -3633,6 +3645,20 @@ 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); + } + } + if (!options?.skipDiscovery) { await registry.discoverAllTools(); } @@ -3641,4 +3667,28 @@ 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()`. The + * callback is stashed and applied lazily inside `createToolRegistry` + * (the only construction site for `McpClientManager`); a late call + * after `initialize()` already ran applies directly to the existing + * manager. + * + * `cb: undefined` clears the registration. `off`-mode managers + * silently drop the callback (their state machine never runs). + */ + setMcpBudgetEventCallback( + cb: ((event: McpBudgetEvent) => void) | undefined, + ): void { + this.pendingMcpBudgetCallback = cb; + if (this.toolRegistry) { + const mgr = this.toolRegistry.getMcpClientManager?.(); + if (mgr && typeof mgr.setOnBudgetEvent === 'function') { + mgr.setOnBudgetEvent(cb); + } + } + } } diff --git a/packages/core/src/tools/mcp-client-manager.test.ts b/packages/core/src/tools/mcp-client-manager.test.ts index a7bac31f214..ea4a8f3e2ca 100644 --- a/packages/core/src/tools/mcp-client-manager.test.ts +++ b/packages/core/src/tools/mcp-client-manager.test.ts @@ -2109,9 +2109,15 @@ describe('McpClientManager — PR 14b push events + hysteresis', () => { (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: 4, + reservedCount: 3, budget: 4, thresholdRatio: 0.75, mode: 'warn', @@ -2480,4 +2486,100 @@ describe('McpClientManager — PR 14b push events + hysteresis', () => { 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 a669cb6cd76..99a00298f0d 100644 --- a/packages/core/src/tools/mcp-client-manager.ts +++ b/packages/core/src/tools/mcp-client-manager.ts @@ -413,6 +413,25 @@ export class McpClientManager { * 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 @@ -513,9 +532,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` @@ -731,6 +773,14 @@ export class McpClientManager { * `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 @@ -787,123 +837,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: 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(); + // 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, config); - 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); - // PR 14b: end-of-pass push events. Order is intentional — - // `refused_batch` first so SDK consumers see refusals before - // any warning the same pass might have crossed (the warning - // can fire on a high reservedCount even with all-success - // connects). Both calls are no-ops when there's nothing to emit - // / `off` mode / no callback registered. - this.emitRefusedBatchIfAny(); - this.evaluateBudgetState(); + 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(); + } } /** @@ -1068,11 +1131,11 @@ export class McpClientManager { // snapshots immediately reflect reality. Mirrors the same // pattern in `readResource`'s late-reserve branch. this.dropRefusalEntry(serverName); - // PR 14b: a successful per-server (re)discover may push the - // ratio past 75% — e.g. operator did `/mcp reconnect` on the - // last server filling the budget. The bulk-pass - // `evaluateBudgetState` won't run here, so fire it inline. - this.evaluateBudgetState(); + // 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` @@ -1112,7 +1175,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. @@ -1206,7 +1269,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); } @@ -1378,178 +1441,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 = []; - // 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); - - // 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); - // PR 14b: end-of-pass push events. Mirrors `discoverAllMcpTools` - // — refused_batch first, then warning. The single-server - // `discoverMcpToolsForServerInternal` calls inside this pass - // already evaluate state on success/refusal individually, but the - // bulk-shape contract is "one batch per pass" so we coalesce - // anything missed by the single-server hooks (e.g. multiple - // discoveries refused but the bulk pass also produced its own - // refused entries via `removeServer` hitting cap edge cases). - this.emitRefusedBatchIfAny(); - this.evaluateBudgetState(); + 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 @@ -1635,7 +1710,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 @@ -1728,7 +1803,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 @@ -1842,11 +1917,9 @@ export class McpClientManager { // the next snapshot reflects the late-reservation success. if (weReservedSlot) { this.dropRefusalEntry(serverName); - // PR 14b: a successful late reservation (lazy-spawn against a - // freshly-freed slot) may push the ratio past 75%. Run the - // state machine inline since this path bypasses the bulk - // `evaluateBudgetState` at end-of-pass. - this.evaluateBudgetState(); + // 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) @@ -1962,7 +2035,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); } From 730e66d6a9e65a7f65d0afb8c51d900413748b47 Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Mon, 18 May 2026 19:05:33 +0800 Subject: [PATCH 3/9] fixup(sdk): address PR 14b review (codex P2 round 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex review round 2 of #4271 flagged two P2 issues on the SDK surface; both folded in: #1 Seed SDK replay for startup guardrail events (DaemonSessionClient.ts:78-94) 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 `DaemonSessionClient.createOrAttach` only seeded `Last-Event-ID: 0` when `req.modelServiceId` was set — every other new session started its first subscription LIVE, missing the startup-window guardrail events the new feature advertises. Unified rule: when `session.attached === false` (newly-created session) OR `req.modelServiceId` is set, seed `lastEventId: 0`. The old `modelServiceId`-only branch is preserved for re-attached sessions that need attach-time switch-event replay; the new `!attached` branch covers PR 14b's startup window. Re-attached sessions without `modelServiceId` still start live (caller may have its own event cursor it doesn't want to reset). The daemon already treats `Last-Event-ID: 0` as "replay from the beginning of the bounded ring"; if older events have been evicted from the ring (subscriber connected past the ring's coverage), the client receives the retained suffix and continues live from there — no behavior regression for late subscribers. #2 Re-export new guardrail event types from SDK barrels (daemon/index.ts + src/index.ts) PR 14b added 6 public types to `events.ts` but didn't thread them through the package's barrel exports. Consumers importing from `@qwen-code/sdk` could narrow `DaemonClientEvictedEvent` / `DaemonSlowClientWarningEvent` (older event types) but not `DaemonMcpBudgetWarningEvent` etc., breaking the SDK's encapsulation contract for the `mcp_guardrail_events` capability tag. Fix: add the 6 missing types to both `daemon/index.ts` and `src/index.ts`: - `DaemonMcpBudgetWarningData` / `DaemonMcpBudgetWarningEvent` - `DaemonMcpRefusedServer` - `DaemonMcpChildRefusedBatchData` / `DaemonMcpChildRefusedBatchEvent` - `DaemonMcpGuardrailEvent` (the union) Tests: 1 new SDK test (`replays from id 0 on freshly-created sessions so startup-window guardrail events are observable`) pins the `attached: false` → `Last-Event-ID: 0` invariant; existing 18 DaemonSessionClient tests stay green (the existing `modelServiceId`-driven replay test still passes because the new rule is OR-merged with the old one, and the `starts live when ... no model service replay need` test uses `attached: true` which bypasses the new branch). Verified: 1309/1309 tests pass across 43 files (manager, serve, acp-integration, sdk all green); typecheck clean across 4 workspaces; lint clean on 4 touched files. --- .../src/daemon/DaemonSessionClient.ts | 36 +++++++++++---- packages/sdk-typescript/src/daemon/index.ts | 8 ++++ packages/sdk-typescript/src/index.ts | 7 +++ .../test/unit/DaemonSessionClient.test.ts | 45 +++++++++++++++++++ 4 files changed, 88 insertions(+), 8 deletions(-) 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/index.ts b/packages/sdk-typescript/src/daemon/index.ts index d2ed7ccfb40..22752c4b764 100644 --- a/packages/sdk-typescript/src/daemon/index.ts +++ b/packages/sdk-typescript/src/daemon/index.ts @@ -40,6 +40,14 @@ export type { DaemonControlEvent, 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 9366752d65b..0508565dda7 100644 --- a/packages/sdk-typescript/src/index.ts +++ b/packages/sdk-typescript/src/index.ts @@ -32,6 +32,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')) { From 84a38e6f93b9312ac9023a2b9e354fe9adca699b Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Mon, 18 May 2026 19:32:35 +0800 Subject: [PATCH 4/9] fixup(serve): address PR 14b review (codex P2 round 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex review round 3 of #4271 surfaced 6 actionable items across DeepSeek / mimo / copilot agents; all 6 folded in. Two suggestions explicitly declined as documented below. Adopted: #1 Debug logging for budget events (DeepSeek Critical) - mcp-client-manager.ts: evaluateBudgetState emits `debugLogger.info` on warning fire AND on re-arm. Pre-fix the manager had ZERO log output for budget events, so oncall could not distinguish "events emitted but dropped downstream" from "events never emitted." Refusal side already had stderr via `refuseAndLog`; warning + re-arm now have parity. - acpAgent.ts: changed `.catch(() => {})` to `.catch((err) => debugLogger.debug(...))`. Pre-fix every extNotification failure was silently swallowed — including real errors (serialization bugs, protocol violations), not just the expected ACP-channel-closed case. `debug` level keeps production quiet but gives oncall a debuggable trail when they flip debug on. #2 Rename mcpRefusedBatchCount → mcpChildRefusedBatchCount (DeepSeek) ViewState fields dropped "child" from the source event name `mcp_child_refused_batch`, breaking the `slow_client_warning → slowClientWarningCount` convention. Now matches: `mcpChildRefusedBatchCount` / `lastMcpChildRefusedBatch`. PR not merged → public API rename is zero-cost. Updated 4 sites (events.ts, daemonEvents.test.ts, capabilities.ts comment, and qwen-serve-protocol.md). #3 `satisfies DaemonEvent` on test fixtures (mimo "Critical") Adds explicit type annotation to mcp_budget_warning + refused- batch test fixtures so literal discriminators (`v: 1`, `type: '...'`) stay narrow. Note: the sdk package's `tsconfig.json` currently scopes `tsc --noEmit` to `src/**/*.ts` — tests aren't gated yet — so the original "TS2345" claim was overstated. The fixture is still better typed for when tests are eventually included. #4 emitRefusedBatchIfAny docstring fix (copilot) Pre-fix the docstring claimed it "Clears both the names list and the transport map after firing", but the implementation only clears `pendingRefusalNames`. The other two are reset at pass-start / `stop()` / `dropRefusalEntry` per the PR 14 contract (snapshot-visible refusals survive between passes). Updated the docstring + the `lastRefusedTransports` field comment to describe reality. #5 Rename pgrep cross-check test (copilot) Test name claimed "in-process subprocessCount matches external pgrep observation" but the assertion uses `clientCount` (the snapshot's exposed field; `subprocessCount` is manager-internal). For stdio-only fixtures the two are numerically equal, so the assertion was correct — but the test name now matches the field. Comment block reworked to make the subprocessCount/clientCount equivalence explicit. #6 emitBudgetEvent try/catch helper (mimo) `evaluateBudgetState` and `emitRefusedBatchIfAny` now route through a private `emitBudgetEvent(event)` helper that try/catches around `this.onBudgetEvent` and debug-logs failures. Pre-fix a synchronously-throwing callback would propagate the exception up through MCP discovery / `readResource` / `disconnectServer` paths and abort unrelated work — budget push events are best-effort telemetry, never critical-path. Production ACP adapter is async + already had its own .catch, but the manager-side guard makes the contract robust against future test fixtures and adapters. Declined (documented in review-thread replies): - Bridge structural validation of extNotification payloads (wenshao Suggestion / DeepSeek): SDK predicate-based validation already routes malformed payloads through `unrecognizedKnownEventCount`; bridge-side validation duplicates SDK logic without protecting against real attacks (the daemon-child is a trust boundary). - Relaxing `thresholdRatio: 0.75` literal (wenshao Suggestion / DeepSeek): premature widening for a hypothetical second threshold (0.5 critical). Better design when actually needed: discriminator union (`{ kind: 'warning_75' } | { kind: 'critical_50' }`) rather than relaxing the literal. Verified: 1309/1309 tests pass across 43 files; typecheck clean across 4 workspaces; lint clean on 6 touched files. --- docs/developers/qwen-serve-protocol.md | 2 +- .../cli/qwen-serve-baseline.test.ts | 55 +++++++----- packages/cli/src/acp-integration/acpAgent.ts | 18 +++- packages/cli/src/serve/capabilities.ts | 2 +- packages/core/src/tools/mcp-client-manager.ts | 83 +++++++++++++++++-- packages/sdk-typescript/src/daemon/events.ts | 12 +-- .../test/unit/daemonEvents.test.ts | 24 ++++-- 7 files changed, 150 insertions(+), 46 deletions(-) diff --git a/docs/developers/qwen-serve-protocol.md b/docs/developers/qwen-serve-protocol.md index 0924830db66..a33fce5d1f6 100644 --- a/docs/developers/qwen-serve-protocol.md +++ b/docs/developers/qwen-serve-protocol.md @@ -132,7 +132,7 @@ registry. Clients **must** gate UI off `features`, not off `mode` (per design - `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`, `mcpRefusedBatchCount`, `lastMcpRefusedBatch` for adapters that want simple lag-style UI. +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 diff --git a/integration-tests/cli/qwen-serve-baseline.test.ts b/integration-tests/cli/qwen-serve-baseline.test.ts index a0704eda19c..97ca6b48d92 100644 --- a/integration-tests/cli/qwen-serve-baseline.test.ts +++ b/integration-tests/cli/qwen-serve-baseline.test.ts @@ -463,18 +463,29 @@ async function measureRssAtSessionCount(sessionCount: number): Promise<{ } }, 120_000); - // PR 14b cross-check: validate the in-process counter against - // external `pgrep -P` measurement. The push-event channel - // (`mcp_budget_warning` / `mcp_child_refused_batch`) reads - // `getMcpClientAccounting().total` for `liveCount` and the - // `subprocessCount` field is the same arithmetic - // (`stdio + websocket`). If the daemon's in-process counter - // diverges from what `pgrep -P` actually observes, the events - // would lie. Skip-gated like the parent describe (POSIX, non- - // sandbox); idle MCP fixtures are stdio-only so - // `subprocessCount` should equal `mcpGrandchildren.length` - // exactly (no amplification slack required). - it('in-process subprocessCount matches external pgrep observation', async () => { + // 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 { @@ -498,23 +509,23 @@ async function measureRssAtSessionCount(sessionCount: number): Promise<{ const snapshot = await daemon.client.workspaceMcp(); // PR 14b invariant: stdio-only fixtures → - // `subprocessCount === mcpGrandchildren.length`. - // The PR 14 amplification slack + // `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` (CONNECTED clients) is allowed to be - // lower than `subprocessCount` if the OS still sees a - // process the daemon already considers disconnected - // (rare race, kept as `<=`); but for fresh-spawn idle - // fixtures we expect equality. + // `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); - // `clientCount` is the snapshot's authoritative live count. - // Validating it against pgrep closes the loop on PR 14b's - // event-source assumption. + // 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, ); diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index ad096a94bc9..640c8ced78f 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -1692,14 +1692,28 @@ class QwenAgent implements Agent { // 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(() => { - // ACP channel closed or peer disconnected — drop event. + .catch((err: unknown) => { + debugLogger.debug( + `MCP budget extNotification dropped ` + + `(session=${sid}, kind=${event.kind}): ` + + `${err instanceof Error ? err.message : String(err)}`, + ); }); }); } diff --git a/packages/cli/src/serve/capabilities.ts b/packages/cli/src/serve/capabilities.ts index 20f7e8317e6..dccf3245b3a 100644 --- a/packages/cli/src/serve/capabilities.ts +++ b/packages/cli/src/serve/capabilities.ts @@ -103,7 +103,7 @@ export const SERVE_CAPABILITY_REGISTRY = { // length-1 per readResource refusal, only in `enforce` mode). SDK // reducer narrows both via `KnownDaemonEvent` (`DaemonSessionViewState` // exposes `mcpBudgetWarningCount`, `lastMcpBudgetWarning`, - // `mcpRefusedBatchCount`, `lastMcpRefusedBatch`). Always-on once + // `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. diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts index 99a00298f0d..097d07455f1 100644 --- a/packages/core/src/tools/mcp-client-manager.ts +++ b/packages/core/src/tools/mcp-client-manager.ts @@ -388,8 +388,17 @@ export class McpClientManager { * 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. Same lifetime as - * `lastRefusedServerNames` — reset per pass, cleared on emit. + * 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(); /** @@ -743,7 +752,18 @@ export class McpClientManager { const ratio = this.reservedSlots.size / this.clientBudget; if (this.warnArmed && ratio >= MCP_BUDGET_WARN_FRACTION) { this.warnArmed = false; - this.onBudgetEvent?.({ + // 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, @@ -753,6 +773,14 @@ export class McpClientManager { }); } 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)`, + ); } } @@ -763,10 +791,20 @@ export class McpClientManager { * spawn refusal path (where it emits a length-1 batch for shape * consistency). * - * Idempotent on no-refusals: an empty `lastRefusedServerNames` - * short-circuits without firing or clearing. Clears both the names - * list and the transport map after firing so the next pass starts - * fresh. + * 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 @@ -810,7 +848,7 @@ export class McpClientManager { transport: this.lastRefusedTransports.get(name) ?? 'unknown', reason: 'budget_exhausted' as const, })); - this.onBudgetEvent?.({ + this.emitBudgetEvent({ kind: 'refused_batch', refusedServers, budget: this.clientBudget, @@ -821,6 +859,35 @@ export class McpClientManager { 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 diff --git a/packages/sdk-typescript/src/daemon/events.ts b/packages/sdk-typescript/src/daemon/events.ts index 3b9ebc94a30..414a469e3e5 100644 --- a/packages/sdk-typescript/src/daemon/events.ts +++ b/packages/sdk-typescript/src/daemon/events.ts @@ -383,8 +383,8 @@ export interface DaemonSessionViewState { * reflects batches not refused-server entries. Mirrors the * snapshot's `disabledReason: 'budget'` per-server tag. */ - mcpRefusedBatchCount: number; - lastMcpRefusedBatch?: DaemonMcpChildRefusedBatchData; + 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 @@ -423,8 +423,8 @@ export function createDaemonSessionViewState( lastSlowClientWarning: seed.lastSlowClientWarning, mcpBudgetWarningCount: seed.mcpBudgetWarningCount ?? 0, lastMcpBudgetWarning: seed.lastMcpBudgetWarning, - mcpRefusedBatchCount: seed.mcpRefusedBatchCount ?? 0, - lastMcpRefusedBatch: seed.lastMcpRefusedBatch, + mcpChildRefusedBatchCount: seed.mcpChildRefusedBatchCount ?? 0, + lastMcpChildRefusedBatch: seed.lastMcpChildRefusedBatch, lastWorkspaceMutation: seed.lastWorkspaceMutation, lastWorkspaceMutationType: seed.lastWorkspaceMutationType, }; @@ -663,8 +663,8 @@ export function reduceDaemonSessionEvent( // session keeps running with a smaller MCP fleet. return { ...base, - mcpRefusedBatchCount: base.mcpRefusedBatchCount + 1, - lastMcpRefusedBatch: event.data, + mcpChildRefusedBatchCount: base.mcpChildRefusedBatchCount + 1, + lastMcpChildRefusedBatch: event.data, }; case 'memory_changed': // Non-terminal: adapters render a "memory just changed" hint and diff --git a/packages/sdk-typescript/test/unit/daemonEvents.test.ts b/packages/sdk-typescript/test/unit/daemonEvents.test.ts index 4205b49f616..f7fecae386c 100644 --- a/packages/sdk-typescript/test/unit/daemonEvents.test.ts +++ b/packages/sdk-typescript/test/unit/daemonEvents.test.ts @@ -750,6 +750,14 @@ describe('daemon event schema', () => { // 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, @@ -761,7 +769,7 @@ describe('daemon event schema', () => { thresholdRatio: 0.75, mode: 'warn', }, - }; + } satisfies DaemonEvent; const known = asKnownDaemonEvent(warning); expect(known?.type).toBe('mcp_budget_warning'); @@ -861,6 +869,10 @@ describe('daemon event schema', () => { }); 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, @@ -875,7 +887,7 @@ describe('daemon event schema', () => { reservedCount: 1, mode: 'enforce', }, - }; + } satisfies DaemonEvent; const known = asKnownDaemonEvent(batch); expect(known?.type).toBe('mcp_child_refused_batch'); @@ -998,8 +1010,8 @@ describe('daemon event schema', () => { }, ]); - expect(state.mcpRefusedBatchCount).toBe(2); - expect(state.lastMcpRefusedBatch).toEqual({ + expect(state.mcpChildRefusedBatchCount).toBe(2); + expect(state.lastMcpChildRefusedBatch).toEqual({ refusedServers: [ { name: 'c', transport: 'http', reason: 'budget_exhausted' }, ], @@ -1040,8 +1052,8 @@ describe('daemon event schema', () => { ); // Refused-batch counter NOT incremented — the malformed payload // didn't reach the typed reducer arm. - expect(state.mcpRefusedBatchCount).toBe(0); - expect(state.lastMcpRefusedBatch).toBeUndefined(); + expect(state.mcpChildRefusedBatchCount).toBe(0); + expect(state.lastMcpChildRefusedBatch).toBeUndefined(); }); it('narrows memory_changed events and rejects malformed payloads', () => { const valid: DaemonEvent = { From 9dfcb5844af895913109ce4c6c419d78e31886ed Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Mon, 18 May 2026 20:54:36 +0800 Subject: [PATCH 5/9] fixup(serve): address PR 14b review (codex P2 round 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex round 4 of #4271 surfaced 1 actionable item from wenshao gpt-5.5: the `Config.setMcpBudgetEventCallback → pendingMcpBudgetCallback → createToolRegistry → registry.getMcpClientManager().setOnBudgetEvent` integration boundary had NO test coverage. The acpAgent test stubs the setter (proves QwenAgent calls it pre-`initialize()`); the manager tests bypass Config by passing `onBudgetEvent` directly to `McpClientManager`. Neither covered 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. Adopted: extend the existing `vi.mock('../tools/tool-registry')` setup so each ToolRegistry instance stamps a per-instance `__mcpManagerMock` with `setOnBudgetEvent` + a stubbed `discoverAllMcpToolsIncremental` (the latter so `Config.startMcpDiscoveryInBackground` doesn't crash on missing method during `initialize`). Added 2 new tests under a new `setMcpBudgetEventCallback handoff to McpClientManager` describe: - `applies pending callback when registry is created during initialize()`: setter called BEFORE `initialize` → callback stashed in `pendingMcpBudgetCallback` → `createToolRegistry`'s apply branch forwards it to the freshly-constructed manager BEFORE discovery fires. Asserts manager's `setOnBudgetEvent` was called exactly once with the same callback function. - `applies callback directly to existing manager when called after initialize()`: setter called AFTER `initialize` → dispatches DIRECTLY to the existing manager via the `if (this.toolRegistry)` branch (the late-call path adapters use when they discover the manager only after Config is up). Also covers `cb: undefined` clearing the registration. Verified: 161/161 Config tests pass; 1470/1470 tests pass across 44 files; typecheck clean across 4 workspaces; lint clean on the 1 touched file. Closes the last open review thread on #4271. --- packages/core/src/config/config.test.ts | 96 +++++++++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 172b3983d19..5d7535841fa 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,74 @@ 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); + }); + }); }); describe('setApprovalMode with folder trust', () => { From 9b3772e5012e7ff946b78cd71488b8a875f15050 Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Mon, 18 May 2026 22:43:29 +0800 Subject: [PATCH 6/9] fixup(serve): address PR 14b review (codex P2 round 5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex round 5 of #4271 surfaced 1 Critical correctness/leak finding from wenshao gpt-5.5: pre-fix, `BridgeClient.extNotification` buffered guardrail events for ANY unknown sessionId. `resolveEntry(sessionId)` returns undefined for both newly-creating sessions (the original buffer target) AND for closed/killed sessions. A late `extNotification` from a dying child for an old sessionId would therefore land in `earlyEvents`. If the SAME sessionId came back via `session/load` or `session/resume` within the 60s TTL, `createSessionEntry`'s `drainEarlyEvents` call would replay stale prior-session telemetry onto the NEW subscriber — false budget warnings, refused server names from the OLD session leaking into a different consumer's view. Fix: tombstone Map for closed/killed session ids on `BridgeClient`. - Marked when the bridge removes a sessionId from `byId` (3 sites: channel.exited handler, closeSession, killSession). - Concurrently purges any in-flight `earlyEvents[id]` so a buffered-but-undrained 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. - Lazy `sweepExpiredTombstones` keeps the Map bounded. 3 byId.delete sites covered: - channel.exited handler (line ~2216): `info.client.markSessionClosed` - closeSession (line ~3367): `ci?.client.markSessionClosed` - killSession (line ~3787): `ci?.client.markSessionClosed` Test: new `tombstones closed sessionIds so late notifications cannot leak into a future load of the same id` exercises the full close-then-stale-notification → load-same-id-no-leak invariant via FakeAgent's loadSessionImpl returning the requested id verbatim. Verified: 1555/1555 tests pass across 46 files; typecheck clean across 4 workspaces; lint clean. --- packages/cli/src/serve/httpAcpBridge.test.ts | 95 ++++++++++++++++++++ packages/cli/src/serve/httpAcpBridge.ts | 81 +++++++++++++++++ 2 files changed, 176 insertions(+) diff --git a/packages/cli/src/serve/httpAcpBridge.test.ts b/packages/cli/src/serve/httpAcpBridge.test.ts index 242b74f83f0..bdd4d4baf6c 100644 --- a/packages/cli/src/serve/httpAcpBridge.test.ts +++ b/packages/cli/src/serve/httpAcpBridge.test.ts @@ -4672,6 +4672,101 @@ describe('createHttpAcpBridge', () => { 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(); + }); }); describe('maxSessions cap (chiga0 Rec 3)', () => { diff --git a/packages/cli/src/serve/httpAcpBridge.ts b/packages/cli/src/serve/httpAcpBridge.ts index 9b7130bee3c..93b92b386ec 100644 --- a/packages/cli/src/serve/httpAcpBridge.ts +++ b/packages/cli/src/serve/httpAcpBridge.ts @@ -1304,6 +1304,34 @@ class BridgeClient implements Client { } >(); + /** + * 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: handle child→bridge ACP `extNotification` calls. Only one * method is recognized today — `qwen/notify/session/mcp-budget-event` @@ -1376,6 +1404,13 @@ class BridgeClient implements Client { 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`. + this.sweepExpiredTombstones(now); + if (this.tombstonedSessionIds.has(sessionId)) return; this.sweepExpiredEarlyEvents(now); let buf = this.earlyEvents.get(sessionId); if (!buf) { @@ -1393,6 +1428,28 @@ class BridgeClient implements Client { } } + 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 { + this.tombstonedSessionIds.set(sessionId, Date.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 #1: drain any frames buffered for `sessionId` onto * `entry.events`. Bridge calls this immediately after @@ -1407,6 +1464,14 @@ class BridgeClient implements Client { * 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); @@ -2146,6 +2211,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(); } @@ -3292,6 +3362,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', @@ -3714,6 +3789,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)) { From 6721f175206746ec968bfc0870ad1e5097845ad4 Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Mon, 18 May 2026 23:32:05 +0800 Subject: [PATCH 7/9] fixup(serve): address PR 14b review (codex P2 round 6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex round 6 of #4271 surfaced 5 actionable items: 1 Critical regression introduced by round 5 + 4 observability/hardening suggestions. All 5 adopted. #1 (Critical) Tombstone-vs-restore window — `httpAcpBridge.ts` The round 5 tombstone (60s post-close) was rejecting LEGITIMATE restore-time guardrail events. `markSessionClosed` sets the tombstone; `drainEarlyEvents` clears it — but `drainEarlyEvents` 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 buffered → tombstone-rejected → lost. Same close → load sequence within 60s lost the entire startup window of the re-opened session. Fix: new `inFlightRestoreIds: Set` allow-list on BridgeClient. `markRestoreInFlight(sid)` called BEFORE the ACP restore call (in the IIFE that wraps `connection.loadSession` / `unstable_resumeSession`). `clearRestoreInFlight(sid)` paired in the IIFE's `finally` so success and failure both release. `bufferEarlyEvent` skips the tombstone check for ids in the allow-list. Idempotent (Set semantics) so coalesced restores work correctly. #2 (Suggestion) `thresholdRatio` predicate relaxed — `packages/sdk-typescript/src/daemon/events.ts` Pre-fix `isMcpBudgetWarningData` required `thresholdRatio === 0.75` exactly; daemon emits `MCP_BUDGET_WARN_FRACTION` from a separate package. A daemon-side bump (e.g. 0.80) would silently route every warning through `unrecognizedKnownEventCount` — a cross-package coordination hazard with no operator-visible failure mode. Relaxed to `isFiniteNumber` with a comment explaining the SDK's role is wire-shape validation, not threshold enforcement (semantics owned by daemon constant + protocol docs). NaN/Infinity rejection preserved. Test updated to assert that 0.5 is now accepted (forward-compat) and NaN is rejected (still wire- sane). #3 (Suggestion) Early-event buffer drops now log — `httpAcpBridge.ts` Pre-fix the 3 drop paths in `bufferEarlyEvent` (tombstone reject, session-cap reached, per-session-cap reached) silently `return`'d while every other drop site in this PR has stderr/debug breadcrumbs. Now each drop emits `writeStderrLine` (not `debugLogger.debug` — these drops can indicate operator-actionable abuse / fanout, so they're at the visible level matching `refuseAndLog`'s pattern). #4 (Suggestion) `emitRefusedBatchIfAny` defensive branches log — `packages/core/src/tools/mcp-client-manager.ts` Two "should never happen" branches (mode/budget invariant violation, names/queue sync gap) silently `clear` and `return` pre-fix. Now each emits `debugLogger.warn` so a future regression that reaches either branch leaves a diagnostic trail. #5 (Suggestion) Clear `pendingMcpBudgetCallback` after apply — `packages/core/src/config/config.ts` Pre-fix the field persisted for Config's lifetime. Subagent override paths (`createApprovalModeOverride` / `buildSubagentContextOverride`) call `createToolRegistry` again and would re-apply the parent session's callback to a fresh manager — routing subagent telemetry through the wrong ACP session. Now cleared to `undefined` after the apply. Late-call setter (`setMcpBudgetEventCallback` after initialize) unaffected — it dispatches directly to the existing manager. Verified: 1580/1580 tests pass across 47 files; typecheck clean across 4 workspaces; lint clean on 5 touched files. --- packages/cli/src/serve/httpAcpBridge.ts | 116 +++++++++++++++++- packages/core/src/config/config.ts | 14 +++ packages/core/src/tools/mcp-client-manager.ts | 24 ++++ packages/sdk-typescript/src/daemon/events.ts | 15 ++- .../test/unit/daemonEvents.test.ts | 28 ++++- 5 files changed, 189 insertions(+), 8 deletions(-) diff --git a/packages/cli/src/serve/httpAcpBridge.ts b/packages/cli/src/serve/httpAcpBridge.ts index 93b92b386ec..2be90e16fc2 100644 --- a/packages/cli/src/serve/httpAcpBridge.ts +++ b/packages/cli/src/serve/httpAcpBridge.ts @@ -1332,6 +1332,33 @@ class BridgeClient implements Client { */ 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` @@ -1409,16 +1436,52 @@ class BridgeClient implements Client { // 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)) return; + 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) return; + 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) return; + 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); } @@ -1450,6 +1513,37 @@ class BridgeClient implements Client { 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 @@ -2766,6 +2860,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 @@ -2899,6 +3002,13 @@ 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 (no future drain will clear it, but no future + // notifications for this id should arrive anyway since the + // child either crashed or never recognized the id). + ci?.client.clearRestoreInFlight(req.sessionId); pendingRestoreEvents.delete(req.sessionId); if (!registeredEntry) { restoreEvents.close(); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index f5f1cd5cbe2..61ebecce6c7 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -3657,6 +3657,20 @@ export class Config { 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) { diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts index 097d07455f1..d7fb4d40165 100644 --- a/packages/core/src/tools/mcp-client-manager.ts +++ b/packages/core/src/tools/mcp-client-manager.ts @@ -825,6 +825,20 @@ export class McpClientManager { // 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; } @@ -840,6 +854,16 @@ export class McpClientManager { // 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; } diff --git a/packages/sdk-typescript/src/daemon/events.ts b/packages/sdk-typescript/src/daemon/events.ts index 7a0318b82a5..6c82c7b0cd0 100644 --- a/packages/sdk-typescript/src/daemon/events.ts +++ b/packages/sdk-typescript/src/daemon/events.ts @@ -1204,12 +1204,25 @@ function isStreamErrorData(value: unknown): value is DaemonStreamErrorData { 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']) && - value['thresholdRatio'] === 0.75 && + isFiniteNumber(value['thresholdRatio']) && (value['mode'] === 'warn' || value['mode'] === 'enforce') ); } diff --git a/packages/sdk-typescript/test/unit/daemonEvents.test.ts b/packages/sdk-typescript/test/unit/daemonEvents.test.ts index 530f618281d..789a93e0e33 100644 --- a/packages/sdk-typescript/test/unit/daemonEvents.test.ts +++ b/packages/sdk-typescript/test/unit/daemonEvents.test.ts @@ -792,6 +792,14 @@ describe('daemon event schema', () => { }, }), ).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, @@ -800,10 +808,22 @@ describe('daemon event schema', () => { liveCount: 4, reservedCount: 4, budget: 4, - // `thresholdRatio` must be the literal 0.75 (the only ratio - // PR 14b emits). A future PR adding 0.50 / 0.95 thresholds - // would extend both the daemon emit + this predicate. - thresholdRatio: 0.5, + 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', }, }), From ae0cf2f01047f7953687a0f143394909a14ca052 Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Tue, 19 May 2026 00:33:37 +0800 Subject: [PATCH 8/9] fixup(serve): address PR 14b review (codex P2 round 7) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex round 7 of #4271 surfaced 3 lifecycle issues — 1 Critical subagent isolation hole + 1 P2 stale-frame leak on restore failure + 1 P3 unbounded tombstone map. All 3 adopted. #1 (P2) Late-call setMcpBudgetEventCallback should not stash — `packages/core/src/config/config.ts` Pre-fix: the late-call branch (after `initialize()`) assigned to `pendingMcpBudgetCallback` BEFORE applying directly to the existing manager. Round 6 #5 cleared the field inside `createToolRegistry`, but the late-call setter re-set it — so a subsequent `createToolRegistry` (subagent override) inherited the parent session's callback and routed subagent telemetry through the wrong ACP session. Fix: restructured to two paths. - Late-call (`toolRegistry` exists): apply directly + clear `pendingMcpBudgetCallback`. Never stash on this path. - Pre-init (no registry yet): stash for `createToolRegistry` to consume — the only way to reach a manager that doesn't exist yet. Test: new `does NOT stash the callback when called after initialize()` asserts that after a late-call setter + a subsequent `config.createToolRegistry({ forSubAgent: true })`, the subagent manager's `setOnBudgetEvent` was never called. #2 (P2) Drop buffered guardrail events on restore failure — `packages/cli/src/serve/httpAcpBridge.ts` Pre-fix: round 6 #1 added `markRestoreInFlight` so `bufferEarlyEvent` accepts frames during restore. Failure path called `clearRestoreInFlight` but did NOT purge `earlyEvents[id]`. A subsequent successful retry (`session/load` of the same id within 60s) would `drainEarlyEvents` those stale frames into the new session — exactly the leak round 5's tombstone was meant to prevent. Fix: in the restore IIFE's `finally` block, when `registeredEntry` is undefined (failure path), call `markSessionClosed(req.sessionId)`. That helper already does both required actions: refresh tombstone + delete `earlyEvents[id]`. Test: new `purges buffered guardrail events when restore fails` — spawn + close + load (fails after child queues guardrail event) + load again (succeeds) → assert no stale `mcp_budget_warning` in the ring. #3 (P3) Sweep tombstones in markSessionClosed — `packages/cli/src/serve/httpAcpBridge.ts` Pre-fix: `sweepExpiredTombstones` was only called inside `bufferEarlyEvent`. On a daemon with high session-churn but few extNotifications (the common production pattern when MCP guardrail mode is `off`), the tombstone map grew monotonically — the documented 60s TTL didn't bound memory. Fix: sweep at the top of `markSessionClosed`. Cheap (one integer compare per entry); under any realistic workload the map stays small. Verified: 1582/1582 tests pass across 47 files; typecheck clean across 4 workspaces; lint clean on 4 touched files. --- packages/cli/src/serve/httpAcpBridge.test.ts | 104 +++++++++++++++++++ packages/cli/src/serve/httpAcpBridge.ts | 30 +++++- packages/core/src/config/config.test.ts | 42 ++++++++ packages/core/src/config/config.ts | 29 ++++-- 4 files changed, 195 insertions(+), 10 deletions(-) diff --git a/packages/cli/src/serve/httpAcpBridge.test.ts b/packages/cli/src/serve/httpAcpBridge.test.ts index bdd4d4baf6c..348b0511817 100644 --- a/packages/cli/src/serve/httpAcpBridge.test.ts +++ b/packages/cli/src/serve/httpAcpBridge.test.ts @@ -4767,6 +4767,110 @@ describe('createHttpAcpBridge', () => { 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)', () => { diff --git a/packages/cli/src/serve/httpAcpBridge.ts b/packages/cli/src/serve/httpAcpBridge.ts index 2be90e16fc2..9296b5239a7 100644 --- a/packages/cli/src/serve/httpAcpBridge.ts +++ b/packages/cli/src/serve/httpAcpBridge.ts @@ -1507,7 +1507,18 @@ class BridgeClient implements Client { * in-flight stale frames to expire. */ markSessionClosed(sessionId: string): void { - this.tombstonedSessionIds.set(sessionId, Date.now() + EARLY_EVENT_TTL_MS); + 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); @@ -3005,13 +3016,24 @@ export function createHttpAcpBridge(opts: BridgeOptions): HttpAcpBridge { // 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 (no future drain will clear it, but no future - // notifications for this id should arrive anyway since the - // child either crashed or never recognized the id). + // 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); } }); diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 5d7535841fa..99ea9eebff5 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2174,6 +2174,48 @@ describe('Server Config (config.ts)', () => { 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(); + }); }); }); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 61ebecce6c7..52f64276b37 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -3685,11 +3685,23 @@ export class Config { /** * 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()`. The - * callback is stashed and applied lazily inside `createToolRegistry` - * (the only construction site for `McpClientManager`); a late call - * after `initialize()` already ran applies directly to the existing - * manager. + * 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). @@ -3697,12 +3709,17 @@ export class Config { setMcpBudgetEventCallback( cb: ((event: McpBudgetEvent) => void) | undefined, ): void { - this.pendingMcpBudgetCallback = cb; 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; } } From 61b1f144a32000d87f16b13fe590c036ac47b67f Mon Sep 17 00:00:00 2001 From: doudouOUC Date: Tue, 19 May 2026 00:46:52 +0800 Subject: [PATCH 9/9] fixup(sdk): satisfies DaemonEvent on slow_client_warning fixture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex round 8 flagged `slow_client_warning` (line 670, this file) as the sibling of PR 14b's `mcp_budget_warning` / `mcp_child_refused_batch` fixtures (rounds 3 #3) — same widening shape (`v: 1` → `number`, `type: '…'` → `string`), same describe block, same `asKnownDaemonEvent` call site. Note: the reviewer's claim that the issue was "caused by the PR's type additions in events.ts which widen TypeScript's inference for existing object literals" is **inaccurate**. The widening is inherent to TypeScript's literal-inference rules for `let`-style object bindings and predates PR 14b. Verified by running explicit `tsc --noEmit test/unit/daemonEvents.test.ts`: 18 TS2345 errors across the file, ~16 of which involve `model_switched` / `permission_request` / `session_update` / `client_evicted` / `stream_error` fixtures that PR 14b never touched. Why fix only one: `slow_client_warning` is the closest topical sibling to PR 14b's two fixtures (the round 3 #3 contract explicitly was "PR 14b's own fixtures"). Extending to the other 17 widening sites would (a) bloat the PR diff with PR 4 / PR 10 / PR 11 era test debt unrelated to PR 14b's review focus, and (b) still leave the underlying root cause (sdk `tsconfig.json` excludes the test directory from typecheck) — `npm run typecheck` passes clean before AND after this commit because tests aren't in scope. A future PR can opt tests into the typecheck scope and fix all 17 in one batch with the actual gating change. Verified: 46/46 sdk daemonEvents tests pass; 1656/1656 across the sweep; typecheck clean across 4 workspaces; lint clean. --- .../test/unit/daemonEvents.test.ts | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/packages/sdk-typescript/test/unit/daemonEvents.test.ts b/packages/sdk-typescript/test/unit/daemonEvents.test.ts index b0837a5b63f..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');