Getting token usage from agent orchestration

[Arno Zeelmaekers (ArnoZeelmaekers)]

Hey all,

I've been trying to figure out how we can get our agent orchestration token usage.
We're using the MagenticOrchestration orchestration and are able to get the agent consumption trough a responsecall back and chat history. But the orchestrator itself is also doing ai calls and using tokens which are no where to be able to be retrieved.
Does anyone have any idea where to get the orchestrators token usage from after the run is done?

    public async Task<AiResponseResult<string>> GetAgentPlainResponse(AiAgentDto ai, string userContent, CancellationToken cancellationToken)
    {
        var chatHistory = new ChatHistory();
        var agentKernel = await CreateAgentKernel(ai, chatHistory, cancellationToken);
        var runTime = new InProcessRuntime();
        await runTime.StartAsync(cancellationToken);

        var result = await agentKernel.InvokeAsync(userContent,
            runTime,
            cancellationToken);
        
        var output = await result.GetValueAsync(TimeSpan.FromSeconds(300), cancellationToken);
        await runTime.RunUntilIdleAsync();
        var aiResponseResult = new AiResponseResult<string>
        {
            AiModelUsed = ai.Ai.AiModelId.ToString(),
            ReasoningTokensUsed = 0,
            Result = output,
            UsedFiles = _usedFiles,
            InputTokensUsed = 0,
            OutputTokensUsed = 0
        };

        return aiResponseResult;
    }
    
    private async Task<MagenticOrchestration> CreateAgentKernel(AiAgentDto aiAgent, ChatHistory chatHistory, CancellationToken cancellationToken = default)
    {
        var agents = new List<Agent>();

        foreach (var a in aiAgent.AiAgentConfigurations)
        {
            var aiModel = await aiService.GetAiModelById(a.Ai.AiModelId, cancellationToken);
            if (aiModel == null)
                throw new AiNotFoundException($"Cannot find AI model with ID {a.Ai.AiModelId}");
            var chatCompletionAgentKernel = Kernel.CreateBuilder()
                .AddAzureOpenAIChatCompletion(aiModel.DeploymentName, azureOpenAiClient);

            await AddAiPluginsInKernelBuilder(chatCompletionAgentKernel, a.Ai, cancellationToken);
            var kernel = chatCompletionAgentKernel.Build();
            await AddAiPluginsInKernel(kernel, a.Ai, cancellationToken);

            agents.Add(new ChatCompletionAgent
            {
                Name = a.Name,
                Kernel = kernel,
                Description = a.Description,
                Instructions = a.Ai.Prompt,
                Arguments = new KernelArguments(new AzureOpenAIPromptExecutionSettings
                {
                    FunctionChoiceBehavior = FunctionChoiceBehavior.Auto(autoInvoke: true),
                    Temperature = a.Ai.Temperature
                })
            });
        }

        var managerAiModel = await aiService.GetAiModelById(aiAgent.Ai.AiModelId, cancellationToken);
        if (managerAiModel == null)
            throw new AiNotFoundException($"Cannot find AI model with ID {aiAgent.Ai.AiModelId}");

        var managerKernel = Kernel.CreateBuilder()
            .AddAzureOpenAIChatCompletion(managerAiModel.DeploymentName, azureOpenAiClient)
            .Build();

        ValueTask ResponseCallback(ChatMessageContent response)
        {
            chatHistory.Add(response);
            return ValueTask.CompletedTask;
        }
        
        var manager = new StandardMagenticManager(managerKernel.GetRequiredService<IChatCompletionService>(), new AzureOpenAIPromptExecutionSettings
        {
            ResponseFormat = aiAgent.Ai.OutputSchema ?? "text",
            ChatSystemPrompt = aiAgent.Ai.Prompt
        })
        {
            MaximumInvocationCount = aiAgent.MaxInvocations
        };
        var orchestration = new MagenticOrchestration(manager, agents.ToArray())
        {
            ResponseCallback = ResponseCallback,
            Description = aiAgent.Ai.Prompt,
            Name = aiAgent.Name
        };
        return orchestration;
    }

[Evan Mattson (moonbox3)]

Adding Chris (@crickman)

[John CSA (jluocsa)]

The StandardMagenticManager is a sealed class that uses the IChatCompletionService you hand it for its own LLM calls — PlanAsync, ReplanAsync, EvaluateTaskProgressAsync, PrepareFinalAnswerAsync, plus the internal facts/plan prompts. You can see this in StandardMagenticManager.cs: every internal call goes through this._service.GetChatMessageContentAsync(...), and the manager keeps response.Content only — the ChatMessageContent.Metadata["Usage"] is never surfaced through the orchestration API.

Two approaches that work today.

1. Decorate the IChatCompletionService you pass to the manager (cleanest)

Wrap your manager's chat service so it accumulates the Usage from every response's metadata, then read the counters after the orchestration finishes:

internal sealed class UsageTrackingChatCompletionService(IChatCompletionService inner) : IChatCompletionService
{
    private long _promptTokens, _completionTokens, _totalTokens;

    public long PromptTokens     => Interlocked.Read(ref _promptTokens);
    public long CompletionTokens => Interlocked.Read(ref _completionTokens);
    public long TotalTokens      => Interlocked.Read(ref _totalTokens);

    public IReadOnlyDictionary<string, object?> Attributes => inner.Attributes;

    public async Task<IReadOnlyList<ChatMessageContent>> GetChatMessageContentsAsync(
        ChatHistory chatHistory, PromptExecutionSettings? executionSettings = null,
        Kernel? kernel = null, CancellationToken cancellationToken = default)
    {
        var results = await inner
            .GetChatMessageContentsAsync(chatHistory, executionSettings, kernel, cancellationToken)
            .ConfigureAwait(false);

        foreach (var r in results)
        {
            // Azure OpenAI / OpenAI surface usage under metadata key "Usage" as OpenAI.Chat.ChatTokenUsage.
            if (r.Metadata is { } md
                && md.TryGetValue("Usage", out var usageObj)
                && usageObj is OpenAI.Chat.ChatTokenUsage usage)
            {
                Interlocked.Add(ref _promptTokens,     usage.InputTokenCount);
                Interlocked.Add(ref _completionTokens, usage.OutputTokenCount);
                Interlocked.Add(ref _totalTokens,      usage.TotalTokenCount);
            }
        }
        return results;
    }

    public async IAsyncEnumerable<StreamingChatMessageContent> GetStreamingChatMessageContentsAsync(
        ChatHistory chatHistory, PromptExecutionSettings? executionSettings = null,
        Kernel? kernel = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        await foreach (var chunk in inner
            .GetStreamingChatMessageContentsAsync(chatHistory, executionSettings, kernel, cancellationToken)
            .ConfigureAwait(false))
        {
            // The final streaming chunk carries Usage when stream_options.include_usage=true.
            if (chunk.Metadata is { } md
                && md.TryGetValue("Usage", out var usageObj)
                && usageObj is OpenAI.Chat.ChatTokenUsage usage)
            {
                Interlocked.Add(ref _promptTokens,     usage.InputTokenCount);
                Interlocked.Add(ref _completionTokens, usage.OutputTokenCount);
                Interlocked.Add(ref _totalTokens,      usage.TotalTokenCount);
            }
            yield return chunk;
        }
    }
}

Wire it in where you build the manager:

var rawManagerChat = managerKernel.GetRequiredService<IChatCompletionService>();
var managerUsage   = new UsageTrackingChatCompletionService(rawManagerChat);

var manager = new StandardMagenticManager(managerUsage, new AzureOpenAIPromptExecutionSettings
{
    ResponseFormat   = aiAgent.Ai.OutputSchema ?? "text",
    ChatSystemPrompt = aiAgent.Ai.Prompt
})
{
    MaximumInvocationCount = aiAgent.MaxInvocations
};

// ... run orchestration ...

aiResponseResult.InputTokensUsed  += (int)managerUsage.PromptTokens;
aiResponseResult.OutputTokensUsed += (int)managerUsage.CompletionTokens;

If you also want one combined total for the whole run, wrap each agent's IChatCompletionService the same way and sum the counters — or keep the per-agent counters you already extract from ResponseCallback's Metadata["Usage"] and just add the manager's counter to them at the end.

2. Use OpenTelemetry (recommended for production)

SK already emits a span per chat completion call with the gen_ai.usage.input_tokens / gen_ai.usage.output_tokens attributes (and the gen_ai.client.token.usage metric). With SK's OTel sources enabled, every manager call and every agent call shows up as a child span under the same orchestration trace, and you can aggregate token counts across all spans of a trace without modifying the manager at all. See the Observability docs and the OTel samples under dotnet/samples/Demos/TelemetryWithAppInsights. This is what I would do if there's already an APM in the picture.

Why subclassing the manager isn't an option

StandardMagenticManager is sealed, so you can't override its internals to surface ChatMessageContent.Metadata["Usage"]. The public base MagenticManager is abstract, so you can write a manager from scratch — but for "I just want token counts" that's far more work than the decorator above.