@theia/ai-chat 1.71.0-next.8 → 1.71.0

package/src/browser/agent-delegation-tool.ts

@@ -15,18 +15,20 @@
 // *****************************************************************************
 
 import { AGENT_DELEGATION_FUNCTION_ID, ToolInvocationContext, ToolProvider, ToolRequest } from '@theia/ai-core';
+import { Disposable } from '@theia/core';
 import { inject, injectable } from '@theia/core/shared/inversify';
 import {
     assertChatContext,
     ChatAgentService,
     ChatAgentServiceFactory,
+    ChatChangeEvent,
     ChatRequest,
     ChatService,
     ChatServiceFactory,
     ChatToolContext,
     MutableChatModel,
     MutableChatRequestModel,
-    ChatSession,
+    MutableChatResponseModel,
     ChatRequestInvocation,
 } from '../common';
 import { TASK_CONTEXT_VARIABLE } from './task-context-variable';
@@ -114,6 +116,7 @@ export class AgentDelegationTool implements ToolProvider {
             }
 
             let newSession;
+            let childModelDisposable: Disposable | undefined;
             try {
                 // FIXME: this creates a new conversation visible in the UI (Panel), which we don't want
                 // It is not possible to start a session without specifying a location (default=Panel)
@@ -146,8 +149,8 @@ export class AgentDelegationTool implements ToolProvider {
                     chatService.setActiveSession(currentActiveSession.id, { focus: false });
                 }
 
-                // Setup ChangeSet bubbling from delegated session to parent session
-                this.setupChangeSetBubbling(newSession, ctx.request.session);
+                // Setup bubbling of child session events to parent session
+                childModelDisposable = this.setupChildSessionBubbling(newSession.model as MutableChatModel, ctx.request.session, ctx.response);
             } catch (sessionError) {
                 const errorMsg = `Failed to create chat session for agent '${agentId}': ${sessionError instanceof Error ? sessionError.message : sessionError}`;
                 console.error(errorMsg, sessionError);
@@ -209,7 +212,8 @@ export class AgentDelegationTool implements ToolProvider {
                     const result = await response.responseCompleted;
                     const stringResult = result.response.asString();
 
-                    // Clean up the session after completion (no need to await)
+                    // Clean up the session and parent-child link after completion
+                    childModelDisposable?.dispose();
                     const chatService = this.getChatService();
                     chatService.deleteSession(newSession.id).catch(error => {
                         console.error('Failed to delete delegated session', error);
@@ -242,33 +246,40 @@ export class AgentDelegationTool implements ToolProvider {
     }
 
     /**
-     * Sets up monitoring of the ChangeSet in the delegated session and bubbles changes to the parent session.
-     * @param delegatedSession The session created for the delegated agent
-     * @param parentModel The parent session model that should receive the bubbled changes
+     * Sets up all event bubbling from a delegated child session to the parent session:
+     * - Interaction forwarding: child interactionNeeded events are forwarded to the parent response
+     * - ChangeSet bubbling: child changeset changes are forwarded to the parent model
      */
-    private setupChangeSetBubbling(
-        delegatedSession: ChatSession,
-        parentModel: MutableChatModel
-    ): void {
-        // Monitor ChangeSet for bubbling
-        delegatedSession.model.changeSet.onDidChange(_event => {
-            this.bubbleChangeSet(delegatedSession, parentModel);
+    private setupChildSessionBubbling(
+        childModel: MutableChatModel,
+        parentModel: MutableChatModel,
+        parentResponse: MutableChatResponseModel
+    ): Disposable {
+
+        // Forward interactionNeeded events to the parent response model
+        // so the UI (which subscribes to response.onInteractionNeeded) can display them.
+        // Also watch for each forwarded interaction's resolution to trigger cleanup.
+        const eventForwarding = childModel.onDidChange(event => {
+            if (ChatChangeEvent.isInteractionNeededEvent(event)) {
+                parentResponse.fireInteractionNeeded(event.contentPart);
+                event.contentPart.whenResolved.then(() => parentResponse.notifyChanged());
+            }
         });
-    }
 
-    /**
-     * Bubbles the ChangeSet from the delegated session to the parent session.
-     * @param delegatedSession The session from which to bubble changes
-     * @param parentModel The parent session model to receive the bubbled changes
-     */
-    private bubbleChangeSet(
-        delegatedSession: ChatSession,
-        parentModel: MutableChatModel
-    ): void {
-        const delegatedElements = delegatedSession.model.changeSet.getElements();
-        if (delegatedElements.length > 0) {
-            parentModel.changeSet.setTitle(delegatedSession.model.changeSet.title);
-            parentModel.changeSet.addElements(...delegatedElements);
-        }
+        // Bubble ChangeSet changes to the parent model
+        const changeSetForwarding = childModel.changeSet.onDidChange(() => {
+            const delegatedElements = childModel.changeSet.getElements();
+            if (delegatedElements.length > 0) {
+                parentModel.changeSet.setTitle(childModel.changeSet.title);
+                parentModel.changeSet.addElements(...delegatedElements);
+            }
+        });
+
+        return {
+            dispose: () => {
+                eventForwarding.dispose();
+                changeSetForwarding.dispose();
+            }
+        };
     }
 }

package/src/browser/chat-tool-preference-bindings.spec.ts

@@ -16,15 +16,24 @@
 
 import { expect } from 'chai';
 import * as sinon from 'sinon';
+import { Container } from '@theia/core/shared/inversify';
 import { ToolConfirmationManager } from './chat-tool-preference-bindings';
-import { ToolConfirmationMode, TOOL_CONFIRMATION_PREFERENCE, ChatToolPreferences } from '../common/chat-tool-preferences';
+import { ToolConfirmationMode } from '../common/chat-tool-preferences';
 import { ToolRequest } from '@theia/ai-core';
 import { PreferenceService } from '@theia/core/lib/common/preferences';
+import { TrustAwarePreferenceReader } from '@theia/ai-core/lib/browser/trust-aware-preference-reader';
 
 describe('ToolConfirmationManager', () => {
     let manager: ToolConfirmationManager;
     let preferenceServiceMock: sinon.SinonStubbedInstance<PreferenceService>;
+    let trustAwareReaderMock: sinon.SinonStubbedInstance<TrustAwarePreferenceReader>;
     let storedPreferences: { [toolId: string]: ToolConfirmationMode };
+    let trusted: boolean;
+    let inspectResult: {
+        defaultValue?: { [toolId: string]: ToolConfirmationMode };
+        globalValue?: { [toolId: string]: ToolConfirmationMode };
+        workspaceValue?: { [toolId: string]: ToolConfirmationMode };
+    } | undefined;
 
     const createToolRequest = (id: string, confirmAlwaysAllow?: boolean | string): ToolRequest => ({
         id,
@@ -34,26 +43,35 @@ describe('ToolConfirmationManager', () => {
         confirmAlwaysAllow
     });
 
-    const getPreferencesMock = (): ChatToolPreferences => ({
-        get [TOOL_CONFIRMATION_PREFERENCE](): { [toolId: string]: ToolConfirmationMode } {
-            return storedPreferences;
-        }
-    }) as unknown as ChatToolPreferences;
-
     beforeEach(() => {
         storedPreferences = {};
+        trusted = true;
+        inspectResult = undefined;
 
         preferenceServiceMock = {
             updateValue: sinon.stub().callsFake((_key: string, value: { [toolId: string]: ToolConfirmationMode }) => {
                 storedPreferences = value;
                 return Promise.resolve();
             }),
-            inspect: sinon.stub().returns(undefined)
+            inspect: sinon.stub().callsFake(() => inspectResult)
         } as unknown as sinon.SinonStubbedInstance<PreferenceService>;
 
-        manager = new ToolConfirmationManager();
-        (manager as unknown as { preferences: ChatToolPreferences }).preferences = getPreferencesMock();
-        (manager as unknown as { preferenceService: PreferenceService }).preferenceService = preferenceServiceMock;
+        trustAwareReaderMock = {
+            get: sinon.stub().callsFake(<T>(_name: string, fallback?: T): T | undefined => {
+                if (trusted) {
+                    return (storedPreferences as unknown as T) ?? fallback;
+                }
+                const insp = inspectResult;
+                const value = insp?.globalValue ?? insp?.defaultValue;
+                return ((value as unknown as T) ?? fallback);
+            })
+        } as unknown as sinon.SinonStubbedInstance<TrustAwarePreferenceReader>;
+
+        const container = new Container();
+        container.bind(ToolConfirmationManager).toSelf().inSingletonScope();
+        container.bind(PreferenceService).toConstantValue(preferenceServiceMock as unknown as PreferenceService);
+        container.bind(TrustAwarePreferenceReader).toConstantValue(trustAwareReaderMock as unknown as TrustAwarePreferenceReader);
+        manager = container.get(ToolConfirmationManager);
     });
 
     describe('getConfirmationMode', () => {
@@ -95,6 +113,48 @@ describe('ToolConfirmationManager', () => {
         });
     });
 
+    describe('workspace trust', () => {
+        it('ignores workspace {"*": "always_allow"} when workspace is untrusted', () => {
+            trusted = false;
+            // Simulate the effective (workspace-merged) preference containing always_allow,
+            // while the user/global scope is not set and the schema default prescribes
+            // CONFIRM. If the workspace override were honoured the mode would be
+            // ALWAYS_ALLOW; because trust filters the workspace scope out, the default
+            // (CONFIRM) wins, which demonstrably proves the override was dropped.
+            storedPreferences['*'] = ToolConfirmationMode.ALWAYS_ALLOW;
+            inspectResult = {
+                defaultValue: { '*': ToolConfirmationMode.CONFIRM },
+                workspaceValue: { '*': ToolConfirmationMode.ALWAYS_ALLOW }
+            };
+
+            const mode = manager.getConfirmationMode('someTool', 'chat-1');
+            expect(mode).to.equal(ToolConfirmationMode.CONFIRM);
+        });
+
+        it('blocks confirmAlwaysAllow bypass via workspace {"*": "always_allow"} when untrusted', () => {
+            trusted = false;
+            storedPreferences['*'] = ToolConfirmationMode.ALWAYS_ALLOW;
+            inspectResult = {
+                defaultValue: {},
+                workspaceValue: { '*': ToolConfirmationMode.ALWAYS_ALLOW }
+            };
+
+            const toolRequest = createToolRequest('dangerousTool', true);
+            const mode = manager.getConfirmationMode('dangerousTool', 'chat-1', toolRequest);
+            // Without the workspace value the map is empty, so the default for a
+            // confirmAlwaysAllow tool (CONFIRM) is used.
+            expect(mode).to.equal(ToolConfirmationMode.CONFIRM);
+        });
+
+        it('honours workspace {"*": "always_allow"} when workspace is trusted', () => {
+            trusted = true;
+            storedPreferences['*'] = ToolConfirmationMode.ALWAYS_ALLOW;
+
+            const mode = manager.getConfirmationMode('regularTool', 'chat-1');
+            expect(mode).to.equal(ToolConfirmationMode.ALWAYS_ALLOW);
+        });
+    });
+
     describe('setConfirmationMode', () => {
         it('should persist ALWAYS_ALLOW for regular tools when different from default', () => {
             storedPreferences['*'] = ToolConfirmationMode.CONFIRM;
@@ -116,34 +176,34 @@ describe('ToolConfirmationManager', () => {
         });
 
         it('should not persist when mode matches the global preference default', () => {
-            (preferenceServiceMock.inspect as sinon.SinonStub).returns({
+            inspectResult = {
                 defaultValue: { '*': ToolConfirmationMode.CONFIRM }
-            });
+            };
             manager.setConfirmationMode('regularTool', ToolConfirmationMode.CONFIRM);
             expect(preferenceServiceMock.updateValue.called).to.be.false;
         });
 
         it('should persist when mode differs from the global preference default', () => {
-            (preferenceServiceMock.inspect as sinon.SinonStub).returns({
+            inspectResult = {
                 defaultValue: { '*': ToolConfirmationMode.CONFIRM }
-            });
+            };
             manager.setConfirmationMode('regularTool', ToolConfirmationMode.DISABLED);
             expect(preferenceServiceMock.updateValue.calledOnce).to.be.true;
             expect(storedPreferences['regularTool']).to.equal(ToolConfirmationMode.DISABLED);
         });
 
         it('should not persist when mode matches the tool-specific preference default', () => {
-            (preferenceServiceMock.inspect as sinon.SinonStub).returns({
+            inspectResult = {
                 defaultValue: { 'myTool': ToolConfirmationMode.DISABLED }
-            });
+            };
             manager.setConfirmationMode('myTool', ToolConfirmationMode.DISABLED);
             expect(preferenceServiceMock.updateValue.called).to.be.false;
         });
 
         it('should remove entry when mode matches the tool-specific preference default and entry exists', () => {
-            (preferenceServiceMock.inspect as sinon.SinonStub).returns({
+            inspectResult = {
                 defaultValue: { 'myTool': ToolConfirmationMode.DISABLED }
-            });
+            };
             storedPreferences['myTool'] = ToolConfirmationMode.ALWAYS_ALLOW;
             manager.setConfirmationMode('myTool', ToolConfirmationMode.DISABLED);
             expect(preferenceServiceMock.updateValue.calledOnce).to.be.true;
@@ -151,18 +211,18 @@ describe('ToolConfirmationManager', () => {
         });
 
         it('should not persist CONFIRM for confirmAlwaysAllow tool when global preference default is CONFIRM', () => {
-            (preferenceServiceMock.inspect as sinon.SinonStub).returns({
+            inspectResult = {
                 defaultValue: { '*': ToolConfirmationMode.CONFIRM }
-            });
+            };
             const toolRequest = createToolRequest('dangerousTool', true);
             manager.setConfirmationMode('dangerousTool', ToolConfirmationMode.CONFIRM, toolRequest);
             expect(preferenceServiceMock.updateValue.called).to.be.false;
         });
 
         it('should persist ALWAYS_ALLOW for confirmAlwaysAllow tool when global preference default is CONFIRM', () => {
-            (preferenceServiceMock.inspect as sinon.SinonStub).returns({
+            inspectResult = {
                 defaultValue: { '*': ToolConfirmationMode.CONFIRM }
-            });
+            };
             const toolRequest = createToolRequest('dangerousTool', true);
             manager.setConfirmationMode('dangerousTool', ToolConfirmationMode.ALWAYS_ALLOW, toolRequest);
             expect(preferenceServiceMock.updateValue.calledOnce).to.be.true;

package/src/browser/chat-tool-preference-bindings.ts

@@ -18,20 +18,21 @@ import { injectable, inject } from '@theia/core/shared/inversify';
 import {
     PreferenceService,
 } from '@theia/core/lib/common/preferences';
-import { ToolConfirmationMode, TOOL_CONFIRMATION_PREFERENCE, ChatToolPreferences } from '../common/chat-tool-preferences';
+import { ToolConfirmationMode, TOOL_CONFIRMATION_PREFERENCE } from '../common/chat-tool-preferences';
 import { ToolRequest } from '@theia/ai-core';
+import { TrustAwarePreferenceReader } from '@theia/ai-core/lib/browser/trust-aware-preference-reader';
 
 /**
  * Utility class to manage tool confirmation settings
  */
 @injectable()
 export class ToolConfirmationManager {
-    @inject(ChatToolPreferences)
-    protected readonly preferences: ChatToolPreferences;
-
     @inject(PreferenceService)
     protected readonly preferenceService: PreferenceService;
 
+    @inject(TrustAwarePreferenceReader)
+    protected readonly trustAwareReader: TrustAwarePreferenceReader;
+
     // In-memory session overrides (not persisted), per chat
     protected sessionOverrides: Map<string, Map<string, ToolConfirmationMode>> = new Map();
 
@@ -51,7 +52,9 @@ export class ToolConfirmationManager {
         if (chatMap && chatMap.has(toolId)) {
             return chatMap.get(toolId)!;
         }
-        const toolConfirmation = this.preferences[TOOL_CONFIRMATION_PREFERENCE];
+        const toolConfirmation = this.trustAwareReader.get<Record<string, ToolConfirmationMode>>(
+            TOOL_CONFIRMATION_PREFERENCE, {}
+        ) ?? {};
         if (toolConfirmation[toolId]) {
             return toolConfirmation[toolId];
         }
@@ -80,7 +83,9 @@ export class ToolConfirmationManager {
         const defaultPref = this.preferenceService.inspect(TOOL_CONFIRMATION_PREFERENCE)?.defaultValue as {
             [toolId: string]: ToolConfirmationMode;
         } || {};
-        const current = this.preferences[TOOL_CONFIRMATION_PREFERENCE] || {};
+        const current = this.trustAwareReader.get<Record<string, ToolConfirmationMode>>(
+            TOOL_CONFIRMATION_PREFERENCE, {}
+        ) ?? {};
         let starMode = current['*'];
         if (starMode === undefined) {
             starMode = defaultPref['*'] ?? ToolConfirmationMode.ALWAYS_ALLOW;
@@ -127,11 +132,15 @@ export class ToolConfirmationManager {
      * Get all tool confirmation settings
      */
     getAllConfirmationSettings(): { [toolId: string]: ToolConfirmationMode } {
-        return this.preferences[TOOL_CONFIRMATION_PREFERENCE] || {};
+        return this.trustAwareReader.get<Record<string, ToolConfirmationMode>>(
+            TOOL_CONFIRMATION_PREFERENCE, {}
+        ) ?? {};
     }
 
     resetAllConfirmationModeSettings(): void {
-        const current = this.preferences[TOOL_CONFIRMATION_PREFERENCE] || {};
+        const current = this.trustAwareReader.get<Record<string, ToolConfirmationMode>>(
+            TOOL_CONFIRMATION_PREFERENCE, {}
+        ) ?? {};
         if ('*' in current) {
             this.preferenceService.updateValue(TOOL_CONFIRMATION_PREFERENCE, { '*': current['*'] });
         } else {

package/src/browser/chat-tool-request-service.ts

@@ -39,12 +39,12 @@ export class FrontendChatToolRequestService extends ChatToolRequestService {
     protected readonly preferences: ChatToolPreferences;
 
     protected override toChatToolRequest(toolRequest: ToolRequest, request: MutableChatRequestModel): ToolRequest {
-        const confirmationMode = this.confirmationManager.getConfirmationMode(toolRequest.id, request.session.id, toolRequest);
-
         return {
             ...toolRequest,
             handler: async (arg_string: string, ctx?: ToolInvocationContext) => {
                 const toolCallId = ctx?.toolCallId;
+                const sessionId = request.session.rootSessionId ?? request.session.id;
+                const confirmationMode = this.confirmationManager.getConfirmationMode(toolRequest.id, sessionId, toolRequest);
 
                 switch (confirmationMode) {
                     case ToolConfirmationMode.DISABLED:
@@ -84,6 +84,7 @@ export class FrontendChatToolRequestService extends ChatToolRequestService {
                                 : this.preferences[TOOL_CONFIRMATION_TIMEOUT_PREFERENCE];
                             toolCallContent.confirmationTimeout = timeoutSeconds;
                             toolCallContent.requestUserConfirmation();
+                            request.response.fireInteractionNeeded(toolCallContent);
                             confirmed = await raceConfirmationWithTimeout(toolCallContent, timeoutSeconds);
                         }
 

package/src/common/chat-agents.spec.ts

@@ -0,0 +1,124 @@
+// *****************************************************************************
+// Copyright (C) 2026 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 'reflect-metadata';
+
+import { expect } from 'chai';
+import { LanguageModelMessage, LanguageModelRequirement } from '@theia/ai-core';
+import { AbstractChatAgent, ChatAgentLocation } from './chat-agents';
+import {
+    MutableChatModel,
+    MutableChatRequestModel,
+    ChatModel,
+    TextChatResponseContentImpl,
+    ThinkingChatResponseContentImpl,
+} from './chat-model';
+import { ParsedChatRequest, ParsedChatRequestTextPart } from './parsed-chat-request';
+
+class TestChatAgent extends AbstractChatAgent {
+    readonly id = 'test-agent';
+    readonly name = 'Test Agent';
+    readonly languageModelRequirements: LanguageModelRequirement[] = [];
+    protected readonly defaultLanguageModelPurpose = 'chat';
+
+    protected addContentsToResponse(): Promise<void> {
+        return Promise.resolve();
+    }
+
+    public async exposeGetMessages(model: ChatModel, includeResponseInProgress = false): Promise<LanguageModelMessage[]> {
+        return this.getMessages(model, includeResponseInProgress);
+    }
+}
+
+function createParsedRequest(text: string): ParsedChatRequest {
+    return {
+        request: { text },
+        parts: [
+            new ParsedChatRequestTextPart({ start: 0, endExclusive: text.length }, text)
+        ],
+        toolRequests: new Map(),
+        variables: []
+    };
+}
+
+describe('AbstractChatAgent.getMessages', () => {
+
+    let agent: TestChatAgent;
+
+    beforeEach(() => {
+        agent = new TestChatAgent();
+    });
+
+    function addThinkingResponse(request: MutableChatRequestModel, content: string, signature: string): void {
+        request.response.response.addContent(new ThinkingChatResponseContentImpl(content, signature));
+    }
+
+    function addTextResponse(request: MutableChatRequestModel, text: string): void {
+        request.response.response.addContent(new TextChatResponseContentImpl(text));
+    }
+
+    it('filters out incomplete thinking blocks (empty signature) from a cancelled stream', async () => {
+        const model = new MutableChatModel(ChatAgentLocation.Panel);
+        const request = model.addRequest(createParsedRequest('Hello'));
+
+        addThinkingResponse(request, 'Some thinking that was cancelled', '');
+        request.response.cancel();
+
+        const messages = await agent.exposeGetMessages(model);
+
+        expect(messages.filter(LanguageModelMessage.isThinkingMessage)).to.have.lengthOf(0);
+        // The user text message should still be included
+        const userTextMessages = messages
+            .filter(LanguageModelMessage.isTextMessage)
+            .filter(m => m.actor === 'user');
+        expect(userTextMessages).to.have.lengthOf(1);
+    });
+
+    it('keeps thinking blocks with a valid signature', async () => {
+        const model = new MutableChatModel(ChatAgentLocation.Panel);
+        const request = model.addRequest(createParsedRequest('Hello'));
+
+        addThinkingResponse(request, 'Complete thought', 'sig-abc');
+        addTextResponse(request, 'Hi there');
+        request.response.complete();
+
+        const messages = await agent.exposeGetMessages(model);
+
+        const thinkingMessages = messages.filter(LanguageModelMessage.isThinkingMessage);
+        expect(thinkingMessages).to.have.lengthOf(1);
+        expect(thinkingMessages[0].thinking).to.equal('Complete thought');
+        expect(thinkingMessages[0].signature).to.equal('sig-abc');
+    });
+
+    it('filters incomplete thinking but preserves following text content from the same response', async () => {
+        const model = new MutableChatModel(ChatAgentLocation.Panel);
+        const request = model.addRequest(createParsedRequest('Hello'));
+
+        addThinkingResponse(request, 'Cancelled thought', '');
+        addTextResponse(request, 'Partial reply before cancel');
+        request.response.cancel();
+
+        const messages = await agent.exposeGetMessages(model);
+
+        expect(messages.filter(LanguageModelMessage.isThinkingMessage)).to.have.lengthOf(0);
+
+        const aiTextMessages = messages
+            .filter(LanguageModelMessage.isTextMessage)
+            .filter(m => m.actor === 'ai');
+        expect(aiTextMessages).to.have.lengthOf(1);
+        expect(aiTextMessages[0].text).to.equal('Partial reply before cancel');
+    });
+});

package/src/common/chat-agents.ts

@@ -33,6 +33,8 @@ import {
     isToolCallResponsePart,
     isUsageResponsePart,
     LanguageModel,
+    TokenUsageService,
+    UsageResponsePart,
     LanguageModelMessage,
     LanguageModelRequirement,
     LanguageModelResponse,
@@ -47,13 +49,14 @@ import {
 } from '@theia/ai-core';
 import {
     Agent,
+    isLanguageModelParsedResponse,
     isLanguageModelStreamResponse,
     isLanguageModelTextResponse,
     LanguageModelRegistry,
     LanguageModelStreamResponsePart
 } from '@theia/ai-core/lib/common';
 import { ContributionProvider, ILogger, isArray, nls } from '@theia/core';
-import { inject, injectable, named, postConstruct } from '@theia/core/shared/inversify';
+import { inject, injectable, named, optional, postConstruct } from '@theia/core/shared/inversify';
 import { ChatAgentService } from './chat-agent-service';
 import {
     ChatModel,
@@ -68,6 +71,8 @@ import {
     ToolCallArgumentsDeltaContent,
     ErrorChatResponseContent,
     InformationalChatResponseContent,
+    ResponseTokenUsage,
+    ThinkingChatResponseContent,
 } from './chat-model';
 import { ChatToolRequestService } from './chat-tool-request-service';
 import { parseContents } from './parse-contents';
@@ -179,6 +184,8 @@ export abstract class AbstractChatAgent implements ChatAgent {
     @inject(DefaultResponseContentFactory)
     protected defaultContentFactory: DefaultResponseContentFactory;
 
+    @inject(TokenUsageService) @optional() protected tokenUsageService: TokenUsageService | undefined;
+
     readonly abstract id: string;
     readonly abstract name: string;
     readonly abstract languageModelRequirements: LanguageModelRequirement[];
@@ -259,6 +266,7 @@ export abstract class AbstractChatAgent implements ChatAgent {
             );
 
             await this.addContentsToResponse(languageModelResponse, request);
+            await this.recordTokenUsageFromResponse(request, languageModel);
             await this.onResponseComplete(request);
 
         } catch (e) {
@@ -418,6 +426,12 @@ export abstract class AbstractChatAgent implements ChatAgent {
                         if (ErrorChatResponseContent.is(c) || InformationalChatResponseContent.is(c)) {
                             return false;
                         }
+                        // skip incomplete thinking blocks (e.g. from a cancelled stream where the
+                        // signature_delta never arrived). Some LLMs (e.g. Anthropic) reject thinking
+                        // blocks without a signature.
+                        if (ThinkingChatResponseContent.is(c) && !c.signature) {
+                            return false;
+                        }
                         // content even has an own converter, definitely include it
                         if (ChatResponseContent.hasToLanguageModelMessage(c)) {
                             return true;
@@ -498,7 +512,7 @@ export abstract class AbstractChatAgent implements ChatAgent {
                 messages,
                 tools,
                 settings,
-                thinkingMode: commonSettings?.thinkingMode,
+                reasoning: commonSettings?.reasoning,
                 agentId: this.id,
                 sessionId: request.session.id,
                 requestId: request.id,
@@ -526,6 +540,34 @@ export abstract class AbstractChatAgent implements ChatAgent {
         return request.response.complete();
     }
 
+    protected mapUsageResponsePart(usage: UsageResponsePart): ResponseTokenUsage {
+        return {
+            inputTokens: usage.input_tokens,
+            outputTokens: usage.output_tokens,
+            cacheCreationInputTokens: usage.cache_creation_input_tokens,
+            cacheReadInputTokens: usage.cache_read_input_tokens,
+        };
+    }
+
+    protected async recordTokenUsageFromResponse(request: MutableChatRequestModel, languageModel: LanguageModel): Promise<void> {
+        if (!this.tokenUsageService) {
+            return;
+        }
+        const entries = request.response.tokenUsageEntries;
+        if (entries.length === 0) {
+            return;
+        }
+        for (const entry of entries) {
+            await this.tokenUsageService.recordTokenUsage(languageModel.id, {
+                inputTokens: entry.inputTokens,
+                outputTokens: entry.outputTokens,
+                cachedInputTokens: entry.cacheCreationInputTokens,
+                readCachedInputTokens: entry.cacheReadInputTokens,
+                requestId: request.id,
+            });
+        }
+    }
+
     protected abstract addContentsToResponse(languageModelResponse: LanguageModelResponse, request: MutableChatRequestModel): Promise<void>;
 }
 
@@ -537,6 +579,9 @@ export abstract class AbstractTextToModelParsingChatAgent<T> extends AbstractCha
         const parsedCommand = await this.parseTextResponse(responseAsText);
         const content = this.createResponseContent(parsedCommand, request);
         request.response.response.addContent(content);
+        if ('usage' in languageModelResponse && languageModelResponse.usage) {
+            request.response.setTokenUsage(this.mapUsageResponsePart(languageModelResponse.usage));
+        }
     }
 
     protected abstract parseTextResponse(text: string): Promise<T>;
@@ -581,6 +626,17 @@ export abstract class AbstractStreamParsingChatAgent extends AbstractChatAgent {
         if (isLanguageModelTextResponse(languageModelResponse)) {
             const contents = this.parseContents(languageModelResponse.text, request);
             request.response.response.addContents(contents);
+            if (languageModelResponse.usage) {
+                request.response.setTokenUsage(this.mapUsageResponsePart(languageModelResponse.usage));
+            }
+            return;
+        }
+        if (isLanguageModelParsedResponse(languageModelResponse)) {
+            const contents = this.parseContents(languageModelResponse.content, request);
+            request.response.response.addContents(contents);
+            if (languageModelResponse.usage) {
+                request.response.setTokenUsage(this.mapUsageResponsePart(languageModelResponse.usage));
+            }
             return;
         }
         if (isLanguageModelStreamResponse(languageModelResponse)) {
@@ -655,6 +711,7 @@ export abstract class AbstractStreamParsingChatAgent extends AbstractChatAgent {
             return new ThinkingChatResponseContentImpl(token.thought, token.signature);
         }
         if (isUsageResponsePart(token)) {
+            request.response.setTokenUsage(this.mapUsageResponsePart(token));
             return [];
         }
         return this.defaultContentFactory.create('', request);