Save input value as IChatTransferData for restoring after transfer (#…

…188543) * Store input value in IChatTransferData * Get widget by chat session id and other changes
microsoft · Jul 25, 2023 · 4767096 · 4767096
1 parent 08235aa
commit 4767096
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 22 deletions.
diff --git a/src/vs/workbench/api/browser/mainThreadChat.ts b/src/vs/workbench/api/browser/mainThreadChat.ts
@@ -58,7 +58,14 @@ export class MainThreadChat extends Disposable implements MainThreadChatShape {
 	}
 
 	$transferChatSession(sessionId: number, toWorkspace: UriComponents): void {
-		this._chatService.transferChatSession(sessionId, URI.revive(toWorkspace));
+		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({ 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');
@@ -102,7 +102,14 @@ export class ChatViewPane extends ViewPane implements IChatViewPane {
 			}));
 			this._widget.render(parent);
 
-			const sessionId = this.chatService.transferredSessionId ?? this.viewState.sessionId;
+			let sessionId: string | undefined;
+			if (this.chatService.transferredSessionData) {
+				sessionId = this.chatService.transferredSessionData.sessionId;
+				this.viewState.inputValue = this.chatService.transferredSessionData.inputValue;
+			} else {
+				sessionId = this.viewState.sessionId;
+			}
+
 			const initialModel = sessionId ? this.chatService.getOrRestoreSession(sessionId) : undefined;
 			this.updateModel(initialModel);
 		} catch (e) {

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,18 +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;
+	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;
 
@@ -221,5 +227,5 @@ export interface IChatService {
 	onDidPerformUserAction: Event<IChatUserActionEvent>;
 	notifyUserAction(event: IChatUserActionEvent): void;
 
-	transferChatSession(sessionProviderId: number, 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';
@@ -32,6 +32,7 @@ interface IChatTransfer {
 	toWorkspace: UriComponents;
 	timestampInMilliseconds: number;
 	chat: ISerializableChatData;
+	inputValue: string;
 }
 const SESSION_TRANSFER_EXPIRATION_IN_MILLISECONDS = 1000 * 60;
 
@@ -127,9 +128,9 @@ export class ChatService extends Disposable implements IChatService {
 	private readonly _persistedSessions: ISerializableChatsData;
 	private readonly _hasProvider: IContextKey<boolean>;
 
-	private _transferred: ISerializableChatData | undefined;
-	public get transferredSessionId(): string | undefined {
-		return this._transferred?.sessionId;
+	private _transferredSessionData: IChatTransferredSessionData | undefined;
+	public get transferredSessionData(): IChatTransferredSessionData | undefined {
+		return this._transferredSessionData;
 	}
 
 	private readonly _onDidPerformUserAction = this._register(new Emitter<IChatUserActionEvent>());
@@ -162,10 +163,12 @@ export class ChatService extends Disposable implements IChatService {
 			this._persistedSessions = {};
 		}
 
-		this._transferred = this.getTransferredSession();
-		if (this._transferred) {
-			this.trace('constructor', `Transferred session ${this._transferred.sessionId}`);
-			this._persistedSessions[this._transferred.sessionId] = this._transferred;
+		const transferredData = this.getTransferredSessionData();
+		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._register(storageService.onWillSaveState(() => this.saveState()));
@@ -252,7 +255,7 @@ export class ChatService extends Disposable implements IChatService {
 		}
 	}
 
-	private getTransferredSession(): ISerializableChatData | undefined {
+	private getTransferredSessionData(): IChatTransfer | undefined {
 		const data: IChatTransfer[] = this.storageService.getObject(globalChatKey, StorageScope.PROFILE, []);
 		const workspaceUri = this.workspaceContextService.getWorkspace().folders[0]?.uri;
 		if (!workspaceUri) {
@@ -266,7 +269,7 @@ export class ChatService extends Disposable implements IChatService {
 		// Keep data that isn't for the current workspace and that hasn't expired yet
 		const filtered = data.filter(item => URI.revive(item.toWorkspace).toString() !== thisWorkspace && (currentTime - item.timestampInMilliseconds < SESSION_TRANSFER_EXPIRATION_IN_MILLISECONDS));
 		this.storageService.store(globalChatKey, JSON.stringify(filtered), StorageScope.PROFILE, StorageTarget.MACHINE);
-		return transferred?.chat;
+		return transferred;
 	}
 
 	getHistory(): IChatDetail[] {
@@ -355,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) {
@@ -366,8 +373,8 @@ export class ChatService extends Disposable implements IChatService {
 			return undefined;
 		}
 
-		if (sessionId === this.transferredSessionId) {
-			this._transferred = undefined;
+		if (sessionId === this.transferredSessionData?.sessionId) {
+			this._transferredSessionData = undefined;
 		}
 
 		return this._startSession(sessionData.providerId, sessionData, CancellationToken.None);
@@ -671,17 +678,18 @@ export class ChatService extends Disposable implements IChatService {
 		});
 	}
 
-	transferChatSession(sessionProviderId: number, toWorkspace: URI): void {
-		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, []);
 		existingRaw.push({
 			chat: model.toJSON(),
 			timestampInMilliseconds: Date.now(),
-			toWorkspace: toWorkspace
+			toWorkspace: toWorkspace,
+			inputValue: transferredSessionData.inputValue,
 		});
 
 		this.storageService.store(globalChatKey, JSON.stringify(existingRaw), StorageScope.PROFILE, StorageTarget.MACHINE);