Wire up agent worktree creation, cleanup, and integration tests

crates/agent/src/agent.rs

@@ -1395,19 +1395,12 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
     fn set_title(
         &self,
         session_id: &acp::SessionId,
-        cx: &App,
+        _cx: &App,
     ) -> Option<Rc<dyn acp_thread::AgentSessionSetTitle>> {
-        self.0.read_with(cx, |agent, _cx| {
-            agent
-                .sessions
-                .get(session_id)
-                .filter(|s| !s.thread.read(cx).is_subagent())
-                .map(|session| {
-                    Rc::new(NativeAgentSessionSetTitle {
-                        thread: session.thread.clone(),
-                    }) as _
-                })
-        })
+        Some(Rc::new(NativeAgentSessionSetTitle {
+            connection: self.clone(),
+            session_id: session_id.clone(),
+        }) as _)
     }
 
     fn session_list(&self, cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
@@ -1466,12 +1459,21 @@ impl NativeAgentSessionList {
     }
 
     fn to_session_info(entry: DbThreadMetadata) -> AgentSessionInfo {
+        let meta = entry.worktree_branch.map(|branch| {
+            let mut map = serde_json::Map::new();
+            map.insert(
+                "worktree_branch".to_string(),
+                serde_json::Value::String(branch),
+            );
+            map
+        });
+
         AgentSessionInfo {
             session_id: entry.id,
             cwd: None,
             title: Some(entry.title),
             updated_at: Some(entry.updated_at),
-            meta: None,
+            meta,
         }
     }
 
@@ -1566,13 +1568,17 @@ impl acp_thread::AgentSessionRetry for NativeAgentSessionRetry {
 }
 
 struct NativeAgentSessionSetTitle {
-    thread: Entity<Thread>,
+    connection: NativeAgentConnection,
+    session_id: acp::SessionId,
 }
 
 impl acp_thread::AgentSessionSetTitle for NativeAgentSessionSetTitle {
     fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>> {
-        self.thread
-            .update(cx, |thread, cx| thread.set_title(title, cx));
+        let Some(session) = self.connection.0.read(cx).sessions.get(&self.session_id) else {
+            return Task::ready(Err(anyhow!("session not found")));
+        };
+        let thread = session.thread.clone();
+        thread.update(cx, |thread, cx| thread.set_title(title, cx));
         Task::ready(Ok(()))
     }
 }

crates/agent/src/db.rs

@@ -23,13 +23,25 @@ pub type DbMessage = crate::Message;
 pub type DbSummary = crate::legacy_thread::DetailedSummaryState;
 pub type DbLanguageModel = crate::legacy_thread::SerializedLanguageModel;
 
+/// Metadata about the git worktree associated with an agent thread.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AgentGitWorktreeInfo {
+    /// The branch name in the git worktree.
+    pub branch: String,
+    /// Absolute path to the git worktree on disk.
+    pub worktree_path: std::path::PathBuf,
+    /// The base branch/commit the worktree was created from.
+    pub base_ref: String,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DbThreadMetadata {
     pub id: acp::SessionId,
     pub parent_session_id: Option<acp::SessionId>,
     #[serde(alias = "summary")]
     pub title: SharedString,
     pub updated_at: DateTime<Utc>,
+    pub worktree_branch: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -53,6 +65,8 @@ pub struct DbThread {
     pub imported: bool,
     #[serde(default)]
     pub subagent_context: Option<crate::SubagentContext>,
+    #[serde(default)]
+    pub git_worktree_info: Option<AgentGitWorktreeInfo>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -91,6 +105,7 @@ impl SharedThread {
             profile: None,
             imported: true,
             subagent_context: None,
+            git_worktree_info: None,
         }
     }
 
@@ -265,6 +280,7 @@ impl DbThread {
             profile: thread.profile,
             imported: false,
             subagent_context: None,
+            git_worktree_info: None,
         })
     }
 }
@@ -369,6 +385,13 @@ impl ThreadsDatabase {
             s().ok();
         }
 
+        if let Ok(mut s) = connection.exec(indoc! {"
+            ALTER TABLE threads ADD COLUMN worktree_branch TEXT
+        "})
+        {
+            s().ok();
+        }
+
         let db = Self {
             executor,
             connection: Arc::new(Mutex::new(connection)),
@@ -397,6 +420,10 @@ impl ThreadsDatabase {
             .subagent_context
             .as_ref()
             .map(|ctx| ctx.parent_thread_id.0.clone());
+        let worktree_branch = thread
+            .git_worktree_info
+            .as_ref()
+            .map(|info| info.branch.clone());
         let json_data = serde_json::to_string(&SerializedThread {
             thread,
             version: DbThread::VERSION,
@@ -408,11 +435,19 @@ impl ThreadsDatabase {
         let data_type = DataType::Zstd;
         let data = compressed;
 
-        let mut insert = connection.exec_bound::<(Arc<str>, Option<Arc<str>>, String, String, DataType, Vec<u8>)>(indoc! {"
-            INSERT OR REPLACE INTO threads (id, parent_id, summary, updated_at, data_type, data) VALUES (?, ?, ?, ?, ?, ?)
+        let mut insert = connection.exec_bound::<(Arc<str>, Option<Arc<str>>, Option<String>, String, String, DataType, Vec<u8>)>(indoc! {"
+            INSERT OR REPLACE INTO threads (id, parent_id, worktree_branch, summary, updated_at, data_type, data) VALUES (?, ?, ?, ?, ?, ?, ?)
         "})?;
 
-        insert((id.0, parent_id, title, updated_at, data_type, data))?;
+        insert((
+            id.0,
+            parent_id,
+            worktree_branch,
+            title,
+            updated_at,
+            data_type,
+            data,
+        ))?;
 
         Ok(())
     }
@@ -424,19 +459,20 @@ impl ThreadsDatabase {
             let connection = connection.lock();
 
             let mut select = connection
-                .select_bound::<(), (Arc<str>, Option<Arc<str>>, String, String)>(indoc! {"
-                SELECT id, parent_id, summary, updated_at FROM threads ORDER BY updated_at DESC
+                .select_bound::<(), (Arc<str>, Option<Arc<str>>, Option<String>, String, String)>(indoc! {"
+                SELECT id, parent_id, worktree_branch, summary, updated_at FROM threads ORDER BY updated_at DESC
             "})?;
 
             let rows = select(())?;
             let mut threads = Vec::new();
 
-            for (id, parent_id, summary, updated_at) in rows {
+            for (id, parent_id, worktree_branch, summary, updated_at) in rows {
                 threads.push(DbThreadMetadata {
                     id: acp::SessionId::new(id),
                     parent_session_id: parent_id.map(acp::SessionId::new),
                     title: summary.into(),
                     updated_at: DateTime::parse_from_rfc3339(&updated_at)?.with_timezone(&Utc),
+                    worktree_branch,
                 });
             }
 
@@ -570,6 +606,7 @@ mod tests {
             profile: None,
             imported: false,
             subagent_context: None,
+            git_worktree_info: None,
         }
     }
 
@@ -713,4 +750,96 @@ mod tests {
             "Regular threads should have no subagent_context"
         );
     }
+
+    #[gpui::test]
+    async fn test_git_worktree_info_roundtrip(cx: &mut TestAppContext) {
+        let database = ThreadsDatabase::new(cx.executor()).unwrap();
+
+        let thread_id = session_id("worktree-thread");
+        let mut thread = make_thread(
+            "Worktree Thread",
+            Utc.with_ymd_and_hms(2024, 6, 15, 12, 0, 0).unwrap(),
+        );
+        thread.git_worktree_info = Some(AgentGitWorktreeInfo {
+            branch: "zed/agent/a4Xiu".to_string(),
+            worktree_path: std::path::PathBuf::from(
+                "/tmp/agent-worktrees/my-project/zed/agent/a4Xiu",
+            ),
+            base_ref: "main".to_string(),
+        });
+
+        database
+            .save_thread(thread_id.clone(), thread)
+            .await
+            .unwrap();
+
+        let loaded = database
+            .load_thread(thread_id)
+            .await
+            .unwrap()
+            .expect("thread should exist");
+
+        let info = loaded
+            .git_worktree_info
+            .expect("git_worktree_info should be restored");
+        assert_eq!(info.branch, "zed/agent/a4Xiu");
+        assert_eq!(
+            info.worktree_path,
+            std::path::PathBuf::from("/tmp/agent-worktrees/my-project/zed/agent/a4Xiu")
+        );
+        assert_eq!(info.base_ref, "main");
+    }
+
+    #[gpui::test]
+    async fn test_session_list_includes_worktree_meta(cx: &mut TestAppContext) {
+        let database = ThreadsDatabase::new(cx.executor()).unwrap();
+
+        // Save a thread with worktree info
+        let worktree_id = session_id("wt-thread");
+        let mut worktree_thread = make_thread(
+            "With Worktree",
+            Utc.with_ymd_and_hms(2024, 6, 15, 12, 0, 0).unwrap(),
+        );
+        worktree_thread.git_worktree_info = Some(AgentGitWorktreeInfo {
+            branch: "zed/agent/bR9kz".to_string(),
+            worktree_path: std::path::PathBuf::from("/tmp/worktrees/bR9kz"),
+            base_ref: "develop".to_string(),
+        });
+
+        database
+            .save_thread(worktree_id.clone(), worktree_thread)
+            .await
+            .unwrap();
+
+        // Save a thread without worktree info
+        let plain_id = session_id("plain-thread");
+        let plain_thread = make_thread(
+            "Without Worktree",
+            Utc.with_ymd_and_hms(2024, 6, 15, 11, 0, 0).unwrap(),
+        );
+
+        database
+            .save_thread(plain_id.clone(), plain_thread)
+            .await
+            .unwrap();
+
+        // List threads and verify worktree_branch is populated correctly
+        let threads = database.list_threads().await.unwrap();
+        assert_eq!(threads.len(), 2);
+
+        let wt_entry = threads
+            .iter()
+            .find(|t| t.id == worktree_id)
+            .expect("should find worktree thread");
+        assert_eq!(wt_entry.worktree_branch.as_deref(), Some("zed/agent/bR9kz"));
+
+        let plain_entry = threads
+            .iter()
+            .find(|t| t.id == plain_id)
+            .expect("should find plain thread");
+        assert!(
+            plain_entry.worktree_branch.is_none(),
+            "plain thread should have no worktree_branch"
+        );
+    }
 }