connection.rs

  1use crate::AcpThread;
  2use agent_client_protocol::{self as acp};
  3use anyhow::Result;
  4use collections::IndexMap;
  5use gpui::{Entity, SharedString, Task};
  6use language_model::LanguageModelProviderId;
  7use project::Project;
  8use serde::{Deserialize, Serialize};
  9use std::{any::Any, error::Error, fmt, path::Path, rc::Rc, sync::Arc};
 10use ui::{App, IconName};
 11use uuid::Uuid;
 12
 13#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
 14pub struct UserMessageId(Arc<str>);
 15
 16impl UserMessageId {
 17    pub fn new() -> Self {
 18        Self(Uuid::new_v4().to_string().into())
 19    }
 20}
 21
 22pub trait AgentConnection {
 23    fn new_thread(
 24        self: Rc<Self>,
 25        project: Entity<Project>,
 26        cwd: &Path,
 27        cx: &mut App,
 28    ) -> Task<Result<Entity<AcpThread>>>;
 29
 30    fn auth_methods(&self) -> &[acp::AuthMethod];
 31
 32    fn authenticate(&self, method: acp::AuthMethodId, cx: &mut App) -> Task<Result<()>>;
 33
 34    fn prompt(
 35        &self,
 36        user_message_id: Option<UserMessageId>,
 37        params: acp::PromptRequest,
 38        cx: &mut App,
 39    ) -> Task<Result<acp::PromptResponse>>;
 40
 41    fn resume(
 42        &self,
 43        _session_id: &acp::SessionId,
 44        _cx: &App,
 45    ) -> Option<Rc<dyn AgentSessionResume>> {
 46        None
 47    }
 48
 49    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App);
 50
 51    fn truncate(
 52        &self,
 53        _session_id: &acp::SessionId,
 54        _cx: &App,
 55    ) -> Option<Rc<dyn AgentSessionTruncate>> {
 56        None
 57    }
 58
 59    fn set_title(
 60        &self,
 61        _session_id: &acp::SessionId,
 62        _cx: &App,
 63    ) -> Option<Rc<dyn AgentSessionSetTitle>> {
 64        None
 65    }
 66
 67    /// Returns this agent as an [Rc<dyn ModelSelector>] if the model selection capability is supported.
 68    ///
 69    /// If the agent does not support model selection, returns [None].
 70    /// This allows sharing the selector in UI components.
 71    fn model_selector(&self) -> Option<Rc<dyn AgentModelSelector>> {
 72        None
 73    }
 74
 75    fn telemetry(&self) -> Option<Rc<dyn AgentTelemetry>> {
 76        None
 77    }
 78
 79    fn list_commands(&self, session_id: &acp::SessionId, cx: &mut App) -> Task<Result<acp::ListCommandsResponse>>;
 80    fn run_command(&self, request: acp::RunCommandRequest, cx: &mut App) -> Task<Result<()>>;
 81
 82    fn into_any(self: Rc<Self>) -> Rc<dyn Any>;
 83}
 84
 85impl dyn AgentConnection {
 86    pub fn downcast<T: 'static + AgentConnection + Sized>(self: Rc<Self>) -> Option<Rc<T>> {
 87        self.into_any().downcast().ok()
 88    }
 89}
 90
 91pub trait AgentSessionTruncate {
 92    fn run(&self, message_id: UserMessageId, cx: &mut App) -> Task<Result<()>>;
 93}
 94
 95pub trait AgentSessionResume {
 96    fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>>;
 97}
 98
 99pub trait AgentSessionSetTitle {
100    fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>>;
101}
102
103pub trait AgentTelemetry {
104    /// The name of the agent used for telemetry.
105    fn agent_name(&self) -> String;
106
107    /// A representation of the current thread state that can be serialized for
108    /// storage with telemetry events.
109    fn thread_data(
110        &self,
111        session_id: &acp::SessionId,
112        cx: &mut App,
113    ) -> Task<Result<serde_json::Value>>;
114}
115
116#[derive(Debug)]
117pub struct AuthRequired {
118    pub description: Option<String>,
119    pub provider_id: Option<LanguageModelProviderId>,
120}
121
122impl AuthRequired {
123    pub fn new() -> Self {
124        Self {
125            description: None,
126            provider_id: None,
127        }
128    }
129
130    pub fn with_description(mut self, description: String) -> Self {
131        self.description = Some(description);
132        self
133    }
134
135    pub fn with_language_model_provider(mut self, provider_id: LanguageModelProviderId) -> Self {
136        self.provider_id = Some(provider_id);
137        self
138    }
139}
140
141impl Error for AuthRequired {}
142impl fmt::Display for AuthRequired {
143    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
144        write!(f, "Authentication required")
145    }
146}
147
148/// Trait for agents that support listing, selecting, and querying language models.
149///
150/// This is an optional capability; agents indicate support via [AgentConnection::model_selector].
151pub trait AgentModelSelector: 'static {
152    /// Lists all available language models for this agent.
153    ///
154    /// # Parameters
155    /// - `cx`: The GPUI app context for async operations and global access.
156    ///
157    /// # Returns
158    /// A task resolving to the list of models or an error (e.g., if no models are configured).
159    fn list_models(&self, cx: &mut App) -> Task<Result<AgentModelList>>;
160
161    /// Selects a model for a specific session (thread).
162    ///
163    /// This sets the default model for future interactions in the session.
164    /// If the session doesn't exist or the model is invalid, it returns an error.
165    ///
166    /// # Parameters
167    /// - `session_id`: The ID of the session (thread) to apply the model to.
168    /// - `model`: The model to select (should be one from [list_models]).
169    /// - `cx`: The GPUI app context.
170    ///
171    /// # Returns
172    /// A task resolving to `Ok(())` on success or an error.
173    fn select_model(
174        &self,
175        session_id: acp::SessionId,
176        model_id: AgentModelId,
177        cx: &mut App,
178    ) -> Task<Result<()>>;
179
180    /// Retrieves the currently selected model for a specific session (thread).
181    ///
182    /// # Parameters
183    /// - `session_id`: The ID of the session (thread) to query.
184    /// - `cx`: The GPUI app context.
185    ///
186    /// # Returns
187    /// A task resolving to the selected model (always set) or an error (e.g., session not found).
188    fn selected_model(
189        &self,
190        session_id: &acp::SessionId,
191        cx: &mut App,
192    ) -> Task<Result<AgentModelInfo>>;
193
194    /// Whenever the model list is updated the receiver will be notified.
195    fn watch(&self, cx: &mut App) -> watch::Receiver<()>;
196}
197
198#[derive(Debug, Clone, PartialEq, Eq, Hash)]
199pub struct AgentModelId(pub SharedString);
200
201impl std::ops::Deref for AgentModelId {
202    type Target = SharedString;
203
204    fn deref(&self) -> &Self::Target {
205        &self.0
206    }
207}
208
209impl fmt::Display for AgentModelId {
210    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
211        self.0.fmt(f)
212    }
213}
214
215#[derive(Debug, Clone, PartialEq, Eq)]
216pub struct AgentModelInfo {
217    pub id: AgentModelId,
218    pub name: SharedString,
219    pub icon: Option<IconName>,
220}
221
222#[derive(Debug, Clone, PartialEq, Eq, Hash)]
223pub struct AgentModelGroupName(pub SharedString);
224
225#[derive(Debug, Clone)]
226pub enum AgentModelList {
227    Flat(Vec<AgentModelInfo>),
228    Grouped(IndexMap<AgentModelGroupName, Vec<AgentModelInfo>>),
229}
230
231impl AgentModelList {
232    pub fn is_empty(&self) -> bool {
233        match self {
234            AgentModelList::Flat(models) => models.is_empty(),
235            AgentModelList::Grouped(groups) => groups.is_empty(),
236        }
237    }
238}
239
240#[cfg(feature = "test-support")]
241mod test_support {
242    use std::sync::Arc;
243
244    use action_log::ActionLog;
245    use collections::HashMap;
246    use futures::{channel::oneshot, future::try_join_all};
247    use gpui::{AppContext as _, WeakEntity};
248    use parking_lot::Mutex;
249
250    use super::*;
251
252    #[derive(Clone, Default)]
253    pub struct StubAgentConnection {
254        sessions: Arc<Mutex<HashMap<acp::SessionId, Session>>>,
255        permission_requests: HashMap<acp::ToolCallId, Vec<acp::PermissionOption>>,
256        next_prompt_updates: Arc<Mutex<Vec<acp::SessionUpdate>>>,
257    }
258
259    struct Session {
260        thread: WeakEntity<AcpThread>,
261        response_tx: Option<oneshot::Sender<acp::StopReason>>,
262    }
263
264    impl StubAgentConnection {
265        pub fn new() -> Self {
266            Self {
267                next_prompt_updates: Default::default(),
268                permission_requests: HashMap::default(),
269                sessions: Arc::default(),
270            }
271        }
272
273        pub fn set_next_prompt_updates(&self, updates: Vec<acp::SessionUpdate>) {
274            *self.next_prompt_updates.lock() = updates;
275        }
276
277        pub fn with_permission_requests(
278            mut self,
279            permission_requests: HashMap<acp::ToolCallId, Vec<acp::PermissionOption>>,
280        ) -> Self {
281            self.permission_requests = permission_requests;
282            self
283        }
284
285        pub fn send_update(
286            &self,
287            session_id: acp::SessionId,
288            update: acp::SessionUpdate,
289            cx: &mut App,
290        ) {
291            assert!(
292                self.next_prompt_updates.lock().is_empty(),
293                "Use either send_update or set_next_prompt_updates"
294            );
295
296            self.sessions
297                .lock()
298                .get(&session_id)
299                .unwrap()
300                .thread
301                .update(cx, |thread, cx| {
302                    thread.handle_session_update(update, cx).unwrap();
303                })
304                .unwrap();
305        }
306
307        pub fn end_turn(&self, session_id: acp::SessionId, stop_reason: acp::StopReason) {
308            self.sessions
309                .lock()
310                .get_mut(&session_id)
311                .unwrap()
312                .response_tx
313                .take()
314                .expect("No pending turn")
315                .send(stop_reason)
316                .unwrap();
317        }
318    }
319
320    impl AgentConnection for StubAgentConnection {
321        fn auth_methods(&self) -> &[acp::AuthMethod] {
322            &[]
323        }
324
325        fn new_thread(
326            self: Rc<Self>,
327            project: Entity<Project>,
328            _cwd: &Path,
329            cx: &mut gpui::App,
330        ) -> Task<gpui::Result<Entity<AcpThread>>> {
331            let session_id = acp::SessionId(self.sessions.lock().len().to_string().into());
332            let action_log = cx.new(|_| ActionLog::new(project.clone()));
333            let thread = cx.new(|cx| {
334                AcpThread::new(
335                    "Test",
336                    self.clone(),
337                    project,
338                    action_log,
339                    session_id.clone(),
340                    watch::Receiver::constant(acp::PromptCapabilities {
341                        image: true,
342                        audio: true,
343                        embedded_context: true,
344                        supports_custom_commands: false,
345                    }),
346                    cx,
347                )
348            });
349            self.sessions.lock().insert(
350                session_id,
351                Session {
352                    thread: thread.downgrade(),
353                    response_tx: None,
354                },
355            );
356            Task::ready(Ok(thread))
357        }
358
359        fn authenticate(
360            &self,
361            _method_id: acp::AuthMethodId,
362            _cx: &mut App,
363        ) -> Task<gpui::Result<()>> {
364            unimplemented!()
365        }
366
367        fn prompt(
368            &self,
369            _id: Option<UserMessageId>,
370            params: acp::PromptRequest,
371            cx: &mut App,
372        ) -> Task<gpui::Result<acp::PromptResponse>> {
373            let mut sessions = self.sessions.lock();
374            let Session {
375                thread,
376                response_tx,
377            } = sessions.get_mut(&params.session_id).unwrap();
378            let mut tasks = vec![];
379            if self.next_prompt_updates.lock().is_empty() {
380                let (tx, rx) = oneshot::channel();
381                response_tx.replace(tx);
382                cx.spawn(async move |_| {
383                    let stop_reason = rx.await?;
384                    Ok(acp::PromptResponse { stop_reason })
385                })
386            } else {
387                for update in self.next_prompt_updates.lock().drain(..) {
388                    let thread = thread.clone();
389                    let update = update.clone();
390                    let permission_request = if let acp::SessionUpdate::ToolCall(tool_call) =
391                        &update
392                        && let Some(options) = self.permission_requests.get(&tool_call.id)
393                    {
394                        Some((tool_call.clone(), options.clone()))
395                    } else {
396                        None
397                    };
398                    let task = cx.spawn(async move |cx| {
399                        if let Some((tool_call, options)) = permission_request {
400                            let permission = thread.update(cx, |thread, cx| {
401                                thread.request_tool_call_authorization(
402                                    tool_call.clone().into(),
403                                    options.clone(),
404                                    cx,
405                                )
406                            })?;
407                            permission?.await?;
408                        }
409                        thread.update(cx, |thread, cx| {
410                            thread.handle_session_update(update.clone(), cx).unwrap();
411                        })?;
412                        anyhow::Ok(())
413                    });
414                    tasks.push(task);
415                }
416
417                cx.spawn(async move |_| {
418                    try_join_all(tasks).await?;
419                    Ok(acp::PromptResponse {
420                        stop_reason: acp::StopReason::EndTurn,
421                    })
422                })
423            }
424        }
425
426        fn cancel(&self, session_id: &acp::SessionId, _cx: &mut App) {
427            if let Some(end_turn_tx) = self
428                .sessions
429                .lock()
430                .get_mut(session_id)
431                .unwrap()
432                .response_tx
433                .take()
434            {
435                end_turn_tx.send(acp::StopReason::Cancelled).unwrap();
436            }
437        }
438
439        fn truncate(
440            &self,
441            _session_id: &agent_client_protocol::SessionId,
442            _cx: &App,
443        ) -> Option<Rc<dyn AgentSessionTruncate>> {
444            Some(Rc::new(StubAgentSessionEditor))
445        }
446
447        fn list_commands(&self, _session_id: &acp::SessionId, _cx: &mut App) -> Task<Result<acp::ListCommandsResponse>> {
448            Task::ready(Ok(acp::ListCommandsResponse { commands: vec![] }))
449        }
450
451        fn run_command(&self, _request: acp::RunCommandRequest, _cx: &mut App) -> Task<Result<()>> {
452            Task::ready(Ok(()))
453        }
454
455        fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
456            self
457        }
458    }
459
460    struct StubAgentSessionEditor;
461
462    impl AgentSessionTruncate for StubAgentSessionEditor {
463        fn run(&self, _: UserMessageId, _: &mut App) -> Task<Result<()>> {
464            Task::ready(Ok(()))
465        }
466    }
467}
468
469#[cfg(feature = "test-support")]
470pub use test_support::*;