Skip to content

Commit

Permalink
Separate IModifiedFileEntry into text and notebook files
Browse files Browse the repository at this point in the history
  • Loading branch information
@DonJayamanne
DonJayamanne committed Nov 21, 2024
1 parent 927f53d commit 97cd534
Show file tree
Hide file tree
Showing 7 changed files with 112 additions and 47 deletions.
11 changes: 6 additions & 5 deletions src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedFileEntry.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,12 @@ import { IUndoRedoService } from '../../../../../platform/undoRedo/common/undoRe
import { SaveReason } from '../../../../common/editor.js';
import { IResolvedTextFileEditorModel, stringToSnapshot } from '../../../../services/textfile/common/textfiles.js';
import { IChatAgentResult } from '../../common/chatAgents.js';
import { ChatEditKind, IModifiedFileEntry, WorkingSetEntryState } from '../../common/chatEditingService.js';
import { ChatEditKind, IModifiedTextFileEntry, ITextSnapshotEntry, WorkingSetEntryState } from '../../common/chatEditingService.js';
import { IChatService } from '../../common/chatService.js';
import { ChatEditingSnapshotTextModelContentProvider, ChatEditingTextModelContentProvider } from './chatEditingTextModelContentProviders.js';

export class ChatEditingModifiedFileEntry extends Disposable implements IModifiedFileEntry {

export class ChatEditingModifiedFileEntry extends Disposable implements IModifiedTextFileEntry {
public readonly kind = 'text';
public static readonly scheme = 'modified-file-entry';
private static lastEntryId = 0;
public readonly entryId = `${ChatEditingModifiedFileEntry.scheme}::${++ChatEditingModifiedFileEntry.lastEntryId}`;
Expand Down Expand Up @@ -179,9 +179,10 @@ export class ChatEditingModifiedFileEntry extends Disposable implements IModifie
this._telemetryInfo = telemetryInfo;
}

createSnapshot(requestId: string | undefined): ISnapshotEntry {
createSnapshot(requestId: string | undefined): ITextSnapshotEntry {
this._isFirstEditAfterStartOrSnapshot = true;
return {
kind: 'text',
resource: this.modifiedURI,
languageId: this.modifiedModel.getLanguageId(),
snapshotUri: ChatEditingSnapshotTextModelContentProvider.getSnapshotFileURI(requestId, this.modifiedURI.path),
Expand All @@ -192,7 +193,7 @@ export class ChatEditingModifiedFileEntry extends Disposable implements IModifie
telemetryInfo: this._telemetryInfo
};
}
restoreFromSnapshot(snapshot: ISnapshotEntry) {
restoreFromSnapshot(snapshot: ITextSnapshotEntry) {
this._stateObs.set(snapshot.state, undefined);
this.docSnapshot.setValue(snapshot.original);
this._setDocValue(snapshot.current);
Expand Down
43 changes: 26 additions & 17 deletions src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingSession.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,10 @@ import { IEditorService } from '../../../../services/editor/common/editorService
import { MultiDiffEditor } from '../../../multiDiffEditor/browser/multiDiffEditor.js';
import { MultiDiffEditorInput } from '../../../multiDiffEditor/browser/multiDiffEditorInput.js';
import { ChatAgentLocation, IChatAgentService } from '../../common/chatAgents.js';
import { ChatEditingSessionChangeType, ChatEditingSessionState, ChatEditKind, IChatEditingSession, WorkingSetDisplayMetadata, WorkingSetEntryRemovalReason, WorkingSetEntryState } from '../../common/chatEditingService.js';
import { ChatEditingSessionChangeType, ChatEditingSessionState, ChatEditKind, IChatEditingSession, IModifiedEntryTelemetryInfo, IModifiedFileEntry, ITextSnapshotEntry, WorkingSetDisplayMetadata, WorkingSetEntryRemovalReason, WorkingSetEntryState } from '../../common/chatEditingService.js';
import { IChatResponseModel } from '../../common/chatModel.js';
import { ChatEditingMultiDiffSourceResolver } from './chatEditingService.js';
import { ChatEditingModifiedFileEntry, IModifiedEntryTelemetryInfo, ISnapshotEntry } from './chatEditingModifiedFileEntry.js';
import { ChatEditingModifiedFileEntry } from './chatEditingModifiedFileEntry.js';
import { ChatEditingTextModelContentProvider } from './chatEditingTextModelContentProviders.js';
import { Schemas } from '../../../../../base/common/network.js';
import { isEqual, joinPath } from '../../../../../base/common/resources.js';
Expand All @@ -58,8 +58,8 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
private readonly _initialFileContents = new ResourceMap<string>();
private readonly _filesToSkipCreating = new ResourceSet();

private readonly _entriesObs = observableValue<readonly ChatEditingModifiedFileEntry[]>(this, []);
public get entries(): IObservable<readonly ChatEditingModifiedFileEntry[]> {
private readonly _entriesObs = observableValue<readonly IModifiedFileEntry[]>(this, []);
public get entries(): IObservable<readonly IModifiedFileEntry[]> {
this._assertNotDisposed();
return this._entriesObs;
}
Expand Down Expand Up @@ -233,9 +233,11 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
for (const [file, state] of this._workingSet) {
workingSet.set(file, state);
}
const entries = new ResourceMap<ISnapshotEntry>();
const entries = new ResourceMap<ITextSnapshotEntry>();
for (const entry of this._entriesObs.get()) {
entries.set(entry.modifiedURI, entry.createSnapshot(requestId));
if (entry.kind !== 'notebook') {
entries.set(entry.modifiedURI, entry.createSnapshot(requestId));
}
}
return {
requestId,
Expand All @@ -251,7 +253,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
}

const snapshotEntry = [...entries.values()].find((e) => isEqual(e.snapshotUri, snapshotUri));
if (!snapshotEntry) {
if (!snapshotEntry || snapshotEntry.kind !== 'text') {
return null;
}

Expand Down Expand Up @@ -297,17 +299,23 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
// but which are not found in the snapshot
for (const entry of this._entriesObs.get()) {
const snapshotEntry = snapshot.entries.get(entry.modifiedURI);
if (!snapshotEntry) {
if (entry.kind === 'text' && !snapshotEntry) {
entry.resetToInitialValue();
entry.dispose();
}
}

const entriesArr: ChatEditingModifiedFileEntry[] = [];
const entriesArr: IModifiedFileEntry[] = [];
// Restore all entries from the snapshot
for (const snapshotEntry of snapshot.entries.values()) {
const entry = await this._getOrCreateModifiedFileEntry(snapshotEntry.resource, snapshotEntry.telemetryInfo);
entry.restoreFromSnapshot(snapshotEntry);
if (entry.kind === 'text' && snapshotEntry.kind === 'text') {
entry.restoreFromSnapshot(snapshotEntry);
// } else if (entry.kind === 'notebook') {
// entry.restoreFromSnapshot(snapshotEntry);
} else {
throw new Error('Unexpected snapshot entry kind');
}
entriesArr.push(entry);
}

Expand Down Expand Up @@ -448,7 +456,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
this._assertNotDisposed();

const entry = this._entriesObs.get().find(e => e.entryId === documentId);
return entry?.originalModel ?? null;
return entry?.kind === 'text' ? entry?.originalModel ?? null : null;
}

acceptStreamingEditsStart(): void {
Expand Down Expand Up @@ -581,7 +589,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
this._onDidChange.fire(ChatEditingSessionChangeType.Other);
}

private async _getOrCreateModifiedFileEntry(resource: URI, responseModel: IModifiedEntryTelemetryInfo): Promise<ChatEditingModifiedFileEntry> {
private async _getOrCreateModifiedFileEntry(resource: URI, responseModel: IModifiedEntryTelemetryInfo): Promise<IModifiedFileEntry> {
const existingEntry = this._entriesObs.get().find(e => isEqual(e.modifiedURI, resource));
if (existingEntry) {
if (responseModel.requestId !== existingEntry.telemetryInfo.requestId) {
Expand All @@ -592,7 +600,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
const initialContent = this._initialFileContents.get(resource);
// This gets manually disposed in .dispose() or in .restoreSnapshot()
const entry = await this._createModifiedFileEntry(resource, responseModel, false, initialContent);
if (!initialContent) {
if (!initialContent && entry.kind === 'text') {
this._initialFileContents.set(resource, entry.initialContent);
}
// If an entry is deleted e.g. reverting a created file,
Expand Down Expand Up @@ -655,7 +663,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
return result;
};
const deserializeChatEditingSessionSnapshot = async (snapshot: IChatEditingSessionSnapshotDTO) => {
const entriesMap = new ResourceMap<ISnapshotEntry>();
const entriesMap = new ResourceMap<ITextSnapshotEntry>();
for (const entryDTO of snapshot.entries) {
const entry = await deserializeSnapshotEntry(entryDTO);
entriesMap.set(entry.resource, entry);
Expand All @@ -668,6 +676,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
};
const deserializeSnapshotEntry = async (entry: ISnapshotEntryDTO) => {
return {
kind: 'text',
resource: URI.parse(entry.resource),
languageId: entry.languageId,
original: await getFileContent(entry.originalHash),
Expand All @@ -676,7 +685,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
state: entry.state,
snapshotUri: URI.parse(entry.snapshotUri),
telemetryInfo: { requestId: entry.telemetryInfo.requestId, agentId: entry.telemetryInfo.agentId, command: entry.telemetryInfo.command, sessionId: this.chatSessionId, result: undefined }
} satisfies ISnapshotEntry;
} satisfies ITextSnapshotEntry;
};
try {
const stateFilePath = joinPath(storageLocation, STORAGE_STATE_FILE);
Expand Down Expand Up @@ -757,7 +766,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
entries: Array.from(snapshot.entries.values()).map(serializeSnapshotEntry)
} satisfies IChatEditingSessionSnapshotDTO);
};
const serializeSnapshotEntry = (entry: ISnapshotEntry) => {
const serializeSnapshotEntry = (entry: ITextSnapshotEntry) => {
return {
resource: entry.resource.toString(),
languageId: entry.languageId,
Expand Down Expand Up @@ -810,7 +819,7 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
export interface IChatEditingSessionSnapshot {
readonly requestId: string | undefined;
readonly workingSet: ResourceMap<WorkingSetDisplayMetadata>;
readonly entries: ResourceMap<ISnapshotEntry>;
readonly entries: ResourceMap<ITextSnapshotEntry>;
}

interface IChatEditingSessionSnapshotDTO {
Expand Down
8 changes: 4 additions & 4 deletions src/vs/workbench/contrib/chat/browser/chatEditorController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ import { InlineDecoration, InlineDecorationType } from '../../../../editor/commo
import { localize } from '../../../../nls.js';
import { IContextKey, IContextKeyService, RawContextKey } from '../../../../platform/contextkey/common/contextkey.js';
import { minimapGutterAddedBackground, minimapGutterDeletedBackground, minimapGutterModifiedBackground, overviewRulerAddedForeground, overviewRulerDeletedForeground, overviewRulerModifiedForeground } from '../../scm/browser/dirtydiffDecorator.js';
import { ChatEditingSessionState, IChatEditingService, IModifiedFileEntry, WorkingSetEntryState } from '../common/chatEditingService.js';
import { ChatEditingSessionState, IChatEditingService, IModifiedTextFileEntry, WorkingSetEntryState } from '../common/chatEditingService.js';
import { Event } from '../../../../base/common/event.js';
import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
import { MenuId } from '../../../../platform/actions/common/actions.js';
Expand Down Expand Up @@ -81,7 +81,7 @@ export class ChatEditorController extends Disposable implements IEditorContribut
const model = modelObs.read(r);

const session = this._chatEditingService.currentEditingSessionObs.read(r);
const entry = session?.entries.read(r).find(e => isEqual(e.modifiedURI, model?.uri));
const entry = session?.entries.read(r).find(e => isEqual(e.modifiedURI, model?.uri) && e.kind === 'text') as IModifiedTextFileEntry | undefined;

if (!entry || entry.state.read(r) !== WorkingSetEntryState.Modified) {
this._clearRendering();
Expand Down Expand Up @@ -144,7 +144,7 @@ export class ChatEditorController extends Disposable implements IEditorContribut
this._ctxHasEditorModification.reset();
}

private _updateWithDiff(entry: IModifiedFileEntry, diff: IDocumentDiff | null | undefined): void {
private _updateWithDiff(entry: IModifiedTextFileEntry, diff: IDocumentDiff | null | undefined): void {
if (!diff) {
this._clearRendering();
return;
Expand Down Expand Up @@ -491,7 +491,7 @@ class DiffHunkWidget implements IOverlayWidget {


constructor(
readonly entry: IModifiedFileEntry,
readonly entry: IModifiedTextFileEntry,
private readonly _undoEdits: ISingleEditOperation[],
private readonly _versionId: number,
private readonly _editor: ICodeEditor,
Expand Down
60 changes: 58 additions & 2 deletions src/vs/workbench/contrib/chat/common/chatEditingService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,15 @@ import { IDisposable } from '../../../../base/common/lifecycle.js';
import { ResourceMap } from '../../../../base/common/map.js';
import { IObservable, ITransaction } from '../../../../base/common/observable.js';
import { URI } from '../../../../base/common/uri.js';
import { OffsetEdit } from '../../../../editor/common/core/offsetEdit.js';
import { IDocumentDiff } from '../../../../editor/common/diff/documentDiffProvider.js';
import { TextEdit } from '../../../../editor/common/languages.js';
import { ITextModel } from '../../../../editor/common/model.js';
import { localize } from '../../../../nls.js';
import { RawContextKey } from '../../../../platform/contextkey/common/contextkey.js';
import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js';
import { ICellEditOperation } from '../../notebook/common/notebookCommon.js';
import { IChatAgentResult } from './chatAgents.js';
import { IChatResponseModel } from './chatModel.js';

export const IChatEditingService = createDecorator<IChatEditingService>('chatEditingService');
Expand Down Expand Up @@ -110,19 +113,72 @@ export const enum ChatEditingSessionChangeType {
Other,
}

export interface IModifiedFileEntry {

export interface IBaseSnapshotEntry {
readonly resource: URI;
readonly snapshotUri: URI;
readonly state: WorkingSetEntryState;
telemetryInfo: IModifiedEntryTelemetryInfo;
}

export interface ITextSnapshotEntry extends IBaseSnapshotEntry {
kind: 'text';
readonly languageId: string;
readonly original: string;
readonly current: string;
readonly originalToCurrentEdit: OffsetEdit;
}

export interface INotebookSnapshotEntry extends IBaseSnapshotEntry {
kind: 'notebook';
}
export type ISnapshotEntry = ITextSnapshotEntry | INotebookSnapshotEntry;

export interface IModifiedEntryTelemetryInfo {
agentId: string | undefined;
command: string | undefined;
sessionId: string;
requestId: string;
result: IChatAgentResult | undefined;
}

interface IModifiedAnyFileEntry extends IDisposable {
readonly entryId: string;
readonly originalURI: URI;
readonly originalModel: ITextModel;
readonly modifiedURI: URI;
readonly onDidDelete: Event<void>;
readonly state: IObservable<WorkingSetEntryState>;
readonly isCurrentlyBeingModified: IObservable<boolean>;
readonly rewriteRatio: IObservable<number>;
readonly diffInfo: IObservable<IDocumentDiff>;
readonly lastModifyingRequestId: string;
readonly telemetryInfo: IModifiedEntryTelemetryInfo;
accept(transaction: ITransaction | undefined): Promise<void>;
reject(transaction: ITransaction | undefined): Promise<void>;
acceptAgentEdits(textEdits: TextEdit[], isLastEdits: boolean): void;
acceptStreamingEditsStart(tx: ITransaction): void;
acceptStreamingEditsEnd(tx: ITransaction): void;
updateTelemetryInfo(telemetryInfo: IModifiedEntryTelemetryInfo): void;
}

export interface IModifiedTextFileEntry extends IModifiedAnyFileEntry {
readonly kind: 'text';
readonly originalModel: ITextModel;
readonly modifiedModel: ITextModel;
resetToInitialValue(): void;
createSnapshot(requestId: string | undefined): ITextSnapshotEntry;
restoreFromSnapshot(snapshot: ITextSnapshotEntry): void;
}

export interface IModifiedNotebookFileEntry extends IModifiedAnyFileEntry {
readonly kind: 'notebook';
acceptAgentEdits(textEdits: (TextEdit | ICellEditOperation)[], isLastEdits: boolean): void;
createSnapshot(requestId: string | undefined): INotebookSnapshotEntry;
restoreFromSnapshot(snapshot: INotebookSnapshotEntry): void;
}

export type IModifiedFileEntry = IModifiedTextFileEntry | IModifiedNotebookFileEntry;

export interface IChatEditingSessionStream {
textEdits(resource: URI, textEdits: TextEdit[], isLastEdits: boolean, responseModel: IChatResponseModel): void;
}
Expand Down
4 changes: 2 additions & 2 deletions src/vs/workbench/contrib/notebook/browser/contrib/chatEdit/notebookChatEditController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import { isEqual } from '../../../../../../base/common/resources.js';
import { Disposable, dispose, IReference, toDisposable } from '../../../../../../base/common/lifecycle.js';
import { autorun, derived, derivedWithStore, observableFromEvent, observableValue } from '../../../../../../base/common/observable.js';
import { IChatEditingService, WorkingSetEntryState } from '../../../../chat/common/chatEditingService.js';
import { IChatEditingService, IModifiedTextFileEntry, WorkingSetEntryState } from '../../../../chat/common/chatEditingService.js';
import { NotebookTextModel } from '../../../common/model/notebookTextModel.js';
import { INotebookEditor, INotebookEditorContribution } from '../../notebookBrowser.js';
import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js';
Expand Down Expand Up @@ -86,7 +86,7 @@ class NotebookChatEditorController extends Disposable {
if (!model || !session) {
return;
}
return session.entries.read(r).find(e => isEqual(e.modifiedURI, model.uri));
return session.entries.read(r).find(e => isEqual(e.modifiedURI, model.uri) && e.kind === 'text') as IModifiedTextFileEntry | undefined;
}).recomputeInitiallyAndOnChange(this._store);


Expand Down
8 changes: 4 additions & 4 deletions ...vs/workbench/contrib/notebook/browser/contrib/chatEdit/notebookOriginalModelRefFactory.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
*--------------------------------------------------------------------------------------------*/

import { AsyncReferenceCollection, IReference, ReferenceCollection } from '../../../../../../base/common/lifecycle.js';
import { IModifiedFileEntry } from '../../../../chat/common/chatEditingService.js';
import { IModifiedTextFileEntry } from '../../../../chat/common/chatEditingService.js';
import { INotebookService } from '../../../common/notebookService.js';
import { bufferToStream, VSBuffer } from '../../../../../../base/common/buffer.js';
import { NotebookTextModel } from '../../../common/model/notebookTextModel.js';
Expand All @@ -15,7 +15,7 @@ export const INotebookOriginalModelReferenceFactory = createDecorator<INotebookO

export interface INotebookOriginalModelReferenceFactory {
readonly _serviceBrand: undefined;
getOrCreate(fileEntry: IModifiedFileEntry, viewType: string): Promise<IReference<NotebookTextModel>>;
getOrCreate(fileEntry: IModifiedTextFileEntry, viewType: string): Promise<IReference<NotebookTextModel>>;
}


Expand All @@ -25,7 +25,7 @@ export class OriginalNotebookModelReferenceCollection extends ReferenceCollectio
super();
}

protected override async createReferencedObject(key: string, fileEntry: IModifiedFileEntry, viewType: string): Promise<NotebookTextModel> {
protected override async createReferencedObject(key: string, fileEntry: IModifiedTextFileEntry, viewType: string): Promise<NotebookTextModel> {
this.modelsToDispose.delete(key);
const uri = fileEntry.originalURI;
const model = this.notebookService.getNotebookTextModel(uri);
Expand Down Expand Up @@ -83,7 +83,7 @@ export class NotebookOriginalModelReferenceFactory implements INotebookOriginalM
constructor(@IInstantiationService private readonly instantiationService: IInstantiationService) {
}

getOrCreate(fileEntry: IModifiedFileEntry, viewType: string): Promise<IReference<NotebookTextModel>> {
getOrCreate(fileEntry: IModifiedTextFileEntry, viewType: string): Promise<IReference<NotebookTextModel>> {
return this.asyncModelCollection.acquire(fileEntry.originalURI.toString(), fileEntry, viewType);
}
}
Expand Down
Loading

0 comments on commit 97cd534

Please sign in to comment.