@theia/ai-core 1.60.2 → 1.61.0

package/src/common/ai-variable-resource.ts

@@ -0,0 +1,86 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclipseSource GmbH and others.
+//
+// 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 deepEqual from 'fast-deep-equal';
+import { inject, injectable, postConstruct } from '@theia/core/shared/inversify';
+import { Resource, URI, generateUuid } from '@theia/core';
+import { AIVariableContext, AIVariableResolutionRequest } from './variable-service';
+import stableJsonStringify = require('fast-json-stable-stringify');
+import { ConfigurableInMemoryResources, ConfigurableMutableReferenceResource } from './configurable-in-memory-resources';
+
+export const AI_VARIABLE_RESOURCE_SCHEME = 'ai-variable';
+export const NO_CONTEXT_AUTHORITY = 'context-free';
+
+@injectable()
+export class AIVariableResourceResolver {
+    @inject(ConfigurableInMemoryResources) protected readonly inMemoryResources: ConfigurableInMemoryResources;
+
+    @postConstruct()
+    protected init(): void {
+        this.inMemoryResources.onWillDispose(resource => this.cache.delete(resource.uri.toString()));
+    }
+
+    protected readonly cache = new Map<string, [Resource, AIVariableContext]>();
+
+    getOrCreate(request: AIVariableResolutionRequest, context: AIVariableContext, value: string): ConfigurableMutableReferenceResource {
+        const uri = this.toUri(request, context);
+        try {
+            const existing = this.inMemoryResources.resolve(uri);
+            existing.update({ contents: value });
+            return existing;
+        } catch { /* No-op */ }
+        const fresh = this.inMemoryResources.add(uri, { contents: value, readOnly: true, initiallyDirty: false });
+        const key = uri.toString();
+        this.cache.set(key, [fresh, context]);
+        return fresh;
+    }
+
+    protected toUri(request: AIVariableResolutionRequest, context: AIVariableContext): URI {
+        return URI.fromComponents({
+            scheme: AI_VARIABLE_RESOURCE_SCHEME,
+            query: stableJsonStringify({ arg: request.arg, name: request.variable.name }),
+            path: '/',
+            authority: this.toAuthority(context),
+            fragment: ''
+        });
+    }
+
+    protected toAuthority(context: AIVariableContext): string {
+        try {
+            if (deepEqual(context, {})) { return NO_CONTEXT_AUTHORITY; }
+            for (const [resource, cachedContext] of this.cache.values()) {
+                if (deepEqual(context, cachedContext)) {
+                    return resource.uri.authority;
+                }
+            }
+        } catch (err) {
+            // Mostly that deep equal could overflow the stack, but it should run into === or inequality before that.
+            console.warn('Problem evaluating context in AIVariableResourceResolver', err);
+        }
+        return generateUuid();
+    }
+
+    fromUri(uri: URI): { variableName: string, arg: string | undefined } | undefined {
+        if (uri.scheme !== AI_VARIABLE_RESOURCE_SCHEME) { return undefined; }
+        try {
+            const { name: variableName, arg } = JSON.parse(uri.query);
+            return variableName ? {
+                variableName,
+                arg,
+            } : undefined;
+        } catch { return undefined; }
+    }
+}

package/src/common/configurable-in-memory-resources.ts

@@ -0,0 +1,156 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclispeSource GmbH and others.
+//
+// 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 { injectable } from '@theia/core/shared/inversify';
+import { SyncReferenceCollection, Reference, ResourceResolver, Resource, Event, Emitter, URI } from '@theia/core';
+import { MarkdownString } from '@theia/core/lib/common/markdown-rendering';
+
+@injectable()
+/** For creating highly configurable in-memory resources */
+export class ConfigurableInMemoryResources implements ResourceResolver {
+
+    protected readonly resources = new SyncReferenceCollection<string, ConfigurableMutableResource>(uri => new ConfigurableMutableResource(new URI(uri)));
+
+    get onWillDispose(): Event<ConfigurableMutableResource> {
+        return this.resources.onWillDispose;
+    }
+
+    add(uri: URI, options: ResourceInitializationOptions): ConfigurableMutableReferenceResource {
+        const resourceUri = uri.toString();
+        if (this.resources.has(resourceUri)) {
+            throw new Error(`Cannot add already existing in-memory resource '${resourceUri}'`);
+        }
+        const resource = this.acquire(resourceUri);
+        resource.update(options);
+        return resource;
+    }
+
+    update(uri: URI, options: ResourceInitializationOptions): Resource {
+        const resourceUri = uri.toString();
+        const resource = this.resources.get(resourceUri);
+        if (!resource) {
+            throw new Error(`Cannot update non-existent in-memory resource '${resourceUri}'`);
+        }
+        resource.update(options);
+        return resource;
+    }
+
+    resolve(uri: URI): ConfigurableMutableReferenceResource {
+        const uriString = uri.toString();
+        if (!this.resources.has(uriString)) {
+            throw new Error(`In memory '${uriString}' resource does not exist.`);
+        }
+        return this.acquire(uriString);
+    }
+
+    protected acquire(uri: string): ConfigurableMutableReferenceResource {
+        const reference = this.resources.acquire(uri);
+        return new ConfigurableMutableReferenceResource(reference);
+    }
+}
+
+export type ResourceInitializationOptions = Pick<Resource, 'autosaveable' | 'initiallyDirty' | 'readOnly'>
+    & { contents?: string | Promise<string>, onSave?: Resource['saveContents'] };
+
+export class ConfigurableMutableResource implements Resource {
+    protected readonly onDidChangeContentsEmitter = new Emitter<void>();
+    readonly onDidChangeContents = this.onDidChangeContentsEmitter.event;
+    protected fireDidChangeContents(): void {
+        this.onDidChangeContentsEmitter.fire();
+    }
+
+    protected readonly onDidChangeReadonlyEmitter = new Emitter<boolean | MarkdownString>();
+    readonly onDidChangeReadOnly = this.onDidChangeReadonlyEmitter.event;
+
+    constructor(readonly uri: URI, protected options?: ResourceInitializationOptions) { }
+
+    get readOnly(): Resource['readOnly'] {
+        return this.options?.readOnly;
+    }
+
+    get autosaveable(): boolean {
+        return this.options?.autosaveable !== false;
+    }
+
+    get initiallyDirty(): boolean {
+        return !!this.options?.initiallyDirty;
+    }
+
+    readContents(): Promise<string> {
+        return Promise.resolve(this.options?.contents ?? '');
+    }
+
+    async saveContents(contents: string): Promise<void> {
+        await this.options?.onSave?.(contents);
+        this.update({ contents });
+    }
+
+    update(options: ResourceInitializationOptions): void {
+        const didContentsChange = 'contents' in options && options.contents !== this.options?.contents;
+        const didReadOnlyChange = 'readOnly' in options && options.readOnly !== this.options?.readOnly;
+        this.options = { ...this.options, ...options };
+        if (didContentsChange) {
+            this.onDidChangeContentsEmitter.fire();
+        }
+        if (didReadOnlyChange) {
+            this.onDidChangeReadonlyEmitter.fire(this.readOnly ?? false);
+        }
+    }
+
+    dispose(): void {
+        this.onDidChangeContentsEmitter.dispose();
+    }
+}
+
+export class ConfigurableMutableReferenceResource implements Resource {
+    constructor(protected reference: Reference<ConfigurableMutableResource>) { }
+
+    get uri(): URI {
+        return this.reference.object.uri;
+    }
+
+    get onDidChangeContents(): Event<void> {
+        return this.reference.object.onDidChangeContents;
+    }
+
+    dispose(): void {
+        this.reference.dispose();
+    }
+
+    readContents(): Promise<string> {
+        return this.reference.object.readContents();
+    }
+
+    saveContents(contents: string): Promise<void> {
+        return this.reference.object.saveContents(contents);
+    }
+
+    update(options: ResourceInitializationOptions): void {
+        this.reference.object.update(options);
+    }
+
+    get readOnly(): Resource['readOnly'] {
+        return this.reference.object.readOnly;
+    }
+
+    get initiallyDirty(): boolean {
+        return this.reference.object.initiallyDirty;
+    }
+
+    get autosaveable(): boolean {
+        return this.reference.object.autosaveable;
+    }
+}

package/src/common/index.ts

@@ -29,3 +29,6 @@ export * from './today-variable-contribution';
 export * from './variable-service';
 export * from './settings-service';
 export * from './language-model-service';
+export * from './token-usage-service';
+export * from './ai-variable-resource';
+export * from './configurable-in-memory-resources';

package/src/common/language-model.ts

@@ -172,7 +172,15 @@ export interface LanguageModelTextResponse {
 export const isLanguageModelTextResponse = (obj: unknown): obj is LanguageModelTextResponse =>
     !!(obj && typeof obj === 'object' && 'text' in obj && typeof (obj as { text: unknown }).text === 'string');
 
-export type LanguageModelStreamResponsePart = TextResponsePart | ToolCallResponsePart | ThinkingResponsePart;
+export type LanguageModelStreamResponsePart = TextResponsePart | ToolCallResponsePart | ThinkingResponsePart | UsageResponsePart;
+export interface UsageResponsePart {
+    input_tokens: number;
+    output_tokens: number;
+}
+export const isUsageResponsePart = (part: unknown): part is UsageResponsePart =>
+    !!(part && typeof part === 'object' &&
+        'input_tokens' in part && typeof part.input_tokens === 'number' &&
+        'output_tokens' in part && typeof part.output_tokens === 'number');
 export interface TextResponsePart {
     content: string;
 }
@@ -242,7 +250,7 @@ export namespace LanguageModelMetaData {
 }
 
 export interface LanguageModel extends LanguageModelMetaData {
-    request(request: LanguageModelRequest, cancellationToken?: CancellationToken): Promise<LanguageModelResponse>;
+    request(request: UserRequest, cancellationToken?: CancellationToken): Promise<LanguageModelResponse>;
 }
 
 export namespace LanguageModel {

package/src/common/prompt-service.ts

@@ -84,7 +84,7 @@ export interface PromptService {
      *
      * @param id the id of the prompt
      * @param args the object with placeholders, mapping the placeholder key to the value
-     * @param context the {@link AIVariableContext} to use during variable resolvement
+     * @param context the {@link AIVariableContext} to use during variable resolution
      * @param resolveVariable the variable resolving method. Fall back to using the {@link AIVariableService} if not given.
      */
     getPromptFragment(
@@ -202,9 +202,19 @@ export interface PromptCustomizationService {
     readonly onDidChangeCustomAgents: Event<void>;
 
     /**
-     * Open the custom agent yaml file.
+     * Returns all locations of existing customAgents.yml files and potential locations where
+     * new customAgents.yml files could be created.
+     *
+     * @returns An array of objects containing the URI and whether the file exists
+     */
+    getCustomAgentsLocations(): Promise<{ uri: URI, exists: boolean }[]>;
+
+    /**
+     * Opens an existing customAgents.yml file at the given URI, or creates a new one if it doesn't exist.
+     *
+     * @param uri The URI of the customAgents.yml file to open or create
      */
-    openCustomAgentYaml(): void;
+    openCustomAgentYaml(uri: URI): Promise<void>;
 }
 
 @injectable()
@@ -331,7 +341,7 @@ export class PromptServiceImpl implements PromptService {
      *
      * @param template the unresolved template text
      * @param args the object with placeholders, mapping the placeholder key to the value
-     * @param context the {@link AIVariableContext} to use during variable resolvement
+     * @param context the {@link AIVariableContext} to use during variable resolution
      * @param resolveVariable the variable resolving method. Fall back to using the {@link AIVariableService} if not given.
      */
     protected async getVariableAndArgReplacements(

package/src/common/prompt-variable-contribution.ts

@@ -13,7 +13,7 @@
 //
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
-import { CommandService, nls } from '@theia/core';
+import { CommandService, ILogger, nls } from '@theia/core';
 import { injectable, inject, optional } from '@theia/core/shared/inversify';
 import * as monaco from '@theia/monaco-editor-core';
 import {
@@ -50,6 +50,9 @@ export class PromptVariableContribution implements AIVariableContribution, AIVar
     @inject(PromptCustomizationService) @optional()
     protected readonly promptCustomizationService: PromptCustomizationService;
 
+    @inject(ILogger)
+    protected logger: ILogger;
+
     registerVariables(service: AIVariableService): void {
         service.registerResolver(PROMPT_VARIABLE, this);
         service.registerArgumentPicker(PROMPT_VARIABLE, this.triggerArgumentPicker.bind(this));
@@ -81,7 +84,12 @@ export class PromptVariableContribution implements AIVariableContribution, AIVar
                 }
             }
         }
-        return undefined;
+        this.logger.debug(`Could not resolve prompt variable '${request.variable.name}' with arg '${request.arg}'. Returning empty string.`);
+        return {
+            variable: request.variable,
+            value: '',
+            allResolvedDependencies: []
+        };
     }
 
     protected async triggerArgumentPicker(): Promise<string | undefined> {

package/src/common/protocol.ts

@@ -14,10 +14,28 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
+import { Event } from '@theia/core';
 import { LanguageModelMetaData } from './language-model';
+import { TokenUsage } from './token-usage-service';
 
 export const LanguageModelRegistryClient = Symbol('LanguageModelRegistryClient');
 export interface LanguageModelRegistryClient {
     languageModelAdded(metadata: LanguageModelMetaData): void;
     languageModelRemoved(id: string): void;
 }
+
+export const TOKEN_USAGE_SERVICE_PATH = '/services/token-usage';
+
+export const TokenUsageServiceClient = Symbol('TokenUsageServiceClient');
+
+export interface TokenUsageServiceClient {
+    /**
+     * Notify the client about new token usage
+     */
+    notifyTokenUsage(usage: TokenUsage): void;
+
+    /**
+     * An event that is fired when token usage data is updated.
+     */
+    readonly onTokenUsageUpdated: Event<TokenUsage>;
+}

package/src/common/token-usage-service.ts

@@ -0,0 +1,56 @@
+// *****************************************************************************
+// 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 { TokenUsageServiceClient } from './protocol';
+
+export const TokenUsageService = Symbol('TokenUsageService');
+
+export interface TokenUsage {
+    /** The input token count */
+    inputTokens: number;
+    /** The output token count */
+    outputTokens: number;
+    /** The model identifier */
+    model: string;
+    /** The timestamp of when the tokens were used */
+    timestamp: Date;
+    /** Request identifier */
+    requestId: string;
+}
+
+export interface TokenUsageParams {
+    /** The input token count */
+    inputTokens: number;
+    /** The output token count */
+    outputTokens: number;
+    /** Request identifier */
+    requestId: string;
+}
+
+export interface TokenUsageService {
+    /**
+     * Records token usage for a model interaction.
+     *
+     * @param model The identifier of the model that was used
+     * @param params Object containing token usage information
+     * @returns A promise that resolves when the token usage has been recorded
+     */
+    recordTokenUsage(model: string, params: TokenUsageParams): Promise<void>;
+
+    getTokenUsages(): Promise<TokenUsage[]>;
+
+    setClient(tokenUsageClient: TokenUsageServiceClient): void;
+}

package/src/common/variable-service.ts

@@ -120,6 +120,13 @@ export namespace AIVariableResolutionRequest {
             'variable' in arg &&
             typeof (arg as { variable: { name: unknown } }).variable.name === 'string';
     }
+
+    export function fromResolved(arg: ResolvedAIContextVariable): AIVariableResolutionRequest {
+        return {
+            variable: arg.variable,
+            arg: arg.arg
+        };
+    }
 }
 
 export interface AIVariableContext {
@@ -136,6 +143,11 @@ export interface AIVariableResolver {
     resolve(request: AIVariableResolutionRequest, context: AIVariableContext): Promise<ResolvedAIVariable | undefined>;
 }
 
+export interface AIVariableOpener {
+    canOpen(request: AIVariableResolutionRequest, context: AIVariableContext): MaybePromise<number>;
+    open(request: AIVariableResolutionRequest, context: AIVariableContext): Promise<void>;
+}
+
 export interface AIVariableResolverWithVariableDependencies extends AIVariableResolver {
     resolve(request: AIVariableResolutionRequest, context: AIVariableContext): Promise<ResolvedAIVariable | undefined>;
     /**
@@ -166,6 +178,7 @@ export interface AIVariableService {
     registerResolver(variable: AIVariable, resolver: AIVariableResolver): Disposable;
     unregisterResolver(variable: AIVariable, resolver: AIVariableResolver): void;
     getResolver(name: string, arg: string | undefined, context: AIVariableContext): Promise<AIVariableResolver | undefined>;
+    resolveVariable(variable: AIVariableArg, context: AIVariableContext, cache?: Map<string, ResolveAIVariableCacheEntry>): Promise<ResolvedAIVariable | undefined>;
 
     registerArgumentPicker(variable: AIVariable, argPicker: AIVariableArgPicker): Disposable;
     unregisterArgumentPicker(variable: AIVariable, argPicker: AIVariableArgPicker): void;
@@ -174,8 +187,6 @@ export interface AIVariableService {
     registerArgumentCompletionProvider(variable: AIVariable, argPicker: AIVariableArgCompletionProvider): Disposable;
     unregisterArgumentCompletionProvider(variable: AIVariable, argPicker: AIVariableArgCompletionProvider): void;
     getArgumentCompletionProvider(name: string): Promise<AIVariableArgCompletionProvider | undefined>;
-
-    resolveVariable(variable: AIVariableArg, context: AIVariableContext, cache?: Map<string, ResolveAIVariableCacheEntry>): Promise<ResolvedAIVariable | undefined>;
 }
 
 /** Contributions on the frontend can optionally implement `FrontendVariableContribution`. */
@@ -191,7 +202,7 @@ export interface ResolveAIVariableCacheEntry {
 
 export type ResolveAIVariableCache = Map<string, ResolveAIVariableCacheEntry>;
 /**
- * Creates a new, empty cache for AI variable resolvement to hand into `AIVariableService.resolveVariable`.
+ * Creates a new, empty cache for AI variable resolution to hand into `AIVariableService.resolveVariable`.
  */
 export function createAIResolveVariableCache(): Map<string, ResolveAIVariableCacheEntry> {
     return new Map();
@@ -216,6 +227,7 @@ export class DefaultAIVariableService implements AIVariableService {
     protected variables = new Map<string, AIVariable>();
     protected resolvers = new Map<string, AIVariableResolver[]>();
     protected argPickers = new Map<string, AIVariableArgPicker>();
+    protected openers = new Map<string, AIVariableOpener[]>();
     protected argCompletionProviders = new Map<string, AIVariableArgCompletionProvider>();
 
     protected readonly onDidChangeVariablesEmitter = new Emitter<void>();
@@ -225,8 +237,7 @@ export class DefaultAIVariableService implements AIVariableService {
         @inject(ContributionProvider) @named(AIVariableContribution)
         protected readonly contributionProvider: ContributionProvider<AIVariableContribution>,
         @inject(ILogger) protected readonly logger: ILogger
-    ) {
-    }
+    ) { }
 
     protected initContributions(): void {
         this.contributionProvider.getContributions().forEach(contribution => contribution.registerVariables(this));
@@ -339,22 +350,27 @@ export class DefaultAIVariableService implements AIVariableService {
         return this.argCompletionProviders.get(this.getKey(name)) ?? undefined;
     }
 
-    async resolveVariable(
-        request: AIVariableArg,
-        context: AIVariableContext,
-        cache: ResolveAIVariableCache = createAIResolveVariableCache()
-    ): Promise<ResolvedAIVariable | undefined> {
-        // Calculate unique variable cache key from variable name and argument
+    protected parseRequest(request: AIVariableArg): { variableName: string, arg: string | undefined } {
         const variableName = typeof request === 'string'
             ? request
             : typeof request.variable === 'string'
                 ? request.variable
                 : request.variable.name;
         const arg = typeof request === 'string' ? undefined : request.arg;
+        return { variableName, arg };
+    }
+
+    async resolveVariable(
+        request: AIVariableArg,
+        context: AIVariableContext,
+        cache: ResolveAIVariableCache = createAIResolveVariableCache()
+    ): Promise<ResolvedAIVariable | undefined> {
+        // Calculate unique variable cache key from variable name and argument
+        const { variableName, arg } = this.parseRequest(request);
         const cacheKey = `${variableName}${PromptText.VARIABLE_SEPARATOR_CHAR}${arg ?? ''}`;
 
         // If the current cache key exists and is still in progress, we reached a cycle.
-        // If we reach it but it has been resolved, it was part of another resolvement branch and we can simply return it.
+        // If we reach it but it has been resolved, it was part of another resolution branch and we can simply return it.
         if (cache.has(cacheKey)) {
             const existingEntry = cache.get(cacheKey)!;
             if (existingEntry.inProgress) {
@@ -364,40 +380,38 @@ export class DefaultAIVariableService implements AIVariableService {
             return existingEntry.promise;
         }
 
-        const entry: ResolveAIVariableCacheEntry = { promise: Promise.resolve(undefined), inProgress: true };
+        const entry: ResolveAIVariableCacheEntry = { promise: this.doResolve(variableName, arg, context, cache), inProgress: true };
+        entry.promise.finally(() => entry.inProgress = false);
         cache.set(cacheKey, entry);
 
-        // Asynchronously resolves a variable, handling its dependencies while preventing cyclical resolution.
-        // Selects the appropriate resolver and resolution strategy based on whether nested dependency resolution is supported.
-        const promise = (async () => {
-            const variable = this.getVariable(variableName);
-            if (!variable) {
-                return undefined;
-            }
-            const resolver = await this.getResolver(variableName, arg, context);
-            let resolved: ResolvedAIVariable | undefined;
-            if (isResolverWithDependencies(resolver)) {
-                // Explicit cast needed because Typescript does not consider the method parameter length of the type guard at compile time
-                resolved = await (resolver as AIVariableResolverWithVariableDependencies).resolve(
-                    { variable, arg },
-                    context,
-                    async (depRequest: AIVariableResolutionRequest) =>
-                        this.resolveVariable(depRequest, context, cache)
-                );
-            } else if (resolver) {
-                // Explicit cast needed because Typescript does not consider the method parameter length of the type guard at compile time
-                resolved = await (resolver as AIVariableResolver).resolve({ variable, arg }, context);
-            } else {
-                resolved = undefined;
-            }
-            return resolved ? { ...resolved, arg } : undefined;
-        })();
-
-        entry.promise = promise;
-        promise.finally(() => {
-            entry.inProgress = false;
-        });
+        return entry.promise;
+    }
 
-        return promise;
+    /**
+     * Asynchronously resolves a variable, handling its dependencies while preventing cyclical resolution.
+     * Selects the appropriate resolver and resolution strategy based on whether nested dependency resolution is supported.
+     */
+    protected async doResolve(variableName: string, arg: string | undefined, context: AIVariableContext, cache: ResolveAIVariableCache): Promise<ResolvedAIVariable | undefined> {
+        const variable = this.getVariable(variableName);
+        if (!variable) {
+            return undefined;
+        }
+        const resolver = await this.getResolver(variableName, arg, context);
+        let resolved: ResolvedAIVariable | undefined;
+        if (isResolverWithDependencies(resolver)) {
+            // Explicit cast needed because Typescript does not consider the method parameter length of the type guard at compile time
+            resolved = await (resolver as AIVariableResolverWithVariableDependencies).resolve(
+                { variable, arg },
+                context,
+                async (depRequest: AIVariableResolutionRequest) =>
+                    this.resolveVariable(depRequest, context, cache)
+            );
+        } else if (resolver) {
+            // Explicit cast needed because Typescript does not consider the method parameter length of the type guard at compile time
+            resolved = await (resolver as AIVariableResolver).resolve({ variable, arg }, context);
+        } else {
+            resolved = undefined;
+        }
+        return resolved ? { ...resolved, arg } : undefined;
     }
 }

package/src/node/ai-core-backend-module.ts

@@ -36,9 +36,13 @@ import {
     LanguageModelRegistryFrontendDelegate,
     languageModelDelegatePath,
     languageModelRegistryDelegatePath,
-    LanguageModelRegistryClient
+    LanguageModelRegistryClient,
+    TokenUsageService,
+    TokenUsageServiceClient,
+    TOKEN_USAGE_SERVICE_PATH
 } from '../common';
 import { BackendLanguageModelRegistry } from './backend-language-model-registry';
+import { TokenUsageServiceImpl } from './token-usage-service-impl';
 
 // We use a connection module to handle AI services separately for each frontend.
 const aiCoreConnectionModule = ConnectionContainerModule.create(({ bind, bindBackendService, bindFrontendService }) => {
@@ -46,6 +50,22 @@ const aiCoreConnectionModule = ConnectionContainerModule.create(({ bind, bindBac
     bind(BackendLanguageModelRegistry).toSelf().inSingletonScope();
     bind(LanguageModelRegistry).toService(BackendLanguageModelRegistry);
 
+    bind(TokenUsageService).to(TokenUsageServiceImpl).inSingletonScope();
+
+    bind(ConnectionHandler)
+        .toDynamicValue(
+            ({ container }) =>
+                new RpcConnectionHandler<TokenUsageServiceClient>(
+                    TOKEN_USAGE_SERVICE_PATH,
+                    client => {
+                        const service = container.get<TokenUsageService>(TokenUsageService);
+                        service.setClient(client);
+                        return service;
+                    }
+                )
+        )
+        .inSingletonScope();
+
     bind(LanguageModelRegistryFrontendDelegate).to(LanguageModelRegistryFrontendDelegateImpl).inSingletonScope();
     bind(ConnectionHandler)
         .toDynamicValue(

package/src/node/language-model-frontend-delegate.ts

@@ -19,7 +19,6 @@ import { CancellationToken, CancellationTokenSource, ILogger, generateUuid } fro
 import {
     LanguageModelMetaData,
     LanguageModelRegistry,
-    LanguageModelRequest,
     isLanguageModelStreamResponse,
     isLanguageModelTextResponse,
     LanguageModelStreamResponsePart,
@@ -29,6 +28,7 @@ import {
     LanguageModelResponseDelegate,
     LanguageModelRegistryClient,
     isLanguageModelParsedResponse,
+    UserRequest,
 } from '../common';
 import { BackendLanguageModelRegistry } from './backend-language-model-registry';
 
@@ -69,7 +69,7 @@ export class LanguageModelFrontendDelegateImpl implements LanguageModelFrontendD
 
     async request(
         modelId: string,
-        request: LanguageModelRequest,
+        request: UserRequest,
         requestId: string,
         cancellationToken?: CancellationToken
     ): Promise<LanguageModelResponseDelegate> {