diff --git a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts index c0efc042f8d..e3c9bfd2c82 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingServiceImpl.ts @@ -33,6 +33,7 @@ import { IExtensionService } from '../../../../services/extensions/common/extens import { ILifecycleService } from '../../../../services/lifecycle/common/lifecycle.js'; import { IMultiDiffSourceResolver, IMultiDiffSourceResolverService, IResolvedMultiDiffSource, MultiDiffEditorItem } from '../../../multiDiffEditor/browser/multiDiffSourceResolverService.js'; import { CellUri } from '../../../notebook/common/notebookCommon.js'; +import { INotebookService } from '../../../notebook/common/notebookService.js'; import { IChatAgentService } from '../../common/chatAgents.js'; import { CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingAgentSupportsReadonlyReferencesContextKey, chatEditingResourceContextKey, ChatEditingSessionState, chatEditingSnapshotScheme, IChatEditingService, IChatEditingSession, IChatRelatedFile, IChatRelatedFilesProvider, IModifiedFileEntry, inChatEditingSessionContextKey, IStreamingEdits, WorkingSetEntryState } from '../../common/chatEditingService.js'; import { IChatResponseModel, isCellTextEditOperation } from '../../common/chatModel.js'; @@ -75,6 +76,7 @@ export class ChatEditingService extends Disposable implements IChatEditingServic @ILogService logService: ILogService, @IExtensionService extensionService: IExtensionService, @IProductService productService: IProductService, + @INotebookService private readonly notebookService: INotebookService ) { super(); this._register(decorationsService.registerDecorationsProvider(_instantiationService.createInstance(ChatDecorationsProvider, this.editingSessionsObs))); @@ -254,7 +256,8 @@ export class ChatEditingService extends Disposable implements IChatEditingServic return; } - editedFilesExist.set(uri, this._fileService.exists(uri).then((e) => { + const fileExists = this.notebookService.getNotebookTextModel(uri) ? Promise.resolve(true) : this._fileService.exists(uri); + editedFilesExist.set(uri, fileExists.then((e) => { if (!e) { return; } diff --git a/src/vs/workbench/contrib/chat/common/tools/editFileTool.ts b/src/vs/workbench/contrib/chat/common/tools/editFileTool.ts index ca087fb21f8..e56d48d5809 100644 --- a/src/vs/workbench/contrib/chat/common/tools/editFileTool.ts +++ b/src/vs/workbench/contrib/chat/common/tools/editFileTool.ts @@ -90,9 +90,10 @@ export class EditTool implements IToolImpl { } const parameters = invocation.parameters as EditToolParams; - const uri = URI.revive(parameters.file); // TODO@roblourens do revive in MainThreadLanguageModelTools + const fileUri = URI.revive(parameters.file); // TODO@roblourens do revive in MainThreadLanguageModelTools + const uri = CellUri.parse(fileUri)?.notebook || fileUri; - if (!this.workspaceContextService.isInsideWorkspace(uri)) { + if (!this.workspaceContextService.isInsideWorkspace(uri) && !this.notebookService.getNotebookTextModel(uri)) { const groupsByLastActive = this.editorGroupsService.getGroups(GroupsOrder.MOST_RECENTLY_ACTIVE); const uriIsOpenInSomeEditor = groupsByLastActive.some((group) => { return group.editors.some((editor) => { @@ -134,13 +135,12 @@ export class EditTool implements IToolImpl { kind: 'markdownContent', content: new MarkdownString(parameters.code + '\n````\n') }); - const notebookUri = CellUri.parse(uri)?.notebook || uri; // Signal start. - if (this.notebookService.hasSupportedNotebooks(notebookUri) && (this.notebookService.getNotebookTextModel(notebookUri))) { + if (this.notebookService.hasSupportedNotebooks(uri) && (this.notebookService.getNotebookTextModel(uri))) { model.acceptResponseProgress(request, { kind: 'notebookEdit', edits: [], - uri: notebookUri + uri }); } else { model.acceptResponseProgress(request, { @@ -169,8 +169,8 @@ export class EditTool implements IToolImpl { }, token); // Signal end. - if (this.notebookService.hasSupportedNotebooks(notebookUri) && (this.notebookService.getNotebookTextModel(notebookUri))) { - model.acceptResponseProgress(request, { kind: 'notebookEdit', uri: notebookUri, edits: [], done: true }); + if (this.notebookService.hasSupportedNotebooks(uri) && (this.notebookService.getNotebookTextModel(uri))) { + model.acceptResponseProgress(request, { kind: 'notebookEdit', uri, edits: [], done: true }); } else { model.acceptResponseProgress(request, { kind: 'textEdit', uri, edits: [], done: true }); } diff --git a/src/vs/workbench/contrib/chat/common/tools/insertNotebookCellsTool.ts b/src/vs/workbench/contrib/chat/common/tools/insertNotebookCellsTool.ts new file mode 100644 index 00000000000..b66a2ace7ba --- /dev/null +++ b/src/vs/workbench/contrib/chat/common/tools/insertNotebookCellsTool.ts @@ -0,0 +1,218 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { CancellationToken } from '../../../../../base/common/cancellation.js'; +import { MarkdownString } from '../../../../../base/common/htmlContent.js'; +import { IDisposable } from '../../../../../base/common/lifecycle.js'; +import { autorun } from '../../../../../base/common/observable.js'; +import { URI, UriComponents } from '../../../../../base/common/uri.js'; +import { generateUuid } from '../../../../../base/common/uuid.js'; +import { localize } from '../../../../../nls.js'; +import { IWorkspaceContextService } from '../../../../../platform/workspace/common/workspace.js'; +import { SaveReason } from '../../../../common/editor.js'; +import { ITextFileService } from '../../../../services/textfile/common/textfiles.js'; +import { CellUri } from '../../../notebook/common/notebookCommon.js'; +import { INotebookService } from '../../../notebook/common/notebookService.js'; +import { ICodeMapperService } from '../chatCodeMapperService.js'; +import { IChatEditingService } from '../chatEditingService.js'; +import { ChatModel } from '../chatModel.js'; +import { IChatService } from '../chatService.js'; +import { ILanguageModelIgnoredFilesService } from '../ignoredFiles.js'; +import { CountTokensCallback, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../languageModelToolsService.js'; +import { IToolInputProcessor } from './tools.js'; + +const codeInstructions = ` +The user is very smart and can understand how to insert cells to their new Notebook files +`; + +export const ExtensionEditToolId = 'vscode_insert_notebook_cells'; +export const InternalEditToolId = 'vscode_insert_notebook_cells_internal'; +export const EditToolData: IToolData = { + id: InternalEditToolId, + displayName: localize('chat.tools.editFile', "Edit File"), + modelDescription: `Insert cells into a new notebook n the workspace. Use this tool once per file that needs to be modified, even if there are multiple changes for a file. Generate the "explanation" property first. ${codeInstructions}`, + inputSchema: { + type: 'object', + properties: { + explanation: { + type: 'string', + description: 'A short explanation of the edit being made. Can be the same as the explanation you showed to the user.', + }, + filePath: { + type: 'string', + description: 'An absolute path to the file to edit, or the URI of a untitled, not yet named, file, such as `untitled:Untitled-1.', + }, + cells: { + type: 'array', + description: 'The cells to insert to apply to the file. ' + codeInstructions + } + }, + required: ['explanation', 'filePath', 'code'] + } +}; + +export class EditTool implements IToolImpl { + + constructor( + @IChatService private readonly chatService: IChatService, + @IChatEditingService private readonly chatEditingService: IChatEditingService, + @ICodeMapperService private readonly codeMapperService: ICodeMapperService, + @IWorkspaceContextService private readonly workspaceContextService: IWorkspaceContextService, + @ILanguageModelIgnoredFilesService private readonly ignoredFilesService: ILanguageModelIgnoredFilesService, + @ITextFileService private readonly textFileService: ITextFileService, + @INotebookService private readonly notebookService: INotebookService, + ) { } + + async invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise { + if (!invocation.context) { + throw new Error('toolInvocationToken is required for this tool'); + } + + const parameters = invocation.parameters as EditToolParams; + const uri = URI.revive(parameters.file); // TODO@roblourens do revive in MainThreadLanguageModelTools + if (!this.workspaceContextService.isInsideWorkspace(uri)) { + throw new Error(`File ${uri.fsPath} can't be edited because it's not inside the current workspace`); + } + + if (await this.ignoredFilesService.fileIsIgnored(uri, token)) { + throw new Error(`File ${uri.fsPath} can't be edited because it is configured to be ignored by Copilot`); + } + + const model = this.chatService.getSession(invocation.context?.sessionId) as ChatModel; + const request = model.getRequests().at(-1)!; + + // Undo stops mark groups of response data in the output. Operations, such + // as text edits, that happen between undo stops are all done or undone together. + if (request.response?.response.getMarkdown().length) { + // slightly hacky way to avoid an extra 'no-op' undo stop at the start of responses that are just edits + model.acceptResponseProgress(request, { + kind: 'undoStop', + id: generateUuid(), + }); + } + + model.acceptResponseProgress(request, { + kind: 'markdownContent', + content: new MarkdownString('\n````\n') + }); + model.acceptResponseProgress(request, { + kind: 'codeblockUri', + uri + }); + model.acceptResponseProgress(request, { + kind: 'markdownContent', + content: new MarkdownString(parameters.code + '\n````\n') + }); + const notebookUri = CellUri.parse(uri)?.notebook || uri; + // Signal start. + if (this.notebookService.hasSupportedNotebooks(notebookUri) && (this.notebookService.getNotebookTextModel(notebookUri))) { + model.acceptResponseProgress(request, { + kind: 'notebookEdit', + edits: [], + uri: notebookUri + }); + } else { + model.acceptResponseProgress(request, { + kind: 'textEdit', + edits: [], + uri + }); + } + + const editSession = this.chatEditingService.getEditingSession(model.sessionId); + if (!editSession) { + throw new Error('This tool must be called from within an editing session'); + } + + const result = await this.codeMapperService.mapCode({ + codeBlocks: [{ code: parameters.code, resource: uri, markdownBeforeBlock: parameters.explanation }], + location: 'tool', + chatRequestId: invocation.chatRequestId + }, { + textEdit: (target, edits) => { + model.acceptResponseProgress(request, { kind: 'textEdit', uri: target, edits }); + }, + notebookEdit(target, edits) { + model.acceptResponseProgress(request, { kind: 'notebookEdit', uri: target, edits }); + }, + }, token); + + // Signal end. + if (this.notebookService.hasSupportedNotebooks(notebookUri) && (this.notebookService.getNotebookTextModel(notebookUri))) { + model.acceptResponseProgress(request, { kind: 'notebookEdit', uri: notebookUri, edits: [], done: true }); + } else { + model.acceptResponseProgress(request, { kind: 'textEdit', uri, edits: [], done: true }); + } + + if (result?.errorMessage) { + throw new Error(result.errorMessage); + } + + let dispose: IDisposable; + await new Promise((resolve) => { + // The file will not be modified until the first edits start streaming in, + // so wait until we see that it _was_ modified before waiting for it to be done. + let wasFileBeingModified = false; + + dispose = autorun((r) => { + + const entries = editSession.entries.read(r); + const currentFile = entries?.find((e) => e.modifiedURI.toString() === uri.toString()); + if (currentFile) { + if (currentFile.isCurrentlyBeingModifiedBy.read(r)) { + wasFileBeingModified = true; + } else if (wasFileBeingModified) { + resolve(true); + } + } + }); + }).finally(() => { + dispose.dispose(); + }); + + await this.textFileService.save(uri, { + reason: SaveReason.AUTO, + skipSaveParticipants: true, + }); + + return { + content: [{ kind: 'text', value: 'The file was edited successfully' }] + }; + } + + async prepareToolInvocation(parameters: any, token: CancellationToken): Promise { + return { + presentation: 'hidden' + }; + } +} + +export interface EditToolParams { + file: UriComponents; + explanation: string; + code: string; +} + +export interface EditToolRawParams { + filePath: string; + explanation: string; + code: string; +} + +export class EditToolInputProcessor implements IToolInputProcessor { + processInput(input: EditToolRawParams): EditToolParams { + if (!input.filePath) { + // Tool name collision, or input wasn't properly validated upstream + return input as any; + } + const filePath = input.filePath; + // Runs in EH, will be mapped + return { + file: filePath.startsWith('untitled:') ? URI.parse(filePath) : URI.file(filePath), + explanation: input.explanation, + code: input.code, + }; + } +} diff --git a/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts b/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts index 5fed1eb581a..17490eb40ef 100644 --- a/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/chatEditingService.test.ts @@ -32,6 +32,7 @@ import { waitForState } from '../../../../../base/common/observable.js'; import { INotebookService } from '../../../notebook/common/notebookService.js'; import { Range } from '../../../../../editor/common/core/range.js'; import { ChatAgentLocation } from '../../common/constants.js'; +import { NotebookTextModel } from '../../../notebook/common/model/notebookTextModel.js'; function getAgentData(id: string) { return { @@ -69,7 +70,10 @@ suite('ChatEditingService', function () { } }); collection.set(INotebookService, new class extends mock() { - override hasSupportedNotebooks(resource: URI): boolean { + override getNotebookTextModel(_uri: URI): NotebookTextModel | undefined { + return undefined; + } + override hasSupportedNotebooks(_resource: URI): boolean { return false; } });