acecode/src/utils/token_tracker.hpp at master · tmoonlight/acecode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
#pragma once

#include "../provider/llm_provider.hpp"
#include <string>
#include <mutex>

namespace acecode {

class TokenTracker {
public:
    struct Snapshot {
        TokenUsage last;
        TokenUsage session;
    };

    // Record usage from a server response.
    // Updates both "last" (current context) and "session" (cumulative) counters.
    void record(const TokenUsage& usage);

    // Estimate and record tokens from character count (fallback when no server usage)
    void record_estimate(int char_count);

    // --- "Last" counters (most recent API call = current context occupancy) ---
    int last_prompt_tokens() const;
    int last_completion_tokens() const;

    // --- Session cumulative counters ---
    int prompt_tokens() const;
    int completion_tokens() const;
    int total_tokens() const;
    int cache_read_tokens() const;
    int cache_write_tokens() const;

    Snapshot snapshot() const;
    void restore(const Snapshot& snapshot);
    void restore(const TokenUsage& last, const TokenUsage& session);

    // Reset all counters
    void reset();

    // Format token count for display (e.g., "1.2k", "45.3k")
    static std::string format_tokens(int count);

    // Format status bar string: "8.0k/128k"
    // Shows current context occupancy (last_prompt_tokens / context_window)
    std::string format_status(int context_window) const;

    // Share of total input tokens served from the provider's prompt cache,
    // for the most recent call and for the session so far. Returns -1 only
    // when the server has reported no usage at all, so callers can hide the
    // indicator instead of guessing. A provider that reports usage but no
    // cache counters reads as 0%, which is the honest answer and the signal
    // that makes a missing-cache regression visible instead of silent.
    int last_cache_hit_percent() const;
    int cache_hit_percent() const;

    // "cache 87%" / empty string when unknown. Session-cumulative.
    std::string format_cache_status() const;
    static std::string format_cache_status_for(int cache_hit_percent);

    // prompt_tokens is the total input across every provider, with
    // cache_read_tokens a subset of it (the Anthropic provider normalizes its
    // split counters to match). Returns -1 when the ratio is unknown.
    static int cache_hit_percent_for(int prompt_tokens, int cache_read_tokens);

    // Context occupancy percentage for the status bar, rounded to nearest int.
    // Returns 0 for invalid context windows and caps overfull contexts at 100.
    int context_percent(int context_window) const;

    static int context_percent_for(int prompt_tokens, int context_window);

private:
    mutable std::mutex mu_;

    // Most recent API call (current context window occupancy)
    int last_prompt_tokens_ = 0;
    int last_completion_tokens_ = 0;
    int last_total_tokens_ = 0;
    int last_cache_read_tokens_ = 0;
    int last_cache_write_tokens_ = 0;
    int last_reasoning_tokens_ = 0;
    bool last_has_data_ = false;

    // Session cumulative
    int session_prompt_tokens_ = 0;
    int session_completion_tokens_ = 0;
    int session_total_tokens_ = 0;
    int session_cache_read_tokens_ = 0;
    int session_cache_write_tokens_ = 0;
    int session_reasoning_tokens_ = 0;
    bool session_has_data_ = false;
};

} // namespace acecode