feat(hive-claude): parse turn telemetry in the lib, return it from run
This commit is contained in:
parent
faa7f982af
commit
487e62a9ca
7 changed files with 260 additions and 227 deletions
181
hive-claude/src/telemetry.rs
Normal file
181
hive-claude/src/telemetry.rs
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
//! What the driver parses out of a turn's stream-json output.
|
||||
//!
|
||||
//! The lib is the single source of truth for reading claude's usage/model
|
||||
//! reporting. A turn's [`Telemetry`] is accumulated by the driver as events
|
||||
//! stream and handed back from [`crate::InfiniteSession::run`]; consumers that
|
||||
//! also want the raw events (for their own SSE / tool-call accounting) still
|
||||
//! get them through their [`crate::Sink`].
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
/// Token counts from one `usage` block. All in tokens; missing fields read `0`.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct TokenUsage {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cache_read_input_tokens: u64,
|
||||
pub cache_creation_input_tokens: u64,
|
||||
}
|
||||
|
||||
impl TokenUsage {
|
||||
/// Context footprint counting against the model window: input + both cache
|
||||
/// classes (not output).
|
||||
#[must_use]
|
||||
pub fn context_tokens(&self) -> u64 {
|
||||
self.input_tokens + self.cache_read_input_tokens + self.cache_creation_input_tokens
|
||||
}
|
||||
|
||||
fn from_obj(u: &Value) -> Self {
|
||||
let field = |k: &str| u.get(k).and_then(Value::as_u64).unwrap_or(0);
|
||||
Self {
|
||||
input_tokens: field("input_tokens"),
|
||||
output_tokens: field("output_tokens"),
|
||||
cache_read_input_tokens: field("cache_read_input_tokens"),
|
||||
cache_creation_input_tokens: field("cache_creation_input_tokens"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Everything the driver tracks from one turn's stream-json output.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct Telemetry {
|
||||
/// Most recent per-inference usage (from `assistant` events) — the live
|
||||
/// context footprint (the number to watch for compaction).
|
||||
pub context: TokenUsage,
|
||||
/// Cumulative usage across the turn (from the terminal `result` event) —
|
||||
/// the cost signal. Sums per-call prompts and can exceed the window.
|
||||
pub cost: TokenUsage,
|
||||
/// Model-reported active context window (`modelUsage.*.contextWindow` on
|
||||
/// the `result` event), if reported.
|
||||
pub context_window: Option<u64>,
|
||||
/// Resolved model id echoed by the API (`assistant.message.model`, e.g.
|
||||
/// `claude-opus-4-8`) — the concrete version, not the requested alias.
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
impl Telemetry {
|
||||
/// Fold one stream-json event into the running telemetry.
|
||||
pub(crate) fn observe(&mut self, event: &Value) {
|
||||
match event.get("type").and_then(Value::as_str) {
|
||||
Some("assistant") => {
|
||||
let Some(message) = event.get("message") else {
|
||||
return;
|
||||
};
|
||||
if let Some(usage) = message.get("usage") {
|
||||
self.context = TokenUsage::from_obj(usage);
|
||||
}
|
||||
if let Some(model) = message
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
self.model = Some(model.to_string());
|
||||
}
|
||||
}
|
||||
Some("result") => {
|
||||
if let Some(usage) = event.get("usage") {
|
||||
self.cost = TokenUsage::from_obj(usage);
|
||||
}
|
||||
if let Some(window) = context_window_from_result(event) {
|
||||
self.context_window = Some(window);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// The minimal signal a [`crate::CompactionPolicy`] needs.
|
||||
#[must_use]
|
||||
pub fn usage(&self) -> Usage {
|
||||
Usage {
|
||||
context_tokens: self.context.context_tokens(),
|
||||
context_window: self.context_window,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// First non-zero `contextWindow` across the `result` event's `modelUsage` map.
|
||||
fn context_window_from_result(event: &Value) -> Option<u64> {
|
||||
for (_model, stats) in event.get("modelUsage")?.as_object()? {
|
||||
if let Some(w) = stats.get("contextWindow").and_then(Value::as_u64)
|
||||
&& w > 0
|
||||
{
|
||||
return Some(w);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// The compaction signal: the live context size and the window it's measured
|
||||
/// against. Derived from [`Telemetry`] via [`Telemetry::usage`].
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct Usage {
|
||||
/// Tokens in the last inference's context. `0` until the first `assistant`
|
||||
/// event.
|
||||
pub context_tokens: u64,
|
||||
/// The model-reported active window, if the turn reported one.
|
||||
pub context_window: Option<u64>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::Telemetry;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn context_tokens_sum_excludes_output() {
|
||||
let mut t = Telemetry::default();
|
||||
t.observe(&json!({
|
||||
"type": "assistant",
|
||||
"message": { "model": "claude-opus-4-8", "usage": {
|
||||
"input_tokens": 100, "output_tokens": 999,
|
||||
"cache_read_input_tokens": 20, "cache_creation_input_tokens": 5,
|
||||
}}
|
||||
}));
|
||||
assert_eq!(t.context.context_tokens(), 125);
|
||||
assert_eq!(t.model.as_deref(), Some("claude-opus-4-8"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn last_assistant_wins() {
|
||||
let mut t = Telemetry::default();
|
||||
for n in [10, 20, 30] {
|
||||
t.observe(&json!({
|
||||
"type": "assistant",
|
||||
"message": { "usage": { "input_tokens": n } }
|
||||
}));
|
||||
}
|
||||
assert_eq!(t.context.input_tokens, 30);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn result_event_sets_cost_and_window() {
|
||||
let mut t = Telemetry::default();
|
||||
t.observe(&json!({
|
||||
"type": "result",
|
||||
"usage": { "input_tokens": 5_000, "output_tokens": 1_000 },
|
||||
"modelUsage": { "claude-opus-4-8": { "contextWindow": 200_000 } }
|
||||
}));
|
||||
assert_eq!(t.cost.input_tokens, 5_000);
|
||||
assert_eq!(t.context_window, Some(200_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_model_ignored_and_non_events_no_op() {
|
||||
let mut t = Telemetry::default();
|
||||
t.observe(&json!({ "type": "assistant", "message": { "model": "" } }));
|
||||
assert_eq!(t.model, None);
|
||||
t.observe(&json!({ "type": "system", "subtype": "init" }));
|
||||
assert_eq!(t, Telemetry::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_view_derives_from_context_and_window() {
|
||||
let mut t = Telemetry::default();
|
||||
t.observe(&json!({ "type": "assistant", "message": { "usage": { "input_tokens": 42 } } }));
|
||||
t.observe(&json!({ "type": "result", "modelUsage": { "m": { "contextWindow": 100 } } }));
|
||||
let u = t.usage();
|
||||
assert_eq!(u.context_tokens, 42);
|
||||
assert_eq!(u.context_window, Some(100));
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue