Improve chat rendering and generation summaries

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,