Use view state to restore InteractiveSessionViewPane session by id (#…

…180732)
microsoft · Apr 24, 2023 · b292560 · b292560
1 parent 14a24d5
commit b292560
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 67 deletions.
diff --git a/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionEditor.ts b/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionEditor.ts
@@ -89,20 +89,20 @@ export class InteractiveSessionEditor extends EditorPane {
 			throw new Error('InteractiveSessionEditor lifecycle issue: no editor widget');
 		}
 
-		this.updateModel(editorModel.model, options);
+		this.updateModel(editorModel.model);
 	}
 
-	private updateModel(model: IInteractiveSessionModel, options: IInteractiveSessionEditorOptions): void {
+	private updateModel(model: IInteractiveSessionModel): void {
 		this._memento = new Memento('interactive-session-editor-' + model.sessionId, this.storageService);
 		this._viewState = this._memento.getMemento(StorageScope.WORKSPACE, StorageTarget.USER) as IViewState;
 		this.widget.setModel(model, { ...this._viewState });
 		const listener = model.onDidDispose(() => {
 			// TODO go back to swapping out the EditorInput when the session is restarted instead of this listener
 			listener.dispose();
-			const newModel = this.interactiveSessionService.startSession(options.providerId, false, CancellationToken.None);
+			const newModel = this.interactiveSessionService.startSession(model.providerId, CancellationToken.None);
 			if (newModel) {
 				(this.input as InteractiveSessionEditorInput).sessionId = newModel.sessionId;
-				this.updateModel(newModel, options);
+				this.updateModel(newModel);
 			}
 		});
 	}

diff --git a/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionEditorInput.ts b/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionEditorInput.ts
@@ -69,7 +69,7 @@ export class InteractiveSessionEditorInput extends EditorInput {
 	override async resolve(): Promise<InteractiveSessionEditorModel | null> {
 		const model = typeof this.sessionId === 'string' ?
 			this.interactiveSessionService.retrieveSession(this.sessionId) :
-			this.interactiveSessionService.startSession(this.options.providerId, false, CancellationToken.None);
+			this.interactiveSessionService.startSession(this.options.providerId, CancellationToken.None);
 
 		if (!model) {
 			return null;

diff --git a/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionViewPane.ts b/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionViewPane.ts
@@ -20,12 +20,17 @@ import { IViewPaneOptions, ViewPane } from 'vs/workbench/browser/parts/views/vie
 import { Memento } from 'vs/workbench/common/memento';
 import { IViewDescriptorService } from 'vs/workbench/common/views';
 import { IViewState, InteractiveSessionWidget } from 'vs/workbench/contrib/interactiveSession/browser/interactiveSessionWidget';
+import { IInteractiveSessionModel } from 'vs/workbench/contrib/interactiveSession/common/interactiveSessionModel';
 import { IInteractiveSessionService } from 'vs/workbench/contrib/interactiveSession/common/interactiveSessionService';
 
 export interface IInteractiveSessionViewOptions {
 	readonly providerId: string;
 }
 
+interface IViewPaneState extends IViewState {
+	sessionId?: string;
+}
+
 export const INTERACTIVE_SIDEBAR_PANEL_ID = 'workbench.panel.interactiveSessionSidebar';
 export class InteractiveSessionViewPane extends ViewPane {
 	static ID = 'workbench.panel.interactiveSession.view';
@@ -35,7 +40,7 @@ export class InteractiveSessionViewPane extends ViewPane {
 
 	private modelDisposables = this._register(new DisposableStore());
 	private memento: Memento;
-	private viewState: IViewState;
+	private viewState: IViewPaneState;
 
 	constructor(
 		private readonly interactiveSessionViewOptions: IInteractiveSessionViewOptions,
@@ -56,18 +61,19 @@ export class InteractiveSessionViewPane extends ViewPane {
 
 		// View state for the ViewPane is currently global per-provider basically, but some other strictly per-model state will require a separate memento.
 		this.memento = new Memento('interactive-session-view-' + this.interactiveSessionViewOptions.providerId, this.storageService);
-		this.viewState = this.memento.getMemento(StorageScope.WORKSPACE, StorageTarget.USER) as IViewState;
+		this.viewState = this.memento.getMemento(StorageScope.WORKSPACE, StorageTarget.USER) as IViewPaneState;
 	}
 
-	private updateModel(initial = false): void {
+	private updateModel(model?: IInteractiveSessionModel | undefined): void {
 		this.modelDisposables.clear();
 
-		const model = this.interactiveSessionService.startSession(this.interactiveSessionViewOptions.providerId, initial, CancellationToken.None);
+		model = model ?? this.interactiveSessionService.startSession(this.interactiveSessionViewOptions.providerId, CancellationToken.None);
 		if (!model) {
 			throw new Error('Could not start interactive session');
 		}
 
 		this._widget.setModel(model, { ...this.viewState });
+		this.viewState.sessionId = model.sessionId;
 		this.modelDisposables.add(model.onDidDispose(() => {
 			this.updateModel();
 		}));
@@ -84,7 +90,8 @@ export class InteractiveSessionViewPane extends ViewPane {
 		}));
 		this._widget.render(parent);
 
-		this.updateModel(true);
+		const initialModel = this.viewState.sessionId ? this.interactiveSessionService.retrieveSession(this.viewState.sessionId) : undefined;
+		this.updateModel(initialModel);
 	}
 
 	acceptInput(query?: string): void {

diff --git a/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionWidget.ts b/src/vs/workbench/contrib/interactiveSession/browser/interactiveSessionWidget.ts
@@ -436,14 +436,6 @@ export class InteractiveSessionWidget extends Disposable implements IInteractive
 		this.inputPart.saveState();
 		return { inputValue: this.inputPart.inputEditor.getValue() };
 	}
-
-	public override dispose(): void {
-		super.dispose();
-
-		if (this.viewModel) {
-			this.interactiveSessionService.releaseSession(this.viewModel.sessionId);
-		}
-	}
 }
 
 export class InteractiveSessionWidgetService implements IInteractiveSessionWidgetService {

diff --git a/src/vs/workbench/contrib/interactiveSession/common/interactiveSessionService.ts b/src/vs/workbench/contrib/interactiveSession/common/interactiveSessionService.ts
@@ -163,7 +163,7 @@ export interface IInteractiveSessionService {
 	_serviceBrand: undefined;
 	registerProvider(provider: IInteractiveProvider): IDisposable;
 	getProviderInfos(): IInteractiveProviderInfo[];
-	startSession(providerId: string, allowRestoringSession: boolean, token: CancellationToken): InteractiveSessionModel | undefined;
+	startSession(providerId: string, token: CancellationToken): InteractiveSessionModel | undefined;
 	retrieveSession(sessionId: string): IInteractiveSessionModel | undefined;
 
 	/**
@@ -176,7 +176,6 @@ export interface IInteractiveSessionService {
 	addInteractiveRequest(context: any): void;
 	addCompleteRequest(sessionId: string, message: string, response: IInteractiveSessionCompleteResponse): void;
 	sendInteractiveRequestToProvider(sessionId: string, message: IInteractiveSessionDynamicRequest): void;
-	releaseSession(sessionId: string): void;
 
 	onDidPerformUserAction: Event<IInteractiveSessionUserActionEvent>;
 	notifyUserAction(event: IInteractiveSessionUserActionEvent): void;

diff --git a/src/vs/workbench/contrib/interactiveSession/common/interactiveSessionServiceImpl.ts b/src/vs/workbench/contrib/interactiveSession/common/interactiveSessionServiceImpl.ts
@@ -109,7 +109,6 @@ export class InteractiveSessionService extends Disposable implements IInteractiv
 
 	private readonly _providers = new Map<string, IInteractiveProvider>();
 	private readonly _sessionModels = new Map<string, InteractiveSessionModel>();
-	private readonly _releasedSessions = new Set<string>();
 	private readonly _pendingRequests = new Map<string, CancelablePromise<void>>();
 	private readonly _unprocessedPersistedSessions: ISerializableInteractiveSessionsData;
 	private readonly _hasProvider: IContextKey<boolean>;
@@ -203,19 +202,9 @@ export class InteractiveSessionService extends Disposable implements IInteractiv
 		}
 	}
 
-	startSession(providerId: string, allowRestoringSession: boolean, token: CancellationToken): InteractiveSessionModel {
-		this.trace('startSession', `providerId=${providerId}, allowRestoringSession=${allowRestoringSession}`);
-
-		const restored = allowRestoringSession ? this.getNextRestoredSession(providerId) : undefined;
-		if (restored instanceof InteractiveSessionModel) {
-			this.trace('startSession', `Restored live session with id ${restored.sessionId}`);
-			return restored;
-		}
-
-		const someSessionHistory = restored;
-		this.trace('startSession', `Has history: ${!!someSessionHistory}. Including provider state: ${!!someSessionHistory?.providerState}`);
-
-		return this._startSession(providerId, someSessionHistory, token);
+	startSession(providerId: string, token: CancellationToken): InteractiveSessionModel {
+		this.trace('startSession', `providerId=${providerId}`);
+		return this._startSession(providerId, undefined, token);
 	}
 
 	private _startSession(providerId: string, someSessionHistory: ISerializableInteractiveSessionData | undefined, token: CancellationToken): InteractiveSessionModel {
@@ -227,7 +216,8 @@ export class InteractiveSessionService extends Disposable implements IInteractiv
 				model.dispose();
 				this._sessionModels.delete(model.sessionId);
 			}
-		}).catch(() => {
+		}).catch(err => {
+			this.trace('startSession', `initializeSession failed: ${err}`);
 			model.dispose();
 			this._sessionModels.delete(model.sessionId);
 		});
@@ -271,27 +261,7 @@ export class InteractiveSessionService extends Disposable implements IInteractiv
 		return model;
 	}
 
-	private getNextRestoredSession(providerId: string): InteractiveSessionModel | ISerializableInteractiveSessionData | undefined {
-		const releasedSessionId = Iterable.find(this._releasedSessions.values(), sessionId => this._sessionModels.get(sessionId)?.providerId === providerId);
-		if (typeof releasedSessionId === 'string') {
-			this._releasedSessions.delete(releasedSessionId);
-			return this._sessionModels.get(releasedSessionId);
-		}
-
-		const providerData = this._unprocessedPersistedSessions[providerId] ?? [];
-		return providerData.shift();
-	}
-
-	releaseSession(sessionId: string): void {
-		this.trace('releaseSession', `sessionId=${sessionId}`);
-		this._releasedSessions.add(sessionId);
-	}
-
 	retrieveSession(sessionId: string): InteractiveSessionModel | undefined {
-		if (this._releasedSessions.has(sessionId)) {
-			this._releasedSessions.delete(sessionId);
-		}
-
 		const model = this._sessionModels.get(sessionId);
 		if (model) {
 			return model;
@@ -514,7 +484,6 @@ export class InteractiveSessionService extends Disposable implements IInteractiv
 		model.dispose();
 		this._sessionModels.delete(sessionId);
 		this._pendingRequests.get(sessionId)?.cancel();
-		this._releasedSessions.delete(sessionId);
 	}
 
 	registerProvider(provider: IInteractiveProvider): IDisposable {

diff --git a/src/vs/workbench/contrib/interactiveSession/test/common/interactiveSessionService.test.ts b/src/vs/workbench/contrib/interactiveSession/test/common/interactiveSessionService.test.ts
@@ -73,18 +73,18 @@ suite('InteractiveSession', () => {
 		testDisposables.clear();
 	});
 
-	test('Restores state for the correct provider', async () => {
+	test('retrieveSession', async () => {
 		const testService = instantiationService.createInstance(InteractiveSessionService);
 		const provider1 = new SimpleTestProvider('provider1');
 		const provider2 = new SimpleTestProvider('provider2');
 		testService.registerProvider(provider1);
 		testService.registerProvider(provider2);
 
-		let session1 = testService.startSession('provider1', true, CancellationToken.None);
+		const session1 = testService.startSession('provider1', CancellationToken.None);
 		await session1.waitForInitialization();
 		session1!.addRequest('request 1');
 
-		let session2 = testService.startSession('provider2', true, CancellationToken.None);
+		const session2 = testService.startSession('provider2', CancellationToken.None);
 		await session2.waitForInitialization();
 		session2!.addRequest('request 2');
 
@@ -97,10 +97,10 @@ suite('InteractiveSession', () => {
 		const testService2 = instantiationService.createInstance(InteractiveSessionService);
 		testService2.registerProvider(provider1);
 		testService2.registerProvider(provider2);
-		session1 = testService2.startSession('provider1', true, CancellationToken.None);
-		await session1.waitForInitialization();
-		session2 = testService2.startSession('provider2', true, CancellationToken.None);
-		await session2.waitForInitialization();
+		const retrieved1 = testService2.retrieveSession(session1.sessionId);
+		await retrieved1!.waitForInitialization();
+		const retrieved2 = testService2.retrieveSession(session2.sessionId);
+		await retrieved2!.waitForInitialization();
 		assert.deepStrictEqual(provider1.lastInitialState, { state: 'provider1_state' });
 		assert.deepStrictEqual(provider2.lastInitialState, { state: 'provider2_state' });
 	});
@@ -127,7 +127,7 @@ suite('InteractiveSession', () => {
 		const provider1 = getFailProvider('provider1');
 		testService.registerProvider(provider1);
 
-		const session1 = testService.startSession('provider1', true, CancellationToken.None);
+		const session1 = testService.startSession('provider1', CancellationToken.None);
 		await assert.rejects(() => session1.waitForInitialization());
 	});
 
@@ -179,7 +179,7 @@ suite('InteractiveSession', () => {
 
 		testService.registerProvider(provider);
 
-		const model = testService.startSession('testProvider', true, CancellationToken.None);
+		const model = testService.startSession('testProvider', CancellationToken.None);
 		const commands = await testService.getSlashCommands(model.sessionId, CancellationToken.None);
 
 		assert.strictEqual(commands?.length, 1);
@@ -192,7 +192,7 @@ suite('InteractiveSession', () => {
 		const testService = instantiationService.createInstance(InteractiveSessionService);
 		testService.registerProvider(new SimpleTestProvider('testProvider'));
 
-		const model = testService.startSession('testProvider', true, CancellationToken.None);
+		const model = testService.startSession('testProvider', CancellationToken.None);
 		assert.strictEqual(model.getRequests().length, 0);
 
 		await testService.sendInteractiveRequestToProvider(model.sessionId, { message: 'test request' });
@@ -203,7 +203,7 @@ suite('InteractiveSession', () => {
 		const testService = instantiationService.createInstance(InteractiveSessionService);
 		testService.registerProvider(new SimpleTestProvider('testProvider'));
 
-		const model = testService.startSession('testProvider', true, CancellationToken.None);
+		const model = testService.startSession('testProvider', CancellationToken.None);
 		assert.strictEqual(model.getRequests().length, 0);
 
 		await testService.addCompleteRequest(model.sessionId, 'test request', { message: 'test response' });