Skip to content

feat(agent): cumulative prompt/completion token-split accessors #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
indykish wants to merge 1 commit into
base: release/v2026.5.29-zmb
Choose a base branch
from
Open

feat(agent): cumulative prompt/completion token-split accessors #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/agent/commands.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3313,6 +3313,12 @@ fn clearSessionState(self: anytype) void {
if (@hasField(@TypeOf(self.*), "total_tokens")) {
self.total_tokens = 0;
}
if (@hasField(@TypeOf(self.*), "prompt_tokens_total")) {
self.prompt_tokens_total = 0;
}
if (@hasField(@TypeOf(self.*), "completion_tokens_total")) {
self.completion_tokens_total = 0;
}
if (@hasField(@TypeOf(self.*), "last_turn_usage")) {
self.last_turn_usage = .{};
}
Expand Down
34 changes: 34 additions & 0 deletions src/agent/root.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,12 @@ pub const Agent = struct {
/// Total tokens used across all turns.
total_tokens: u64 = 0,

/// Cumulative prompt-side tokens across all turns (split of `total_tokens`).
prompt_tokens_total: u64 = 0,

/// Cumulative completion-side tokens across all turns (split of `total_tokens`).
completion_tokens_total: u64 = 0,

/// Total cost in USD across all turns.
total_cost_usd: f64 = 0,

Expand Down Expand Up @@ -2435,6 +2441,8 @@ pub const Agent = struct {
response.usage = normalized_usage;

self.total_tokens += normalized_usage.total_tokens;
self.prompt_tokens_total += normalized_usage.prompt_tokens;
self.completion_tokens_total += normalized_usage.completion_tokens;
Comment on lines 2443 to +2445

@greptile-apps greptile-apps Bot Jun 12, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Split counters silently zero-out when provider omits per-side data

When a provider sets total_tokens > 0 but leaves prompt_tokens = 0 and completion_tokens = 0 (neither normalization block fires because total_tokens is already non-zero), the accumulators for prompt_tokens_total and completion_tokens_total receive 0 for that turn while total_tokens advances correctly. Over multiple such turns, promptTokensUsed() + completionTokensUsed() falls short of tokensUsed() with no signal to the caller. Since the stated purpose of these accessors is by-side billing, an embedder that relies on them for cost metering will silently under-charge prompt-side and completion-side spend for any provider that only surfaces a total — a hard-to-detect billing gap.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: src/agent/root.zig
Line: 2443-2445

Comment:
**Split counters silently zero-out when provider omits per-side data**

When a provider sets `total_tokens > 0` but leaves `prompt_tokens = 0` and `completion_tokens = 0` (neither normalization block fires because `total_tokens` is already non-zero), the accumulators for `prompt_tokens_total` and `completion_tokens_total` receive 0 for that turn while `total_tokens` advances correctly. Over multiple such turns, `promptTokensUsed() + completionTokensUsed()` falls short of `tokensUsed()` with no signal to the caller. Since the stated purpose of these accessors is by-side billing, an embedder that relies on them for cost metering will silently under-charge prompt-side and completion-side spend for any provider that only surfaces a total — a hard-to-detect billing gap.

How can I resolve this? If you propose a fix, please make it concise.

self.total_cost_usd += cost_mod.TokenUsage.fromProviders(turn_model_name, normalized_usage).cost();
self.last_turn_usage = normalized_usage;
if (normalized_usage.total_tokens > 0) {
Expand Down Expand Up @@ -2853,6 +2861,8 @@ pub const Agent = struct {
}
summary_response.usage = normalized_summary_usage;
self.total_tokens += normalized_summary_usage.total_tokens;
self.prompt_tokens_total += normalized_summary_usage.prompt_tokens;
self.completion_tokens_total += normalized_summary_usage.completion_tokens;
self.total_cost_usd += cost_mod.TokenUsage.fromProviders(self.model_name, normalized_summary_usage).cost();
self.last_turn_usage = normalized_summary_usage;
if (normalized_summary_usage.total_tokens > 0) {
Expand Down Expand Up @@ -3812,6 +3822,16 @@ pub const Agent = struct {
return self.total_tokens;
}

/// Cumulative prompt-side tokens across all turns.
pub fn promptTokensUsed(self: *const Agent) u64 {
return self.prompt_tokens_total;
}

/// Cumulative completion-side tokens across all turns.
pub fn completionTokensUsed(self: *const Agent) u64 {
Comment thread
greptile-apps[bot] marked this conversation as resolved.
return self.completion_tokens_total;
}

/// Get current history length.
pub fn historyLen(self: *const Agent) usize {
return self.history.items.len;
Expand Down Expand Up @@ -4093,6 +4113,14 @@ test "Agent tokens tracking" {
try std.testing.expectEqual(@as(u64, 100), agent.tokensUsed());
agent.total_tokens += 50;
try std.testing.expectEqual(@as(u64, 150), agent.tokensUsed());

// Split accessors mirror their cumulative fields and default to zero.
try std.testing.expectEqual(@as(u64, 0), agent.promptTokensUsed());
try std.testing.expectEqual(@as(u64, 0), agent.completionTokensUsed());
agent.prompt_tokens_total = 90;
agent.completion_tokens_total = 60;
try std.testing.expectEqual(@as(u64, 90), agent.promptTokensUsed());
try std.testing.expectEqual(@as(u64, 60), agent.completionTokensUsed());
}

test "Agent trimHistory no-op when under limit" {
Expand Down Expand Up @@ -5183,6 +5211,8 @@ test "slash /new clears history" {
});
agent.has_system_prompt = true;
agent.total_tokens = 42;
agent.prompt_tokens_total = 30;
agent.completion_tokens_total = 12;
agent.last_turn_usage = .{ .prompt_tokens = 10, .completion_tokens = 5, .total_tokens = 15 };

const response = (try agent.handleSlashCommand("/new")).?;
Expand All @@ -5192,6 +5222,10 @@ test "slash /new clears history" {
try std.testing.expectEqual(@as(usize, 0), agent.historyLen());
try std.testing.expect(!agent.has_system_prompt);
try std.testing.expectEqual(@as(u64, 0), agent.total_tokens);
// The split counters reset with the session too — else by-side metering
// carries stale per-session figures after /new.
try std.testing.expectEqual(@as(u64, 0), agent.promptTokensUsed());
try std.testing.expectEqual(@as(u64, 0), agent.completionTokensUsed());
try std.testing.expectEqual(@as(u32, 0), agent.last_turn_usage.total_tokens);
}

Expand Down