1use std::borrow::Cow;
2use std::cell::{Ref, RefCell};
3use std::path::{Path, PathBuf};
4use std::rc::Rc;
5use std::sync::Arc;
6
7use anyhow::{Context as _, Result, anyhow};
8use assistant_settings::{AgentProfile, AgentProfileId, AssistantSettings};
9use assistant_tool::{ToolId, ToolSource, ToolWorkingSet};
10use chrono::{DateTime, Utc};
11use collections::HashMap;
12use context_server::manager::ContextServerManager;
13use context_server::{ContextServerFactoryRegistry, ContextServerTool};
14use fs::Fs;
15use futures::FutureExt as _;
16use futures::future::{self, BoxFuture, Shared};
17use gpui::{
18 App, BackgroundExecutor, Context, Entity, EventEmitter, Global, ReadGlobal, SharedString,
19 Subscription, Task, prelude::*,
20};
21use heed::Database;
22use heed::types::SerdeBincode;
23use language_model::{LanguageModelToolUseId, Role, TokenUsage};
24use project::{Project, Worktree};
25use prompt_store::{ProjectContext, PromptBuilder, RulesFileContext, WorktreeContext};
26use serde::{Deserialize, Serialize};
27use settings::{Settings as _, SettingsStore};
28use util::ResultExt as _;
29
30use crate::thread::{
31 DetailedSummaryState, ExceededWindowError, MessageId, ProjectSnapshot, Thread, ThreadId,
32};
33
34const RULES_FILE_NAMES: [&'static str; 6] = [
35 ".rules",
36 ".cursorrules",
37 ".windsurfrules",
38 ".clinerules",
39 ".github/copilot-instructions.md",
40 "CLAUDE.md",
41];
42
43pub fn init(cx: &mut App) {
44 ThreadsDatabase::init(cx);
45}
46
47/// A system prompt shared by all threads created by this ThreadStore
48#[derive(Clone, Default)]
49pub struct SharedProjectContext(Rc<RefCell<Option<ProjectContext>>>);
50
51impl SharedProjectContext {
52 pub fn borrow(&self) -> Ref<Option<ProjectContext>> {
53 self.0.borrow()
54 }
55}
56
57pub struct ThreadStore {
58 project: Entity<Project>,
59 tools: Arc<ToolWorkingSet>,
60 prompt_builder: Arc<PromptBuilder>,
61 context_server_manager: Entity<ContextServerManager>,
62 context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,
63 threads: Vec<SerializedThreadMetadata>,
64 project_context: SharedProjectContext,
65 _subscriptions: Vec<Subscription>,
66}
67
68pub struct RulesLoadingError {
69 pub message: SharedString,
70}
71
72impl EventEmitter<RulesLoadingError> for ThreadStore {}
73
74impl ThreadStore {
75 pub fn load(
76 project: Entity<Project>,
77 tools: Arc<ToolWorkingSet>,
78 prompt_builder: Arc<PromptBuilder>,
79 cx: &mut App,
80 ) -> Task<Entity<Self>> {
81 let thread_store = cx.new(|cx| Self::new(project, tools, prompt_builder, cx));
82 let reload = thread_store.update(cx, |store, cx| store.reload_system_prompt(cx));
83 cx.foreground_executor().spawn(async move {
84 reload.await;
85 thread_store
86 })
87 }
88
89 fn new(
90 project: Entity<Project>,
91 tools: Arc<ToolWorkingSet>,
92 prompt_builder: Arc<PromptBuilder>,
93 cx: &mut Context<Self>,
94 ) -> Self {
95 let context_server_factory_registry = ContextServerFactoryRegistry::default_global(cx);
96 let context_server_manager = cx.new(|cx| {
97 ContextServerManager::new(context_server_factory_registry, project.clone(), cx)
98 });
99 let settings_subscription =
100 cx.observe_global::<SettingsStore>(move |this: &mut Self, cx| {
101 this.load_default_profile(cx);
102 });
103 let project_subscription = cx.subscribe(&project, Self::handle_project_event);
104
105 let this = Self {
106 project,
107 tools,
108 prompt_builder,
109 context_server_manager,
110 context_server_tool_ids: HashMap::default(),
111 threads: Vec::new(),
112 project_context: SharedProjectContext::default(),
113 _subscriptions: vec![settings_subscription, project_subscription],
114 };
115 this.load_default_profile(cx);
116 this.register_context_server_handlers(cx);
117 this.reload(cx).detach_and_log_err(cx);
118 this
119 }
120
121 fn handle_project_event(
122 &mut self,
123 _project: Entity<Project>,
124 event: &project::Event,
125 cx: &mut Context<Self>,
126 ) {
127 match event {
128 project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {
129 self.reload_system_prompt(cx).detach();
130 }
131 project::Event::WorktreeUpdatedEntries(_, items) => {
132 if items.iter().any(|(path, _, _)| {
133 RULES_FILE_NAMES
134 .iter()
135 .any(|name| path.as_ref() == Path::new(name))
136 }) {
137 self.reload_system_prompt(cx).detach();
138 }
139 }
140 _ => {}
141 }
142 }
143
144 pub fn reload_system_prompt(&self, cx: &mut Context<Self>) -> Task<()> {
145 let project = self.project.read(cx);
146 let tasks = project
147 .visible_worktrees(cx)
148 .map(|worktree| {
149 Self::load_worktree_info_for_system_prompt(
150 project.fs().clone(),
151 worktree.read(cx),
152 cx,
153 )
154 })
155 .collect::<Vec<_>>();
156
157 cx.spawn(async move |this, cx| {
158 let results = futures::future::join_all(tasks).await;
159 let worktrees = results
160 .into_iter()
161 .map(|(worktree, rules_error)| {
162 if let Some(rules_error) = rules_error {
163 this.update(cx, |_, cx| cx.emit(rules_error)).ok();
164 }
165 worktree
166 })
167 .collect::<Vec<_>>();
168 this.update(cx, |this, _cx| {
169 *this.project_context.0.borrow_mut() = Some(ProjectContext::new(worktrees));
170 })
171 .ok();
172 })
173 }
174
175 fn load_worktree_info_for_system_prompt(
176 fs: Arc<dyn Fs>,
177 worktree: &Worktree,
178 cx: &App,
179 ) -> Task<(WorktreeContext, Option<RulesLoadingError>)> {
180 let root_name = worktree.root_name().into();
181 let abs_path = worktree.abs_path();
182
183 let rules_task = Self::load_worktree_rules_file(fs, worktree, cx);
184 let Some(rules_task) = rules_task else {
185 return Task::ready((
186 WorktreeContext {
187 root_name,
188 abs_path,
189 rules_file: None,
190 },
191 None,
192 ));
193 };
194
195 cx.spawn(async move |_| {
196 let (rules_file, rules_file_error) = match rules_task.await {
197 Ok(rules_file) => (Some(rules_file), None),
198 Err(err) => (
199 None,
200 Some(RulesLoadingError {
201 message: format!("{err}").into(),
202 }),
203 ),
204 };
205 let worktree_info = WorktreeContext {
206 root_name,
207 abs_path,
208 rules_file,
209 };
210 (worktree_info, rules_file_error)
211 })
212 }
213
214 fn load_worktree_rules_file(
215 fs: Arc<dyn Fs>,
216 worktree: &Worktree,
217 cx: &App,
218 ) -> Option<Task<Result<RulesFileContext>>> {
219 let selected_rules_file = RULES_FILE_NAMES
220 .into_iter()
221 .filter_map(|name| {
222 worktree
223 .entry_for_path(name)
224 .filter(|entry| entry.is_file())
225 .map(|entry| (entry.path.clone(), worktree.absolutize(&entry.path)))
226 })
227 .next();
228
229 // Note that Cline supports `.clinerules` being a directory, but that is not currently
230 // supported. This doesn't seem to occur often in GitHub repositories.
231 selected_rules_file.map(|(path_in_worktree, abs_path)| {
232 let fs = fs.clone();
233 cx.background_spawn(async move {
234 let abs_path = abs_path?;
235 let text = fs.load(&abs_path).await.with_context(|| {
236 format!("Failed to load assistant rules file {:?}", abs_path)
237 })?;
238 anyhow::Ok(RulesFileContext {
239 path_in_worktree,
240 abs_path: abs_path.into(),
241 text: text.trim().to_string(),
242 })
243 })
244 })
245 }
246
247 pub fn context_server_manager(&self) -> Entity<ContextServerManager> {
248 self.context_server_manager.clone()
249 }
250
251 pub fn tools(&self) -> Arc<ToolWorkingSet> {
252 self.tools.clone()
253 }
254
255 /// Returns the number of threads.
256 pub fn thread_count(&self) -> usize {
257 self.threads.len()
258 }
259
260 pub fn threads(&self) -> Vec<SerializedThreadMetadata> {
261 let mut threads = self.threads.iter().cloned().collect::<Vec<_>>();
262 threads.sort_unstable_by_key(|thread| std::cmp::Reverse(thread.updated_at));
263 threads
264 }
265
266 pub fn recent_threads(&self, limit: usize) -> Vec<SerializedThreadMetadata> {
267 self.threads().into_iter().take(limit).collect()
268 }
269
270 pub fn create_thread(&mut self, cx: &mut Context<Self>) -> Entity<Thread> {
271 cx.new(|cx| {
272 Thread::new(
273 self.project.clone(),
274 self.tools.clone(),
275 self.prompt_builder.clone(),
276 self.project_context.clone(),
277 cx,
278 )
279 })
280 }
281
282 pub fn open_thread(
283 &self,
284 id: &ThreadId,
285 cx: &mut Context<Self>,
286 ) -> Task<Result<Entity<Thread>>> {
287 let id = id.clone();
288 let database_future = ThreadsDatabase::global_future(cx);
289 cx.spawn(async move |this, cx| {
290 let database = database_future.await.map_err(|err| anyhow!(err))?;
291 let thread = database
292 .try_find_thread(id.clone())
293 .await?
294 .ok_or_else(|| anyhow!("no thread found with ID: {id:?}"))?;
295
296 let thread = this.update(cx, |this, cx| {
297 cx.new(|cx| {
298 Thread::deserialize(
299 id.clone(),
300 thread,
301 this.project.clone(),
302 this.tools.clone(),
303 this.prompt_builder.clone(),
304 this.project_context.clone(),
305 cx,
306 )
307 })
308 })?;
309
310 Ok(thread)
311 })
312 }
313
314 pub fn save_thread(&self, thread: &Entity<Thread>, cx: &mut Context<Self>) -> Task<Result<()>> {
315 let (metadata, serialized_thread) =
316 thread.update(cx, |thread, cx| (thread.id().clone(), thread.serialize(cx)));
317
318 let database_future = ThreadsDatabase::global_future(cx);
319 cx.spawn(async move |this, cx| {
320 let serialized_thread = serialized_thread.await?;
321 let database = database_future.await.map_err(|err| anyhow!(err))?;
322 database.save_thread(metadata, serialized_thread).await?;
323
324 this.update(cx, |this, cx| this.reload(cx))?.await
325 })
326 }
327
328 pub fn delete_thread(&mut self, id: &ThreadId, cx: &mut Context<Self>) -> Task<Result<()>> {
329 let id = id.clone();
330 let database_future = ThreadsDatabase::global_future(cx);
331 cx.spawn(async move |this, cx| {
332 let database = database_future.await.map_err(|err| anyhow!(err))?;
333 database.delete_thread(id.clone()).await?;
334
335 this.update(cx, |this, cx| {
336 this.threads.retain(|thread| thread.id != id);
337 cx.notify();
338 })
339 })
340 }
341
342 pub fn reload(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
343 let database_future = ThreadsDatabase::global_future(cx);
344 cx.spawn(async move |this, cx| {
345 let threads = database_future
346 .await
347 .map_err(|err| anyhow!(err))?
348 .list_threads()
349 .await?;
350
351 this.update(cx, |this, cx| {
352 this.threads = threads;
353 cx.notify();
354 })
355 })
356 }
357
358 fn load_default_profile(&self, cx: &Context<Self>) {
359 let assistant_settings = AssistantSettings::get_global(cx);
360
361 self.load_profile_by_id(&assistant_settings.default_profile, cx);
362 }
363
364 pub fn load_profile_by_id(&self, profile_id: &AgentProfileId, cx: &Context<Self>) {
365 let assistant_settings = AssistantSettings::get_global(cx);
366
367 if let Some(profile) = assistant_settings.profiles.get(profile_id) {
368 self.load_profile(profile, cx);
369 }
370 }
371
372 pub fn load_profile(&self, profile: &AgentProfile, cx: &Context<Self>) {
373 self.tools.disable_all_tools();
374 self.tools.enable(
375 ToolSource::Native,
376 &profile
377 .tools
378 .iter()
379 .filter_map(|(tool, enabled)| enabled.then(|| tool.clone()))
380 .collect::<Vec<_>>(),
381 );
382
383 if profile.enable_all_context_servers {
384 for context_server in self.context_server_manager.read(cx).all_servers() {
385 self.tools.enable_source(
386 ToolSource::ContextServer {
387 id: context_server.id().into(),
388 },
389 cx,
390 );
391 }
392 } else {
393 for (context_server_id, preset) in &profile.context_servers {
394 self.tools.enable(
395 ToolSource::ContextServer {
396 id: context_server_id.clone().into(),
397 },
398 &preset
399 .tools
400 .iter()
401 .filter_map(|(tool, enabled)| enabled.then(|| tool.clone()))
402 .collect::<Vec<_>>(),
403 )
404 }
405 }
406 }
407
408 fn register_context_server_handlers(&self, cx: &mut Context<Self>) {
409 cx.subscribe(
410 &self.context_server_manager.clone(),
411 Self::handle_context_server_event,
412 )
413 .detach();
414 }
415
416 fn handle_context_server_event(
417 &mut self,
418 context_server_manager: Entity<ContextServerManager>,
419 event: &context_server::manager::Event,
420 cx: &mut Context<Self>,
421 ) {
422 let tool_working_set = self.tools.clone();
423 match event {
424 context_server::manager::Event::ServerStarted { server_id } => {
425 if let Some(server) = context_server_manager.read(cx).get_server(server_id) {
426 let context_server_manager = context_server_manager.clone();
427 cx.spawn({
428 let server = server.clone();
429 let server_id = server_id.clone();
430 async move |this, cx| {
431 let Some(protocol) = server.client() else {
432 return;
433 };
434
435 if protocol.capable(context_server::protocol::ServerCapability::Tools) {
436 if let Some(tools) = protocol.list_tools().await.log_err() {
437 let tool_ids = tools
438 .tools
439 .into_iter()
440 .map(|tool| {
441 log::info!(
442 "registering context server tool: {:?}",
443 tool.name
444 );
445 tool_working_set.insert(Arc::new(
446 ContextServerTool::new(
447 context_server_manager.clone(),
448 server.id(),
449 tool,
450 ),
451 ))
452 })
453 .collect::<Vec<_>>();
454
455 this.update(cx, |this, cx| {
456 this.context_server_tool_ids.insert(server_id, tool_ids);
457 this.load_default_profile(cx);
458 })
459 .log_err();
460 }
461 }
462 }
463 })
464 .detach();
465 }
466 }
467 context_server::manager::Event::ServerStopped { server_id } => {
468 if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
469 tool_working_set.remove(&tool_ids);
470 self.load_default_profile(cx);
471 }
472 }
473 }
474 }
475}
476
477#[derive(Debug, Clone, Serialize, Deserialize)]
478pub struct SerializedThreadMetadata {
479 pub id: ThreadId,
480 pub summary: SharedString,
481 pub updated_at: DateTime<Utc>,
482}
483
484#[derive(Serialize, Deserialize, Debug)]
485pub struct SerializedThread {
486 pub version: String,
487 pub summary: SharedString,
488 pub updated_at: DateTime<Utc>,
489 pub messages: Vec<SerializedMessage>,
490 #[serde(default)]
491 pub initial_project_snapshot: Option<Arc<ProjectSnapshot>>,
492 #[serde(default)]
493 pub cumulative_token_usage: TokenUsage,
494 #[serde(default)]
495 pub detailed_summary_state: DetailedSummaryState,
496 #[serde(default)]
497 pub exceeded_window_error: Option<ExceededWindowError>,
498}
499
500impl SerializedThread {
501 pub const VERSION: &'static str = "0.1.0";
502
503 pub fn from_json(json: &[u8]) -> Result<Self> {
504 let saved_thread_json = serde_json::from_slice::<serde_json::Value>(json)?;
505 match saved_thread_json.get("version") {
506 Some(serde_json::Value::String(version)) => match version.as_str() {
507 SerializedThread::VERSION => Ok(serde_json::from_value::<SerializedThread>(
508 saved_thread_json,
509 )?),
510 _ => Err(anyhow!(
511 "unrecognized serialized thread version: {}",
512 version
513 )),
514 },
515 None => {
516 let saved_thread =
517 serde_json::from_value::<LegacySerializedThread>(saved_thread_json)?;
518 Ok(saved_thread.upgrade())
519 }
520 version => Err(anyhow!(
521 "unrecognized serialized thread version: {:?}",
522 version
523 )),
524 }
525 }
526}
527
528#[derive(Debug, Serialize, Deserialize)]
529pub struct SerializedMessage {
530 pub id: MessageId,
531 pub role: Role,
532 #[serde(default)]
533 pub segments: Vec<SerializedMessageSegment>,
534 #[serde(default)]
535 pub tool_uses: Vec<SerializedToolUse>,
536 #[serde(default)]
537 pub tool_results: Vec<SerializedToolResult>,
538 #[serde(default)]
539 pub context: String,
540}
541
542#[derive(Debug, Serialize, Deserialize)]
543#[serde(tag = "type")]
544pub enum SerializedMessageSegment {
545 #[serde(rename = "text")]
546 Text { text: String },
547 #[serde(rename = "thinking")]
548 Thinking { text: String },
549}
550
551#[derive(Debug, Serialize, Deserialize)]
552pub struct SerializedToolUse {
553 pub id: LanguageModelToolUseId,
554 pub name: SharedString,
555 pub input: serde_json::Value,
556}
557
558#[derive(Debug, Serialize, Deserialize)]
559pub struct SerializedToolResult {
560 pub tool_use_id: LanguageModelToolUseId,
561 pub is_error: bool,
562 pub content: Arc<str>,
563}
564
565#[derive(Serialize, Deserialize)]
566struct LegacySerializedThread {
567 pub summary: SharedString,
568 pub updated_at: DateTime<Utc>,
569 pub messages: Vec<LegacySerializedMessage>,
570 #[serde(default)]
571 pub initial_project_snapshot: Option<Arc<ProjectSnapshot>>,
572}
573
574impl LegacySerializedThread {
575 pub fn upgrade(self) -> SerializedThread {
576 SerializedThread {
577 version: SerializedThread::VERSION.to_string(),
578 summary: self.summary,
579 updated_at: self.updated_at,
580 messages: self.messages.into_iter().map(|msg| msg.upgrade()).collect(),
581 initial_project_snapshot: self.initial_project_snapshot,
582 cumulative_token_usage: TokenUsage::default(),
583 detailed_summary_state: DetailedSummaryState::default(),
584 exceeded_window_error: None,
585 }
586 }
587}
588
589#[derive(Debug, Serialize, Deserialize)]
590struct LegacySerializedMessage {
591 pub id: MessageId,
592 pub role: Role,
593 pub text: String,
594 #[serde(default)]
595 pub tool_uses: Vec<SerializedToolUse>,
596 #[serde(default)]
597 pub tool_results: Vec<SerializedToolResult>,
598}
599
600impl LegacySerializedMessage {
601 fn upgrade(self) -> SerializedMessage {
602 SerializedMessage {
603 id: self.id,
604 role: self.role,
605 segments: vec![SerializedMessageSegment::Text { text: self.text }],
606 tool_uses: self.tool_uses,
607 tool_results: self.tool_results,
608 context: String::new(),
609 }
610 }
611}
612
613struct GlobalThreadsDatabase(
614 Shared<BoxFuture<'static, Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>>,
615);
616
617impl Global for GlobalThreadsDatabase {}
618
619pub(crate) struct ThreadsDatabase {
620 executor: BackgroundExecutor,
621 env: heed::Env,
622 threads: Database<SerdeBincode<ThreadId>, SerializedThread>,
623}
624
625impl heed::BytesEncode<'_> for SerializedThread {
626 type EItem = SerializedThread;
627
628 fn bytes_encode(item: &Self::EItem) -> Result<Cow<[u8]>, heed::BoxedError> {
629 serde_json::to_vec(item).map(Cow::Owned).map_err(Into::into)
630 }
631}
632
633impl<'a> heed::BytesDecode<'a> for SerializedThread {
634 type DItem = SerializedThread;
635
636 fn bytes_decode(bytes: &'a [u8]) -> Result<Self::DItem, heed::BoxedError> {
637 // We implement this type manually because we want to call `SerializedThread::from_json`,
638 // instead of the Deserialize trait implementation for `SerializedThread`.
639 SerializedThread::from_json(bytes).map_err(Into::into)
640 }
641}
642
643impl ThreadsDatabase {
644 fn global_future(
645 cx: &mut App,
646 ) -> Shared<BoxFuture<'static, Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>> {
647 GlobalThreadsDatabase::global(cx).0.clone()
648 }
649
650 fn init(cx: &mut App) {
651 let executor = cx.background_executor().clone();
652 let database_future = executor
653 .spawn({
654 let executor = executor.clone();
655 let database_path = paths::data_dir().join("threads/threads-db.1.mdb");
656 async move { ThreadsDatabase::new(database_path, executor) }
657 })
658 .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new)))
659 .boxed()
660 .shared();
661
662 cx.set_global(GlobalThreadsDatabase(database_future));
663 }
664
665 pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result<Self> {
666 std::fs::create_dir_all(&path)?;
667
668 const ONE_GB_IN_BYTES: usize = 1024 * 1024 * 1024;
669 let env = unsafe {
670 heed::EnvOpenOptions::new()
671 .map_size(ONE_GB_IN_BYTES)
672 .max_dbs(1)
673 .open(path)?
674 };
675
676 let mut txn = env.write_txn()?;
677 let threads = env.create_database(&mut txn, Some("threads"))?;
678 txn.commit()?;
679
680 Ok(Self {
681 executor,
682 env,
683 threads,
684 })
685 }
686
687 pub fn list_threads(&self) -> Task<Result<Vec<SerializedThreadMetadata>>> {
688 let env = self.env.clone();
689 let threads = self.threads;
690
691 self.executor.spawn(async move {
692 let txn = env.read_txn()?;
693 let mut iter = threads.iter(&txn)?;
694 let mut threads = Vec::new();
695 while let Some((key, value)) = iter.next().transpose()? {
696 threads.push(SerializedThreadMetadata {
697 id: key,
698 summary: value.summary,
699 updated_at: value.updated_at,
700 });
701 }
702
703 Ok(threads)
704 })
705 }
706
707 pub fn try_find_thread(&self, id: ThreadId) -> Task<Result<Option<SerializedThread>>> {
708 let env = self.env.clone();
709 let threads = self.threads;
710
711 self.executor.spawn(async move {
712 let txn = env.read_txn()?;
713 let thread = threads.get(&txn, &id)?;
714 Ok(thread)
715 })
716 }
717
718 pub fn save_thread(&self, id: ThreadId, thread: SerializedThread) -> Task<Result<()>> {
719 let env = self.env.clone();
720 let threads = self.threads;
721
722 self.executor.spawn(async move {
723 let mut txn = env.write_txn()?;
724 threads.put(&mut txn, &id, &thread)?;
725 txn.commit()?;
726 Ok(())
727 })
728 }
729
730 pub fn delete_thread(&self, id: ThreadId) -> Task<Result<()>> {
731 let env = self.env.clone();
732 let threads = self.threads;
733
734 self.executor.spawn(async move {
735 let mut txn = env.write_txn()?;
736 threads.delete(&mut txn, &id)?;
737 txn.commit()?;
738 Ok(())
739 })
740 }
741}