chatEditingServiceImpl.ts

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import { compareBy, delta } from '../../../../../base/common/arrays.js';
import { CancellationToken } from '../../../../../base/common/cancellation.js';
import { Codicon } from '../../../../../base/common/codicons.js';
import { groupBy } from '../../../../../base/common/collections.js';
import { ErrorNoTelemetry } from '../../../../../base/common/errors.js';
import { Emitter, Event } from '../../../../../base/common/event.js';
import { Iterable } from '../../../../../base/common/iterator.js';
import { Disposable, DisposableStore, dispose, IDisposable } from '../../../../../base/common/lifecycle.js';
import { LinkedList } from '../../../../../base/common/linkedList.js';
import { ResourceMap } from '../../../../../base/common/map.js';
import { Schemas } from '../../../../../base/common/network.js';
import { derived, IObservable, observableValueOpts, runOnChange, ValueWithChangeEventFromObservable } from '../../../../../base/common/observable.js';
import { isEqual } from '../../../../../base/common/resources.js';
import { compare } from '../../../../../base/common/strings.js';
import { ThemeIcon } from '../../../../../base/common/themables.js';
import { assertType } from '../../../../../base/common/types.js';
import { URI } from '../../../../../base/common/uri.js';
import { TextEdit } from '../../../../../editor/common/languages.js';
import { ITextModelService } from '../../../../../editor/common/services/resolverService.js';
import { localize } from '../../../../../nls.js';
import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
import { IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js';
import { IFileService } from '../../../../../platform/files/common/files.js';
import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
import { ILogService } from '../../../../../platform/log/common/log.js';
import { IProductService } from '../../../../../platform/product/common/productService.js';
import { IStorageService } from '../../../../../platform/storage/common/storage.js';
import { IDecorationData, IDecorationsProvider, IDecorationsService } from '../../../../services/decorations/common/decorations.js';
import { IEditorService } from '../../../../services/editor/common/editorService.js';
import { IExtensionService } from '../../../../services/extensions/common/extensions.js';
import { ILifecycleService } from '../../../../services/lifecycle/common/lifecycle.js';
import { IMultiDiffSourceResolver, IMultiDiffSourceResolverService, IResolvedMultiDiffSource, MultiDiffEditorItem } from '../../../multiDiffEditor/browser/multiDiffSourceResolverService.js';
import { CellUri, ICellEditOperation } from '../../../notebook/common/notebookCommon.js';
import { INotebookService } from '../../../notebook/common/notebookService.js';
import { CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingAgentSupportsReadonlyReferencesContextKey, chatEditingResourceContextKey, ChatEditingSessionState, IChatEditingService, IChatEditingSession, IModifiedFileEntry, inChatEditingSessionContextKey, IStreamingEdits, ModifiedFileEntryState, parseChatMultiDiffUri } from '../../common/editing/chatEditingService.js';
import { ChatModel, ICellTextEditOperation, IChatResponseModel, isCellTextEditOperationArray } from '../../common/model/chatModel.js';
import { IChatService } from '../../common/chatService/chatService.js';
import { ChatEditorInput } from '../widgetHosts/editor/chatEditorInput.js';
import { AbstractChatEditingModifiedFileEntry } from './chatEditingModifiedFileEntry.js';
import { ChatEditingSession } from './chatEditingSession.js';
import { ChatEditingSnapshotTextModelContentProvider, ChatEditingTextModelContentProvider } from './chatEditingTextModelContentProviders.js';

export class ChatEditingService extends Disposable implements IChatEditingService {

	_serviceBrand: undefined;


	private readonly _sessionsObs = observableValueOpts<LinkedList<ChatEditingSession>>({ equalsFn: (a, b) => false }, new LinkedList());

	readonly editingSessionsObs: IObservable<readonly IChatEditingSession[]> = derived(r => {
		const result = Array.from(this._sessionsObs.read(r));
		return result;
	});

	constructor(
		@IInstantiationService private readonly _instantiationService: IInstantiationService,
		@IMultiDiffSourceResolverService multiDiffSourceResolverService: IMultiDiffSourceResolverService,
		@ITextModelService textModelService: ITextModelService,
		@IContextKeyService contextKeyService: IContextKeyService,
		@IChatService private readonly _chatService: IChatService,
		@IEditorService private readonly _editorService: IEditorService,
		@IDecorationsService decorationsService: IDecorationsService,
		@IFileService private readonly _fileService: IFileService,
		@ILifecycleService private readonly lifecycleService: ILifecycleService,
		@IStorageService storageService: IStorageService,
		@ILogService logService: ILogService,
		@IExtensionService extensionService: IExtensionService,
		@IProductService productService: IProductService,
		@INotebookService private readonly notebookService: INotebookService,
		@IConfigurationService private readonly _configurationService: IConfigurationService,
	) {
		super();
		this._register(decorationsService.registerDecorationsProvider(_instantiationService.createInstance(ChatDecorationsProvider, this.editingSessionsObs)));
		this._register(multiDiffSourceResolverService.registerResolver(_instantiationService.createInstance(ChatEditingMultiDiffSourceResolver, this.editingSessionsObs)));

		// TODO@jrieken
		// some ugly casting so that this service can pass itself as argument instad as service dependeny
		// eslint-disable-next-line local/code-no-any-casts, @typescript-eslint/no-explicit-any
		this._register(textModelService.registerTextModelContentProvider(ChatEditingTextModelContentProvider.scheme, _instantiationService.createInstance(ChatEditingTextModelContentProvider as any, this)));
		// eslint-disable-next-line local/code-no-any-casts, @typescript-eslint/no-explicit-any
		this._register(textModelService.registerTextModelContentProvider(Schemas.chatEditingSnapshotScheme, _instantiationService.createInstance(ChatEditingSnapshotTextModelContentProvider as any, this)));

		this._register(this._chatService.onDidDisposeSession((e) => {
			if (e.reason === 'cleared') {
				for (const resource of e.sessionResource) {
					this.getEditingSession(resource)?.stop();
				}
			}
		}));

		// todo@connor4312: temporary until chatReadonlyPromptReference proposal is finalized
		const readonlyEnabledContextKey = chatEditingAgentSupportsReadonlyReferencesContextKey.bindTo(contextKeyService);
		const setReadonlyFilesEnabled = () => {
			const enabled = productService.quality !== 'stable' && extensionService.extensions.some(e => e.enabledApiProposals?.includes('chatReadonlyPromptReference'));
			readonlyEnabledContextKey.set(enabled);
		};
		setReadonlyFilesEnabled();
		this._register(extensionService.onDidRegisterExtensions(setReadonlyFilesEnabled));
		this._register(extensionService.onDidChangeExtensions(setReadonlyFilesEnabled));


		// eslint-disable-next-line @typescript-eslint/no-explicit-any
		let storageTask: Promise<any> | undefined;

		this._register(storageService.onWillSaveState(() => {
			// eslint-disable-next-line @typescript-eslint/no-explicit-any
			const tasks: Promise<any>[] = [];

			for (const session of this.editingSessionsObs.get()) {
				if (!session.isGlobalEditingSession) {
					continue;
				}
				tasks.push((session as ChatEditingSession).storeState());
			}

			storageTask = Promise.resolve(storageTask)
				.then(() => Promise.all(tasks))
				.finally(() => storageTask = undefined);
		}));

		this._register(this.lifecycleService.onWillShutdown(e => {
			if (!storageTask) {
				return;
			}
			e.join(storageTask, {
				id: 'join.chatEditingSession',
				label: localize('join.chatEditingSession', "Saving chat edits history")
			});
		}));
	}

	override dispose(): void {
		dispose(this._sessionsObs.get());
		super.dispose();
	}

	startOrContinueGlobalEditingSession(chatModel: ChatModel): IChatEditingSession {
		return this.getEditingSession(chatModel.sessionResource) || this.createEditingSession(chatModel, true);
	}

	private _lookupEntry(uri: URI): AbstractChatEditingModifiedFileEntry | undefined {

		for (const item of Iterable.concat(this.editingSessionsObs.get())) {
			const candidate = item.getEntry(uri);
			if (candidate instanceof AbstractChatEditingModifiedFileEntry) {
				// make sure to ref-count this object
				return candidate.acquire();
			}
		}
		return undefined;
	}

	getEditingSession(chatSessionResource: URI): IChatEditingSession | undefined {
		return this.editingSessionsObs.get()
			.find(candidate => isEqual(candidate.chatSessionResource, chatSessionResource));
	}

	createEditingSession(chatModel: ChatModel, global: boolean = false): IChatEditingSession {
		return this._createEditingSession(chatModel, global, undefined);
	}

	transferEditingSession(chatModel: ChatModel, session: IChatEditingSession): IChatEditingSession {
		return this._createEditingSession(chatModel, session.isGlobalEditingSession, session);
	}

	private _createEditingSession(chatModel: ChatModel, global: boolean, initFrom: IChatEditingSession | undefined): IChatEditingSession {

		assertType(this.getEditingSession(chatModel.sessionResource) === undefined, 'CANNOT have more than one editing session per chat session');

		const session = this._instantiationService.createInstance(ChatEditingSession, chatModel.sessionResource, global, this._lookupEntry.bind(this), initFrom);

		const list = this._sessionsObs.get();
		const removeSession = list.unshift(session);

		const store = new DisposableStore();
		this._store.add(store);

		store.add(this.installAutoApplyObserver(session, chatModel));

		store.add(session.onDidDispose(e => {
			removeSession();
			this._sessionsObs.set(list, undefined);
			this._store.delete(store);
		}));

		this._sessionsObs.set(list, undefined);

		return session;
	}

	private installAutoApplyObserver(session: ChatEditingSession, chatModel: ChatModel): IDisposable {
		if (!chatModel) {
			throw new ErrorNoTelemetry(`Edit session was created for a non-existing chat session: ${session.chatSessionResource}`);
		}

		const observerDisposables = new DisposableStore();

		observerDisposables.add(chatModel.onDidChange(async e => {
			if (e.kind !== 'addRequest') {
				return;
			}
			session.createSnapshot(e.request.id, undefined);
			const responseModel = e.request.response;
			if (responseModel) {
				this.observerEditsInResponse(e.request.id, responseModel, session, observerDisposables);
			}
		}));
		observerDisposables.add(chatModel.onDidDispose(() => observerDisposables.dispose()));
		return observerDisposables;
	}

	private observerEditsInResponse(requestId: string, responseModel: IChatResponseModel, session: ChatEditingSession, observerDisposables: DisposableStore) {
		// Sparse array: the indicies are indexes of `responseModel.response.value`
		// that are edit groups, and then this tracks the edit application for
		// each of them. Note that text edit groups can be updated
		// multiple times during the process of response streaming.
		const enum K { Stream, Workspace }
		const editsSeen: ({ kind: K.Stream; seen: number; stream: IStreamingEdits } | { kind: K.Workspace })[] = [];

		let editorDidChange = false;
		const editorListener = Event.once(this._editorService.onDidActiveEditorChange)(() => {
			editorDidChange = true;
		});
		const editorOpenPromises = new ResourceMap<Promise<void>>();
		const openChatEditedFiles = this._configurationService.getValue('accessibility.openChatEditedFiles');

		const ensureEditorOpen = (partUri: URI) => {
			const uri = CellUri.parse(partUri)?.notebook ?? partUri;
			if (editorOpenPromises.has(uri)) {
				return;
			}
			editorOpenPromises.set(uri, (async () => {
				if (this.notebookService.getNotebookTextModel(uri) || uri.scheme === Schemas.untitled || await this._fileService.exists(uri).catch(() => false)) {
					const activeUri = this._editorService.activeEditorPane?.input.resource;
					const inactive = editorDidChange
						|| this._editorService.activeEditorPane?.input instanceof ChatEditorInput && isEqual(this._editorService.activeEditorPane.input.sessionResource, session.chatSessionResource)
						|| Boolean(activeUri && session.entries.get().find(entry => isEqual(activeUri, entry.modifiedURI)));

					this._editorService.openEditor({ resource: uri, options: { inactive, preserveFocus: true, pinned: true } });
				}
			})());
		};

		const onResponseComplete = () => {
			for (const remaining of editsSeen) {
				if (remaining?.kind === K.Stream) {
					remaining.stream.complete();
				}
			}

			editsSeen.length = 0;
			editorOpenPromises.clear();
			editorListener.dispose();
		};

		const handleResponseParts = async () => {
			if (responseModel.isCanceled) {
				return;
			}

			let undoStop: undefined | string;
			for (let i = 0; i < responseModel.response.value.length; i++) {
				const part = responseModel.response.value[i];

				if (part.kind === 'undoStop') {
					undoStop = part.id;
					continue;
				}

				if (part.kind === 'workspaceEdit') {
					// Track if we've already started processing this workspace edit
					if (!editsSeen[i]) {
						editsSeen[i] = { kind: K.Workspace };
						session.applyWorkspaceEdit(part, responseModel, undoStop ?? responseModel.requestId);
					}
					continue;
				}

				if (part.kind !== 'textEditGroup' && part.kind !== 'notebookEditGroup') {
					continue;
				}

				// Skip external edits - they're already applied on disk
				if (part.isExternalEdit) {
					continue;
				}

				if (openChatEditedFiles) {
					ensureEditorOpen(part.uri);
				}

				// get new edits and start editing session
				let entry = editsSeen[i];
				if (!entry) {
					entry = { kind: K.Stream, seen: 0, stream: session.startStreamingEdits(CellUri.parse(part.uri)?.notebook ?? part.uri, responseModel, undoStop) };
					editsSeen[i] = entry;
				}

				if (entry.kind !== K.Stream) {
					continue;
				}

				const isFirst = entry.seen === 0;
				const newEdits = part.edits.slice(entry.seen);
				entry.seen = part.edits.length;

				if (newEdits.length > 0 || isFirst) {
					for (let i = 0; i < newEdits.length; i++) {
						const edit = newEdits[i];
						const done = part.done ? i === newEdits.length - 1 : false;

						if (isTextEditOperationArray(edit)) {
							entry.stream.pushText(edit, done);
						} else if (isCellTextEditOperationArray(edit)) {
							for (const edits of Object.values(groupBy(edit, e => e.uri.toString()))) {
								if (edits) {
									entry.stream.pushNotebookCellText(edits[0].uri, edits.map(e => e.edit), done);
								}
							}
						} else {
							entry.stream.pushNotebook(edit, done);
						}
					}
				}

				if (part.done) {
					entry.stream.complete();
				}
			}
		};

		if (responseModel.isComplete) {
			handleResponseParts().then(() => {
				onResponseComplete();
			});
		} else {
			const disposable = observerDisposables.add(responseModel.onDidChange(e2 => {
				if (e2.reason === 'undoStop') {
					session.createSnapshot(requestId, e2.id);
				} else {
					handleResponseParts().then(() => {
						if (responseModel.isComplete) {
							onResponseComplete();
							observerDisposables.delete(disposable);
						}
					});
				}
			}));
		}
	}
}

/**
 * Emits an event containing the added or removed elements of the observable.
 */
function observeArrayChanges<T>(obs: IObservable<T[]>, compare: (a: T, b: T) => number, store: DisposableStore): Event<T[]> {
	const emitter = store.add(new Emitter<T[]>());
	store.add(runOnChange(obs, (newArr, oldArr) => {
		const change = delta(oldArr || [], newArr, compare);
		const changedElements = ([] as T[]).concat(change.added).concat(change.removed);
		emitter.fire(changedElements);
	}));
	return emitter.event;
}

class ChatDecorationsProvider extends Disposable implements IDecorationsProvider {

	readonly label: string = localize('chat', "Chat Editing");

	private readonly _currentEntries = derived<readonly IModifiedFileEntry[]>(this, (r) => {
		const sessions = this._sessions.read(r);
		if (!sessions) {
			return [];
		}
		const result: IModifiedFileEntry[] = [];
		for (const session of sessions) {
			if (session.state.read(r) !== ChatEditingSessionState.Disposed) {
				const entries = session.entries.read(r);
				result.push(...entries);
			}
		}
		return result;
	});

	private readonly _currentlyEditingUris = derived<URI[]>(this, (r) => {
		const uri = this._currentEntries.read(r);
		return uri.filter(entry => entry.isCurrentlyBeingModifiedBy.read(r)).map(entry => entry.modifiedURI);
	});

	private readonly _modifiedUris = derived<URI[]>(this, (r) => {
		const uri = this._currentEntries.read(r);
		return uri.filter(entry => !entry.isCurrentlyBeingModifiedBy.read(r) && entry.state.read(r) === ModifiedFileEntryState.Modified).map(entry => entry.modifiedURI);
	});

	readonly onDidChange: Event<URI[]>;

	constructor(
		private readonly _sessions: IObservable<readonly IChatEditingSession[]>
	) {
		super();
		this.onDidChange = Event.any(
			observeArrayChanges(this._currentlyEditingUris, compareBy(uri => uri.toString(), compare), this._store),
			observeArrayChanges(this._modifiedUris, compareBy(uri => uri.toString(), compare), this._store),
		);
	}

	provideDecorations(uri: URI, _token: CancellationToken): IDecorationData | undefined {
		const isCurrentlyBeingModified = this._currentlyEditingUris.get().some(e => e.toString() === uri.toString());
		if (isCurrentlyBeingModified) {
			return {
				weight: 1000,
				letter: ThemeIcon.modify(Codicon.loading, 'spin'),
				bubble: false
			};
		}
		const isModified = this._modifiedUris.get().some(e => e.toString() === uri.toString());
		if (isModified) {
			return {
				weight: 1000,
				letter: Codicon.diffModified,
				tooltip: localize('chatEditing.modified2', "Pending changes from chat"),
				bubble: true
			};
		}
		return undefined;
	}
}

export class ChatEditingMultiDiffSourceResolver implements IMultiDiffSourceResolver {

	constructor(
		private readonly _editingSessionsObs: IObservable<readonly IChatEditingSession[]>,
		@IInstantiationService private readonly _instantiationService: IInstantiationService,
	) { }

	canHandleUri(uri: URI): boolean {
		return uri.scheme === CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME;
	}

	async resolveDiffSource(uri: URI): Promise<IResolvedMultiDiffSource> {

		const parsed = parseChatMultiDiffUri(uri);
		const thisSession = derived(this, r => {
			return this._editingSessionsObs.read(r).find(candidate => isEqual(candidate.chatSessionResource, parsed.chatSessionResource));
		});

		return this._instantiationService.createInstance(ChatEditingMultiDiffSource, thisSession, parsed.showPreviousChanges);
	}
}

class ChatEditingMultiDiffSource implements IResolvedMultiDiffSource {
	private readonly _resources = derived<readonly MultiDiffEditorItem[]>(this, (reader) => {
		const currentSession = this._currentSession.read(reader);
		if (!currentSession) {
			return [];
		}
		const entries = currentSession.entries.read(reader);
		return entries.map((entry) => {
			if (this._showPreviousChanges) {
				const entryDiffObs = currentSession.getEntryDiffBetweenStops(entry.modifiedURI, undefined, undefined);
				const entryDiff = entryDiffObs?.read(reader);
				if (entryDiff) {
					return new MultiDiffEditorItem(
						entryDiff.originalURI,
						entryDiff.modifiedURI,
						undefined,
						undefined,
						{
							[chatEditingResourceContextKey.key]: entry.entryId,
						},
					);
				}
			}

			return new MultiDiffEditorItem(
				entry.originalURI,
				entry.modifiedURI,
				undefined,
				undefined,
				{
					[chatEditingResourceContextKey.key]: entry.entryId,
					// [inChatEditingSessionContextKey.key]: true
				},
			);
		});
	});
	readonly resources = new ValueWithChangeEventFromObservable(this._resources);

	readonly contextKeys = {
		[inChatEditingSessionContextKey.key]: true
	};

	constructor(
		private readonly _currentSession: IObservable<IChatEditingSession | undefined>,
		private readonly _showPreviousChanges: boolean
	) { }
}

function isTextEditOperationArray(value: TextEdit[] | ICellTextEditOperation[] | ICellEditOperation[]): value is TextEdit[] {
	return value.some(e => TextEdit.isTextEdit(e));
}