@theia/ai-chat 1.62.1 → 1.63.0-next.24

package/lib/common/chat-model.js

@@ -20,16 +20,19 @@
  *--------------------------------------------------------------------------------------------*/
 // Partially copied from https://github.com/microsoft/vscode/blob/a2cab7255c0df424027be05d58e1b7b941f4ea60/src/vs/workbench/contrib/chat/common/chatModel.ts
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.ProgressChatResponseContentImpl = exports.ErrorChatResponseModel = exports.MutableChatResponseModel = exports.QuestionResponseContentImpl = exports.HorizontalLayoutChatResponseContentImpl = exports.CommandChatResponseContentImpl = exports.COMMAND_CHAT_RESPONSE_COMMAND = exports.ToolCallChatResponseContentImpl = exports.CodeChatResponseContentImpl = exports.InformationalChatResponseContentImpl = exports.MarkdownChatResponseContentImpl = exports.ThinkingChatResponseContentImpl = exports.TextChatResponseContentImpl = exports.ErrorChatResponseContentImpl = exports.MutableChatRequestModel = exports.ChatContextManagerImpl = exports.ChangeSetImpl = exports.ChatRequestHierarchyBranchImpl = exports.ChatRequestHierarchyImpl = exports.MutableChatModel = exports.QuestionResponseContent = exports.ProgressChatResponseContent = exports.ThinkingChatResponseContent = exports.ErrorChatResponseContent = exports.ToolCallChatResponseContent = exports.HorizontalLayoutChatResponseContent = exports.CodeChatResponseContent = exports.CommandChatResponseContent = exports.InformationalChatResponseContent = exports.MarkdownChatResponseContent = exports.TextChatResponseContent = exports.Location = exports.ChatResponseContent = exports.EditableChatRequestModel = exports.ChatRequestModel = exports.ChatSuggestionCallback = exports.ChatChangeEvent = void 0;
+exports.ProgressChatResponseContentImpl = exports.ErrorChatResponseModel = exports.MutableChatResponseModel = exports.QuestionResponseContentImpl = exports.HorizontalLayoutChatResponseContentImpl = exports.CommandChatResponseContentImpl = exports.COMMAND_CHAT_RESPONSE_COMMAND = exports.ToolCallChatResponseContentImpl = exports.CodeChatResponseContentImpl = exports.InformationalChatResponseContentImpl = exports.MarkdownChatResponseContentImpl = exports.ThinkingChatResponseContentImpl = exports.TextChatResponseContentImpl = exports.ErrorChatResponseContentImpl = exports.MutableChatRequestModel = exports.ChatContextManagerImpl = exports.ChatRequestHierarchyBranchImpl = exports.ChatRequestHierarchyImpl = exports.ChatTreeChangeSet = exports.MutableChatModel = exports.QuestionResponseContent = exports.ProgressChatResponseContent = exports.ThinkingChatResponseContent = exports.ErrorChatResponseContent = exports.ToolCallChatResponseContent = exports.HorizontalLayoutChatResponseContent = exports.CodeChatResponseContent = exports.CommandChatResponseContent = exports.InformationalChatResponseContent = exports.MarkdownChatResponseContent = exports.TextChatResponseContent = exports.Location = exports.ChatResponseContent = exports.EditableChatRequestModel = exports.ChatRequestModel = exports.ChatSuggestionCallback = exports.ChatChangeEvent = exports.ChangeSetImpl = void 0;
 const ai_core_1 = require("@theia/ai-core");
 const core_1 = require("@theia/core");
 const markdown_rendering_1 = require("@theia/core/lib/common/markdown-rendering");
 const vscode_languageserver_protocol_1 = require("@theia/core/shared/vscode-languageserver-protocol");
 const chat_agents_1 = require("./chat-agents");
+const change_set_1 = require("./change-set");
+Object.defineProperty(exports, "ChangeSetImpl", { enumerable: true, get: function () { return change_set_1.ChangeSetImpl; } });
+const debounce = require("@theia/core/shared/lodash.debounce");
 var ChatChangeEvent;
 (function (ChatChangeEvent) {
     function isChangeSetEvent(event) {
-        return event.kind === 'setChangeSet' || event.kind === 'removeChangeSet' || event.kind === 'updateChangeSet';
+        return event.kind === 'updateChangeSet';
     }
     ChatChangeEvent.isChangeSetEvent = isChangeSetEvent;
 })(ChatChangeEvent || (exports.ChatChangeEvent = ChatChangeEvent = {}));
@@ -252,16 +255,27 @@ class MutableChatModel {
         this._contextManager = new ChatContextManagerImpl();
         // TODO accept serialized data as a parameter to restore a previously saved ChatModel
         this._hierarchy = new ChatRequestHierarchyImpl();
+        this._changeSet = new ChatTreeChangeSet(this._hierarchy);
+        this.toDispose.push(this._changeSet);
+        this._changeSet.onDidChange(this._onDidChangeEmitter.fire, this._onDidChangeEmitter, this.toDispose);
         this._id = (0, core_1.generateUuid)();
         this.toDispose.pushAll([
             this._onDidChangeEmitter,
-            this._contextManager.onDidChange(e => this._onDidChangeEmitter.fire(e)),
-            this._hierarchy.onDidChange(event => this._onDidChangeEmitter.fire({
-                kind: 'changeHierarchyBranch',
-                branch: event.branch,
-            })),
+            this._contextManager.onDidChange(this._onDidChangeEmitter.fire, this._onDidChangeEmitter),
+            this._hierarchy.onDidChange(event => {
+                this._onDidChangeEmitter.fire({
+                    kind: 'changeHierarchyBranch',
+                    branch: event.branch,
+                });
+            }),
         ]);
     }
+    get id() {
+        return this._id;
+    }
+    get changeSet() {
+        return this._changeSet;
+    }
     getBranches() {
         return this._hierarchy.activeBranches();
     }
@@ -274,12 +288,6 @@ class MutableChatModel {
     getRequest(id) {
         return this.getRequests().find(request => request.id === id);
     }
-    get id() {
-        return this._id;
-    }
-    get changeSet() {
-        return this._changeSet;
-    }
     get suggestions() {
         return this._suggestions;
     }
@@ -292,49 +300,30 @@ class MutableChatModel {
     setSettings(settings) {
         this._settings = settings;
     }
-    setChangeSet(changeSet) {
-        if (!changeSet) {
-            return this.removeChangeSet();
-        }
-        const oldChangeSet = this._changeSet;
-        oldChangeSet === null || oldChangeSet === void 0 ? void 0 : oldChangeSet.dispose();
-        this._changeSet = changeSet;
-        this._onDidChangeEmitter.fire({
-            kind: 'setChangeSet',
-            changeSet,
-            oldChangeSet,
-        });
-        changeSet.onDidChange(() => {
-            this._onDidChangeEmitter.fire({
-                kind: 'updateChangeSet',
-                changeSet,
-            });
-        });
-    }
-    removeChangeSet() {
-        if (this._changeSet) {
-            const oldChangeSet = this._changeSet;
-            this._changeSet = undefined;
-            oldChangeSet.dispose();
-            this._onDidChangeEmitter.fire({
-                kind: 'removeChangeSet',
-                changeSet: oldChangeSet,
-            });
-        }
-    }
     addRequest(parsedChatRequest, agentId, context = { variables: [] }) {
-        if (parsedChatRequest.request.referencedRequestId) {
-            return this.applyEdit(parsedChatRequest, agentId, context);
-        }
+        const add = this.getTargetForRequestAddition(parsedChatRequest);
         const requestModel = new MutableChatRequestModel(this, parsedChatRequest, agentId, context);
-        this.toDispose.push(requestModel);
-        this._hierarchy.append(requestModel);
+        requestModel.onDidChange(event => {
+            if (!ChatChangeEvent.isChangeSetEvent(event)) {
+                this._onDidChangeEmitter.fire(event);
+            }
+        }, this, this.toDispose);
+        add(requestModel);
+        this._changeSet.registerRequest(requestModel);
         this._onDidChangeEmitter.fire({
             kind: 'addRequest',
             request: requestModel,
         });
         return requestModel;
     }
+    getTargetForRequestAddition(request) {
+        const requestId = request.request.referencedRequestId;
+        const branch = requestId !== undefined && this._hierarchy.findBranch(requestId);
+        if (requestId !== undefined && !branch) {
+            throw new Error(`Cannot find branch for requestId: ${requestId}`);
+        }
+        return branch ? branch.add.bind(branch) : this._hierarchy.append.bind(this._hierarchy);
+    }
     setSuggestions(suggestions) {
         this._suggestions = Object.freeze(suggestions);
         this._onDidChangeEmitter.fire({
@@ -345,31 +334,127 @@ class MutableChatModel {
     isEmpty() {
         return this.getRequests().length === 0;
     }
-    applyEdit(parsedChatRequest, agentId, context = { variables: [] }) {
-        const requestId = parsedChatRequest.request.referencedRequestId;
-        const branch = this._hierarchy.findBranch(requestId);
-        if (!branch) {
-            throw new Error(`Cannot find branch for requestId: ${requestId}`);
+    dispose() {
+        this.toDispose.dispose();
+    }
+}
+exports.MutableChatModel = MutableChatModel;
+class ChatTreeChangeSet {
+    get onDidChange() {
+        return this.onDidChangeEmitter.event;
+    }
+    constructor(hierarchy) {
+        this.hierarchy = hierarchy;
+        this.onDidChangeEmitter = new core_1.Emitter();
+        this.toDispose = new core_1.DisposableCollection();
+        this.currentElements = [];
+        this.handleChangeSetChange = debounce(this.doHandleChangeSetChange.bind(this), 100, { leading: false, trailing: true });
+        this.toDisposeOnRequestAdded = new core_1.DisposableCollection();
+        hierarchy.onDidChange(this.handleChangeSetChange, this, this.toDispose);
+    }
+    get title() {
+        var _a, _b;
+        return (_b = (_a = this.getCurrentChangeSet()) === null || _a === void 0 ? void 0 : _a.title) !== null && _b !== void 0 ? _b : '';
+    }
+    removeElements(...uris) {
+        return this.getMutableChangeSet().removeElements(...uris);
+    }
+    addElements(...elements) {
+        return this.getMutableChangeSet().addElements(...elements);
+    }
+    setElements(...elements) {
+        this.getMutableChangeSet().setElements(...elements);
+    }
+    setTitle(title) {
+        this.getMutableChangeSet().setTitle(title);
+    }
+    getElementByURI(uri) {
+        return this.currentElements.find(candidate => candidate.uri.isEqual(uri));
+    }
+    doHandleChangeSetChange() {
+        var _a;
+        const newElements = this.computeChangeSetElements();
+        this.handleElementChange(newElements);
+        this.currentElements = newElements;
+        this.onDidChangeEmitter.fire({ kind: 'updateChangeSet', elements: this.currentElements, title: (_a = this.getCurrentChangeSet()) === null || _a === void 0 ? void 0 : _a.title });
+    }
+    getElements() {
+        return this.currentElements;
+    }
+    computeChangeSetElements() {
+        const allElements = change_set_1.ChangeSetImpl.combine((function* (requests) {
+            for (let i = requests.length - 1; i >= 0; i--) {
+                const changeSet = requests[i].changeSet;
+                if (changeSet) {
+                    yield changeSet;
+                }
+            }
+        })(this.hierarchy.activeRequests()));
+        return core_1.ArrayUtils.coalesce(Array.from(allElements.values()));
+    }
+    handleElementChange(newElements) {
+        var _a, _b;
+        const old = new Set(this.currentElements);
+        for (const element of newElements) {
+            if (!old.delete(element)) {
+                (_a = element.onShow) === null || _a === void 0 ? void 0 : _a.call(element);
+            }
         }
-        const requestModel = new MutableChatRequestModel(this, parsedChatRequest, agentId, context);
-        this.toDispose.push(requestModel);
-        branch.add(requestModel);
-        this.removeChangeSet();
-        this._onDidChangeEmitter.fire({
-            kind: 'addRequest',
-            request: requestModel,
-        });
-        return requestModel;
+        for (const element of old) {
+            (_b = element.onHide) === null || _b === void 0 ? void 0 : _b.call(element);
+        }
+    }
+    registerRequest(request) {
+        request.onDidChange(event => event.kind === 'updateChangeSet' && this.handleChangeSetChange(), this, this.toDispose);
+        if (this.localChangeSet) {
+            request.changeSet = this.localChangeSet;
+            this.localChangeSet = undefined;
+        }
+        this.toDisposeOnRequestAdded.dispose();
+    }
+    getMutableChangeSet() {
+        const tipRequest = this.hierarchy.activeRequests().at(-1);
+        const existingChangeSet = tipRequest === null || tipRequest === void 0 ? void 0 : tipRequest.changeSet;
+        if (existingChangeSet) {
+            return existingChangeSet;
+        }
+        if (this.localChangeSet && tipRequest) {
+            throw new Error('Non-empty chat model retained reference to own change set. This is unexpected!');
+        }
+        if (this.localChangeSet) {
+            return this.localChangeSet;
+        }
+        const newChangeSet = new change_set_1.ChangeSetImpl();
+        if (tipRequest) {
+            tipRequest.changeSet = newChangeSet;
+        }
+        else {
+            this.localChangeSet = newChangeSet;
+            newChangeSet.onDidChange(this.handleChangeSetChange, this, this.toDisposeOnRequestAdded);
+        }
+        return newChangeSet;
+    }
+    getCurrentChangeSet() {
+        var _a;
+        const holder = this.getBranchParent(candidate => !!candidate.get().changeSet);
+        return (_a = holder === null || holder === void 0 ? void 0 : holder.get().changeSet) !== null && _a !== void 0 ? _a : this.localChangeSet;
+    }
+    /** Returns the lowest node among active nodes that satisfies {@link criterion} */
+    getBranchParent(criterion) {
+        const branches = this.hierarchy.activeBranches();
+        for (let i = branches.length - 1; i >= 0; i--) {
+            const branch = branches[i];
+            if (criterion === null || criterion === void 0 ? void 0 : criterion(branch)) {
+                return branch;
+            }
+        }
+        return branches.at(0);
     }
     dispose() {
-        this.removeChangeSet(); // Signal disposal of last change set.
         this.toDispose.dispose();
     }
-    emit(event) {
-        this._onDidChangeEmitter.fire(event);
-    }
 }
-exports.MutableChatModel = MutableChatModel;
+exports.ChatTreeChangeSet = ChatTreeChangeSet;
 class ChatRequestHierarchyImpl {
     constructor() {
         this.onDidChangeActiveBranchEmitter = new core_1.Emitter();
@@ -429,6 +514,10 @@ class ChatRequestHierarchyImpl {
     notifyChange(event) {
         this.onDidChangeActiveBranchEmitter.fire(event);
     }
+    dispose() {
+        this.onDidChangeActiveBranchEmitter.dispose();
+        this.branch.dispose();
+    }
 }
 exports.ChatRequestHierarchyImpl = ChatRequestHierarchyImpl;
 class ChatRequestHierarchyBranchImpl {
@@ -516,72 +605,14 @@ class ChatRequestHierarchyBranchImpl {
         }
         return branches;
     }
-}
-exports.ChatRequestHierarchyBranchImpl = ChatRequestHierarchyBranchImpl;
-class ChangeSetImpl {
-    constructor(title, elements = []) {
-        this.title = title;
-        this._onDidChangeEmitter = new core_1.Emitter();
-        this.onDidChange = this._onDidChangeEmitter.event;
-        this._elements = [];
-        this.addElements(...elements);
-    }
-    getElements() {
-        return this._elements;
-    }
-    /**
-     * Find an element by URI.
-     * @param uri The URI to look for.
-     * @returns The element with the given URI, or undefined if not found.
-     */
-    getElementByURI(uri) {
-        const uriString = uri.toString();
-        for (const element of this._elements) {
-            if (element.uri.toString() === uriString) {
-                return element;
-            }
-        }
-        return undefined;
-    }
-    /** Will replace any element that is already present, using URI as identity criterion. */
-    addElements(...elements) {
-        const added = [];
-        const modified = [];
-        const toDispose = [];
-        const current = new Map(this.getElements().map((element, index) => [element.uri.toString(), index]));
-        elements.forEach(element => {
-            var _a;
-            const existingIndex = current.get(element.uri.toString());
-            if (existingIndex !== undefined) {
-                modified.push(element.uri);
-                toDispose.push(this._elements[existingIndex]);
-                this._elements[existingIndex] = element;
-            }
-            else {
-                added.push(element.uri);
-                this._elements.push(element);
-            }
-            (_a = element.onDidChange) === null || _a === void 0 ? void 0 : _a.call(element, () => this.notifyChange({ state: [element.uri] }));
-        });
-        toDispose.forEach(element => { var _a; return (_a = element.dispose) === null || _a === void 0 ? void 0 : _a.call(element); });
-        this.notifyChange({ added, modified });
-    }
-    removeElements(...indices) {
-        // From highest to lowest so that we don't affect lower indices with our splicing.
-        const sorted = indices.slice().sort((left, right) => left - right);
-        const deletions = sorted.flatMap(index => this._elements.splice(index, 1));
-        deletions.forEach(deleted => { var _a; return (_a = deleted.dispose) === null || _a === void 0 ? void 0 : _a.call(deleted); });
-        this.notifyChange({ removed: deletions.map(element => element.uri) });
-    }
-    notifyChange(change) {
-        this._onDidChangeEmitter.fire(change);
-    }
     dispose() {
-        this._onDidChangeEmitter.dispose();
-        this._elements.forEach(element => { var _a; return (_a = element.dispose) === null || _a === void 0 ? void 0 : _a.call(element); });
+        if (core_1.Disposable.is(this.get())) {
+            this.items.forEach(({ element }) => core_1.Disposable.is(element) && element.dispose());
+        }
+        this.items.length = 0;
     }
 }
-exports.ChangeSetImpl = ChangeSetImpl;
+exports.ChatRequestHierarchyBranchImpl = ChatRequestHierarchyBranchImpl;
 class ChatContextManagerImpl {
     get onDidChange() {
         return this.onDidChangeEmitter.event;
@@ -641,6 +672,8 @@ exports.ChatContextManagerImpl = ChatContextManagerImpl;
 class MutableChatRequestModel {
     constructor(session, message, agentId, context = { variables: [] }, data = {}) {
         this.message = message;
+        this._onDidChangeEmitter = new core_1.Emitter();
+        this.onDidChange = this._onDidChangeEmitter.event;
         this._isEditing = false;
         this.toDispose = new core_1.DisposableCollection();
         // TODO accept serialized data as a parameter to restore a previously saved ChatRequestModel
@@ -652,9 +685,19 @@ class MutableChatRequestModel {
         this._agentId = agentId;
         this._data = data;
         this.editContextManager = new ChatContextManagerImpl(context);
-        this.toDispose.pushAll([
-            this.editContextManager.onDidChange(e => this.session.emit(e))
-        ]);
+        this.editContextManager.onDidChange(this._onDidChangeEmitter.fire, this._onDidChangeEmitter, this.toDispose);
+        this.toDispose.push(this._onDidChangeEmitter);
+    }
+    get changeSet() {
+        return this._changeSet;
+    }
+    set changeSet(changeSet) {
+        var _a;
+        (_a = this._changeSet) === null || _a === void 0 ? void 0 : _a.dispose();
+        this._changeSet = changeSet;
+        this.toDispose.push(changeSet);
+        changeSet.onDidChange(() => this._onDidChangeEmitter.fire({ kind: 'updateChangeSet', elements: changeSet.getElements(), title: changeSet.title }), this, this.toDispose);
+        this._onDidChangeEmitter.fire({ kind: 'updateChangeSet', elements: changeSet.getElements(), title: changeSet.title });
     }
     get isEditing() {
         return this._isEditing;
@@ -727,7 +770,7 @@ class MutableChatRequestModel {
         if (!branch) {
             throw new Error(`Cannot find hierarchy for requestId: ${request.id}`);
         }
-        this.session.emit({
+        this._onDidChangeEmitter.fire({
             kind: 'enableEdit',
             request,
             branch,
@@ -738,7 +781,7 @@ class MutableChatRequestModel {
         if (!branch) {
             throw new Error(`Cannot find branch for requestId: ${request.id}`);
         }
-        this.session.emit({
+        this._onDidChangeEmitter.fire({
             kind: 'cancelEdit',
             request,
             branch,
@@ -749,7 +792,7 @@ class MutableChatRequestModel {
         if (!branch) {
             throw new Error(`Cannot find branch for requestId: ${request.id}`);
         }
-        this.session.emit({
+        this._onDidChangeEmitter.fire({
             kind: 'submitEdit',
             request,
             branch,
@@ -914,6 +957,8 @@ class ToolCallChatResponseContentImpl {
         this._arguments = arg_string;
         this._finished = finished;
         this._result = result;
+        // Initialize the confirmation promise immediately
+        this._confirmed = this.createConfirmationPromise();
     }
     get id() {
         return this._id;
@@ -930,6 +975,48 @@ class ToolCallChatResponseContentImpl {
     get result() {
         return this._result;
     }
+    get confirmed() {
+        return this._confirmed;
+    }
+    /**
+     * Create a confirmation promise that can be resolved/rejected later
+     */
+    createConfirmationPromise() {
+        // The promise is always created, just ensure we have resolution handlers
+        if (!this._confirmationResolver) {
+            this._confirmed = new Promise((resolve, reject) => {
+                this._confirmationResolver = resolve;
+                this._confirmationRejecter = reject;
+            });
+        }
+        return this._confirmed;
+    }
+    /**
+     * Confirm the tool execution
+     */
+    confirm() {
+        if (this._confirmationResolver) {
+            this._confirmationResolver(true);
+        }
+    }
+    /**
+     * Deny the tool execution
+     */
+    deny() {
+        if (this._confirmationResolver) {
+            this._confirmationResolver(false);
+            this._finished = true;
+            this._result = 'Tool execution denied by user';
+        }
+    }
+    /**
+     * Cancel the confirmation (reject the promise)
+     */
+    cancelConfirmation(reason) {
+        if (this._confirmationRejecter) {
+            this._confirmationRejecter(reason);
+        }
+    }
     asString() {
         return '';
     }
@@ -943,6 +1030,7 @@ class ToolCallChatResponseContentImpl {
             this._result = nextChatResponseContent.result;
             const args = nextChatResponseContent.arguments;
             this._arguments = (args && args.length > 0) ? args : this._arguments;
+            // Don't merge confirmation promises - they should be managed separately
             return true;
         }
         if (nextChatResponseContent.name !== undefined) {