diff --git a/crates/hi-ai/src/fallback.rs b/crates/hi-ai/src/fallback.rs index 5474db78..9f760d6b 100644 --- a/crates/hi-ai/src/fallback.rs +++ b/crates/hi-ai/src/fallback.rs @@ -129,11 +129,8 @@ mod tests { usage: Usage { input_tokens: input, output_tokens: output, - cache_read_tokens: 0, - cache_creation_tokens: 0, - input_includes_cache: false, context_occupancy: input, - rate_limits: None, + ..Usage::default() }, stop_reason: None, } @@ -145,11 +142,8 @@ mod tests { usage: Usage { input_tokens: input, output_tokens: output, - cache_read_tokens: 0, - cache_creation_tokens: 0, - input_includes_cache: false, context_occupancy: input, - rate_limits: None, + ..Usage::default() }, stop_reason: None, } diff --git a/crates/hi-ai/src/moa.rs b/crates/hi-ai/src/moa.rs index f2feeca6..02bf593f 100644 --- a/crates/hi-ai/src/moa.rs +++ b/crates/hi-ai/src/moa.rs @@ -430,6 +430,7 @@ fn add_reference_usage(aggregate: &mut Usage, reference: Usage) { aggregate.context_occupancy = aggregate_context; aggregate.input_includes_cache = aggregate_includes_cache; aggregate.rate_limits = aggregate_rate_limits.or(reference_rate_limits); + aggregate.estimated |= reference.estimated; } #[cfg(test)] @@ -519,11 +520,8 @@ mod tests { usage: Usage { input_tokens: input, output_tokens: output, - cache_read_tokens: 0, - cache_creation_tokens: 0, - input_includes_cache: false, context_occupancy: input, - rate_limits: None, + ..Usage::default() }, stop_reason: Some("stop".into()), } diff --git a/crates/hi-ai/src/openai.rs b/crates/hi-ai/src/openai.rs index 88415a52..b25b76ea 100644 --- a/crates/hi-ai/src/openai.rs +++ b/crates/hi-ai/src/openai.rs @@ -177,6 +177,7 @@ impl Provider for OpenAiProvider { input_includes_cache: true, context_occupancy: estimated_input_tokens, rate_limits, + estimated: true, }) })?; stream::backfill_missing_usage(&mut completion, &request); diff --git a/crates/hi-ai/src/openai/stream.rs b/crates/hi-ai/src/openai/stream.rs index e8dbe842..2f1ef20a 100644 --- a/crates/hi-ai/src/openai/stream.rs +++ b/crates/hi-ai/src/openai/stream.rs @@ -17,6 +17,15 @@ use crate::types::{Completion, Content, StreamEvent, Usage, estimate_messages_to const SPECIAL_TOKEN_PREFIX: &str = "<|"; const SPECIAL_TOKEN_SUFFIX: &str = "|>"; +/// How long to keep reading after `finish_reason` for the trailing usage frame. +/// With `stream_options.include_usage`, spec-compliant providers send usage in +/// a separate chunk AFTER the finish chunk (empty `choices`) — breaking on +/// finish discarded it, leaving streamed calls with zero provider usage and the +/// token counters running on chars/4 estimates. Bounded so a provider that +/// holds the connection open after finish (no `[DONE]`, no close) still can't +/// leave a completed answer spinning — the guarantee break-on-finish gave. +const POST_FINISH_USAGE_GRACE: std::time::Duration = std::time::Duration::from_secs(2); + /// Check if `inner` (the text between `<|` and `|>`) looks like a special /// token identifier: non-empty, alphanumeric + underscore, no spaces or /// newlines. Real tokens: `im_start`, `im_end`, `endoftext`, `system`, etc. @@ -599,9 +608,11 @@ fn strip_text_tool_protocol_artifact(text: &str) -> String { /// [`Completion`], forwarding text/reasoning/tool tokens to `sink`. /// /// The collector reads until the provider sends `[DONE]`, closes the stream, -/// returns a stream error, or reports a `finish_reason`. A `finish_reason` is -/// treated as completion so providers that omit `[DONE]` cannot leave a -/// completed answer spinning forever. +/// or returns a stream error. After a `finish_reason` it keeps reading — but +/// only within [`POST_FINISH_USAGE_GRACE`] — for the trailing usage-only chunk +/// that `stream_options.include_usage` providers send after the finish chunk, +/// so providers that omit `[DONE]` still cannot leave a completed answer +/// spinning forever. pub(crate) async fn collect_completion( mut stream: S, sink: &mut (dyn FnMut(StreamEvent) + Send), @@ -625,8 +636,22 @@ where // borrow-conflict chaining of two `&mut dyn FnMut` filters. let mut filter = StreamingTextFilter::new(sink); + let mut usage_seen = false; + loop { - let data = match stream.next().await { + // After the finish chunk, only wait a bounded grace for the trailing + // usage frame; before it, block on the stream as usual. + let next = if stream_complete { + match tokio::time::timeout(POST_FINISH_USAGE_GRACE, stream.next()).await { + Ok(item) => item, + // Provider holds the stream open after finish without sending + // the usage frame: give up on it, keep the completed answer. + Err(_) => break, + } + } else { + stream.next().await + }; + let data = match next { Some(Ok(data)) => data, // A stream *read* error mid-flight: the provider reset the connection // or sent a truncated frame instead of closing cleanly / sending @@ -658,6 +683,7 @@ where } if let Some(usage) = chunk.usage { + usage_seen = true; let cached = usage .prompt_tokens_details .map(|d| d.cached_tokens) @@ -673,6 +699,7 @@ where input_includes_cache: true, context_occupancy: usage.prompt_tokens, rate_limits: None, + estimated: false, }; } for choice in chunk.choices { @@ -721,7 +748,11 @@ where stream_complete = true; } } - if stream_complete { + // Once both the finish marker and a usage frame have arrived there is + // nothing left to wait for (providers that attach usage to the finish + // chunk break here immediately, as before). With finish but no usage + // yet, loop back into the bounded drain above for the trailing frame. + if stream_complete && usage_seen { break; } } @@ -744,6 +775,7 @@ where } if completion.usage.output_tokens == 0 && output_chars > 0 { completion.usage.output_tokens = output_chars.div_ceil(4) as u64; + completion.usage.estimated = true; } Ok(completion) } @@ -827,10 +859,18 @@ pub(crate) fn backfill_missing_usage( ) { if completion.usage.input_tokens == 0 { completion.usage.input_tokens = estimate_messages_tokens(&request.messages); + completion.usage.estimated = true; } if completion.usage.output_tokens == 0 { completion.usage.output_tokens = crate::types::estimate_completion_output_tokens(&completion.content); + completion.usage.estimated = true; + } + // Keep the occupancy gauge alive on the estimate path too — it previously + // stayed 0 here, leaving consumers of `context_occupancy` blind whenever + // the provider reported no usage. + if completion.usage.context_occupancy == 0 { + completion.usage.context_occupancy = completion.usage.input_tokens; } } @@ -1017,8 +1057,9 @@ mod tests { #[tokio::test(start_paused = true)] async fn stops_after_finish_reason_without_done() { // Some providers send `finish_reason` and then neither `[DONE]` nor a - // socket close. Treating the finish marker as terminal returns promptly - // without any timer. + // socket close. The finish marker is still terminal — after the bounded + // post-finish usage grace, the completed answer is returned rather than + // spinning forever. let stream = never_ending(vec![ r#"{"choices":[{"delta":{"content":"the answer"},"finish_reason":null}]}"#, r#"{"choices":[{"delta":{},"finish_reason":"stop"}]}"#, @@ -1035,8 +1076,8 @@ mod tests { #[tokio::test(start_paused = true)] async fn usage_on_finish_chunk_is_captured() { - // Usage included on the same chunk as `finish_reason` is captured without - // requiring a trailing grace timer. + // Usage included on the same chunk as `finish_reason` is captured and + // terminates collection immediately — no trailing grace timer runs. let stream = never_ending(vec![ r#"{"choices":[{"delta":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":2}}"#, ]); @@ -1044,6 +1085,44 @@ mod tests { let completion = collect_completion(stream, &mut sink).await.unwrap(); assert_eq!(completion.usage.input_tokens, 10); assert_eq!(completion.usage.output_tokens, 2); + assert!(!completion.usage.estimated); + } + + #[tokio::test(start_paused = true)] + async fn usage_in_post_finish_chunk_is_captured() { + // The spec shape for `stream_options.include_usage`: usage arrives in a + // separate chunk AFTER the finish chunk, with empty `choices`. Breaking + // on finish used to discard it, leaving streamed calls with zero + // provider usage (and the counters running on chars/4 estimates). + let stream = never_ending(vec![ + r#"{"choices":[{"delta":{"content":"the answer"},"finish_reason":null}]}"#, + r#"{"choices":[{"delta":{},"finish_reason":"stop"}]}"#, + r#"{"choices":[],"usage":{"prompt_tokens":1234,"completion_tokens":56}}"#, + ]); + let mut sink = |_: StreamEvent| {}; + let completion = collect_completion(stream, &mut sink).await.unwrap(); + assert_eq!(completion.stop_reason.as_deref(), Some("stop")); + assert_eq!(completion.usage.input_tokens, 1234); + assert_eq!(completion.usage.output_tokens, 56); + assert_eq!(completion.usage.context_occupancy, 1234); + assert!(!completion.usage.estimated); + } + + #[tokio::test(start_paused = true)] + async fn all_zero_usage_frame_marks_output_estimate() { + // Observed on pipenetwork: the post-finish usage frame exists but is all + // zeros. The in-stream chars/4 output fallback kicks in and the usage is + // marked estimated, so downstream consumers know the numbers are guesses. + let stream = never_ending(vec![ + r#"{"choices":[{"delta":{"content":"four ch"},"finish_reason":null}]}"#, + r#"{"choices":[{"delta":{},"finish_reason":"stop"}]}"#, + r#"{"choices":[],"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}}"#, + ]); + let mut sink = |_: StreamEvent| {}; + let completion = collect_completion(stream, &mut sink).await.unwrap(); + assert_eq!(completion.usage.input_tokens, 0, "input backfills upstream"); + assert!(completion.usage.output_tokens > 0, "chars/4 fallback"); + assert!(completion.usage.estimated); } #[tokio::test(start_paused = true)] diff --git a/crates/hi-ai/src/types.rs b/crates/hi-ai/src/types.rs index f101d004..91de61f3 100644 --- a/crates/hi-ai/src/types.rs +++ b/crates/hi-ai/src/types.rs @@ -284,6 +284,13 @@ pub struct Usage { /// usage so frontends can show whether failures are route/provider throttles. #[serde(default, skip_serializing_if = "Option::is_none")] pub rate_limits: Option, + /// True when a token field was backfilled from a chars/4 estimate rather + /// than provider-reported usage (the provider sent no usage frame, or an + /// all-zeros one). Sticky across [`Usage::add`], so session totals disclose + /// that they contain guessed numbers — surfaced as `usage_estimated` in + /// `--report`. + #[serde(default)] + pub estimated: bool, } impl Usage { @@ -303,7 +310,13 @@ impl Usage { self.output_tokens += other.output_tokens; self.cache_read_tokens += other.cache_read_tokens; self.cache_creation_tokens += other.cache_creation_tokens; - self.rate_limits = other.rate_limits; + // "Latest observed": a booking that carries no rate-limit snapshot + // (side-calls, error usage, estimates) must not wipe the last real one — + // that made the rate-limit display blank out mid-session. + if other.rate_limits.is_some() { + self.rate_limits = other.rate_limits; + } + self.estimated |= other.estimated; } /// Deprecated: prefer [`Usage::context_occupancy`], which is set by the @@ -388,3 +401,60 @@ pub struct ToolCall<'a> { pub name: &'a str, pub arguments: &'a str, } + +#[cfg(test)] +mod tests { + use super::{RateLimitBucket, RateLimitState, Usage}; + + #[test] + fn add_preserves_last_observed_rate_limits_and_sticks_estimated() { + let mut totals = Usage { + input_tokens: 100, + output_tokens: 10, + rate_limits: Some(RateLimitState { + requests_min: RateLimitBucket { + limit: 10, + remaining: 8, + reset_seconds: 1, + }, + ..RateLimitState::default() + }), + ..Usage::default() + }; + + // A booking with no rate-limit snapshot (side-call, error usage, + // estimate) must not wipe the last observed one. + totals.add(Usage { + input_tokens: 50, + output_tokens: 5, + estimated: true, + ..Usage::default() + }); + assert_eq!(totals.input_tokens, 150); + assert!( + totals.rate_limits.is_some(), + "zero-snapshot add wiped rate limits" + ); + // Estimated is sticky: once any component was guessed, totals say so. + assert!(totals.estimated); + totals.add(Usage { + input_tokens: 5, + ..Usage::default() + }); + assert!(totals.estimated, "estimated must not reset"); + + // A booking that carries a fresh snapshot replaces the old one. + totals.add(Usage { + rate_limits: Some(RateLimitState { + requests_min: RateLimitBucket { + limit: 10, + remaining: 3, + reset_seconds: 2, + }, + ..RateLimitState::default() + }), + ..Usage::default() + }); + assert_eq!(totals.rate_limits.unwrap().requests_min.remaining, 3); + } +} diff --git a/crates/hi-cli/src/main.rs b/crates/hi-cli/src/main.rs index 8236ac87..bfd9954b 100644 --- a/crates/hi-cli/src/main.rs +++ b/crates/hi-cli/src/main.rs @@ -1113,6 +1113,10 @@ fn write_report( "usage_scope": "session_cumulative_full_context", "input_token_scope": "full_request_context_not_user_prompt", "output_token_scope": "generated_output", + // True when any counted call fell back to a chars/4 estimate because + // the provider reported no (or all-zeros) usage — the numbers above are + // then approximate, not provider-metered. + "usage_estimated": totals.estimated, "session_input_tokens": totals.input_tokens, "session_output_tokens": totals.output_tokens, "session_total_tokens": totals.total(), diff --git a/crates/hi-cli/src/session.rs b/crates/hi-cli/src/session.rs index 204e61f4..9bca5509 100644 --- a/crates/hi-cli/src/session.rs +++ b/crates/hi-cli/src/session.rs @@ -20,6 +20,15 @@ enum SessionMeta { Usage { input_tokens: u64, output_tokens: u64, + /// Cache counters and the estimated marker ride along so a resumed + /// session's totals keep full fidelity. `#[serde(default)]` so session + /// files written before these fields load as zero/false. + #[serde(default)] + cache_read_tokens: u64, + #[serde(default)] + cache_creation_tokens: u64, + #[serde(default)] + estimated: bool, }, Checkpoints { refs: Vec, @@ -113,6 +122,9 @@ impl SessionSink for JsonlSession { let line = serde_json::to_string(&SessionMeta::Usage { input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, + cache_read_tokens: usage.cache_read_tokens, + cache_creation_tokens: usage.cache_creation_tokens, + estimated: usage.estimated, })?; writeln!(writer, "{line}")?; writer.flush()?; @@ -525,15 +537,19 @@ pub fn load_history(path: &Path) -> Result { SessionMeta::Usage { input_tokens, output_tokens, + cache_read_tokens, + cache_creation_tokens, + estimated, } => { usage = Usage { input_tokens, output_tokens, - cache_read_tokens: 0, - cache_creation_tokens: 0, + cache_read_tokens, + cache_creation_tokens, input_includes_cache: false, context_occupancy: input_tokens, rate_limits: None, + estimated, }; } SessionMeta::Checkpoints { refs } => { @@ -786,11 +802,8 @@ mod tests { Usage { input_tokens: 123, output_tokens: 45, - cache_read_tokens: 0, - cache_creation_tokens: 0, - input_includes_cache: false, context_occupancy: 123, - rate_limits: None, + ..Usage::default() }, ) .unwrap(); @@ -958,11 +971,8 @@ mod tests { Usage { input_tokens: 1, output_tokens: 1, - cache_read_tokens: 0, - cache_creation_tokens: 0, - input_includes_cache: false, context_occupancy: 1, - rate_limits: None, + ..Usage::default() }, ) .unwrap(); @@ -996,6 +1006,45 @@ mod tests { ); } + #[test] + fn usage_round_trips_cache_tokens_and_estimated_marker() { + // Session totals must keep full fidelity across resume: cache counters + // and the estimated marker used to be dropped (only input/output were + // persisted), silently shrinking a resumed session's numbers. + let mut path = std::env::temp_dir(); + path.push(format!( + "hi-session-usage-fidelity-{}-{}.jsonl", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let mut session = JsonlSession::new(path.clone()); + session + .record( + &[Message::user("go")], + Usage { + input_tokens: 100, + output_tokens: 20, + cache_read_tokens: 60, + cache_creation_tokens: 7, + estimated: true, + ..Usage::default() + }, + ) + .unwrap(); + + let loaded = load_history(&path).unwrap(); + let _ = std::fs::remove_file(&path); + + assert_eq!(loaded.usage.input_tokens, 100); + assert_eq!(loaded.usage.output_tokens, 20); + assert_eq!(loaded.usage.cache_read_tokens, 60); + assert_eq!(loaded.usage.cache_creation_tokens, 7); + assert!(loaded.usage.estimated, "estimated marker survives resume"); + } + #[test] fn load_history_skips_corrupted_lines() { // A partially-written last line (from a crash mid-flush) must not make