uses /edit editor UX for "Apply In Editor" (#237944)

* WIP

* uses /edit editor UX for "Apply In Editor"

fixes microsoft/vscode-copilot#8577

src/vs/base/common/iterator.ts

@@ -3,6 +3,8 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { isIterable } from './types.js';
+
 export namespace Iterable {
 
 	export function is<T = any>(thing: any): thing is Iterable<T> {
@@ -90,9 +92,13 @@ export namespace Iterable {
 		}
 	}
 
-	export function* concat<T>(...iterables: Iterable<T>[]): Iterable<T> {
-		for (const iterable of iterables) {
-			yield* iterable;
+	export function* concat<T>(...iterables: (Iterable<T> | T)[]): Iterable<T> {
+		for (const item of iterables) {
+			if (isIterable(item)) {
+				yield* item;
+			} else {
+				yield item;
+			}
 		}
 	}
 

src/vs/workbench/contrib/chat/browser/actions/chatCodeblockActions.ts

@@ -558,7 +558,7 @@ export function registerChatCodeCompareBlockActions() {
 				const inlineChatController = InlineChatController.get(editorToApply);
 				if (inlineChatController) {
 					editorToApply.revealLineInCenterIfOutsideViewport(firstEdit.range.startLineNumber);
-					inlineChatController.reviewEdits(firstEdit.range, textEdits, CancellationToken.None);
+					inlineChatController.reviewEdits(textEdits, CancellationToken.None);
 					response.setEditApplied(item, 1);
 					return true;
 				}

src/vs/workbench/contrib/chat/browser/actions/codeBlockOperations.ts

@@ -266,7 +266,7 @@ export class ApplyCodeBlockOperation {
 		const inlineChatController = InlineChatController.get(codeEditor);
 		if (inlineChatController) {
 			let isOpen = true;
-			const promise = inlineChatController.reviewEdits(codeEditor.getSelection(), edits, tokenSource.token);
+			const promise = inlineChatController.reviewEdits(edits, tokenSource.token);
 			promise.finally(() => {
 				isOpen = false;
 				tokenSource.dispose();

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedFileEntry.ts

@@ -134,6 +134,8 @@ export class ChatEditingModifiedFileEntry extends Disposable implements IModifie
 
 	private readonly _diffTrimWhitespace: IObservable<boolean>;
 
+	private _refCounter: number = 1;
+
 	constructor(
 		resourceRef: IReference<IResolvedTextEditorModel>,
 		private readonly _multiDiffEntryDelegate: { collapse: (transaction: ITransaction | undefined) => void },
@@ -199,6 +201,17 @@ export class ChatEditingModifiedFileEntry extends Disposable implements IModifie
 		}));
 	}
 
+	override dispose(): void {
+		if (--this._refCounter === 0) {
+			super.dispose();
+		}
+	}
+
+	acquire() {
+		this._refCounter++;
+		return this;
+	}
+
 	private _clearCurrentEditLineDecoration() {
 		this._editDecorations = this.doc.deltaDecorations(this._editDecorations, []);
 	}

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingService.ts

@@ -9,9 +9,11 @@ import { CancellationToken, CancellationTokenSource } from '../../../../../base/
 import { Codicon } from '../../../../../base/common/codicons.js';
 import { BugIndicatingError } from '../../../../../base/common/errors.js';
 import { Emitter, Event } from '../../../../../base/common/event.js';
+import { Iterable } from '../../../../../base/common/iterator.js';
 import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../../base/common/lifecycle.js';
+import { LinkedList } from '../../../../../base/common/linkedList.js';
 import { ResourceMap } from '../../../../../base/common/map.js';
-import { derived, IObservable, observableValue, runOnChange, ValueWithChangeEventFromObservable } from '../../../../../base/common/observable.js';
+import { derived, IObservable, observableValue, observableValueOpts, runOnChange, ValueWithChangeEventFromObservable } from '../../../../../base/common/observable.js';
 import { compare } from '../../../../../base/common/strings.js';
 import { ThemeIcon } from '../../../../../base/common/themables.js';
 import { isString } from '../../../../../base/common/types.js';
@@ -37,6 +39,7 @@ import { ChatContextKeys } from '../../common/chatContextKeys.js';
 import { applyingChatEditsContextKey, applyingChatEditsFailedContextKey, CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingAgentSupportsReadonlyReferencesContextKey, chatEditingMaxFileAssignmentName, chatEditingResourceContextKey, ChatEditingSessionState, decidedChatEditingResourceContextKey, defaultChatEditingMaxFileLimit, hasAppliedChatEditsContextKey, hasUndecidedChatEditingResourceContextKey, IChatEditingService, IChatEditingSession, IChatEditingSessionStream, IChatRelatedFile, IChatRelatedFilesProvider, IModifiedFileEntry, inChatEditingSessionContextKey, WorkingSetEntryState } from '../../common/chatEditingService.js';
 import { IChatResponseModel, IChatTextEditGroup } from '../../common/chatModel.js';
 import { IChatService } from '../../common/chatService.js';
+import { ChatEditingModifiedFileEntry } from './chatEditingModifiedFileEntry.js';
 import { ChatEditingSession } from './chatEditingSession.js';
 import { ChatEditingSnapshotTextModelContentProvider, ChatEditingTextModelContentProvider } from './chatEditingTextModelContentProviders.js';
 
@@ -50,6 +53,17 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 	private readonly _currentSessionObs = observableValue<ChatEditingSession | null>(this, null);
 	private readonly _currentSessionDisposables = this._register(new DisposableStore());
 
+	private readonly _adhocSessionsObs = observableValueOpts<LinkedList<ChatEditingSession>>({ equalsFn: (a, b) => false }, new LinkedList());
+
+	readonly editingSessionsObs: IObservable<readonly IChatEditingSession[]> = derived(r => {
+		const result = Array.from(this._adhocSessionsObs.read(r));
+		const globalSession = this._currentSessionObs.read(r);
+		if (globalSession) {
+			result.push(globalSession);
+		}
+		return result;
+	});
+
 	private readonly _currentAutoApplyOperationObs = observableValue<CancellationTokenSource | null>(this, null);
 	get currentAutoApplyOperation(): CancellationTokenSource | null {
 		return this._currentAutoApplyOperationObs.get();
@@ -63,9 +77,6 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 		return this._currentSessionObs;
 	}
 
-	private readonly _onDidChangeEditingSession = this._register(new Emitter<void>());
-	public readonly onDidChangeEditingSession = this._onDidChangeEditingSession.event;
-
 	private _editingSessionFileLimitPromise: Promise<number>;
 	private _editingSessionFileLimit: number | undefined;
 	get editingSessionFileLimit() {
@@ -111,13 +122,14 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 			return decidedEntries.map(entry => entry.entryId);
 		}));
 		this._register(bindContextKey(hasUndecidedChatEditingResourceContextKey, contextKeyService, (reader) => {
-			const currentSession = this._currentSessionObs.read(reader);
-			if (!currentSession) {
-				return;
+
+			for (const session of this.editingSessionsObs.read(reader)) {
+				const entries = session.entries.read(reader);
+				const decidedEntries = entries.filter(entry => entry.state.read(reader) === WorkingSetEntryState.Modified);
+				return decidedEntries.length > 0;
 			}
-			const entries = currentSession.entries.read(reader);
-			const decidedEntries = entries.filter(entry => entry.state.read(reader) === WorkingSetEntryState.Modified);
-			return decidedEntries.length > 0;
+
+			return false;
 		}));
 		this._register(bindContextKey(hasAppliedChatEditsContextKey, contextKeyService, (reader) => {
 			const currentSession = this._currentSessionObs.read(reader);
@@ -211,14 +223,26 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 	}
 
 
+	private _lookupEntry(uri: URI): ChatEditingModifiedFileEntry | undefined {
+
+		for (const item of Iterable.concat(this.editingSessionsObs.get())) {
+			const candidate = item.getEntry(uri);
+			if (candidate instanceof ChatEditingModifiedFileEntry) {
+				// make sure to ref-count this object
+				return candidate.acquire();
+			}
+		}
+		return undefined;
+	}
+
 	private async _createEditingSession(chatSessionId: string): Promise<IChatEditingSession> {
 		if (this._currentSessionObs.get()) {
 			throw new BugIndicatingError('Cannot have more than one active editing session');
 		}
 
 		this._currentSessionDisposables.clear();
 
-		const session = this._instantiationService.createInstance(ChatEditingSession, chatSessionId, this._editingSessionFileLimitPromise);
+		const session = this._instantiationService.createInstance(ChatEditingSession, chatSessionId, this._editingSessionFileLimitPromise, this._lookupEntry.bind(this));
 		await session.init();
 
 		// listen for completed responses, run the code mapper and apply the edits to this edit session
@@ -227,14 +251,33 @@ export class ChatEditingService extends Disposable implements IChatEditingServic
 		this._currentSessionDisposables.add(session.onDidDispose(() => {
 			this._currentSessionDisposables.clear();
 			this._currentSessionObs.set(null, undefined);
-			this._onDidChangeEditingSession.fire();
-		}));
-		this._currentSessionDisposables.add(session.onDidChange(() => {
-			this._onDidChangeEditingSession.fire();
 		}));
 
 		this._currentSessionObs.set(session, undefined);
-		this._onDidChangeEditingSession.fire();
+		return session;
+	}
+
+	async createAdhocEditingSession(chatSessionId: string): Promise<IChatEditingSession & IDisposable> {
+		const session = this._instantiationService.createInstance(ChatEditingSession, chatSessionId, this._editingSessionFileLimitPromise, this._lookupEntry.bind(this));
+		await session.init();
+
+		const list = this._adhocSessionsObs.get();
+		const removeSession = list.unshift(session);
+
+		const store = new DisposableStore();
+		this._store.add(store);
+
+		store.add(this.installAutoApplyObserver(session));
+
+		store.add(session.onDidDispose(e => {
+			removeSession();
+			this._adhocSessionsObs.set(list, undefined);
+			this._store.deleteAndLeak(store);
+			store.dispose();
+		}));
+
+		this._adhocSessionsObs.set(list, undefined);
+
 		return session;
 	}
 

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts

@@ -163,6 +163,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
 	constructor(
 		public readonly chatSessionId: string,
 		private editingSessionFileLimitPromise: Promise<number>,
+		private _lookupExternalEntry: (uri: URI) => ChatEditingModifiedFileEntry | undefined,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@IModelService private readonly _modelService: IModelService,
 		@ILanguageService private readonly _languageService: ILanguageService,
@@ -670,23 +671,39 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
 			}
 			return existingEntry;
 		}
-		const initialContent = this._initialFileContents.get(resource);
-		// This gets manually disposed in .dispose() or in .restoreSnapshot()
-		const entry = await this._createModifiedFileEntry(resource, responseModel, false, initialContent);
-		if (!initialContent) {
-			this._initialFileContents.set(resource, entry.initialContent);
+
+		let entry: ChatEditingModifiedFileEntry;
+		const existingExternalEntry = this._lookupExternalEntry(resource);
+		if (existingExternalEntry) {
+			entry = existingExternalEntry;
+		} else {
+			const initialContent = this._initialFileContents.get(resource);
+			// This gets manually disposed in .dispose() or in .restoreSnapshot()
+			entry = await this._createModifiedFileEntry(resource, responseModel, false, initialContent);
+			if (!initialContent) {
+				this._initialFileContents.set(resource, entry.initialContent);
+			}
 		}
+
 		// If an entry is deleted e.g. reverting a created file,
 		// remove it from the entries and don't show it in the working set anymore
 		// so that it can be recreated e.g. through retry
-		this._register(entry.onDidDelete(() => {
+		const listener = entry.onDidDelete(() => {
 			const newEntries = this._entriesObs.get().filter(e => !isEqual(e.modifiedURI, entry.modifiedURI));
 			this._entriesObs.set(newEntries, undefined);
 			this._workingSet.delete(entry.modifiedURI);
 			this._editorService.closeEditors(this._editorService.findEditors(entry.modifiedURI));
-			entry.dispose();
+
+			if (!existingExternalEntry) {
+				// don't dispose entries that are not yours!
+				entry.dispose();
+			}
+
+			this._store.delete(listener);
 			this._onDidChange.fire(ChatEditingSessionChangeType.WorkingSet);
-		}));
+		});
+		this._store.add(listener);
+
 		const entriesArr = [...this._entriesObs.get(), entry];
 		this._entriesObs.set(entriesArr, undefined);
 		this._onDidChange.fire(ChatEditingSessionChangeType.WorkingSet);

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

@@ -61,14 +61,15 @@ abstract class NavigateAction extends Action2 {
 		if (!isCodeEditor(editor) || !editor.hasModel()) {
 			return;
 		}
-
-		const session = chatEditingService.currentEditingSession;
-		if (!session) {
+		const ctrl = ChatEditorController.get(editor);
+		if (!ctrl) {
 			return;
 		}
 
-		const ctrl = ChatEditorController.get(editor);
-		if (!ctrl) {
+		const session = chatEditingService.editingSessionsObs.get()
+			.find(candidate => candidate.getEntry(editor.getModel().uri));
+
+		if (!session) {
 			return;
 		}
 
@@ -167,7 +168,9 @@ abstract class AcceptDiscardAction extends Action2 {
 			return;
 		}
 
-		const session = chatEditingService.currentEditingSession;
+		const session = chatEditingService.editingSessionsObs.get()
+			.find(candidate => candidate.getEntry(uri));
+
 		if (!session) {
 			return;
 		}

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

@@ -73,8 +73,8 @@ export class ChatEditorController extends Disposable implements IEditorContribut
 
 	constructor(
 		private readonly _editor: ICodeEditor,
-		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@IChatEditingService private readonly _chatEditingService: IChatEditingService,
+		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@IChatAgentService private readonly _chatAgentService: IChatAgentService,
 		@IEditorService private readonly _editorService: IEditorService,
 		@IContextKeyService contextKeyService: IContextKeyService,
@@ -92,27 +92,31 @@ export class ChatEditorController extends Disposable implements IEditorContribut
 
 
 		this._store.add(autorun(r => {
-			const session = this._chatEditingService.currentEditingSessionObs.read(r);
-			this._ctxRequestInProgress.set(session?.state.read(r) === ChatEditingSessionState.StreamingEdits);
+			let isStreamingEdits = false;
+			for (const session of _chatEditingService.editingSessionsObs.read(r)) {
+				isStreamingEdits ||= session.state.read(r) === ChatEditingSessionState.StreamingEdits;
+			}
+			this._ctxRequestInProgress.set(isStreamingEdits);
 		}));
 
-
 		const entryForEditor = derived(r => {
 			const model = modelObs.read(r);
-			const session = this._chatEditingService.currentEditingSessionObs.read(r);
-			if (!session) {
-				return undefined;
+			if (!model) {
+				return;
 			}
 
-			const entries = session.entries.read(r);
-			const idx = model?.uri
-				? entries.findIndex(e => isEqual(e.modifiedURI, model.uri))
-				: -1;
+			for (const session of _chatEditingService.editingSessionsObs.read(r)) {
+				const entries = session.entries.read(r);
+				const idx = model?.uri
+					? entries.findIndex(e => isEqual(e.modifiedURI, model.uri))
+					: -1;
 
-			if (idx < 0) {
-				return undefined;
+				if (idx >= 0) {
+					return { session, entry: entries[idx], entries, idx };
+				}
 			}
-			return { session, entry: entries[idx], entries, idx };
+
+			return undefined;
 		});
 
 
@@ -185,13 +189,17 @@ export class ChatEditorController extends Disposable implements IEditorContribut
 		// ---- readonly while streaming
 
 		const shouldBeReadOnly = derived(this, r => {
-			const value = this._chatEditingService.currentEditingSessionObs.read(r);
-			if (!value || value.state.read(r) !== ChatEditingSessionState.StreamingEdits) {
-				return false;
-			}
 
 			const model = modelObs.read(r);
-			return model ? value.readEntry(model.uri, r) : undefined;
+			if (!model) {
+				return undefined;
+			}
+			for (const session of _chatEditingService.editingSessionsObs.read(r)) {
+				if (session.readEntry(model.uri, r) && session.state.read(r) === ChatEditingSessionState.StreamingEdits) {
+					return true;
+				}
+			}
+			return false;
 		});
 
 
@@ -570,7 +578,14 @@ export class ChatEditorController extends Disposable implements IEditorContribut
 			return;
 		}
 
-		const entry = this._chatEditingService.currentEditingSessionObs.get()?.getEntry(this._editor.getModel().uri);
+		let entry: IModifiedFileEntry | undefined;
+		for (const session of this._chatEditingService.editingSessionsObs.get()) {
+			entry = session.getEntry(this._editor.getModel().uri);
+			if (entry) {
+				break;
+			}
+		}
+
 		if (!entry) {
 			return;
 		}