1use std::path::PathBuf;
2use std::sync::Arc;
3
4use anyhow::{anyhow, Result};
5use assistant_tool::{ToolId, ToolWorkingSet};
6use chrono::{DateTime, Utc};
7use collections::HashMap;
8use context_server::manager::ContextServerManager;
9use context_server::{ContextServerFactoryRegistry, ContextServerTool};
10use futures::future::{self, BoxFuture, Shared};
11use futures::FutureExt as _;
12use gpui::{
13 prelude::*, App, BackgroundExecutor, Context, Entity, Global, ReadGlobal, SharedString, Task,
14};
15use heed::types::{SerdeBincode, SerdeJson};
16use heed::Database;
17use language_model::{LanguageModelToolUseId, Role};
18use project::Project;
19use prompt_store::PromptBuilder;
20use serde::{Deserialize, Serialize};
21use util::ResultExt as _;
22
23use crate::thread::{MessageId, ProjectSnapshot, Thread, ThreadId};
24
25pub fn init(cx: &mut App) {
26 ThreadsDatabase::init(cx);
27}
28
29pub struct ThreadStore {
30 project: Entity<Project>,
31 tools: Arc<ToolWorkingSet>,
32 prompt_builder: Arc<PromptBuilder>,
33 context_server_manager: Entity<ContextServerManager>,
34 context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,
35 threads: Vec<SerializedThreadMetadata>,
36}
37
38impl ThreadStore {
39 pub fn new(
40 project: Entity<Project>,
41 tools: Arc<ToolWorkingSet>,
42 prompt_builder: Arc<PromptBuilder>,
43 cx: &mut App,
44 ) -> Result<Entity<Self>> {
45 let this = cx.new(|cx| {
46 let context_server_factory_registry = ContextServerFactoryRegistry::default_global(cx);
47 let context_server_manager = cx.new(|cx| {
48 ContextServerManager::new(context_server_factory_registry, project.clone(), cx)
49 });
50
51 let this = Self {
52 project,
53 tools,
54 prompt_builder,
55 context_server_manager,
56 context_server_tool_ids: HashMap::default(),
57 threads: Vec::new(),
58 };
59 this.register_context_server_handlers(cx);
60 this.reload(cx).detach_and_log_err(cx);
61
62 this
63 });
64
65 Ok(this)
66 }
67
68 pub fn context_server_manager(&self) -> Entity<ContextServerManager> {
69 self.context_server_manager.clone()
70 }
71
72 pub fn tools(&self) -> Arc<ToolWorkingSet> {
73 self.tools.clone()
74 }
75
76 /// Returns the number of threads.
77 pub fn thread_count(&self) -> usize {
78 self.threads.len()
79 }
80
81 pub fn threads(&self) -> Vec<SerializedThreadMetadata> {
82 let mut threads = self.threads.iter().cloned().collect::<Vec<_>>();
83 threads.sort_unstable_by_key(|thread| std::cmp::Reverse(thread.updated_at));
84 threads
85 }
86
87 pub fn recent_threads(&self, limit: usize) -> Vec<SerializedThreadMetadata> {
88 self.threads().into_iter().take(limit).collect()
89 }
90
91 pub fn create_thread(&mut self, cx: &mut Context<Self>) -> Entity<Thread> {
92 cx.new(|cx| {
93 Thread::new(
94 self.project.clone(),
95 self.tools.clone(),
96 self.prompt_builder.clone(),
97 cx,
98 )
99 })
100 }
101
102 pub fn open_thread(
103 &self,
104 id: &ThreadId,
105 cx: &mut Context<Self>,
106 ) -> Task<Result<Entity<Thread>>> {
107 let id = id.clone();
108 let database_future = ThreadsDatabase::global_future(cx);
109 cx.spawn(async move |this, cx| {
110 let database = database_future.await.map_err(|err| anyhow!(err))?;
111 let thread = database
112 .try_find_thread(id.clone())
113 .await?
114 .ok_or_else(|| anyhow!("no thread found with ID: {id:?}"))?;
115
116 this.update(cx, |this, cx| {
117 cx.new(|cx| {
118 Thread::deserialize(
119 id.clone(),
120 thread,
121 this.project.clone(),
122 this.tools.clone(),
123 this.prompt_builder.clone(),
124 cx,
125 )
126 })
127 })
128 })
129 }
130
131 pub fn save_thread(&self, thread: &Entity<Thread>, cx: &mut Context<Self>) -> Task<Result<()>> {
132 let (metadata, serialized_thread) =
133 thread.update(cx, |thread, cx| (thread.id().clone(), thread.serialize(cx)));
134
135 let database_future = ThreadsDatabase::global_future(cx);
136 cx.spawn(async move |this, cx| {
137 let serialized_thread = serialized_thread.await?;
138 let database = database_future.await.map_err(|err| anyhow!(err))?;
139 database.save_thread(metadata, serialized_thread).await?;
140
141 this.update(cx, |this, cx| this.reload(cx))?.await
142 })
143 }
144
145 pub fn delete_thread(&mut self, id: &ThreadId, cx: &mut Context<Self>) -> Task<Result<()>> {
146 let id = id.clone();
147 let database_future = ThreadsDatabase::global_future(cx);
148 cx.spawn(async move |this, cx| {
149 let database = database_future.await.map_err(|err| anyhow!(err))?;
150 database.delete_thread(id.clone()).await?;
151
152 this.update(cx, |this, _cx| {
153 this.threads.retain(|thread| thread.id != id)
154 })
155 })
156 }
157
158 pub fn reload(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
159 let database_future = ThreadsDatabase::global_future(cx);
160 cx.spawn(async move |this, cx| {
161 let threads = database_future
162 .await
163 .map_err(|err| anyhow!(err))?
164 .list_threads()
165 .await?;
166
167 this.update(cx, |this, cx| {
168 this.threads = threads;
169 cx.notify();
170 })
171 })
172 }
173
174 fn register_context_server_handlers(&self, cx: &mut Context<Self>) {
175 cx.subscribe(
176 &self.context_server_manager.clone(),
177 Self::handle_context_server_event,
178 )
179 .detach();
180 }
181
182 fn handle_context_server_event(
183 &mut self,
184 context_server_manager: Entity<ContextServerManager>,
185 event: &context_server::manager::Event,
186 cx: &mut Context<Self>,
187 ) {
188 let tool_working_set = self.tools.clone();
189 match event {
190 context_server::manager::Event::ServerStarted { server_id } => {
191 if let Some(server) = context_server_manager.read(cx).get_server(server_id) {
192 let context_server_manager = context_server_manager.clone();
193 cx.spawn({
194 let server = server.clone();
195 let server_id = server_id.clone();
196 async move |this, cx| {
197 let Some(protocol) = server.client() else {
198 return;
199 };
200
201 if protocol.capable(context_server::protocol::ServerCapability::Tools) {
202 if let Some(tools) = protocol.list_tools().await.log_err() {
203 let tool_ids = tools
204 .tools
205 .into_iter()
206 .map(|tool| {
207 log::info!(
208 "registering context server tool: {:?}",
209 tool.name
210 );
211 tool_working_set.insert(Arc::new(
212 ContextServerTool::new(
213 context_server_manager.clone(),
214 server.id(),
215 tool,
216 ),
217 ))
218 })
219 .collect::<Vec<_>>();
220
221 this.update(cx, |this, _cx| {
222 this.context_server_tool_ids.insert(server_id, tool_ids);
223 })
224 .log_err();
225 }
226 }
227 }
228 })
229 .detach();
230 }
231 }
232 context_server::manager::Event::ServerStopped { server_id } => {
233 if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
234 tool_working_set.remove(&tool_ids);
235 }
236 }
237 }
238 }
239}
240
241#[derive(Debug, Clone, Serialize, Deserialize)]
242pub struct SerializedThreadMetadata {
243 pub id: ThreadId,
244 pub summary: SharedString,
245 pub updated_at: DateTime<Utc>,
246}
247
248#[derive(Serialize, Deserialize)]
249pub struct SerializedThread {
250 pub summary: SharedString,
251 pub updated_at: DateTime<Utc>,
252 pub messages: Vec<SerializedMessage>,
253 #[serde(default)]
254 pub initial_project_snapshot: Option<Arc<ProjectSnapshot>>,
255}
256
257#[derive(Debug, Serialize, Deserialize)]
258pub struct SerializedMessage {
259 pub id: MessageId,
260 pub role: Role,
261 pub text: String,
262 #[serde(default)]
263 pub tool_uses: Vec<SerializedToolUse>,
264 #[serde(default)]
265 pub tool_results: Vec<SerializedToolResult>,
266}
267
268#[derive(Debug, Serialize, Deserialize)]
269pub struct SerializedToolUse {
270 pub id: LanguageModelToolUseId,
271 pub name: SharedString,
272 pub input: serde_json::Value,
273}
274
275#[derive(Debug, Serialize, Deserialize)]
276pub struct SerializedToolResult {
277 pub tool_use_id: LanguageModelToolUseId,
278 pub is_error: bool,
279 pub content: Arc<str>,
280}
281
282struct GlobalThreadsDatabase(
283 Shared<BoxFuture<'static, Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>>,
284);
285
286impl Global for GlobalThreadsDatabase {}
287
288pub(crate) struct ThreadsDatabase {
289 executor: BackgroundExecutor,
290 env: heed::Env,
291 threads: Database<SerdeBincode<ThreadId>, SerdeJson<SerializedThread>>,
292}
293
294impl ThreadsDatabase {
295 fn global_future(
296 cx: &mut App,
297 ) -> Shared<BoxFuture<'static, Result<Arc<ThreadsDatabase>, Arc<anyhow::Error>>>> {
298 GlobalThreadsDatabase::global(cx).0.clone()
299 }
300
301 fn init(cx: &mut App) {
302 let executor = cx.background_executor().clone();
303 let database_future = executor
304 .spawn({
305 let executor = executor.clone();
306 let database_path = paths::support_dir().join("threads/threads-db.1.mdb");
307 async move { ThreadsDatabase::new(database_path, executor) }
308 })
309 .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new)))
310 .boxed()
311 .shared();
312
313 cx.set_global(GlobalThreadsDatabase(database_future));
314 }
315
316 pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result<Self> {
317 std::fs::create_dir_all(&path)?;
318
319 const ONE_GB_IN_BYTES: usize = 1024 * 1024 * 1024;
320 let env = unsafe {
321 heed::EnvOpenOptions::new()
322 .map_size(ONE_GB_IN_BYTES)
323 .max_dbs(1)
324 .open(path)?
325 };
326
327 let mut txn = env.write_txn()?;
328 let threads = env.create_database(&mut txn, Some("threads"))?;
329 txn.commit()?;
330
331 Ok(Self {
332 executor,
333 env,
334 threads,
335 })
336 }
337
338 pub fn list_threads(&self) -> Task<Result<Vec<SerializedThreadMetadata>>> {
339 let env = self.env.clone();
340 let threads = self.threads;
341
342 self.executor.spawn(async move {
343 let txn = env.read_txn()?;
344 let mut iter = threads.iter(&txn)?;
345 let mut threads = Vec::new();
346 while let Some((key, value)) = iter.next().transpose()? {
347 threads.push(SerializedThreadMetadata {
348 id: key,
349 summary: value.summary,
350 updated_at: value.updated_at,
351 });
352 }
353
354 Ok(threads)
355 })
356 }
357
358 pub fn try_find_thread(&self, id: ThreadId) -> Task<Result<Option<SerializedThread>>> {
359 let env = self.env.clone();
360 let threads = self.threads;
361
362 self.executor.spawn(async move {
363 let txn = env.read_txn()?;
364 let thread = threads.get(&txn, &id)?;
365 Ok(thread)
366 })
367 }
368
369 pub fn save_thread(&self, id: ThreadId, thread: SerializedThread) -> Task<Result<()>> {
370 let env = self.env.clone();
371 let threads = self.threads;
372
373 self.executor.spawn(async move {
374 let mut txn = env.write_txn()?;
375 threads.put(&mut txn, &id, &thread)?;
376 txn.commit()?;
377 Ok(())
378 })
379 }
380
381 pub fn delete_thread(&self, id: ThreadId) -> Task<Result<()>> {
382 let env = self.env.clone();
383 let threads = self.threads;
384
385 self.executor.spawn(async move {
386 let mut txn = env.write_txn()?;
387 threads.delete(&mut txn, &id)?;
388 txn.commit()?;
389 Ok(())
390 })
391 }
392}