@theia/ai-chat-ui 1.63.0-next.0 → 1.63.0-next.52

package/src/browser/chat-node-toolbar-action-contribution.ts

@@ -78,6 +78,11 @@ export namespace ChatNodeToolbarCommands {
         id: 'chat:node:toolbar:cancel-request',
         category: CHAT_NODE_TOOLBAR_CATEGORY,
     }, '', CHAT_NODE_TOOLBAR_CATEGORY_KEY);
+
+    export const RETRY = Command.toLocalizedCommand({
+        id: 'chat:node:toolbar:retry-message',
+        category: CHAT_NODE_TOOLBAR_CATEGORY,
+    }, 'Retry', CHAT_NODE_TOOLBAR_CATEGORY_KEY);
 }
 
 export class DefaultChatNodeToolbarActionContribution implements ChatNodeToolbarActionContribution {
@@ -96,6 +101,15 @@ export class DefaultChatNodeToolbarActionContribution implements ChatNodeToolbar
                 tooltip: nls.localize('theia/ai/chat-ui/node/toolbar/edit', 'Edit'),
             }];
         } else {
+            const shouldShowRetry = node.response.isError || node.response.isCanceled;
+            if (shouldShowRetry) {
+                return [{
+                    commandId: ChatNodeToolbarCommands.RETRY.id,
+                    icon: codicon('refresh'),
+                    tooltip: nls.localize('theia/ai/chat-ui/node/toolbar/retry', 'Retry'),
+                    priority: -1 // Higher priority to show it first
+                }];
+            }
             return [];
         }
     }

package/src/browser/chat-response-renderer/delegation-response-renderer.tsx

@@ -0,0 +1,177 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclipseSource GmbH.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+import { inject, injectable } from '@theia/core/shared/inversify';
+import { ChatRequestInvocation, ChatResponseContent, ChatResponseModel } from '@theia/ai-chat';
+import { ChatResponsePartRenderer } from '../chat-response-part-renderer';
+import * as React from '@theia/core/shared/react';
+import { DelegationResponseContent, isDelegationResponseContent } from '@theia/ai-chat/lib/browser/delegation-response-content';
+import { ResponseNode } from '../chat-tree-view';
+import { CompositeTreeNode } from '@theia/core/lib/browser';
+import { SubChatWidgetFactory } from '../chat-tree-view/sub-chat-widget';
+import { DisposableCollection } from '@theia/core';
+
+@injectable()
+export class DelegationResponseRenderer implements ChatResponsePartRenderer<DelegationResponseContent> {
+
+    @inject(SubChatWidgetFactory)
+    subChatWidgetFactory: SubChatWidgetFactory;
+
+    canHandle(response: ChatResponseContent): number {
+        if (isDelegationResponseContent(response)) {
+            return 10;
+        }
+        return -1;
+    }
+    render(response: DelegationResponseContent, parentNode: ResponseNode): React.ReactNode {
+        return this.renderExpandableNode(response, parentNode);
+    }
+
+    private renderExpandableNode(response: DelegationResponseContent, parentNode: ResponseNode): React.ReactNode {
+        return <DelegatedChat
+            response={response.response}
+            agentId={response.agentId}
+            prompt={response.prompt}
+            parentNode={parentNode}
+            subChatWidgetFactory={this.subChatWidgetFactory} />;
+    }
+}
+
+interface DelegatedChatProps {
+    response: ChatRequestInvocation;
+    agentId: string;
+    prompt: string;
+    parentNode: ResponseNode;
+    subChatWidgetFactory: SubChatWidgetFactory;
+}
+
+interface DelegatedChatState {
+    node?: ResponseNode;
+}
+
+class DelegatedChat extends React.Component<DelegatedChatProps, DelegatedChatState> {
+    private widget: ReturnType<SubChatWidgetFactory>;
+    private readonly toDispose = new DisposableCollection();
+
+    constructor(props: DelegatedChatProps) {
+        super(props);
+        this.state = {
+            node: undefined
+        };
+        this.widget = props.subChatWidgetFactory();
+    }
+
+    override componentDidMount(): void {
+        // Start rendering as soon as the response is created (streaming mode)
+        this.props.response.responseCreated.then(chatModel => {
+            const node = mapResponseToNode(chatModel, this.props.parentNode);
+            this.setState({ node });
+
+            // Listen for changes to update the rendering as the response streams in
+            const changeListener = () => {
+                // Force re-render when the response content changes
+                this.forceUpdate();
+            };
+            this.toDispose.push(chatModel.onDidChange(changeListener));
+        }).catch(error => {
+            console.error('Failed to create delegated chat response:', error);
+            // Still try to handle completion in case of partial success
+        });
+
+        // Keep the completion handling for final cleanup if needed
+        this.props.response.responseCompleted.then(() => {
+            // Final update when response is complete
+            this.forceUpdate();
+        }).catch(error => {
+            console.error('Error in delegated chat response completion:', error);
+            // Force update anyway to show any partial content or error state
+            this.forceUpdate();
+        });
+    }
+
+    override componentWillUnmount(): void {
+        this.toDispose.dispose();
+    }
+
+    override render(): React.ReactNode {
+        const { agentId, prompt } = this.props;
+        const hasNode = !!this.state.node;
+        const isComplete = this.state.node?.response.isComplete ?? false;
+        const isCanceled = this.state.node?.response.isCanceled ?? false;
+        const isError = this.state.node?.response.isError ?? false;
+
+        let statusIcon = '';
+        let statusText = '';
+        if (hasNode) {
+            if (isComplete) {
+                statusIcon = 'codicon-check';
+                statusText = 'completed';
+            } else if (isCanceled) {
+                statusIcon = 'codicon-cancel';
+                statusText = 'canceled';
+            } else if (isError) {
+                statusIcon = 'codicon-error';
+                statusText = 'error';
+            } else {
+                statusIcon = 'codicon-loading';
+                statusText = 'generating...';
+            }
+        } else {
+            statusIcon = 'codicon-loading';
+            statusText = 'starting...';
+        }
+
+        return (
+            <div className="theia-delegation-container">
+                <details className="delegation-response-details">
+                    <summary className="delegation-summary">
+                        <div className="delegation-header">
+                            <span className="delegation-agent">
+                                <strong>Agent:</strong> {agentId}
+                            </span>
+                            <span className="delegation-status">
+                                <span className={`codicon ${statusIcon} delegation-status-icon`}></span>
+                                <span className="delegation-status-text">{statusText}</span>
+                            </span>
+                        </div>
+                    </summary>
+                    <div className="delegation-content">
+                        <div className="delegation-prompt-section">
+                            <strong>Delegated prompt:</strong>
+                            <div className="delegation-prompt">{prompt}</div>
+                        </div>
+                        <div className="delegation-response-section">
+                            <strong>Response:</strong>
+                            <div className='delegation-response-placeholder'>
+                                {hasNode && this.state.node ? this.widget.renderChatResponse(this.state.node) :
+                                    <div className="theia-ChatContentInProgress">Starting delegation...</div>
+                                }
+                            </div>
+                        </div>
+                    </div>
+                </details>
+            </div>
+        );
+    }
+}
+
+function mapResponseToNode(response: ChatResponseModel, parentNode: ResponseNode): ResponseNode {
+    return {
+        id: response.id,
+        parent: parentNode as unknown as CompositeTreeNode,
+        response,
+        sessionId: parentNode.sessionId
+    };
+}

package/src/browser/chat-response-renderer/index.ts

@@ -23,3 +23,5 @@ export * from './text-part-renderer';
 export * from './toolcall-part-renderer';
 export * from './thinking-part-renderer';
 export * from './progress-part-renderer';
+export * from './tool-confirmation';
+export * from './delegation-response-renderer';

package/src/browser/chat-response-renderer/tool-confirmation.tsx

@@ -0,0 +1,173 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclipseSource GmbH.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+
+import * as React from '@theia/core/shared/react';
+import { nls } from '@theia/core/lib/common/nls';
+import { codicon } from '@theia/core/lib/browser';
+import { ToolCallChatResponseContent } from '@theia/ai-chat/lib/common';
+
+/**
+ * States the tool confirmation component can be in
+ */
+export type ToolConfirmationState = 'waiting' | 'allowed' | 'denied';
+
+export interface ToolConfirmationProps {
+    response: ToolCallChatResponseContent;
+    onAllow: (mode?: 'once' | 'session' | 'forever') => void;
+    onDeny: (mode?: 'once' | 'session' | 'forever') => void;
+}
+
+/**
+ * Component that displays approval/denial buttons for tool execution
+ */
+export const ToolConfirmation: React.FC<ToolConfirmationProps> = ({ response, onAllow, onDeny }) => {
+    const [state, setState] = React.useState<ToolConfirmationState>('waiting');
+    // Track selected mode for each action
+    const [allowMode, setAllowMode] = React.useState<'once' | 'session' | 'forever'>('once');
+    const [denyMode, setDenyMode] = React.useState<'once' | 'session' | 'forever'>('once');
+    const [dropdownOpen, setDropdownOpen] = React.useState<'allow' | 'deny' | undefined>(undefined);
+
+    const handleAllow = React.useCallback(() => {
+        setState('allowed');
+        onAllow(allowMode);
+    }, [onAllow, allowMode]);
+
+    const handleDeny = React.useCallback(() => {
+        setState('denied');
+        onDeny(denyMode);
+    }, [onDeny, denyMode]);
+
+    if (state === 'allowed') {
+        return (
+            <div className="theia-tool-confirmation-status allowed">
+                <span className={codicon('check')}></span> {nls.localize('theia/ai/chat-ui/toolconfirmation/allowed', 'Tool execution allowed')}
+            </div>
+        );
+    }
+
+    if (state === 'denied') {
+        return (
+            <div className="theia-tool-confirmation-status denied">
+                <span className={codicon('close')}></span> {nls.localize('theia/ai/chat-ui/toolconfirmation/denied', 'Tool execution denied')}
+            </div>
+        );
+    }
+
+    // Helper for dropdown options
+    const MODES: Array<'once' | 'session' | 'forever'> = ['once', 'session', 'forever'];
+    // Unified labels for both main button and dropdown, as requested
+    const modeLabel = (type: 'allow' | 'deny', mode: 'once' | 'session' | 'forever') => {
+        if (type === 'allow') {
+            switch (mode) {
+                case 'once': return nls.localize('theia/ai/chat-ui/toolconfirmation/allow', 'Allow');
+                case 'session': return nls.localize('theia/ai/chat-ui/toolconfirmation/allow-session', 'Allow for this Chat');
+                case 'forever': return nls.localize('theia/ai/chat-ui/toolconfirmation/allow-forever', 'Always Allow');
+            }
+        } else {
+            switch (mode) {
+                case 'once': return nls.localize('theia/ai/chat-ui/toolconfirmation/deny', 'Deny');
+                case 'session': return nls.localize('theia/ai/chat-ui/toolconfirmation/deny-session', 'Deny for this Chat');
+                case 'forever': return nls.localize('theia/ai/chat-ui/toolconfirmation/deny-forever', 'Always Deny');
+            }
+        }
+    };
+    // Main button label is always the same as the dropdown label for the selected mode
+    const mainButtonLabel = modeLabel; // Use the same function for both
+
+    // Tooltips for dropdown options
+    const modeTooltip = (type: 'allow' | 'deny', mode: 'once' | 'session' | 'forever') => {
+        if (type === 'allow') {
+            switch (mode) {
+                case 'once': return nls.localize('theia/ai/chat-ui/toolconfirmation/allow-tooltip', 'Allow this tool call once');
+                case 'session': return nls.localize('theia/ai/chat-ui/toolconfirmation/allow-session-tooltip', 'Allow all calls of this tool for this chat session');
+                case 'forever': return nls.localize('theia/ai/chat-ui/toolconfirmation/allow-forever-tooltip', 'Always allow this tool');
+            }
+        } else {
+            switch (mode) {
+                case 'once': return nls.localize('theia/ai/chat-ui/toolconfirmation/deny-tooltip', 'Deny this tool call once');
+                case 'session': return nls.localize('theia/ai/chat-ui/toolconfirmation/deny-session-tooltip', 'Deny all calls of this tool for this chat session');
+                case 'forever': return nls.localize('theia/ai/chat-ui/toolconfirmation/deny-forever-tooltip', 'Always deny this tool');
+            }
+        }
+    };
+
+    // Split button for approve/deny
+    const renderSplitButton = (type: 'allow' | 'deny') => {
+        const selectedMode = type === 'allow' ? allowMode : denyMode;
+        const setMode = type === 'allow' ? setAllowMode : setDenyMode;
+        const handleMain = type === 'allow' ? handleAllow : handleDeny;
+        const otherModes = MODES.filter(m => m !== selectedMode);
+        return (
+            <div className={`theia-tool-confirmation-split-button ${type}`}
+                style={{ display: 'inline-flex', position: 'relative' }}>
+                <button
+                    className={`theia-button ${type === 'allow' ? 'primary' : 'secondary'} theia-tool-confirmation-main-btn`}
+                    onClick={handleMain}
+                >
+                    {mainButtonLabel(type, selectedMode)}
+                </button>
+                <button
+                    className={`theia-button ${type === 'allow' ? 'primary' : 'secondary'} theia-tool-confirmation-chevron-btn`}
+                    onClick={() => setDropdownOpen(dropdownOpen === type ? undefined : type)}
+                    aria-haspopup="true"
+                    aria-expanded={dropdownOpen === type}
+                    tabIndex={0}
+                    title={type === 'allow' ? 'More Allow Options' : 'More Deny Options'}
+                >
+                    <span className={codicon('chevron-down')}></span>
+                </button>
+                {dropdownOpen === type && (
+                    <ul
+                        className="theia-tool-confirmation-dropdown-menu"
+                        onMouseLeave={() => setDropdownOpen(undefined)}
+                    >
+                        {otherModes.map(mode => (
+                            <li
+                                key={mode}
+                                className="theia-tool-confirmation-dropdown-item"
+                                onClick={() => {
+                                    setMode(mode);
+                                    setDropdownOpen(undefined);
+                                }}
+                                title={modeTooltip(type, mode)}
+                            >
+                                {modeLabel(type, mode)}
+                            </li>
+                        ))}
+                    </ul>
+                )}
+            </div>
+        );
+    };
+
+    return (
+        <div className="theia-tool-confirmation">
+            <div className="theia-tool-confirmation-header">
+                <span className={codicon('shield')}></span> {nls.localize('theia/ai/chat-ui/toolconfirmation/header', 'Confirm Tool Execution')}
+            </div>
+            <div className="theia-tool-confirmation-info">
+                <div className="theia-tool-confirmation-name">
+                    <span className="label">{nls.localize('theia/ai/chat-ui/toolconfirmation/tool', 'Tool')}:</span>
+                    <span className="value">{response.name}</span>
+                </div>
+            </div>
+            <div className="theia-tool-confirmation-actions">
+                {renderSplitButton('deny')}
+                {renderSplitButton('allow')}
+            </div>
+        </div>
+    );
+};

package/src/browser/chat-response-renderer/toolcall-part-renderer.tsx

@@ -15,15 +15,22 @@
 // *****************************************************************************
 
 import { ChatResponsePartRenderer } from '../chat-response-part-renderer';
-import { injectable } from '@theia/core/shared/inversify';
+import { inject, injectable } from '@theia/core/shared/inversify';
 import { ChatResponseContent, ToolCallChatResponseContent } from '@theia/ai-chat/lib/common';
 import { ReactNode } from '@theia/core/shared/react';
 import { nls } from '@theia/core/lib/common/nls';
+import { codicon } from '@theia/core/lib/browser';
 import * as React from '@theia/core/shared/react';
+import { ToolConfirmation, ToolConfirmationState } from './tool-confirmation';
+import { ToolConfirmationManager, ToolConfirmationMode } from '@theia/ai-chat/lib/browser/chat-tool-preferences';
+import { ResponseNode } from '../chat-tree-view';
 
 @injectable()
 export class ToolCallPartRenderer implements ChatResponsePartRenderer<ToolCallChatResponseContent> {
 
+    @inject(ToolConfirmationManager)
+    protected toolConfirmationManager: ToolConfirmationManager;
+
     canHandle(response: ChatResponseContent): number {
         if (ToolCallChatResponseContent.is(response)) {
             return 10;
@@ -31,23 +38,20 @@ export class ToolCallPartRenderer implements ChatResponsePartRenderer<ToolCallCh
         return -1;
     }
 
-    render(response: ToolCallChatResponseContent): ReactNode {
-        return (
-            <h4 className='theia-toolCall'>
-                {response.finished ? (
-                    <details>
-                        <summary>{nls.localize('theia/ai/chat-ui/toolcall-part-renderer/finished', 'Ran')} {response.name}
-                            ({this.renderCollapsibleArguments(response.arguments)})
-                        </summary>
-                        <pre>{this.tryPrettyPrintJson(response)}</pre>
-                    </details>
-                ) : (
-                    <span>
-                        <Spinner /> {nls.localizeByDefault('Running')} {response.name}({this.renderCollapsibleArguments(response.arguments)})
-                    </span>
-                )}
-            </h4>
-        );
+    render(response: ToolCallChatResponseContent, parentNode: ResponseNode): ReactNode {
+        const chatId = parentNode.sessionId;
+        const confirmationMode = response.name ? this.getToolConfirmationSettings(response.name, chatId) : ToolConfirmationMode.DISABLED;
+        return <ToolCallContent
+            response={response}
+            confirmationMode={confirmationMode}
+            toolConfirmationManager={this.toolConfirmationManager}
+            chatId={chatId}
+            renderCollapsibleArguments={this.renderCollapsibleArguments.bind(this)}
+            tryPrettyPrintJson={this.tryPrettyPrintJson.bind(this)} />;
+    }
+
+    protected getToolConfirmationSettings(responseId: string, chatId: string): ToolConfirmationMode {
+        return this.toolConfirmationManager.getConfirmationMode(responseId, chatId);
     }
 
     protected renderCollapsibleArguments(args: string | undefined): ReactNode {
@@ -97,5 +101,97 @@ export class ToolCallPartRenderer implements ChatResponsePartRenderer<ToolCallCh
 }
 
 const Spinner = () => (
-    <i className="fa fa-spinner fa-spin"></i>
+    <span className={`${codicon('loading')} theia-animation-spin`}></span>
 );
+
+interface ToolCallContentProps {
+    response: ToolCallChatResponseContent;
+    confirmationMode: ToolConfirmationMode;
+    toolConfirmationManager: ToolConfirmationManager;
+    chatId: string;
+    renderCollapsibleArguments: (args: string | undefined) => ReactNode;
+    tryPrettyPrintJson: (response: ToolCallChatResponseContent) => string | undefined;
+}
+
+/**
+ * A function component to handle tool call rendering and confirmation
+ */
+const ToolCallContent: React.FC<ToolCallContentProps> = ({ response, confirmationMode, toolConfirmationManager, chatId, tryPrettyPrintJson, renderCollapsibleArguments }) => {
+    const [confirmationState, setConfirmationState] = React.useState<ToolConfirmationState>('waiting');
+
+    React.useEffect(() => {
+        if (confirmationMode === ToolConfirmationMode.ALWAYS_ALLOW) {
+            response.confirm();
+            setConfirmationState('allowed');
+            return;
+        } else if (confirmationMode === ToolConfirmationMode.DISABLED) {
+            response.deny();
+            setConfirmationState('denied');
+            return;
+        }
+        response.confirmed.then(
+            confirmed => {
+                if (confirmed === true) {
+                    setConfirmationState('allowed');
+                } else {
+                    setConfirmationState('denied');
+                }
+            }
+        )
+            .catch(() => {
+                setConfirmationState('denied');
+            });
+    }, [response, confirmationMode]);
+
+    const handleAllow = React.useCallback((mode: 'once' | 'session' | 'forever' = 'once') => {
+        if (mode === 'forever' && response.name) {
+            toolConfirmationManager.setConfirmationMode(response.name, ToolConfirmationMode.ALWAYS_ALLOW);
+        } else if (mode === 'session' && response.name) {
+            toolConfirmationManager.setSessionConfirmationMode(response.name, ToolConfirmationMode.ALWAYS_ALLOW, chatId);
+        }
+        response.confirm();
+    }, [response, toolConfirmationManager, chatId]);
+
+    const handleDeny = React.useCallback((mode: 'once' | 'session' | 'forever' = 'once') => {
+        if (mode === 'forever' && response.name) {
+            toolConfirmationManager.setConfirmationMode(response.name, ToolConfirmationMode.DISABLED);
+        } else if (mode === 'session' && response.name) {
+            toolConfirmationManager.setSessionConfirmationMode(response.name, ToolConfirmationMode.DISABLED, chatId);
+        }
+        response.deny();
+    }, [response, toolConfirmationManager, chatId]);
+
+    return (
+        <div className='theia-toolCall'>
+            <h4>
+                {confirmationState === 'denied' ? (
+                    <span className="theia-tool-denied">
+                        <span className={codicon('error')}></span> {nls.localize('theia/ai/chat-ui/toolcall-part-renderer/denied', 'Execution denied')}: {response.name}
+                    </span>
+                ) : response.finished ? (
+                    <details>
+                        <summary>{nls.localize('theia/ai/chat-ui/toolcall-part-renderer/finished', 'Ran')} {response.name}
+                            ({renderCollapsibleArguments(response.arguments)})
+                        </summary>
+                        <pre>{tryPrettyPrintJson(response)}</pre>
+                    </details>
+                ) : (
+                    confirmationState === 'allowed' && (
+                        <span>
+                            <Spinner /> {nls.localizeByDefault('Running')} {response.name}
+                        </span>
+                    )
+                )}
+            </h4>
+
+            {/* Show confirmation UI when waiting for allow */}
+            {confirmationState === 'waiting' && (
+                <ToolConfirmation
+                    response={response}
+                    onAllow={handleAllow}
+                    onDeny={handleDeny}
+                />
+            )}
+        </div>
+    );
+};

package/src/browser/chat-tree-view/chat-view-tree-input-widget.tsx

@@ -19,7 +19,7 @@ import { AIChatInputWidget, type AIChatInputConfiguration } from '../chat-input-
 import type { EditableRequestNode } from './chat-view-tree-widget';
 import { URI } from '@theia/core';
 import { CHAT_VIEW_LANGUAGE_EXTENSION } from '../chat-view-language-contribution';
-import type { ChatRequestModel, EditableChatRequestModel } from '@theia/ai-chat';
+import type { ChatRequestModel, EditableChatRequestModel, ChatHierarchyBranch } from '@theia/ai-chat';
 import type { AIVariableResolutionRequest } from '@theia/ai-core';
 import { Key } from '@theia/core/lib/browser';
 
@@ -29,6 +29,10 @@ export interface AIChatTreeInputConfiguration extends AIChatInputConfiguration {
 export const AIChatTreeInputArgs = Symbol('AIChatTreeInputArgs');
 export interface AIChatTreeInputArgs {
     node: EditableRequestNode;
+    /**
+     * The branch of the chat tree for this request node (used by the input widget for state tracking).
+     */
+    branch?: ChatHierarchyBranch;
     initialValue?: string;
     onQuery: (query: string) => Promise<void>;
     onUnpin?: () => void;
@@ -60,6 +64,13 @@ export class AIChatTreeInputWidget extends AIChatInputWidget {
     @postConstruct()
     protected override init(): void {
         super.init();
+        this.updateBranch();
+
+        const request = this.requestNode.request;
+        this.toDispose.push(request.session.onDidChange(() => {
+            this.updateBranch();
+        }));
+
         this.addKeyListener(this.node, Key.ESCAPE, () => {
             this.request.cancelEdit();
         });
@@ -71,6 +82,10 @@ export class AIChatTreeInputWidget extends AIChatInputWidget {
         });
     }
 
+    protected updateBranch(): void {
+        this.branch = this.args.branch ?? this.requestNode.branch;
+    }
+
     protected override getResourceUri(): URI {
         return new URI(`ai-chat:/${this.requestNode.id}-input.${CHAT_VIEW_LANGUAGE_EXTENSION}`);
     }