@theia/ai-chat 1.60.2 → 1.61.0

package/src/common/chat-model.ts

@@ -19,12 +19,12 @@
  *--------------------------------------------------------------------------------------------*/
 // Partially copied from https://github.com/microsoft/vscode/blob/a2cab7255c0df424027be05d58e1b7b941f4ea60/src/vs/workbench/contrib/chat/common/chatModel.ts
 
-import { CancellationToken, CancellationTokenSource, Command, Disposable, Emitter, Event, generateUuid, URI } from '@theia/core';
+import { AIVariableResolutionRequest, LanguageModelMessage, ResolvedAIContextVariable, TextMessage, ThinkingMessage, ToolResultMessage, ToolUseMessage } from '@theia/ai-core';
+import { CancellationToken, CancellationTokenSource, Command, Disposable, DisposableCollection, Emitter, Event, generateUuid, URI } from '@theia/core';
 import { MarkdownString, MarkdownStringImpl } from '@theia/core/lib/common/markdown-rendering';
 import { Position } from '@theia/core/shared/vscode-languageserver-protocol';
 import { ChatAgentLocation } from './chat-agents';
 import { ParsedChatRequest } from './parsed-chat-request';
-import { AIVariableResolutionRequest, LanguageModelMessage, ResolvedAIContextVariable, TextMessage, ThinkingMessage, ToolResultMessage, ToolUseMessage } from '@theia/ai-core';
 
 /**********************
  * INTERFACES AND TYPE GUARDS
@@ -35,16 +35,46 @@ export type ChatChangeEvent =
     | ChatAddResponseEvent
     | ChatAddVariableEvent
     | ChatRemoveVariableEvent
+    | ChatSetVariablesEvent
     | ChatRemoveRequestEvent
     | ChatSetChangeSetEvent
+    | ChatSuggestionsChangedEvent
     | ChatUpdateChangeSetEvent
-    | ChatRemoveChangeSetEvent;
+    | ChatRemoveChangeSetEvent
+    | ChatEditRequestEvent
+    | ChatEditCancelEvent
+    | ChatEditSubmitEvent
+    | ChatChangeHierarchyBranchEvent;
 
 export interface ChatAddRequestEvent {
     kind: 'addRequest';
     request: ChatRequestModel;
 }
 
+export interface ChatEditRequestEvent {
+    kind: 'enableEdit';
+    request: EditableChatRequestModel;
+    branch: ChatHierarchyBranch<ChatRequestModel>;
+}
+
+export interface ChatEditCancelEvent {
+    kind: 'cancelEdit';
+    request: EditableChatRequestModel;
+    branch: ChatHierarchyBranch<ChatRequestModel>;
+}
+
+export interface ChatEditSubmitEvent {
+    kind: 'submitEdit';
+    request: EditableChatRequestModel;
+    branch: ChatHierarchyBranch<ChatRequestModel>;
+    newRequest: ChatRequest;
+}
+
+export interface ChatChangeHierarchyBranchEvent {
+    kind: 'changeHierarchyBranch';
+    branch: ChatHierarchyBranch<ChatRequestModel>;
+}
+
 export interface ChatAddResponseEvent {
     kind: 'addResponse';
     response: ChatResponseModel;
@@ -74,6 +104,15 @@ export interface ChatRemoveVariableEvent {
     kind: 'removeVariable';
 }
 
+export interface ChatSetVariablesEvent {
+    kind: 'setVariables';
+}
+
+export interface ChatSuggestionsChangedEvent {
+    kind: 'suggestionsChanged';
+    suggestions: ChatSuggestion[];
+}
+
 export namespace ChatChangeEvent {
     export function isChangeSetEvent(event: ChatChangeEvent): event is ChatSetChangeSetEvent | ChatUpdateChangeSetEvent | ChatRemoveChangeSetEvent {
         return event.kind === 'setChangeSet' || event.kind === 'removeChangeSet' || event.kind === 'updateChangeSet';
@@ -89,14 +128,76 @@ export interface ChatRemoveRequestEvent {
     reason: ChatRequestRemovalReason;
 }
 
+/**
+ * A model that contains information about a chat request that may branch off.
+ *
+ * The hierarchy of requests is represented by a tree structure.
+ * - The root of the tree is the initial request
+ * - Within each branch, the requests are stored in a list. Those requests are the alternatives to the original request.
+ *   Each of those items can have a next branch, which is the next request in the hierarchy.
+ */
+export interface ChatRequestHierarchy<TRequest extends ChatRequestModel = ChatRequestModel> {
+    readonly branch: ChatHierarchyBranch<TRequest>
+
+    onDidChange: Event<ChangeActiveBranchEvent<TRequest>>;
+
+    append(request: TRequest): ChatHierarchyBranch<TRequest>;
+    activeRequests(): TRequest[];
+    activeBranches(): ChatHierarchyBranch<TRequest>[];
+    findRequest(requestId: string): TRequest | undefined;
+    findBranch(requestId: string): ChatHierarchyBranch<TRequest> | undefined;
+
+    notifyChange(event: ChangeActiveBranchEvent<TRequest>): void
+}
+
+export interface ChangeActiveBranchEvent<TRequest extends ChatRequestModel = ChatRequestModel> {
+    branch: ChatHierarchyBranch<TRequest>,
+    item: ChatHierarchyBranchItem<TRequest>
+}
+
+/**
+ * A branch of the chat request hierarchy.
+ * It contains a list of items, each representing a request.
+ * Those items can have a next branch, which is the next request in the hierarchy.
+ */
+export interface ChatHierarchyBranch<TRequest extends ChatRequestModel = ChatRequestModel> {
+    readonly id: string;
+    readonly hierarchy: ChatRequestHierarchy<TRequest>;
+    readonly previous?: ChatHierarchyBranch<TRequest>;
+    readonly items: ChatHierarchyBranchItem<TRequest>[];
+    readonly activeBranchIndex: number;
+
+    next(): ChatHierarchyBranch<TRequest> | undefined;
+    get(): TRequest;
+    add(request: TRequest): void;
+    remove(request: TRequest | string): void;
+    /**
+     * Create a new branch by inserting it as the next branch of the active item.
+     */
+    continue(request: TRequest): ChatHierarchyBranch<TRequest>;
+
+    enable(request: TRequest): ChatHierarchyBranchItem<TRequest>;
+    enablePrevious(): ChatHierarchyBranchItem<TRequest>;
+    enableNext(): ChatHierarchyBranchItem<TRequest>;
+
+    succeedingBranches(): ChatHierarchyBranch<TRequest>[];
+}
+
+export interface ChatHierarchyBranchItem<TRequest extends ChatRequestModel = ChatRequestModel> {
+    readonly element: TRequest;
+    readonly next?: ChatHierarchyBranch<TRequest>;
+}
+
 export interface ChatModel {
     readonly onDidChange: Event<ChatChangeEvent>;
     readonly id: string;
     readonly location: ChatAgentLocation;
     readonly changeSet?: ChangeSet;
     readonly context: ChatContextManager;
+    readonly suggestions: readonly ChatSuggestion[];
     readonly settings?: { [key: string]: unknown };
     getRequests(): ChatRequestModel[];
+    getBranches(): ChatHierarchyBranch<ChatRequestModel>[];
     isEmpty(): boolean;
 }
 
@@ -107,8 +208,26 @@ export interface ChangeSet extends Disposable {
     dispose(): void;
 }
 
+export interface ChatSuggestionCallback {
+    kind: 'callback',
+    callback: () => unknown;
+    content: string | MarkdownString;
+}
+export namespace ChatSuggestionCallback {
+    export function is(candidate: ChatSuggestion): candidate is ChatSuggestionCallback {
+        return typeof candidate === 'object' && 'callback' in candidate;
+    }
+    export function containsCallbackLink(candidate: ChatSuggestion): candidate is ChatSuggestionCallback {
+        if (!is(candidate)) { return false; }
+        const text = typeof candidate.content === 'string' ? candidate.content : candidate.content.value;
+        return text.includes('](_callback)');
+    }
+}
+
+export type ChatSuggestion = | string | MarkdownString | ChatSuggestionCallback;
+
 export interface ChatContextManager {
-    onDidChange: Event<ChatAddVariableEvent | ChatRemoveVariableEvent>;
+    onDidChange: Event<ChatAddVariableEvent | ChatRemoveVariableEvent | ChatSetVariablesEvent>;
     getVariables(): readonly AIVariableResolutionRequest[]
     addVariables(...variables: AIVariableResolutionRequest[]): void;
     deleteVariables(...indices: number[]): void;
@@ -134,9 +253,21 @@ export interface ChangeSetElement {
     dispose?(): void;
 }
 
+export interface ChangeSetDecoration {
+    readonly priority?: number;
+    readonly additionalInfoSuffixIcon?: string[];
+}
+
 export interface ChatRequest {
     readonly text: string;
     readonly displayText?: string;
+    /**
+     * If the request has been triggered in the context of
+     * an existing request, this id will be set to the id of the
+     * referenced request.
+     */
+    readonly referencedRequestId?: string;
+    readonly variables?: readonly AIVariableResolutionRequest[];
 }
 
 export interface ChatContext {
@@ -180,6 +311,29 @@ export namespace ChatRequestModel {
     }
 }
 
+export interface EditableChatRequestModel extends ChatRequestModel {
+    readonly isEditing: boolean;
+    editContextManager: ChatContextManagerImpl;
+    enableEdit(): void;
+    cancelEdit(): void;
+    submitEdit(newRequest: ChatRequest): void;
+}
+
+export namespace EditableChatRequestModel {
+    export function is(request: unknown): request is EditableChatRequestModel {
+        return !!(
+            ChatRequestModel.is(request) &&
+            'enableEdit' in request &&
+            'cancelEdit' in request &&
+            'submitEdit' in request
+        );
+    }
+
+    export function isEditing(request: unknown): request is EditableChatRequestModel {
+        return is(request) && request.isEditing;
+    }
+}
+
 export interface ChatProgressMessage {
     kind: 'progressMessage';
     id: string;
@@ -278,6 +432,12 @@ export interface ThinkingChatResponseContent
     signature: string;
 }
 
+export interface ProgressChatResponseContent
+    extends Required<ChatResponseContent> {
+    kind: 'progress';
+    message: string;
+}
+
 export interface Location {
     uri: URI;
     position: Position;
@@ -414,6 +574,17 @@ export namespace ThinkingChatResponseContent {
     }
 }
 
+export namespace ProgressChatResponseContent {
+    export function is(obj: unknown): obj is ProgressChatResponseContent {
+        return (
+            ChatResponseContent.is(obj) &&
+            obj.kind === 'progress' &&
+            'message' in obj &&
+            typeof obj.message === 'string'
+        );
+    }
+}
+
 export type QuestionResponseHandler = (
     selectedOption: { text: string, value?: string },
 ) => void;
@@ -519,25 +690,44 @@ export class MutableChatModel implements ChatModel, Disposable {
     protected readonly _onDidChangeEmitter = new Emitter<ChatChangeEvent>();
     onDidChange: Event<ChatChangeEvent> = this._onDidChangeEmitter.event;
 
-    protected _requests: MutableChatRequestModel[];
+    protected readonly toDispose = new DisposableCollection();
+
+    protected _hierarchy: ChatRequestHierarchy<MutableChatRequestModel>;
     protected _id: string;
     protected _changeSet?: ChangeSetImpl;
+    protected _suggestions: readonly ChatSuggestion[] = [];
     protected readonly _contextManager = new ChatContextManagerImpl();
     protected _settings: { [key: string]: unknown };
 
     constructor(public readonly location = ChatAgentLocation.Panel) {
         // TODO accept serialized data as a parameter to restore a previously saved ChatModel
-        this._requests = [];
+        this._hierarchy = new ChatRequestHierarchyImpl<MutableChatRequestModel>();
         this._id = generateUuid();
-        this._contextManager.onDidChange(e => this._onDidChangeEmitter.fire(e));
+
+        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,
+            })),
+        ]);
+    }
+
+    getBranches(): ChatHierarchyBranch<ChatRequestModel>[] {
+        return this._hierarchy.activeBranches();
+    }
+
+    getBranch(requestId: string): ChatHierarchyBranch<ChatRequestModel> | undefined {
+        return this._hierarchy.findBranch(requestId);
     }
 
     getRequests(): MutableChatRequestModel[] {
-        return this._requests;
+        return this._hierarchy.activeRequests();
     }
 
     getRequest(id: string): MutableChatRequestModel | undefined {
-        return this._requests.find(request => request.id === id);
+        return this.getRequests().find(request => request.id === id);
     }
 
     get id(): string {
@@ -548,6 +738,10 @@ export class MutableChatModel implements ChatModel, Disposable {
         return this._changeSet;
     }
 
+    get suggestions(): readonly ChatSuggestion[] {
+        return this._suggestions;
+    }
+
     get context(): ChatContextManager {
         return this._contextManager;
     }
@@ -593,8 +787,13 @@ export class MutableChatModel implements ChatModel, Disposable {
     }
 
     addRequest(parsedChatRequest: ParsedChatRequest, agentId?: string, context: ChatContext = { variables: [] }): MutableChatRequestModel {
+        if (parsedChatRequest.request.referencedRequestId) {
+            return this.applyEdit(parsedChatRequest, agentId, context);
+        }
+
         const requestModel = new MutableChatRequestModel(this, parsedChatRequest, agentId, context);
-        this._requests.push(requestModel);
+        this.toDispose.push(requestModel);
+        this._hierarchy.append(requestModel);
         this._onDidChangeEmitter.fire({
             kind: 'addRequest',
             request: requestModel,
@@ -602,13 +801,217 @@ export class MutableChatModel implements ChatModel, Disposable {
         return requestModel;
     }
 
+    setSuggestions(suggestions: ChatSuggestion[]): void {
+        this._suggestions = Object.freeze(suggestions);
+        this._onDidChangeEmitter.fire({
+            kind: 'suggestionsChanged',
+            suggestions
+        });
+    }
+
     isEmpty(): boolean {
-        return this._requests.length === 0;
+        return this.getRequests().length === 0;
+    }
+
+    applyEdit(parsedChatRequest: ParsedChatRequest, agentId?: string, context: ChatContext = { variables: [] }): MutableChatRequestModel {
+        const requestId = parsedChatRequest.request.referencedRequestId!;
+        const branch = this._hierarchy.findBranch(requestId);
+        if (!branch) {
+            throw new Error(`Cannot find branch for requestId: ${requestId}`);
+        }
+
+        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;
     }
 
     dispose(): void {
         this.removeChangeSet(); // Signal disposal of last change set.
-        this._onDidChangeEmitter.dispose();
+        this.toDispose.dispose();
+    }
+
+    emit(event: ChatChangeEvent): void {
+        this._onDidChangeEmitter.fire(event);
+    }
+}
+
+export class ChatRequestHierarchyImpl<TRequest extends ChatRequestModel = ChatRequestModel> implements ChatRequestHierarchy<TRequest> {
+    protected readonly onDidChangeActiveBranchEmitter = new Emitter<ChangeActiveBranchEvent<TRequest>>();
+    readonly onDidChange = this.onDidChangeActiveBranchEmitter.event;
+
+    readonly branch: ChatHierarchyBranch<TRequest> = new ChatRequestHierarchyBranchImpl<TRequest>(this);
+
+    append(request: TRequest): ChatHierarchyBranch<TRequest> {
+        const branches = this.activeBranches();
+
+        if (branches.length === 0) {
+            this.branch.add(request);
+            return this.branch;
+        }
+
+        return branches.at(-1)!.continue(request);
+    }
+
+    activeRequests(): TRequest[] {
+        return this.activeBranches().map(h => h.get());
+    }
+
+    activeBranches(): ChatHierarchyBranch<TRequest>[] {
+        return Array.from(this.iterateBranches());
+    }
+
+    protected *iterateBranches(): Generator<ChatHierarchyBranch<TRequest>> {
+        let current: ChatHierarchyBranch<TRequest> | undefined = this.branch;
+        while (current) {
+            if (current.items.length > 0) {
+                yield current;
+                current = current.next();
+            } else {
+                break;
+            }
+        }
+    }
+
+    findRequest(requestId: string): TRequest | undefined {
+        const branch = this.findInBranch(this.branch, requestId);
+        return branch?.items.find(item => item.element.id === requestId)?.element;
+    }
+
+    findBranch(requestId: string): ChatHierarchyBranch<TRequest> | undefined {
+        return this.findInBranch(this.branch, requestId);
+    }
+
+    protected findInBranch(branch: ChatHierarchyBranch<TRequest>, requestId: string): ChatHierarchyBranch<TRequest> | undefined {
+        for (const item of branch.items) {
+            if (item.element.id === requestId) {
+                return branch;
+            }
+        }
+        for (const item of branch.items) {
+            if (item.next) {
+                const found = this.findInBranch(item.next, requestId);
+                if (found) {
+                    return found;
+                }
+            }
+        }
+        return undefined;
+    }
+
+    notifyChange(event: ChangeActiveBranchEvent<TRequest>): void {
+        this.onDidChangeActiveBranchEmitter.fire(event);
+    }
+}
+
+export class ChatRequestHierarchyBranchImpl<TRequest extends ChatRequestModel> implements ChatHierarchyBranch<TRequest> {
+    readonly id = generateUuid();
+
+    constructor(
+        readonly hierarchy: ChatRequestHierarchy<TRequest>,
+        readonly previous?: ChatHierarchyBranch<TRequest>,
+        readonly items: ChatHierarchyBranchItem<TRequest>[] = [],
+        protected _activeIndex = -1
+    ) { }
+
+    get activeBranchIndex(): number {
+        return this._activeIndex;
+    }
+
+    protected set activeBranchIndex(value: number) {
+        this._activeIndex = value;
+        this.hierarchy.notifyChange({
+            branch: this,
+            item: this.items[this._activeIndex]
+        });
+    }
+
+    next(): ChatHierarchyBranch<TRequest> | undefined {
+        return this.items[this.activeBranchIndex]?.next;
+    }
+
+    get(): TRequest {
+        return this.items[this.activeBranchIndex].element;
+    }
+
+    add(request: TRequest): void {
+        const branch: ChatHierarchyBranchItem<TRequest> = {
+            element: request
+        };
+        this.items.push(branch);
+        this.activeBranchIndex = this.items.length - 1;
+    }
+
+    remove(request: TRequest | string): void {
+        const requestId = typeof request === 'string' ? request : request.id;
+        const index = this.items.findIndex(version => version.element.id === requestId);
+        if (index !== -1) {
+            this.items.splice(index, 1);
+            if (this.activeBranchIndex >= index) {
+                this.activeBranchIndex--;
+            }
+        }
+    }
+
+    continue(request: TRequest): ChatHierarchyBranch<TRequest> {
+        if (this.items.length === 0) {
+            this.add(request);
+            return this;
+        }
+
+        const item = this.items[this.activeBranchIndex];
+
+        if (item) {
+            const next = new ChatRequestHierarchyBranchImpl(this.hierarchy, this, [{ element: request }], 0);
+            this.items[this.activeBranchIndex] = {
+                ...item,
+                next
+            };
+            return next;
+        }
+
+        throw new Error(`No current branch to continue from. Active Index: ${this.activeBranchIndex}`);
+    }
+
+    enable(request: TRequest): ChatHierarchyBranchItem<TRequest> {
+        this.activeBranchIndex = this.items.findIndex(pred => pred.element.id === request.id);
+        return this.items[this.activeBranchIndex];
+    }
+
+    enablePrevious(): ChatHierarchyBranchItem<TRequest> {
+        if (this.activeBranchIndex > 0) {
+            this.activeBranchIndex--;
+            return this.items[this.activeBranchIndex];
+        }
+        return this.items[0];
+    }
+
+    enableNext(): ChatHierarchyBranchItem<TRequest> {
+        if (this.activeBranchIndex < this.items.length - 1) {
+            this.activeBranchIndex++;
+            return this.items[this.activeBranchIndex];
+        }
+
+        return this.items[this.activeBranchIndex];
+    }
+
+    succeedingBranches(): ChatHierarchyBranch<TRequest>[] {
+        const branches: ChatHierarchyBranch<TRequest>[] = [];
+
+        let current: ChatHierarchyBranch<TRequest> | undefined = this;
+        while (current !== undefined) {
+            branches.push(current);
+            current = current.next();
+        }
+
+        return branches;
     }
 }
 
@@ -676,11 +1079,17 @@ export class ChangeSetImpl implements ChangeSet {
 
 export class ChatContextManagerImpl implements ChatContextManager {
     protected readonly variables = new Array<AIVariableResolutionRequest>();
-    protected readonly onDidChangeEmitter = new Emitter<ChatAddVariableEvent | ChatRemoveVariableEvent>();
-    get onDidChange(): Event<ChatAddVariableEvent | ChatRemoveVariableEvent> {
+    protected readonly onDidChangeEmitter = new Emitter<ChatAddVariableEvent | ChatRemoveVariableEvent | ChatSetVariablesEvent>();
+    get onDidChange(): Event<ChatAddVariableEvent | ChatRemoveVariableEvent | ChatSetVariablesEvent> {
         return this.onDidChangeEmitter.event;
     }
 
+    constructor(context?: ChatContext) {
+        if (context) {
+            this.variables.push(...context.variables.map(AIVariableResolutionRequest.fromResolved));
+        }
+    }
+
     getVariables(): readonly AIVariableResolutionRequest[] {
         const result = this.variables.slice();
         Object.freeze(result);
@@ -711,6 +1120,17 @@ export class ChatContextManagerImpl implements ChatContextManager {
         }
     }
 
+    setVariables(variables: AIVariableResolutionRequest[]): void {
+        this.variables.length = 0;
+        variables.forEach(variable => {
+            if (this.variables.some(existing => existing.variable.id === variable.variable.id && existing.arg === variable.arg)) {
+                return;
+            }
+            this.variables.push(variable);
+        });
+        this.onDidChangeEmitter.fire({ kind: 'setVariables' });
+    }
+
     clear(): void {
         if (this.variables.length) {
             this.variables.length = 0;
@@ -719,7 +1139,7 @@ export class ChatContextManagerImpl implements ChatContextManager {
     }
 }
 
-export class MutableChatRequestModel implements ChatRequestModel {
+export class MutableChatRequestModel implements ChatRequestModel, EditableChatRequestModel, Disposable {
     protected readonly _id: string;
     protected _session: MutableChatModel;
     protected _request: ChatRequest;
@@ -727,6 +1147,10 @@ export class MutableChatRequestModel implements ChatRequestModel {
     protected _context: ChatContext;
     protected _agentId?: string;
     protected _data: { [key: string]: unknown };
+    protected _isEditing = false;
+
+    protected readonly toDispose = new DisposableCollection();
+    readonly editContextManager: ChatContextManagerImpl;
 
     constructor(session: MutableChatModel, public readonly message: ParsedChatRequest, agentId?: string,
         context: ChatContext = { variables: [] }, data: { [key: string]: unknown } = {}) {
@@ -738,6 +1162,20 @@ export class MutableChatRequestModel implements ChatRequestModel {
         this._context = context;
         this._agentId = agentId;
         this._data = data;
+
+        this.editContextManager = new ChatContextManagerImpl(context);
+        this.toDispose.pushAll([
+            this.editContextManager.onDidChange(e => this.session.emit(e))
+        ]);
+    }
+
+    get isEditing(): boolean {
+        return this._isEditing;
+    }
+
+    enableEdit(): void {
+        this._isEditing = true;
+        this.emitEditRequest(this);
     }
 
     get data(): { [key: string]: unknown } | undefined {
@@ -776,9 +1214,78 @@ export class MutableChatRequestModel implements ChatRequestModel {
         return this._agentId;
     }
 
+    cancelEdit(): void {
+        if (this.isEditing) {
+            this._isEditing = false;
+            this.emitCancelEdit(this);
+
+            this.clearEditContext();
+        }
+    }
+
+    submitEdit(newRequest: ChatRequest): void {
+        if (this.isEditing) {
+            this._isEditing = false;
+            const variables = this.editContextManager.getVariables() ?? [];
+
+            this.emitSubmitEdit(this, {
+                ...newRequest,
+                referencedRequestId: this.id,
+                variables
+            });
+
+            this.clearEditContext();
+        }
+    }
+
     cancel(): void {
         this.response.cancel();
     }
+
+    dispose(): void {
+        this.toDispose.dispose();
+    }
+
+    protected clearEditContext(): void {
+        this.editContextManager.setVariables(this.context.variables.map(AIVariableResolutionRequest.fromResolved));
+    }
+
+    protected emitEditRequest(request: MutableChatRequestModel): void {
+        const branch = this.session.getBranch(request.id);
+        if (!branch) {
+            throw new Error(`Cannot find hierarchy for requestId: ${request.id}`);
+        }
+        this.session.emit({
+            kind: 'enableEdit',
+            request,
+            branch,
+        });
+    }
+
+    protected emitCancelEdit(request: MutableChatRequestModel): void {
+        const branch = this.session.getBranch(request.id);
+        if (!branch) {
+            throw new Error(`Cannot find branch for requestId: ${request.id}`);
+        }
+        this.session.emit({
+            kind: 'cancelEdit',
+            request,
+            branch,
+        });
+    }
+
+    protected emitSubmitEdit(request: MutableChatRequestModel, newRequest: ChatRequest): void {
+        const branch = this.session.getBranch(request.id);
+        if (!branch) {
+            throw new Error(`Cannot find branch for requestId: ${request.id}`);
+        }
+        this.session.emit({
+            kind: 'submitEdit',
+            request,
+            branch,
+            newRequest
+        });
+    }
 }
 
 export class ErrorChatResponseContentImpl implements ErrorChatResponseContent {
@@ -1044,8 +1551,7 @@ export const COMMAND_CHAT_RESPONSE_COMMAND: Command = {
 export class CommandChatResponseContentImpl implements CommandChatResponseContent {
     readonly kind = 'command';
 
-    constructor(public command?: Command, public customCallback?: CustomCallback, protected args?: unknown[]) {
-    }
+    constructor(public command?: Command, public customCallback?: CustomCallback, protected args?: unknown[]) { }
 
     get arguments(): unknown[] {
         return this.args ?? [];
@@ -1127,6 +1633,12 @@ class ChatResponseImpl implements ChatResponse {
         return this._content;
     }
 
+    clearContent(): void {
+        this._content = [];
+        this._updateResponseRepresentation();
+        this._onDidChangeEmitter.fire();
+    }
+
     addContents(contents: ChatResponseContent[]): void {
         contents.forEach(c => this.doAddContent(c));
         this._onDidChangeEmitter.fire();
@@ -1353,3 +1865,40 @@ export class ErrorChatResponseModel extends MutableChatResponseModel {
         this.error(error);
     }
 }
+
+export class ProgressChatResponseContentImpl implements ProgressChatResponseContent {
+    readonly kind = 'progress';
+    protected _message: string;
+
+    constructor(message: string) {
+        this._message = message;
+    }
+
+    get message(): string {
+        return this._message;
+    }
+
+    asString(): string {
+        return JSON.stringify({
+            type: 'progress',
+            message: this.message
+        });
+    }
+
+    asDisplayString(): string | undefined {
+        return `<Progress>${this.message}</Progress>`;
+    }
+
+    merge(nextChatResponseContent: ProgressChatResponseContent): boolean {
+        this._message = nextChatResponseContent.message;
+        return true;
+    }
+
+    toLanguageModelMessage(): TextMessage {
+        return {
+            actor: 'ai',
+            type: 'text',
+            text: this.message
+        };
+    }
+}