Complete long-context chat support

src/app.rs

@@ -403,6 +403,8 @@ pub(crate) struct App {
     pub(super) composer: String,
     pub(super) conversation: Vec<ChatMessage>,
     pub(super) generating: bool,
+    pub(super) context_used: u32,
+    pub(super) context_limit: u32,
     #[cfg(target_os = "macos")]
     generation_worker: Option<GenerationWorker>,
     error: Option<String>,
@@ -464,6 +466,7 @@ enum GenerationCommand {
 enum GenerationEvent {
     Loading,
     Chunk { reasoning: bool, content: String },
+    Context { used: u32, limit: u32 },
     Finished(Result<(), String>),
 }
 
@@ -572,6 +575,7 @@ impl App {
                         Ok(draft) => draft,
                         Err(error) => return Self::failed(error, main_window),
                     };
+                    let context_limit = preferences.context_tokens.max(0) as u32;
                     Self {
                         main_window,
                         model_manager_window: None,
@@ -593,6 +597,8 @@ impl App {
                         composer: String::new(),
                         conversation: Vec::new(),
                         generating: false,
+                        context_used: 0,
+                        context_limit,
                         #[cfg(target_os = "macos")]
                         generation_worker: None,
                         error: None,
@@ -608,6 +614,7 @@ impl App {
         let preferences = AppPreferences::default();
         let preference_draft = PreferenceDraft::from_saved(&preferences)
             .expect("default preferences must use a supported model");
+        let context_limit = preferences.context_tokens.max(0) as u32;
         Self {
             main_window,
             model_manager_window: None,
@@ -629,6 +636,8 @@ impl App {
             composer: String::new(),
             conversation: Vec::new(),
             generating: false,
+            context_used: 0,
+            context_limit,
             #[cfg(target_os = "macos")]
             generation_worker: None,
             error: Some(format!("Could not open the project database: {error}")),
@@ -953,6 +962,8 @@ impl App {
                 self.selected_session = None;
                 self.conversation.clear();
                 self.composer.clear();
+                self.context_used = 0;
+                self.context_limit = self.preferences.context_tokens.max(0) as u32;
                 self.error = None;
             }
             Message::DeleteProject(project_id) => {
@@ -969,6 +980,7 @@ impl App {
                                 self.selected_session = None;
                                 self.conversation.clear();
                                 self.composer.clear();
+                                self.context_used = 0;
                             }
                             self.reload_projects();
                         }
@@ -986,6 +998,12 @@ impl App {
                 if self.selected_session == Some(session_id) {
                     return Task::none();
                 }
+                let saved_context = self
+                    .projects
+                    .iter()
+                    .flat_map(|project| &project.sessions)
+                    .find(|session| session.id == session_id)
+                    .map(|session| (session.context_used, session.context_limit));
                 let Some(database) = &mut self.database else {
                     return Task::none();
                 };
@@ -995,6 +1013,13 @@ impl App {
                         self.composer.clear();
                         self.selected_project = Some(project_id);
                         self.selected_session = Some(session_id);
+                        let (used, limit) = saved_context.unwrap_or_default();
+                        self.context_used = used.max(0) as u32;
+                        self.context_limit = if limit > 0 {
+                            limit as u32
+                        } else {
+                            self.preferences.context_tokens.max(0) as u32
+                        };
                         self.error = None;
                     }
                     Err(error) => {
@@ -1015,6 +1040,7 @@ impl App {
                                 self.selected_session = None;
                                 self.conversation.clear();
                                 self.composer.clear();
+                                self.context_used = 0;
                             }
                             self.reload_projects();
                         }
@@ -1277,6 +1303,8 @@ impl App {
                 self.selected_session = Some(session.id);
                 self.conversation.clear();
                 self.composer.clear();
+                self.context_used = 0;
+                self.context_limit = self.preferences.context_tokens.max(0) as u32;
                 self.error = None;
                 self.reload_projects();
             }
@@ -1473,6 +1501,8 @@ impl App {
         #[cfg(target_os = "macos")]
         let mut transcript_changed = false;
         #[cfg(target_os = "macos")]
+        let mut context_changed = false;
+        #[cfg(target_os = "macos")]
         loop {
             match worker.events.try_recv() {
                 Ok(GenerationEvent::Loading) => {}
@@ -1484,6 +1514,11 @@ impl App {
                         transcript_changed = true;
                     }
                 }
+                Ok(GenerationEvent::Context { used, limit }) => {
+                    self.context_used = used;
+                    self.context_limit = limit;
+                    context_changed = true;
+                }
                 Ok(GenerationEvent::Finished(result)) => {
                     self.generating = false;
                     worker.cancel = None;
@@ -1521,6 +1556,25 @@ impl App {
             }
             self.error = Some(format!("Could not save generated chat text: {error}"));
         }
+        #[cfg(target_os = "macos")]
+        if context_changed
+            && let Some(session_id) = self.selected_session
+            && let Some(database) = &mut self.database
+        {
+            if let Err(error) =
+                database.update_session_context(session_id, self.context_used, self.context_limit)
+            {
+                self.error = Some(format!("Could not save context usage: {error}"));
+            } else if let Some(session) = self
+                .projects
+                .iter_mut()
+                .flat_map(|project| &mut project.sessions)
+                .find(|session| session.id == session_id)
+            {
+                session.context_used = self.context_used as i32;
+                session.context_limit = self.context_limit as i32;
+            }
+        }
     }
 }
 
@@ -1567,6 +1621,10 @@ fn spawn_generation_worker() -> Result<GenerationWorker, String> {
                                     let _ = event_sender
                                         .send(GenerationEvent::Chunk { reasoning, content });
                                 },
+                                |used, limit| {
+                                    let _ =
+                                        event_sender.send(GenerationEvent::Context { used, limit });
+                                },
                             );
                             let _ = event_sender.send(GenerationEvent::Finished(result));
                             last_used = Instant::now();