thread_metadata_store.rs

  1use std::{path::Path, sync::Arc};
  2
  3use acp_thread::AgentSessionInfo;
  4use agent::{ThreadStore, ZED_AGENT_ID};
  5use agent_client_protocol as acp;
  6use anyhow::Result;
  7use chrono::{DateTime, Utc};
  8use collections::HashMap;
  9use db::{
 10    sqlez::{
 11        bindable::Column, domain::Domain, statement::Statement,
 12        thread_safe_connection::ThreadSafeConnection,
 13    },
 14    sqlez_macros::sql,
 15};
 16use feature_flags::{AgentV2FeatureFlag, FeatureFlagAppExt};
 17use gpui::{AppContext as _, Entity, Global, Subscription, Task};
 18use project::AgentId;
 19use ui::{App, Context, SharedString};
 20use workspace::PathList;
 21
 22pub fn init(cx: &mut App) {
 23    SidebarThreadMetadataStore::init_global(cx);
 24
 25    if cx.has_flag::<AgentV2FeatureFlag>() {
 26        migrate_thread_metadata(cx);
 27    }
 28    cx.observe_flag::<AgentV2FeatureFlag, _>(|has_flag, cx| {
 29        if has_flag {
 30            migrate_thread_metadata(cx);
 31        }
 32    })
 33    .detach();
 34}
 35
 36/// Migrate existing thread metadata from native agent thread store to the new metadata storage.
 37///
 38/// TODO: Remove this after N weeks of shipping the sidebar
 39fn migrate_thread_metadata(cx: &mut App) {
 40    SidebarThreadMetadataStore::global(cx).update(cx, |store, cx| {
 41        let list = store.list(cx);
 42        cx.spawn(async move |this, cx| {
 43            let Ok(list) = list.await else {
 44                return;
 45            };
 46            if list.is_empty() {
 47                this.update(cx, |this, cx| {
 48                    let metadata = ThreadStore::global(cx)
 49                        .read(cx)
 50                        .entries()
 51                        .map(|entry| ThreadMetadata {
 52                            session_id: entry.id,
 53                            agent_id: None,
 54                            title: entry.title,
 55                            updated_at: entry.updated_at,
 56                            created_at: entry.created_at,
 57                            folder_paths: entry.folder_paths,
 58                        })
 59                        .collect::<Vec<_>>();
 60                    for entry in metadata {
 61                        this.save(entry, cx).detach_and_log_err(cx);
 62                    }
 63                })
 64                .ok();
 65            }
 66        })
 67        .detach();
 68    });
 69}
 70
 71struct GlobalThreadMetadataStore(Entity<SidebarThreadMetadataStore>);
 72impl Global for GlobalThreadMetadataStore {}
 73
 74/// Lightweight metadata for any thread (native or ACP), enough to populate
 75/// the sidebar list and route to the correct load path when clicked.
 76#[derive(Debug, Clone)]
 77pub struct ThreadMetadata {
 78    pub session_id: acp::SessionId,
 79    /// `None` for native Zed threads, `Some("claude-code")` etc. for ACP agents.
 80    pub agent_id: Option<AgentId>,
 81    pub title: SharedString,
 82    pub updated_at: DateTime<Utc>,
 83    pub created_at: Option<DateTime<Utc>>,
 84    pub folder_paths: PathList,
 85}
 86
 87impl ThreadMetadata {
 88    pub fn from_session_info(agent_id: AgentId, session: &AgentSessionInfo) -> Self {
 89        let session_id = session.session_id.clone();
 90        let title = session.title.clone().unwrap_or_default();
 91        let updated_at = session.updated_at.unwrap_or_else(|| Utc::now());
 92        let created_at = session.created_at.unwrap_or(updated_at);
 93        let folder_paths = session.work_dirs.clone().unwrap_or_default();
 94        let agent_id = if agent_id.as_ref() == ZED_AGENT_ID.as_ref() {
 95            None
 96        } else {
 97            Some(agent_id)
 98        };
 99        Self {
100            session_id,
101            agent_id,
102            title,
103            updated_at,
104            created_at: Some(created_at),
105            folder_paths,
106        }
107    }
108
109    pub fn from_thread(thread: &Entity<acp_thread::AcpThread>, cx: &App) -> Self {
110        let thread_ref = thread.read(cx);
111        let session_id = thread_ref.session_id().clone();
112        let title = thread_ref.title();
113        let updated_at = Utc::now();
114
115        let agent_id = thread_ref.connection().agent_id();
116
117        let agent_id = if agent_id.as_ref() == ZED_AGENT_ID.as_ref() {
118            None
119        } else {
120            Some(agent_id)
121        };
122
123        let folder_paths = {
124            let project = thread_ref.project().read(cx);
125            let paths: Vec<Arc<Path>> = project
126                .visible_worktrees(cx)
127                .map(|worktree| worktree.read(cx).abs_path())
128                .collect();
129            PathList::new(&paths)
130        };
131
132        Self {
133            session_id,
134            agent_id,
135            title,
136            created_at: Some(updated_at), // handled by db `ON CONFLICT`
137            updated_at,
138            folder_paths,
139        }
140    }
141}
142
143/// The store holds all metadata needed to show threads in the sidebar.
144/// Effectively, all threads stored in here are "non-archived".
145///
146/// Automatically listens to AcpThread events and updates metadata if it has changed.
147pub struct SidebarThreadMetadataStore {
148    db: ThreadMetadataDb,
149    session_subscriptions: HashMap<acp::SessionId, Subscription>,
150}
151
152impl SidebarThreadMetadataStore {
153    #[cfg(not(any(test, feature = "test-support")))]
154    pub fn init_global(cx: &mut App) {
155        if cx.has_global::<Self>() {
156            return;
157        }
158
159        let db = ThreadMetadataDb::global(cx);
160        let thread_store = cx.new(|cx| Self::new(db, cx));
161        cx.set_global(GlobalThreadMetadataStore(thread_store));
162    }
163
164    #[cfg(any(test, feature = "test-support"))]
165    pub fn init_global(cx: &mut App) {
166        let thread = std::thread::current();
167        let test_name = thread.name().unwrap_or("unknown_test");
168        let db_name = format!("THREAD_METADATA_DB_{}", test_name);
169        let db = smol::block_on(db::open_test_db::<ThreadMetadataDb>(&db_name));
170        let thread_store = cx.new(|cx| Self::new(ThreadMetadataDb(db), cx));
171        cx.set_global(GlobalThreadMetadataStore(thread_store));
172    }
173
174    pub fn try_global(cx: &App) -> Option<Entity<Self>> {
175        cx.try_global::<GlobalThreadMetadataStore>()
176            .map(|store| store.0.clone())
177    }
178
179    pub fn global(cx: &App) -> Entity<Self> {
180        cx.global::<GlobalThreadMetadataStore>().0.clone()
181    }
182
183    pub fn list_ids(&self, cx: &App) -> Task<Result<Vec<acp::SessionId>>> {
184        let db = self.db.clone();
185        cx.background_spawn(async move {
186            let s = db.list_ids()?;
187            Ok(s)
188        })
189    }
190
191    pub fn list(&self, cx: &App) -> Task<Result<Vec<ThreadMetadata>>> {
192        let db = self.db.clone();
193        cx.background_spawn(async move {
194            let s = db.list()?;
195            Ok(s)
196        })
197    }
198
199    pub fn save(&mut self, metadata: ThreadMetadata, cx: &mut Context<Self>) -> Task<Result<()>> {
200        if !cx.has_flag::<AgentV2FeatureFlag>() {
201            return Task::ready(Ok(()));
202        }
203
204        let db = self.db.clone();
205        cx.spawn(async move |this, cx| {
206            db.save(metadata).await?;
207            this.update(cx, |_this, cx| cx.notify())
208        })
209    }
210
211    pub fn delete(
212        &mut self,
213        session_id: acp::SessionId,
214        cx: &mut Context<Self>,
215    ) -> Task<Result<()>> {
216        if !cx.has_flag::<AgentV2FeatureFlag>() {
217            return Task::ready(Ok(()));
218        }
219
220        let db = self.db.clone();
221        cx.spawn(async move |this, cx| {
222            db.delete(session_id).await?;
223            this.update(cx, |_this, cx| cx.notify())
224        })
225    }
226
227    fn new(db: ThreadMetadataDb, cx: &mut Context<Self>) -> Self {
228        let weak_store = cx.weak_entity();
229
230        cx.observe_new::<acp_thread::AcpThread>(move |thread, _window, cx| {
231            // Don't track subagent threads in the sidebar.
232            if thread.parent_session_id().is_some() {
233                return;
234            }
235
236            let thread_entity = cx.entity();
237
238            cx.on_release({
239                let weak_store = weak_store.clone();
240                move |thread, cx| {
241                    weak_store
242                        .update(cx, |store, _cx| {
243                            store.session_subscriptions.remove(thread.session_id());
244                        })
245                        .ok();
246                }
247            })
248            .detach();
249
250            weak_store
251                .update(cx, |this, cx| {
252                    let subscription = cx.subscribe(&thread_entity, Self::handle_thread_update);
253                    this.session_subscriptions
254                        .insert(thread.session_id().clone(), subscription);
255                })
256                .ok();
257        })
258        .detach();
259
260        Self {
261            db,
262            session_subscriptions: HashMap::default(),
263        }
264    }
265
266    fn handle_thread_update(
267        &mut self,
268        thread: Entity<acp_thread::AcpThread>,
269        event: &acp_thread::AcpThreadEvent,
270        cx: &mut Context<Self>,
271    ) {
272        // Don't track subagent threads in the sidebar.
273        if thread.read(cx).parent_session_id().is_some() {
274            return;
275        }
276
277        match event {
278            acp_thread::AcpThreadEvent::NewEntry
279            | acp_thread::AcpThreadEvent::EntryUpdated(_)
280            | acp_thread::AcpThreadEvent::TitleUpdated => {
281                let metadata = ThreadMetadata::from_thread(&thread, cx);
282                self.save(metadata, cx).detach_and_log_err(cx);
283            }
284            _ => {}
285        }
286    }
287}
288
289impl Global for SidebarThreadMetadataStore {}
290
291struct ThreadMetadataDb(ThreadSafeConnection);
292
293impl Domain for ThreadMetadataDb {
294    const NAME: &str = stringify!(ThreadMetadataDb);
295
296    const MIGRATIONS: &[&str] = &[sql!(
297        CREATE TABLE IF NOT EXISTS sidebar_threads(
298            session_id TEXT PRIMARY KEY,
299            agent_id TEXT,
300            title TEXT NOT NULL,
301            updated_at TEXT NOT NULL,
302            created_at TEXT,
303            folder_paths TEXT,
304            folder_paths_order TEXT
305        ) STRICT;
306    )];
307}
308
309db::static_connection!(ThreadMetadataDb, []);
310
311impl ThreadMetadataDb {
312    /// List all sidebar thread session IDs.
313    pub fn list_ids(&self) -> anyhow::Result<Vec<acp::SessionId>> {
314        self.select::<Arc<str>>("SELECT session_id FROM sidebar_threads")?()
315            .map(|ids| ids.into_iter().map(|id| acp::SessionId::new(id)).collect())
316    }
317
318    /// List all sidebar thread metadata, ordered by updated_at descending.
319    pub fn list(&self) -> anyhow::Result<Vec<ThreadMetadata>> {
320        self.select::<ThreadMetadata>(
321            "SELECT session_id, agent_id, title, updated_at, created_at, folder_paths, folder_paths_order \
322             FROM sidebar_threads \
323             ORDER BY updated_at DESC"
324        )?()
325    }
326
327    /// Upsert metadata for a thread.
328    pub async fn save(&self, row: ThreadMetadata) -> anyhow::Result<()> {
329        let id = row.session_id.0.clone();
330        let agent_id = row.agent_id.as_ref().map(|id| id.0.to_string());
331        let title = row.title.to_string();
332        let updated_at = row.updated_at.to_rfc3339();
333        let created_at = row.created_at.map(|dt| dt.to_rfc3339());
334        let serialized = row.folder_paths.serialize();
335        let (folder_paths, folder_paths_order) = if row.folder_paths.is_empty() {
336            (None, None)
337        } else {
338            (Some(serialized.paths), Some(serialized.order))
339        };
340
341        self.write(move |conn| {
342            let sql = "INSERT INTO sidebar_threads(session_id, agent_id, title, updated_at, created_at, folder_paths, folder_paths_order) \
343                       VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) \
344                       ON CONFLICT(session_id) DO UPDATE SET \
345                           agent_id = excluded.agent_id, \
346                           title = excluded.title, \
347                           updated_at = excluded.updated_at, \
348                           folder_paths = excluded.folder_paths, \
349                           folder_paths_order = excluded.folder_paths_order";
350            let mut stmt = Statement::prepare(conn, sql)?;
351            let mut i = stmt.bind(&id, 1)?;
352            i = stmt.bind(&agent_id, i)?;
353            i = stmt.bind(&title, i)?;
354            i = stmt.bind(&updated_at, i)?;
355            i = stmt.bind(&created_at, i)?;
356            i = stmt.bind(&folder_paths, i)?;
357            stmt.bind(&folder_paths_order, i)?;
358            stmt.exec()
359        })
360        .await
361    }
362
363    /// Delete metadata for a single thread.
364    pub async fn delete(&self, session_id: acp::SessionId) -> anyhow::Result<()> {
365        let id = session_id.0.clone();
366        self.write(move |conn| {
367            let mut stmt =
368                Statement::prepare(conn, "DELETE FROM sidebar_threads WHERE session_id = ?")?;
369            stmt.bind(&id, 1)?;
370            stmt.exec()
371        })
372        .await
373    }
374}
375
376impl Column for ThreadMetadata {
377    fn column(statement: &mut Statement, start_index: i32) -> anyhow::Result<(Self, i32)> {
378        let (id, next): (Arc<str>, i32) = Column::column(statement, start_index)?;
379        let (agent_id, next): (Option<String>, i32) = Column::column(statement, next)?;
380        let (title, next): (String, i32) = Column::column(statement, next)?;
381        let (updated_at_str, next): (String, i32) = Column::column(statement, next)?;
382        let (created_at_str, next): (Option<String>, i32) = Column::column(statement, next)?;
383        let (folder_paths_str, next): (Option<String>, i32) = Column::column(statement, next)?;
384        let (folder_paths_order_str, next): (Option<String>, i32) =
385            Column::column(statement, next)?;
386
387        let updated_at = DateTime::parse_from_rfc3339(&updated_at_str)?.with_timezone(&Utc);
388        let created_at = created_at_str
389            .as_deref()
390            .map(DateTime::parse_from_rfc3339)
391            .transpose()?
392            .map(|dt| dt.with_timezone(&Utc));
393
394        let folder_paths = folder_paths_str
395            .map(|paths| {
396                PathList::deserialize(&util::path_list::SerializedPathList {
397                    paths,
398                    order: folder_paths_order_str.unwrap_or_default(),
399                })
400            })
401            .unwrap_or_default();
402
403        Ok((
404            ThreadMetadata {
405                session_id: acp::SessionId::new(id),
406                agent_id: agent_id.map(|id| AgentId::new(id)),
407                title: title.into(),
408                updated_at,
409                created_at,
410                folder_paths,
411            },
412            next,
413        ))
414    }
415}
416
417#[cfg(test)]
418mod tests {
419    use super::*;
420    use acp_thread::{AgentConnection, StubAgentConnection};
421    use action_log::ActionLog;
422    use agent::DbThread;
423    use agent_client_protocol as acp;
424    use feature_flags::FeatureFlagAppExt;
425    use gpui::TestAppContext;
426    use project::FakeFs;
427    use project::Project;
428    use std::path::Path;
429    use std::rc::Rc;
430    use util::path_list::PathList;
431
432    fn make_db_thread(title: &str, updated_at: DateTime<Utc>) -> DbThread {
433        DbThread {
434            title: title.to_string().into(),
435            messages: Vec::new(),
436            updated_at,
437            detailed_summary: None,
438            initial_project_snapshot: None,
439            cumulative_token_usage: Default::default(),
440            request_token_usage: Default::default(),
441            model: None,
442            profile: None,
443            imported: false,
444            subagent_context: None,
445            speed: None,
446            thinking_enabled: false,
447            thinking_effort: None,
448            draft_prompt: None,
449            ui_scroll_position: None,
450        }
451    }
452
453    #[gpui::test]
454    async fn test_migrate_thread_metadata(cx: &mut TestAppContext) {
455        cx.update(|cx| {
456            ThreadStore::init_global(cx);
457            SidebarThreadMetadataStore::init_global(cx);
458        });
459
460        // Verify the list is empty before migration
461        let metadata_list = cx.update(|cx| {
462            let store = SidebarThreadMetadataStore::global(cx);
463            store.read(cx).list(cx)
464        });
465
466        let list = metadata_list.await.unwrap();
467        assert_eq!(list.len(), 0);
468
469        let now = Utc::now();
470
471        // Populate the native ThreadStore via save_thread
472        let save1 = cx.update(|cx| {
473            let thread_store = ThreadStore::global(cx);
474            thread_store.update(cx, |store, cx| {
475                store.save_thread(
476                    acp::SessionId::new("session-1"),
477                    make_db_thread("Thread 1", now),
478                    PathList::default(),
479                    cx,
480                )
481            })
482        });
483        save1.await.unwrap();
484        cx.run_until_parked();
485
486        let save2 = cx.update(|cx| {
487            let thread_store = ThreadStore::global(cx);
488            thread_store.update(cx, |store, cx| {
489                store.save_thread(
490                    acp::SessionId::new("session-2"),
491                    make_db_thread("Thread 2", now),
492                    PathList::default(),
493                    cx,
494                )
495            })
496        });
497        save2.await.unwrap();
498        cx.run_until_parked();
499
500        // Run migration
501        cx.update(|cx| {
502            migrate_thread_metadata(cx);
503        });
504
505        cx.run_until_parked();
506
507        // Verify the metadata was migrated
508        let metadata_list = cx.update(|cx| {
509            let store = SidebarThreadMetadataStore::global(cx);
510            store.read(cx).list(cx)
511        });
512
513        let list = metadata_list.await.unwrap();
514        assert_eq!(list.len(), 2);
515
516        let metadata1 = list
517            .iter()
518            .find(|m| m.session_id.0.as_ref() == "session-1")
519            .expect("session-1 should be in migrated metadata");
520        assert_eq!(metadata1.title.as_ref(), "Thread 1");
521        assert!(metadata1.agent_id.is_none());
522
523        let metadata2 = list
524            .iter()
525            .find(|m| m.session_id.0.as_ref() == "session-2")
526            .expect("session-2 should be in migrated metadata");
527        assert_eq!(metadata2.title.as_ref(), "Thread 2");
528        assert!(metadata2.agent_id.is_none());
529    }
530
531    #[gpui::test]
532    async fn test_migrate_thread_metadata_skips_when_data_exists(cx: &mut TestAppContext) {
533        cx.update(|cx| {
534            ThreadStore::init_global(cx);
535            SidebarThreadMetadataStore::init_global(cx);
536        });
537
538        // Pre-populate the metadata store with existing data
539        let existing_metadata = ThreadMetadata {
540            session_id: acp::SessionId::new("existing-session"),
541            agent_id: None,
542            title: "Existing Thread".into(),
543            updated_at: Utc::now(),
544            created_at: Some(Utc::now()),
545            folder_paths: PathList::default(),
546        };
547
548        cx.update(|cx| {
549            let store = SidebarThreadMetadataStore::global(cx);
550            store.update(cx, |store, cx| {
551                store.save(existing_metadata, cx).detach();
552            });
553        });
554
555        cx.run_until_parked();
556
557        // Add an entry to native thread store that should NOT be migrated
558        let save_task = cx.update(|cx| {
559            let thread_store = ThreadStore::global(cx);
560            thread_store.update(cx, |store, cx| {
561                store.save_thread(
562                    acp::SessionId::new("native-session"),
563                    make_db_thread("Native Thread", Utc::now()),
564                    PathList::default(),
565                    cx,
566                )
567            })
568        });
569        save_task.await.unwrap();
570        cx.run_until_parked();
571
572        // Run migration - should skip because metadata store is not empty
573        cx.update(|cx| {
574            migrate_thread_metadata(cx);
575        });
576
577        cx.run_until_parked();
578
579        // Verify only the existing metadata is present (migration was skipped)
580        let metadata_list = cx.update(|cx| {
581            let store = SidebarThreadMetadataStore::global(cx);
582            store.read(cx).list(cx)
583        });
584
585        let list = metadata_list.await.unwrap();
586        assert_eq!(list.len(), 1);
587        assert_eq!(list[0].session_id.0.as_ref(), "existing-session");
588    }
589
590    #[gpui::test]
591    async fn test_subagent_threads_excluded_from_sidebar_metadata(cx: &mut TestAppContext) {
592        cx.update(|cx| {
593            let settings_store = settings::SettingsStore::test(cx);
594            cx.set_global(settings_store);
595            cx.update_flags(true, vec!["agent-v2".to_string()]);
596            ThreadStore::init_global(cx);
597            SidebarThreadMetadataStore::init_global(cx);
598        });
599
600        let fs = FakeFs::new(cx.executor());
601        let project = Project::test(fs, None::<&Path>, cx).await;
602        let connection = Rc::new(StubAgentConnection::new());
603
604        // Create a regular (non-subagent) AcpThread.
605        let regular_thread = cx
606            .update(|cx| {
607                connection
608                    .clone()
609                    .new_session(project.clone(), PathList::default(), cx)
610            })
611            .await
612            .unwrap();
613
614        let regular_session_id = cx.read(|cx| regular_thread.read(cx).session_id().clone());
615
616        // Set a title on the regular thread to trigger a save via handle_thread_update.
617        cx.update(|cx| {
618            regular_thread.update(cx, |thread, cx| {
619                thread.set_title("Regular Thread".into(), cx).detach();
620            });
621        });
622        cx.run_until_parked();
623
624        // Create a subagent AcpThread
625        let subagent_session_id = acp::SessionId::new("subagent-session");
626        let subagent_thread = cx.update(|cx| {
627            let action_log = cx.new(|_| ActionLog::new(project.clone()));
628            cx.new(|cx| {
629                acp_thread::AcpThread::new(
630                    Some(regular_session_id.clone()),
631                    "Subagent Thread",
632                    None,
633                    connection.clone(),
634                    project.clone(),
635                    action_log,
636                    subagent_session_id.clone(),
637                    watch::Receiver::constant(acp::PromptCapabilities::new()),
638                    cx,
639                )
640            })
641        });
642
643        // Set a title on the subagent thread to trigger handle_thread_update.
644        cx.update(|cx| {
645            subagent_thread.update(cx, |thread, cx| {
646                thread
647                    .set_title("Subagent Thread Title".into(), cx)
648                    .detach();
649            });
650        });
651        cx.run_until_parked();
652
653        // List all metadata from the store.
654        let metadata_list = cx.update(|cx| {
655            let store = SidebarThreadMetadataStore::global(cx);
656            store.read(cx).list(cx)
657        });
658
659        let list = metadata_list.await.unwrap();
660
661        // The subagent thread should NOT appear in the sidebar metadata.
662        // Only the regular thread should be listed.
663        assert_eq!(
664            list.len(),
665            1,
666            "Expected only the regular thread in sidebar metadata, \
667             but found {} entries (subagent threads are leaking into the sidebar)",
668            list.len(),
669        );
670        assert_eq!(list[0].session_id, regular_session_id);
671        assert_eq!(list[0].title.as_ref(), "Regular Thread");
672    }
673}