server.rs

  1use crate::{AcpThread, AgentThreadEntryContent, ThreadEntryId, ThreadId, ToolCallId};
  2use agentic_coding_protocol as acp;
  3use anyhow::{Context as _, Result};
  4use async_trait::async_trait;
  5use collections::HashMap;
  6use futures::channel::oneshot;
  7use gpui::{App, AppContext, AsyncApp, Context, Entity, Task, WeakEntity};
  8use parking_lot::Mutex;
  9use project::Project;
 10use smol::process::Child;
 11use std::{io::Write as _, path::Path, sync::Arc};
 12use util::ResultExt;
 13
 14pub struct AcpServer {
 15    connection: Arc<acp::AgentConnection>,
 16    threads: Arc<Mutex<HashMap<ThreadId, WeakEntity<AcpThread>>>>,
 17    project: Entity<Project>,
 18    _handler_task: Task<()>,
 19    _io_task: Task<()>,
 20}
 21
 22struct AcpClientDelegate {
 23    project: Entity<Project>,
 24    threads: Arc<Mutex<HashMap<ThreadId, WeakEntity<AcpThread>>>>,
 25    cx: AsyncApp,
 26    // sent_buffer_versions: HashMap<Entity<Buffer>, HashMap<u64, BufferSnapshot>>,
 27}
 28
 29impl AcpClientDelegate {
 30    fn new(
 31        project: Entity<Project>,
 32        threads: Arc<Mutex<HashMap<ThreadId, WeakEntity<AcpThread>>>>,
 33        cx: AsyncApp,
 34    ) -> Self {
 35        Self {
 36            project,
 37            threads,
 38            cx: cx,
 39        }
 40    }
 41
 42    fn update_thread<R>(
 43        &self,
 44        thread_id: &ThreadId,
 45        cx: &mut App,
 46        callback: impl FnOnce(&mut AcpThread, &mut Context<AcpThread>) -> R,
 47    ) -> Option<R> {
 48        let thread = self.threads.lock().get(&thread_id)?.clone();
 49        let Some(thread) = thread.upgrade() else {
 50            self.threads.lock().remove(&thread_id);
 51            return None;
 52        };
 53        Some(thread.update(cx, callback))
 54    }
 55}
 56
 57#[async_trait(?Send)]
 58impl acp::Client for AcpClientDelegate {
 59    async fn stat(&self, params: acp::StatParams) -> Result<acp::StatResponse> {
 60        let cx = &mut self.cx.clone();
 61        self.project.update(cx, |project, cx| {
 62            let path = project
 63                .project_path_for_absolute_path(Path::new(&params.path), cx)
 64                .context("Failed to get project path")?;
 65
 66            match project.entry_for_path(&path, cx) {
 67                // todo! refresh entry?
 68                None => Ok(acp::StatResponse {
 69                    exists: false,
 70                    is_directory: false,
 71                }),
 72                Some(entry) => Ok(acp::StatResponse {
 73                    exists: entry.is_created(),
 74                    is_directory: entry.is_dir(),
 75                }),
 76            }
 77        })?
 78    }
 79
 80    async fn stream_message_chunk(
 81        &self,
 82        params: acp::StreamMessageChunkParams,
 83    ) -> Result<acp::StreamMessageChunkResponse> {
 84        let cx = &mut self.cx.clone();
 85
 86        cx.update(|cx| {
 87            self.update_thread(&params.thread_id.into(), cx, |thread, cx| {
 88                thread.push_assistant_chunk(params.chunk, cx)
 89            });
 90        })?;
 91
 92        Ok(acp::StreamMessageChunkResponse)
 93    }
 94
 95    async fn read_text_file(
 96        &self,
 97        request: acp::ReadTextFileParams,
 98    ) -> Result<acp::ReadTextFileResponse> {
 99        let cx = &mut self.cx.clone();
100        let buffer = self
101            .project
102            .update(cx, |project, cx| {
103                let path = project
104                    .project_path_for_absolute_path(Path::new(&request.path), cx)
105                    .context("Failed to get project path")?;
106                anyhow::Ok(project.open_buffer(path, cx))
107            })??
108            .await?;
109
110        buffer.update(cx, |buffer, cx| {
111            let start = language::Point::new(request.line_offset.unwrap_or(0), 0);
112            let end = match request.line_limit {
113                None => buffer.max_point(),
114                Some(limit) => start + language::Point::new(limit + 1, 0),
115            };
116
117            let content: String = buffer.text_for_range(start..end).collect();
118            self.update_thread(&request.thread_id.into(), cx, |thread, cx| {
119                thread.push_entry(
120                    AgentThreadEntryContent::ReadFile {
121                        path: request.path.clone(),
122                        content: content.clone(),
123                    },
124                    cx,
125                );
126            });
127
128            acp::ReadTextFileResponse {
129                content,
130                version: acp::FileVersion(0),
131            }
132        })
133    }
134
135    async fn read_binary_file(
136        &self,
137        request: acp::ReadBinaryFileParams,
138    ) -> Result<acp::ReadBinaryFileResponse> {
139        let cx = &mut self.cx.clone();
140        let file = self
141            .project
142            .update(cx, |project, cx| {
143                let (worktree, path) = project
144                    .find_worktree(Path::new(&request.path), cx)
145                    .context("Failed to get project path")?;
146
147                let task = worktree.update(cx, |worktree, cx| worktree.load_binary_file(&path, cx));
148                anyhow::Ok(task)
149            })??
150            .await?;
151
152        // todo! test
153        let content = cx
154            .background_spawn(async move {
155                let start = request.byte_offset.unwrap_or(0) as usize;
156                let end = request
157                    .byte_limit
158                    .map(|limit| (start + limit as usize).min(file.content.len()))
159                    .unwrap_or(file.content.len());
160
161                let range_content = &file.content[start..end];
162
163                let mut base64_content = Vec::new();
164                let mut base64_encoder = base64::write::EncoderWriter::new(
165                    std::io::Cursor::new(&mut base64_content),
166                    &base64::engine::general_purpose::STANDARD,
167                );
168                base64_encoder.write_all(range_content)?;
169                drop(base64_encoder);
170
171                // SAFETY: The base64 encoder should not produce non-UTF8.
172                unsafe { anyhow::Ok(String::from_utf8_unchecked(base64_content)) }
173            })
174            .await?;
175
176        Ok(acp::ReadBinaryFileResponse {
177            content,
178            // todo!
179            version: acp::FileVersion(0),
180        })
181    }
182
183    async fn glob_search(
184        &self,
185        _request: acp::GlobSearchParams,
186    ) -> Result<acp::GlobSearchResponse> {
187        todo!()
188    }
189
190    async fn request_tool_call(
191        &self,
192        request: acp::RequestToolCallParams,
193    ) -> Result<acp::RequestToolCallResponse> {
194        let (tx, rx) = oneshot::channel();
195
196        let cx = &mut self.cx.clone();
197        let entry_id = cx
198            .update(|cx| {
199                self.update_thread(&request.thread_id.into(), cx, |thread, cx| {
200                    // todo! tools that don't require confirmation
201                    thread.push_tool_call(request.tool_name, request.description, tx, cx)
202                })
203            })?
204            .context("Failed to update thread")?;
205
206        if dbg!(rx.await)? {
207            Ok(acp::RequestToolCallResponse::Allowed {
208                id: entry_id.into(),
209            })
210        } else {
211            Ok(acp::RequestToolCallResponse::Rejected)
212        }
213    }
214}
215
216impl AcpServer {
217    pub fn stdio(mut process: Child, project: Entity<Project>, cx: &mut AsyncApp) -> Arc<Self> {
218        let stdin = process.stdin.take().expect("process didn't have stdin");
219        let stdout = process.stdout.take().expect("process didn't have stdout");
220
221        let threads: Arc<Mutex<HashMap<ThreadId, WeakEntity<AcpThread>>>> = Default::default();
222        let (connection, handler_fut, io_fut) = acp::AgentConnection::connect_to_agent(
223            AcpClientDelegate::new(project.clone(), threads.clone(), cx.clone()),
224            stdin,
225            stdout,
226        );
227
228        let io_task = cx.background_spawn(async move {
229            io_fut.await.log_err();
230            process.status().await.log_err();
231        });
232
233        Arc::new(Self {
234            project,
235            connection: Arc::new(connection),
236            threads,
237            _handler_task: cx.foreground_executor().spawn(handler_fut),
238            _io_task: io_task,
239        })
240    }
241}
242
243impl AcpServer {
244    pub async fn create_thread(self: Arc<Self>, cx: &mut AsyncApp) -> Result<Entity<AcpThread>> {
245        let response = self.connection.request(acp::CreateThreadParams).await?;
246        let thread_id: ThreadId = response.thread_id.into();
247        let server = self.clone();
248        let thread = cx.new(|_| AcpThread {
249            title: "The agent2 thread".into(),
250            id: thread_id.clone(),
251            next_entry_id: ThreadEntryId(0),
252            entries: Vec::default(),
253            project: self.project.clone(),
254            server,
255        })?;
256        self.threads.lock().insert(thread_id, thread.downgrade());
257        Ok(thread)
258    }
259
260    pub async fn send_message(
261        &self,
262        thread_id: ThreadId,
263        message: acp::Message,
264        _cx: &mut AsyncApp,
265    ) -> Result<()> {
266        self.connection
267            .request(acp::SendMessageParams {
268                thread_id: thread_id.clone().into(),
269                message,
270            })
271            .await?;
272        Ok(())
273    }
274}
275
276impl From<acp::ThreadId> for ThreadId {
277    fn from(thread_id: acp::ThreadId) -> Self {
278        Self(thread_id.0.into())
279    }
280}
281
282impl From<ThreadId> for acp::ThreadId {
283    fn from(thread_id: ThreadId) -> Self {
284        acp::ThreadId(thread_id.0.to_string())
285    }
286}
287
288impl From<acp::ToolCallId> for ToolCallId {
289    fn from(tool_call_id: acp::ToolCallId) -> Self {
290        Self(ThreadEntryId(tool_call_id.0.into()))
291    }
292}
293
294impl From<ToolCallId> for acp::ToolCallId {
295    fn from(tool_call_id: ToolCallId) -> Self {
296        acp::ToolCallId(tool_call_id.0.0)
297    }
298}