1use std::sync::Arc;
2
3use anyhow::Result;
4use assistant_tool::{ToolId, ToolWorkingSet};
5use collections::HashMap;
6use context_server::manager::ContextServerManager;
7use context_server::{ContextServerFactoryRegistry, ContextServerTool};
8use gpui::{prelude::*, AppContext, Model, ModelContext, Task};
9use project::Project;
10use unindent::Unindent;
11use util::ResultExt as _;
12
13use crate::thread::{Thread, ThreadId};
14
15pub struct ThreadStore {
16 #[allow(unused)]
17 project: Model<Project>,
18 tools: Arc<ToolWorkingSet>,
19 context_server_manager: Model<ContextServerManager>,
20 context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,
21 threads: Vec<Model<Thread>>,
22}
23
24impl ThreadStore {
25 pub fn new(
26 project: Model<Project>,
27 tools: Arc<ToolWorkingSet>,
28 cx: &mut AppContext,
29 ) -> Task<Result<Model<Self>>> {
30 cx.spawn(|mut cx| async move {
31 let this = cx.new_model(|cx: &mut ModelContext<Self>| {
32 let context_server_factory_registry =
33 ContextServerFactoryRegistry::default_global(cx);
34 let context_server_manager = cx.new_model(|cx| {
35 ContextServerManager::new(context_server_factory_registry, project.clone(), cx)
36 });
37
38 let mut this = Self {
39 project,
40 tools,
41 context_server_manager,
42 context_server_tool_ids: HashMap::default(),
43 threads: Vec::new(),
44 };
45 this.mock_recent_threads(cx);
46 this.register_context_server_handlers(cx);
47
48 this
49 })?;
50
51 Ok(this)
52 })
53 }
54
55 pub fn recent_threads(&self, limit: usize, _cx: &ModelContext<Self>) -> Vec<Model<Thread>> {
56 self.threads.iter().take(limit).cloned().collect()
57 }
58
59 pub fn create_thread(&mut self, cx: &mut ModelContext<Self>) -> Model<Thread> {
60 let thread = cx.new_model(|cx| Thread::new(self.tools.clone(), cx));
61 self.threads.push(thread.clone());
62 thread
63 }
64
65 pub fn open_thread(&self, id: &ThreadId, cx: &mut ModelContext<Self>) -> Option<Model<Thread>> {
66 self.threads
67 .iter()
68 .find(|thread| thread.read(cx).id() == id)
69 .cloned()
70 }
71
72 fn register_context_server_handlers(&self, cx: &mut ModelContext<Self>) {
73 cx.subscribe(
74 &self.context_server_manager.clone(),
75 Self::handle_context_server_event,
76 )
77 .detach();
78 }
79
80 fn handle_context_server_event(
81 &mut self,
82 context_server_manager: Model<ContextServerManager>,
83 event: &context_server::manager::Event,
84 cx: &mut ModelContext<Self>,
85 ) {
86 let tool_working_set = self.tools.clone();
87 match event {
88 context_server::manager::Event::ServerStarted { server_id } => {
89 if let Some(server) = context_server_manager.read(cx).get_server(server_id) {
90 let context_server_manager = context_server_manager.clone();
91 cx.spawn({
92 let server = server.clone();
93 let server_id = server_id.clone();
94 |this, mut cx| async move {
95 let Some(protocol) = server.client() else {
96 return;
97 };
98
99 if protocol.capable(context_server::protocol::ServerCapability::Tools) {
100 if let Some(tools) = protocol.list_tools().await.log_err() {
101 let tool_ids = tools
102 .tools
103 .into_iter()
104 .map(|tool| {
105 log::info!(
106 "registering context server tool: {:?}",
107 tool.name
108 );
109 tool_working_set.insert(Arc::new(
110 ContextServerTool::new(
111 context_server_manager.clone(),
112 server.id(),
113 tool,
114 ),
115 ))
116 })
117 .collect::<Vec<_>>();
118
119 this.update(&mut cx, |this, _cx| {
120 this.context_server_tool_ids.insert(server_id, tool_ids);
121 })
122 .log_err();
123 }
124 }
125 }
126 })
127 .detach();
128 }
129 }
130 context_server::manager::Event::ServerStopped { server_id } => {
131 if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
132 tool_working_set.remove(&tool_ids);
133 }
134 }
135 }
136 }
137}
138
139impl ThreadStore {
140 /// Creates some mocked recent threads for testing purposes.
141 fn mock_recent_threads(&mut self, cx: &mut ModelContext<Self>) {
142 use language_model::Role;
143
144 self.threads.push(cx.new_model(|cx| {
145 let mut thread = Thread::new(self.tools.clone(), cx);
146 thread.insert_user_message("Hello! Can you help me understand quantum computing?", cx);
147 thread.insert_message(Role::Assistant, "Of course! I'd be happy to help you understand quantum computing. Quantum computing is a fascinating field that uses the principles of quantum mechanics to process information. Unlike classical computers that use bits (0s and 1s), quantum computers use quantum bits or 'qubits'. These qubits can exist in multiple states simultaneously, a property called superposition. This allows quantum computers to perform certain calculations much faster than classical computers. What specific aspect of quantum computing would you like to know more about?", cx);
148 thread.insert_user_message("That's interesting! Can you explain how quantum entanglement is used in quantum computing?", cx);
149 thread.insert_message(Role::Assistant, "Certainly! Quantum entanglement is a key principle used in quantum computing. When two qubits become entangled, the state of one qubit is directly related to the state of the other, regardless of the distance between them. This property is used in quantum computing to create complex quantum states and to perform operations on multiple qubits simultaneously. Entanglement allows quantum computers to process information in ways that classical computers cannot, potentially solving certain problems much more efficiently. For example, it's crucial in quantum error correction and in algorithms like quantum teleportation, which is important for quantum communication.", cx);
150 thread
151 }));
152
153 self.threads.push(cx.new_model(|cx| {
154 let mut thread = Thread::new(self.tools.clone(), cx);
155 thread.insert_user_message("Can you show me an example of Rust code for a simple web server?", cx);
156 thread.insert_message(Role::Assistant, "Certainly! Here's an example of a simple web server in Rust using the `actix-web` framework:
157
158 ```rust
159 use actix_web::{web, App, HttpResponse, HttpServer, Responder};
160
161 async fn hello() -> impl Responder {
162 HttpResponse::Ok().body(\"Hello, World!\")
163 }
164
165 #[actix_web::main]
166 async fn main() -> std::io::Result<()> {
167 HttpServer::new(|| {
168 App::new()
169 .route(\"/\", web::get().to(hello))
170 })
171 .bind(\"127.0.0.1:8080\")?
172 .run()
173 .await
174 }
175 ```
176
177 This code creates a basic web server that responds with 'Hello, World!' when you access the root URL. Here's a breakdown of what's happening:
178
179 1. We import necessary items from the `actix-web` crate.
180 2. We define an async `hello` function that returns a simple HTTP response.
181 3. In the `main` function, we set up the server to listen on `127.0.0.1:8080`.
182 4. We configure the app to respond to GET requests on the root path with our `hello` function.
183
184 To run this, you'd need to add `actix-web` to your `Cargo.toml` dependencies:
185
186 ```toml
187 [dependencies]
188 actix-web = \"4.0\"
189 ```
190
191 Then you can run the server with `cargo run` and access it at `http://localhost:8080`.".unindent(), cx);
192 thread.insert_user_message("That's great! Can you explain more about async functions in Rust?", cx);
193 thread.insert_message(Role::Assistant, "Certainly! Async functions are a key feature in Rust for writing efficient, non-blocking code, especially for I/O-bound operations. Here's an overview:
194
195 1. **Syntax**: Async functions are declared using the `async` keyword:
196
197 ```rust
198 async fn my_async_function() -> Result<(), Error> {
199 // Asynchronous code here
200 }
201 ```
202
203 2. **Futures**: Async functions return a `Future`. A `Future` represents a value that may not be available yet but will be at some point.
204
205 3. **Await**: Inside an async function, you can use the `.await` syntax to wait for other async operations to complete:
206
207 ```rust
208 async fn fetch_data() -> Result<String, Error> {
209 let response = make_http_request().await?;
210 let data = process_response(response).await?;
211 Ok(data)
212 }
213 ```
214
215 4. **Non-blocking**: Async functions allow the runtime to work on other tasks while waiting for I/O or other operations to complete, making efficient use of system resources.
216
217 5. **Runtime**: To execute async code, you need a runtime like `tokio` or `async-std`. Actix-web, which we used in the previous example, includes its own runtime.
218
219 6. **Error Handling**: Async functions work well with Rust's `?` operator for error handling.
220
221 Async programming in Rust provides a powerful way to write concurrent code that's both safe and efficient. It's particularly useful for servers, network programming, and any application that deals with many concurrent operations.".unindent(), cx);
222 thread
223 }));
224 }
225}