TokenCapChatClient.cs

using System.Runtime.CompilerServices;
using Microsoft.Extensions.AI;

namespace BlogWriter;

/// <summary>
/// Chat-client middleware that enforces a hard cap on cumulative token usage
/// across every model round-trip in the process (including the extra calls made
/// during tool invocation). When the running total exceeds the cap, the
/// application is terminated with an explanatory message rather than continuing
/// to spend tokens.
/// </summary>
public sealed class TokenCapChatClient : DelegatingChatClient
{
    // Key used by the OpenAI connector to report reasoning tokens inside
    // UsageDetails.AdditionalCounts (there is no dedicated top-level property).
    private const string ReasoningTokenCountKey = "OutputTokenDetails.ReasoningTokenCount";

    private readonly TokenBudget _budget;

    public TokenCapChatClient(IChatClient innerClient, long maxTotalTokens)
        : this(innerClient, new TokenBudget(maxTotalTokens))
    {
    }

    private TokenCapChatClient(IChatClient innerClient, TokenBudget budget) : base(innerClient) =>
        _budget = budget;

    /// <summary>
    /// Creates a MAF chat-client middleware factory whose clients share one
    /// cumulative process-wide token budget.
    /// </summary>
    public static Func<IChatClient, IChatClient> CreateSharedFactory(long maxTotalTokens)
    {
        var budget = new TokenBudget(maxTotalTokens);
        return innerClient => new TokenCapChatClient(innerClient, budget);
    }

    /// <summary>Cumulative token usage observed across every model round-trip so far.</summary>
    public TokenUsageSnapshot UsageSnapshot => _budget.UsageSnapshot;

    public override async Task<ChatResponse> GetResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        ChatResponse response = await base.GetResponseAsync(messages, options, cancellationToken);
        Track(response.Usage);
        return response;
    }

    public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        await foreach (ChatResponseUpdate update in
            base.GetStreamingResponseAsync(messages, options, cancellationToken))
        {
            foreach (AIContent content in update.Contents)
            {
                if (content is UsageContent usageContent)
                {
                    Track(usageContent.Details);
                }
            }

            yield return update;
        }
    }

    private void Track(UsageDetails? usage)
    {
        _budget.Track(usage);
    }

    private sealed class TokenBudget
    {
        private readonly long _maxTotalTokens;
        private long _totalTokens;
        private long _inputTokens;
        private long _outputTokens;
        private long _reasoningTokens;

        public TokenBudget(long maxTotalTokens)
        {
            _maxTotalTokens = maxTotalTokens > 0
                ? maxTotalTokens
                : throw new ArgumentOutOfRangeException(nameof(maxTotalTokens), maxTotalTokens, "Token cap must be a positive number.");
        }

        public TokenUsageSnapshot UsageSnapshot => new(
            Interlocked.Read(ref _inputTokens),
            Interlocked.Read(ref _outputTokens),
            Interlocked.Read(ref _reasoningTokens),
            Interlocked.Read(ref _totalTokens));

        public void Track(UsageDetails? usage)
        {
            if (usage is null)
            {
                return;
            }

            Interlocked.Add(ref _inputTokens, usage.InputTokenCount ?? 0);
            Interlocked.Add(ref _outputTokens, usage.OutputTokenCount ?? 0);
            if (usage.AdditionalCounts is { } additionalCounts &&
                additionalCounts.TryGetValue(ReasoningTokenCountKey, out long reasoningTokens))
            {
                Interlocked.Add(ref _reasoningTokens, reasoningTokens);
            }

            long used = usage.TotalTokenCount ?? 0;
            if (used == 0)
            {
                return;
            }

            long total = Interlocked.Add(ref _totalTokens, used);
            if (total > _maxTotalTokens)
            {
                throw new TokenCapExceededException(total, _maxTotalTokens);
            }
        }
    }
}

/// <summary>Point-in-time totals of tokens consumed across all model round-trips.</summary>
public readonly record struct TokenUsageSnapshot(long InputTokens, long OutputTokens, long ReasoningTokens, long TotalTokens);

/// <summary>
/// Thrown when cumulative model token usage exceeds the configured cap. Callers
/// catch this to shut down gracefully instead of continuing to spend tokens.
/// </summary>
public sealed class TokenCapExceededException(long tokensUsed, long tokenLimit)
    : Exception($"Token cap exceeded: consumed {tokensUsed} tokens, limit is {tokenLimit}.")
{
    public long TokensUsed { get; } = tokensUsed;

    public long TokenLimit { get; } = tokenLimit;
}