server.rs

  1use crate::{AcpThread, 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
119            acp::ReadTextFileResponse {
120                content,
121                version: acp::FileVersion(0),
122            }
123        })
124    }
125
126    async fn read_binary_file(
127        &self,
128        request: acp::ReadBinaryFileParams,
129    ) -> Result<acp::ReadBinaryFileResponse> {
130        let cx = &mut self.cx.clone();
131        let file = self
132            .project
133            .update(cx, |project, cx| {
134                let (worktree, path) = project
135                    .find_worktree(Path::new(&request.path), cx)
136                    .context("Failed to get project path")?;
137
138                let task = worktree.update(cx, |worktree, cx| worktree.load_binary_file(&path, cx));
139                anyhow::Ok(task)
140            })??
141            .await?;
142
143        // todo! test
144        let content = cx
145            .background_spawn(async move {
146                let start = request.byte_offset.unwrap_or(0) as usize;
147                let end = request
148                    .byte_limit
149                    .map(|limit| (start + limit as usize).min(file.content.len()))
150                    .unwrap_or(file.content.len());
151
152                let range_content = &file.content[start..end];
153
154                let mut base64_content = Vec::new();
155                let mut base64_encoder = base64::write::EncoderWriter::new(
156                    std::io::Cursor::new(&mut base64_content),
157                    &base64::engine::general_purpose::STANDARD,
158                );
159                base64_encoder.write_all(range_content)?;
160                drop(base64_encoder);
161
162                // SAFETY: The base64 encoder should not produce non-UTF8.
163                unsafe { anyhow::Ok(String::from_utf8_unchecked(base64_content)) }
164            })
165            .await?;
166
167        Ok(acp::ReadBinaryFileResponse {
168            content,
169            // todo!
170            version: acp::FileVersion(0),
171        })
172    }
173
174    async fn glob_search(
175        &self,
176        _request: acp::GlobSearchParams,
177    ) -> Result<acp::GlobSearchResponse> {
178        todo!()
179    }
180
181    async fn request_tool_call(
182        &self,
183        request: acp::RequestToolCallParams,
184    ) -> Result<acp::RequestToolCallResponse> {
185        let (tx, rx) = oneshot::channel();
186
187        let cx = &mut self.cx.clone();
188        let entry_id = cx
189            .update(|cx| {
190                self.update_thread(&request.thread_id.into(), cx, |thread, cx| {
191                    // todo! tools that don't require confirmation
192                    thread.push_tool_call(request.tool_name, request.description, tx, cx)
193                })
194            })?
195            .context("Failed to update thread")?;
196
197        if rx.await? {
198            Ok(acp::RequestToolCallResponse::Allowed {
199                id: entry_id.into(),
200            })
201        } else {
202            Ok(acp::RequestToolCallResponse::Rejected)
203        }
204    }
205}
206
207impl AcpServer {
208    pub fn stdio(mut process: Child, project: Entity<Project>, cx: &mut AsyncApp) -> Arc<Self> {
209        let stdin = process.stdin.take().expect("process didn't have stdin");
210        let stdout = process.stdout.take().expect("process didn't have stdout");
211
212        let threads: Arc<Mutex<HashMap<ThreadId, WeakEntity<AcpThread>>>> = Default::default();
213        let (connection, handler_fut, io_fut) = acp::AgentConnection::connect_to_agent(
214            AcpClientDelegate::new(project.clone(), threads.clone(), cx.clone()),
215            stdin,
216            stdout,
217        );
218
219        let io_task = cx.background_spawn(async move {
220            io_fut.await.log_err();
221            process.status().await.log_err();
222        });
223
224        Arc::new(Self {
225            project,
226            connection: Arc::new(connection),
227            threads,
228            _handler_task: cx.foreground_executor().spawn(handler_fut),
229            _io_task: io_task,
230        })
231    }
232}
233
234impl AcpServer {
235    pub async fn create_thread(self: Arc<Self>, cx: &mut AsyncApp) -> Result<Entity<AcpThread>> {
236        let response = self.connection.request(acp::CreateThreadParams).await?;
237        let thread_id: ThreadId = response.thread_id.into();
238        let server = self.clone();
239        let thread = cx.new(|_| AcpThread {
240            title: "The agent2 thread".into(),
241            id: thread_id.clone(),
242            next_entry_id: ThreadEntryId(0),
243            entries: Vec::default(),
244            project: self.project.clone(),
245            server,
246        })?;
247        self.threads.lock().insert(thread_id, thread.downgrade());
248        Ok(thread)
249    }
250
251    pub async fn send_message(
252        &self,
253        thread_id: ThreadId,
254        message: acp::Message,
255        _cx: &mut AsyncApp,
256    ) -> Result<()> {
257        self.connection
258            .request(acp::SendMessageParams {
259                thread_id: thread_id.clone().into(),
260                message,
261            })
262            .await?;
263        Ok(())
264    }
265}
266
267impl From<acp::ThreadId> for ThreadId {
268    fn from(thread_id: acp::ThreadId) -> Self {
269        Self(thread_id.0.into())
270    }
271}
272
273impl From<ThreadId> for acp::ThreadId {
274    fn from(thread_id: ThreadId) -> Self {
275        acp::ThreadId(thread_id.0.to_string())
276    }
277}
278
279impl From<acp::ToolCallId> for ToolCallId {
280    fn from(tool_call_id: acp::ToolCallId) -> Self {
281        Self(ThreadEntryId(tool_call_id.0.into()))
282    }
283}
284
285impl From<ToolCallId> for acp::ToolCallId {
286    fn from(tool_call_id: ToolCallId) -> Self {
287        acp::ToolCallId(tool_call_id.0.0)
288    }
289}