Improve AI diff UX

packages/ai-chat-ui/src/browser/chat-input-widget.tsx

@@ -378,8 +378,8 @@ const buildChangeSetUI = (changeSet: ChangeSet, labelProvider: LabelProvider, on
         name: element.name ?? labelProvider.getName(element.uri),
         additionalInfo: element.additionalInfo ?? labelProvider.getDetails(element.uri),
         openChange: element?.openChange?.bind(element),
-        accept: element.state !== 'applied' ? element?.accept?.bind(element) : undefined,
-        discard: element.state === 'applied' ? element?.discard?.bind(element) : undefined,
+        accept: element.state !== 'applied' ? element?.apply?.bind(element) : undefined,
+        discard: element.state === 'applied' ? element?.revert?.bind(element) : undefined,
         delete: () => onDeleteChangeSetElement(changeSet.getElements().indexOf(element))
     }))
 });
@@ -484,15 +484,15 @@ const ChatInputOptions: React.FunctionComponent<ChatInputOptionsProps> = ({ left
 );
 
 function acceptAllPendingElements(changeSet: ChangeSet): void {
-    acceptablePendingElements(changeSet).forEach(e => e.accept!());
+    acceptablePendingElements(changeSet).forEach(e => e.apply!());
 }
 
 function hasPendingElementsToAccept(changeSet: ChangeSet): boolean | undefined {
     return acceptablePendingElements(changeSet).length > 0;
 }
 
 function acceptablePendingElements(changeSet: ChangeSet): ChangeSetElement[] {
-    return changeSet.getElements().filter(e => e.accept && (e.state === undefined || e.state === 'pending'));
+    return changeSet.getElements().filter(e => e.apply && (e.state === undefined || e.state === 'pending'));
 }
 
 function getLatestRequest(chatModel: ChatModel): ChatRequestModel | undefined {

packages/ai-chat/src/browser/change-set-file-element.ts

@@ -14,7 +14,7 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { URI } from '@theia/core';
+import { MEMORY_TEXT, URI } from '@theia/core';
 import { inject, injectable, postConstruct } from '@theia/core/shared/inversify';
 import { ChangeSetElement, ChangeSetImpl } from '../common';
 import { createChangeSetFileUri } from './change-set-file-resource';
@@ -47,7 +47,7 @@ export class ChangeSetFileElement implements ChangeSetElement {
     @inject(ChangeSetFileService)
     protected readonly changeSetFileService: ChangeSetFileService;
 
-    protected _state: 'pending' | 'applied' | 'discarded' | undefined;
+    protected _state: 'pending' | 'applied' | undefined;
 
     protected originalContent: string | undefined;
 
@@ -64,6 +64,10 @@ export class ChangeSetFileElement implements ChangeSetElement {
         return this.elementProps.uri;
     }
 
+    get readOnlyUri(): URI {
+        return this.elementProps.uri.withScheme(MEMORY_TEXT).withQuery(this.originalContent ?? '');
+    }
+
     get changedUri(): URI {
         return createChangeSetFileUri(this.elementProps.chatSessionId, this.uri);
     }
@@ -80,11 +84,11 @@ export class ChangeSetFileElement implements ChangeSetElement {
         return this.changeSetFileService.getAdditionalInfo(this.uri);
     }
 
-    get state(): 'pending' | 'applied' | 'discarded' | undefined {
+    get state(): 'pending' | 'applied' | undefined {
         return this._state ?? this.elementProps.state;
     }
 
-    protected set state(value: 'pending' | 'applied' | 'discarded' | undefined) {
+    protected set state(value: 'pending' | 'applied' | undefined) {
         this._state = value;
         this.elementProps.changeSet.notifyChange();
     }
@@ -107,31 +111,36 @@ export class ChangeSetFileElement implements ChangeSetElement {
 
     async openChange(): Promise<void> {
         this.changeSetFileService.openDiff(
-            this.uri,
+            this.readOnlyUri,
             this.changedUri
         );
     }
 
-    async accept(contents?: string): Promise<void> {
-        this.state = 'applied';
-        if (this.type === 'delete') {
+    async apply(contents?: string): Promise<void> {
+        if (await this.changeSetFileService.trySave(this.changedUri)) { /** Continue */ } else if (this.type === 'delete') {
             await this.changeSetFileService.delete(this.uri);
             this.state = 'applied';
-            return;
+        } else {
+            await this.writeChanges(contents);
         }
+        this.changeSetFileService.closeDiffsFor(this.readOnlyUri);
+    }
 
-        await this.changeSetFileService.write(this.uri, contents !== undefined ? contents : this.targetState);
+    async writeChanges(contents?: string): Promise<void> {
+        await this.changeSetFileService.writeFrom(this.changedUri, this.uri, contents ?? this.targetState);
+        this.state = 'applied';
     }
 
-    async discard(): Promise<void> {
-        this.state = 'discarded';
+    async revert(): Promise<void> {
+        this.state = 'pending';
         if (this.type === 'add') {
             await this.changeSetFileService.delete(this.uri);
-            return;
-        }
-        if (this.originalContent) {
+        } else if (this.originalContent) {
             await this.changeSetFileService.write(this.uri, this.originalContent);
         }
     }
 
+    dispose(): void {
+        this.changeSetFileService.closeDiffsFor(this.readOnlyUri);
+    }
 }

packages/ai-chat/src/browser/change-set-file-resource.ts

@@ -18,6 +18,7 @@ import { Resource, ResourceResolver, ResourceSaveOptions, URI } from '@theia/cor
 import { inject, injectable } from '@theia/core/shared/inversify';
 import { ChatService } from '../common';
 import { ChangeSetFileElement } from './change-set-file-element';
+import { ChangeSetFileService } from './change-set-file-service';
 
 export const CHANGE_SET_FILE_RESOURCE_SCHEME = 'changeset-file';
 const QUERY = 'uri=';
@@ -36,6 +37,9 @@ export class ChangeSetFileResourceResolver implements ResourceResolver {
     @inject(ChatService)
     protected readonly chatService: ChatService;
 
+    @inject(ChangeSetFileService)
+    protected readonly changeSetFileService: ChangeSetFileService;
+
     async resolve(uri: URI): Promise<Resource> {
         if (uri.scheme !== CHANGE_SET_FILE_RESOURCE_SCHEME) {
             throw new Error('The given uri is not a change set file uri: ' + uri);
@@ -53,22 +57,36 @@ export class ChangeSetFileResourceResolver implements ResourceResolver {
         }
 
         const fileUri = decodeURIComponent(uri.query.toString().replace(QUERY, ''));
-        const element = changeSet.getElements().find(e => e.uri.path.toString() === fileUri);
-        if (!(element instanceof ChangeSetFileElement)) {
-            throw new Error('Change set element not found: ' + fileUri);
-        }
+        let current: ChangeSetFileElement | undefined = undefined;
+        const refreshElement = () => {
+            const element = changeSet.getElements().find(e => e.uri.path.toString() === fileUri);
+            if (element instanceof ChangeSetFileElement) {
+                current = element;
+            } else if (!current) {
+                throw new Error('Change set element not found: ' + fileUri);
+            }
+            return current;
+        };
+        refreshElement();
+
+        const changeListener = changeSet.onDidChange(() => {
+            try {
+                const activeElement = refreshElement();
+                this.changeSetFileService.updateEditorsWithNewSuggestion(activeElement.changedUri, activeElement.targetState);
+            } catch { /** No op - we should be closing any editors anyway. */ }
+        });
 
         return {
             uri,
             readOnly: false,
             initiallyDirty: true,
-            readContents: async () => element.targetState ?? '',
-            saveContents: async (content: string, options?: ResourceSaveOptions): Promise<void> => {
-                element.accept(content);
-            },
-            dispose: () => { }
+            autosaveable: false,
+            readContents: async () => refreshElement().targetState ?? '',
+            saveContents: async (content: string, options?: ResourceSaveOptions): Promise<void> => refreshElement().writeChanges(),
+            dispose: () => {
+                changeListener.dispose();
+            }
         };
     }
 
 }
-

packages/ai-chat/src/browser/change-set-file-service.ts

@@ -14,8 +14,8 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { ILogger, UNTITLED_SCHEME, URI } from '@theia/core';
-import { DiffUris, LabelProvider, OpenerService, open } from '@theia/core/lib/browser';
+import { ILogger, URI } from '@theia/core';
+import { ApplicationShell, DiffUris, LabelProvider, NavigatableWidget, OpenerService, open } from '@theia/core/lib/browser';
 import { inject, injectable } from '@theia/core/shared/inversify';
 import { EditorManager } from '@theia/editor/lib/browser';
 import { FileService } from '@theia/filesystem/lib/browser/file-service';
@@ -40,6 +40,9 @@ export class ChangeSetFileService {
     @inject(EditorManager)
     protected readonly editorManager: EditorManager;
 
+    @inject(ApplicationShell)
+    protected readonly shell: ApplicationShell;
+
     @inject(MonacoWorkspace)
     protected readonly monacoWorkspace: MonacoWorkspace;
 
@@ -95,15 +98,14 @@ export class ChangeSetFileService {
     }
 
     async openDiff(originalUri: URI, suggestedUri: URI): Promise<void> {
-        const exists = await this.fileService.exists(originalUri);
-        const openedUri = exists ? originalUri : originalUri.withScheme(UNTITLED_SCHEME);
-        // Currently we don't have a great way to show the suggestions in a diff editor with accept/reject buttons
-        // So we just use plain diffs with the suggestions as original and the current state as modified, so users can apply changes in their current state
-        // But this leads to wrong colors and wrong label (revert change instead of accept change)
-        const diffUri = DiffUris.encode(openedUri, suggestedUri,
+        const diffUri = this.getDiffUri(originalUri, suggestedUri);
+        open(this.openerService, diffUri);
+    }
+
+    protected getDiffUri(originalUri: URI, suggestedUri: URI): URI {
+        return DiffUris.encode(originalUri, suggestedUri,
             `AI Changes: ${this.labelProvider.getName(originalUri)}`,
         );
-        open(this.openerService, diffUri);
     }
 
     async delete(uri: URI): Promise<void> {
@@ -113,24 +115,42 @@ export class ChangeSetFileService {
         }
     }
 
-    async write(uri: URI, targetState: string): Promise<void> {
-        const exists = await this.fileService.exists(uri);
-        if (!exists) {
-            await this.fileService.create(uri, targetState);
+    /** Returns true if there was a document available to save for the specified URI. */
+    async trySave(suggestedUri: URI): Promise<boolean> {
+        const openModel = this.monacoWorkspace.getTextDocument(suggestedUri.toString());
+        if (openModel) {
+            await openModel.save();
+            return true;
+        } else {
+            return false;
         }
-        await this.doWrite(uri, targetState);
     }
 
-    protected async doWrite(uri: URI, text: string): Promise<void> {
+    async writeFrom(from: URI, to: URI, fallbackContent: string): Promise<void> {
+        const authoritativeContent = this.monacoWorkspace.getTextDocument(from.toString())?.getText() ?? fallbackContent;
+        await this.write(to, authoritativeContent);
+    }
+
+    async write(uri: URI, text: string): Promise<void> {
         const document = this.monacoWorkspace.getTextDocument(uri.toString());
         if (document) {
             await this.monacoWorkspace.applyBackgroundEdit(document, [{
                 range: document.textEditorModel.getFullModelRange(),
                 text
-            }], (editor, wasDirty) => editor === undefined || !wasDirty);
+            }], () => true);
         } else {
             await this.fileService.write(uri, text);
         }
     }
 
+    updateEditorsWithNewSuggestion(uri: URI, content: string): void {
+        const existingModel = this.monacoWorkspace.getTextDocument(uri.toString());
+        existingModel?.textEditorModel.setValue(content);
+    }
+
+    closeDiffsFor(originalUri: URI): void {
+        // Diff editors return the URI of the left side as their URI.
+        const openEditors = this.shell.widgets.filter(widget => NavigatableWidget.getUri(widget)?.isEqual(originalUri));
+        openEditors.forEach(editor => editor.close());
+    }
 }

packages/ai-chat/src/common/chat-model.ts

@@ -95,6 +95,7 @@ export interface ChatModel {
 export interface ChangeSet {
     readonly title: string;
     getElements(): ChangeSetElement[];
+    onDidChange: Event<void>;
 }
 
 export interface ChangeSetElement {
@@ -104,14 +105,15 @@ export interface ChangeSetElement {
     readonly icon?: string;
     readonly additionalInfo?: string;
 
-    readonly state?: 'pending' | 'applied' | 'discarded';
+    readonly state?: 'pending' | 'applied';
     readonly type?: 'add' | 'modify' | 'delete';
     readonly data?: { [key: string]: unknown };
 
     open?(): Promise<void>;
     openChange?(): Promise<void>;
-    accept?(): Promise<void>;
-    discard?(): Promise<void>;
+    apply?(): Promise<void>;
+    revert?(): Promise<void>;
+    dispose?(): void;
 }
 
 export interface ChatRequest {
@@ -577,7 +579,8 @@ export class ChangeSetImpl implements ChangeSet {
     }
 
     removeElement(index: number): void {
-        this._elements.splice(index, 1);
+        const deletion = this._elements.splice(index, 1);
+        deletion.forEach(deleted => deleted.dispose?.());
         this.notifyChange();
     }
 

packages/core/src/browser/saveable-service.ts

@@ -162,7 +162,7 @@ export class SaveableService implements FrontendApplicationContribution {
             // Never auto-save untitled documents
             return false;
         } else {
-            return saveable.dirty;
+            return saveable.autosaveable !== false && saveable.dirty;
         }
     }
 

packages/core/src/browser/saveable.ts

@@ -28,6 +28,8 @@ export type AutoSaveMode = 'off' | 'afterDelay' | 'onFocusChange' | 'onWindowCha
 
 export interface Saveable {
     readonly dirty: boolean;
+    /** If false, the saveable will not participate in autosaving. */
+    readonly autosaveable?: boolean;
     /**
      * This event is fired when the content of the `dirty` variable changes.
      */

packages/core/src/common/resource.ts

@@ -62,6 +62,8 @@ export interface Resource extends Disposable {
     readonly readOnly?: boolean | MarkdownString;
 
     readonly initiallyDirty?: boolean;
+    /** If false, the application should not attempt to auto-save this resource. */
+    readonly autosaveable?: boolean;
     /**
      * Reads latest content of this resource.
      *
@@ -380,7 +382,8 @@ export class UntitledResourceResolver implements ResourceResolver {
 export class UntitledResource implements Resource {
 
     protected readonly onDidChangeContentsEmitter = new Emitter<void>();
-    initiallyDirty: boolean;
+    readonly initiallyDirty: boolean;
+    readonly autosaveable = false;
     get onDidChangeContents(): Event<void> {
         return this.onDidChangeContentsEmitter.event;
     }

packages/monaco/src/browser/monaco-editor-model.ts

@@ -259,6 +259,10 @@ export class MonacoEditorModel implements IResolvedTextEditorModel, TextEditorDo
         return this.resource.uri.toString();
     }
 
+    get autosaveable(): boolean | undefined {
+        return this.resource.autosaveable;
+    }
+
     protected _languageId: string | undefined;
     get languageId(): string {
         return this._languageId !== undefined ? this._languageId : this.model.getLanguageId();