AnchorCli/OpenRouter/TokenTracker.cs

namespace AnchorCli.OpenRouter;

/// <summary>
/// Tracks token usage and calculates costs for the session.
/// </summary>
internal sealed class TokenTracker
{
    private ChatSession _session;
    
    public TokenTracker(ChatSession session)
    {
        _session = session;
    }

    /// <summary>
    /// Gets or sets the session. Allows setting the session after construction
    /// to support dependency injection patterns.
    /// </summary>
    public ChatSession Session
    {
        get => _session;
        set => _session = value;
    }
    
    public string Provider { get; set; } = "Unknown";
    public long SessionInputTokens => _session.SessionInputTokens;
    public long SessionOutputTokens => _session.SessionOutputTokens;
    public int RequestCount => _session.RequestCount;

    /// <summary>Maximum context window for the model (tokens). 0 = unknown.</summary>
    public int ContextLength { get; set; }

    /// <summary>Input tokens from the most recent API response — approximates current context size.</summary>
    public int LastInputTokens { get; private set; }

    /// <summary>USD per input token.</summary>
    public decimal InputPrice { get; set; }

    /// <summary>USD per output token.</summary>
    public decimal OutputPrice { get; set; }

    /// <summary>Fixed USD per API request.</summary>
    public decimal RequestPrice { get; set; }
    /// <summary>
    /// Record usage from one response (may span multiple LLM rounds).
    /// </summary>
    public void AddUsage(int inputTokens, int outputTokens)
    {
        _session.SessionInputTokens += inputTokens;
        _session.SessionOutputTokens += outputTokens;
        LastInputTokens = inputTokens;
        _session.RequestCount++;
    }
    public void Reset()
    {
        _session.SessionInputTokens = 0;
        _session.SessionOutputTokens = 0;
        _session.RequestCount = 0;
        LastInputTokens = 0;
    }


    private const int MaxContextReserve = 150_000;

    /// <summary>
    /// Returns true if the context is getting too large and should be compacted.
    /// Triggers at min(75% of model context, 150K tokens).
    /// </summary>
    public bool ShouldCompact()
    {
        if (LastInputTokens <= 0) return false;

        int threshold = ContextLength > 0
            ? Math.Min((int)(ContextLength * 0.75), MaxContextReserve)
            : MaxContextReserve;

        return LastInputTokens >= threshold;
    }

    /// <summary>Context usage as a percentage (0-100). Returns -1 if context length is unknown.</summary>
    public double ContextUsagePercent =>
        ContextLength > 0 && LastInputTokens > 0
            ? (double)LastInputTokens / ContextLength * 100.0
            : -1;

    /// <summary>
    /// Calculate cost for a single response.
    /// </summary>
    public decimal CalculateCost(int inputTokens, int outputTokens) =>
        inputTokens * InputPrice +
        outputTokens * OutputPrice +
        RequestPrice;

    /// <summary>
    /// Total session cost.
    /// </summary>
    public decimal SessionCost =>
        SessionInputTokens * InputPrice +
        SessionOutputTokens * OutputPrice +
        RequestCount * RequestPrice;

    public static string FormatTokens(long count) =>
        count >= 1_000 ? $"{count / 1_000.0:F1}k" : count.ToString("N0");

    public static string FormatCost(decimal cost) =>
        cost < 0.01m ? $"${cost:F4}" : $"${cost:F2}";
}