prev/next A2UI surface and better dismissal

src/database.rs

@@ -145,6 +145,7 @@ pub struct StoredA2uiMessage {
     pub session_id: i32,
     pub message_id: i32,
     pub json: String,
+    pub dismissed: bool,
 }
 
 #[derive(Insertable)]
@@ -153,6 +154,7 @@ struct NewA2uiMessage<'a> {
     session_id: i32,
     message_id: i32,
     json: &'a str,
+    dismissed: bool,
 }
 
 #[derive(Debug)]
@@ -344,12 +346,51 @@ impl Database {
                 session_id,
                 message_id,
                 json,
+                dismissed: false,
             })
             .returning(StoredA2uiMessage::as_returning())
             .get_result(&mut self.connection)
             .map_err(|error| error.to_string())
     }
 
+    pub fn dismiss_a2ui_surface(
+        &mut self,
+        session_id: i32,
+        surface_id: &str,
+    ) -> Result<StoredMessage, String> {
+        let content = format!(
+            "A2UI surface `{surface_id}` was dismissed by the user. There is no active A2UI surface. If the user asks for UI again, create a complete new surface with createSurface and a new surfaceId; do not update any earlier surface."
+        );
+        let json = serde_json::json!({"surfaceId": surface_id}).to_string();
+        self.connection
+            .transaction(|connection| {
+                let message = diesel::insert_into(messages::table)
+                    .values(NewMessage {
+                        session_id,
+                        user: false,
+                        tool: false,
+                        reasoning: None,
+                        reasoning_complete: true,
+                        content: &content,
+                        system: true,
+                        compaction: false,
+                        compaction_tail_start: None,
+                    })
+                    .returning(StoredMessage::as_returning())
+                    .get_result(connection)?;
+                diesel::insert_into(a2ui_messages::table)
+                    .values(NewA2uiMessage {
+                        session_id,
+                        message_id: message.id,
+                        json: &json,
+                        dismissed: true,
+                    })
+                    .execute(connection)?;
+                Ok(message)
+            })
+            .map_err(|error: diesel::result::Error| error.to_string())
+    }
+
     pub fn update_session_context(
         &mut self,
         session_id: i32,
@@ -672,6 +713,86 @@ mod tests {
         fs::remove_file(path).unwrap();
     }
 
+    #[test]
+    fn a2ui_dismissal_persists_a_fresh_surface_boundary() {
+        let id = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let path = std::env::temp_dir().join(format!("ds4-a2ui-dismiss-{id}.sqlite3"));
+        let mut database = Database::open(&path).unwrap();
+        let project = database
+            .create_project("DS4", "/tmp/ds4-a2ui-dismiss")
+            .unwrap();
+        let session = database.create_session(project.id, "A2UI").unwrap();
+        let first = database
+            .start_chat_turn(session.id, "First", &[], false)
+            .unwrap()
+            .pop()
+            .unwrap();
+        database
+            .insert_a2ui_message(
+                session.id,
+                first.id,
+                r#"{"version":"v1.0","createSurface":{"surfaceId":"first","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
+            )
+            .unwrap();
+        let dismissal = database.dismiss_a2ui_surface(session.id, "first").unwrap();
+        assert!(dismissal.system);
+        assert!(dismissal.content.contains("createSurface"));
+        let protocol = database.load_a2ui_messages(session.id).unwrap();
+        let (_, active, errors) = crate::a2ui::replay_epochs(protocol.iter().map(|message| {
+            (
+                message.id,
+                message.message_id,
+                message.dismissed,
+                message.json.as_str(),
+            )
+        }));
+        assert!(errors.is_empty());
+        assert!(active.active_surface().is_none());
+
+        let second = database
+            .start_chat_turn(session.id, "Second", &[], false)
+            .unwrap()
+            .pop()
+            .unwrap();
+        database
+            .insert_a2ui_message(
+                session.id,
+                second.id,
+                r#"{"version":"v1.0","createSurface":{"surfaceId":"second","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
+            )
+            .unwrap();
+        drop(database);
+
+        let mut reopened = Database::open(&path).unwrap();
+        let protocol = reopened.load_a2ui_messages(session.id).unwrap();
+        assert_eq!(protocol.len(), 3);
+        assert!(protocol[1].dismissed);
+        assert_eq!(protocol[1].message_id, dismissal.id);
+        let (_, active, errors) = crate::a2ui::replay_epochs(protocol.iter().map(|message| {
+            (
+                message.id,
+                message.message_id,
+                message.dismissed,
+                message.json.as_str(),
+            )
+        }));
+        assert!(errors.is_empty());
+        assert_eq!(active.active_surface().unwrap().id, "second");
+        let messages = reopened.load_messages(session.id).unwrap();
+        assert!(
+            messages
+                .iter()
+                .position(|message| message.id == dismissal.id)
+                .is_some_and(|index| messages[index + 1].content == "Second")
+        );
+        reopened.delete_project(project.id).unwrap();
+        drop(reopened);
+        fs::remove_file(path).unwrap();
+    }
+
     #[test]
     fn chat_and_compaction_history_survive_reopen_and_session_deletion() {
         let id = SystemTime::now()