diff --git a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts index bb09aae9aab..695ec6c3035 100644 --- a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts @@ -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(kind: StateComponents, channel: string, owner: string): IReference> { - return this._subscriptionManager.getSubscription(kind, channel, owner); + return this._subscriptionManager.getSubscriptionByChannel(kind, channel, owner); } getSubscriptionUnmanaged(_kind: StateComponents, resource: URI): IAgentSubscription | undefined { diff --git a/src/vs/platform/agentHost/common/state/agentSubscription.ts b/src/vs/platform/agentHost/common/state/agentSubscription.ts index 5ef0ee99912..ec93dce9f74 100644 --- a/src/vs/platform/agentHost/common/state/agentSubscription.ts +++ b/src/vs/platform/agentHost/common/state/agentSubscription.ts @@ -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(resource: URI | string): IAgentSubscription | undefined { + getSubscriptionUnmanaged(resource: URI): IAgentSubscription | undefined { const entry = this._subscriptions.get(this._subscriptionResource(resource).key); return entry?.sub as IAgentSubscription | undefined; } @@ -970,8 +970,16 @@ export class AgentSubscriptionManager extends Disposable { * subscription. Use a stable, human-readable identifier such as the * acquiring class name. */ - getSubscription(kind: StateComponents, resource: URI | string, owner: string): IReference> { - const resolved = this._subscriptionResource(resource); + getSubscription(kind: StateComponents, resource: URI, owner: string): IReference> { + return this._getSubscription(kind, this._subscriptionResource(resource), owner); + } + + /** Get or create a subscription using an exact protocol channel string. */ + getSubscriptionByChannel(kind: StateComponents, channel: string, owner: string): IReference> { + return this._getSubscription(kind, this._subscriptionChannel(channel), owner); + } + + private _getSubscription(kind: StateComponents, resolved: Pick, owner: string): IReference> { 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 { - const channel = typeof resource === 'string' ? resource : resource.toString(); - const uri = typeof resource === 'string' ? URI.parse(resource) : resource; + private _subscriptionResource(resource: URI): Pick { + return { + resource, + channel: resource.toString(), + key: `resource:${getComparisonKey(resource)}`, + }; + } + + private _subscriptionChannel(channel: string): Pick { + 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 }; } } diff --git a/src/vs/platform/agentHost/test/common/agentSubscription.test.ts b/src/vs/platform/agentHost/test/common/agentSubscription.test.ts index d830b14b574..eaa0f98c891 100644 --- a/src/vs/platform/agentHost/test/common/agentSubscription.test.ts +++ b/src/vs/platform/agentHost/test/common/agentSubscription.test.ts @@ -930,7 +930,7 @@ suite('AgentSubscriptionManager', () => { subscribedResources.push(channel); return { resource: channel, state: { automations: [] }, fromSeq: 0 }; }); - const ref = mgr.getSubscription(StateComponents.AutomationCatalog, AUTOMATION_CATALOG_URI, 'AutomationHolder'); + const ref = mgr.getSubscriptionByChannel(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(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); diff --git a/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts b/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts index 231557b2244..2664a4e1c68 100644 --- a/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts @@ -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 () => {