diff --git a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts index 9fedd760188282..4ece652ca710b4 100644 --- a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts +++ b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts @@ -2627,6 +2627,7 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement * the provider itself is disposed. */ private readonly _newSessions = this._register(new DisposableMap()); + private readonly _firstSendModelReferences = this._register(new DisposableMap()); /** The in-flight new session with the given id, if any. */ protected _getNewSession(sessionId: string): NewSession | undefined { @@ -4446,8 +4447,7 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement cached.setChatModelId(chat.resource, selectedModelId, ChatModelSource.CarriedOver); cached.setChatAgent(chat.resource, selectedAgentUri ? { uri: selectedAgentUri, name: '' } : undefined); - await this._chatSessionsService.getOrCreateChatSession(chat.resource, CancellationToken.None); - await this._updateChatSessionState(chat.resource, selectedModelId, selectedAgentUri); + await this._prepareFirstSendChatModel(chat.resource, selectedModelId, selectedAgentUri); return chat; } @@ -4510,7 +4510,8 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement hideFromTranscript: options.hideFromTranscript, }; - const modelRef = await this._chatService.acquireOrLoadSession(chatResource, ChatAgentLocation.Chat, CancellationToken.None); + const modelRef = this._firstSendModelReferences.deleteAndLeak(chatResource.toString()) + ?? await this._chatService.acquireOrLoadSession(chatResource, ChatAgentLocation.Chat, CancellationToken.None); if (!modelRef) { throw new Error(`[${this.id}] Unable to load chat session ${chatResource.toString()}`); } @@ -4534,6 +4535,15 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement return cached; } + private async _prepareFirstSendChatModel(chatResource: URI, modelId: string | undefined, agentUri: string | undefined): Promise { + const modelRef = await this._chatService.acquireOrLoadSession(chatResource, ChatAgentLocation.Chat, CancellationToken.None); + if (!modelRef) { + return; + } + this._applyChatSessionState(modelRef, modelId, agentUri); + this._firstSendModelReferences.set(chatResource.toString(), modelRef); + } + private async _updateChatSessionState(chatResource: URI, modelId: string | undefined, agentUri: string | undefined, options?: { readonly clearDraft?: boolean }): Promise { const modelRef = await this._chatService.acquireOrLoadSession(chatResource, ChatAgentLocation.Chat, CancellationToken.None); if (!modelRef) { diff --git a/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts b/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts index 285b890295ca89..10b1584e187343 100644 --- a/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts +++ b/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts @@ -4884,6 +4884,49 @@ suite('LocalAgentHostSessionsProvider', () => { }); })); + test('createSideChat retains its prepared model through the first send', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + agentHost.setAgents([{ provider: 'copilotcli', displayName: 'Copilot', description: '', models: [], capabilities: { multipleChats: { fork: true, sideChat: true } } } as AgentInfo]); + let acquireCount = 0; + let disposeCount = 0; + const provider = createProvider(disposables, agentHost, undefined, { + acquireOrLoadSession: async () => { + acquireCount++; + const inputModel = new class extends mock() { + override readonly state = constObservable(undefined); + override setState(): void { } + override clearState(): void { } + override toJSON(): undefined { return undefined; } + }(); + return { + object: new class extends mock() { + override readonly inputModel = inputModel; + }(), + dispose: () => { disposeCount++; }, + }; + }, + }); + const session = setupMultiChatSession(provider, 'retained-side-chat'); + const sessionUri = AgentSession.uri('copilotcli', 'retained-side-chat').toString(); + const defaultChat = buildDefaultChatUri(sessionUri); + agentHost.setSessionState('retained-side-chat', 'copilotcli', makeState([ + makeChatSummary(defaultChat, ''), + ], { defaultChat })); + + const sideChat = await provider.createSideChat(session.sessionId, session.resource, 'turn-1'); + const disposedBeforeSend = disposeCount; + await provider.sendRequest(session.sessionId, sideChat.resource, { query: 'Side question' }); + + assert.deepStrictEqual({ + acquireCount, + disposedBeforeSend, + disposeCount, + }, { + acquireCount: 1, + disposedBeforeSend: 0, + disposeCount: 1, + }); + })); + test('createSideChat rejects when the session capability is not advertised', async () => { const provider = createProvider(disposables, agentHost); const session = setupMultiChatSession(provider, 'multi-side-chat-unsupported');