agentHost: retain side chat through first send

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<string, NewSession>());
+	private readonly _firstSendModelReferences = this._register(new DisposableMap<string, IChatModelReference>());
 
 	/** 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<void> {
+		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<void> {
 		const modelRef = await this._chatService.acquireOrLoadSession(chatResource, ChatAgentLocation.Chat, CancellationToken.None);
 		if (!modelRef) {

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<void>({ 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<IInputModel>() {
+						override readonly state = constObservable<IChatModelInputState | undefined>(undefined);
+						override setState(): void { }
+						override clearState(): void { }
+						override toJSON(): undefined { return undefined; }
+					}();
+					return {
+						object: new class extends mock<IChatModel>() {
+							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');