Skip to content

Commit

Permalink
feat: track open editors as transient working set entries
Browse files Browse the repository at this point in the history
  • Loading branch information
@joyceerhl
joyceerhl committed Oct 12, 2024
1 parent d91b89d commit f4d5254
Show file tree
Hide file tree
Showing 9 changed files with 157 additions and 47 deletions.
70 changes: 42 additions & 28 deletions src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,24 +27,24 @@ 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';
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);
Expand Down Expand Up @@ -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({
Expand All @@ -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();
Expand Down Expand Up @@ -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) {
Expand Down
8 changes: 6 additions & 2 deletions src/vs/workbench/contrib/chat/browser/chatAttachmentModel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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[]) {
Expand Down
24 changes: 17 additions & 7 deletions src/vs/workbench/contrib/chat/browser/chatContentParts/chatReferencesContentPart.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -385,13 +385,23 @@ class CollapsibleListRenderer implements IListRenderer<IChatCollapsibleListItem,
} else if (matchesSomeScheme(uri, Schemas.mailto, Schemas.http, Schemas.https)) {
templateData.label.setResource({ resource: uri, name: uri.toString() }, { icon: icon ?? Codicon.globe, title: data.options?.status?.description ?? data.title ?? uri.toString() });
} else {
templateData.label.setFile(uri, {
fileKind: FileKind.FILE,
// Should not have this live-updating data on a historical reference
fileDecorations: undefined,
range: 'range' in reference ? reference.range : undefined,
title: data.options?.status?.description ?? data.title
});
if (data.state === WorkingSetEntryState.Transient) {
templateData.label.setResource(
{
resource: uri,
name: basenameOrAuthority(uri),
description: localize('chat.openEditor', 'Open Editor'),
range: 'range' in reference ? reference.range : undefined,
}, { icon, title: data.options?.status?.description ?? data.title });
} else {
templateData.label.setFile(uri, {
fileKind: FileKind.FILE,
// Should not have this live-updating data on a historical reference
fileDecorations: undefined,
range: 'range' in reference ? reference.range : undefined,
title: data.options?.status?.description ?? data.title
});
}
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/vs/workbench/contrib/chat/browser/chatEditingActions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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'
}],
Expand Down
70 changes: 63 additions & 7 deletions src/vs/workbench/contrib/chat/browser/chatEditingService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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';
Expand Down Expand Up @@ -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();
}
}

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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();
}
}
Expand Down
14 changes: 13 additions & 1 deletion src/vs/workbench/contrib/chat/browser/chatInputPart.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<WorkbenchList<IChatCollapsibleListItem>> | undefined;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand Down
12 changes: 11 additions & 1 deletion src/vs/workbench/contrib/chat/browser/chatWidget.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions src/vs/workbench/contrib/chat/common/chatEditingService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ export const enum WorkingSetEntryState {
Modified,
Accepted,
Rejected,
Transient,
Attached,
Sent,
}
Expand Down
3 changes: 3 additions & 0 deletions src/vs/workbench/contrib/chat/common/chatWidgetHistoryService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand All @@ -20,6 +22,7 @@ export interface IChatHistoryEntry {
export interface IChatInputState {
[key: string]: any;
chatContextAttachments?: ReadonlyArray<IChatRequestVariableEntry>;
chatWorkingSet?: ReadonlyArray<{ uri: URI; state: WorkingSetEntryState }>;
}

export const IChatWidgetHistoryService = createDecorator<IChatWidgetHistoryService>('IChatWidgetHistoryService');
Expand Down

0 comments on commit f4d5254

Please sign in to comment.