From 55822477358f63227a3413b4139107c3643f13cb Mon Sep 17 00:00:00 2001 From: Joyce Er Date: Fri, 11 Oct 2024 20:13:27 -0700 Subject: [PATCH] feat: track open editors as transient working set entries --- .../browser/actions/chatContextActions.ts | 70 +++++++++++-------- .../chat/browser/chatAttachmentModel.ts | 8 ++- .../chat/browser/chatEditingActions.ts | 2 +- .../chat/browser/chatEditingService.ts | 70 +++++++++++++++++-- .../contrib/chat/browser/chatInputPart.ts | 14 +++- .../contrib/chat/browser/chatWidget.ts | 12 +++- .../contrib/chat/common/chatEditingService.ts | 1 + .../chat/common/chatWidgetHistoryService.ts | 3 + 8 files changed, 140 insertions(+), 40 deletions(-) diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts index c540bd29d33be..83fa8df35c83a 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts @@ -27,11 +27,12 @@ import { AnythingQuickAccessProviderRunOptions } from '../../../../../platform/q import { IQuickInputService, IQuickPickItem, IQuickPickItemWithResource, IQuickPickSeparator, QuickPickItem } from '../../../../../platform/quickinput/common/quickInput.js'; import { ActiveEditorContext } from '../../../../common/contextkeys.js'; import { IEditorService } from '../../../../services/editor/common/editorService.js'; +import { VIEW_ID as SEARCH_VIEW_ID } from '../../../../services/search/common/search.js'; import { IViewsService } from '../../../../services/views/common/viewsService.js'; import { AnythingQuickAccessProvider } from '../../../search/browser/anythingQuickAccess.js'; +import { SearchView } from '../../../search/browser/searchView.js'; import { ISymbolQuickPickItem, SymbolsQuickAccessProvider } from '../../../search/browser/symbolsQuickAccess.js'; import { SearchContext } from '../../../search/common/constants.js'; -import { VIEW_ID as SEARCH_VIEW_ID } from '../../../../services/search/common/search.js'; import { ChatAgentLocation, IChatAgentService } from '../../common/chatAgents.js'; import { CONTEXT_CHAT_LOCATION, CONTEXT_IN_CHAT_INPUT } from '../../common/chatContextKeys.js'; import { IChatEditingService } from '../../common/chatEditingService.js'; @@ -39,12 +40,11 @@ import { IChatRequestVariableEntry } from '../../common/chatModel.js'; import { ChatRequestAgentPart } from '../../common/chatParserTypes.js'; import { IChatVariableData, IChatVariablesService } from '../../common/chatVariables.js'; import { ILanguageModelToolsService } from '../../common/languageModelToolsService.js'; -import { imageToHash, isImage } from '../chatPasteProviders.js'; import { IChatWidget, IChatWidgetService, IQuickChatService, showChatView } from '../chat.js'; +import { imageToHash, isImage } from '../chatPasteProviders.js'; import { isQuickChat } from '../chatWidget.js'; -import { CHAT_CATEGORY } from './chatActions.js'; -import { SearchView } from '../../../search/browser/searchView.js'; import { getScreenshotAsVariable, ScreenshotVariableId } from '../contrib/screenshot.js'; +import { CHAT_CATEGORY } from './chatActions.js'; export function registerChatContextActions() { registerAction2(AttachContextAction); @@ -296,14 +296,17 @@ export class AttachContextAction extends Action2 { }); } else { // file attachment - toAttach.push({ - id: this._getFileContextId({ resource: pick.resource }), - value: pick.resource, - name: pick.label, - isFile: true, - isDynamic: true, - }); - chatEditingService?.addFileToWorkingSet(pick.resource); + if (chatEditingService) { + chatEditingService.addFileToWorkingSet(pick.resource); + } else { + toAttach.push({ + id: this._getFileContextId({ resource: pick.resource }), + value: pick.resource, + name: pick.label, + isFile: true, + isDynamic: true, + }); + } } } else if (isIGotoSymbolQuickPickItem(pick) && pick.uri && pick.range) { toAttach.push({ @@ -317,27 +320,33 @@ export class AttachContextAction extends Action2 { } else if (isIOpenEditorsQuickPickItem(pick)) { for (const editor of editorService.editors) { if (editor.resource) { - toAttach.push({ - id: this._getFileContextId({ resource: editor.resource }), - value: editor.resource, - name: labelService.getUriBasenameLabel(editor.resource), - isFile: true, - isDynamic: true - }); - chatEditingService?.addFileToWorkingSet(editor.resource); + if (chatEditingService) { + chatEditingService.addFileToWorkingSet(editor.resource); + } else { + toAttach.push({ + id: this._getFileContextId({ resource: editor.resource }), + value: editor.resource, + name: labelService.getUriBasenameLabel(editor.resource), + isFile: true, + isDynamic: true + }); + } } } } else if (isISearchResultsQuickPickItem(pick)) { const searchView = viewsService.getViewWithId(SEARCH_VIEW_ID) as SearchView; for (const result of searchView.model.searchResult.matches()) { - toAttach.push({ - id: this._getFileContextId({ resource: result.resource }), - value: result.resource, - name: labelService.getUriBasenameLabel(result.resource), - isFile: true, - isDynamic: true - }); - chatEditingService?.addFileToWorkingSet(result.resource); + if (chatEditingService) { + chatEditingService.addFileToWorkingSet(result.resource); + } else { + toAttach.push({ + id: this._getFileContextId({ resource: result.resource }), + value: result.resource, + name: labelService.getUriBasenameLabel(result.resource), + isFile: true, + isDynamic: true + }); + } } } else if (isScreenshotQuickPickItem(pick)) { const variable = await getScreenshotAsVariable(); @@ -580,6 +589,11 @@ export class AttachContextAction extends Action2 { filter: (item: IChatContextQuickPickItem | IQuickPickSeparator) => { // Avoid attaching the same context twice const attachedContext = widget.attachmentModel.getAttachmentIDs(); + if (chatEditingService) { + for (const file of chatEditingService.currentEditingSessionObs.get()?.workingSet.keys() ?? []) { + attachedContext.add(this._getFileContextId({ resource: file })); + } + } if (isIOpenEditorsQuickPickItem(item)) { for (const editor of editorService.editors) { diff --git a/src/vs/workbench/contrib/chat/browser/chatAttachmentModel.ts b/src/vs/workbench/contrib/chat/browser/chatAttachmentModel.ts index 65a2f92965f33..3fae7db826c51 100644 --- a/src/vs/workbench/contrib/chat/browser/chatAttachmentModel.ts +++ b/src/vs/workbench/contrib/chat/browser/chatAttachmentModel.ts @@ -38,13 +38,17 @@ export class ChatAttachmentModel extends Disposable { } addFile(uri: URI, range?: IRange) { - this.addContext({ + this.addContext(this.asVariableEntry(uri, range)); + } + + asVariableEntry(uri: URI, range?: IRange) { + return { value: uri, id: uri.toString() + (range?.toString() ?? ''), name: basename(uri), isFile: true, isDynamic: true - }); + }; } addContext(...attachments: IChatRequestVariableEntry[]) { diff --git a/src/vs/workbench/contrib/chat/browser/chatEditingActions.ts b/src/vs/workbench/contrib/chat/browser/chatEditingActions.ts index a04234f9c486e..14e2a1c23932f 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditingActions.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditingActions.ts @@ -59,7 +59,7 @@ registerAction2(class RemoveFileFromWorkingSet extends WorkingSetAction { icon: Codicon.close, menu: [{ id: MenuId.ChatEditingSessionWidgetToolbar, - when: ContextKeyExpr.equals(chatEditingWidgetFileStateContextKey.key, WorkingSetEntryState.Attached), + when: ContextKeyExpr.or(ContextKeyExpr.equals(chatEditingWidgetFileStateContextKey.key, WorkingSetEntryState.Attached), ContextKeyExpr.equals(chatEditingWidgetFileStateContextKey.key, WorkingSetEntryState.Transient)), order: 0, group: 'navigation' }], diff --git a/src/vs/workbench/contrib/chat/browser/chatEditingService.ts b/src/vs/workbench/contrib/chat/browser/chatEditingService.ts index d291d927dacd0..5781761fcc14b 100644 --- a/src/vs/workbench/contrib/chat/browser/chatEditingService.ts +++ b/src/vs/workbench/contrib/chat/browser/chatEditingService.ts @@ -8,7 +8,7 @@ import { CancellationToken, CancellationTokenSource } from '../../../../base/com import { BugIndicatingError } from '../../../../base/common/errors.js'; import { Emitter } from '../../../../base/common/event.js'; import { Disposable, DisposableStore, IDisposable, IReference } from '../../../../base/common/lifecycle.js'; -import { ResourceMap } from '../../../../base/common/map.js'; +import { ResourceMap, ResourceSet } from '../../../../base/common/map.js'; import { derived, IObservable, ITransaction, observableValue, ValueWithChangeEventFromObservable } from '../../../../base/common/observable.js'; import { themeColorFromId } from '../../../../base/common/themables.js'; import { URI } from '../../../../base/common/uri.js'; @@ -33,6 +33,7 @@ import { IInstantiationService } from '../../../../platform/instantiation/common import { bindContextKey } from '../../../../platform/observable/common/platformObservableUtils.js'; import { IProgressService, ProgressLocation } from '../../../../platform/progress/common/progress.js'; import { editorSelectionBackground } from '../../../../platform/theme/common/colorRegistry.js'; +import { IEditorCloseEvent } from '../../../common/editor.js'; import { DiffEditorInput } from '../../../common/editor/diffEditorInput.js'; import { EditorInput } from '../../../common/editor/editorInput.js'; import { IEditorGroup, IEditorGroupsService } from '../../../services/editor/common/editorGroupsService.js'; @@ -533,17 +534,64 @@ class ChatEditingSession extends Disposable implements IChatEditingSession { return; // Shouldn't happen } - // Add the currently active editor to the working set - let activeEditorControl = this._editorService.activeTextEditorControl; - if (activeEditorControl) { + // Add the currently active editors to the working set + this._trackCurrentEditorsInWorkingSet(); + this._register(this._editorService.onDidActiveEditorChange(() => { + this._trackCurrentEditorsInWorkingSet(); + })); + this._register(this._editorService.onDidCloseEditor((e) => { + this._trackCurrentEditorsInWorkingSet(e); + })); + } + + private _trackCurrentEditorsInWorkingSet(e?: IEditorCloseEvent) { + const closedEditor = e?.editor.resource?.toString(); + + const existingTransientEntries = new ResourceSet(); + for (const file of this._workingSet.keys()) { + if (this._workingSet.get(file) === WorkingSetEntryState.Transient) { + existingTransientEntries.add(file); + } + } + if (existingTransientEntries.size === 0 && this._workingSet.size > 0) { + // The user manually added or removed attachments, don't inherit the visible editors + return; + } + + const activeEditors = new ResourceSet(); + this._editorGroupsService.groups.forEach((group) => { + if (!group.activeEditorPane) { + return; + } + let activeEditorControl = group.activeEditorPane.getControl(); if (isDiffEditor(activeEditorControl)) { activeEditorControl = activeEditorControl.getOriginalEditor().hasTextFocus() ? activeEditorControl.getOriginalEditor() : activeEditorControl.getModifiedEditor(); } if (isCodeEditor(activeEditorControl) && activeEditorControl.hasModel()) { const uri = activeEditorControl.getModel().uri; - this._workingSet.set(uri, WorkingSetEntryState.Attached); - widget.attachmentModel.addFile(uri); + if (closedEditor === uri.toString()) { + // The editor group service sees recently closed editors? + // Continue, since we want this to be deleted from the working set + } else if (existingTransientEntries.has(uri)) { + existingTransientEntries.delete(uri); + } else { + activeEditors.add(uri); + } } + }); + + let didChange = false; + for (const entry of existingTransientEntries) { + didChange ||= this._workingSet.delete(entry); + } + + for (const entry of activeEditors) { + this._workingSet.set(entry, WorkingSetEntryState.Transient); + didChange = true; + } + + if (didChange) { + this._onDidChange.fire(); } } @@ -655,7 +703,7 @@ class ChatEditingSession extends Disposable implements IChatEditingSession { let didRemoveUris = false; for (const uri of uris) { - didRemoveUris = didRemoveUris || this._workingSet.delete(uri); + didRemoveUris ||= this._workingSet.delete(uri); } if (!didRemoveUris) { @@ -791,6 +839,14 @@ class ChatEditingSession extends Disposable implements IChatEditingSession { addFileToWorkingSet(resource: URI) { if (!this._workingSet.has(resource)) { this._workingSet.set(resource, WorkingSetEntryState.Attached); + + // Convert all transient entries to attachments + for (const file of this._workingSet.keys()) { + if (this._workingSet.get(file) === WorkingSetEntryState.Transient) { + this._workingSet.set(file, WorkingSetEntryState.Attached); + } + } + this._onDidChange.fire(); } } diff --git a/src/vs/workbench/contrib/chat/browser/chatInputPart.ts b/src/vs/workbench/contrib/chat/browser/chatInputPart.ts index 287e9ca96f2cd..2d663d02790cc 100644 --- a/src/vs/workbench/contrib/chat/browser/chatInputPart.ts +++ b/src/vs/workbench/contrib/chat/browser/chatInputPart.ts @@ -186,6 +186,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge private readonly _chatEditsActionsDisposables = this._register(new DisposableStore()); private readonly _chatEditsDisposables = this._register(new DisposableStore()); + private _chatEditingSession: IChatEditingSession | undefined; private _chatEditsProgress: ProgressBar | undefined; private _chatEditsListPool: CollapsibleListPool; private _chatEditList: IDisposableReference> | undefined; @@ -227,10 +228,15 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge this._attachmentModel = this._register(new ChatAttachmentModel()); this.getInputState = (): IChatInputState => { - // Get input state from widget contribs, merge with attachments + // Get input state from widget contribs, merge with attachments and working set + const chatWorkingSet: { uri: URI; state: WorkingSetEntryState }[] = []; + for (const [uri, state] of this._chatEditingSession?.workingSet.entries() ?? []) { + chatWorkingSet.push({ uri, state }); + } return { ...getContribsInputState(), chatContextAttachments: this._attachmentModel.attachments, + chatWorkingSet: chatWorkingSet }; }; this.inputEditorMaxHeight = this.options.renderStyle === 'compact' ? INPUT_EDITOR_MAX_HEIGHT / 3 : INPUT_EDITOR_MAX_HEIGHT; @@ -368,6 +374,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge this.history.previous() : this.history.next(); const historyAttachments = historyEntry.state?.chatContextAttachments ?? []; + this._chatEditingSession?.workingSet.clear(); + for (const entry of historyEntry.state?.chatWorkingSet ?? []) { + this._chatEditingSession?.workingSet.set(entry.uri, entry.state); + } this._attachmentModel.clearAndSetContext(...historyAttachments); aria.status(historyEntry.text); @@ -846,9 +856,11 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge this._chatEditsDisposables.clear(); this._chatEditList = undefined; this._chatEditsProgress?.dispose(); + this._chatEditingSession = undefined; return; } + this._chatEditingSession = chatEditingSession; const currentChatEditingState = chatEditingSession.state.get(); if (this._chatEditList && !chatWidget?.viewModel?.requestInProgress && (currentChatEditingState === ChatEditingSessionState.Idle || currentChatEditingState === ChatEditingSessionState.Initial)) { this._chatEditsProgress?.stop(); diff --git a/src/vs/workbench/contrib/chat/browser/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/chatWidget.ts index c5ed324e18981..9f7b14e1e4100 100644 --- a/src/vs/workbench/contrib/chat/browser/chatWidget.ts +++ b/src/vs/workbench/contrib/chat/browser/chatWidget.ts @@ -931,12 +931,22 @@ export class ChatWidget extends Disposable implements IChatWidget { } } + const attachedContext = [...this.attachmentModel.attachments]; + if (this.location === ChatAgentLocation.EditingSession) { + const currentEditingSession = this.chatEditingService.currentEditingSessionObs.get(); + if (currentEditingSession?.workingSet) { + for (const [file, _] of currentEditingSession?.workingSet) { + attachedContext.push(this.attachmentModel.asVariableEntry(file)); + } + } + } + const result = await this.chatService.sendRequest(this.viewModel.sessionId, input, { userSelectedModelId: this.inputPart.currentLanguageModel, location: this.location, locationData: this._location.resolveData?.(), parserContext: { selectedAgent: this._lastSelectedAgent }, - attachedContext: [...this.attachmentModel.attachments] + attachedContext: attachedContext }); if (result) { diff --git a/src/vs/workbench/contrib/chat/common/chatEditingService.ts b/src/vs/workbench/contrib/chat/common/chatEditingService.ts index b2f94d2ab8c02..eeb55ebdbf38f 100644 --- a/src/vs/workbench/contrib/chat/common/chatEditingService.ts +++ b/src/vs/workbench/contrib/chat/common/chatEditingService.ts @@ -63,6 +63,7 @@ export const enum WorkingSetEntryState { Modified, Accepted, Rejected, + Transient, Attached, Sent, } diff --git a/src/vs/workbench/contrib/chat/common/chatWidgetHistoryService.ts b/src/vs/workbench/contrib/chat/common/chatWidgetHistoryService.ts index b5a209b8278c7..2ea99b4bcfacf 100644 --- a/src/vs/workbench/contrib/chat/common/chatWidgetHistoryService.ts +++ b/src/vs/workbench/contrib/chat/common/chatWidgetHistoryService.ts @@ -4,10 +4,12 @@ *--------------------------------------------------------------------------------------------*/ import { Emitter, Event } from '../../../../base/common/event.js'; +import { URI } from '../../../../base/common/uri.js'; import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js'; import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js'; import { Memento } from '../../../common/memento.js'; import { ChatAgentLocation } from './chatAgents.js'; +import { WorkingSetEntryState } from './chatEditingService.js'; import { IChatRequestVariableEntry } from './chatModel.js'; import { CHAT_PROVIDER_ID } from './chatParticipantContribTypes.js'; @@ -20,6 +22,7 @@ export interface IChatHistoryEntry { export interface IChatInputState { [key: string]: any; chatContextAttachments?: ReadonlyArray; + chatWorkingSet?: ReadonlyArray<{ uri: URI; state: WorkingSetEntryState }>; } export const IChatWidgetHistoryService = createDecorator('IChatWidgetHistoryService');