src/vs/platform/agentHost/common/state/agentSubscription.ts
1201 LOC · 992 covered · 209 uncovered · 210 ranges · 253 concepts · 60 introducers · 123 tests
File neighbourhood
The centred file is linked to every concept that introduces one of its ranges, every test that runs code from the file, and the gray connector concepts standing between those tests and the file's own introducer concepts. Undirected links join concepts to every file where they introduce source and concepts to the tests they introduce; arrows show specialization between the displayed concepts and bridge only concepts omitted from this view. Concept colors match the source ranges below; connector concepts have no source color and are shown in gray.
Focused file, its introducer and connector concepts, their introduced files, and tests that run code from the file
In the embedded map, ordinary wheel input scrolls the page; use the visible controls to zoom and drag to pan. Open the full-screen map for canvas navigation: wheel pans, Ctrl/Command plus wheel zooms, and arrow keys pan when this region is focused. On touch screens, open the full-screen map to pan or pinch. If JavaScript or WebGL is unavailable, use the related-file, concept, and source links on this page.
Graph controls are ready.
Interactive rendering requires JavaScript and WebGL. Use the related-file, concept, and source links on this page while the interactive map is unavailable.
1
>
/*---------------------------------------------------------------------------------------------
agentSubscription.ts ×81
2
>
* Copyright (c) Microsoft Corporation. All rights reserved.
3
>
* Licensed under the MIT License. See License.txt in the project root for license information.
4
>
*--------------------------------------------------------------------------------------------*/
5
>
6
>
import { assertNever } from '../../../../base/common/assert.js';
7
>
import { Emitter, Event } from '../../../../base/common/event.js';
8
>
import { Disposable, IReference } from '../../../../base/common/lifecycle.js';
9
>
import { ResourceMap } from '../../../../base/common/map.js';
10
>
import { IObservable, observableFromEvent } from '../../../../base/common/observable.js';
11
>
import { URI } from '../../../../base/common/uri.js';
12
>
import { ActionEnvelope, ActionType, ChangesetAction, ChatAction, AnnotationsAction, ClientAnnotationsAction, ClientChangesetAction, IRootConfigChangedAction, SessionAction, StateAction, isChangesetAction, isChatAction, isAnnotationsAction, isSessionAction } from './sessionActions.js';
13
>
import { changesetReducer, chatReducer, annotationsReducer, rootReducer, sessionReducer } from './sessionReducers.js';
14
>
import { terminalReducer } from './protocol/reducers.js';
15
>
import type { RootAction, SessionAction as IProtocolSessionAction, ChatAction as IProtocolChatAction, TerminalAction } from './protocol/action-origin.generated.js';
16
>
import type { AnnotationsState, ChangesetState, ChatState, RootState, SessionState, TerminalState } from './protocol/state.js';
17
>
import type { IStateSnapshot } from './sessionProtocol.js';
18
>
import { isAhpRootChannel, ROOT_STATE_URI, StateComponents } from './sessionState.js';
19
>
20
>
// --- Public API --------------------------------------------------------------
21
>
22
>
/**
23
>
* A read-only subscription to an agent host resource (root, session, or terminal).
24
>
*
25
>
* Subscriptions are hydrated from an initial server snapshot and kept in sync
26
>
* via action envelopes. Session subscriptions support write-ahead
27
>
* reconciliation — optimistic state is layered on top of confirmed state.
28
>
*/
29
>
export interface IAgentSubscription<T> {
30
>
/**
31
>
* The current state value. For write-ahead subscriptions (sessions) this
32
>
* reflects the optimistic state (confirmed + pending replayed). For
33
>
* server-only subscriptions (root, terminal) this equals `verifiedValue`.
34
>
*
35
>
* `undefined` until the first snapshot arrives. An `Error` if subscription
36
>
* failed.
37
>
*/
38
>
readonly value: T | Error | undefined;
39
>
40
>
/**
41
>
* The server-confirmed state with no pending optimistic actions applied.
42
>
* `undefined` until the first snapshot arrives.
43
>
*/
44
>
readonly verifiedValue: T | undefined;
45
>
46
>
/** Fires when {@link value} changes (optimistic or confirmed). */
47
>
readonly onDidChange: Event<T>;
48
>
49
>
/** Fires when the subscription enters an error state. */
50
>
readonly onDidError?: Event<Error>;
51
>
52
>
/** Fires before a server-originated action is applied to this subscription's state. */
53
>
readonly onWillApplyAction: Event<ActionEnvelope>;
54
>
55
>
/** Fires after a server-originated action is applied to this subscription's state. */
56
>
readonly onDidApplyAction: Event<ActionEnvelope>;
57
>
}
58
>
59
>
/**
60
>
* Read-only snapshot describing a single active resource subscription. Used by
61
>
* inspection/debug surfaces that enumerate everything a connection is currently
62
>
* subscribed to. Does not include the always-live root state.
63
>
*/
64
>
export interface IActiveSubscriptionInfo {
65
>
/** The protocol resource URI subscribed to. */
66
>
readonly resource: URI;
67
>
/** Which state component this subscription tracks. */
68
>
readonly kind: StateComponents;
69
>
/** Number of outstanding {@link IReference} holders. */
70
>
readonly refCount: number;
71
>
/**
72
>
* The named owners currently holding a reference to this subscription,
73
>
* with how many references each holds. Names come from the `owner`
74
>
* argument passed to {@link AgentSubscriptionManager.getSubscription}.
75
>
*/
76
>
readonly holders: readonly IActiveSubscriptionHolder[];
77
>
/**
78
>
* Lifecycle status derived from the subscription's value:
79
>
* `pending` before the first snapshot, `error` if it failed, otherwise
80
>
* `snapshot`.
81
>
*/
82
>
readonly status: 'pending' | 'snapshot' | 'error';
83
>
}
84
>
85
>
/** A named owner holding one or more references to a subscription. */
86
>
export interface IActiveSubscriptionHolder {
87
>
readonly owner: string;
88
>
readonly count: number;
89
>
}
90
>
91
>
// --- Base Implementation -----------------------------------------------------
92
>
93
>
/**
94
>
* Base class for agent subscriptions. Handles envelope reception, confirmed
95
>
* state management, and action event emission.
96
>
*
97
>
* Subclasses provide the reducer and optionally override reconciliation
98
>
* behavior.
99
>
*/
100
>
abstract class BaseAgentSubscription<T> extends Disposable implements IAgentSubscription<T> {
101
>
102
>
protected _confirmedState: T | undefined;
103
>
private _error: Error | undefined;
104
>
private _bufferedEnvelopes: ActionEnvelope[] | undefined;
105
>
106
>
protected readonly _onDidChange = this._register(new Emitter<T>());
107
>
readonly onDidChange: Event<T> = this._onDidChange.event;
108
>
109
>
protected readonly _onDidError = this._register(new Emitter<Error>());
110
>
readonly onDidError: Event<Error> = this._onDidError.event;
111
>
112
>
protected readonly _onWillApplyAction = this._register(new Emitter<ActionEnvelope>());
113
>
readonly onWillApplyAction: Event<ActionEnvelope> = this._onWillApplyAction.event;
114
>
115
>
protected readonly _onDidApplyAction = this._register(new Emitter<ActionEnvelope>());
116
>
readonly onDidApplyAction: Event<ActionEnvelope> = this._onDidApplyAction.event;
117
>
118
>
protected readonly _clientId: string;
119
>
protected readonly _log: (msg: string) => void;
120
>
121
>
constructor(clientId: string, log: (msg: string) => void) {
123
>
this._clientId = clientId;
124
>
this._log = log;
125
>
}
127
>
get value(): T | Error | undefined {
130
>
}
134
>
get verifiedValue(): T | undefined {
136
>
}
138
>
/**
139
>
* Apply an initial snapshot from the server.
140
>
*/
141
>
handleSnapshot(state: T, fromSeq: number): void {
143
>
this._error = undefined;
144
>
this._onSnapshotApplied(fromSeq);
145
>
this._onDidChange.fire(this.value as T);
146
>
}
148
>
/**
149
>
* Mark this subscription as failed.
150
>
*/
151
>
setError(error: Error): void {
153
>
this._onDidError.fire(error);
154
>
}
156
>
/**
157
>
* Process an incoming action envelope. The subscription determines
158
>
* whether the action is relevant via {@link _isRelevantEnvelope}.
159
>
*/
160
>
receiveEnvelope(envelope: ActionEnvelope): void {
163
>
}
165
>
// Buffer actions that arrive before the snapshot has been applied.
166
>
// They're replayed in _onSnapshotApplied().
167
>
if (this._confirmedState === undefined) {
169
>
this._bufferedEnvelopes = [];
170
>
}
171
>
this._bufferedEnvelopes.push(envelope);
172
>
return;
173
>
}
176
>
this._onWillApplyAction.fire(envelope);
177
>
178
>
this._reconcile(envelope, isOwnAction);
179
>
180
>
this._onDidApplyAction.fire(envelope);
181
>
}
183
>
/** Apply the reducer to confirmed state. Subclasses must implement. */
184
>
protected abstract _applyReducer(state: T, action: StateAction): T;
185
>
186
>
/** Whether the given envelope targets this subscription. */
187
>
protected abstract _isRelevantEnvelope(envelope: ActionEnvelope): boolean;
188
>
189
>
/** Return optimistic state if write-ahead is active, otherwise `undefined`. */
190
>
protected _getOptimisticState(): T | undefined {
192
>
}
194
>
/** Hook called after a snapshot is applied. Replays buffered actions. */
195
>
protected _onSnapshotApplied(_fromSeq: number): void {
197
>
const buffered = this._bufferedEnvelopes;
198
>
if (buffered) {
200
>
for (const envelope of buffered) {
201
>
// Only replay actions with serverSeq > fromSeq (snapshot is authoritative up to fromSeq)
202
>
if (envelope.serverSeq > _fromSeq) {
204
>
this._reconcile(envelope, isOwnAction);
205
>
}
207
>
}
210
>
/**
211
>
* Default reconciliation: apply to confirmed, fire change event.
212
>
* Session subscriptions override this for write-ahead.
213
>
*/
214
>
protected _reconcile(envelope: ActionEnvelope, _isOwnAction: boolean): void {
215
>
this._confirmedState = this._applyReducer(this._confirmedState!, envelope.action);
agentSubscription.ts ×1
216
>
this._onDidChange.fire(this.value as T);
217
>
}
219
>
220
>
// --- Root State Subscription -------------------------------------------------
221
>
222
>
/**
223
>
* Subscription to the root state at `agenthost:/root`.
224
>
* Server-only mutations — no write-ahead.
225
>
*/
226
>
export class RootStateSubscription extends BaseAgentSubscription<RootState> {
227
>
228
>
protected override _applyReducer(state: RootState, action: StateAction): RootState {
230
>
}
232
>
protected override _isRelevantEnvelope(envelope: ActionEnvelope): boolean {
233
>
return isAhpRootChannel(envelope.channel) && envelope.action.type.startsWith('root/');
agentSubscription.ts ×1
234
>
}
236
>
237
>
// --- Session State Subscription ----------------------------------------------
238
>
239
>
interface IPendingAction {
240
>
readonly clientSeq: number;
241
>
readonly action: SessionAction;
242
>
}
243
>
244
>
/**
245
>
* A pending optimistic action awaiting server confirmation, paired with the
246
>
* channel it was dispatched to so it can be replayed across a reconnect. The
247
>
* channel is a session channel for {@link SessionStateSubscription} actions and
248
>
* a chat channel for {@link ChatStateSubscription} actions.
249
>
*/
250
>
export interface IPendingDispatchAction {
251
>
readonly clientSeq: number;
252
>
/** The optimistic action awaiting confirmation. */
253
>
readonly action: SessionAction | ChatAction;
254
>
/** URI of the channel this action targets, as stored on the subscription. */
255
>
readonly channel: string;
256
>
}
257
>
258
>
/**
259
>
* Subscription to a session at `copilot:/<uuid>`.
260
>
* Supports write-ahead reconciliation for client-dispatchable actions.
261
>
*/
262
>
export class SessionStateSubscription extends BaseAgentSubscription<SessionState> {
263
>
264
>
private readonly _pendingActions: IPendingAction[] = [];
265
>
private _optimisticState: SessionState | undefined;
266
>
private readonly _sessionUri: string;
267
>
private readonly _seqAllocator: () => number;
268
>
269
>
constructor(
271
>
clientId: string,
272
>
seqAllocator: () => number,
273
>
log: (msg: string) => void,
274
>
) {
275
>
super(clientId, log);
276
>
this._sessionUri = sessionUri;
277
>
this._seqAllocator = seqAllocator;
278
>
}
280
>
/**
281
>
* Optimistically apply a session action. Returns the clientSeq to send
282
>
* to the server so it can echo back for reconciliation.
283
>
*/
284
>
applyOptimistic(action: SessionAction): number {
286
>
this._pendingActions.push({ clientSeq, action });
287
>
// Apply on top of current optimistic
288
>
const base = this._optimisticState ?? this.verifiedValue;
289
>
if (base) {
290
>
this._optimisticState = sessionReducer(base, action as IProtocolSessionAction, this._log);
291
>
this._onDidChange.fire(this._optimisticState);
292
>
}
293
>
return clientSeq;
294
>
}
296
>
protected override _getOptimisticState(): SessionState | undefined {
298
>
}
300
>
protected override _applyReducer(state: SessionState, action: StateAction): SessionState {
301
>
return sessionReducer(state, action as IProtocolSessionAction, this._log);
agentSubscription.ts ×2
302
>
}
304
>
protected override _isRelevantEnvelope(envelope: ActionEnvelope): boolean {
305
>
return isSessionAction(envelope.action) && envelope.channel === this._sessionUri;
agentSubscription.ts ×1
306
>
}
308
>
protected override _onSnapshotApplied(fromSeq: number): void {
310
>
super._onSnapshotApplied(fromSeq);
311
>
// Re-apply pending actions on top of new confirmed state
312
>
this._recomputeOptimistic();
313
>
}
315
>
protected override _reconcile(envelope: ActionEnvelope, isOwnAction: boolean): void {
316
>
// A rejected envelope must never mutate confirmed state — it only rolls
agentSubscription.ts ×3
317
>
// back the originating client's matching optimistic action. Guarding all
318
>
// apply branches also prevents a broadcast rejection from leaking the
319
>
// rejected action into a non-origin client's state.
320
>
if (isOwnAction && envelope.origin) {
321
>
const idx = this._pendingActions.findIndex(p => p.clientSeq === envelope.origin!.clientSeq);
agentSubscription.ts ×2
322
>
if (idx !== -1) {
323
>
if (!envelope.rejectionReason) {
325
>
}
327
>
} else if (!envelope.rejectionReason) {
328
this._confirmedApply(envelope.action);
329
}
332
>
}
334
>
}
336
>
private _confirmedApply(action: StateAction): void {
338
>
this._confirmedState = this._applyReducer(this._confirmedState, action);
339
>
}
340
>
}
342
>
private _recomputeOptimistic(): void {
344
>
if (!confirmed) {
345
this._optimisticState = undefined;
346
return;
347
}
349
>
if (this._pendingActions.length === 0) {
350
>
this._optimisticState = undefined; // No pending → value falls through to confirmed
351
>
this._onDidChange.fire(confirmed);
352
>
return;
353
>
}
355
>
let state = confirmed;
356
>
for (const pending of this._pendingActions) {
357
>
state = sessionReducer(state, pending.action as IProtocolSessionAction, this._log);
358
>
}
359
>
this._optimisticState = state;
360
>
this._onDidChange.fire(state);
363
>
/**
364
>
* Clear pending actions for this session (e.g., on unsubscribe).
365
>
*/
366
>
clearPending(): void {
368
>
this._optimisticState = undefined;
369
>
}
371
>
/**
372
>
* Snapshot of the currently-pending optimistic actions, with the session
373
>
* URI included so callers can re-issue them across a reconnect. The
374
>
* actions remain in the subscription so the optimistic state continues
375
>
* to reflect them — the client must explicitly drop entries echoed back
376
>
* by the server.
377
>
*/
378
>
getPendingActions(): IPendingDispatchAction[] {
379
return this._pendingActions.map(p => ({ clientSeq: p.clientSeq, action: p.action, channel: this._sessionUri }));
380
}
382
>
/**
383
>
* Drop the pending entry whose `clientSeq` matches the supplied value.
384
>
* Used during reconnect to evict actions the server already echoed back
385
>
* in the replay buffer so they're not resent.
386
>
*/
387
>
dropPendingByClientSeq(clientSeq: number): boolean {
388
const idx = this._pendingActions.findIndex(p => p.clientSeq === clientSeq);
389
if (idx === -1) {
390
return false;
391
}
392
this._pendingActions.splice(idx, 1);
393
return true;
394
}
396
>
397
>
// --- Chat State Subscription -------------------------------------------------
398
>
399
>
interface IPendingChatAction {
400
>
readonly clientSeq: number;
401
>
readonly action: ChatAction;
402
>
}
403
>
404
>
/**
405
>
* Subscription to a chat channel (e.g. a session's default chat URI). Turns,
406
>
* tool calls and pending/input state moved off the session onto the chat
407
>
* channel in the multi-chat protocol, so this subscription carries the
408
>
* conversation contents. Supports write-ahead reconciliation for
409
>
* client-dispatchable chat actions (turn starts, confirmations, etc.).
410
>
*/
411
>
export class ChatStateSubscription extends BaseAgentSubscription<ChatState> {
412
>
413
>
private readonly _pendingActions: IPendingChatAction[] = [];
414
>
private _optimisticState: ChatState | undefined;
415
>
private readonly _chatUri: string;
416
>
private readonly _seqAllocator: () => number;
417
>
418
>
constructor(
420
>
clientId: string,
421
>
seqAllocator: () => number,
422
>
log: (msg: string) => void,
423
>
) {
424
>
super(clientId, log);
425
>
this._chatUri = chatUri;
426
>
this._seqAllocator = seqAllocator;
427
>
}
429
>
/**
430
>
* Optimistically apply a chat action. Returns the clientSeq to send to
431
>
* the server so it can echo back for reconciliation.
432
>
*/
433
>
applyOptimistic(action: ChatAction): number {
435
>
this._pendingActions.push({ clientSeq, action });
436
>
const base = this._optimisticState ?? this.verifiedValue;
437
>
if (base) {
438
>
this._optimisticState = chatReducer(base, action as IProtocolChatAction, this._log);
439
>
this._onDidChange.fire(this._optimisticState);
440
>
}
441
>
return clientSeq;
442
>
}
444
>
protected override _getOptimisticState(): ChatState | undefined {
446
>
}
448
>
protected override _applyReducer(state: ChatState, action: StateAction): ChatState {
450
>
}
452
>
protected override _isRelevantEnvelope(envelope: ActionEnvelope): boolean {
453
>
return isChatAction(envelope.action) && envelope.channel === this._chatUri;
agentSubscription.ts ×16
454
>
}
456
>
protected override _onSnapshotApplied(fromSeq: number): void {
458
>
this._recomputeOptimistic();
459
>
}
461
>
protected override _reconcile(envelope: ActionEnvelope, isOwnAction: boolean): void {
462
>
// A rejected envelope must never mutate confirmed state — it only rolls
agentSubscription.ts ×16
463
>
// back the originating client's matching optimistic action. Guarding all
464
>
// apply branches also prevents a broadcast rejection from leaking the
465
>
// rejected action into a non-origin client's state.
466
>
if (isOwnAction && envelope.origin) {
467
const idx = this._pendingActions.findIndex(p => p.clientSeq === envelope.origin!.clientSeq);
468
if (idx !== -1) {
469
if (!envelope.rejectionReason) {
470
this._confirmedApply(envelope.action);
471
}
472
this._pendingActions.splice(idx, 1);
473
} else if (!envelope.rejectionReason) {
474
this._confirmedApply(envelope.action);
475
}
477
>
this._promotePendingTurnStartIfTerminal(envelope.action);
478
>
this._confirmedApply(envelope.action);
479
>
}
480
>
this._recomputeOptimistic();
481
>
}
483
>
private _promotePendingTurnStartIfTerminal(action: StateAction): void {
484
>
// A backend-originated terminal turn action may arrive without the clientSeq
agentSubscription.ts ×16
485
>
// that would normally confirm our optimistic turn start. Promote that start
486
>
// first so the terminal action can close it instead of leaving it pending.
487
>
if (!isChatAction(action)) {
488
return;
489
}
490
>
if (action.type !== ActionType.ChatTurnComplete && action.type !== ActionType.ChatTurnCancelled && action.type !== ActionType.ChatError) {
agentSubscription.ts ×16
491
return;
492
}
493
>
const index = this._pendingActions.findIndex(p => p.action.type === ActionType.ChatTurnStarted && p.action.turnId === action.turnId);
agentSubscription.ts ×16
494
>
if (index === -1) {
495
return;
496
}
497
>
const [{ action: pendingAction }] = this._pendingActions.splice(index, 1);
agentSubscription.ts ×16
498
>
if (this._confirmedState && (!this._confirmedState.activeTurn || this._confirmedState.activeTurn.id !== action.turnId)) {
499
>
this._confirmedState = this._applyReducer(this._confirmedState, pendingAction);
500
>
}
501
>
}
503
>
private _confirmedApply(action: StateAction): void {
505
>
this._confirmedState = this._applyReducer(this._confirmedState, action);
506
>
}
507
>
}
509
>
private _recomputeOptimistic(): void {
511
>
if (!confirmed) {
512
this._optimisticState = undefined;
513
return;
514
}
516
>
this._optimisticState = undefined;
517
>
this._onDidChange.fire(confirmed);
518
>
return;
519
>
}
520
let state = confirmed;
521
for (const pending of this._pendingActions) {
522
state = chatReducer(state, pending.action as IProtocolChatAction, this._log);
523
}
524
this._optimisticState = state;
525
this._onDidChange.fire(state);
528
>
clearPending(): void {
529
this._pendingActions.length = 0;
530
this._optimisticState = undefined;
531
}
533
>
getPendingActions(): IPendingDispatchAction[] {
534
return this._pendingActions.map(p => ({ clientSeq: p.clientSeq, action: p.action, channel: this._chatUri }));
535
}
537
>
dropPendingByClientSeq(clientSeq: number): boolean {
538
const idx = this._pendingActions.findIndex(p => p.clientSeq === clientSeq);
539
if (idx === -1) {
540
return false;
541
}
542
this._pendingActions.splice(idx, 1);
543
return true;
544
}
546
>
547
>
// --- Terminal State Subscription ---------------------------------------------
548
>
549
>
/**
550
>
* Subscription to a terminal at an agent-host terminal URI.
551
>
* Server-only mutations — no write-ahead (terminal I/O is side-effect-only).
552
>
*/
553
>
export class TerminalStateSubscription extends BaseAgentSubscription<TerminalState> {
554
>
555
>
private readonly _terminalUri: string;
556
>
557
>
constructor(terminalUri: string, clientId: string, log: (msg: string) => void) {
559
>
this._terminalUri = terminalUri;
560
>
}
562
>
protected override _applyReducer(state: TerminalState, action: StateAction): TerminalState {
564
>
}
566
>
protected override _isRelevantEnvelope(envelope: ActionEnvelope): boolean {
567
>
return envelope.action.type.startsWith('terminal/') && envelope.channel === this._terminalUri;
agentSubscription.ts ×1
568
>
}
570
>
571
>
// --- Changeset State Subscription --------------------------------------------
572
>
573
>
/**
574
>
* Subscription to a changeset at an expanded changeset URI (e.g.
575
>
* `<sessionUri>/changeset/session`).
576
>
*
577
>
* Changeset review actions are client-dispatchable, so this subscription
578
>
* supports write-ahead reconciliation.
579
>
*/
580
>
export class ChangesetStateSubscription extends BaseAgentSubscription<ChangesetState> {
581
>
582
>
private readonly _pendingActions: { readonly clientSeq: number; readonly action: ClientChangesetAction }[] = [];
583
>
private _optimisticState: ChangesetState | undefined;
584
>
private readonly _changesetUri: string;
585
>
private readonly _seqAllocator: () => number;
586
>
587
>
constructor(changesetUri: string, clientId: string, seqAllocator: () => number, log: (msg: string) => void) {
589
>
this._changesetUri = changesetUri;
590
>
this._seqAllocator = seqAllocator;
591
>
}
593
>
/**
594
>
* Optimistically apply a changeset action and return its client sequence.
595
>
*/
596
>
applyOptimistic(action: ClientChangesetAction): number {
598
>
this._pendingActions.push({ clientSeq, action });
599
>
const base = this._optimisticState ?? this.verifiedValue;
600
>
if (base) {
601
>
this._optimisticState = changesetReducer(base, action, this._log);
602
>
this._onDidChange.fire(this._optimisticState);
603
>
}
604
>
return clientSeq;
605
>
}
607
>
protected override _getOptimisticState(): ChangesetState | undefined {
609
>
}
611
>
protected override _applyReducer(state: ChangesetState, action: StateAction): ChangesetState {
613
>
}
615
>
protected override _isRelevantEnvelope(envelope: ActionEnvelope): boolean {
616
>
return isChangesetAction(envelope.action) && envelope.channel === this._changesetUri;
agentSubscription.ts ×6
617
>
}
619
>
protected override _onSnapshotApplied(fromSeq: number): void {
621
>
this._recomputeOptimistic();
622
>
}
624
>
protected override _reconcile(envelope: ActionEnvelope, isOwnAction: boolean): void {
626
>
const index = this._pendingActions.findIndex(pending => pending.clientSeq === envelope.origin!.clientSeq);
627
>
if (index !== -1) {
628
>
if (!envelope.rejectionReason) {
629
>
this._confirmedApply(envelope.action);
630
>
}
631
>
this._pendingActions.splice(index, 1);
632
>
} else {
633
this._confirmedApply(envelope.action);
634
}
636
this._confirmedApply(envelope.action);
637
}
639
>
}
641
>
private _confirmedApply(action: StateAction): void {
643
>
this._confirmedState = this._applyReducer(this._confirmedState, action);
644
>
}
645
>
}
647
>
private _recomputeOptimistic(): void {
649
>
if (!confirmed) {
650
this._optimisticState = undefined;
651
return;
652
}
654
>
if (this._pendingActions.length === 0) {
655
>
this._optimisticState = undefined;
656
>
this._onDidChange.fire(confirmed);
657
>
return;
658
>
}
659
660
let state = confirmed;
661
for (const pending of this._pendingActions) {
662
state = changesetReducer(state, pending.action, this._log);
663
}
664
this._optimisticState = state;
665
this._onDidChange.fire(state);
668
>
669
>
type ManagedSubscription = SessionStateSubscription | ChatStateSubscription | TerminalStateSubscription | ChangesetStateSubscription | AnnotationsStateSubscription;
670
>
671
>
// --- Annotations State Subscription ------------------------------------------
672
>
673
>
interface IPendingAnnotationsAction {
674
>
readonly clientSeq: number;
675
>
readonly action: AnnotationsAction;
676
>
}
677
>
678
>
/**
679
>
* Subscription to a session's annotations channel (e.g.
680
>
* `<sessionUri>/annotations`).
681
>
*
682
>
* Annotations actions are client-dispatchable, so this subscription supports
683
>
* write-ahead reconciliation: optimistic state is layered on top of confirmed
684
>
* state and reconciled as the server echoes the client's own actions back.
685
>
*
686
>
* Like {@link ChangesetStateSubscription}, the subscription does NOT
687
>
* self-tear-down on lifecycle events; cleanup is driven externally by the
688
>
* holder releasing its `IReference`.
689
>
*/
690
>
export class AnnotationsStateSubscription extends BaseAgentSubscription<AnnotationsState> {
691
>
692
>
private readonly _pendingActions: IPendingAnnotationsAction[] = [];
693
>
private _optimisticState: AnnotationsState | undefined;
694
>
private readonly _annotationsUri: string;
695
>
private readonly _seqAllocator: () => number;
696
>
697
>
constructor(annotationsUri: string, clientId: string, seqAllocator: () => number, log: (msg: string) => void) {
698
super(clientId, log);
699
this._annotationsUri = annotationsUri;
700
this._seqAllocator = seqAllocator;
701
}
703
>
/**
704
>
* Optimistically apply an annotations action. Returns the clientSeq to
705
>
* send to the server so it can echo back for reconciliation.
706
>
*/
707
>
applyOptimistic(action: AnnotationsAction): number {
708
const clientSeq = this._seqAllocator();
709
this._pendingActions.push({ clientSeq, action });
710
const base = this._optimisticState ?? this.verifiedValue;
711
if (base) {
712
this._optimisticState = annotationsReducer(base, action, this._log);
713
this._onDidChange.fire(this._optimisticState);
714
}
715
return clientSeq;
716
}
718
>
protected override _getOptimisticState(): AnnotationsState | undefined {
719
return this._optimisticState;
720
}
722
>
protected override _applyReducer(state: AnnotationsState, action: StateAction): AnnotationsState {
723
return annotationsReducer(state, action as AnnotationsAction, this._log);
724
}
726
>
protected override _isRelevantEnvelope(envelope: ActionEnvelope): boolean {
727
return isAnnotationsAction(envelope.action) && envelope.channel === this._annotationsUri;
728
}
730
>
protected override _onSnapshotApplied(fromSeq: number): void {
731
super._onSnapshotApplied(fromSeq);
732
this._recomputeOptimistic();
733
}
735
>
protected override _reconcile(envelope: ActionEnvelope, isOwnAction: boolean): void {
736
if (isOwnAction && envelope.origin) {
737
const idx = this._pendingActions.findIndex(p => p.clientSeq === envelope.origin!.clientSeq);
738
if (idx !== -1) {
739
if (!envelope.rejectionReason) {
740
this._confirmedApply(envelope.action);
741
}
742
this._pendingActions.splice(idx, 1);
743
} else {
744
this._confirmedApply(envelope.action);
745
}
746
} else {
747
this._confirmedApply(envelope.action);
748
}
749
this._recomputeOptimistic();
750
}
752
>
private _confirmedApply(action: StateAction): void {
753
if (this._confirmedState) {
754
this._confirmedState = this._applyReducer(this._confirmedState, action);
755
}
756
}
758
>
private _recomputeOptimistic(): void {
759
const confirmed = this._confirmedState;
760
if (!confirmed) {
761
this._optimisticState = undefined;
762
return;
763
}
764
765
if (this._pendingActions.length === 0) {
766
this._optimisticState = undefined; // No pending → value falls through to confirmed
767
this._onDidChange.fire(confirmed);
768
return;
769
}
770
771
let state = confirmed;
772
for (const pending of this._pendingActions) {
773
state = annotationsReducer(state, pending.action, this._log);
774
}
775
this._optimisticState = state;
776
this._onDidChange.fire(state);
777
}
779
>
780
>
type ManagedSubscriptionEntry = { sub: ManagedSubscription; kind: StateComponents; refCount: number; holders: Map<number, string> };
781
>
782
>
// --- Subscription Manager ----------------------------------------------------
783
>
784
>
785
>
/**
786
>
* Manages the lifecycle of resource subscriptions for an agent connection.
787
>
*
788
>
* Provides refcounted access via {@link getSubscription} — the subscription
789
>
* is created on first acquire, subscribes to the server, and stays alive
790
>
* until the last reference is disposed.
791
>
*
792
>
* The connection feeds action envelopes to all active subscriptions via
793
>
* {@link receiveEnvelope}.
794
>
*/
795
>
export class AgentSubscriptionManager extends Disposable {
796
>
797
>
private readonly _subscriptions = new ResourceMap<ManagedSubscriptionEntry>();
798
>
private readonly _inflightCreates = new ResourceMap<Promise<unknown>>();
799
>
private _referenceOwnerIds = 0;
800
>
private readonly _rootState: RootStateSubscription;
801
>
private readonly _clientId: string;
802
>
private readonly _seqAllocator: () => number;
803
>
private readonly _log: (msg: string) => void;
804
>
private readonly _subscribe: (resource: URI) => Promise<IStateSnapshot>;
805
>
private readonly _unsubscribe: (resource: URI) => void;
806
>
807
>
constructor(
809
>
seqAllocator: () => number,
810
>
log: (msg: string) => void,
811
>
subscribe: (resource: URI) => Promise<IStateSnapshot>,
812
>
unsubscribe: (resource: URI) => void,
813
>
) {
814
>
super();
815
>
this._clientId = clientId;
816
>
this._seqAllocator = seqAllocator;
817
>
this._log = log;
818
>
this._subscribe = subscribe;
819
>
this._unsubscribe = unsubscribe;
820
>
this._rootState = this._register(new RootStateSubscription(clientId, log));
821
>
}
823
>
/** The always-live root state subscription. */
824
>
get rootState(): IAgentSubscription<RootState> {
826
>
}
828
>
/**
829
>
* Initialize the root state from a snapshot received during the
830
>
* connection handshake.
831
>
*/
832
>
handleRootSnapshot(state: RootState, fromSeq: number): void {
834
>
}
836
>
/**
837
>
* Returns an existing subscription without affecting its refcount.
838
>
* Returns `undefined` if no subscription is active for the given resource.
839
>
*/
840
>
getSubscriptionUnmanaged<T>(resource: URI): IAgentSubscription<T> | undefined {
842
>
return entry?.sub as IAgentSubscription<T> | undefined;
843
>
}
845
>
/**
846
>
* Returns the in-flight `createSession` Promise for this URI, or `undefined` if no create is pending. Used by
847
>
* callers that need to gate their own work on a still-running eager `createSession` (e.g. the chat handler awaits
848
>
* this before deciding whether the sessions provider's eager-create raced first send).
849
>
*/
850
>
getInflightSessionCreate(resource: URI): Promise<unknown> | undefined {
851
return this._inflightCreates.get(resource);
852
}
854
>
/**
855
>
* Register an in-flight `createSession` Promise for a session URI. Any
856
>
* subscribe issued for this resource while the create is pending waits
857
>
* for the Promise before issuing the wire-level subscribe.
858
>
*/
859
>
trackSessionCreate(resource: URI, promise: Promise<unknown>): void {
860
this._inflightCreates.set(resource, promise);
861
// This branch only observes settlement to evict the inflight entry; the
862
// `createSession` caller (and the server, via logService.error) owns the
863
// result. `finally` re-raises a rejection, so without this trailing
864
// `catch` an expected create failure (e.g. AHP_AUTH_REQUIRED) would be
865
// reported a second time as an unhandled rejection.
866
void promise.finally(() => {
867
if (this._inflightCreates.get(resource) === promise) {
868
this._inflightCreates.delete(resource);
869
}
870
}).catch(() => { });
871
}
873
>
/**
874
>
* Get or create a refcounted subscription to any resource. Disposing
875
>
* the returned reference decrements the refcount; when it reaches zero
876
>
* the subscription is torn down and the server is notified.
877
>
*
878
>
* `owner` names the caller holding the reference so inspection surfaces
879
>
* (see {@link getActiveSubscriptions}) can attribute who is retaining a
880
>
* subscription. Use a stable, human-readable identifier such as the
881
>
* acquiring class name.
882
>
*/
883
>
getSubscription<T>(kind: StateComponents, resource: URI, owner: string): IReference<IAgentSubscription<T>> {
885
>
if (existing) {
888
>
// the errored entry so this acquire performs a fresh subscribe.
889
>
this._subscriptions.delete(resource);
890
>
this._disposeSubscriptionEntry(resource, existing);
893
>
return this._acquireReference<T>(resource, existing, owner);
894
>
}
897
>
// Create new subscription based on caller-specified kind
898
>
const key = resource.toString();
899
>
const sub = this._createSubscription(kind, key);
900
>
const entry: ManagedSubscriptionEntry = { sub, kind, refCount: 1, holders: new Map() };
901
>
this._subscriptions.set(resource, entry);
902
>
903
>
// Kick off server subscription asynchronously.
904
>
// Capture the entry reference so we can validate it hasn't been
905
>
// replaced by a new subscription for the same key (race guard).
906
>
void (async () => {
907
>
const inflight = this._inflightCreates.get(resource);
908
>
if (inflight) {
909
try {
910
await inflight;
911
} catch {
912
// Swallow — fall through to subscribe so the error
913
// surfaces consistently via setError() on the
914
// subscription, matching the no-inflight path.
915
}
916
}
918
>
const snapshot = await this._subscribe(resource);
920
>
sub.handleSnapshot(snapshot.state as never, snapshot.fromSeq);
921
>
}
924
>
sub.setError(err instanceof Error ? err : new Error(String(err)));
925
>
}
926
>
}
928
>
929
>
return this._acquireReference<T>(resource, entry, owner);
930
>
}
932
>
/**
933
>
* Register `owner` as a holder of `entry` and return a reference whose
934
>
* disposal removes that holder and releases the subscription. The
935
>
* caller is responsible for the matching refcount increment (a fresh
936
>
* entry starts at 1; an existing entry is bumped before calling this).
937
>
*/
938
>
private _acquireReference<T>(resource: URI, entry: ManagedSubscriptionEntry, owner: string): IReference<IAgentSubscription<T>> {
940
>
entry.holders.set(ownerId, owner);
941
>
942
>
let isDisposed = false;
943
>
return {
944
>
object: entry.sub as unknown as IAgentSubscription<T>,
945
>
dispose: () => {
946
>
if (isDisposed) {
948
>
}
950
>
entry.holders.delete(ownerId);
951
>
this._releaseSubscription(resource, entry);
952
>
},
953
>
};
954
>
}
956
>
private _disposeSubscriptionEntry(resource: URI, entry: ManagedSubscriptionEntry): void {
958
>
if (entry.sub instanceof SessionStateSubscription || entry.sub instanceof ChatStateSubscription) {
960
>
}
962
>
}
964
>
private _tryUnsubscribe(resource: URI): void {
966
>
this._unsubscribe(resource);
967
>
} catch (error) {
968
const message = error instanceof Error ? error.message : String(error);
969
this._log(`Failed to unsubscribe ${resource.toString()}: ${message}`);
970
}
973
>
/**
974
>
* Route an incoming action envelope to all active subscriptions.
975
>
*/
976
>
receiveEnvelope(envelope: ActionEnvelope): void {
978
>
this._rootState.receiveEnvelope(envelope);
979
>
// Other subscriptions get filtered actions
980
>
for (const { sub } of this._subscriptions.values()) {
981
>
sub.receiveEnvelope(envelope);
982
>
}
983
>
}
985
>
/**
986
>
* Dispatch a client action. Applies optimistically to the relevant
987
>
* subscription if applicable, then returns the clientSeq.
988
>
*
989
>
* `channel` is the protocol URI string identifying the channel the
990
>
* action targets (a session URI for session actions, etc.).
991
>
*/
992
>
dispatchOptimistic(channel: string, action: SessionAction | ChatAction | TerminalAction | ClientChangesetAction | ClientAnnotationsAction | IRootConfigChangedAction): number {
995
>
if (entry?.sub instanceof SessionStateSubscription) {
996
>
return entry.sub.applyOptimistic(action);
997
>
}
999
const entry = this._subscriptions.get(URI.parse(channel));
1000
if (entry?.sub instanceof ChatStateSubscription) {
1001
return entry.sub.applyOptimistic(action);
1002
}
1004
>
const entry = this._subscriptions.get(URI.parse(channel));
1005
>
if (entry?.sub instanceof ChangesetStateSubscription) {
1006
>
return entry.sub.applyOptimistic(action);
1007
>
}
1008
>
} else if (isAnnotationsAction(action)) {
1009
const entry = this._subscriptions.get(URI.parse(channel));
1010
if (entry?.sub instanceof AnnotationsStateSubscription) {
1011
return entry.sub.applyOptimistic(action);
1012
}
1013
}
1014
return this._seqAllocator();
1017
>
/**
1018
>
* URIs currently subscribed to via {@link getSubscription}. Used to
1019
>
* build the `subscriptions` payload for a `reconnect` RPC so the
1020
>
* server can restore them in one round-trip.
1021
>
*
1022
>
* Does NOT include the always-live root state, which the protocol
1023
>
* client manages separately.
1024
>
*/
1025
>
currentSubscriptionUris(): URI[] {
1026
return [...this._subscriptions.keys()];
1027
}
1029
>
/**
1030
>
* Read-only descriptors of every active resource subscription, for
1031
>
* inspection/debug surfaces. Does NOT include the always-live root
1032
>
* state, which the connection exposes separately via {@link rootState}.
1033
>
*/
1034
>
getActiveSubscriptions(): readonly IActiveSubscriptionInfo[] {
1036
>
for (const [resource, entry] of this._subscriptions) {
1037
>
const value = entry.sub.value;
1038
>
const status = value === undefined ? 'pending' : value instanceof Error ? 'error' : 'snapshot';
1039
>
out.push({ resource, kind: entry.kind, refCount: entry.refCount, holders: this._summarizeHolders(entry), status });
1040
>
}
1041
>
return out;
1042
>
}
1044
>
/** Group an entry's holders by owner name, sorted by descending count. */
1045
>
private _summarizeHolders(entry: ManagedSubscriptionEntry): IActiveSubscriptionHolder[] {
1047
>
for (const owner of entry.holders.values()) {
1048
>
counts.set(owner, (counts.get(owner) ?? 0) + 1);
1049
>
}
1050
>
return [...counts.entries()]
1051
>
.map(([owner, count]) => ({ owner, count }))
1052
>
.sort((a, b) => b.count - a.count);
1053
>
}
1055
>
/**
1056
>
* Snapshot of every pending optimistic action across all session
1057
>
* subscriptions. Callers use this to replay actions after a transport
1058
>
* reconnect; entries are kept on their subscriptions until they're
1059
>
* either echoed back by the server or explicitly dropped via
1060
>
* {@link dropPendingSessionAction}.
1061
>
*/
1062
>
getPendingSessionActions(): IPendingDispatchAction[] {
1063
const out: IPendingDispatchAction[] = [];
1064
for (const { sub } of this._subscriptions.values()) {
1065
if (sub instanceof SessionStateSubscription || sub instanceof ChatStateSubscription) {
1066
out.push(...sub.getPendingActions());
1067
}
1068
}
1069
return out;
1070
}
1072
>
/**
1073
>
* Remove a single pending optimistic action for a session by its
1074
>
* `clientSeq`. Used during reconnect to evict actions the server
1075
>
* already processed (and replayed back to us) so they're not resent.
1076
>
*/
1077
>
dropPendingSessionAction(sessionUri: string, clientSeq: number): void {
1078
const entry = this._subscriptions.get(URI.parse(sessionUri));
1079
if (entry?.sub instanceof SessionStateSubscription || entry?.sub instanceof ChatStateSubscription) {
1080
entry.sub.dropPendingByClientSeq(clientSeq);
1081
}
1082
}
1084
>
/**
1085
>
* Apply a fresh snapshot to a subscribed resource — used when the server
1086
>
* responds to a `reconnect` request with `type: 'snapshot'` because the
1087
>
* replay buffer no longer covers the client's gap. Routes to the root
1088
>
* subscription when {@link ROOT_STATE_URI} matches, otherwise reseats the
1089
>
* matching entry in {@link _subscriptions}. Unknown resources are ignored.
1090
>
*/
1091
>
applyReconnectSnapshot(resource: string, state: unknown, fromSeq: number): void {
1092
if (isAhpRootChannel(resource)) {
1093
this._rootState.handleSnapshot(state as RootState, fromSeq);
1094
return;
1095
}
1096
const entry = this._subscriptions.get(URI.parse(resource));
1097
if (!entry) {
1098
return;
1099
}
1100
// Clear any pending optimistic actions before reseating confirmed
1101
// state \u2014 they were predicated on the pre-disconnect confirmed
1102
// state and won't reconcile correctly against a fresh snapshot.
1103
if (entry.sub instanceof SessionStateSubscription || entry.sub instanceof ChatStateSubscription) {
1104
entry.sub.clearPending();
1105
}
1106
entry.sub.handleSnapshot(state as never, fromSeq);
1107
}
1109
>
/**
1110
>
* Mark a set of subscriptions as no longer resumable on the server
1111
>
* (reported via `ReconnectReplayResult.missing`). The subscriptions
1112
>
* themselves stay alive so consumers continue to hold valid references,
1113
>
* but their value transitions to an `Error` until they're recreated.
1114
>
*/
1115
>
markSubscriptionsMissing(missing: readonly URI[]): void {
1116
for (const resource of missing) {
1117
const entry = this._subscriptions.get(resource);
1118
if (entry) {
1119
if (entry.sub instanceof SessionStateSubscription || entry.sub instanceof ChatStateSubscription) {
1120
entry.sub.clearPending();
1121
}
1122
entry.sub.setError(new Error(`Subscription no longer available after reconnect: ${resource.toString()}`));
1123
}
1124
}
1125
}
1127
>
private _createSubscription(kind: StateComponents, key: string): ManagedSubscription {
1129
>
case StateComponents.Session:
1130
>
return new SessionStateSubscription(key, this._clientId, this._seqAllocator, this._log);
agentSubscription.ts ×1
1132
return new ChatStateSubscription(key, this._clientId, this._seqAllocator, this._log);
1134
>
return new TerminalStateSubscription(key, this._clientId, this._log);
agentSubscription.ts ×1
1136
>
return new ChangesetStateSubscription(key, this._clientId, this._seqAllocator, this._log);
agentSubscription.ts ×2
1138
return new AnnotationsStateSubscription(key, this._clientId, this._seqAllocator, this._log);
1140
throw new Error('_createSubscription: root subscription is managed separately');
1142
assertNever(kind, `_createSubscription: unsupported StateComponents kind: ${kind}`);
1144
>
}
1146
>
private _releaseSubscription(resource: URI, expected?: ManagedSubscriptionEntry): void {
1148
>
// A failed subscription can be evicted and replaced while old references
1149
>
// still exist; stale disposals must not release the replacement entry.
1150
>
if (!entry || (expected && entry !== expected)) {
1152
>
}
1154
>
if (entry.refCount <= 0) {
1155
>
this._subscriptions.delete(resource);
1156
>
this._disposeSubscriptionEntry(resource, entry);
1157
>
}
1160
>
override dispose(): void {
1163
>
entry.sub.dispose();
1164
>
}
1166
>
super.dispose();
1167
>
}
1169
>
1170
>
/** Returns whether an action envelope targets one of the subscribed channel URIs. */
1171
>
export function isActionEnvelopeRelevantToSubscriptionUris(envelope: ActionEnvelope, subscribedUris: Iterable<string>): boolean {
1173
>
for (const uri of subscribedUris) {
1174
>
if (isAhpRootChannel(uri)) {
1176
>
}
1179
>
}
1181
>
if (uri === envelope.channel) {
1182
>
return true;
1183
>
}
1184
>
}
1185
>
return false;
1186
>
}
1188
>
// --- Observable Adapter ------------------------------------------------------
1189
>
1190
>
/**
1191
>
* Adapts an {@link IAgentSubscription} into an {@link IObservable} of the
1192
>
* subscription's value. Errors and the pre-snapshot phase are surfaced as
1193
>
* `undefined`; consumers that need the error itself should read
1194
>
* {@link IAgentSubscription.value} directly.
1195
>
*/
1196
>
export function observableFromSubscription<T>(owner: object | undefined, sub: IAgentSubscription<T>): IObservable<T | undefined> {
1197
return observableFromEvent(owner, sub.onDidChange, () => {
1198
const v = sub.value;
1199
return v instanceof Error ? undefined : v;
1200
});
1201
}