Skip to content

feat(agent-metrics): add encrypted usage snapshots #7443

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
timmstudio wants to merge 6 commits into
base: main
Choose a base branch
from
Open

feat(agent-metrics): add encrypted usage snapshots #7443

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5021a4f
feat(agent-metrics): persist encrypted context snapshots
Sep 5, 2026
bf788c4
feat(desktop): surface compact agent usage status
Sep 6, 2026
627286b
feat(mobile): show agent usage indicator on bot avatars
Sep 5, 2026
0889202
fix(mobile): verify and align agent usage snapshots
Sep 6, 2026
be8f27e
test(desktop): use generic agent usage fixtures
Sep 6, 2026
4a3d363
fix(agent-metrics): harden usage snapshot handling
Sep 7, 2026
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
40 changes: 39 additions & 1 deletion crates/buzz-acp/src/pool.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2127,11 +2127,15 @@ pub async fn run_prompt_task(
turn_id: String,
) {
// Is this a channel prompt or a heartbeat?
let source = match &batch {
let source = match batch.as_ref() {
Some(b) => PromptSource::Channel(b.scope.clone()),
None => PromptSource::Heartbeat,
};
let observer_channel_id = source.channel_id();
let observer_thread_root_id = match &source {
PromptSource::Channel(scope) => scope.root_event_id().map(str::to_string),
PromptSource::Heartbeat => None,
};
let turn_started_at = chrono::Utc::now().to_rfc3339();
agent.acp.set_observer_context(observer::context_for_turn(
observer_channel_id,
Expand Down Expand Up @@ -2576,6 +2580,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
Some(*cid),
scope.root_event_id(),
&session_id,
&format!("{turn_id}:initial"),
Some(acp_stop_to_core(&stop_reason)),
Expand Down Expand Up @@ -2611,6 +2616,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
Some(*cid),
scope.root_event_id(),
&session_id,
&format!("{turn_id}:initial"),
Some(acp_stop_to_core(&stop_reason)),
Expand Down Expand Up @@ -2924,6 +2930,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Cancelled),
Expand Down Expand Up @@ -2960,6 +2967,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Error),
Expand Down Expand Up @@ -3025,6 +3033,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::EndTurn),
Expand Down Expand Up @@ -3099,6 +3108,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(core_stop),
Expand All @@ -3122,6 +3132,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Error),
Expand Down Expand Up @@ -3154,6 +3165,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Cancelled),
Expand Down Expand Up @@ -3182,6 +3194,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Error),
Expand All @@ -3207,6 +3220,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Error),
Expand Down Expand Up @@ -3236,6 +3250,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Error),
Expand Down Expand Up @@ -3263,6 +3278,7 @@ pub async fn run_prompt_task(
&ctx,
usage,
observer_channel_id,
observer_thread_root_id.as_deref(),
&session_id,
&turn_id,
Some(buzz_core::agent_turn_metric::StopReason::Error),
Expand Down Expand Up @@ -4911,6 +4927,7 @@ async fn publish_agent_turn_metric(
ctx: &PromptContext,
usage: Option<crate::usage::TurnUsage>,
channel_id: Option<uuid::Uuid>,
thread_root_id: Option<&str>,
session_id: &str,
turn_id: &str,
stop_reason: Option<buzz_core::agent_turn_metric::StopReason>,
Expand All @@ -4929,6 +4946,7 @@ async fn publish_agent_turn_metric(
harness: ctx.harness_name.clone(),
model: usage.model.clone(),
channel_id: channel_id.map(|id| id.to_string()),
thread_root_id: thread_root_id.map(str::to_string),
session_id: Some(usage.session_id.clone()),
turn_id: Some(turn_id.to_string()),
turn_seq: Some(usage.turn_seq),
Expand All @@ -4937,6 +4955,9 @@ async fn publish_agent_turn_metric(
cumulative: cumulative_counts,
delta_reliable: usage.delta_reliable,
stop_reason,
context_used_tokens: usage.context_used_tokens,
context_limit_tokens: usage.context_limit_tokens,
account_usage_windows: Vec::new(),
pricing_identity: usage.pricing_identity.clone(),
};
let ciphertext = match buzz_core::agent_turn_metric::encrypt_agent_turn_metric(
Expand Down Expand Up @@ -8551,6 +8572,7 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
&ctx,
None,
None,
None,
"sess-1",
"turn-1",
Some(buzz_core::agent_turn_metric::StopReason::EndTurn),
Expand Down Expand Up @@ -8579,13 +8601,16 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
cumulative_cache_read_tokens: None,
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};
// owner_pubkey = None → early return, no panic.
publish_agent_turn_metric(
&ctx,
Some(usage),
None,
None,
"sess-1",
"turn-1",
Some(buzz_core::agent_turn_metric::StopReason::EndTurn),
Expand Down Expand Up @@ -8618,13 +8643,16 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
cumulative_cache_read_tokens: None,
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};
// Will try to publish and fail (no real relay) but must not panic.
publish_agent_turn_metric(
&ctx,
Some(usage),
Some(uuid::Uuid::new_v4()),
None,
"sess-1",
"turn-1",
Some(buzz_core::agent_turn_metric::StopReason::EndTurn),
Expand Down Expand Up @@ -8658,13 +8686,16 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
cumulative_cache_read_tokens: None,
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};
// Must not panic; HTTP submit will fail (no real relay) — that's fine.
publish_agent_turn_metric(
&ctx,
Some(usage),
Some(uuid::Uuid::new_v4()),
None,
"sess-cancel",
"turn-cancel",
Some(buzz_core::agent_turn_metric::StopReason::Cancelled),
Expand Down Expand Up @@ -8698,13 +8729,16 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
cumulative_cache_read_tokens: None,
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};
// Will try to publish (encrypt succeeds) and fail HTTP (no relay) — must not panic.
publish_agent_turn_metric(
&ctx,
Some(usage),
Some(uuid::Uuid::new_v4()),
None,
"sess-ba",
"turn-ba",
Some(buzz_core::agent_turn_metric::StopReason::EndTurn),
Expand Down Expand Up @@ -8735,6 +8769,8 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
cumulative_cache_read_tokens: None,
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};

Expand Down Expand Up @@ -8787,6 +8823,8 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
cumulative_cache_read_tokens: None,
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};

Expand Down
41 changes: 41 additions & 0 deletions crates/buzz-acp/src/usage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,10 @@ pub struct TurnUsage {
/// Effective model id for this turn (maps to NIP-AM `model`). `None` if the
/// harness did not include the model in its usage notification.
pub model: Option<String>,
/// Input-side context used by the last successful model request this turn.
pub context_used_tokens: Option<u64>,
/// Context-window capacity paired with `context_used_tokens`.
pub context_limit_tokens: Option<u64>,
/// Billing identity for this turn, as received from the publisher.
/// `None` when the publisher omitted it (unrecognised endpoint, mixed
/// identities, old harness). Per-turn only — not session-cumulative.
Expand Down Expand Up @@ -391,6 +395,8 @@ impl StandardUsageTracker {
cumulative_cache_read_tokens: None,
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
})
}
Expand Down Expand Up @@ -729,6 +735,12 @@ impl UsageTracker {
// Already poisoned: stays poisoned regardless of this notification.
poisoned @ Some(None) => poisoned,
};
let (context_used_tokens, context_limit_tokens) =
if payload.used > 0 && payload.context_limit > 0 {
(Some(payload.used), Some(payload.context_limit))
} else {
(None, None)
};
self.pending = Some(TurnUsage {
session_id: session_id.to_string(),
turn_seq,
Expand All @@ -746,6 +758,8 @@ impl UsageTracker {
cumulative_cache_read_tokens: current_cached_input,
cumulative_cache_write_tokens: current_cache_write,
model: payload.model.clone(),
context_used_tokens,
context_limit_tokens,
// The folded identity is written in take() — use a placeholder
// here and replace it before returning the record.
pricing_identity: None,
Expand Down Expand Up @@ -1009,6 +1023,29 @@ mod tests {
}
}

// ── Context-window snapshot threading ───────────────────────────────────

#[test]
fn context_snapshot_is_last_write_wins_and_invalid_pair_means_unknown() {
let mut tracker = UsageTracker::default();
tracker.seed_zero_baseline("sess-context");
tracker.begin_turn("sess-context");

let mut first = payload(100, 20, None);
first.used = 90;
first.context_limit = 200_000;
tracker.record("sess-context", &first);

let mut final_snapshot = payload(180, 40, None);
final_snapshot.used = 210_000;
final_snapshot.context_limit = 0;
tracker.record("sess-context", &final_snapshot);

let usage = tracker.take().expect("turn usage");
assert_eq!(usage.context_used_tokens, None);
assert_eq!(usage.context_limit_tokens, None);
}

// ── Turn scoping: setup notifications must not pollute the first real turn ─

#[test]
Expand Down Expand Up @@ -1988,6 +2025,8 @@ mod tests {
cumulative_cache_read_tokens: None, // harness did not report the field
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};

Expand Down Expand Up @@ -2029,6 +2068,8 @@ mod tests {
cumulative_cache_read_tokens: Some(600),
cumulative_cache_write_tokens: None,
model: None,
context_used_tokens: None,
context_limit_tokens: None,
pricing_identity: None,
};

Expand Down
23 changes: 14 additions & 9 deletions crates/buzz-agent/src/agent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,22 +292,27 @@ impl RunCtx<'_> {
let write_total = base
.cache_write_tokens
.merge_session(*self.turn_cache_write_tokens);
let payload = wire::usage_update_payload(
base.input_tokens
let payload = wire::usage_update_payload_with_context(wire::UsageUpdateSnapshot {
accumulated_input_tokens: base
.input_tokens
.merge_session(*self.turn_input_tokens)
.exact_value(),
base.output_tokens
accumulated_output_tokens: base
.output_tokens
.merge_session(*self.turn_output_tokens)
.exact_value(),
cached_total.exact_value(),
write_total.exact_value(),
base.total_state.merge_session(*self.turn_total_state),
self.effective_model,
accumulated_cached_input_tokens: cached_total.exact_value(),
accumulated_cache_write_tokens: write_total.exact_value(),
accumulated_total: base.total_state.merge_session(*self.turn_total_state),
model: self.effective_model,
// Extract the proven identity if this turn is consistent so far.
self.turn_pricing_identity
pricing_identity: self
.turn_pricing_identity
.as_ref()
.and_then(|inner| inner.as_ref()),
);
context_used_tokens: *self.last_request_input_tokens,
context_limit_tokens: Some(self.cfg.max_context_tokens),
});
wire::send(
self.wire,
wire::goose_session_update(self.session_id, payload),
Expand Down
Loading