fix: report agent usage per provider round, not once per turn (#4545)

## The bug

buzz-agent emitted its `usage_update` notification in exactly one place:
after `ctx.run()` returned. Until that moment a turn's token counters
lived only in the prompt task's stack frame. **A turn killed mid-flight
reported nothing at all** — the provider had already billed every round
it completed, and no consumer ever saw any of it.

That is not a corner case for anything that ends a turn on a clock. It
is the normal case for a long-horizon benchmark run that relaunches its
agent between phases.

## How big

Measured against a provider's own billing ledger over one run's window:

| | provider ledger | what we recorded |
|---|---|---|
| the relaunched lead seat | $485 / 348M tok | $98.99 / 90.3M tok |
| the two seats that were not relaunched | $29.90 / 856M | $25.81 / 765M
— reconciles |

97% of that run's usage rows came back all zeros, against 1–4% for
comparable runs that never relaunch. In one 450-phase trial exactly 7
phases recorded any usage — and each of those carries 177k–437k input
tokens, a whole session's worth landing in the one phase that happened
to end gracefully.

Worth being precise about what was *not* wrong, since both were
plausible and both were checked:

- **Not pricing.** The rates were verified against the provider's
endpoints API and match what we charge.
- **Not a truncation bug.** The usage files were intact and internally
consistent. The tokens were never captured in the first place.

## The fix

The run loop now emits a session-cumulative `usage_update` after every
usage-bearing provider response, so an interrupted turn has reported
everything but its single in-flight request.

- **Emitting more than once per turn is already part of the contract.**
buzz-acp's `UsageTracker` advances its committed baseline only at
publish time, and goose behaves the same way — which is why the tracker
was written to tolerate it.
- **The turn-start session baseline is snapshotted into `RunCtx`** so
the mid-turn figure stays *session*-cumulative. A turn-local number
would be discarded by a high-water-mark consumer and lose the turn
entirely; there is a test for exactly that.
- **Snapshot by value, not a session handle.** The loop reports once per
round, and taking the sessions lock on each would serialise concurrent
sessions behind one another's provider round-trips. Nothing else
advances those counters while the turn holds `busy`, so it cannot go
stale.
- **One shared `wire::usage_update_payload`** for both call sites, so
the mid-turn and end-of-turn shapes cannot drift. A drift there would
present as tokens silently vanishing, which is the failure this
reporting exists to prevent.

## Why not a SIGTERM handler

That was the obvious shape and it does not work. At signal time the
counters are not sitting anywhere a handler could reach — they are in
the turn's stack frame, and the value the handler would need has not
been folded into the session yet. Making usage durable *during* the turn
is what actually fixes it; once it is, a handler adds nothing beyond the
in-flight request, whose cost is unknown until its response lands.

## Tests

- `usage_is_reported_after_each_round_not_only_at_turn_end` — two
rounds; asserts the **first** notification carries round 1's counts
alone, proving it went out before round 2 returned.
- `mid_turn_usage_includes_earlier_turns` — a mid-turn report must be
session-cumulative, not turn-local.

buzz-agent 18/18 on the `fake_llm` suite, 382 unit. `cargo fmt` /
`clippy` / `cargo check --workspace --all-targets` clean.

## Scope

Agent-side only, against `main`. The matching harness change — settling
usage on the timeout path, which was skipped on the reasoning that an
incomplete turn has nothing to flush — is **#4553**, against the
benchmark branch, since that harness does not exist on `main`.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: Atish Patel <atish@squareup.com>
Co-authored-by: Claude Code <noreply@anthropic.com>

crates/buzz-agent/src/agent.rs

@@ -13,8 +13,8 @@ use crate::mcp::McpRegistry;
 use crate::mcp::ResultBudget;
 
 use crate::types::{
-    AgentError, ContentBlock, HistoryItem, ProviderStop, StopReason, ToolCall, ToolResult,
-    ToolResultContent, TurnTotalState,
+    AgentError, ContentBlock, HistoryItem, ProviderStop, SessionUsageBaseline, StopReason,
+    ToolCall, ToolResult, ToolResultContent, TurnTotalState,
 };
 use crate::wire::{self, WireSender};
 
@@ -150,9 +150,40 @@ pub struct RunCtx<'a> {
     /// Reset to `Unseen` at turn start in `run()`. Callers must not derive a
     /// total by summing input+output — that is the UI display approximation only.
     pub turn_total_state: &'a mut TurnTotalState,
+    /// Session-cumulative counters as they stood when this turn began. Added to
+    /// the `turn_*` accumulators above to report a cumulative figure mid-turn;
+    /// the session's own copy is only advanced once, after the turn returns.
+    pub usage_baseline: SessionUsageBaseline,
 }
 
 impl RunCtx<'_> {
+    /// Send a session-cumulative `usage_update` reflecting everything observed
+    /// up to and including the most recent LLM response.
+    ///
+    /// The figure is the turn-start baseline plus this turn's running
+    /// accumulators, which is exactly what `session/prompt` will fold into the
+    /// session once the turn returns — so a mid-turn notification and the
+    /// end-of-turn one agree, and a turn that never returns has still reported
+    /// everything but its final in-flight request.
+    async fn emit_usage_update(&self) {
+        let base = self.usage_baseline;
+        let payload = wire::usage_update_payload(
+            base.input_tokens
+                .saturating_add(self.turn_input_tokens.unwrap_or(0)),
+            base.output_tokens
+                .saturating_add(self.turn_output_tokens.unwrap_or(0)),
+            base.cached_input_tokens
+                .saturating_add(self.turn_cached_input_tokens.unwrap_or(0)),
+            base.total_state.merge_session(*self.turn_total_state),
+            self.effective_model,
+        );
+        wire::send(
+            self.wire,
+            wire::goose_session_update(self.session_id, payload),
+        )
+        .await;
+    }
+
     pub async fn run(&mut self, prompt: Vec<ContentBlock>) -> Result<StopReason, AgentError> {
         let user_text = prompt_to_text(prompt)?;
         if user_text.len() > MAX_PROMPT_BYTES {
@@ -299,6 +330,23 @@ impl RunCtx<'_> {
             // this gate rather than representing absent categories as zero.
             if response.input_tokens.is_some() || response.output_tokens.is_some() {
                 *self.turn_total_state = self.turn_total_state.fold(response.total_tokens);
+                // Report what the turn has burned SO FAR, before running the
+                // next round. A turn is many provider round-trips over many
+                // minutes, and until this point the only report was the one
+                // `session/prompt` sends after the turn returns — so a turn
+                // that was cancelled, timed out, or whose process was killed
+                // reported nothing at all, and its tokens (already billed)
+                // existed only in this stack frame. Reporting per round bounds
+                // the loss to the single request in flight.
+                //
+                // Emitting more than one `usage_update` per turn is expected by
+                // the consumer: buzz-acp's UsageTracker advances its committed
+                // baseline only when the turn's metric is published, so every
+                // notification within a turn measures from the same frozen
+                // baseline and the last one seen is the turn's true total.
+                // goose behaves the same way, which is why the tracker was
+                // written to tolerate it.
+                self.emit_usage_update().await;
             }
 
             if !response.reasoning.is_empty() {

crates/buzz-agent/src/lib.rs

@@ -658,6 +658,7 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
         effective_model_override,
         run_id,
         mut steer_rx,
+        usage_baseline,
     ) = match acquire_session(&app, &p.session_id).await {
         Ok(v) => v,
         Err(reason) => {
@@ -709,6 +710,7 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
         turn_output_tokens: &mut turn_output_tokens,
         turn_cached_input_tokens: &mut turn_cached_input_tokens,
         turn_total_state: &mut turn_total_state,
+        usage_baseline,
     };
     let result = ctx.run(p.prompt).await;
     if let Some(s) = app.sessions.lock().await.get_mut(&sid) {
@@ -766,28 +768,16 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
         if let Some((accumulated_in, accumulated_out, accumulated_cached, accumulated_total)) =
             accumulated
         {
-            // Build the usage_update payload. `accumulatedTotalTokens` is only
-            // included when the cumulative is exactly known — never when Unseen
-            // (no total ever observed) or Unknown (at least one turn lacked a
-            // total). A goose consumer that doesn't recognise the field ignores it.
-            let mut update = serde_json::json!({
-                "sessionUpdate": "usage_update",
-                // used: total tokens as a context-usage proxy;
-                // contextLimit: 0 (buzz-agent has no context limit tracking).
-                "used": accumulated_in.saturating_add(accumulated_out),
-                "contextLimit": 0u64,
-                "accumulatedInputTokens": accumulated_in,
-                "accumulatedOutputTokens": accumulated_out,
-                // A subset of accumulatedInputTokens, not an addition to
-                // it. Extends goose's usage_update shape; a consumer that
-                // does not know the field ignores it and prices exactly as
-                // it did before.
-                "accumulatedCachedInputTokens": accumulated_cached,
-                "model": effective_model_str,
-            });
-            if let crate::types::TurnTotalState::Exact(total) = accumulated_total {
-                update["accumulatedTotalTokens"] = serde_json::json!(total);
-            }
+            // Same builder the run loop uses for its per-round reports, so the
+            // final notification is shape-identical to the ones that preceded
+            // it and a consumer taking the high-water mark lands on this one.
+            let update = wire::usage_update_payload(
+                accumulated_in,
+                accumulated_out,
+                accumulated_cached,
+                accumulated_total,
+                effective_model_str,
+            );
             wire::send(&wire_tx, goose_session_update(&sid, update)).await;
         }
     }
@@ -821,6 +811,7 @@ async fn acquire_session(
         Option<String>,
         String,
         mpsc::UnboundedReceiver<Vec<ContentBlock>>,
+        crate::types::SessionUsageBaseline,
     ),
     &'static str,
 > {
@@ -857,6 +848,17 @@ async fn acquire_session(
         effective_model,
         run_id,
         steer_rx,
+        // Snapshot rather than a handle: the run loop reports cumulative usage
+        // after every LLM round, and taking the sessions lock on each of those
+        // would serialise concurrent sessions behind one another's provider
+        // round-trips. Nothing else advances these counters while this turn
+        // holds `busy`, so the snapshot cannot go stale under it.
+        crate::types::SessionUsageBaseline {
+            input_tokens: s.accumulated_input_tokens,
+            output_tokens: s.accumulated_output_tokens,
+            cached_input_tokens: s.accumulated_cached_input_tokens,
+            total_state: s.accumulated_total_state,
+        },
     ))
 }
 

crates/buzz-agent/src/types.rs

@@ -308,6 +308,30 @@ impl TurnTotalState {
     }
 }
 
+/// The session-cumulative usage counters as of the START of a turn.
+///
+/// Copied out of the session under the lock when a turn begins and handed to
+/// `RunCtx` by value, so the run loop can emit a cumulative `usage_update`
+/// after every LLM round without reaching back into `App.sessions` (which it
+/// holds no handle to, and which is locked by the turn's own bookkeeping at
+/// both ends).
+///
+/// This exists so that usage is durable *during* a turn rather than only after
+/// it. The counters a turn accrues live in the prompt task's stack frame until
+/// the turn returns; a process killed mid-turn takes them with it and the
+/// tokens are billed by the provider but recorded nowhere. That is not
+/// hypothetical — it silently under-reported a long-horizon benchmark's cost by
+/// several-fold, because every phase of a `continue_until_timeout` run is
+/// terminated mid-turn by design.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct SessionUsageBaseline {
+    pub input_tokens: u64,
+    pub output_tokens: u64,
+    /// The cache-served subset of `input_tokens`, not an addition to it.
+    pub cached_input_tokens: u64,
+    pub total_state: TurnTotalState,
+}
+
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum StopReason {
     EndTurn,

crates/buzz-agent/src/wire.rs

@@ -148,6 +148,48 @@ pub fn goose_session_update(sid: &str, update: Value) -> Value {
     })
 }
 
+/// Build the `usage_update` payload for a `_goose/unstable/session/update`.
+///
+/// Shared by the two places that report usage — after each LLM round inside a
+/// turn, and once more when the turn completes — so the wire shape cannot drift
+/// between them. A consumer takes the high-water mark per session, so the
+/// mid-turn payloads are supersets of each other and the final one wins; a
+/// divergence in field names or units between the two call sites would instead
+/// show up as tokens silently vanishing, which is the failure this reporting
+/// exists to prevent.
+///
+/// All counts are SESSION-cumulative, matching goose, so buzz-acp's
+/// `UsageTracker` can compute per-turn deltas symmetrically for both agents.
+pub fn usage_update_payload(
+    accumulated_input_tokens: u64,
+    accumulated_output_tokens: u64,
+    accumulated_cached_input_tokens: u64,
+    accumulated_total: crate::types::TurnTotalState,
+    model: &str,
+) -> Value {
+    let mut update = json!({
+        "sessionUpdate": "usage_update",
+        // used: total tokens as a context-usage proxy;
+        // contextLimit: 0 (buzz-agent has no context limit tracking).
+        "used": accumulated_input_tokens.saturating_add(accumulated_output_tokens),
+        "contextLimit": 0u64,
+        "accumulatedInputTokens": accumulated_input_tokens,
+        "accumulatedOutputTokens": accumulated_output_tokens,
+        // A subset of accumulatedInputTokens, not an addition to it. Extends
+        // goose's usage_update shape; a consumer that does not know the field
+        // ignores it and prices exactly as it did before.
+        "accumulatedCachedInputTokens": accumulated_cached_input_tokens,
+        "model": model,
+    });
+    // Only when the cumulative is exactly known — never when Unseen (no total
+    // ever observed) or Unknown (at least one turn lacked a total). A goose
+    // consumer that doesn't recognise the field ignores it.
+    if let Some(total) = accumulated_total.exact_value() {
+        update["accumulatedTotalTokens"] = json!(total);
+    }
+    update
+}
+
 /// A `session/update` notification carrying a `update._meta.goose.<key>` field.
 /// Used to advertise `activeRunId` (so steer-capable clients can target the
 /// in-flight run) and `queuedSteer` (so they can correlate an accepted steer

crates/buzz-agent/tests/fake_llm.rs

@@ -933,6 +933,135 @@ async fn no_usage_turn_emits_no_usage_notification() {
     h.shutdown().await;
 }
 
+/// Usage must be reported after EVERY provider round, not only once the turn
+/// returns.
+///
+/// A turn is many provider round-trips over many minutes. While the only report
+/// was the one `session/prompt` sends after the turn returns, a turn whose
+/// process was killed mid-flight reported nothing at all: its counters lived in
+/// the prompt task's stack frame, the provider had already billed them, and no
+/// consumer ever saw them. That is not a corner case for a long-horizon
+/// benchmark — every phase of a `continue_until_timeout` run is terminated
+/// mid-turn by design, which under-reported one measured run's cost several-fold.
+///
+/// Two rounds with distinct usage. The assertion that matters is the FIRST
+/// notification: it must carry round 1's counts alone, proving it was sent
+/// before round 2 had returned, so a kill between the rounds would still have
+/// left round 1 on the wire.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn usage_is_reported_after_each_round_not_only_at_turn_end() {
+    let url = spawn_fake_llm(vec![
+        openai_tool_call_with_usage("call_round1", "fake__noop", json!({}), 15, 6),
+        openai_text_with_usage("done", 20, 8),
+    ])
+    .await;
+    let mut h = Harness::spawn(&url).await;
+    let sid = init_session(&mut h).await;
+
+    let p_id = h
+        .send(
+            "session/prompt",
+            json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}),
+        )
+        .await;
+
+    let (frames_before, response) = recv_until_with_drain(&mut h, |v| v["id"] == p_id).await;
+    assert_eq!(
+        response["result"]["stopReason"], "end_turn",
+        "turn must complete with end_turn"
+    );
+
+    let usage: Vec<&Value> = frames_before
+        .iter()
+        .filter(|v| is_usage_update(v))
+        .collect();
+    assert!(
+        usage.len() >= 2,
+        "expected a usage_update per round (2 rounds), got {}; frames: {frames_before:#?}",
+        usage.len()
+    );
+
+    // Round 1 alone — emitted while round 2 was still outstanding.
+    assert_eq!(
+        usage[0]["params"]["update"]["accumulatedInputTokens"],
+        json!(15u64),
+        "first notification must carry round 1's input tokens only"
+    );
+    assert_eq!(
+        usage[0]["params"]["update"]["accumulatedOutputTokens"],
+        json!(6u64),
+        "first notification must carry round 1's output tokens only"
+    );
+
+    // The last one is the turn total and is what a high-water-mark consumer keeps.
+    let last = usage[usage.len() - 1];
+    assert_eq!(
+        last["params"]["update"]["accumulatedInputTokens"],
+        json!(35u64),
+        "final notification must carry the turn total 15+20=35"
+    );
+    assert_eq!(
+        last["params"]["update"]["accumulatedOutputTokens"],
+        json!(14u64),
+        "final notification must carry the turn total 6+8=14"
+    );
+
+    h.shutdown().await;
+}
+
+/// A mid-turn report must be SESSION-cumulative, not turn-local.
+///
+/// The baseline handed to the run loop is a snapshot taken when the turn began;
+/// if it were dropped, a consumer taking the high-water mark per session would
+/// see turn 2's first round (a small number) arrive after turn 1's total and
+/// discard it, silently losing turn 2 for any turn that never completed.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn mid_turn_usage_includes_earlier_turns() {
+    let url = spawn_fake_llm(vec![
+        openai_text_with_usage("turn one", 10, 5),
+        openai_tool_call_with_usage("call_t2", "fake__noop", json!({}), 20, 8),
+        openai_text_with_usage("turn two done", 30, 9),
+    ])
+    .await;
+    let mut h = Harness::spawn(&url).await;
+    let sid = init_session(&mut h).await;
+
+    let p1 = h
+        .send(
+            "session/prompt",
+            json!({"sessionId": sid, "prompt": [{"type":"text","text":"turn 1"}]}),
+        )
+        .await;
+    let (_, _) = recv_until_with_drain(&mut h, |v| v["id"] == p1).await;
+
+    let p2 = h
+        .send(
+            "session/prompt",
+            json!({"sessionId": sid, "prompt": [{"type":"text","text":"turn 2"}]}),
+        )
+        .await;
+    let (frames_before, _) = recv_until_with_drain(&mut h, |v| v["id"] == p2).await;
+
+    let first = frames_before
+        .iter()
+        .find(|v| is_usage_update(v))
+        .unwrap_or_else(|| {
+            panic!("expected a usage_update during turn 2; frames: {frames_before:#?}")
+        });
+    assert_eq!(
+        first["params"]["update"]["accumulatedInputTokens"],
+        json!(30u64),
+        "turn 2 round 1 must report 10 (turn 1) + 20 (this round), not 20"
+    );
+    assert_eq!(
+        first["params"]["update"]["accumulatedOutputTokens"],
+        json!(13u64),
+        "turn 2 round 1 must report 5 (turn 1) + 8 (this round), not 8"
+    );
+
+    h.shutdown().await;
+}
+
 /// When a turn is cancelled AFTER the provider has already returned a response
 /// (so token counts are observed), buzz-agent must still emit the usage
 /// notification before the cancelled `session/prompt` response.