Register a subscription before awaiting its snapshot

[Ryan Ewen (RyanEwen)]

A chat response intermittently renders as an empty bubble: the request shows, the response header shows the right duration and model, and no content ever appears. Reloading the window shows the entire response, so the state was always correct.

Cause

subscribe registers the client's subscription only after awaiting IAgentService.subscribe:

const snapshot = await this._agentService.subscribe(...);
client.subscriptions.set(classified.uri, classified);   // too late

A client subscribes and then dispatches on the same connection without waiting for the response. Any envelope emitted during that await is routed to no subscriber, so the client never receives the echo of its own write-ahead action.

For a chat turn the lost echo is chat/turnStarted. The renderer installs its turn observer when that action comes back, so it is never installed: no markdown part is ever set up, no tool call is ever rendered, and every later envelope arrives and is ignored. The agent host runs the turn normally throughout.

Evidence

Captured on a dev container over websocket. Same client, same prompt, two consecutive turns:

painted     00:16:11.074 c2s subscribe
            00:16:11.076 c2s chat/turnStarted     (+2ms)
            00:16:11.076 s2c chat/turnStarted     <- echo delivered

unpainted   00:17:21.442 c2s subscribe
            00:17:21.443 c2s chat/turnStarted     (+1ms)
                             (no echo, ever)
            00:17:23.333 s2c chat/delta           <- later envelopes arrive

Instrumenting _setupMarkdownPart in the workbench confirmed the consequence: on the painted turn it logged setup and 16 emits; on the unpainted turn it was never called at all, while 11 chat/delta events crossed the wire. Measured DOM text was 1591 chars vs 319 (the user's prompt alone), reproducing in 3 of 8 unattended runs.

Fix

Register before the await, and remove the registration if the subscribe rejects. The OtlpLogs and ResourceWatch branches directly above already register synchronously before returning; this branch was the outlier.

Testing

• Added an action dispatched while subscribe is still resolving is still echoed, which gates the fake service's subscribe on a promise, dispatches chat/turnStarted during the await, then releases it. Without the fix it fails with "turnStarted must be echoed to the client that dispatched it".
• ProtocolServerHandler 94 passing, AgentSideEffects 219 passing.
• npm run typecheck-client clean.

[Copilot]

Pull request overview

Prevents lost action echoes while state snapshots are resolving.

Changes:

• Registers subscriptions before awaiting snapshots.
• Rolls back registrations on failure.
• Adds regression coverage for concurrent dispatch.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
protocolServerHandler.ts	Adjusts subscription registration ordering.
protocolServerHandler.test.ts	Tests dispatch during pending subscription.

Suppressed comments (1)

src/vs/platform/agentHost/node/protocolServerHandler.ts:1325

• A stale failed subscribe can delete a newer subscription for the same URI. Request handlers run concurrently, and the client subscription manager can unsubscribe/reacquire while the first request is pending; the old catch then removes the replacement map entry, while AgentService.subscribe also rolls back the shared client-id registration. Make subscribe cleanup generation-aware (and preserve/re-register the current agent-service subscription) and cover the unsubscribe/resubscribe race.

				client.subscriptions.delete(classified.uri);

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ryan Ewen (RyanEwen)]

Both comment findings were already addressed in 55e28eb, pushed about a minute after this review ran. The inline comment in protocolServerHandler.ts and the one in the test are each a single line now.

That commit also covers two things this review did not flag, applying patterns raised on the earlier PRs in this series: removing the subscription on a failed subscribe destroyed a pre-existing valid subscription when a client re-subscribes to a channel it already had (now only the entry this call added is removed), and the test matched on action type alone, so it now asserts the echoed turn ids.