From f660ad2c640de2e9ff51405770952985e179c029 Mon Sep 17 00:00:00 2001 From: juan <2930882+juacker@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:23:25 +0200 Subject: [PATCH] fix(assistant): accumulate run token usage across turns The engine overwrote the run's usage on every provider usage event, so a multi-turn run persisted only its final turn. In a tool-using run that is the shortest turn of all, which made the reported token cost of a run arbitrary and far too low: a delegated worker run in this repo's own history logged 11 tool calls but reported 227 input tokens. Track usage per turn instead. Reports inside one turn restate that turn's running counters rather than adding to them (Anthropic sends the prompt count in `message_start` and the cumulative output in `message_delta`), so they collapse field-wise; finished turns are then summed into the run total. `usage` keeps holding that derived total, so the cancel and failure paths continue to report the tokens spent by earlier turns. Also derive `total_tokens` when the provider omits it. Anthropic never sends one and several OpenAI-compatible endpoints (MiniMax among them) do not either, which left runs with input and output counts but no total -- 93 of 93 such runs in this repo's history. The derived total excludes reasoning tokens because both native providers already bill those inside the output count. --- src-tauri/src/assistant/engine.rs | 39 ++++++- src-tauri/src/assistant/types.rs | 180 +++++++++++++++++++++++++++++- 2 files changed, 217 insertions(+), 2 deletions(-) diff --git a/src-tauri/src/assistant/engine.rs b/src-tauri/src/assistant/engine.rs index 3c82ec37..c49d92c9 100644 --- a/src-tauri/src/assistant/engine.rs +++ b/src-tauri/src/assistant/engine.rs @@ -168,7 +168,12 @@ pub async fn run_session_turn( let session_allowed_command_prefixes = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); let session_blocked_command_prefixes = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + // `usage` is the run total reported to every terminal path (complete, + // cancel, fail). `completed_usage` holds only the turns that have finished, + // so the in-flight turn can be re-folded as it restates its counters + // without double-counting earlier turns. let mut usage: Option = None; + let mut completed_usage: Option = None; // === Tool execution loop === // Build system prompt (prepended to every API call, not persisted). @@ -422,6 +427,9 @@ pub async fn run_session_turn( // is tracked separately because the execution loop below needs it. let mut content_parts: Vec = Vec::new(); let mut tool_calls: Vec = Vec::new(); + // Usage reported by this turn so far. Reset every iteration because a + // provider request bills its own prompt and completion. + let mut turn_usage: Option = None; loop { match tokio::select! { @@ -473,7 +481,24 @@ pub async fn run_session_turn( set_thinking_signature(&mut content_parts, signature); } ProviderEvent::Usage { usage: u } => { - usage = Some(u); + // A turn may report usage several times, each report + // restating that turn's running counters rather than + // adding to them, so collapse them field-wise. The run + // total is the finished turns plus this turn's latest + // report; keeping `usage` in that derived form means the + // cancel and failure paths below already see the tokens + // spent by every earlier turn. + let turn = turn_usage.get_or_insert_with(RunUsage::default); + turn.merge_turn_report(&u); + // Derive the total on a copy: a later report in this + // same turn may still raise input or output, and a total + // derived from a partial report must not survive the + // field-wise merge above as if the provider had sent it. + let mut turn_so_far = turn.clone(); + turn_so_far.ensure_total(); + let mut total = completed_usage.clone().unwrap_or_default(); + total.add_turn(&turn_so_far); + usage = Some(total); } ProviderEvent::ToolCallReady { tool_call } => { // Record it in content (in position) and in tool_calls @@ -533,6 +558,18 @@ pub async fn run_session_turn( } } + // This turn is over: fold its final report into the completed-turn + // total so the next turn accumulates on top of it. Before this the run + // kept only the last turn's usage, which in a tool-using run is the + // shortest turn of all. + if let Some(mut turn) = turn_usage.take() { + turn.ensure_total(); + let mut total = completed_usage.take().unwrap_or_default(); + total.add_turn(&turn); + completed_usage = Some(total); + usage = completed_usage.clone(); + } + // Finalize the assistant message from whatever we accumulated, even if // the provider never emitted MessageComplete. This prevents the orphan- // tool case: tool_calls captured via `finish_reason: tool_calls` but diff --git a/src-tauri/src/assistant/types.rs b/src-tauri/src/assistant/types.rs index 4deafd7c..89b575ea 100644 --- a/src-tauri/src/assistant/types.rs +++ b/src-tauri/src/assistant/types.rs @@ -294,7 +294,7 @@ pub struct AssistantMessagePage { pub total_count: u32, } -#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)] #[serde(rename_all = "camelCase")] #[ts(export, export_to = "bindings.ts")] pub struct RunUsage { @@ -308,6 +308,76 @@ pub struct RunUsage { pub total_tokens: Option, } +impl RunUsage { + /// Collapse one more usage report belonging to the *same* provider turn. + /// + /// A single turn can report usage more than once: Anthropic sends + /// `input_tokens` in `message_start` and then a cumulative `output_tokens` + /// in `message_delta`, and `message_start` already carries a small + /// `output_tokens` preview. Each report restates the running total for the + /// fields it carries rather than adding to it, so the correct fold is a + /// field-wise maximum: summing would double-count that preview. + pub fn merge_turn_report(&mut self, next: &RunUsage) { + self.input_tokens = max_optional(self.input_tokens, next.input_tokens); + self.output_tokens = max_optional(self.output_tokens, next.output_tokens); + self.reasoning_tokens = max_optional(self.reasoning_tokens, next.reasoning_tokens); + self.total_tokens = max_optional(self.total_tokens, next.total_tokens); + } + + /// Fold a finished turn into the run total. + /// + /// An agentic run issues one provider request per turn, and each request + /// bills its own prompt and completion independently, so the run total is + /// the sum over turns. Without this the persisted usage would describe only + /// the final turn, which is the shortest one in a typical tool-using run. + pub fn add_turn(&mut self, turn: &RunUsage) { + self.input_tokens = add_optional(self.input_tokens, turn.input_tokens); + self.output_tokens = add_optional(self.output_tokens, turn.output_tokens); + self.reasoning_tokens = add_optional(self.reasoning_tokens, turn.reasoning_tokens); + self.total_tokens = add_optional(self.total_tokens, turn.total_tokens); + } + + /// Fill in `total_tokens` from the components when the provider never sent + /// one. Anthropic never sends a total, and several OpenAI-compatible + /// endpoints (MiniMax among them) omit it too, which previously left the + /// run with input and output counts but no total at all. + /// + /// Reasoning tokens are deliberately excluded: both native providers bill + /// them inside the completion/output count, so adding them again would + /// overstate the total. Call this once per turn, after the turn's reports + /// have been merged, so the derived total covers the whole turn. + pub fn ensure_total(&mut self) { + if self.total_tokens.is_some() { + return; + } + self.total_tokens = match (self.input_tokens, self.output_tokens) { + (None, None) => None, + _ => Some( + self.input_tokens + .unwrap_or(0) + .saturating_add(self.output_tokens.unwrap_or(0)), + ), + }; + } +} + +/// Sum of two optional counters, where absent means "not reported" rather than +/// zero: the result is only absent when neither side reported anything. +fn add_optional(left: Option, right: Option) -> Option { + match (left, right) { + (None, None) => None, + _ => Some(left.unwrap_or(0).saturating_add(right.unwrap_or(0))), + } +} + +/// Larger of two optional counters, with the same absent-vs-zero distinction. +fn max_optional(left: Option, right: Option) -> Option { + match (left, right) { + (None, None) => None, + _ => Some(left.unwrap_or(0).max(right.unwrap_or(0))), + } +} + #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[serde(rename_all = "camelCase")] #[ts(export, export_to = "bindings.ts")] @@ -612,4 +682,112 @@ mod tests { other => panic!("expected Image, got {other:?}"), } } + + /// Regression: the engine used to overwrite `usage` on every provider + /// usage event, so a multi-turn run persisted only its final turn — the + /// shortest one in a tool-using run. + #[test] + fn run_usage_sums_completed_turns() { + let mut run = RunUsage::default(); + + let mut first = RunUsage { + input_tokens: Some(12_000), + output_tokens: Some(400), + ..RunUsage::default() + }; + first.ensure_total(); + run.add_turn(&first); + + let mut second = RunUsage { + input_tokens: Some(13_500), + output_tokens: Some(120), + ..RunUsage::default() + }; + second.ensure_total(); + run.add_turn(&second); + + assert_eq!(run.input_tokens, Some(25_500)); + assert_eq!(run.output_tokens, Some(520)); + assert_eq!(run.total_tokens, Some(26_020)); + } + + /// Anthropic reports twice per turn: `message_start` carries the prompt + /// count plus a tiny output preview, `message_delta` the cumulative output. + /// Each report restates the turn's counters, so they must collapse + /// field-wise rather than add. + #[test] + fn run_usage_collapses_repeated_reports_within_one_turn() { + let mut turn = RunUsage::default(); + + turn.merge_turn_report(&RunUsage { + input_tokens: Some(1_000), + output_tokens: Some(1), + ..RunUsage::default() + }); + turn.merge_turn_report(&RunUsage { + output_tokens: Some(500), + ..RunUsage::default() + }); + + assert_eq!(turn.input_tokens, Some(1_000)); + // 500, not 501: the preview is restated, not additional. + assert_eq!(turn.output_tokens, Some(500)); + + turn.ensure_total(); + assert_eq!(turn.total_tokens, Some(1_500)); + } + + /// Anthropic sends no total and several OpenAI-compatible endpoints + /// (MiniMax) omit it, which used to persist a run with no total at all. + #[test] + fn run_usage_derives_absent_total() { + let mut usage = RunUsage { + input_tokens: Some(700), + output_tokens: Some(80), + ..RunUsage::default() + }; + usage.ensure_total(); + assert_eq!(usage.total_tokens, Some(780)); + } + + /// An authoritative provider total always wins over the derived one. + #[test] + fn run_usage_keeps_provider_supplied_total() { + let mut usage = RunUsage { + input_tokens: Some(700), + output_tokens: Some(80), + reasoning_tokens: None, + total_tokens: Some(999), + }; + usage.ensure_total(); + assert_eq!(usage.total_tokens, Some(999)); + } + + /// Both native providers bill reasoning inside the output count, so the + /// derived total must not add it a second time. + #[test] + fn run_usage_derived_total_excludes_reasoning() { + let mut usage = RunUsage { + input_tokens: Some(100), + output_tokens: Some(60), + reasoning_tokens: Some(45), + total_tokens: None, + }; + usage.ensure_total(); + assert_eq!(usage.total_tokens, Some(160)); + } + + /// Absent must stay distinguishable from zero: a provider that reports + /// nothing should not manufacture a zeroed usage record. + #[test] + fn run_usage_absent_stays_absent() { + let mut usage = RunUsage::default(); + usage.ensure_total(); + assert_eq!(usage.total_tokens, None); + + let mut run = RunUsage::default(); + run.add_turn(&RunUsage::default()); + assert_eq!(run.input_tokens, None); + assert_eq!(run.total_tokens, None); + } }