Merge pull request microsoft#183538 from microsoft/joh/available-roundworm

joh/available roundworm

Author: jrieken

src/vs/editor/browser/editorExtensions.ts

@@ -450,9 +450,13 @@ export abstract class EditorAction2 extends Action2 {
 		// precondition does hold
 		return editor.invokeWithinContext((editorAccessor) => {
 			const kbService = editorAccessor.get(IContextKeyService);
-			if (kbService.contextMatchesRules(withNullAsUndefined(this.desc.precondition))) {
-				return this.runEditorCommand(editorAccessor, editor!, ...args);
+			const logService = editorAccessor.get(ILogService);
+			const enabled = kbService.contextMatchesRules(withNullAsUndefined(this.desc.precondition));
+			if (!enabled) {
+				logService.debug(`[EditorAction2] NOT running command because its precondition is FALSE`, this.desc.id, this.desc.precondition?.serialize());
+				return;
 			}
+			return this.runEditorCommand(editorAccessor, editor!, ...args);
 		});
 	}
 

src/vs/workbench/contrib/interactiveEditor/browser/interactiveEditor.contribution.ts

@@ -11,11 +11,15 @@ import { IInteractiveEditorService, INTERACTIVE_EDITOR_ID } from 'vs/workbench/c
 import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/common/extensions';
 import { InteractiveEditorServiceImpl } from 'vs/workbench/contrib/interactiveEditor/common/interactiveEditorServiceImpl';
 import { IInteractiveEditorSessionService, InteractiveEditorSessionService } from 'vs/workbench/contrib/interactiveEditor/browser/interactiveEditorSession';
+import { Registry } from 'vs/platform/registry/common/platform';
+import { IWorkbenchContributionsRegistry, Extensions } from 'vs/workbench/common/contributions';
+import { LifecyclePhase } from 'vs/workbench/services/lifecycle/common/lifecycle';
+import { InteractiveEditorNotebookContribution } from 'vs/workbench/contrib/interactiveEditor/browser/interactiveEditorNotebook';
 
 registerSingleton(IInteractiveEditorService, InteractiveEditorServiceImpl, InstantiationType.Delayed);
 registerSingleton(IInteractiveEditorSessionService, InteractiveEditorSessionService, InstantiationType.Delayed);
 
-registerEditorContribution(INTERACTIVE_EDITOR_ID, InteractiveEditorController, EditorContributionInstantiation.Lazy);
+registerEditorContribution(INTERACTIVE_EDITOR_ID, InteractiveEditorController, EditorContributionInstantiation.Eager); // EAGER because of notebook dispose/create of editors
 
 registerAction2(interactiveEditorActions.StartSessionAction);
 registerAction2(interactiveEditorActions.UnstashSessionAction);
@@ -42,3 +46,7 @@ registerAction2(interactiveEditorActions.FeebackUnhelpfulCommand);
 registerAction2(interactiveEditorActions.ApplyPreviewEdits);
 
 registerAction2(interactiveEditorActions.CopyRecordings);
+
+
+Registry.as<IWorkbenchContributionsRegistry>(Extensions.Workbench)
+	.registerWorkbenchContribution(InteractiveEditorNotebookContribution, LifecyclePhase.Restored);

src/vs/workbench/contrib/interactiveEditor/browser/interactiveEditorController.ts

@@ -130,16 +130,26 @@ export class InteractiveEditorController implements IEditorContribution {
 				return;
 			}
 
-			this._logService.trace('[IE] session RESUMING');
+			this._log('session RESUMING', e);
 			await this._nextState(State.CREATE_SESSION, { existingSession });
-			this._logService.trace('[IE] session done or paused');
+			this._log('session done or paused');
 		}));
+		this._log('NEW controller');
 	}
 
 	dispose(): void {
 		this._stashedSession.clear();
 		this._finishExistingSession();
 		this._store.dispose();
+		this._log('controller disposed');
+	}
+
+	private _log(message: string | Error, ...more: any[]): void {
+		if (message instanceof Error) {
+			this._logService.error(message, ...more);
+		} else {
+			this._logService.trace(`[IE] (editor:${this._editor.getId()})${message}`, ...more);
+		}
 	}
 
 	getId(): string {
@@ -161,21 +171,21 @@ export class InteractiveEditorController implements IEditorContribution {
 	}
 
 	async run(options: InteractiveEditorRunOptions | undefined): Promise<void> {
-		this._logService.trace('[IE] session starting');
+		this._log('session starting');
 		await this._finishExistingSession();
 		this._stashedSession.clear();
 
 		await this._nextState(State.CREATE_SESSION, options);
-		this._logService.trace('[IE] session done or paused');
+		this._log('session done or paused');
 	}
 
 	private async _finishExistingSession(): Promise<void> {
 		if (this._activeSession) {
 			if (this._activeSession.editMode === EditMode.Preview) {
-				this._logService.trace('[IE] finishing existing session, using CANCEL', this._activeSession.editMode);
+				this._log('finishing existing session, using CANCEL', this._activeSession.editMode);
 				await this.cancelSession();
 			} else {
-				this._logService.trace('[IE] finishing existing session, using APPLY', this._activeSession.editMode);
+				this._log('finishing existing session, using APPLY', this._activeSession.editMode);
 				await this.applyChanges();
 			}
 		}
@@ -184,7 +194,7 @@ export class InteractiveEditorController implements IEditorContribution {
 	// ---- state machine
 
 	protected async _nextState(state: State, options: InteractiveEditorRunOptions | undefined): Promise<void> {
-		this._logService.trace('[IE] setState to ', state);
+		this._log('setState to ', state);
 		const nextState = await this[state](options);
 		if (nextState) {
 			await this._nextState(nextState, options);
@@ -200,7 +210,7 @@ export class InteractiveEditorController implements IEditorContribution {
 		if (!session) {
 			const createSessionCts = new CancellationTokenSource();
 			const msgListener = Event.once(this._messages.event)(m => {
-				this._logService.trace('[IE](state=_createSession) message received', m);
+				this._log('state=_createSession) message received', m);
 				createSessionCts.cancel();
 			});
 
@@ -261,11 +271,12 @@ export class InteractiveEditorController implements IEditorContribution {
 		this._zone.widget.updateInfo(this._activeSession.session.message ?? localize('welcome.1', "AI-generated code may be incorrect"));
 		this._zone.show(this._activeSession.wholeRange.getEndPosition());
 
-		this._sessionStore.add(this._editor.onDidChangeModel(() => {
-			this._messages.fire(this._activeSession?.lastExchange
+		this._sessionStore.add(this._editor.onDidChangeModel((e) => {
+			const msg = this._activeSession?.lastExchange
 				? Message.PAUSE_SESSION // pause when switching models/tabs and when having a previous exchange
-				: Message.CANCEL_SESSION
-			);
+				: Message.CANCEL_SESSION;
+			this._log('model changed, pause or cancel session', msg, e);
+			this._messages.fire(msg);
 		}));
 
 		this._sessionStore.add(this._editor.onDidChangeModelContent(e => {
@@ -282,7 +293,7 @@ export class InteractiveEditorController implements IEditorContribution {
 			this._activeSession!.recordExternalEditOccurred(editIsOutsideOfWholeRange);
 
 			if (editIsOutsideOfWholeRange) {
-				this._logService.info('[IE] text changed outside of whole range, FINISH session');
+				this._log('text changed outside of whole range, FINISH session');
 				this._finishExistingSession();
 			}
 		}));
@@ -363,7 +374,7 @@ export class InteractiveEditorController implements IEditorContribution {
 		} else {
 			const barrier = new Barrier();
 			const msgListener = Event.once(this._messages.event)(m => {
-				this._logService.trace('[IE](state=_waitForInput) message received', m);
+				this._log('state=_waitForInput) message received', m);
 				message = m;
 				barrier.open();
 			});
@@ -397,7 +408,7 @@ export class InteractiveEditorController implements IEditorContribution {
 
 		const refer = this._activeSession.session.slashCommands?.some(value => value.refer && input!.startsWith(`/${value.command}`));
 		if (refer) {
-			this._logService.info('[IE] seeing refer command, continuing outside editor', this._activeSession.provider.debugName);
+			this._log('[IE] seeing refer command, continuing outside editor', this._activeSession.provider.debugName);
 			this._editor.setSelection(this._activeSession.wholeRange);
 			this._instaService.invokeFunction(sendRequest, input);
 
@@ -421,7 +432,7 @@ export class InteractiveEditorController implements IEditorContribution {
 
 		let message = Message.NONE;
 		const msgListener = Event.once(this._messages.event)(m => {
-			this._logService.trace('[IE](state=_makeRequest) message received', m);
+			this._log('state=_makeRequest) message received', m);
 			message = m;
 			requestCts.cancel();
 		});
@@ -438,7 +449,7 @@ export class InteractiveEditorController implements IEditorContribution {
 			attempt: 0,
 		};
 		const task = this._activeSession.provider.provideResponse(this._activeSession.session, request, requestCts.token);
-		this._logService.trace('[IE] request started', this._activeSession.provider.debugName, this._activeSession.session, request);
+		this._log('request started', this._activeSession.provider.debugName, this._activeSession.session, request);
 
 		let response: EditResponse | MarkdownResponse | ErrorResponse | EmptyResponse;
 		let reply: IInteractiveEditorResponse | null | undefined;
@@ -463,7 +474,7 @@ export class InteractiveEditorController implements IEditorContribution {
 			this._ctxHasActiveRequest.set(false);
 			this._zone.widget.updateProgress(false);
 			this._zone.widget.updateInfo('');
-			this._logService.trace('[IE] request took', sw.elapsed(), this._activeSession.provider.debugName);
+			this._log('request took', sw.elapsed(), this._activeSession.provider.debugName);
 
 		}
 
@@ -497,7 +508,7 @@ export class InteractiveEditorController implements IEditorContribution {
 			}
 			const moreMinimalEdits = (await this._editorWorkerService.computeHumanReadableDiff(this._activeSession.textModelN.uri, response.localEdits));
 			const editOperations = (moreMinimalEdits ?? response.localEdits).map(edit => EditOperation.replace(Range.lift(edit.range), edit.text));
-			this._logService.trace('[IE] edits from PROVIDER and after making them MORE MINIMAL', this._activeSession.provider.debugName, response.localEdits, moreMinimalEdits);
+			this._log('edits from PROVIDER and after making them MORE MINIMAL', this._activeSession.provider.debugName, response.localEdits, moreMinimalEdits);
 
 			const textModelNplus1 = this._modelService.createModel(createTextBufferFactoryFromSnapshot(this._activeSession.textModelN.createSnapshot()), null, undefined, true);
 			textModelNplus1.applyEdits(editOperations);
@@ -682,8 +693,8 @@ export class InteractiveEditorController implements IEditorContribution {
 				await strategy?.apply();
 			} catch (err) {
 				this._dialogService.error(localize('err.apply', "Failed to apply changes.", toErrorMessage(err)));
-				this._logService.error('[IE] FAILED to apply changes');
-				this._logService.error(err);
+				this._log('FAILED to apply changes');
+				this._log(err);
 			}
 			strategy?.dispose();
 			this._messages.fire(Message.ACCEPT_SESSION);
@@ -702,8 +713,8 @@ export class InteractiveEditorController implements IEditorContribution {
 			await strategy?.cancel();
 		} catch (err) {
 			this._dialogService.error(localize('err.discard', "Failed to discard changes.", toErrorMessage(err)));
-			this._logService.error('[IE] FAILED to discard changes');
-			this._logService.error(err);
+			this._log('FAILED to discard changes');
+			this._log(err);
 		}
 		strategy?.dispose();
 		this._messages.fire(Message.CANCEL_SESSION);

src/vs/workbench/contrib/interactiveEditor/browser/interactiveEditorNotebook.ts

@@ -0,0 +1,35 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { illegalState } from 'vs/base/common/errors';
+import { Schemas } from 'vs/base/common/network';
+import { isEqual } from 'vs/base/common/resources';
+import { IInteractiveEditorSessionService } from 'vs/workbench/contrib/interactiveEditor/browser/interactiveEditorSession';
+import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/services/notebookEditorService';
+import { CellUri } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+
+export class InteractiveEditorNotebookContribution {
+
+	constructor(
+		@IInteractiveEditorSessionService sessionService: IInteractiveEditorSessionService,
+		@INotebookEditorService notebookEditorService: INotebookEditorService,
+	) {
+
+		sessionService.registerSessionKeyComputer(Schemas.vscodeNotebookCell, {
+			getComparisonKey: (_editor, uri) => {
+				const data = CellUri.parse(uri);
+				if (!data) {
+					throw illegalState('Expected notebook');
+				}
+				for (const editor of notebookEditorService.listNotebookEditors()) {
+					if (isEqual(editor.textModel?.uri, data.notebook)) {
+						return `<notebook>${editor.getId()}#${uri}`;
+					}
+				}
+				throw illegalState('Expected notebook');
+			}
+		});
+	}
+}

src/vs/workbench/contrib/interactiveEditor/browser/interactiveEditorSession.ts

@@ -13,7 +13,6 @@ import { EditMode, IInteractiveEditorSessionProvider, IInteractiveEditorSession,
 import { IRange, Range } from 'vs/editor/common/core/range';
 import { IActiveCodeEditor, ICodeEditor } from 'vs/editor/browser/editorBrowser';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
-import { ResourceMap } from 'vs/base/common/map';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { IModelService } from 'vs/editor/common/services/model';
 import { ITextModelService } from 'vs/editor/common/services/resolverService';
@@ -266,6 +265,10 @@ export class EditResponse {
 	}
 }
 
+export interface ISessionKeyComputer {
+	getComparisonKey(editor: ICodeEditor, uri: URI): string;
+}
+
 export const IInteractiveEditorSessionService = createDecorator<IInteractiveEditorSessionService>('IInteractiveEditorSessionService');
 
 export interface IInteractiveEditorSessionService {
@@ -277,6 +280,8 @@ export interface IInteractiveEditorSessionService {
 
 	releaseSession(session: Session): void;
 
+	registerSessionKeyComputer(scheme: string, value: ISessionKeyComputer): IDisposable;
+
 	//
 
 	recordings(): readonly Recording[];
@@ -291,7 +296,8 @@ export class InteractiveEditorSessionService implements IInteractiveEditorSessio
 
 	declare _serviceBrand: undefined;
 
-	private readonly _sessions = new Map<ICodeEditor, ResourceMap<SessionData>>();
+	private readonly _sessions = new Map<string, SessionData>();
+	private readonly _keyComputers = new Map<string, ISessionKeyComputer>();
 	private _recordings: Recording[] = [];
 
 	constructor(
@@ -360,35 +366,27 @@ export class InteractiveEditorSessionService implements IInteractiveEditorSessio
 
 		const session = new Session(options.editMode, editor, textModel0, textModel, provider, raw, wholeRangeDecorationId);
 
-		// store: editor -> uri -> session
-		let map = this._sessions.get(editor);
-		if (!map) {
-			map = new ResourceMap<SessionData>();
-			this._sessions.set(editor, map);
+		// store: key -> session
+		const key = this._key(editor, textModel.uri);
+		if (this._sessions.has(key)) {
+			store.dispose();
+			throw new Error(`Session already stored for ${key}`);
 		}
-		if (map.has(textModel.uri)) {
-			throw new Error(`Session already stored for ${textModel.uri}`);
-		}
-		map.set(textModel.uri, { session, store });
+		this._sessions.set(key, { session, store });
 		return session;
 	}
 
 	releaseSession(session: Session): void {
 
-		const { editor, textModelN } = session;
+		const { editor } = session;
 
 		// cleanup
-		const map = this._sessions.get(editor);
-		if (map) {
-			const data = map.get(textModelN.uri);
-			if (data) {
-				data.store.dispose();
-				data.session.session.dispose?.();
-
-				map.delete(textModelN.uri);
-			}
-			if (map.size === 0) {
-				this._sessions.delete(editor);
+		for (const [key, value] of this._sessions) {
+			if (value.session === session) {
+				value.store.dispose();
+				this._sessions.delete(key);
+				this._logService.trace(`[IE] did RELEASED session for ${editor.getId()}, ${session.provider.debugName}`);
+				break;
 			}
 		}
 
@@ -403,7 +401,21 @@ export class InteractiveEditorSessionService implements IInteractiveEditorSessio
 	}
 
 	getSession(editor: ICodeEditor, uri: URI): Session | undefined {
-		return this._sessions.get(editor)?.get(uri)?.session;
+		const key = this._key(editor, uri);
+		return this._sessions.get(key)?.session;
+	}
+
+	private _key(editor: ICodeEditor, uri: URI): string {
+		const item = this._keyComputers.get(uri.scheme);
+		return item
+			? item.getComparisonKey(editor, uri)
+			: `${editor.getId()}@${uri.toString()}`;
+
+	}
+
+	registerSessionKeyComputer(scheme: string, value: ISessionKeyComputer): IDisposable {
+		this._keyComputers.set(scheme, value);
+		return toDisposable(() => this._keyComputers.delete(scheme));
 	}
 
 	// --- debug