feat: a2ui interface to enablee the LLM to give structured information

src/database.rs

@@ -4,7 +4,7 @@ use std::collections::HashMap;
 use std::fs;
 use std::path::Path;
 
-use crate::schema::{messages, projects, sessions};
+use crate::schema::{a2ui_messages, messages, projects, sessions};
 
 pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
 
@@ -137,6 +137,24 @@ struct NewMessage<'a> {
     compaction_tail_start: Option<i32>,
 }
 
+#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
+#[diesel(table_name = a2ui_messages)]
+#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
+pub struct StoredA2uiMessage {
+    pub id: i32,
+    pub session_id: i32,
+    pub message_id: i32,
+    pub json: String,
+}
+
+#[derive(Insertable)]
+#[diesel(table_name = a2ui_messages)]
+struct NewA2uiMessage<'a> {
+    session_id: i32,
+    message_id: i32,
+    json: &'a str,
+}
+
 #[derive(Debug)]
 pub struct ProjectWithSessions {
     pub project: Project,
@@ -219,6 +237,16 @@ impl Database {
     pub fn delete_project(&mut self, project_id: i32) -> Result<(), String> {
         self.connection
             .transaction(|connection| {
+                diesel::delete(
+                    a2ui_messages::table.filter(
+                        a2ui_messages::session_id.eq_any(
+                            sessions::table
+                                .filter(sessions::project_id.eq(project_id))
+                                .select(sessions::id),
+                        ),
+                    ),
+                )
+                .execute(connection)?;
                 diesel::delete(
                     messages::table.filter(
                         messages::session_id.eq_any(
@@ -272,6 +300,10 @@ impl Database {
     pub fn delete_session(&mut self, session_id: i32) -> Result<(), String> {
         self.connection
             .transaction(|connection| {
+                diesel::delete(
+                    a2ui_messages::table.filter(a2ui_messages::session_id.eq(session_id)),
+                )
+                .execute(connection)?;
                 diesel::delete(messages::table.filter(messages::session_id.eq(session_id)))
                     .execute(connection)?;
                 diesel::delete(sessions::table.find(session_id)).execute(connection)?;
@@ -289,6 +321,35 @@ impl Database {
             .map_err(|error| error.to_string())
     }
 
+    pub fn load_a2ui_messages(
+        &mut self,
+        session_id: i32,
+    ) -> Result<Vec<StoredA2uiMessage>, String> {
+        a2ui_messages::table
+            .filter(a2ui_messages::session_id.eq(session_id))
+            .order(a2ui_messages::id.asc())
+            .select(StoredA2uiMessage::as_select())
+            .load(&mut self.connection)
+            .map_err(|error| error.to_string())
+    }
+
+    pub fn insert_a2ui_message(
+        &mut self,
+        session_id: i32,
+        message_id: i32,
+        json: &str,
+    ) -> Result<StoredA2uiMessage, String> {
+        diesel::insert_into(a2ui_messages::table)
+            .values(NewA2uiMessage {
+                session_id,
+                message_id,
+                json,
+            })
+            .returning(StoredA2uiMessage::as_returning())
+            .get_result(&mut self.connection)
+            .map_err(|error| error.to_string())
+    }
+
     pub fn update_session_context(
         &mut self,
         session_id: i32,
@@ -628,6 +689,13 @@ mod tests {
         database
             .update_message(assistant.id, Some("Reasoning"), true, "Answer")
             .unwrap();
+        database
+            .insert_a2ui_message(
+                session.id,
+                assistant.id,
+                r#"{"version":"v1.0","createSurface":{"surfaceId":"saved","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
+            )
+            .unwrap();
         database
             .continue_tool_turn(
                 session.id,
@@ -662,6 +730,10 @@ mod tests {
         assert_eq!(messages[5].content, "Tool reminder");
         assert!(!messages[6].user);
         assert!(!messages[6].tool);
+        let a2ui = reopened.load_a2ui_messages(session.id).unwrap();
+        assert_eq!(a2ui.len(), 1);
+        assert_eq!(a2ui[0].message_id, messages[2].id);
+        assert!(a2ui[0].json.contains("createSurface"));
         let first = reopened
             .record_compaction(
                 session.id,
@@ -744,6 +816,7 @@ mod tests {
         assert_eq!(history[15].content, "After third compaction");
         reopened.delete_session(session.id).unwrap();
         assert!(reopened.load_messages(session.id).unwrap().is_empty());
+        assert!(reopened.load_a2ui_messages(session.id).unwrap().is_empty());
         drop(reopened);
         fs::remove_file(path).unwrap();
     }