Support transferring interactive sessions across workspaces (microsoft#185702)

---------

Co-authored-by: Rob Lourens <[email protected]>

Author: joyceerhl

src/vs/workbench/api/browser/mainThreadChat.ts

@@ -5,7 +5,7 @@
 
 import { Emitter } from 'vs/base/common/event';
 import { Disposable, DisposableMap } from 'vs/base/common/lifecycle';
-import { URI } from 'vs/base/common/uri';
+import { URI, UriComponents } from 'vs/base/common/uri';
 import { ExtHostChatShape, ExtHostContext, IChatRequestDto, MainContext, MainThreadChatShape } from 'vs/workbench/api/common/extHost.protocol';
 import { IChatWidgetService } from 'vs/workbench/contrib/chat/browser/chat';
 import { IChatContributionService } from 'vs/workbench/contrib/chat/common/chatContributionService';
@@ -53,6 +53,10 @@ export class MainThreadChat extends Disposable implements MainThreadChatShape {
 		this._providerRegistrations.deleteAndDispose(handle);
 	}
 
+	$transferChatSession(sessionId: number, toWorkspace: UriComponents): void {
+		this._chatService.transferChatSession(sessionId, URI.revive(toWorkspace));
+	}
+
 	async $registerChatProvider(handle: number, id: string): Promise<void> {
 		const registration = this.chatContribService.registeredProviders.find(staticProvider => staticProvider.id === id);
 		if (!registration) {

src/vs/workbench/api/common/extHost.api.impl.ts

@@ -1298,6 +1298,10 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			get onDidPerformUserAction() {
 				checkProposedApiEnabled(extension, 'interactiveUserActions');
 				return extHostChat.onDidPerformUserAction;
+			},
+			transferChatSession(session: vscode.InteractiveSession, toWorkspace: vscode.Uri) {
+				checkProposedApiEnabled(extension, 'interactive');
+				return extHostChat.transferChatSession(session, toWorkspace);
 			}
 		};
 

src/vs/workbench/api/common/extHost.protocol.ts

@@ -1167,6 +1167,7 @@ export interface MainThreadChatShape extends IDisposable {
 	$sendRequestToProvider(providerId: string, message: IChatDynamicRequest): void;
 	$unregisterChatProvider(handle: number): Promise<void>;
 	$acceptResponseProgress(handle: number, sessionId: number, progress: IChatProgress): void;
+	$transferChatSession(sessionId: number, toWorkspace: UriComponents): void;
 
 	$registerSlashCommandProvider(handle: number, chatProviderId: string): Promise<void>;
 	$unregisterSlashCommandProvider(handle: number): Promise<void>;

src/vs/workbench/api/common/extHostChat.ts

@@ -5,6 +5,7 @@
 
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { Emitter } from 'vs/base/common/event';
+import { Iterable } from 'vs/base/common/iterator';
 import { toDisposable } from 'vs/base/common/lifecycle';
 import { StopWatch } from 'vs/base/common/stopwatch';
 import { withNullAsUndefined } from 'vs/base/common/types';
@@ -60,6 +61,15 @@ export class ExtHostChat implements ExtHostChatShape {
 		});
 	}
 
+	transferChatSession(session: vscode.InteractiveSession, newWorkspace: vscode.Uri): void {
+		const sessionId = Iterable.find(this._chatSessions.keys(), key => this._chatSessions.get(key) === session) ?? 0;
+		if (typeof sessionId !== 'number') {
+			return;
+		}
+
+		this._proxy.$transferChatSession(sessionId, newWorkspace);
+	}
+
 	addChatRequest(context: vscode.InteractiveSessionRequestArgs): void {
 		this._proxy.$addRequest(context);
 	}

src/vs/workbench/contrib/chat/browser/chatViewPane.ts

@@ -71,7 +71,9 @@ export class ChatViewPane extends ViewPane implements IChatViewPane {
 	private updateModel(model?: IChatModel | undefined): void {
 		this.modelDisposables.clear();
 
-		model = model ?? this.chatService.startSession(this.chatViewOptions.providerId, CancellationToken.None);
+		model = model ?? (this.chatService.transferredSessionId
+			? this.chatService.getOrRestoreSession(this.chatService.transferredSessionId)
+			: this.chatService.startSession(this.chatViewOptions.providerId, CancellationToken.None));
 		if (!model) {
 			throw new Error('Could not start chat session');
 		}
@@ -100,7 +102,8 @@ export class ChatViewPane extends ViewPane implements IChatViewPane {
 			}));
 			this._widget.render(parent);
 
-			const initialModel = this.viewState.sessionId ? this.chatService.getOrRestoreSession(this.viewState.sessionId) : undefined;
+			const sessionId = this.chatService.transferredSessionId ?? this.viewState.sessionId;
+			const initialModel = sessionId ? this.chatService.getOrRestoreSession(sessionId) : undefined;
 			this.updateModel(initialModel);
 		} catch (e) {
 			this.logService.error(e);

src/vs/workbench/contrib/chat/common/chatService.ts

@@ -175,6 +175,7 @@ export const IChatService = createDecorator<IChatService>('IChatService');
 
 export interface IChatService {
 	_serviceBrand: undefined;
+	transferredSessionId: string | undefined;
 	registerProvider(provider: IChatProvider): IDisposable;
 	registerSlashCommandProvider(provider: ISlashCommandProvider): IDisposable;
 	getProviderInfos(): IChatProviderInfo[];
@@ -199,4 +200,6 @@ export interface IChatService {
 
 	onDidPerformUserAction: Event<IChatUserActionEvent>;
 	notifyUserAction(event: IChatUserActionEvent): void;
+
+	transferChatSession(sessionProviderId: number, toWorkspace: URI): void;
 }

src/vs/workbench/contrib/chat/common/chatServiceImpl.ts

@@ -11,20 +11,30 @@ import { Iterable } from 'vs/base/common/iterator';
 import { Disposable, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { StopWatch } from 'vs/base/common/stopwatch';
 import { withNullAsUndefined } from 'vs/base/common/types';
+import { URI, UriComponents } from 'vs/base/common/uri';
 import { localize } from 'vs/nls';
 import { CommandsRegistry } from 'vs/platform/commands/common/commands';
 import { IContextKey, IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ILogService } from 'vs/platform/log/common/log';
 import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage';
 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 { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 
 const serializedChatKey = 'interactive.sessions';
 
+const globalChatKey = 'chat.workspaceTransfer';
+interface IChatTransfer {
+	toWorkspace: UriComponents;
+	timestampInMilliseconds: number;
+	chat: ISerializableChatData;
+}
+const SESSION_TRANSFER_EXPIRATION_IN_MILLISECONDS = 1000 * 60;
+
 type ChatProviderInvokedEvent = {
 	providerId: string;
 	timeToFirstProgress: number;
@@ -115,6 +125,11 @@ 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 readonly _onDidPerformUserAction = this._register(new Emitter<IChatUserActionEvent>());
 	public readonly onDidPerformUserAction: Event<IChatUserActionEvent> = this._onDidPerformUserAction.event;
 
@@ -125,6 +140,7 @@ export class ChatService extends Disposable implements IChatService {
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@ITelemetryService private readonly telemetryService: ITelemetryService,
 		@IContextKeyService private readonly contextKeyService: IContextKeyService,
+		@IWorkspaceContextService private readonly workspaceContextService: IWorkspaceContextService
 	) {
 		super();
 
@@ -140,6 +156,12 @@ export class ChatService extends Disposable implements IChatService {
 			this.trace('constructor', 'No persisted sessions');
 		}
 
+		this._transferred = this.getTransferredSession();
+		if (this._transferred) {
+			this.trace('constructor', `Transferred session ${this._transferred.sessionId}`);
+			this._persistedSessions[this._transferred.sessionId] = this._transferred;
+		}
+
 		this._register(storageService.onWillSaveState(() => this.saveState()));
 	}
 
@@ -218,6 +240,23 @@ export class ChatService extends Disposable implements IChatService {
 		}
 	}
 
+	private getTransferredSession(): ISerializableChatData | undefined {
+		const data: IChatTransfer[] = this.storageService.getObject(globalChatKey, StorageScope.PROFILE, []);
+		const workspaceUri = this.workspaceContextService.getWorkspace().folders[0]?.uri;
+		if (!workspaceUri) {
+			return;
+		}
+
+		const thisWorkspace = workspaceUri.toString();
+		const currentTime = Date.now();
+		// Only use transferred data if it was created recently
+		const transferred = data.find(item => URI.revive(item.toWorkspace).toString() === thisWorkspace && (currentTime - item.timestampInMilliseconds < SESSION_TRANSFER_EXPIRATION_IN_MILLISECONDS));
+		// 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;
+	}
+
 	getHistory(): IChatDetail[] {
 		const sessions = Object.values(this._persistedSessions)
 			.filter(session => session.requests.length > 0);
@@ -307,6 +346,10 @@ export class ChatService extends Disposable implements IChatService {
 			return undefined;
 		}
 
+		if (sessionId === this.transferredSessionId) {
+			this._transferred = undefined;
+		}
+
 		return this._startSession(sessionData.providerId, sessionData, CancellationToken.None);
 	}
 
@@ -592,4 +635,21 @@ 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);
+		if (!model) {
+			throw new Error(`Failed to transfer session. Unknown session provider ID: ${sessionProviderId}`);
+		}
+
+		const existingRaw: IChatTransfer[] = this.storageService.getObject(globalChatKey, StorageScope.PROFILE, []);
+		existingRaw.push({
+			chat: model.toJSON(),
+			timestampInMilliseconds: Date.now(),
+			toWorkspace: toWorkspace
+		});
+
+		this.storageService.store(globalChatKey, JSON.stringify(existingRaw), StorageScope.PROFILE, StorageTarget.MACHINE);
+		this.trace('transferChatSession', `Transferred session ${model.sessionId} to workspace ${toWorkspace.toString()}`);
+	}
 }

src/vs/workbench/contrib/chat/test/common/chatService.test.ts

@@ -18,7 +18,8 @@ import { IChatContributionService } from 'vs/workbench/contrib/chat/common/chatC
 import { IChatProgress, IChatProvider, IChatRequest, IChatResponse, IChat, ISlashCommand, IPersistedChatState } from 'vs/workbench/contrib/chat/common/chatService';
 import { ChatService } from 'vs/workbench/contrib/chat/common/chatServiceImpl';
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
-import { TestExtensionService, TestStorageService } from 'vs/workbench/test/common/workbenchTestServices';
+import { TestContextService, TestExtensionService, TestStorageService } from 'vs/workbench/test/common/workbenchTestServices';
+import { IWorkspaceContextService } from 'vs/platform/workspace/common/workspace';
 
 class SimpleTestProvider extends Disposable implements IChatProvider {
 	private static sessionId = 0;
@@ -67,6 +68,7 @@ suite('Chat', () => {
 		instantiationService.stub(IContextKeyService, new MockContextKeyService());
 		instantiationService.stub(IViewsService, new TestExtensionService());
 		instantiationService.stub(IChatContributionService, new TestExtensionService());
+		instantiationService.stub(IWorkspaceContextService, new TestContextService());
 	});
 
 	teardown(() => {

src/vscode-dts/vscode.proposed.interactive.d.ts

@@ -181,5 +181,7 @@ declare module 'vscode' {
 		export function sendInteractiveRequestToProvider(providerId: string, message: InteractiveSessionDynamicRequest): void;
 
 		export function registerInteractiveEditorSessionProvider(provider: InteractiveEditorSessionProvider): Disposable;
+
+		export function transferChatSession(session: InteractiveSession, toWorkspace: Uri): void;
 	}
 }