Fix token counters: post-finish usage frame, rate-limit wipe, session fidelity, estimate disclosure

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,
         }

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()),
         }

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);

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<S>(
     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,15 +1076,53 @@ 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}}"#,
         ]);
         let mut sink = |_: StreamEvent| {};
         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)]

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<RateLimitState>,
+    /// 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);
+    }
+}

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(),