Microsoft/vscode
1
0
Fork 0
mirror of https://github.com/microsoft/vscode.git synced 2026-08-27 12:56:39 +01:00
Code Issues Packages Projects Releases Wiki Activity

agentHost: separate subscription resources and channels

Browse Source 
Keep URI-based subscription APIs narrow while preserving exact AHP catalogue channels. Mark failed reconnect restorations by channel and cover the exact-channel path with regressions.
This commit is contained in:
Ben Villalobos
2026-08-24 16:43:45 -07:00
parent ae94bb8ada
commit 5ebf08a3f5
4 changed files with 111 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/vs/platform/agentHost/browser/agentHostProtocolClient.ts
+3 -3
View File
@@ -759,8 +759,8 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect
if (error instanceof ProtocolError && error.code === AHP_CLIENT_CONNECTION_CLOSED) {
throw error;
}
this._logService.warn(`[AgentHostProtocolClient] Failed to restore subscription ${subscription.resource.toString()} after host restart: ${error instanceof Error ? error.message : String(error)}`);
this._subscriptionManager.markSubscriptionsMissing([subscription.resource]);
this._logService.warn(`[AgentHostProtocolClient] Failed to restore subscription ${subscription.channel} after host restart: ${error instanceof Error ? error.message : String(error)}`);
this._subscriptionManager.markSubscriptionsMissing([subscription.channel]);
}
}));
};
@@ -939,7 +939,7 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect
}
getSubscriptionByChannel<T>(kind: StateComponents, channel: string, owner: string): IReference<IAgentSubscription<T>> {
return this._subscriptionManager.getSubscription<T>(kind, channel, owner);
return this._subscriptionManager.getSubscriptionByChannel<T>(kind, channel, owner);
}
getSubscriptionUnmanaged<T>(_kind: StateComponents, resource: URI): IAgentSubscription<T> | undefined {
src/vs/platform/agentHost/common/state/agentSubscription.ts
+32 -17
View File
@@ -927,7 +927,7 @@ export class AgentSubscriptionManager extends Disposable {
* Returns an existing subscription without affecting its refcount.
* Returns `undefined` if no subscription is active for the given resource.
*/
getSubscriptionUnmanaged<T>(resource: URI | string): IAgentSubscription<T> | undefined {
getSubscriptionUnmanaged<T>(resource: URI): IAgentSubscription<T> | undefined {
const entry = this._subscriptions.get(this._subscriptionResource(resource).key);
return entry?.sub as IAgentSubscription<T> | undefined;
}
@@ -970,8 +970,16 @@ export class AgentSubscriptionManager extends Disposable {
* subscription. Use a stable, human-readable identifier such as the
* acquiring class name.
*/
getSubscription<T>(kind: StateComponents, resource: URI | string, owner: string): IReference<IAgentSubscription<T>> {
const resolved = this._subscriptionResource(resource);
getSubscription<T>(kind: StateComponents, resource: URI, owner: string): IReference<IAgentSubscription<T>> {
return this._getSubscription(kind, this._subscriptionResource(resource), owner);
}
/** Get or create a subscription using an exact protocol channel string. */
getSubscriptionByChannel<T>(kind: StateComponents, channel: string, owner: string): IReference<IAgentSubscription<T>> {
return this._getSubscription(kind, this._subscriptionChannel(channel), owner);
}
private _getSubscription<T>(kind: StateComponents, resolved: Pick<ManagedSubscriptionEntry, 'resource' | 'channel' | 'key'>, owner: string): IReference<IAgentSubscription<T>> {
const existing = this._subscriptions.get(resolved.key);
if (existing) {
if (existing.sub.value instanceof Error) {
@@ -1081,22 +1089,22 @@ export class AgentSubscriptionManager extends Disposable {
*/
dispatchOptimistic(channel: string, action: SessionAction | ChatAction | TerminalAction | ClientChangesetAction | ClientAnnotationsAction | ClientAutomationAction | ClientAutomationRunAction | IRootConfigChangedAction): number {
if (isSessionAction(action)) {
const entry = this._subscriptions.get(this._subscriptionResource(channel).key);
const entry = this._subscriptions.get(this._subscriptionChannel(channel).key);
if (entry?.sub instanceof SessionStateSubscription) {
return entry.sub.applyOptimistic(action);
}
} else if (isChatAction(action)) {
const entry = this._subscriptions.get(this._subscriptionResource(channel).key);
const entry = this._subscriptions.get(this._subscriptionChannel(channel).key);
if (entry?.sub instanceof ChatStateSubscription) {
return entry.sub.applyOptimistic(action);
}
} else if (isChangesetAction(action)) {
const entry = this._subscriptions.get(this._subscriptionResource(channel).key);
const entry = this._subscriptions.get(this._subscriptionChannel(channel).key);
if (entry?.sub instanceof ChangesetStateSubscription) {
return entry.sub.applyOptimistic(action);
}
} else if (isAnnotationsAction(action)) {
const entry = this._subscriptions.get(this._subscriptionResource(channel).key);
const entry = this._subscriptions.get(this._subscriptionChannel(channel).key);
if (entry?.sub instanceof AnnotationsStateSubscription) {
return entry.sub.applyOptimistic(action);
}
@@ -1165,7 +1173,7 @@ export class AgentSubscriptionManager extends Disposable {
* already processed (and replayed back to us) so they're not resent.
*/
dropPendingAction(resource: string, clientSeq: number): void {
const entry = this._subscriptions.get(this._subscriptionResource(resource).key);
const entry = this._subscriptions.get(this._subscriptionChannel(resource).key);
if (entry?.sub instanceof SessionStateSubscription || entry?.sub instanceof ChatStateSubscription || entry?.sub instanceof AnnotationsStateSubscription) {
entry.sub.dropPendingByClientSeq(clientSeq);
}
@@ -1183,7 +1191,7 @@ export class AgentSubscriptionManager extends Disposable {
this._rootState.handleSnapshot(state as RootState, fromSeq);
return;
}
const entry = this._subscriptions.get(this._subscriptionResource(resource).key);
const entry = this._subscriptions.get(this._subscriptionChannel(resource).key);
if (!entry) {
return;
}
@@ -1202,9 +1210,9 @@ export class AgentSubscriptionManager extends Disposable {
* themselves stay alive so consumers continue to hold valid references,
* but their value transitions to an `Error` until they're recreated.
*/
markSubscriptionsMissing(missing: readonly (URI | string)[]): void {
for (const resource of missing) {
const entry = this._subscriptions.get(this._subscriptionResource(resource).key);
markSubscriptionsMissing(missing: readonly string[]): void {
for (const channel of missing) {
const entry = this._subscriptions.get(this._subscriptionChannel(channel).key);
if (entry) {
if (entry.sub instanceof SessionStateSubscription || entry.sub instanceof ChatStateSubscription || entry.sub instanceof AnnotationsStateSubscription) {
entry.sub.clearPending();
@@ -1260,13 +1268,20 @@ export class AgentSubscriptionManager extends Disposable {
super.dispose();
}
private _subscriptionResource(resource: URI | string): Pick<ManagedSubscriptionEntry, 'resource' | 'channel' | 'key'> {
const channel = typeof resource === 'string' ? resource : resource.toString();
const uri = typeof resource === 'string' ? URI.parse(resource) : resource;
private _subscriptionResource(resource: URI): Pick<ManagedSubscriptionEntry, 'resource' | 'channel' | 'key'> {
return {
resource,
channel: resource.toString(),
key: `resource:${getComparisonKey(resource)}`,
};
}
private _subscriptionChannel(channel: string): Pick<ManagedSubscriptionEntry, 'resource' | 'channel' | 'key'> {
const resource = URI.parse(channel);
const key = isAhpAutomationCatalogChannel(channel)
? `channel:${channel}`
: `resource:${getComparisonKey(uri)}`;
return { resource: uri, channel, key };
: `resource:${getComparisonKey(resource)}`;
return { resource, channel, key };
}
}
src/vs/platform/agentHost/test/common/agentSubscription.test.ts
+23 -2
View File
@@ -930,7 +930,7 @@ suite('AgentSubscriptionManager', () => {
subscribedResources.push(channel);
return { resource: channel, state: { automations: [] }, fromSeq: 0 };
});
const ref = mgr.getSubscription<AutomationCatalogState>(StateComponents.AutomationCatalog, AUTOMATION_CATALOG_URI, 'AutomationHolder');
const ref = mgr.getSubscriptionByChannel<AutomationCatalogState>(StateComponents.AutomationCatalog, AUTOMATION_CATALOG_URI, 'AutomationHolder');
await Event.toPromise(ref.object.onDidChange);
assert.deepStrictEqual({
@@ -1184,13 +1184,34 @@ suite('AgentSubscriptionManager', () => {
mgr.dispatchOptimistic(sessionUri, { type: ActionType.SessionWorkingDirectorySet, directory: 'file:///ws2' });
mgr.markSubscriptionsMissing([URI.parse(sessionUri)]);
mgr.markSubscriptionsMissing([sessionUri]);
assert.ok(ref.object.value instanceof Error);
assert.deepStrictEqual(mgr.getPendingActions(), []);
ref.dispose();
});
test('markSubscriptionsMissing preserves exact protocol channels', async () => {
const mgr = createManager(async channel => ({
resource: channel,
state: { automations: [] },
fromSeq: 0,
}));
const ref = mgr.getSubscriptionByChannel<AutomationCatalogState>(StateComponents.AutomationCatalog, AUTOMATION_CATALOG_URI, 'test');
await Event.toPromise(ref.object.onDidChange);
mgr.markSubscriptionsMissing([AUTOMATION_CATALOG_URI]);
assert.deepStrictEqual({
valueIsError: ref.object.value instanceof Error,
channel: mgr.getActiveSubscriptions()[0].channel,
}, {
valueIsError: true,
channel: AUTOMATION_CATALOG_URI,
});
ref.dispose();
});
test('fresh reconnect snapshots preserve pending annotation actions for replay', async () => {
const mgr = createManager();
const annotationsUri = buildAnnotationsUri(sessionUri);
src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts
+53 -1
View File
@@ -27,7 +27,7 @@ import { ActionType, type ChatTurnStartedAction, type SessionActiveClientSetActi
import { ProtocolError, type AhpServerNotification, type JsonRpcNotification, type JsonRpcRequest, type JsonRpcResponse, type ProtocolMessage } from '../../common/state/sessionProtocol.js';
import { hasKey } from '../../../../base/common/types.js';
import { mainWindow } from '../../../../base/browser/window.js';
import { CustomizationType, MessageAttachmentKind, MessageKind, PendingMessageKind, readSessionExternal, readSessionWorkspaceless, ROOT_STATE_URI, SessionStatus, StateComponents, customizationId, withSessionExternal, withSessionWorkspaceless } from '../../common/state/sessionState.js';
import { AUTOMATION_CATALOG_URI, CustomizationType, MessageAttachmentKind, MessageKind, PendingMessageKind, readSessionExternal, readSessionWorkspaceless, ROOT_STATE_URI, SessionStatus, StateComponents, customizationId, withSessionExternal, withSessionWorkspaceless } from '../../common/state/sessionState.js';
import { NonReconnectableTransportError, type IClientTransport, type IProtocolTransport } from '../../common/state/sessionTransport.js';
import { TestConfigurationService } from '../../../configuration/test/common/testConfigurationService.js';
import { ITelemetryService, TelemetryConfiguration, TelemetryLevel, TELEMETRY_SETTING_ID } from '../../../telemetry/common/telemetry.js';
@@ -2472,6 +2472,58 @@ suite('AgentHostProtocolClient', () => {
client.dispose();
});
test('marks an exact-channel subscription missing when restore fails', async function () {
this.timeout(10_000);
const { client, transports } = createFactoryClient();
const connectPromise = client.connect();
await completeHandshake(transports[0], connectPromise);
const catalogRef = client.getSubscriptionByChannel(StateComponents.AutomationCatalog, AUTOMATION_CATALOG_URI, 'test');
const initialSubscribe = await waitForRequest(transports[0], 'subscribe');
transports[0].fireMessage({
jsonrpc: '2.0', id: initialSubscribe.id,
result: { snapshot: { resource: AUTOMATION_CATALOG_URI, state: { automations: [] }, fromSeq: 5 } },
});
await flushMicrotasks();
transports[0].fireClose();
await waitForReconnecting(client);
const reconnectTransport = await waitForTransport(transports, 1);
reconnectTransport.connectDeferred.complete();
const reconnect = await waitForRequest(reconnectTransport, 'reconnect');
reconnectTransport.fireMessage({
jsonrpc: '2.0', id: reconnect.id,
error: { code: AhpErrorCodes.NotFound, message: 'Reconnect client not found' },
});
const initialize = await waitForRequest(reconnectTransport, 'initialize');
reconnectTransport.fireMessage({
jsonrpc: '2.0', id: initialize.id,
result: {
protocolVersion: PROTOCOL_VERSION,
serverSeq: 0,
snapshots: [{ resource: ROOT_STATE_URI, state: { agents: [], activeSessions: 0 }, fromSeq: 0 }],
},
});
const restoredSubscribe = await waitForRequest(reconnectTransport, 'subscribe');
reconnectTransport.fireMessage({
jsonrpc: '2.0', id: restoredSubscribe.id,
error: { code: JsonRpcErrorCodes.InternalError, message: 'Catalogue unavailable' },
});
await flushMicrotasks();
assert.deepStrictEqual({
channel: (restoredSubscribe.params as { channel: string }).channel,
valueIsError: catalogRef.object.value instanceof Error,
}, {
channel: AUTOMATION_CATALOG_URI,
valueIsError: true,
});
catalogRef.dispose();
client.dispose();
});
test('replays pending optimistic actions after reconnect', async function () {
this.timeout(10_000);
return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => {