feat: better chat support for thinking

src/database.rs

@@ -4,7 +4,7 @@ use std::collections::HashMap;
 use std::fs;
 use std::path::Path;
 
-use crate::schema::{preferences, projects, sessions};
+use crate::schema::{messages, preferences, projects, sessions};
 use crate::settings::{
     DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, ReasoningMode,
     RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences,
@@ -269,6 +269,28 @@ struct NewSession<'a> {
     title: &'a str,
 }
 
+#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
+#[diesel(table_name = messages)]
+#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
+pub struct StoredMessage {
+    pub id: i32,
+    pub session_id: i32,
+    pub user: bool,
+    pub reasoning: Option<String>,
+    pub reasoning_complete: bool,
+    pub content: String,
+}
+
+#[derive(Insertable)]
+#[diesel(table_name = messages)]
+struct NewMessage<'a> {
+    session_id: i32,
+    user: bool,
+    reasoning: Option<&'a str>,
+    reasoning_complete: bool,
+    content: &'a str,
+}
+
 #[derive(Debug)]
 pub struct ProjectWithSessions {
     pub project: Project,
@@ -408,6 +430,16 @@ impl Database {
     pub fn delete_project(&mut self, project_id: i32) -> Result<(), String> {
         self.connection
             .transaction(|connection| {
+                diesel::delete(
+                    messages::table.filter(
+                        messages::session_id.eq_any(
+                            sessions::table
+                                .filter(sessions::project_id.eq(project_id))
+                                .select(sessions::id),
+                        ),
+                    ),
+                )
+                .execute(connection)?;
                 diesel::delete(sessions::table.filter(sessions::project_id.eq(project_id)))
                     .execute(connection)?;
                 diesel::delete(projects::table.find(project_id)).execute(connection)?;
@@ -425,7 +457,71 @@ impl Database {
     }
 
     pub fn delete_session(&mut self, session_id: i32) -> Result<(), String> {
-        diesel::delete(sessions::table.find(session_id))
+        self.connection
+            .transaction(|connection| {
+                diesel::delete(messages::table.filter(messages::session_id.eq(session_id)))
+                    .execute(connection)?;
+                diesel::delete(sessions::table.find(session_id)).execute(connection)?;
+                Ok(())
+            })
+            .map_err(|error: diesel::result::Error| error.to_string())
+    }
+
+    pub fn load_messages(&mut self, session_id: i32) -> Result<Vec<StoredMessage>, String> {
+        messages::table
+            .filter(messages::session_id.eq(session_id))
+            .order(messages::id.asc())
+            .select(StoredMessage::as_select())
+            .load(&mut self.connection)
+            .map_err(|error| error.to_string())
+    }
+
+    pub fn start_chat_turn(
+        &mut self,
+        session_id: i32,
+        prompt: &str,
+        reasoning: bool,
+    ) -> Result<(StoredMessage, StoredMessage), String> {
+        self.connection
+            .transaction(|connection| {
+                let user = diesel::insert_into(messages::table)
+                    .values(NewMessage {
+                        session_id,
+                        user: true,
+                        reasoning: None,
+                        reasoning_complete: true,
+                        content: prompt,
+                    })
+                    .returning(StoredMessage::as_returning())
+                    .get_result(connection)?;
+                let assistant = diesel::insert_into(messages::table)
+                    .values(NewMessage {
+                        session_id,
+                        user: false,
+                        reasoning: reasoning.then_some(""),
+                        reasoning_complete: !reasoning,
+                        content: "",
+                    })
+                    .returning(StoredMessage::as_returning())
+                    .get_result(connection)?;
+                Ok((user, assistant))
+            })
+            .map_err(|error: diesel::result::Error| error.to_string())
+    }
+
+    pub fn update_message(
+        &mut self,
+        id: i32,
+        reasoning: Option<&str>,
+        reasoning_complete: bool,
+        content: &str,
+    ) -> Result<(), String> {
+        diesel::update(messages::table.find(id))
+            .set((
+                messages::reasoning.eq(reasoning),
+                messages::reasoning_complete.eq(reasoning_complete),
+                messages::content.eq(content),
+            ))
             .execute(&mut self.connection)
             .map(|_| ())
             .map_err(|error| error.to_string())
@@ -535,4 +631,36 @@ mod tests {
         drop(reopened);
         fs::remove_file(path).unwrap();
     }
+
+    #[test]
+    fn chat_messages_survive_reopen_and_follow_session_deletion() {
+        let id = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let path = std::env::temp_dir().join(format!("ds4-chat-{id}.sqlite3"));
+        let mut database = Database::open(&path).unwrap();
+        let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap();
+        let session = database.create_session(project.id, "Chat").unwrap();
+        let (_, assistant) = database
+            .start_chat_turn(session.id, "Question", true)
+            .unwrap();
+        database
+            .update_message(assistant.id, Some("Reasoning"), true, "Answer")
+            .unwrap();
+        drop(database);
+
+        let mut reopened = Database::open(&path).unwrap();
+        let messages = reopened.load_messages(session.id).unwrap();
+        assert_eq!(messages.len(), 2);
+        assert!(messages[0].user);
+        assert_eq!(messages[0].content, "Question");
+        assert_eq!(messages[1].reasoning.as_deref(), Some("Reasoning"));
+        assert!(messages[1].reasoning_complete);
+        assert_eq!(messages[1].content, "Answer");
+        reopened.delete_session(session.id).unwrap();
+        assert!(reopened.load_messages(session.id).unwrap().is_empty());
+        drop(reopened);
+        fs::remove_file(path).unwrap();
+    }
 }