Complete long-context chat support

src/database.rs

@@ -260,6 +260,8 @@ pub struct Session {
     pub id: i32,
     pub project_id: i32,
     pub title: String,
+    pub context_used: i32,
+    pub context_limit: i32,
 }
 
 #[derive(Insertable)]
@@ -476,6 +478,24 @@ impl Database {
             .map_err(|error| error.to_string())
     }
 
+    pub fn update_session_context(
+        &mut self,
+        session_id: i32,
+        used: u32,
+        limit: u32,
+    ) -> Result<(), String> {
+        let used = i32::try_from(used).map_err(|_| "Used context is too large to save")?;
+        let limit = i32::try_from(limit).map_err(|_| "Context limit is too large to save")?;
+        diesel::update(sessions::table.find(session_id))
+            .set((
+                sessions::context_used.eq(used),
+                sessions::context_limit.eq(limit),
+            ))
+            .execute(&mut self.connection)
+            .map(|_| ())
+            .map_err(|error| error.to_string())
+    }
+
     pub fn start_chat_turn(
         &mut self,
         session_id: i32,
@@ -648,9 +668,15 @@ mod tests {
         database
             .update_message(assistant.id, Some("Reasoning"), true, "Answer")
             .unwrap();
+        database
+            .update_session_context(session.id, 1_234, 65_536)
+            .unwrap();
         drop(database);
 
         let mut reopened = Database::open(&path).unwrap();
+        let projects = reopened.load_projects().unwrap();
+        assert_eq!(projects[0].sessions[0].context_used, 1_234);
+        assert_eq!(projects[0].sessions[0].context_limit, 65_536);
         let messages = reopened.load_messages(session.id).unwrap();
         assert_eq!(messages.len(), 2);
         assert!(messages[0].user);