Skip to content

Commit 792929f

Browse files
roblourensroblourens
authored
Get chat session transferring working on the ChatSessionStore (#283512)
* Get chat session transferring working on the ChatSessionStore * fix comment * Fix tests, validate location * Tests * IChatTransferredSessionData is just a URI * Fix leak
1 parent f6686c6 commit 792929f

19 files changed

+721
-267
lines changed

‎src/vs/workbench/api/browser/mainThreadChatAgents2.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -148,16 +148,15 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA
148148
this._agents.deleteAndDispose(handle);
149149
}
150150

151-
$transferActiveChatSession(toWorkspace: UriComponents): void {
151+
async $transferActiveChatSession(toWorkspace: UriComponents): Promise<void> {
152152
const widget = this._chatWidgetService.lastFocusedWidget;
153153
const model = widget?.viewModel?.model;
154154
if (!model) {
155155
this._logService.error(`MainThreadChat#$transferActiveChatSession: No active chat session found`);
156156
return;
157157
}
158158

159-
const location = widget.location;
160-
this._chatService.transferChatSession({ sessionId: model.sessionId, inputState: model.inputModel.state.get(), location }, URI.revive(toWorkspace));
159+
await this._chatService.transferChatSession(model.sessionResource, URI.revive(toWorkspace));
161160
}
162161

163162
async $registerAgent(handle: number, extension: ExtensionIdentifier, id: string, metadata: IExtensionChatAgentMetadata, dynamicProps: IDynamicChatAgentProps | undefined): Promise<void> {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1469,7 +1469,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
14691469

14701470
// namespace: interactive
14711471
const interactive: typeof vscode.interactive = {
1472-
transferActiveChat(toWorkspace: vscode.Uri) {
1472+
transferActiveChat(toWorkspace: vscode.Uri): Thenable<void> {
14731473
checkProposedApiEnabled(extension, 'interactive');
14741474
return extHostChatAgents2.transferActiveChat(toWorkspace);
14751475
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1402,7 +1402,7 @@ export interface MainThreadChatAgentsShape2 extends IChatAgentProgressShape, IDi
14021402
$updateAgent(handle: number, metadataUpdate: IExtensionChatAgentMetadata): void;
14031403
$unregisterAgent(handle: number): void;
14041404

1405-
$transferActiveChatSession(toWorkspace: UriComponents): void;
1405+
$transferActiveChatSession(toWorkspace: UriComponents): Promise<void>;
14061406
}
14071407

14081408
export interface ICodeMapperTextEdit {

‎src/vs/workbench/api/common/extHostChatAgents2.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -437,8 +437,8 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS
437437
});
438438
}
439439

440-
transferActiveChat(newWorkspace: vscode.Uri): void {
441-
this._proxy.$transferActiveChatSession(newWorkspace);
440+
async transferActiveChat(newWorkspace: vscode.Uri): Promise<void> {
441+
await this._proxy.$transferActiveChatSession(newWorkspace);
442442
}
443443

444444
createChatAgent(extension: IExtensionDescription, id: string, handler: vscode.ChatExtendedRequestHandler): vscode.ChatParticipant {

‎src/vs/workbench/contrib/chat/browser/actions/chatActions.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -310,13 +310,15 @@ abstract class OpenChatGlobalAction extends Action2 {
310310
let resp: Promise<IChatResponseModel | undefined> | undefined;
311311

312312
if (opts?.query) {
313-
chatWidget.setInput(opts.query);
314313

315-
if (!opts.isPartialQuery) {
314+
if (opts.isPartialQuery) {
315+
chatWidget.setInput(opts.query);
316+
} else {
316317
if (!chatWidget.viewModel) {
317318
await Event.toPromise(chatWidget.onDidChangeViewModel);
318319
}
319320
await waitForDefaultAgent(chatAgentService, chatWidget.input.currentModeKind);
321+
chatWidget.setInput(opts.query); // wait until the model is restored before setting the input, or it will be cleared when the model is restored
320322
resp = chatWidget.acceptInput();
321323
}
322324
}

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

Lines changed: 10 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -215,9 +215,9 @@ export class ChatViewPane extends ViewPane implements IViewWelcomeDelegate {
215215
private onDidChangeAgents(): void {
216216
if (this.chatAgentService.getDefaultAgent(ChatAgentLocation.Chat)) {
217217
if (!this._widget?.viewModel && !this.restoringSession) {
218-
const info = this.getTransferredOrPersistedSessionInfo();
218+
const sessionResource = this.getTransferredOrPersistedSessionInfo();
219219
this.restoringSession =
220-
(info.sessionId ? this.chatService.getOrRestoreSession(LocalChatSessionUri.forSession(info.sessionId)) : Promise.resolve(undefined)).then(async modelRef => {
220+
(sessionResource ? this.chatService.getOrRestoreSession(sessionResource) : Promise.resolve(undefined)).then(async modelRef => {
221221
if (!this._widget) {
222222
return; // renderBody has not been called yet
223223
}
@@ -228,9 +228,6 @@ export class ChatViewPane extends ViewPane implements IViewWelcomeDelegate {
228228
const wasVisible = this._widget.visible;
229229
try {
230230
this._widget.setVisible(false);
231-
if (info.inputState && modelRef) {
232-
modelRef.object.inputModel.setState(info.inputState);
233-
}
234231

235232
await this.showModel(modelRef);
236233
} finally {
@@ -245,16 +242,12 @@ export class ChatViewPane extends ViewPane implements IViewWelcomeDelegate {
245242
this._onDidChangeViewWelcomeState.fire();
246243
}
247244

248-
private getTransferredOrPersistedSessionInfo(): { sessionId?: string; inputState?: IChatModelInputState; mode?: ChatModeKind } {
249-
if (this.chatService.transferredSessionData?.location === ChatAgentLocation.Chat) {
250-
const sessionId = this.chatService.transferredSessionData.sessionId;
251-
return {
252-
sessionId,
253-
inputState: this.chatService.transferredSessionData.inputState,
254-
};
245+
private getTransferredOrPersistedSessionInfo(): URI | undefined {
246+
if (this.chatService.transferredSessionResource) {
247+
return this.chatService.transferredSessionResource;
255248
}
256249

257-
return { sessionId: this.viewState.sessionId };
250+
return this.viewState.sessionId ? LocalChatSessionUri.forSession(this.viewState.sessionId) : undefined;
258251
}
259252

260253
protected override renderBody(parent: HTMLElement): void {
@@ -658,12 +651,8 @@ export class ChatViewPane extends ViewPane implements IViewWelcomeDelegate {
658651
//#region Model Management
659652

660653
private async applyModel(): Promise<void> {
661-
const info = this.getTransferredOrPersistedSessionInfo();
662-
const modelRef = info.sessionId ? await this.chatService.getOrRestoreSession(LocalChatSessionUri.forSession(info.sessionId)) : undefined;
663-
if (modelRef && info.inputState) {
664-
modelRef.object.inputModel.setState(info.inputState);
665-
}
666-
654+
const sessionResource = this.getTransferredOrPersistedSessionInfo();
655+
const modelRef = sessionResource ? await this.chatService.getOrRestoreSession(sessionResource) : undefined;
667656
await this.showModel(modelRef);
668657
}
669658

@@ -673,8 +662,8 @@ export class ChatViewPane extends ViewPane implements IViewWelcomeDelegate {
673662

674663
let ref: IChatModelReference | undefined;
675664
if (startNewSession) {
676-
ref = modelRef ?? (this.chatService.transferredSessionData?.sessionId && this.chatService.transferredSessionData?.location === ChatAgentLocation.Chat
677-
? await this.chatService.getOrRestoreSession(LocalChatSessionUri.forSession(this.chatService.transferredSessionData.sessionId))
665+
ref = modelRef ?? (this.chatService.transferredSessionResource
666+
? await this.chatService.getOrRestoreSession(this.chatService.transferredSessionResource)
678667
: this.chatService.startSession(ChatAgentLocation.Chat));
679668
if (!ref) {
680669
throw new Error('Could not start chat session');

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

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import { ICellEditOperation } from '../../notebook/common/notebookCommon.js';
2323
import { IWorkspaceSymbol } from '../../search/common/search.js';
2424
import { IChatAgentCommand, IChatAgentData, IChatAgentResult, UserSelectedTools } from './chatAgents.js';
2525
import { IChatEditingSession } from './chatEditingService.js';
26-
import { IChatModel, IChatModelInputState, IChatRequestModeInfo, IChatRequestModel, IChatRequestVariableData, IChatResponseModel, IExportableChatData, ISerializableChatData } from './chatModel.js';
26+
import { IChatModel, IChatRequestModeInfo, IChatRequestModel, IChatRequestVariableData, IChatResponseModel, IExportableChatData, ISerializableChatData } from './chatModel.js';
2727
import { IParsedChatRequest } from './chatParserTypes.js';
2828
import { IChatParserContext } from './chatRequestParser.js';
2929
import { IChatRequestVariableEntry } from './chatVariableEntries.js';
@@ -934,12 +934,6 @@ export interface IChatProviderInfo {
934934
id: string;
935935
}
936936

937-
export interface IChatTransferredSessionData {
938-
sessionId: string;
939-
location: ChatAgentLocation;
940-
inputState: IChatModelInputState | undefined;
941-
}
942-
943937
export interface IChatSendRequestResponseState {
944938
responseCreatedPromise: Promise<IChatResponseModel>;
945939
responseCompletePromise: Promise<void>;
@@ -1006,7 +1000,7 @@ export const IChatService = createDecorator<IChatService>('IChatService');
10061000

10071001
export interface IChatService {
10081002
_serviceBrand: undefined;
1009-
transferredSessionData: IChatTransferredSessionData | undefined;
1003+
transferredSessionResource: URI | undefined;
10101004

10111005
readonly onDidSubmitRequest: Event<{ readonly chatSessionResource: URI }>;
10121006

@@ -1066,7 +1060,7 @@ export interface IChatService {
10661060
notifyUserAction(event: IChatUserActionEvent): void;
10671061
readonly onDidDisposeSession: Event<{ readonly sessionResource: URI[]; readonly reason: 'cleared' }>;
10681062

1069-
transferChatSession(transferredSessionData: IChatTransferredSessionData, toWorkspace: URI): void;
1063+
transferChatSession(transferredSessionResource: URI, toWorkspace: URI): Promise<void>;
10701064

10711065
activateDefaultAgent(location: ChatAgentLocation): Promise<void>;
10721066

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

Lines changed: 31 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import { Disposable, DisposableResourceMap, DisposableStore, IDisposable, Mutabl
1414
import { revive } from '../../../../base/common/marshalling.js';
1515
import { Schemas } from '../../../../base/common/network.js';
1616
import { autorun, derived, IObservable } from '../../../../base/common/observable.js';
17+
import { isEqual } from '../../../../base/common/resources.js';
1718
import { StopWatch } from '../../../../base/common/stopwatch.js';
1819
import { isDefined } from '../../../../base/common/types.js';
1920
import { URI } from '../../../../base/common/uri.js';
@@ -24,7 +25,7 @@ import { IConfigurationService } from '../../../../platform/configuration/common
2425
import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
2526
import { ILogService } from '../../../../platform/log/common/log.js';
2627
import { Progress } from '../../../../platform/progress/common/progress.js';
27-
import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
28+
import { IStorageService, StorageScope } from '../../../../platform/storage/common/storage.js';
2829
import { IWorkspaceContextService } from '../../../../platform/workspace/common/workspace.js';
2930
import { IExtensionService } from '../../../services/extensions/common/extensions.js';
3031
import { InlineChatConfigKeys } from '../../inlineChat/common/inlineChat.js';
@@ -36,10 +37,10 @@ import { ChatModel, ChatRequestModel, ChatRequestRemovalReason, IChatModel, ICha
3637
import { ChatModelStore, IStartSessionProps } from './chatModelStore.js';
3738
import { chatAgentLeader, ChatRequestAgentPart, ChatRequestAgentSubcommandPart, ChatRequestSlashCommandPart, ChatRequestTextPart, chatSubcommandLeader, getPromptText, IParsedChatRequest } from './chatParserTypes.js';
3839
import { ChatRequestParser } from './chatRequestParser.js';
39-
import { ChatMcpServersStarting, IChatCompleteResponse, IChatDetail, IChatFollowup, IChatModelReference, IChatProgress, IChatSendRequestData, IChatSendRequestOptions, IChatSendRequestResponseState, IChatService, IChatSessionContext, IChatSessionStartOptions, IChatTransferredSessionData, IChatUserActionEvent, ResponseModelState } from './chatService.js';
40+
import { ChatMcpServersStarting, IChatCompleteResponse, IChatDetail, IChatFollowup, IChatModelReference, IChatProgress, IChatSendRequestData, IChatSendRequestOptions, IChatSendRequestResponseState, IChatService, IChatSessionContext, IChatSessionStartOptions, IChatUserActionEvent, ResponseModelState } from './chatService.js';
4041
import { ChatRequestTelemetry, ChatServiceTelemetry } from './chatServiceTelemetry.js';
4142
import { IChatSessionsService } from './chatSessionsService.js';
42-
import { ChatSessionStore, IChatSessionEntryMetadata, IChatTransfer2 } from './chatSessionStore.js';
43+
import { ChatSessionStore, IChatSessionEntryMetadata } from './chatSessionStore.js';
4344
import { IChatSlashCommandService } from './chatSlashCommands.js';
4445
import { IChatTransferService } from './chatTransferService.js';
4546
import { LocalChatSessionUri } from './chatUri.js';
@@ -50,10 +51,6 @@ import { ILanguageModelToolsService } from './languageModelToolsService.js';
5051

5152
const serializedChatKey = 'interactive.sessions';
5253

53-
const TransferredGlobalChatKey = 'chat.workspaceTransfer';
54-
55-
const SESSION_TRANSFER_EXPIRATION_IN_MILLISECONDS = 1000 * 60;
56-
5754
class CancellableRequest implements IDisposable {
5855
constructor(
5956
public readonly cancellationTokenSource: CancellationTokenSource,
@@ -82,9 +79,9 @@ export class ChatService extends Disposable implements IChatService {
8279
private _persistedSessions: ISerializableChatsData;
8380
private _saveModelsEnabled = true;
8481

85-
private _transferredSessionData: IChatTransferredSessionData | undefined;
86-
public get transferredSessionData(): IChatTransferredSessionData | undefined {
87-
return this._transferredSessionData;
82+
private _transferredSessionResource: URI | undefined;
83+
public get transferredSessionResource(): URI | undefined {
84+
return this._transferredSessionResource;
8885
}
8986

9087
private readonly _onDidSubmitRequest = this._register(new Emitter<{ readonly chatSessionResource: URI }>());
@@ -128,7 +125,7 @@ export class ChatService extends Disposable implements IChatService {
128125
}
129126

130127
constructor(
131-
@IStorageService private readonly storageService: IStorageService,
128+
@IStorageService storageService: IStorageService,
132129
@ILogService private readonly logService: ILogService,
133130
@IExtensionService private readonly extensionService: IExtensionService,
134131
@IInstantiationService private readonly instantiationService: IInstantiationService,
@@ -175,21 +172,15 @@ export class ChatService extends Disposable implements IChatService {
175172
this._persistedSessions = {};
176173
}
177174

178-
const transferredData = this.getTransferredSessionData();
179-
const transferredChat = transferredData?.chat;
180-
if (transferredChat) {
181-
this.trace('constructor', `Transferred session ${transferredChat.sessionId}`);
182-
this._persistedSessions[transferredChat.sessionId] = transferredChat;
183-
this._transferredSessionData = {
184-
sessionId: transferredChat.sessionId,
185-
location: transferredData.location,
186-
inputState: transferredData.inputState
187-
};
188-
}
189-
190175
this._chatSessionStore = this._register(this.instantiationService.createInstance(ChatSessionStore));
191176
this._chatSessionStore.migrateDataIfNeeded(() => this._persistedSessions);
192177

178+
const transferredData = this._chatSessionStore.getTransferredSessionData();
179+
if (transferredData) {
180+
this.trace('constructor', `Transferred session ${transferredData}`);
181+
this._transferredSessionResource = transferredData;
182+
}
183+
193184
// When using file storage, populate _persistedSessions with session metadata from the index
194185
// This ensures that getPersistedSessionTitle() can find titles for inactive sessions
195186
this.initializePersistedSessionsFromFileStorage().then(() => {
@@ -309,23 +300,6 @@ export class ChatService extends Disposable implements IChatService {
309300
}
310301
}
311302

312-
private getTransferredSessionData(): IChatTransfer2 | undefined {
313-
const data: IChatTransfer2[] = this.storageService.getObject(TransferredGlobalChatKey, StorageScope.PROFILE, []);
314-
const workspaceUri = this.workspaceContextService.getWorkspace().folders[0]?.uri;
315-
if (!workspaceUri) {
316-
return;
317-
}
318-
319-
const thisWorkspace = workspaceUri.toString();
320-
const currentTime = Date.now();
321-
// Only use transferred data if it was created recently
322-
const transferred = data.find(item => URI.revive(item.toWorkspace).toString() === thisWorkspace && (currentTime - item.timestampInMilliseconds < SESSION_TRANSFER_EXPIRATION_IN_MILLISECONDS));
323-
// Keep data that isn't for the current workspace and that hasn't expired yet
324-
const filtered = data.filter(item => URI.revive(item.toWorkspace).toString() !== thisWorkspace && (currentTime - item.timestampInMilliseconds < SESSION_TRANSFER_EXPIRATION_IN_MILLISECONDS));
325-
this.storageService.store(TransferredGlobalChatKey, JSON.stringify(filtered), StorageScope.PROFILE, StorageTarget.MACHINE);
326-
return transferred;
327-
}
328-
329303
/**
330304
* todo@connor4312 This will be cleaned up with the globalization of edits.
331305
*/
@@ -540,8 +514,9 @@ export class ChatService extends Disposable implements IChatService {
540514
}
541515

542516
let sessionData: ISerializableChatData | undefined;
543-
if (this.transferredSessionData?.sessionId === sessionId) {
544-
sessionData = revive(this._persistedSessions[sessionId]);
517+
if (isEqual(this.transferredSessionResource, sessionResource)) {
518+
this._transferredSessionResource = undefined;
519+
sessionData = revive(await this._chatSessionStore.readTransferredSession(sessionResource));
545520
} else {
546521
sessionData = revive(await this._chatSessionStore.readSession(sessionId));
547522
}
@@ -558,11 +533,6 @@ export class ChatService extends Disposable implements IChatService {
558533
canUseTools: true,
559534
});
560535

561-
const isTransferred = this.transferredSessionData?.sessionId === sessionId;
562-
if (isTransferred) {
563-
this._transferredSessionData = undefined;
564-
}
565-
566536
return sessionRef;
567537
}
568538

@@ -1309,22 +1279,25 @@ export class ChatService extends Disposable implements IChatService {
13091279
return this._chatSessionStore.hasSessions();
13101280
}
13111281

1312-
transferChatSession(transferredSessionData: IChatTransferredSessionData, toWorkspace: URI): void {
1313-
const model = Iterable.find(this._sessionModels.values(), model => model.sessionId === transferredSessionData.sessionId);
1282+
async transferChatSession(transferredSessionResource: URI, toWorkspace: URI): Promise<void> {
1283+
if (!LocalChatSessionUri.isLocalSession(transferredSessionResource)) {
1284+
throw new Error(`Can only transfer local chat sessions. Invalid session: ${transferredSessionResource}`);
1285+
}
1286+
1287+
const model = this._sessionModels.get(transferredSessionResource) as ChatModel | undefined;
13141288
if (!model) {
1315-
throw new Error(`Failed to transfer session. Unknown session ID: ${transferredSessionData.sessionId}`);
1289+
throw new Error(`Failed to transfer session. Unknown session: ${transferredSessionResource}`);
1290+
}
1291+
1292+
if (model.initialLocation !== ChatAgentLocation.Chat) {
1293+
throw new Error(`Can only transfer chat sessions located in the Chat view. Session ${transferredSessionResource} has location=${model.initialLocation}`);
13161294
}
13171295

1318-
const existingRaw: IChatTransfer2[] = this.storageService.getObject(TransferredGlobalChatKey, StorageScope.PROFILE, []);
1319-
existingRaw.push({
1320-
chat: model.toJSON(),
1296+
await this._chatSessionStore.storeTransferSession({
1297+
sessionResource: model.sessionResource,
13211298
timestampInMilliseconds: Date.now(),
13221299
toWorkspace: toWorkspace,
1323-
inputState: transferredSessionData.inputState,
1324-
location: transferredSessionData.location,
1325-
});
1326-
1327-
this.storageService.store(TransferredGlobalChatKey, JSON.stringify(existingRaw), StorageScope.PROFILE, StorageTarget.MACHINE);
1300+
}, model);
13281301
this.chatTransferService.addWorkspaceToTransferred(toWorkspace);
13291302
this.trace('transferChatSession', `Transferred session ${model.sessionResource} to workspace ${toWorkspace.toString()}`);
13301303
}

0 commit comments

Comments
 (0)