diff --git a/src/app.rs b/src/app.rs index 4b023ea..a25c86b 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1708,17 +1708,18 @@ impl App { msg.role == crate::session::types::MessageRole::Assistant && msg.is_complete })?; + // Upstream formula (opencode #46108): billed output / decode time, + // 250ms floor, no inter-token adjustment. let format_tps = |precomputed: Option, tokens: usize, decode_ms: u64| -> Option { if let Some(tps) = precomputed { if tps.is_finite() && tps > 0.0 { return Some(tps); } } - // OpenCode inter-token: (n - 1) / duration; need >1 token. - if decode_ms == 0 || tokens < 2 { + if tokens == 0 || decode_ms < 250 { return None; } - let tps = ((tokens - 1) as f64) / (decode_ms as f64 / 1000.0); + let tps = tokens as f64 / (decode_ms as f64 / 1000.0); if tps.is_finite() && tps > 0.0 { Some(tps) } else { @@ -1727,7 +1728,9 @@ impl App { }; if let (Some(t0), Some(t1), Some(tn)) = (message.t0_ms, message.t1_ms, message.tn_ms) { - let output_tokens = message.output_tokens.or(message.token_count).unwrap_or(0); + // t/s inputs are output tokens only — `token_count` is the billed + // total and would inflate the rate on reloaded sessions. + let output_tokens = message.output_tokens.unwrap_or(0); let ttft_ms = t1.saturating_sub(t0); let decode_ms = message.duration_ms.unwrap_or_else(|| tn.saturating_sub(t1)); let total_ms = ttft_ms.saturating_add(decode_ms); @@ -1741,16 +1744,24 @@ impl App { return Some(format!("{:.1}s", total_sec)); } - if let (Some(token_count), Some(duration_ms)) = (message.token_count, message.duration_ms) { + if let (Some(output_tokens), Some(duration_ms)) = + (message.output_tokens, message.duration_ms) + { let duration_sec = duration_ms as f64 / 1000.0; if let Some(tokens_per_sec) = - format_tps(message.tokens_per_sec, token_count, duration_ms) + format_tps(message.tokens_per_sec, output_tokens, duration_ms) { return Some(format!("{:.1}s | {:.0}t/s", duration_sec, tokens_per_sec)); } return Some(format!("{:.1}s", duration_sec)); } + if message.duration_ms.is_some() { + // Total-only legacy row: duration without t/s. + let duration_sec = message.duration_ms.unwrap_or(0) as f64 / 1000.0; + return Some(format!("{:.1}s", duration_sec)); + } + None } diff --git a/src/persistence/conversions.rs b/src/persistence/conversions.rs index b5076eb..6ebf61b 100644 --- a/src/persistence/conversions.rs +++ b/src/persistence/conversions.rs @@ -110,6 +110,7 @@ impl From for Message { t1_ms: msg.t1_ms.map(|v| v as i64), tn_ms: msg.tn_ms.map(|v| v as i64), output_tokens: msg.output_tokens.map(|v| v as i64), + tokens_per_sec: msg.tokens_per_sec, } } } @@ -185,6 +186,31 @@ impl TryFrom for SessionMessage { _ => return Err(anyhow::anyhow!("Unknown role: {}", msg.role)), }; + // Billed output buckets are exact; prefer them over the persisted + // text estimate when backfilling rows stored before output_tokens + // existed. Never derive output tokens from `tokens_used` (total). + let billed_output: Option = { + let mut total = 0u64; + let mut found = false; + for part in &session_parts { + if part.part_type == "usage" { + if let Some(output) = part.data.get("output").and_then(|v| v.as_u64()) { + total = total.saturating_add(output); + found = true; + } + } + } + if found && total > 0 && total <= usize::MAX as u64 { + Some(total as usize) + } else { + None + } + }; + let persisted_output_tokens: Option = + msg.output_tokens + .and_then(|v| if v > 0 { Some(v as usize) } else { None }); + let output_tokens = persisted_output_tokens.or(billed_output); + Ok(SessionMessage { role, content, @@ -213,10 +239,8 @@ impl TryFrom for SessionMessage { tn_ms: msg .tn_ms .and_then(|v| if v > 0 { Some(v as u64) } else { None }), - output_tokens: msg - .output_tokens - .and_then(|v| if v > 0 { Some(v as usize) } else { None }), - tokens_per_sec: None, + output_tokens, + tokens_per_sec: msg.tokens_per_sec.filter(|v| v.is_finite() && *v > 0.0), model: msg.model.clone(), provider: msg.provider.clone(), local_image_paths, @@ -358,4 +382,55 @@ mod tests { assert_eq!(usage.input, 80_000); assert_eq!(usage.output, 400); } + + #[test] + fn precomputed_tps_round_trips_through_persistence() { + let mut session_message = SessionMessage::assistant("done"); + session_message.output_tokens = Some(390); + session_message.tokens_per_sec = Some(145.0); + + let persistence_message: Message = session_message.into(); + assert_eq!(persistence_message.tokens_per_sec, Some(145.0)); + + let restored = SessionMessage::try_from(persistence_message).unwrap(); + assert_eq!(restored.tokens_per_sec, Some(145.0)); + assert_eq!(restored.output_tokens, Some(390)); + } + + #[test] + fn billed_output_backfills_output_tokens_not_total() { + // Legacy row stored before output_tokens existed: tokens_used is the + // billed total (in+out+cache), usage part carries exact buckets. + let mut legacy = Message { + id: "legacy".to_string(), + session_id: 1, + role: "assistant".to_string(), + parts: vec![PersistenceMessagePart { + part_type: "text".to_string(), + data: serde_json::json!({ "text": "done" }), + }], + timestamp: 0, + tokens_used: 8000, + model: None, + provider: None, + agent_mode: None, + duration_ms: 2600, + t0_ms: Some(1000), + t1_ms: Some(10_000), + tn_ms: Some(12_600), + output_tokens: None, + tokens_per_sec: None, + }; + legacy.parts.push(PersistenceMessagePart { + part_type: "usage".to_string(), + data: serde_json::json!({ + "input": 7000, "output": 390, + "cache_read": 500, "cache_write": 110, "cost": 0.01, + }), + }); + + let restored = SessionMessage::try_from(legacy).unwrap(); + // Output bucket (390), never the billed total (8000). + assert_eq!(restored.output_tokens, Some(390)); + } } diff --git a/src/persistence/history.rs b/src/persistence/history.rs index ff9998f..c111bc3 100644 --- a/src/persistence/history.rs +++ b/src/persistence/history.rs @@ -62,6 +62,7 @@ pub struct Message { pub t1_ms: Option, pub tn_ms: Option, pub output_tokens: Option, + pub tokens_per_sec: Option, } pub struct HistoryDAO { @@ -438,9 +439,9 @@ impl HistoryDAO { self.conn.execute( "INSERT INTO messages ( id, session_id, role, parts, timestamp, tokens_used, model, provider, agent_mode, duration_ms, - t0_ms, t1_ms, tn_ms, output_tokens + t0_ms, t1_ms, tn_ms, output_tokens, tokens_per_sec ) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)", + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)", params![ &msg.id, msg.session_id, @@ -456,6 +457,7 @@ impl HistoryDAO { msg.t1_ms, msg.tn_ms, msg.output_tokens, + msg.tokens_per_sec, ], )?; @@ -488,9 +490,9 @@ impl HistoryDAO { let mut insert = tx.prepare_cached( "INSERT INTO messages ( id, session_id, role, parts, timestamp, tokens_used, model, provider, agent_mode, duration_ms, - t0_ms, t1_ms, tn_ms, output_tokens + t0_ms, t1_ms, tn_ms, output_tokens, tokens_per_sec ) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)", + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)", )?; for msg in messages { @@ -514,6 +516,7 @@ impl HistoryDAO { msg.t1_ms, msg.tn_ms, msg.output_tokens, + msg.tokens_per_sec, ])?; } } @@ -555,7 +558,7 @@ impl HistoryDAO { pub fn get_messages(&self, session_id: i64) -> Result> { let mut stmt = self.conn.prepare( "SELECT id, session_id, role, parts, timestamp, tokens_used, model, provider, agent_mode, duration_ms, - t0_ms, t1_ms, tn_ms, output_tokens + t0_ms, t1_ms, tn_ms, output_tokens, tokens_per_sec FROM messages WHERE session_id = ?1 ORDER BY timestamp ASC, rowid ASC", )?; @@ -578,6 +581,7 @@ impl HistoryDAO { t1_ms: row.get(11)?, tn_ms: row.get(12)?, output_tokens: row.get(13)?, + tokens_per_sec: row.get(14).unwrap_or(None), }) })?; diff --git a/src/persistence/migrations.rs b/src/persistence/migrations.rs index 1fedbbf..e292687 100644 --- a/src/persistence/migrations.rs +++ b/src/persistence/migrations.rs @@ -16,6 +16,10 @@ pub fn run_migrations(db: &mut Connection) -> Result<()> { migrate_to_v3(db)?; } + if current_version < 4 { + migrate_to_v4(db)?; + } + Ok(()) } @@ -63,6 +67,7 @@ fn migrate_to_v1(db: &mut Connection) -> Result<()> { t1_ms INTEGER, tn_ms INTEGER, output_tokens INTEGER, + tokens_per_sec REAL, FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE ); @@ -181,3 +186,19 @@ fn migrate_to_v3(db: &mut Connection) -> Result<()> { tx.commit()?; Ok(()) } + +fn migrate_to_v4(db: &mut Connection) -> Result<()> { + let tx = db.transaction()?; + + // Precomputed inter-token TPS so a reloaded session shows the same t/s + // as the live stream instead of recomputing from token estimates. + let _ = tx.execute("ALTER TABLE messages ADD COLUMN tokens_per_sec REAL", []); + + tx.execute( + "INSERT OR IGNORE INTO migrations (version, applied_at) VALUES (4, strftime('%s', 'now'))", + params![], + )?; + + tx.commit()?; + Ok(()) +} diff --git a/src/ui/components/chat.rs b/src/ui/components/chat.rs index eab851c..a307630 100644 --- a/src/ui/components/chat.rs +++ b/src/ui/components/chat.rs @@ -455,7 +455,26 @@ impl GenerationSample { } } -/// Prefer precomputed OpenCode TPS; fall back to inter-token formula. +/// Upstream TPS formula (opencode #46108 `packages/core/src/session/tokens.ts`): +/// billed output tokens over decode time (`completed - firstToken`), +/// 250ms floor, no inter-token adjustment. Provider `output` already +/// counts reasoning tokens, so no separate bucket is needed. +fn upstream_tokens_per_sec(generated_tokens: u64, decode_ms: u64) -> Option { + if generated_tokens == 0 { + return None; + } + if decode_ms < MIN_TOKENS_PER_SECOND_ELAPSED_MS as u64 { + return None; + } + let tps = generated_tokens as f64 / (decode_ms as f64 / 1000.0); + if tps.is_finite() && tps > 0.0 { + Some(tps) + } else { + None + } +} + +/// Prefer precomputed TPS; fall back to the upstream formula. fn message_tokens_per_sec( precomputed: Option, output_tokens: usize, @@ -466,15 +485,7 @@ fn message_tokens_per_sec( return Some(tps); } } - if decode_ms == 0 || output_tokens < MIN_TPS_SAMPLE_TOKENS { - return None; - } - let tps = ((output_tokens - 1) as f64) / (decode_ms as f64 / 1000.0); - if tps.is_finite() && tps > 0.0 { - Some(tps) - } else { - None - } + upstream_tokens_per_sec(output_tokens as u64, decode_ms) } const MIN_MOUSE_WHEEL_LINES: usize = 1; @@ -2646,15 +2657,26 @@ impl Chat { } }; - // Final TPS from completed samples only. - let final_tps = Self::aggregate_generation_tps(&self.generation_samples, None); - self.cached_tokens_per_sec = final_tps; + // Sample-based estimate kept as fallback when the provider reported + // no billed usage. + let sample_tps = Self::aggregate_generation_tps(&self.generation_samples, None); if let Some(idx) = self .messages .iter() .rposition(|m| m.role == MessageRole::Assistant) { + // Upstream-style final rate (opencode #46108): billed output + // over decode time. Billed usage accumulates across every + // provider step (tool-call steps included), so reloads + // recompute identically instead of drifting from estimates. + let billed_output = self.messages[idx] + .recorded_usage() + .map(|usage| usage.output) + .unwrap_or(0); + let final_tps = + upstream_tokens_per_sec(billed_output, decode_duration_ms).or(sample_tps); + self.cached_tokens_per_sec = final_tps; if let Some(msg) = self.messages.get_mut(idx) { msg.output_tokens = Some(msg.output_tokens.unwrap_or(token_count)); msg.token_count = msg.output_tokens; @@ -6623,11 +6645,15 @@ impl Chat { )); // Timing + throughput metrics are shown only once the stream is done. - // TPS uses OpenCode inter-token rate: (tokens - 1) / generation_duration, - // preferring the precomputed sample aggregate on the message. + // TPS uses the upstream rate (billed output / decode time), + // preferring the precomputed value persisted on the message. if include_metrics { if let (Some(t0), Some(t1), Some(tn)) = (message.t0_ms, message.t1_ms, message.tn_ms) { - let output_tokens = message.output_tokens.or(message.token_count).unwrap_or(0); + // t/s inputs are output tokens only. `token_count` is the + // billed total (prompt + completion + cache) and must never + // feed the throughput fallback — that inflated reloaded + // sessions (e.g. 145t/s live vs 3292t/s on reopen). + let output_tokens = message.output_tokens.unwrap_or(0); let ttft_ms = t1.saturating_sub(t0); let decode_ms = message.duration_ms.unwrap_or_else(|| tn.saturating_sub(t1)); @@ -6653,8 +6679,8 @@ impl Chat { Style::default().fg(colors.text_weak), )); } - } else if let (Some(token_count), Some(duration_ms)) = - (message.token_count, message.duration_ms) + } else if let (Some(output_tokens), Some(duration_ms)) = + (message.output_tokens, message.duration_ms) { // Backward-compatible fallback: duration_ms reflects decode time. let duration_sec = duration_ms as f64 / 1000.0; @@ -6663,13 +6689,21 @@ impl Chat { Style::default().fg(colors.text_weak), )); if let Some(tokens_per_sec) = - message_tokens_per_sec(message.tokens_per_sec, token_count, duration_ms) + message_tokens_per_sec(message.tokens_per_sec, output_tokens, duration_ms) { spans.push(Span::styled( format!(" • {:.0}t/s", tokens_per_sec), Style::default().fg(colors.text_weak), )); } + } else if let Some(duration_ms) = message.duration_ms { + // Total-only legacy row: show duration without t/s rather + // than dividing the billed total by decode time. + let duration_sec = duration_ms as f64 / 1000.0; + spans.push(Span::styled( + format!(" • {:.1}s", duration_sec), + Style::default().fg(colors.text_weak), + )); } } @@ -10015,8 +10049,8 @@ codex exec --skip-git-repo-check \ assert!(metadata.contains("1.0s")); assert!(metadata.contains("ttft 0.2s")); - // OpenCode inter-token: (40 - 1) / 0.8s = 48.75 → rounds to 49t/s - assert!(metadata.contains("49t/s")); + // Upstream formula: 40 / 0.8s = 50t/s + assert!(metadata.contains("50t/s")); } #[test] @@ -10103,7 +10137,7 @@ codex exec --skip-git-repo-check \ message.tn_ms = Some(12_000); message.output_tokens = Some(100); message.token_count = Some(100); - // 1s decode, OpenCode inter-token: (100 - 1) / 1s = 99 t/s + // 1s decode, upstream formula: 100 / 1s = 100 t/s message.duration_ms = Some(1_000); } @@ -10112,8 +10146,8 @@ codex exec --skip-git-repo-check \ let rendered = lines.iter().map(line_text).collect::>().join("\n"); assert!( - rendered.contains("99t/s"), - "metadata should use inter-token TPS (n-1)/duration:\n{}", + rendered.contains("100t/s"), + "metadata should use upstream TPS tokens/duration:\n{}", rendered ); assert!( @@ -10233,6 +10267,38 @@ codex exec --skip-git-repo-check \ ); } + #[test] + fn test_finalize_prefers_billed_output_over_sample_estimate() { + use std::time::Duration; + + let mut chat = Chat::new(); + chat.add_assistant_message(""); + if let Some(last) = chat.messages.last_mut() { + last.is_complete = false; + } + + chat.begin_streaming_turn(); + // Tiny visible text: the sample estimate would be ~10t/s. + chat.append_to_last_assistant("hi"); + // ...but billed output says 300 tokens over ~0.3s ≈ 1000t/s. + chat.record_usage(0, 300, 0, 0, 0.0); + std::thread::sleep(Duration::from_millis(300)); + chat.mark_streaming_end(); + chat.finalize_streaming_metrics(); + + let msg = chat + .messages + .iter() + .rev() + .find(|m| m.role == MessageRole::Assistant) + .expect("assistant message"); + let tps = msg.tokens_per_sec.expect("expected billed TPS"); + assert!( + tps > 500.0, + "final TPS should come from billed output, not the text estimate: got {tps}" + ); + } + #[test] fn test_reasoning_does_not_open_generation_or_set_ttft() { use std::time::Duration;