diff --git a/migrations/20260727150000_add_message_generation_stats/down.sql b/migrations/20260727150000_add_message_generation_stats/down.sql new file mode 100644 index 0000000..fd73026 --- /dev/null +++ b/migrations/20260727150000_add_message_generation_stats/down.sql @@ -0,0 +1,4 @@ +ALTER TABLE messages DROP COLUMN output_tokens; +ALTER TABLE messages DROP COLUMN cached_tokens; +ALTER TABLE messages DROP COLUMN input_tokens; +ALTER TABLE messages DROP COLUMN generation_duration_ms; diff --git a/migrations/20260727150000_add_message_generation_stats/up.sql b/migrations/20260727150000_add_message_generation_stats/up.sql new file mode 100644 index 0000000..9da1a5d --- /dev/null +++ b/migrations/20260727150000_add_message_generation_stats/up.sql @@ -0,0 +1,4 @@ +ALTER TABLE messages ADD COLUMN generation_duration_ms INTEGER; +ALTER TABLE messages ADD COLUMN input_tokens INTEGER; +ALTER TABLE messages ADD COLUMN cached_tokens INTEGER; +ALTER TABLE messages ADD COLUMN output_tokens INTEGER; diff --git a/src/app.rs b/src/app.rs index 29fb0f2..605b45e 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1957,6 +1957,7 @@ mod tests { system: false, compaction: false, compaction_tail_start: None, + generation_stats: None, reasoning: Some(String::new()), reasoning_complete: false, reasoning_open: true, diff --git a/src/app/generation.rs b/src/app/generation.rs index 805cc0f..6fc4a19 100644 --- a/src/app/generation.rs +++ b/src/app/generation.rs @@ -55,6 +55,7 @@ pub(crate) struct ChatMessage { pub(super) system: bool, pub(super) compaction: bool, pub(super) compaction_tail_start: Option, + pub(super) generation_stats: Option, pub(super) reasoning: Option, pub(super) reasoning_complete: bool, pub(super) reasoning_open: bool, @@ -66,6 +67,14 @@ pub(crate) struct ChatMessage { pub(super) a2ui_open_urls: Vec, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct GenerationStats { + pub(super) duration_ms: u64, + pub(super) input_tokens: u32, + pub(super) cached_tokens: u32, + pub(super) output_tokens: u32, +} + const TOOL_PROTOCOL_CORRECTION: &str = "Tool protocol correction (no tool was executed): you emitted a complete tool call inside private reasoning. End reasoning, then emit the tool call again as assistant content using the required syntax; do not merely discuss it."; const EMPTY_RESPONSE_CORRECTION: &str = "Protocol correction: your previous response was empty. Continue the task now with either a valid tool call or a final answer."; @@ -122,6 +131,27 @@ fn correction_already_sent(conversation: &[ChatMessage], correction: &str) -> bo impl From for ChatMessage { fn from(message: StoredMessage) -> Self { + let generation_stats = match ( + message.generation_duration_ms, + message.input_tokens, + message.cached_tokens, + message.output_tokens, + ) { + (Some(duration_ms), Some(input_tokens), Some(cached_tokens), Some(output_tokens)) + if duration_ms >= 0 + && input_tokens >= 0 + && cached_tokens >= 0 + && output_tokens >= 0 => + { + Some(GenerationStats { + duration_ms: duration_ms as u64, + input_tokens: input_tokens as u32, + cached_tokens: cached_tokens as u32, + output_tokens: output_tokens as u32, + }) + } + _ => None, + }; let mut message = Self { id: message.id, user: message.user, @@ -129,6 +159,7 @@ impl From for ChatMessage { system: message.system, compaction: message.compaction, compaction_tail_start: message.compaction_tail_start, + generation_stats, reasoning: message.reasoning, reasoning_complete: message.reasoning_complete, reasoning_open: false, @@ -661,6 +692,30 @@ impl App { context_changed = true; } Ok(GenerationEvent::Finished(result)) => { + if let Ok(output) = &result { + let stats = GenerationStats { + duration_ms: u64::try_from(active.started_at.elapsed().as_millis()) + .unwrap_or(u64::MAX), + input_tokens: output.prompt_tokens, + cached_tokens: output.cached_tokens, + output_tokens: output.completion_tokens, + }; + if let Some(message) = self.conversation.last_mut() { + message.generation_stats = Some(stats); + if let Some(database) = &mut self.database + && let Err(error) = database.update_message_generation_stats( + message.id, + stats.duration_ms, + stats.input_tokens, + stats.cached_tokens, + stats.output_tokens, + ) + { + self.error = + Some(format!("Could not save generation summary: {error}")); + } + } + } match result { Ok(_) if self.stop_requested => { self.generating = false; @@ -1605,6 +1660,7 @@ mod tests { system: false, compaction: false, compaction_tail_start: None, + generation_stats: None, reasoning: reasoning.map(str::to_owned), reasoning_complete: false, reasoning_open: false, @@ -1664,6 +1720,7 @@ mod tests { system: false, compaction: false, compaction_tail_start: None, + generation_stats: None, reasoning: None, reasoning_complete: true, reasoning_open: false, @@ -1729,6 +1786,7 @@ mod tests { system: compaction, compaction, compaction_tail_start: tail, + generation_stats: None, reasoning: None, reasoning_complete: true, reasoning_open: false, @@ -1765,6 +1823,7 @@ mod tests { system, compaction, compaction_tail_start: None, + generation_stats: None, reasoning: None, reasoning_complete: true, reasoning_open: false, diff --git a/src/app/view.rs b/src/app/view.rs index 917906b..b505081 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -947,10 +947,7 @@ fn preference_group_style(_: &Theme) -> container::Style { } } -fn chat_message_style(theme: &Theme, user: bool) -> container::Style { - if !user { - return overview_style(theme); - } +fn chat_message_style(_: &Theme) -> container::Style { container::Style { background: Some(Background::Color(Color::from_rgb8(27, 34, 44))), border: Border { @@ -1052,8 +1049,8 @@ mod tests { assert_eq!(palette.background.weak.color, Color::from_rgb8(38, 38, 40)); assert_eq!(palette.secondary.base.color, Color::from_rgb8(47, 47, 50)); assert_ne!( - chat_message_style(&theme, true).background, - chat_message_style(&theme, false).background + chat_message_style(&theme).background, + overview_style(&theme).background ); assert_eq!( preference_group_style(&theme).background, diff --git a/src/app/view/chat.rs b/src/app/view/chat.rs index 1227120..dd533ec 100644 --- a/src/app/view/chat.rs +++ b/src/app/view/chat.rs @@ -180,13 +180,15 @@ impl App { body = body.push(tool_cards(cards)); } } - let user = message.user; - messages = messages.push( - container(body) - .padding(14) - .width(Length::Fill) - .style(move |theme| chat_message_style(theme, user)), - ); + let message_body = container(body).padding(14).width(Length::Fill); + messages = messages.push(if message.user { + message_body.style(chat_message_style) + } else { + message_body + }); + if let Some(stats) = message.generation_stats { + messages = messages.push(generation_summary(stats)); + } } if let Some(activity) = &self.activity { messages = messages.push( @@ -335,6 +337,26 @@ impl App { } } +fn generation_summary(stats: crate::app::generation::GenerationStats) -> Element<'static, Message> { + column![ + row![ + Space::new().width(Length::Fill), + text(format!( + "{} · {} input · {} cached · {} output", + format_duration(stats.duration_ms as f64 / 1_000.0), + stats.input_tokens, + stats.cached_tokens, + stats.output_tokens, + )) + .size(11) + .color(muted_text()), + ], + rule::horizontal(1), + ] + .spacing(5) + .into() +} + fn tool_cards(cards: Vec) -> Element<'static, Message> { let mut rows = column![].spacing(0); for (index, card) in cards.into_iter().enumerate() { diff --git a/src/database.rs b/src/database.rs index 5489767..2b15d52 100644 --- a/src/database.rs +++ b/src/database.rs @@ -121,6 +121,10 @@ pub struct StoredMessage { pub system: bool, pub compaction: bool, pub compaction_tail_start: Option, + pub generation_duration_ms: Option, + pub input_tokens: Option, + pub cached_tokens: Option, + pub output_tokens: Option, } #[derive(Insertable)] @@ -577,6 +581,34 @@ impl Database { .map_err(|error| error.to_string()) } + pub fn update_message_generation_stats( + &mut self, + id: i32, + duration_ms: u64, + input_tokens: u32, + cached_tokens: u32, + output_tokens: u32, + ) -> Result<(), String> { + let duration_ms = + i32::try_from(duration_ms).map_err(|_| "Generation duration is too large to save")?; + let input_tokens = + i32::try_from(input_tokens).map_err(|_| "Input token count is too large to save")?; + let cached_tokens = + i32::try_from(cached_tokens).map_err(|_| "Cached token count is too large to save")?; + let output_tokens = + i32::try_from(output_tokens).map_err(|_| "Output token count is too large to save")?; + diesel::update(messages::table.find(id)) + .set(( + messages::generation_duration_ms.eq(duration_ms), + messages::input_tokens.eq(input_tokens), + messages::cached_tokens.eq(cached_tokens), + messages::output_tokens.eq(output_tokens), + )) + .execute(&mut self.connection) + .map(|_| ()) + .map_err(|error| error.to_string()) + } + pub fn record_compaction( &mut self, session_id: i32, @@ -810,6 +842,9 @@ mod tests { database .update_message(assistant.id, Some("Reasoning"), true, "Answer") .unwrap(); + database + .update_message_generation_stats(assistant.id, 12_345, 1_024, 768, 256) + .unwrap(); database .insert_a2ui_message( session.id, @@ -843,6 +878,10 @@ mod tests { assert_eq!(messages[2].reasoning.as_deref(), Some("Reasoning")); assert!(messages[2].reasoning_complete); assert_eq!(messages[2].content, "Answer"); + assert_eq!(messages[2].generation_duration_ms, Some(12_345)); + assert_eq!(messages[2].input_tokens, Some(1_024)); + assert_eq!(messages[2].cached_tokens, Some(768)); + assert_eq!(messages[2].output_tokens, Some(256)); assert!(messages[3].tool); assert_eq!(messages[3].content, "Tool result"); assert!(messages[4].user); diff --git a/src/runtime.rs b/src/runtime.rs index 46fe5f9..a402e8e 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -19,6 +19,7 @@ pub(crate) struct GenerationService { pub(crate) struct ActiveGeneration { pub(crate) events: Receiver, pub(crate) cancel: Arc, + pub(crate) started_at: Instant, } impl Drop for ActiveGeneration { @@ -107,6 +108,7 @@ impl GenerationService { checkpoint: CheckpointTarget, idle_timeout: Duration, ) -> Result { + let started_at = Instant::now(); let source = checkpoint.source(); let cancel = Arc::new(AtomicBool::new(false)); let (events, receiver) = mpsc::channel(); @@ -131,6 +133,7 @@ impl GenerationService { Ok(ActiveGeneration { events: receiver, cancel, + started_at, }) } @@ -145,6 +148,7 @@ impl GenerationService { checkpoint: PathBuf, idle_timeout: Duration, ) -> Result { + let started_at = Instant::now(); let cancel = Arc::new(AtomicBool::new(false)); let (events, receiver) = mpsc::channel(); self.commands @@ -165,6 +169,7 @@ impl GenerationService { Ok(ActiveGeneration { events: receiver, cancel, + started_at, }) } @@ -175,6 +180,7 @@ impl GenerationService { messages: Vec, idle_timeout: Duration, ) -> Result { + let started_at = Instant::now(); let cancel = Arc::new(AtomicBool::new(false)); let (events, receiver) = mpsc::channel(); self.metrics.request_queued(WorkSource::LocalChat); @@ -193,6 +199,7 @@ impl GenerationService { Ok(ActiveGeneration { events: receiver, cancel, + started_at, }) } } @@ -430,6 +437,7 @@ mod tests { drop(ActiveGeneration { events, cancel: Arc::clone(&cancel), + started_at: Instant::now(), }); assert!(cancel.load(Ordering::Relaxed)); diff --git a/src/schema.rs b/src/schema.rs index 28623d5..bbf0e82 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -20,6 +20,10 @@ diesel::table! { system -> Bool, compaction -> Bool, compaction_tail_start -> Nullable, + generation_duration_ms -> Nullable, + input_tokens -> Nullable, + cached_tokens -> Nullable, + output_tokens -> Nullable, } } diff --git a/src/server.rs b/src/server.rs index 9c39996..7fcc8ca 100644 --- a/src/server.rs +++ b/src/server.rs @@ -945,6 +945,7 @@ mod tests { let active = crate::runtime::ActiveGeneration { events: receiver, cancel: Arc::new(AtomicBool::new(false)), + started_at: std::time::Instant::now(), }; let producer = thread::spawn(move || { thread::sleep(Duration::from_millis(25));