Improve chat rendering and generation summaries

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