diff --git a/src-tauri/src/assistant/engine.rs b/src-tauri/src/assistant/engine.rs index 3c82ec3..c49d92c 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 4deafd7..89b575e 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); + } }