From 59f7592b5d1cb4e765236be453bae07de2b6f2af Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Thu, 14 May 2026 14:25:31 -0700 Subject: [PATCH] chat: recover failed agent host subscriptions (#316499) Avoid reusing a session subscription that previously failed while the backend session was still being created. - Drops cached Agent Host session subscriptions whose current value is an Error before subscribing again. - Adds a regression test for the first-send path when a provisional session has a stale failed subscription. - Lets the existing create-then-subscribe flow recover and sync the initial pending message normally. (Commit message generated by Copilot) --- .../agentHost/agentHostSessionHandler.ts | 5 ++ .../agentHostChatContribution.test.ts | 52 +++++++++++++++++-- 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index 97ac49f46fc..2ea7264b0c6 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -2832,6 +2832,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC */ private _ensureSessionSubscription(sessionUri: string): IAgentSubscription { let ref = this._sessionSubscriptions.get(sessionUri); + if (ref?.object.value instanceof Error) { + this._sessionSubscriptions.delete(sessionUri); + ref.dispose(); + ref = undefined; + } if (!ref) { ref = this._config.connection.getSubscription(StateComponents.Session, URI.parse(sessionUri)); this._sessionSubscriptions.set(sessionUri, ref); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts index 9647c0d5ed8..0eaec4380fa 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts @@ -90,6 +90,7 @@ class MockAgentHostService extends mock() { private readonly _sessions = new Map(); public createSessionCalls: IAgentCreateSessionConfig[] = []; public disposedSessions: URI[] = []; + public failNextSubscriptionFor = new Set(); public agents = [{ provider: 'copilot' as const, displayName: 'Agent Host - Copilot', description: 'test', requiresAuth: true }]; /** @@ -221,6 +222,24 @@ class MockAgentHostService extends mock() { const onWillApply = new Emitter(); const onDidApply = new Emitter(); + if (this.failNextSubscriptionFor.delete(resourceStr)) { + const error = new Error(`Session not found on backend: ${resourceStr}`); + return { + object: { + get value() { return error; }, + get verifiedValue() { return undefined; }, + onDidChange: emitter.event, + onWillApplyAction: onWillApply.event, + onDidApplyAction: onDidApply.event, + }, + dispose: () => { + emitter.dispose(); + onWillApply.dispose(); + onDidApply.dispose(); + }, + }; + } + // Hydrate synchronously with a default state const existingState = this.sessionStates.get(resourceStr); let initialState: SessionState; @@ -355,7 +374,7 @@ class MockChatWidgetService extends mock() { // ---- Helpers ---------------------------------------------------------------- -function createTestServices(disposables: DisposableStore, workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }, authServiceOverride?: Partial, languageModels?: ReadonlyMap) { +function createTestServices(disposables: DisposableStore, workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }, authServiceOverride?: Partial, languageModels?: ReadonlyMap, provisionalServiceOverride?: Partial) { const instantiationService = disposables.add(new TestInstantiationService()); const agentHostService = new MockAgentHostService(); @@ -471,14 +490,15 @@ function createTestServices(disposables: DisposableStore, workingDirectoryResolv getOrCreate: async () => undefined, tryRebind: async () => undefined, disposeSession: async () => { }, + ...provisionalServiceOverride, } as Partial as IAgentHostUntitledProvisionalSessionService); instantiationService.stub(IOpenerService, openerService as IOpenerService); return { instantiationService, agentHostService, chatAgentService, chatWidgetService, chatService, openerService }; } -function createContribution(disposables: DisposableStore, opts?: { authServiceOverride?: Partial; workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }; languageModels?: ReadonlyMap }) { - const { instantiationService, agentHostService, chatAgentService, chatWidgetService, chatService, openerService } = createTestServices(disposables, opts?.workingDirectoryResolver, opts?.authServiceOverride, opts?.languageModels); +function createContribution(disposables: DisposableStore, opts?: { authServiceOverride?: Partial; workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }; languageModels?: ReadonlyMap; provisionalServiceOverride?: Partial }) { + const { instantiationService, agentHostService, chatAgentService, chatWidgetService, chatService, openerService } = createTestServices(disposables, opts?.workingDirectoryResolver, opts?.authServiceOverride, opts?.languageModels, opts?.provisionalServiceOverride); const listController = disposables.add(instantiationService.createInstance(AgentHostSessionListController, 'agent-host-copilot', 'copilot', agentHostService, undefined, 'local')); const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { @@ -977,6 +997,32 @@ suite('AgentHostChatContribution', () => { assert.strictEqual(AgentSession.id(URI.parse(session)), 'existing-session-42'); })); + test('recovers from stale failed subscription before first send', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const sessionResource = URI.from({ scheme: 'agent-host-copilot', path: '/existing-subscribe-retry' }); + const backendSession = AgentSession.uri('copilot', 'existing-subscribe-retry'); + const { agentHostService, chatAgentService } = createContribution(disposables, { + provisionalServiceOverride: { + get: resource => resource.toString() === sessionResource.toString() ? backendSession : undefined, + }, + }); + agentHostService.failNextSubscriptionFor.add(backendSession.toString()); + + const registered = chatAgentService.registeredAgents.get('agent-host-copilot')!; + const turnPromise = registered.impl.invoke( + makeRequest({ message: 'Recovered', sessionResource }), + () => { }, [], CancellationToken.None, + ); + await timeout(10); + + const dispatch = agentHostService.turnActions[0]; + const action = dispatch.action as ITurnStartedAction; + agentHostService.fireAction({ action: dispatch.action, serverSeq: 1, origin: { clientId: agentHostService.clientId, clientSeq: dispatch.clientSeq } }); + agentHostService.fireAction({ action: { type: 'session/turnComplete', session: action.session, turnId: action.turnId } as SessionAction, serverSeq: 2, origin: undefined }); + await turnPromise; + + assert.deepStrictEqual(agentHostService.turnActions.map(d => (d.action as ITurnStartedAction).userMessage.text), ['Recovered']); + })); + test('rejects generic contributed-chat untitled resource', async () => { const { sessionHandler, agentHostService, chatAgentService } = createContribution(disposables);