Improve chat rendering and generation summaries

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;

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;

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,

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<i32>,
+    pub(super) generation_stats: Option<GenerationStats>,
     pub(super) reasoning: Option<String>,
     pub(super) reasoning_complete: bool,
     pub(super) reasoning_open: bool,
@@ -66,6 +67,14 @@ pub(crate) struct ChatMessage {
     pub(super) a2ui_open_urls: Vec<String>,
 }
 
+#[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<StoredMessage> 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<StoredMessage> 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,

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 {
+fn chat_message_style(_: &Theme) -> container::Style {
-    if !user {
-        return overview_style(theme);
-    }
     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).background,
-            chat_message_style(&theme, false).background
+            overview_style(&theme).background
         );
         assert_eq!(
             preference_group_style(&theme).background,

src/app/view/chat.rs

@@ -180,13 +180,15 @@ impl App {
                             body = body.push(tool_cards(cards));
                         }
                     }
-                    let user = message.user;
+                    let message_body = container(body).padding(14).width(Length::Fill);
-                    messages = messages.push(
+                    messages = messages.push(if message.user {
-                        container(body)
+                        message_body.style(chat_message_style)
-                            .padding(14)
+                    } else {
-                            .width(Length::Fill)
+                        message_body
-                            .style(move |theme| chat_message_style(theme, user)),
+                    });
-                    );
+                    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<crate::agent::ToolCard>) -> Element<'static, Message> {
     let mut rows = column![].spacing(0);
     for (index, card) in cards.into_iter().enumerate() {

src/database.rs

@@ -121,6 +121,10 @@ pub struct StoredMessage {
     pub system: bool,
     pub compaction: bool,
     pub compaction_tail_start: Option<i32>,
+    pub generation_duration_ms: Option<i32>,
+    pub input_tokens: Option<i32>,
+    pub cached_tokens: Option<i32>,
+    pub output_tokens: Option<i32>,
 }
 
 #[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);

src/runtime.rs

@@ -19,6 +19,7 @@ pub(crate) struct GenerationService {
 pub(crate) struct ActiveGeneration {
     pub(crate) events: Receiver<GenerationEvent>,
     pub(crate) cancel: Arc<AtomicBool>,
+    pub(crate) started_at: Instant,
 }
 
 impl Drop for ActiveGeneration {
@@ -107,6 +108,7 @@ impl GenerationService {
         checkpoint: CheckpointTarget,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
+        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<ActiveGeneration, String> {
+        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<ChatTurn>,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
+        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));

src/schema.rs

@@ -20,6 +20,10 @@ diesel::table! {
         system -> Bool,
         compaction -> Bool,
         compaction_tail_start -> Nullable<Integer>,
+        generation_duration_ms -> Nullable<Integer>,
+        input_tokens -> Nullable<Integer>,
+        cached_tokens -> Nullable<Integer>,
+        output_tokens -> Nullable<Integer>,
     }
 }
 

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));