From b1de123f06846b3835af064af14210a0c1b19b60 Mon Sep 17 00:00:00 2001 From: Blankeos Date: Fri, 14 Aug 2026 02:09:24 +0800 Subject: [PATCH 1/2] fix(chat): switch completion TPS calculation to OpenCode-style generation samples + live still Track assistant text generation as discrete samples from first text token to completion/tool-call boundary, then compute aggregate throughput via `(tokens - 1) / duration` with proper exclusion of tool-call segments and pause intervals; fall back when no valid rate exists. Persist `tokens_per_sec` support in session message conversion and message type defaults. --- src/app.rs | 82 ++-- src/persistence/conversions.rs | 1 + src/session/types.rs | 5 + src/ui/components/chat.rs | 689 ++++++++++++++++++++++++++++----- src/ui/components/input.rs | 11 + 5 files changed, 643 insertions(+), 145 deletions(-) diff --git a/src/app.rs b/src/app.rs index 6ade4b0..59fe0fc 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1498,38 +1498,7 @@ impl App { } fn completion_notification_stats(&self) -> Option { - let message = self.chat_state.chat.messages.iter().rev().find(|msg| { - msg.role == crate::session::types::MessageRole::Assistant && msg.is_complete - })?; - - 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); - 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); - - let total_sec = total_ms as f64 / 1000.0; - let tokens_per_sec = if decode_ms > 0 && output_tokens > 0 { - (output_tokens as f64) / (decode_ms as f64 / 1000.0) - } else { - 0.0 - }; - - return Some(format!("{:.1}s | {:.0}t/s", total_sec, tokens_per_sec)); - } - - if let (Some(token_count), Some(duration_ms)) = (message.token_count, message.duration_ms) { - let duration_sec = duration_ms as f64 / 1000.0; - let tokens_per_sec = if duration_ms > 0 { - (token_count as f64) / (duration_ms as f64 / 1000.0) - } else { - 0.0 - }; - - return Some(format!("{:.1}s | {:.0}t/s", duration_sec, tokens_per_sec)); - } - - None + Self::completion_notification_stats_for_chat(&self.chat_state.chat) } fn completion_notification_stats_for_chat(chat: &Chat) -> Option { @@ -1537,31 +1506,47 @@ impl App { msg.role == crate::session::types::MessageRole::Assistant && msg.is_complete })?; + 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 { + return None; + } + let tps = ((tokens - 1) as f64) / (decode_ms as f64 / 1000.0); + if tps.is_finite() && tps > 0.0 { + Some(tps) + } else { + None + } + }; + 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); 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); - let total_sec = total_ms as f64 / 1000.0; - let tokens_per_sec = if decode_ms > 0 && output_tokens > 0 { - (output_tokens as f64) / (decode_ms as f64 / 1000.0) - } else { - 0.0 - }; - return Some(format!("{:.1}s | {:.0}t/s", total_sec, tokens_per_sec)); + if let Some(tokens_per_sec) = + format_tps(message.tokens_per_sec, output_tokens, decode_ms) + { + return Some(format!("{:.1}s | {:.0}t/s", total_sec, tokens_per_sec)); + } + return Some(format!("{:.1}s", total_sec)); } if let (Some(token_count), Some(duration_ms)) = (message.token_count, message.duration_ms) { let duration_sec = duration_ms as f64 / 1000.0; - let tokens_per_sec = if duration_ms > 0 { - (token_count as f64) / (duration_ms as f64 / 1000.0) - } else { - 0.0 - }; - - return Some(format!("{:.1}s | {:.0}t/s", duration_sec, tokens_per_sec)); + if let Some(tokens_per_sec) = + format_tps(message.tokens_per_sec, token_count, duration_ms) + { + return Some(format!("{:.1}s | {:.0}t/s", duration_sec, tokens_per_sec)); + } + return Some(format!("{:.1}s", duration_sec)); } None @@ -8940,6 +8925,11 @@ impl App { crate::llm::ChunkMessage::Metrics { .. } => true, crate::llm::ChunkMessage::ToolCalls(tool_calls) => { self.set_session_retry_status(session_id, None); + // Close the generation sample as a tool-calls finish (excluded from + // TPS) and pause timing for the duration of tool execution. + if let Some(chat) = self.chat_for_session_mut(session_id) { + chat.end_generation_for_tool_calls(); + } self.add_tool_calls_to_session(session_id, tool_calls); true } diff --git a/src/persistence/conversions.rs b/src/persistence/conversions.rs index 33c92a6..54dbe02 100644 --- a/src/persistence/conversions.rs +++ b/src/persistence/conversions.rs @@ -200,6 +200,7 @@ impl TryFrom for SessionMessage { output_tokens: msg .output_tokens .and_then(|v| if v > 0 { Some(v as usize) } else { None }), + tokens_per_sec: None, model: msg.model.clone(), provider: msg.provider.clone(), local_image_paths, diff --git a/src/session/types.rs b/src/session/types.rs index 449121d..2ace65a 100644 --- a/src/session/types.rs +++ b/src/session/types.rs @@ -174,6 +174,9 @@ pub struct Message { pub t1_ms: Option, pub tn_ms: Option, pub output_tokens: Option, + /// Precomputed tokens/s (OpenCode inter-token aggregate). Prefer over + /// recomputing `output_tokens / duration_ms`. + pub tokens_per_sec: Option, pub model: Option, pub provider: Option, pub local_image_paths: Vec, @@ -205,6 +208,7 @@ impl Message { t1_ms: None, tn_ms: None, output_tokens: None, + tokens_per_sec: None, model: None, provider: None, local_image_paths: Vec::new(), @@ -252,6 +256,7 @@ impl Message { t1_ms: None, tn_ms: None, output_tokens: None, + tokens_per_sec: None, model: None, provider: None, local_image_paths: Vec::new(), diff --git a/src/ui/components/chat.rs b/src/ui/components/chat.rs index dc12370..560f3aa 100644 --- a/src/ui/components/chat.rs +++ b/src/ui/components/chat.rs @@ -264,7 +264,15 @@ pub struct Chat { streaming_pause_started_at: Option, streaming_paused_duration: std::time::Duration, streaming_decode_paused_duration: std::time::Duration, + /// Completed generation samples for OpenCode-style TPS aggregation. + generation_samples: Vec, + /// Active generation sample (opened on first text token of a step). + active_generation: Option, + /// Text-only token counter for the active generation sample. + generation_token_counter: Option, streaming_token_counter: Option, + /// Model id used to build token counters for this turn. + streaming_model: Option, /// Whether to autoscroll to bottom when new content arrives /// Only autoscrolls if user is already near the bottom pub autoscroll_enabled: bool, @@ -377,6 +385,78 @@ struct RenderedDiffLocationState { // Minimum elapsed time before showing tokens/s (250ms) const MIN_TOKENS_PER_SECOND_ELAPSED_MS: u128 = 250; +/// OpenCode-style floor: need more than one token for inter-token rate. +const MIN_TPS_SAMPLE_TOKENS: usize = 2; + +/// One LLM generation step's timing for TPS (OpenCode tps-tally model). +/// +/// - `started` = first *text* token of the step (not reasoning, not request start) +/// - `generated` = step end (stream finish or tool-call boundary) +/// - Tool-call finishes are excluded from TPS; only normal completions count +/// - Rate uses `(tokens - 1) / duration` (inter-token / TPOT style) +#[derive(Debug, Clone)] +struct GenerationSample { + started: std::time::Instant, + started_ms: u64, + generated: Option, + generated_ms: Option, + /// Visible assistant *text* tokens in this step (reasoning excluded). + tokens: usize, + /// Pause time while this sample was open (permission/question overlays). + paused_duration: std::time::Duration, + /// True when the step ended because the model requested tools. + tool_calls_finish: bool, +} + +impl GenerationSample { + fn generation_duration_ms(&self) -> Option { + let generated = self.generated?; + let raw_ms = generated + .saturating_duration_since(self.started) + .as_millis() as u64; + let paused_ms = self.paused_duration.as_millis() as u64; + Some(raw_ms.saturating_sub(paused_ms)) + } + + /// OpenCode eligibility: not tool-calls finish, tokens > 1, duration > 0. + fn tps_contribution(&self) -> Option<(usize, u64)> { + if self.tool_calls_finish { + return None; + } + if self.tokens < MIN_TPS_SAMPLE_TOKENS { + return None; + } + let duration_ms = self.generation_duration_ms()?; + if duration_ms == 0 { + return None; + } + // Inter-token units: (n - 1) tokens after the first. + Some((self.tokens - 1, duration_ms)) + } +} + +/// Prefer precomputed OpenCode TPS; fall back to inter-token formula. +fn message_tokens_per_sec( + precomputed: Option, + output_tokens: usize, + decode_ms: u64, +) -> Option { + if let Some(tps) = precomputed { + if tps.is_finite() && tps > 0.0 { + 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 + } +} + const MIN_MOUSE_WHEEL_LINES: usize = 1; const MAX_MOUSE_WHEEL_LINES: usize = 3; const MOUSE_WHEEL_VIEWPORT_FRACTION: usize = 8; @@ -1609,7 +1689,11 @@ impl Chat { streaming_pause_started_at: None, streaming_paused_duration: std::time::Duration::default(), streaming_decode_paused_duration: std::time::Duration::default(), + generation_samples: Vec::new(), + active_generation: None, + generation_token_counter: None, streaming_token_counter: None, + streaming_model: None, autoscroll_enabled: true, user_scrolled_up: false, cached_tokens_per_sec: None, @@ -1676,7 +1760,11 @@ impl Chat { streaming_pause_started_at: None, streaming_paused_duration: std::time::Duration::default(), streaming_decode_paused_duration: std::time::Duration::default(), + generation_samples: Vec::new(), + active_generation: None, + generation_token_counter: None, streaming_token_counter: None, + streaming_model: None, autoscroll_enabled: true, user_scrolled_up: false, cached_tokens_per_sec: None, @@ -1845,6 +1933,9 @@ impl Chat { pub fn append_to_last_assistant(&mut self, chunk: impl AsRef) { let chunk_str = chunk.as_ref(); + if chunk_str.is_empty() { + return; + } let appended_idx; // Append only if the last message is the current streaming assistant segment. @@ -1871,10 +1962,9 @@ impl Chat { self.streaming_start_time = Some(now); self.streaming_t0_ms = Some(now_epoch_ms()); } - if self.streaming_first_token_time.is_none() { - self.streaming_first_token_time = Some(now); - self.streaming_t1_ms = Some(now_epoch_ms()); - } + + // First *text* token opens a generation sample (OpenCode: time.started). + self.ensure_active_generation(); self.update_streaming_token_count(chunk_str); if self.should_autoscroll() { self.scroll_offset = usize::MAX; @@ -1884,6 +1974,9 @@ impl Chat { pub fn append_reasoning_to_last_assistant(&mut self, chunk: impl AsRef) { let chunk_str = chunk.as_ref(); + if chunk_str.is_empty() { + return; + } let appended_idx; if self @@ -1913,11 +2006,10 @@ impl Chat { self.streaming_start_time = Some(now); self.streaming_t0_ms = Some(now_epoch_ms()); } - if self.streaming_first_token_time.is_none() { - self.streaming_first_token_time = Some(now); - self.streaming_t1_ms = Some(now_epoch_ms()); - } - self.update_streaming_token_count(chunk_str); + // Reasoning does NOT open a generation sample and does NOT set TTFT + // (OpenCode only counts text-start / text-delta for TPS). + // Still track turn-level token count for display totals. + self.update_streaming_token_count_turn_only(chunk_str); if self.should_autoscroll() { self.scroll_offset = usize::MAX; self.user_scrolled_up = false; @@ -1974,7 +2066,11 @@ impl Chat { self.streaming_pause_started_at = None; self.streaming_paused_duration = std::time::Duration::default(); self.streaming_decode_paused_duration = std::time::Duration::default(); + self.generation_samples.clear(); + self.active_generation = None; + self.generation_token_counter = None; self.streaming_token_counter = None; + self.streaming_model = None; self.selection.reset(); self.pending_click_anchor = None; self.hovered_image = None; @@ -2170,6 +2266,9 @@ impl Chat { self.streaming_pause_started_at = None; self.streaming_paused_duration = std::time::Duration::default(); self.streaming_decode_paused_duration = std::time::Duration::default(); + self.generation_samples.clear(); + self.active_generation = None; + self.generation_token_counter = None; self.cached_tokens_per_sec = None; self.last_tps_calculated = None; @@ -2190,6 +2289,8 @@ impl Chat { let now = std::time::Instant::now(); self.streaming_end_time = Some(now); self.streaming_tn_ms = Some(now_epoch_ms()); + // Close the active generation sample as a normal (non-tool) finish. + self.close_active_generation(false); } pub fn get_streaming_tokens_per_sec(&self) -> Option { @@ -2200,6 +2301,7 @@ impl Chat { self.streaming_token_count } + /// Pause TPS timing (permission overlays, tool execution, questions). pub fn pause_streaming_tps_timer(&mut self) { if self.streaming_start_time.is_none() { return; @@ -2210,10 +2312,21 @@ impl Chat { } } + /// Close the active generation sample as a tool-calls finish (OpenCode: + /// tool-call steps are excluded from TPS). Tool-execution wall time is + /// also paused via [`pause_streaming_tps_timer`]. + pub fn end_generation_for_tool_calls(&mut self) { + self.close_active_generation(true); + self.pause_streaming_tps_timer(); + self.cached_tokens_per_sec = None; + self.last_tps_calculated = None; + } + pub fn resume_streaming_tps_timer(&mut self) { if let Some(started) = self.streaming_pause_started_at.take() { let ended = std::time::Instant::now(); - self.streaming_paused_duration += ended.duration_since(started); + let pause = ended.duration_since(started); + self.streaming_paused_duration += pause; if let Some(first_token_time) = self.streaming_first_token_time { if ended > first_token_time { let decode_pause_started = if started > first_token_time { @@ -2225,6 +2338,10 @@ impl Chat { ended.duration_since(decode_pause_started); } } + // Attribute overlay pauses to the active generation sample. + if let Some(sample) = self.active_generation.as_mut() { + sample.paused_duration += pause; + } self.last_tps_calculated = None; } } @@ -2268,29 +2385,179 @@ impl Chat { self.streaming_start_time.is_some() && self.streaming_assistant_idx().is_some() } + /// OpenCode-style aggregate: sum eligible `(tokens - 1)` and durations + /// across generation samples (tool-call finishes excluded). + fn aggregate_generation_tps( + samples: &[GenerationSample], + active: Option<&GenerationSample>, + ) -> Option { + let mut total_units: u64 = 0; + let mut total_duration_ms: u64 = 0; + + for sample in samples { + if let Some((units, duration_ms)) = sample.tps_contribution() { + total_units = total_units.saturating_add(units as u64); + total_duration_ms = total_duration_ms.saturating_add(duration_ms); + } + } + + if let Some(sample) = active { + // Live sample: use current instant as provisional `generated`. + let mut provisional = sample.clone(); + provisional.generated = Some(std::time::Instant::now()); + provisional.generated_ms = Some(now_epoch_ms()); + if let Some((units, duration_ms)) = provisional.tps_contribution() { + total_units = total_units.saturating_add(units as u64); + total_duration_ms = total_duration_ms.saturating_add(duration_ms); + } + } + + if total_units == 0 || total_duration_ms == 0 { + return None; + } + if total_duration_ms < MIN_TOKENS_PER_SECOND_ELAPSED_MS as u64 { + return None; + } + + let tps = (total_units as f64) / (total_duration_ms as f64 / 1000.0); + if tps.is_finite() && tps > 0.0 { + Some(tps) + } else { + None + } + } + + /// Sum of generation (decode) durations across samples. + /// Includes tool-call-ending steps (that was still LLM generation time) but + /// never includes tool *execution* wall time (samples end before tools run). + fn aggregate_generation_duration_ms( + samples: &[GenerationSample], + active: Option<&GenerationSample>, + ) -> u64 { + let mut total: u64 = 0; + for sample in samples { + if let Some(duration_ms) = sample.generation_duration_ms() { + total = total.saturating_add(duration_ms); + } + } + if let Some(sample) = active { + let mut provisional = sample.clone(); + provisional.generated = Some(std::time::Instant::now()); + // Attribute in-progress pause to the provisional sample when present. + if let Some(duration_ms) = provisional.generation_duration_ms() { + total = total.saturating_add(duration_ms); + } + } + total + } + + fn close_active_generation(&mut self, tool_calls_finish: bool) { + // Flush any pending overlay pause into the sample before closing. + if self.streaming_pause_started_at.is_some() && self.active_generation.is_some() { + // Don't clear the pause flag — tool execution may still be paused — + // but attribute elapsed pause so far into the sample. + if let (Some(started), Some(sample)) = ( + self.streaming_pause_started_at, + self.active_generation.as_mut(), + ) { + sample.paused_duration += started.elapsed(); + // Reset pause start so we don't double-count on resume. + self.streaming_pause_started_at = Some(std::time::Instant::now()); + } + } + + if let Some(mut sample) = self.active_generation.take() { + let now = std::time::Instant::now(); + sample.generated = Some(now); + sample.generated_ms = Some(now_epoch_ms()); + sample.tool_calls_finish = tool_calls_finish; + // Prefer the dedicated text-only counter for sample token count. + if let Some(counter) = self.generation_token_counter.as_ref() { + sample.tokens = counter.total_tokens(); + } + self.generation_samples.push(sample); + } + self.generation_token_counter = None; + } + + /// Open a new generation sample on the first *text* token of a step. + fn ensure_active_generation(&mut self) { + if self.active_generation.is_some() { + return; + } + // Ending an overlay pause before starting a new step keeps tool time out. + if self.streaming_pause_started_at.is_some() { + self.resume_streaming_tps_timer(); + } + + let now = std::time::Instant::now(); + let started_ms = now_epoch_ms(); + + // First text token of the whole turn → TTFT (t1). + if self.streaming_first_token_time.is_none() { + self.streaming_first_token_time = Some(now); + self.streaming_t1_ms = Some(started_ms); + if let Some(msg) = self + .messages + .last_mut() + .filter(|m| m.role == MessageRole::Assistant && !m.is_complete) + { + msg.t1_ms = Some(started_ms); + } + } + + // Seed the text-only counter for this step. + if self.generation_token_counter.is_none() { + let model = self.streaming_model.as_deref().unwrap_or(""); + self.generation_token_counter = Some(StreamingTokenCounter::new(model)); + } + + self.active_generation = Some(GenerationSample { + started: now, + started_ms, + generated: None, + generated_ms: None, + tokens: 0, + paused_duration: std::time::Duration::default(), + tool_calls_finish: false, + }); + } + pub fn finalize_streaming_metrics(&mut self) { let finalized_at = std::time::Instant::now(); + + // Close any still-open generation as a normal finish. + if self.active_generation.is_some() { + self.close_active_generation(false); + } + let token_count = self.streaming_token_count; let t0_ms = self.streaming_t0_ms; let t1_ms = self.streaming_t1_ms; - let tn_ms = self.streaming_tn_ms.or_else(|| { - // Fallback: if caller didn't mark end, compute an end timestamp now. - Some(now_epoch_ms()) - }); - - let paused_ms = self.total_decode_paused_duration().as_millis(); + let tn_ms = self.streaming_tn_ms.or_else(|| Some(now_epoch_ms())); - let decode_duration_ms = if let (Some(t1), Some(tn)) = - (self.streaming_first_token_time, self.streaming_end_time) - { - tn.duration_since(t1).as_millis().saturating_sub(paused_ms) as u64 - } else if let Some(t1) = self.streaming_first_token_time { - t1.elapsed().as_millis().saturating_sub(paused_ms) as u64 + // Prefer OpenCode sample-based decode duration; fall back to wall decode. + let sample_duration_ms = + Self::aggregate_generation_duration_ms(&self.generation_samples, None); + let decode_duration_ms = if sample_duration_ms > 0 { + sample_duration_ms } else { - 0 + let paused_ms = self.total_decode_paused_duration().as_millis(); + if let (Some(t1), Some(tn)) = (self.streaming_first_token_time, self.streaming_end_time) + { + tn.duration_since(t1).as_millis().saturating_sub(paused_ms) as u64 + } else if let Some(t1) = self.streaming_first_token_time { + t1.elapsed().as_millis().saturating_sub(paused_ms) as u64 + } else { + 0 + } }; + // Final TPS from completed samples only. + let final_tps = Self::aggregate_generation_tps(&self.generation_samples, None); + self.cached_tokens_per_sec = final_tps; + if let Some(idx) = self .messages .iter() @@ -2300,6 +2567,7 @@ impl Chat { msg.output_tokens = Some(token_count); msg.token_count = Some(token_count); msg.duration_ms = Some(decode_duration_ms); + msg.tokens_per_sec = final_tps; msg.finish_reasoning_timer(finalized_at); msg.t0_ms = t0_ms; msg.t1_ms = t1_ms; @@ -2318,6 +2586,9 @@ impl Chat { self.streaming_pause_started_at = None; self.streaming_paused_duration = std::time::Duration::default(); self.streaming_decode_paused_duration = std::time::Duration::default(); + self.generation_samples.clear(); + self.active_generation = None; + self.generation_token_counter = None; self.streaming_renderer = None; self.streaming_message_idx = None; self.streaming_renderer_content_len = 0; @@ -2325,6 +2596,7 @@ impl Chat { self.streaming_reasoning_message_idx = None; self.streaming_reasoning_renderer_content_len = 0; self.streaming_token_counter = None; + self.streaming_model = None; self.invalidate_cache(); } @@ -2379,10 +2651,13 @@ impl Chat { } pub fn prepare_streaming_token_counter(&mut self, model: &str) { + self.streaming_model = Some(model.to_string()); self.streaming_token_counter = Some(StreamingTokenCounter::new(model)); } - fn update_streaming_token_count(&mut self, chunk: &str) { + /// Count tokens for turn totals only (reasoning). Does not open a generation + /// sample or feed the text-only per-step TPS counter. + fn update_streaming_token_count_turn_only(&mut self, chunk: &str) { if let Some(counter) = self.streaming_token_counter.as_mut() { self.streaming_token_count = counter.add_text(chunk); } else { @@ -2390,6 +2665,20 @@ impl Chat { .streaming_token_count .saturating_add(estimate_tokens(chunk)); } + } + + /// Count tokens for both turn totals and the active generation sample (text). + fn update_streaming_token_count(&mut self, chunk: &str) { + self.update_streaming_token_count_turn_only(chunk); + + if let Some(counter) = self.generation_token_counter.as_mut() { + let tokens = counter.add_text(chunk); + if let Some(sample) = self.active_generation.as_mut() { + sample.tokens = tokens; + } + } else if let Some(sample) = self.active_generation.as_mut() { + sample.tokens = sample.tokens.saturating_add(estimate_tokens(chunk)); + } self.update_streaming_tokens_per_sec(); } @@ -2405,28 +2694,16 @@ impl Chat { } self.last_tps_calculated = Some(now); - let result = if let Some(first_token_time) = self.streaming_first_token_time { - let paused_ms = self.total_decode_paused_duration().as_millis(); - let elapsed_ms = first_token_time - .elapsed() - .as_millis() - .saturating_sub(paused_ms); - if elapsed_ms >= MIN_TOKENS_PER_SECOND_ELAPSED_MS && self.streaming_token_count > 0 { - let tokens_per_sec = - (self.streaming_token_count as f64) / (elapsed_ms as f64 / 1000.0); - if tokens_per_sec.is_finite() { - Some(tokens_per_sec) - } else { - None - } - } else { - None - } - } else { - None - }; - - self.cached_tokens_per_sec = result; + // Live aggregate: completed samples + provisional active sample. + // Attribute in-progress pause time to the active sample for accuracy. + let mut active = self.active_generation.clone(); + if let (Some(sample), Some(pause_started)) = + (active.as_mut(), self.streaming_pause_started_at) + { + sample.paused_duration += pause_started.elapsed(); + } + self.cached_tokens_per_sec = + Self::aggregate_generation_tps(&self.generation_samples, active.as_ref()); } /// Update the streaming markdown renderer for the current streaming message @@ -6011,6 +6288,8 @@ 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. 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); @@ -6031,15 +6310,14 @@ impl Chat { Style::default().fg(colors.text_weak), )); - let tokens_per_sec = if decode_ms > 0 && output_tokens > 0 { - (output_tokens as f64) / (decode_ms as f64 / 1000.0) - } else { - 0.0 - }; - spans.push(Span::styled( - format!(" • {:.0}t/s", tokens_per_sec), - Style::default().fg(colors.text_weak), - )); + if let Some(tokens_per_sec) = + message_tokens_per_sec(message.tokens_per_sec, output_tokens, decode_ms) + { + spans.push(Span::styled( + format!(" • {:.0}t/s", tokens_per_sec), + Style::default().fg(colors.text_weak), + )); + } } else if let (Some(token_count), Some(duration_ms)) = (message.token_count, message.duration_ms) { @@ -6049,15 +6327,14 @@ impl Chat { format!(" • {:.1}s", duration_sec), Style::default().fg(colors.text_weak), )); - let tokens_per_sec = if duration_ms > 0 { - (token_count as f64) / (duration_ms as f64 / 1000.0) - } else { - 0.0 - }; - spans.push(Span::styled( - format!(" • {:.0}t/s", tokens_per_sec), - Style::default().fg(colors.text_weak), - )); + if let Some(tokens_per_sec) = + message_tokens_per_sec(message.tokens_per_sec, token_count, duration_ms) + { + spans.push(Span::styled( + format!(" • {:.0}t/s", tokens_per_sec), + Style::default().fg(colors.text_weak), + )); + } } } @@ -9392,7 +9669,148 @@ codex exec --skip-git-repo-check \ } #[test] - fn test_streaming_pause_excluded_from_decode_duration() { + fn test_metadata_tps_uses_pause_adjusted_decode_duration() { + let mut chat = Chat::new(); + chat.add_assistant_message("hello"); + if let Some(message) = chat.messages.last_mut() { + message.is_complete = true; + message.t0_ms = Some(1_000); + message.t1_ms = Some(2_000); + 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 + message.duration_ms = Some(1_000); + } + + let colors = test_colors(); + let lines = chat.build_all_lines(100, "model", &colors); + 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 + ); + assert!( + rendered.contains("2.0s"), + "total duration should be ttft + adjusted decode duration:\n{}", + rendered + ); + } + + #[test] + fn test_metadata_tps_prefers_precomputed_sample_aggregate() { + let mut chat = Chat::new(); + chat.add_assistant_message("hello"); + if let Some(message) = chat.messages.last_mut() { + message.is_complete = true; + message.t0_ms = Some(1_000); + message.t1_ms = Some(2_000); + message.tn_ms = Some(3_000); + message.output_tokens = Some(50); + message.token_count = Some(50); + message.duration_ms = Some(1_000); + // Precomputed OpenCode multi-step aggregate should win over naive math. + message.tokens_per_sec = Some(42.0); + } + + let colors = test_colors(); + let lines = chat.build_all_lines(100, "model", &colors); + let rendered = lines.iter().map(line_text).collect::>().join("\n"); + + assert!( + rendered.contains("42t/s"), + "metadata should prefer precomputed tokens_per_sec:\n{}", + rendered + ); + } + + #[test] + fn test_tool_call_finish_excluded_from_tps() { + use std::time::{Duration, Instant}; + + 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(); + let turn_start = Instant::now(); + chat.append_to_last_assistant("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"); + std::thread::sleep(Duration::from_millis(250)); + chat.end_generation_for_tool_calls(); + let tool_start = Instant::now(); + std::thread::sleep(Duration::from_millis(400)); // tool execution wall time + let tool_elapsed_ms = tool_start.elapsed().as_millis() as u64; + chat.append_to_last_assistant("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"); + std::thread::sleep(Duration::from_millis(250)); + chat.mark_streaming_end(); + let wall_ms = turn_start.elapsed().as_millis() as u64; + + assert_eq!( + chat.generation_samples.len(), + 2, + "expected tool-finish sample + normal-finish sample" + ); + assert!( + chat.generation_samples[0].tool_calls_finish, + "first sample should be tool-calls finish" + ); + assert!( + !chat.generation_samples[1].tool_calls_finish, + "second sample should be normal finish" + ); + + // Tool-call finish must not contribute TPS units. + assert!( + chat.generation_samples[0].tps_contribution().is_none(), + "tool-calls finish sample must be excluded from TPS" + ); + // Normal finish with enough tokens should contribute. + assert!( + chat.generation_samples[1].tps_contribution().is_some(), + "normal finish sample should contribute to TPS" + ); + + let sample_sum = chat + .generation_samples + .iter() + .filter_map(|s| s.generation_duration_ms()) + .sum::(); + + chat.finalize_streaming_metrics(); + + let msg = chat + .messages + .iter() + .rev() + .find(|m| m.role == MessageRole::Assistant) + .expect("assistant message"); + + assert!( + msg.tokens_per_sec.is_some(), + "expected TPS from post-tool generation step" + ); + assert_eq!( + msg.duration_ms.unwrap_or(0), + sample_sum, + "duration_ms should equal sum of generation samples" + ); + // Sample sum must exclude tool-execution wall time (400ms+). + assert!( + sample_sum + tool_elapsed_ms <= wall_ms + 50, + "sample_sum ({sample_sum}) + tool ({tool_elapsed_ms}) should be ~wall ({wall_ms})" + ); + assert!( + sample_sum + 200 < wall_ms, + "sample_sum ({sample_sum}) should be well below wall ({wall_ms}) by tool time" + ); + } + + #[test] + fn test_reasoning_does_not_open_generation_or_set_ttft() { use std::time::Duration; let mut chat = Chat::new(); @@ -9402,16 +9820,95 @@ codex exec --skip-git-repo-check \ } chat.begin_streaming_turn(); + chat.append_reasoning_to_last_assistant("thinking hard about stuff..."); + std::thread::sleep(Duration::from_millis(80)); + assert!( + chat.streaming_first_token_time.is_none(), + "reasoning must not set TTFT" + ); + assert!( + chat.active_generation.is_none(), + "reasoning must not open a generation sample" + ); + + chat.append_to_last_assistant("answer"); + assert!( + chat.streaming_first_token_time.is_some(), + "first text token should set TTFT" + ); + assert!( + chat.active_generation.is_some(), + "first text token should open a generation sample" + ); + } + + #[test] + fn test_short_sample_rejected_from_tps() { + 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(); + // Single short chunk → likely 1 estimated token → rejected (need >1). + chat.append_to_last_assistant("x"); + 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"); + + // With only ~1 token the sample is ineligible (OpenCode: tokens > 1). + if msg.output_tokens.unwrap_or(0) < 2 { + assert!( + msg.tokens_per_sec.is_none(), + "single-token sample must not produce TPS" + ); + } + } + + #[test] + fn test_streaming_pause_excluded_from_decode_duration() { + use std::time::{Duration, Instant}; + + 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(); + let wall_start = Instant::now(); chat.append_to_last_assistant("hello"); std::thread::sleep(Duration::from_millis(40)); chat.pause_streaming_tps_timer(); + let pause_start = Instant::now(); std::thread::sleep(Duration::from_millis(320)); + let pause_ms = pause_start.elapsed().as_millis() as u64; chat.resume_streaming_tps_timer(); std::thread::sleep(Duration::from_millis(40)); chat.mark_streaming_end(); + assert_eq!(chat.generation_samples.len(), 1); + let sample = &chat.generation_samples[0]; + let sample_paused_ms = sample.paused_duration.as_millis() as u64; + let sample_duration_ms = sample.generation_duration_ms().unwrap_or(0); + assert!( + sample_paused_ms + 50 >= pause_ms, + "sample should attribute overlay pause (sample={sample_paused_ms}ms, actual={pause_ms}ms)" + ); + chat.finalize_streaming_metrics(); + let wall_ms = wall_start.elapsed().as_millis() as u64; let duration_ms = chat .messages @@ -9421,36 +9918,15 @@ codex exec --skip-git-repo-check \ .and_then(|m| m.duration_ms) .unwrap_or(0); - assert!(duration_ms < 250, "duration was {}ms", duration_ms); - } - - #[test] - fn test_metadata_tps_uses_pause_adjusted_decode_duration() { - let mut chat = Chat::new(); - chat.add_assistant_message("hello"); - if let Some(message) = chat.messages.last_mut() { - message.is_complete = true; - message.t0_ms = Some(1_000); - message.t1_ms = Some(2_000); - message.tn_ms = Some(12_000); - message.output_tokens = Some(100); - message.token_count = Some(100); - message.duration_ms = Some(1_000); - } - - let colors = test_colors(); - let lines = chat.build_all_lines(100, "model", &colors); - let rendered = lines.iter().map(line_text).collect::>().join("\n"); - + assert_eq!(duration_ms, sample_duration_ms); + // Decode duration must exclude the long overlay pause. assert!( - rendered.contains("100t/s"), - "metadata did not use adjusted decode duration:\n{}", - rendered + duration_ms + pause_ms <= wall_ms + 50, + "duration ({duration_ms}) + pause ({pause_ms}) should be ~wall ({wall_ms})" ); assert!( - rendered.contains("2.0s"), - "total duration should be ttft + adjusted decode duration:\n{}", - rendered + duration_ms + 200 < wall_ms, + "duration ({duration_ms}) should be well below wall ({wall_ms}) by pause time" ); } @@ -9493,7 +9969,7 @@ codex exec --skip-git-repo-check \ #[test] fn test_pre_token_pause_does_not_reduce_decode_duration() { - use std::time::Duration; + use std::time::{Duration, Instant}; let mut chat = Chat::new(); chat.add_assistant_message(""); @@ -9505,9 +9981,24 @@ codex exec --skip-git-repo-check \ chat.pause_streaming_tps_timer(); std::thread::sleep(Duration::from_millis(120)); chat.resume_streaming_tps_timer(); + + let gen_start = Instant::now(); chat.append_to_last_assistant("hello"); std::thread::sleep(Duration::from_millis(80)); chat.mark_streaming_end(); + let gen_wall_ms = gen_start.elapsed().as_millis() as u64; + + assert_eq!(chat.generation_samples.len(), 1); + let sample = &chat.generation_samples[0]; + // Generation sample opens on first text token — pre-token pause is not + // part of the sample window, so paused_duration should be ~0. + assert!( + sample.paused_duration.as_millis() < 30, + "pre-token pause leaked into sample: {}ms", + sample.paused_duration.as_millis() + ); + let sample_duration_ms = sample.generation_duration_ms().unwrap_or(0); + chat.finalize_streaming_metrics(); let duration_ms = chat @@ -9518,15 +10009,15 @@ codex exec --skip-git-repo-check \ .and_then(|m| m.duration_ms) .unwrap_or(0); + assert_eq!(duration_ms, sample_duration_ms); + // Duration should track post-token generation wall, not the 120ms pre-token pause. assert!( - duration_ms >= 40, - "pre-token pause should not be subtracted from decode duration: {}ms", - duration_ms + duration_ms + 50 >= gen_wall_ms.saturating_sub(30), + "decode duration too short ({duration_ms}ms vs gen wall {gen_wall_ms}ms)" ); assert!( - duration_ms < 180, - "decode duration unexpectedly included pre-token pause: {}ms", - duration_ms + duration_ms <= gen_wall_ms + 50, + "decode duration ({duration_ms}ms) exceeded post-token wall ({gen_wall_ms}ms)" ); } diff --git a/src/ui/components/input.rs b/src/ui/components/input.rs index 8cb4ed8..223a971 100644 --- a/src/ui/components/input.rs +++ b/src/ui/components/input.rs @@ -2523,6 +2523,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -2545,6 +2546,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -2576,6 +2578,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -2621,6 +2624,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -2784,6 +2788,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -2817,6 +2822,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -2845,6 +2851,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -3040,6 +3047,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -3079,6 +3087,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -3104,6 +3113,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); @@ -3294,6 +3304,7 @@ mod tests { "provider", None, &colors, + false, ); }) .unwrap(); From 0813fa654cc2ba8169e1914e2c1b7bbad98f6613 Mon Sep 17 00:00:00 2001 From: Blankeos Date: Fri, 14 Aug 2026 02:56:58 +0800 Subject: [PATCH 2/2] fix(chat): reuse token counter and defer streaming layout refresh MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - reuse the previous turn’s streaming token counter when available to avoid repeated tiktoken warm-up after tool calls - gate markdown stream cache invalidation behind layout interval cadence and keep pending dirty state for later flush - update streaming and token-rate tests for revised collapsed-reasoning label and layout-timing behavior --- src/app.rs | 19 +++++++++++-- src/ui/components/chat.rs | 56 ++++++++++++++++++++++++++++++--------- 2 files changed, 60 insertions(+), 15 deletions(-) diff --git a/src/app.rs b/src/app.rs index 59fe0fc..d16becf 100644 --- a/src/app.rs +++ b/src/app.rs @@ -2102,6 +2102,9 @@ impl App { fn set_session_retry_status(&mut self, session_id: &str, status: Option) { self.ensure_session_view_state(session_id); if let Some(state) = self.session_view_states.get_mut(session_id) { + if state.retry_status == status { + return; + } state.retry_status = status; } } @@ -13005,6 +13008,9 @@ mod tests { .chat .add_message(crate::session::types::Message::incomplete("")); app.chat_state.chat.begin_streaming_turn(); + app.chat_state + .chat + .prepare_streaming_token_counter("test-model"); let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); sender @@ -13129,14 +13135,23 @@ mod tests { app.overlay_focus = OverlayFocus::SessionsDialog; app.sessions_dialog_state.dialog.show(); app.refresh_sessions_dialog(); - app.sessions_dialog_live_dirty = false; - let probe_before = app.last_sessions_dialog_metadata_probe; app.base_focus = BaseFocus::Chat; app.chat_state .chat .add_message(crate::session::types::Message::incomplete("")); app.chat_state.chat.begin_streaming_turn(); + // Warm tiktoken before probing so load cost is not charged to the + // sessions-dialog probe interval during process_streaming_chunks. + app.chat_state + .chat + .prepare_streaming_token_counter("test-model"); + + // Reset probe after any setup cost (e.g. tiktoken warm-up) so the + // assertion only covers process_streaming_chunks itself. + app.sessions_dialog_live_dirty = false; + app.last_sessions_dialog_metadata_probe = std::time::Instant::now(); + let probe_before = app.last_sessions_dialog_metadata_probe; let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); sender diff --git a/src/ui/components/chat.rs b/src/ui/components/chat.rs index 560f3aa..247bd0c 100644 --- a/src/ui/components/chat.rs +++ b/src/ui/components/chat.rs @@ -2507,9 +2507,17 @@ impl Chat { } // Seed the text-only counter for this step. + // Reuse the turn counter when present so the first text chunk after a + // tool call does not re-load tiktoken (can take hundreds of ms). if self.generation_token_counter.is_none() { - let model = self.streaming_model.as_deref().unwrap_or(""); - self.generation_token_counter = Some(StreamingTokenCounter::new(model)); + if let Some(turn_counter) = self.streaming_token_counter.clone() { + let mut sample_counter = turn_counter; + sample_counter.reset(); + self.generation_token_counter = Some(sample_counter); + } else { + let model = self.streaming_model.as_deref().unwrap_or(""); + self.generation_token_counter = Some(StreamingTokenCounter::new(model)); + } } self.active_generation = Some(GenerationSample { @@ -2839,13 +2847,24 @@ impl Chat { } if refreshed { - let dirty_from = self - .pending_streaming_render_dirty_from - .take() - .unwrap_or(last_idx); - self.pending_streaming_content_dirty = false; - self.last_streaming_cache_refresh_at = Some(std::time::Instant::now()); - self.invalidate_cache_from(dirty_from); + let dirty_from = self.pending_streaming_render_dirty_from.unwrap_or(last_idx); + // Markdown may re-render on a shorter cadence than tool-heavy layout. + // Only bump render_revision / drop line caches when the adaptive + // layout interval has elapsed. + let layout_due = layout_min_interval.is_zero() + || self + .last_streaming_cache_refresh_at + .map_or(true, |last| last.elapsed() >= layout_min_interval); + if layout_due { + self.pending_streaming_render_dirty_from = None; + self.pending_streaming_content_dirty = false; + self.last_streaming_cache_refresh_at = Some(std::time::Instant::now()); + self.invalidate_cache_from(dirty_from); + } else { + // Keep content marked dirty so a later frame flushes layout. + self.pending_streaming_content_dirty = true; + self.pending_streaming_render_dirty_from = Some(dirty_from); + } } self.flush_non_content_streaming_render_pending(); @@ -7499,9 +7518,13 @@ mod tests { .map(line_text) .collect::>(); - assert!(collapsed - .iter() - .any(|line| line.contains("Thinking collapsed"))); + // Collapsed reasoning shows the duration/status label only (not "Thinking collapsed"). + assert!( + collapsed + .iter() + .any(|line| { line.contains("Thought for") || line.contains("Thinking") }), + "expected collapsed thinking label, got: {collapsed:?}" + ); assert!(!collapsed .iter() .any(|line| line.contains("Private reasoning"))); @@ -7606,6 +7629,8 @@ mod tests { chat.update_streaming_renderer(100, &colors); let first_line_count = chat.build_all_lines(100, "model", &colors).len(); + // Freeze layout clock so deferred append cannot race wall-clock expiry. + chat.last_streaming_cache_refresh_at = Some(std::time::Instant::now()); let first_revision = chat.render_revision; let cached_prefix_len = chat .ordered_tool_prefix_cache @@ -7614,7 +7639,11 @@ mod tests { .map(|cache| cache.lines.len()) .expect("tool prefix cache"); assert!(cached_prefix_len > 0); + chat.append_to_last_assistant(" continues"); + // Re-freeze immediately before the deferred update so wall-clock from + // tool-prefix warm-up cannot open the layout interval. + chat.last_streaming_cache_refresh_at = Some(std::time::Instant::now()); chat.update_streaming_renderer(100, &colors); assert_eq!(chat.render_revision, first_revision); @@ -9592,7 +9621,8 @@ codex exec --skip-git-repo-check \ assert!(metadata.contains("1.0s")); assert!(metadata.contains("ttft 0.2s")); - assert!(metadata.contains("50t/s")); + // OpenCode inter-token: (40 - 1) / 0.8s = 48.75 → rounds to 49t/s + assert!(metadata.contains("49t/s")); } #[test]