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)
This commit is contained in:
Connor Peet
2026-05-14 14:25:31 -07:00
committed by GitHub
parent ce2ab59a8a
commit 59f7592b5d
2 changed files with 54 additions and 3 deletions

View File

@@ -2832,6 +2832,11 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
*/
private _ensureSessionSubscription(sessionUri: string): IAgentSubscription<SessionState> {
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);

View File

@@ -90,6 +90,7 @@ class MockAgentHostService extends mock<IAgentHostService>() {
private readonly _sessions = new Map<string, IAgentSessionMetadata>();
public createSessionCalls: IAgentCreateSessionConfig[] = [];
public disposedSessions: URI[] = [];
public failNextSubscriptionFor = new Set<string>();
public agents = [{ provider: 'copilot' as const, displayName: 'Agent Host - Copilot', description: 'test', requiresAuth: true }];
/**
@@ -221,6 +222,24 @@ class MockAgentHostService extends mock<IAgentHostService>() {
const onWillApply = new Emitter<ActionEnvelope>();
const onDidApply = new Emitter<ActionEnvelope>();
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<IChatWidgetService>() {
// ---- Helpers ----------------------------------------------------------------
function createTestServices(disposables: DisposableStore, workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }, authServiceOverride?: Partial<IAuthenticationService>, languageModels?: ReadonlyMap<string, ILanguageModelChatMetadata>) {
function createTestServices(disposables: DisposableStore, workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }, authServiceOverride?: Partial<IAuthenticationService>, languageModels?: ReadonlyMap<string, ILanguageModelChatMetadata>, provisionalServiceOverride?: Partial<IAgentHostUntitledProvisionalSessionService>) {
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<IAgentHostUntitledProvisionalSessionService> as IAgentHostUntitledProvisionalSessionService);
instantiationService.stub(IOpenerService, openerService as IOpenerService);
return { instantiationService, agentHostService, chatAgentService, chatWidgetService, chatService, openerService };
}
function createContribution(disposables: DisposableStore, opts?: { authServiceOverride?: Partial<IAuthenticationService>; workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }; languageModels?: ReadonlyMap<string, ILanguageModelChatMetadata> }) {
const { instantiationService, agentHostService, chatAgentService, chatWidgetService, chatService, openerService } = createTestServices(disposables, opts?.workingDirectoryResolver, opts?.authServiceOverride, opts?.languageModels);
function createContribution(disposables: DisposableStore, opts?: { authServiceOverride?: Partial<IAuthenticationService>; workingDirectoryResolver?: { resolve(sessionResource: URI): URI | undefined; isNewSession?: (sessionResource: URI) => boolean }; languageModels?: ReadonlyMap<string, ILanguageModelChatMetadata>; provisionalServiceOverride?: Partial<IAgentHostUntitledProvisionalSessionService> }) {
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);