From ceba479e724aa24e4f343cd93d9a2809a4fa4dc3 Mon Sep 17 00:00:00 2001 From: Martin Aeschlimann Date: Fri, 27 Sep 2024 17:49:15 +0200 Subject: [PATCH] Move the triggering of auto apply from extension to UI (#229976) * auto invoke autoApply, get code blocks from chat model * update * update --- .../browser/actions/chatCodeblockActions.ts | 38 ++------ src/vs/workbench/contrib/chat/browser/chat.ts | 1 + .../chat/browser/chatEditingService.ts | 89 +++++++++++++++---- .../contrib/chat/browser/chatWidget.ts | 4 + .../chat/common/chatCodeMapperService.ts | 77 ++++++++++++++++ .../contrib/chat/common/chatEditingService.ts | 4 +- .../chat/test/browser/mockChatWidget.ts | 5 ++ 7 files changed, 169 insertions(+), 49 deletions(-) diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatCodeblockActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatCodeblockActions.ts index c0870c8811057..3425dad368f04 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatCodeblockActions.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatCodeblockActions.ts @@ -3,15 +3,12 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { CancellationTokenSource } from '../../../../../base/common/cancellation.js'; import { Codicon } from '../../../../../base/common/codicons.js'; import { KeyCode, KeyMod } from '../../../../../base/common/keyCodes.js'; -import { URI } from '../../../../../base/common/uri.js'; import { ICodeEditor } from '../../../../../editor/browser/editorBrowser.js'; import { ServicesAccessor } from '../../../../../editor/browser/editorExtensions.js'; import { ICodeEditorService } from '../../../../../editor/browser/services/codeEditorService.js'; import { EditorContextKeys } from '../../../../../editor/common/editorContextKeys.js'; -import { TextEdit } from '../../../../../editor/common/languages.js'; import { CopyAction } from '../../../../../editor/contrib/clipboard/browser/clipboard.js'; import { localize2 } from '../../../../../nls.js'; import { Action2, MenuId, registerAction2 } from '../../../../../platform/actions/common/actions.js'; @@ -24,7 +21,6 @@ import { IUntitledTextResourceEditorInput } from '../../../../common/editor.js'; import { IEditorService } from '../../../../services/editor/common/editorService.js'; import { accessibleViewInCodeBlock } from '../../../accessibility/browser/accessibilityConfiguration.js'; import { ITerminalEditorService, ITerminalGroupService, ITerminalService } from '../../../terminal/browser/terminal.js'; -import { ICodeMapperCodeBlock, ICodeMapperService } from '../../common/chatCodeMapperService.js'; import { CONTEXT_CHAT_EDIT_APPLIED, CONTEXT_CHAT_ENABLED, CONTEXT_IN_CHAT_INPUT, CONTEXT_IN_CHAT_SESSION } from '../../common/chatContextKeys.js'; import { IChatEditingService } from '../../common/chatEditingService.js'; import { ChatCopyKind, IChatService } from '../../common/chatService.js'; @@ -235,41 +231,21 @@ export function registerChatCodeBlockActions() { override async run(accessor: ServicesAccessor, ...args: any[]) { const chatWidgetService = accessor.get(IChatWidgetService); - const codemapperService = accessor.get(ICodeMapperService); const chatEditingService = accessor.get(IChatEditingService); const widget = chatWidgetService.lastFocusedWidget; - if (!widget) { + if (!widget || !widget.viewModel) { return; } - const items = widget.viewModel?.getItems() ?? []; - const item = widget.getFocus() ?? items[items.length - 1]; - if (!isResponseVM(item)) { - return; - } + const applyEditsId = args[0]; - const codeblocks = widget.getCodeBlockInfosForResponse(item); - const request: ICodeMapperCodeBlock[] = []; - for (const codeblock of codeblocks) { - if (codeblock.codemapperUri && codeblock.uri) { - const code = codeblock.getContent(); - request.push({ resource: codeblock.codemapperUri, code }); - } + const chatModel = widget.viewModel.model; + const request = chatModel.getRequests().find(request => request.response?.result?.metadata?.applyEditsId === applyEditsId); + if (request && request.response) { + await chatEditingService.startOrContinueEditingSession(widget.viewModel.sessionId, { silent: true }); // make sure we have an editing session + await chatEditingService.triggerEditComputation(request.response); } - - await chatEditingService.startOrContinueEditingSession(item.sessionId, async (stream) => { - - const response = { - textEdit: (resource: URI, textEdits: TextEdit[]) => { - stream.textEdits(resource, textEdits); - } - }; - - // Invoke the code mapper for all the code blocks in this response - const tokenSource = new CancellationTokenSource(); - await codemapperService.mapCode({ codeBlocks: request, conversation: [] }, response, tokenSource.token); - }, { silent: true }); } }); diff --git a/src/vs/workbench/contrib/chat/browser/chat.ts b/src/vs/workbench/contrib/chat/browser/chat.ts index 45955c0bb5523..62943602abcf7 100644 --- a/src/vs/workbench/contrib/chat/browser/chat.ts +++ b/src/vs/workbench/contrib/chat/browser/chat.ts @@ -36,6 +36,7 @@ export interface IChatWidgetService { getWidgetByInputUri(uri: URI): IChatWidget | undefined; getWidgetBySessionId(sessionId: string): IChatWidget | undefined; + getWidgetByLocation(location: ChatAgentLocation): IChatWidget[]; } export async function showChatView(viewsService: IViewsService): Promise { diff --git a/src/vs/workbench/contrib/chat/browser/chatEditingService.ts b/src/vs/workbench/contrib/chat/browser/chatEditingService.ts index 23878a50926a2..2df4a1d6fdee2 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditingService.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditingService.ts @@ -2,12 +2,11 @@ * Copyright (c) Microsoft Corporation. All rights reserved. * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ - import { Sequencer } from '../../../../base/common/async.js'; import { Codicon } from '../../../../base/common/codicons.js'; import { BugIndicatingError } from '../../../../base/common/errors.js'; import { Emitter } from '../../../../base/common/event.js'; -import { Disposable, IReference, MutableDisposable } from '../../../../base/common/lifecycle.js'; +import { Disposable, DisposableStore, IDisposable, IReference, MutableDisposable } from '../../../../base/common/lifecycle.js'; import { Schemas } from '../../../../base/common/network.js'; import { derived, IObservable, ITransaction, observableValue, ValueWithChangeEventFromObservable } from '../../../../base/common/observable.js'; import { URI } from '../../../../base/common/uri.js'; @@ -37,10 +36,13 @@ import { IMultiDiffSourceResolver, IMultiDiffSourceResolverService, IResolvedMul import { ChatAgentLocation } from '../common/chatAgents.js'; import { CONTEXT_CHAT_EDITING_ENABLED, CONTEXT_CHAT_ENABLED } from '../common/chatContextKeys.js'; import { ChatEditingSessionState, IChatEditingService, IChatEditingSession, IChatEditingSessionStream, IModifiedFileEntry, ModifiedFileEntryState } from '../common/chatEditingService.js'; +import { IChatResponseModel } from '../common/chatModel.js'; import { IChatService } from '../common/chatService.js'; import { IChatVariablesService } from '../common/chatVariables.js'; import { CHAT_CATEGORY } from './actions/chatActions.js'; import { CHAT_VIEW_ID, IChatWidgetService, showChatView } from './chat.js'; +import { ICodeMapperResponse, ICodeMapperService } from '../common/chatCodeMapperService.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js'; const decidedChatEditingResourceContextKey = new RawContextKey('decidedChatEditingResource', []); const chatEditingResourceContextKey = new RawContextKey('chatEditingResource', undefined); @@ -70,6 +72,7 @@ export class ChatEditingService extends Disposable implements IChatEditingServic @IContextKeyService contextKeyService: IContextKeyService, @IChatService private readonly _chatService: IChatService, @IProgressService private readonly _progressService: IProgressService, + @ICodeMapperService private readonly _codeMapperService: ICodeMapperService, ) { super(); this._register(multiDiffSourceResolverService.registerResolver(_instantiationService.createInstance(ChatEditingMultiDiffSourceResolver, this._currentSessionObs))); @@ -95,24 +98,24 @@ export class ChatEditingService extends Disposable implements IChatEditingServic super.dispose(); } - async startOrContinueEditingSession(chatSessionId: string, builder?: (stream: IChatEditingSessionStream) => Promise, options?: { silent: boolean }): Promise { + async startOrContinueEditingSession(chatSessionId: string, options?: { silent: boolean }): Promise { const session = this._currentSessionObs.get(); if (session) { if (session.chatSessionId !== chatSessionId) { throw new BugIndicatingError('Cannot start new session while another session is active'); } - if (builder) { - return this._continueEditingSession(builder, options); - } } - return this._createEditingSession(chatSessionId, builder, options); + return this._createEditingSession(chatSessionId, options); } - private async _createEditingSession(chatSessionId: string, builder?: (stream: IChatEditingSessionStream) => Promise, options?: { silent: boolean }): Promise { + private async _createEditingSession(chatSessionId: string, options?: { silent: boolean }): Promise { if (this._currentSessionObs.get()) { throw new BugIndicatingError('Cannot have more than one active editing session'); } + // listen for completed responses, run the code mapper and apply the edits to this edit session + this._register(this.installAutoApplyObserver(chatSessionId)); + const input = MultiDiffEditorInput.fromResourceMultiDiffEditorInput({ multiDiffSource: ChatEditingMultiDiffSourceResolver.getMultiDiffSourceUri(), label: localize('multiDiffEditorInput.name', "Suggested Edits") @@ -128,13 +131,56 @@ export class ChatEditingService extends Disposable implements IChatEditingServic this._currentSessionObs.set(session, undefined); this._onDidCreateEditingSession.fire(session); + } + + public triggerEditComputation(responseModel: IChatResponseModel): Promise { + return this._continueEditingSession(async (builder, token) => { + const codeMapperResponse: ICodeMapperResponse = { + textEdit: (resource, edits) => builder.textEdits(resource, edits), + }; + await this._codeMapperService.mapCodeFromResponse(responseModel, codeMapperResponse, token); + }, { silent: true }); + } - if (builder) { - return this._continueEditingSession(builder, options); + private installAutoApplyObserver(sessionId: string): IDisposable { + + const chatModel = this._chatService.getSession(sessionId); + if (!chatModel) { + throw new Error(`Edit session was created for a non-existing chat session: ${sessionId}`); } + + const observerDisposables = new DisposableStore(); + + const onResponseComplete = (responseModel: IChatResponseModel) => { + if (responseModel.result?.metadata?.autoApplyEdits) { + this.triggerEditComputation(responseModel); + } + }; + + observerDisposables.add(chatModel.onDidChange(e => { + if (e.kind === 'addRequest') { + const responseModel = e.request.response; + if (responseModel) { + if (responseModel.isComplete) { + onResponseComplete(responseModel); + } else { + const disposable = responseModel.onDidChange(() => { + if (responseModel.isComplete) { + onResponseComplete(responseModel); + disposable.dispose(); + } else if (responseModel.isCanceled || responseModel.isStale) { + disposable.dispose(); + } + }); + } + } + } + })); + observerDisposables.add(chatModel.onDidDispose(() => observerDisposables.dispose())); + return observerDisposables; } - private async _continueEditingSession(builder: (stream: IChatEditingSessionStream) => Promise, options?: { silent?: boolean }): Promise { + private async _continueEditingSession(builder: (stream: IChatEditingSessionStream, token: CancellationToken) => Promise, options?: { silent?: boolean }): Promise { const session = this._currentSessionObs.get(); if (!session) { throw new BugIndicatingError('Cannot continue missing session'); @@ -161,18 +207,22 @@ export class ChatEditingService extends Disposable implements IChatEditingServic } }; session.acceptStreamingEditsStart(); + const cancellationTokenSource = new CancellationTokenSource(); try { if (editorPane) { - await editorPane?.showWhile(builder(stream)); + await editorPane?.showWhile(builder(stream, cancellationTokenSource.token)); } else { await this._progressService.withProgress({ location: ProgressLocation.Window, title: localize2('chatEditing.startingSession', 'Generating edits...').value, }, async () => { - await builder(stream); - }); + await builder(stream, cancellationTokenSource.token); + }, + () => cancellationTokenSource.cancel() + ); } } finally { + cancellationTokenSource.dispose(); session.resolve(); } } @@ -448,12 +498,14 @@ export class ChatEditingStartSessionAction extends Action2 { const chatEditingService = accessor.get(IChatEditingService); const chatWidgetService = accessor.get(IChatWidgetService); const viewsService = accessor.get(IViewsService); - const currentEditingSession = chatEditingService.currentEditingSession; if (currentEditingSession) { return; } - const widget = chatWidgetService.lastFocusedWidget ?? await showChatView(viewsService); + + const panelWidgets = chatWidgetService.getWidgetByLocation(ChatAgentLocation.Panel); + + const widget = panelWidgets[0] ?? await showChatView(viewsService); if (!widget?.viewModel) { return; } @@ -465,9 +517,12 @@ export class ChatEditingStartSessionAction extends Action2 { variablesService.attachContext('file', { uri: activeUri, }, ChatAgentLocation.Panel); } }); - await chatEditingService.startOrContinueEditingSession(widget.viewModel.sessionId, undefined, { silent: true }); + await chatEditingService.startOrContinueEditingSession(widget.viewModel.sessionId, { silent: true }); } } + + + registerAction2(ChatEditingStartSessionAction); export class ChatEditingStopSessionAction extends Action2 { static readonly ID = 'chatEditing.stopSession'; diff --git a/src/vs/workbench/contrib/chat/browser/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/chatWidget.ts index 2d211eb939a0a..b388ffd59d5da 100644 --- a/src/vs/workbench/contrib/chat/browser/chatWidget.ts +++ b/src/vs/workbench/contrib/chat/browser/chatWidget.ts @@ -1126,6 +1126,10 @@ export class ChatWidgetService implements IChatWidgetService { return this._widgets.find(w => isEqual(w.inputUri, uri)); } + getWidgetByLocation(location: ChatAgentLocation): ChatWidget[] { + return this._widgets.filter(w => w.location === location); + } + getWidgetBySessionId(sessionId: string): ChatWidget | undefined { return this._widgets.find(w => w.viewModel?.sessionId === sessionId); } diff --git a/src/vs/workbench/contrib/chat/common/chatCodeMapperService.ts b/src/vs/workbench/contrib/chat/common/chatCodeMapperService.ts index e64fea9e5db0b..150cda7efd83e 100644 --- a/src/vs/workbench/contrib/chat/common/chatCodeMapperService.ts +++ b/src/vs/workbench/contrib/chat/common/chatCodeMapperService.ts @@ -4,10 +4,14 @@ *--------------------------------------------------------------------------------------------*/ import { CancellationToken } from '../../../../base/common/cancellation.js'; +import { CharCode } from '../../../../base/common/charCode.js'; import { IDisposable } from '../../../../base/common/lifecycle.js'; +import { splitLinesIncludeSeparators } from '../../../../base/common/strings.js'; +import { isString } from '../../../../base/common/types.js'; import { URI } from '../../../../base/common/uri.js'; import { TextEdit } from '../../../../editor/common/languages.js'; import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js'; +import { IChatResponseModel } from './chatModel.js'; export interface ICodeMapperResponse { @@ -49,6 +53,7 @@ export interface ICodeMapperService { readonly _serviceBrand: undefined; registerCodeMapperProvider(handle: number, provider: ICodeMapperProvider): IDisposable; mapCode(request: ICodeMapperRequest, response: ICodeMapperResponse, token: CancellationToken): Promise; + mapCodeFromResponse(responseModel: IChatResponseModel, response: ICodeMapperResponse, token: CancellationToken): Promise; } export class CodeMapperService implements ICodeMapperService { @@ -77,4 +82,76 @@ export class CodeMapperService implements ICodeMapperService { } return undefined; } + + async mapCodeFromResponse(responseModel: IChatResponseModel, response: ICodeMapperResponse, token: CancellationToken) { + const fenceLanguageRegex = /^`{3,}/; + const codeBlocks: ICodeMapperCodeBlock[] = []; + + const currentBlock = []; + const markdownBeforeBlock = []; + let currentBlockUri = undefined; + + let fence = undefined; // if set, we are in a block + + for (const lineOrUri of iterateLinesOrUris(responseModel)) { + if (isString(lineOrUri)) { + const fenceLanguageIdMatch = lineOrUri.match(fenceLanguageRegex); + if (fenceLanguageIdMatch) { + // we found a line that starts with a fence + if (fence !== undefined && fenceLanguageIdMatch[0] === fence) { + // we are in a code block and the fence matches the opening fence: Close the code block + fence = undefined; + if (currentBlockUri) { + // report the code block if we have a URI + codeBlocks.push({ code: currentBlock.join(''), resource: currentBlockUri }); + currentBlock.length = 0; + markdownBeforeBlock.length = 0; + currentBlockUri = undefined; + } + } else { + // we are not in a code block. Open the block + fence = fenceLanguageIdMatch[0]; + } + } else { + if (fence !== undefined) { + currentBlock.push(lineOrUri); + } else { + markdownBeforeBlock.push(lineOrUri); + } + } + } else { + currentBlockUri = lineOrUri; + } + } + return this.mapCode({ codeBlocks, conversation: [] }, response, token); + } +} + +function iterateLinesOrUris(responseModel: IChatResponseModel): Iterable { + return { + *[Symbol.iterator](): Iterator { + let lastIncompleteLine = undefined; + for (const part of responseModel.response.value) { + if (part.kind === 'markdownContent' || part.kind === 'markdownVuln') { + const lines = splitLinesIncludeSeparators(part.content.value); + if (lines.length > 0) { + if (lastIncompleteLine !== undefined) { + lines[0] = lastIncompleteLine + lines[0]; // merge the last incomplete line with the first markdown line + } + lastIncompleteLine = isLineIncomplete(lines[lines.length - 1]) ? lines.pop() : undefined; + for (const line of lines) { + yield line; + } + } + } else if (part.kind === 'codeblockUri') { + yield part.uri; + } + } + } + }; +} + +function isLineIncomplete(line: string) { + const lastChar = line.charCodeAt(line.length - 1); + return lastChar !== CharCode.LineFeed && lastChar !== CharCode.CarriageReturn; } diff --git a/src/vs/workbench/contrib/chat/common/chatEditingService.ts b/src/vs/workbench/contrib/chat/common/chatEditingService.ts index 85df0983324ac..cb9b41173d22d 100644 --- a/src/vs/workbench/contrib/chat/common/chatEditingService.ts +++ b/src/vs/workbench/contrib/chat/common/chatEditingService.ts @@ -8,6 +8,7 @@ import { IObservable, ITransaction } from '../../../../base/common/observable.js import { URI } from '../../../../base/common/uri.js'; import { TextEdit } from '../../../../editor/common/languages.js'; import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js'; +import { IChatResponseModel } from './chatModel.js'; export const IChatEditingService = createDecorator('chatEditingService'); @@ -18,7 +19,8 @@ export interface IChatEditingService { readonly currentEditingSession: IChatEditingSession | null; - startOrContinueEditingSession(chatSessionId: string, builder?: (stream: IChatEditingSessionStream) => Promise, options?: { silent?: boolean }): Promise; + startOrContinueEditingSession(chatSessionId: string, options?: { silent: boolean }): Promise; + triggerEditComputation(responseModel: IChatResponseModel): Promise; } export interface IChatEditingSession { diff --git a/src/vs/workbench/contrib/chat/test/browser/mockChatWidget.ts b/src/vs/workbench/contrib/chat/test/browser/mockChatWidget.ts index c5e26d45f3ec5..8bf8cbebd77fc 100644 --- a/src/vs/workbench/contrib/chat/test/browser/mockChatWidget.ts +++ b/src/vs/workbench/contrib/chat/test/browser/mockChatWidget.ts @@ -5,6 +5,7 @@ import { URI } from '../../../../../base/common/uri.js'; import { IChatWidget, IChatWidgetService } from '../../browser/chat.js'; +import { ChatAgentLocation } from '../../common/chatAgents.js'; export class MockChatWidgetService implements IChatWidgetService { readonly _serviceBrand: undefined; @@ -21,4 +22,8 @@ export class MockChatWidgetService implements IChatWidgetService { getWidgetBySessionId(sessionId: string): IChatWidget | undefined { return undefined; } + + getWidgetByLocation(location: ChatAgentLocation): IChatWidget[] { + return []; + } }