From 0ebfc317d9d9e28b34caa4f38409e66ee4d84faa Mon Sep 17 00:00:00 2001
From: Sandeep Somavarapu <sasomava@microsoft.com>
Date: Tue, 1 Sep 2026 19:03:51 +0200
Subject: [PATCH] [cherry-pick] sessions: Refresh chat status after background
 completion (#333808)

sessions: Refresh chat status after background completion

Reacquire the owning session state when a chat turn completes so main and hidden chat statuses do not remain stale after the idle subscription expires.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../browser/baseAgentHostSessionsProvider.ts  | 12 +++++
 .../localAgentHostSessionsProvider.test.ts    | 46 +++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts
index cbe9bcb0250..828b0f07a77 100644
--- a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts
+++ b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts
@@ -4934,6 +4934,17 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement
 		);
 	}
 
+	private _keepChatSessionStateAlive(chatChannel: string): void {
+		const parsedChat = parseChatUri(chatChannel);
+		if (!parsedChat) {
+			return;
+		}
+		const cached = this._sessionCache.get(AgentSession.id(parsedChat.session));
+		if (cached) {
+			this._keepSessionStateAlive(cached.sessionId);
+		}
+	}
+
 	/**
 	 * Lazily acquire a session-state subscription for `sessionId` so that
 	 * `_runningSessionConfigs` is seeded from the AHP `SessionState.config`
@@ -5625,6 +5636,7 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement
 				return;
 			}
 			if (e.action.type === ActionType.ChatTurnComplete && isChatAction(e.action)) {
+				this._keepChatSessionStateAlive(e.channel);
 				this._refreshSessions();
 			} else if (e.action.type === ActionType.SessionTitleChanged && isSessionAction(e.action)) {
 				this._handleTitleChanged(e.channel, e.action.title);
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 6fe855b413d..867d8484262 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
@@ -5467,6 +5467,52 @@ suite('LocalAgentHostSessionsProvider', () => {
 				defaultChatTitle: 'Renamed Default',
 			});
 		});
+
+		test('turn completion refreshes main chat status after the session-state subscription expires', () => runWithFakedTimers<void>({ useFakeTimers: true }, async () => {
+			const provider = createProvider(disposables, agentHost);
+			const session = setupMultiChatSession(provider, 'multi-complete');
+			const sessionUri = AgentSession.uri('copilotcli', 'multi-complete').toString();
+			const defaultChat = buildDefaultChatUri(sessionUri);
+			const subagentChat = buildSubagentChatUri(sessionUri, 'tc-1');
+			const stateWithStatus = (status: ProtocolSessionStatus): SessionState => makeState([
+				makeChatSummary(defaultChat, '', status),
+				{
+					...makeChatSummary(subagentChat, 'Reviewer', status),
+					origin: { kind: ProtocolChatOriginKind.Tool, chat: defaultChat, toolCallId: 'tc-1' },
+				},
+			], { defaultChat });
+
+			agentHost.setSessionState('multi-complete', 'copilotcli', stateWithStatus(ProtocolSessionStatus.InProgress));
+			await timeout(31_000);
+			agentHost.setSessionState('multi-complete', 'copilotcli', stateWithStatus(ProtocolSessionStatus.Idle));
+			const staleStatus = session.mainChat.get().status.get();
+
+			agentHost.fireAction({
+				channel: subagentChat,
+				action: {
+					type: ActionType.ChatTurnComplete,
+					turnId: 'turn-1',
+					duration: 1000,
+				},
+				serverSeq: 1,
+				origin: undefined,
+			} as ActionEnvelope);
+			await timeout(0);
+
+			assert.deepStrictEqual({
+				staleStatus,
+				updatedStatus: session.mainChat.get().status.get(),
+				subagentStatus: session.chats.get().find(chat => chat.origin?.kind === ChatOriginKind.Tool)?.status.get(),
+				subscribeCount: agentHost.sessionSubscribeCounts.get(sessionUri),
+				unsubscribeCount: agentHost.sessionUnsubscribeCounts.get(sessionUri),
+			}, {
+				staleStatus: SessionStatus.InProgress,
+				updatedStatus: SessionStatus.Completed,
+				subagentStatus: SessionStatus.Completed,
+				subscribeCount: 2,
+				unsubscribeCount: 1,
+			});
+		}));
 	});
 
 	// ---- Title change from server -------