Compaction overhaul (#5186)

Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: David Katz <dkatz@squareup.com>
Co-authored-by: Alex Hancock <alexhancock@block.xyz>

crates/goose/src/session/session_manager.rs

@@ -18,7 +18,7 @@ use tokio::sync::OnceCell;
 use tracing::{info, warn};
 use utoipa::ToSchema;
 
-const CURRENT_SCHEMA_VERSION: i32 = 2;
+const CURRENT_SCHEMA_VERSION: i32 = 3;
 
 static SESSION_STORAGE: OnceCell<Arc<SessionStorage>> = OnceCell::const_new();
 
@@ -425,7 +425,8 @@ impl SessionStorage {
                 content_json TEXT NOT NULL,
                 created_timestamp INTEGER NOT NULL,
                 timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                tokens INTEGER
+                tokens INTEGER,
+                metadata_json TEXT
             )
         "#,
         )
@@ -610,6 +611,15 @@ impl SessionStorage {
                 .execute(&self.pool)
                 .await?;
             }
+            3 => {
+                sqlx::query(
+                    r#"
+                    ALTER TABLE messages ADD COLUMN metadata_json TEXT
+                "#,
+                )
+                .execute(&self.pool)
+                .await?;
+            }
             _ => {
                 anyhow::bail!("Unknown migration version: {}", version);
             }
@@ -768,15 +778,15 @@ impl SessionStorage {
     }
 
     async fn get_conversation(&self, session_id: &str) -> Result<Conversation> {
-        let rows = sqlx::query_as::<_, (String, String, i64)>(
-            "SELECT role, content_json, created_timestamp FROM messages WHERE session_id = ? ORDER BY timestamp",
+        let rows = sqlx::query_as::<_, (String, String, i64, Option<String>)>(
+            "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp",
         )
             .bind(session_id)
             .fetch_all(&self.pool)
             .await?;
 
         let mut messages = Vec::new();
-        for (role_str, content_json, created_timestamp) in rows {
+        for (role_str, content_json, created_timestamp, metadata_json) in rows {
             let role = match role_str.as_str() {
                 "user" => Role::User,
                 "assistant" => Role::Assistant,
@@ -784,7 +794,12 @@ impl SessionStorage {
             };
 
             let content = serde_json::from_str(&content_json)?;
-            let message = Message::new(role, created_timestamp, content);
+            let metadata = metadata_json
+                .and_then(|json| serde_json::from_str(&json).ok())
+                .unwrap_or_default();
+
+            let mut message = Message::new(role, created_timestamp, content);
+            message.metadata = metadata;
             messages.push(message);
         }
 
@@ -792,16 +807,19 @@ impl SessionStorage {
     }
 
     async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> {
+        let metadata_json = serde_json::to_string(&message.metadata)?;
+
         sqlx::query(
             r#"
-            INSERT INTO messages (session_id, role, content_json, created_timestamp)
-            VALUES (?, ?, ?, ?)
+            INSERT INTO messages (session_id, role, content_json, created_timestamp, metadata_json)
+            VALUES (?, ?, ?, ?, ?)
         "#,
         )
         .bind(session_id)
         .bind(role_to_string(&message.role))
         .bind(serde_json::to_string(&message.content)?)
         .bind(message.created)
+        .bind(metadata_json)
         .execute(&self.pool)
         .await?;
 
@@ -826,16 +844,19 @@ impl SessionStorage {
             .await?;
 
         for message in conversation.messages() {
+            let metadata_json = serde_json::to_string(&message.metadata)?;
+
             sqlx::query(
                 r#"
-            INSERT INTO messages (session_id, role, content_json, created_timestamp)
-            VALUES (?, ?, ?, ?)
+            INSERT INTO messages (session_id, role, content_json, created_timestamp, metadata_json)
+            VALUES (?, ?, ?, ?, ?)
         "#,
             )
             .bind(session_id)
             .bind(role_to_string(&message.role))
             .bind(serde_json::to_string(&message.content)?)
             .bind(message.created)
+            .bind(metadata_json)
             .execute(&mut *tx)
             .await?;
         }