Get widget by chat session id and other changes

microsoft · Jul 23, 2023 · 4311d68 · 4311d68
1 parent 0ca38fe
commit 4311d68
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 30 deletions.
diff --git a/src/vs/workbench/api/browser/mainThreadChat.ts b/src/vs/workbench/api/browser/mainThreadChat.ts
@@ -54,9 +54,14 @@ export class MainThreadChat extends Disposable implements MainThreadChatShape {
 	}
 
 	$transferChatSession(sessionId: number, toWorkspace: UriComponents): void {
-		const widget = this._chatWidgetService.lastFocusedWidget;
+		const sessionIdStr = this._chatService.getSessionId(sessionId);
+		if (!sessionIdStr) {
+			throw new Error(`Failed to transfer session. Unknown session provider ID: ${sessionId}`);
+		}
+
+		const widget = this._chatWidgetService.getWidgetBySessionId(sessionIdStr);
 		const inputValue = widget?.inputEditor.getValue() ?? '';
-		this._chatService.transferChatSession({ sessionProviderId: sessionId, inputValue: inputValue }, URI.revive(toWorkspace));
+		this._chatService.transferChatSession({ sessionId: sessionIdStr, inputValue: inputValue }, URI.revive(toWorkspace));
 	}
 
 	async $registerChatProvider(handle: number, id: string): Promise<void> {

diff --git a/src/vs/workbench/contrib/chat/browser/chat.ts b/src/vs/workbench/contrib/chat/browser/chat.ts
@@ -28,6 +28,8 @@ export interface IChatWidgetService {
 	revealViewForProvider(providerId: string): Promise<IChatWidget | undefined>;
 
 	getWidgetByInputUri(uri: URI): IChatWidget | undefined;
+
+	getWidgetBySessionId(sessionId: string): IChatWidget | undefined;
 }
 
 

diff --git a/src/vs/workbench/contrib/chat/browser/chatViewPane.ts b/src/vs/workbench/contrib/chat/browser/chatViewPane.ts
@@ -71,8 +71,8 @@ export class ChatViewPane extends ViewPane implements IChatViewPane {
 	private updateModel(model?: IChatModel | undefined): void {
 		this.modelDisposables.clear();
 
-		model = model ?? (this.chatService.transferredSessionId
-			? this.chatService.getOrRestoreSession(this.chatService.transferredSessionId)
+		model = model ?? (this.chatService.transferredSessionData?.sessionId
+			? this.chatService.getOrRestoreSession(this.chatService.transferredSessionData.sessionId)
 			: this.chatService.startSession(this.chatViewOptions.providerId, CancellationToken.None));
 		if (!model) {
 			throw new Error('Could not start chat session');
@@ -103,9 +103,9 @@ export class ChatViewPane extends ViewPane implements IChatViewPane {
 			this._widget.render(parent);
 
 			let sessionId: string | undefined;
-			if (this.chatService.transferredSessionId) {
-				sessionId = this.chatService.transferredSessionId;
-				this.viewState.inputValue = this.chatService.transferredInputValue;
+			if (this.chatService.transferredSessionData) {
+				sessionId = this.chatService.transferredSessionData.sessionId;
+				this.viewState.inputValue = this.chatService.transferredSessionData.inputValue;
 			} else {
 				sessionId = this.viewState.sessionId;
 			}

diff --git a/src/vs/workbench/contrib/chat/browser/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/chatWidget.ts
@@ -574,6 +574,10 @@ export class ChatWidgetService implements IChatWidgetService {
 		return this._widgets.find(w => isEqual(w.inputUri, uri));
 	}
 
+	getWidgetBySessionId(sessionId: string): ChatWidget | undefined {
+		return this._widgets.find(w => w.viewModel?.sessionId === sessionId);
+	}
+
 	async revealViewForProvider(providerId: string): Promise<ChatWidget | undefined> {
 		const viewId = this.chatContributionService.getViewIdForProvider(providerId);
 		const view = await this.viewsService.openView<ChatViewPane>(viewId);

diff --git a/src/vs/workbench/contrib/chat/common/chatService.ts b/src/vs/workbench/contrib/chat/common/chatService.ts
@@ -189,19 +189,24 @@ export interface IChatProviderInfo {
 	displayName: string;
 }
 
+export interface IChatTransferredSessionData {
+	sessionId: string;
+	inputValue: string;
+}
+
 export const IChatService = createDecorator<IChatService>('IChatService');
 
 export interface IChatService {
 	_serviceBrand: undefined;
-	transferredSessionId: string | undefined;
-	transferredInputValue: string | undefined;
+	transferredSessionData: IChatTransferredSessionData | undefined;
 
 	onDidSubmitSlashCommand: Event<{ slashCommand: string; sessionId: string }>;
 	registerProvider(provider: IChatProvider): IDisposable;
 	registerSlashCommandProvider(provider: ISlashCommandProvider): IDisposable;
 	getProviderInfos(): IChatProviderInfo[];
 	startSession(providerId: string, token: CancellationToken): ChatModel | undefined;
 	getSession(sessionId: string): IChatModel | undefined;
+	getSessionId(sessionProviderId: number): string | undefined;
 	getOrRestoreSession(sessionId: string): IChatModel | undefined;
 	loadSessionFromContent(data: ISerializableChatData): IChatModel | undefined;
 
@@ -222,5 +227,5 @@ export interface IChatService {
 	onDidPerformUserAction: Event<IChatUserActionEvent>;
 	notifyUserAction(event: IChatUserActionEvent): void;
 
-	transferChatSession(transferredSessionData: { sessionProviderId: number; inputValue: string }, toWorkspace: URI): void;
+	transferChatSession(transferredSessionData: IChatTransferredSessionData, toWorkspace: URI): void;
 }
diff --git a/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts b/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts
@@ -22,7 +22,7 @@ import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { IWorkspaceContextService } from 'vs/platform/workspace/common/workspace';
 import { CONTEXT_PROVIDER_EXISTS } from 'vs/workbench/contrib/chat/common/chatContextKeys';
 import { ChatModel, ChatWelcomeMessageModel, IChatModel, ISerializableChatData, ISerializableChatsData } from 'vs/workbench/contrib/chat/common/chatModel';
-import { IChat, IChatCompleteResponse, IChatDetail, IChatDynamicRequest, IChatProgress, IChatProvider, IChatProviderInfo, IChatReplyFollowup, IChatService, IChatUserActionEvent, ISlashCommand, ISlashCommandProvider, InteractiveSessionCopyKind, InteractiveSessionVoteDirection } from 'vs/workbench/contrib/chat/common/chatService';
+import { IChat, IChatCompleteResponse, IChatDetail, IChatDynamicRequest, IChatProgress, IChatProvider, IChatProviderInfo, IChatReplyFollowup, IChatService, IChatTransferredSessionData, IChatUserActionEvent, ISlashCommand, ISlashCommandProvider, InteractiveSessionCopyKind, InteractiveSessionVoteDirection } from 'vs/workbench/contrib/chat/common/chatService';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 
 const serializedChatKey = 'interactive.sessions';
@@ -128,15 +128,11 @@ export class ChatService extends Disposable implements IChatService {
 	private readonly _persistedSessions: ISerializableChatsData;
 	private readonly _hasProvider: IContextKey<boolean>;
 
-	private _transferredChat: ISerializableChatData | undefined;
-	public get transferredSessionId(): string | undefined {
-		return this._transferredChat?.sessionId;
+	private _transferredSessionData: IChatTransferredSessionData | undefined;
+	public get transferredSessionData(): IChatTransferredSessionData | undefined {
+		return this._transferredSessionData;
 	}
 
-	private _transferredInputValue: string | undefined;
-	public get transferredInputValue(): string | undefined {
-		return this._transferredInputValue;
-	}
 	private readonly _onDidPerformUserAction = this._register(new Emitter<IChatUserActionEvent>());
 	public readonly onDidPerformUserAction: Event<IChatUserActionEvent> = this._onDidPerformUserAction.event;
 
@@ -168,12 +164,12 @@ export class ChatService extends Disposable implements IChatService {
 		}
 
 		const transferredData = this.getTransferredSessionData();
-		this._transferredChat = transferredData?.chat;
-		if (this._transferredChat) {
-			this.trace('constructor', `Transferred session ${this._transferredChat.sessionId}`);
-			this._persistedSessions[this._transferredChat.sessionId] = this._transferredChat;
+		const transferredChat = transferredData?.chat;
+		if (transferredChat) {
+			this.trace('constructor', `Transferred session ${transferredChat.sessionId}`);
+			this._persistedSessions[transferredChat.sessionId] = transferredChat;
+			this._transferredSessionData = { sessionId: transferredChat.sessionId, inputValue: transferredData.inputValue };
 		}
-		this._transferredInputValue = transferredData?.inputValue;
 
 		this._register(storageService.onWillSaveState(() => this.saveState()));
 	}
@@ -362,6 +358,10 @@ export class ChatService extends Disposable implements IChatService {
 		return this._sessionModels.get(sessionId);
 	}
 
+	getSessionId(sessionProviderId: number): string | undefined {
+		return Iterable.find(this._sessionModels.values(), model => model.session?.id === sessionProviderId)?.sessionId;
+	}
+
 	getOrRestoreSession(sessionId: string): ChatModel | undefined {
 		const model = this._sessionModels.get(sessionId);
 		if (model) {
@@ -373,9 +373,8 @@ export class ChatService extends Disposable implements IChatService {
 			return undefined;
 		}
 
-		if (sessionId === this.transferredSessionId) {
-			this._transferredChat = undefined;
-			this._transferredInputValue = undefined;
+		if (sessionId === this.transferredSessionData?.sessionId) {
+			this._transferredSessionData = undefined;
 		}
 
 		return this._startSession(sessionData.providerId, sessionData, CancellationToken.None);
@@ -677,11 +676,10 @@ export class ChatService extends Disposable implements IChatService {
 		});
 	}
 
-	transferChatSession(transferredSessionData: { sessionProviderId: number; inputValue: string }, toWorkspace: URI): void {
-		const sessionProviderId = transferredSessionData.sessionProviderId;
-		const model = Iterable.find(this._sessionModels.values(), model => model.session?.id === sessionProviderId);
+	transferChatSession(transferredSessionData: IChatTransferredSessionData, toWorkspace: URI): void {
+		const model = Iterable.find(this._sessionModels.values(), model => model.sessionId === transferredSessionData.sessionId);
 		if (!model) {
-			throw new Error(`Failed to transfer session. Unknown session provider ID: ${sessionProviderId}`);
+			throw new Error(`Failed to transfer session. Unknown session ID: ${transferredSessionData.sessionId}`);
 		}
 
 		const existingRaw: IChatTransfer[] = this.storageService.getObject(globalChatKey, StorageScope.PROFILE, []);