connection.rs

  1use crate::AcpThread;
  2use agent_client_protocol::{self as acp};
  3use anyhow::Result;
  4use chrono::{DateTime, Utc};
  5use collections::IndexMap;
  6use gpui::{Entity, SharedString, Task};
  7use language_model::LanguageModelProviderId;
  8use project::Project;
  9use serde::{Deserialize, Serialize};
 10use std::{
 11    any::Any,
 12    error::Error,
 13    fmt,
 14    path::{Path, PathBuf},
 15    rc::Rc,
 16    sync::Arc,
 17};
 18use ui::{App, IconName};
 19use uuid::Uuid;
 20
 21#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
 22pub struct UserMessageId(Arc<str>);
 23
 24impl UserMessageId {
 25    pub fn new() -> Self {
 26        Self(Uuid::new_v4().to_string().into())
 27    }
 28}
 29
 30pub trait AgentConnection {
 31    fn telemetry_id(&self) -> SharedString;
 32
 33    fn new_session(
 34        self: Rc<Self>,
 35        project: Entity<Project>,
 36        cwd: &Path,
 37        cx: &mut App,
 38    ) -> Task<Result<Entity<AcpThread>>>;
 39
 40    /// Whether this agent supports loading existing sessions.
 41    fn supports_load_session(&self) -> bool {
 42        false
 43    }
 44
 45    /// Load an existing session by ID.
 46    fn load_session(
 47        self: Rc<Self>,
 48        _session_id: acp::SessionId,
 49        _project: Entity<Project>,
 50        _cwd: &Path,
 51        _title: Option<SharedString>,
 52        _cx: &mut App,
 53    ) -> Task<Result<Entity<AcpThread>>> {
 54        Task::ready(Err(anyhow::Error::msg("Loading sessions is not supported")))
 55    }
 56
 57    /// Whether this agent supports closing existing sessions.
 58    fn supports_close_session(&self) -> bool {
 59        false
 60    }
 61
 62    /// Close an existing session. Allows the agent to free the session from memory.
 63    fn close_session(&self, _session_id: &acp::SessionId, _cx: &mut App) -> Task<Result<()>> {
 64        Task::ready(Err(anyhow::Error::msg("Closing sessions is not supported")))
 65    }
 66
 67    /// Whether this agent supports resuming existing sessions without loading history.
 68    fn supports_resume_session(&self) -> bool {
 69        false
 70    }
 71
 72    /// Resume an existing session by ID without replaying previous messages.
 73    fn resume_session(
 74        self: Rc<Self>,
 75        _session_id: acp::SessionId,
 76        _project: Entity<Project>,
 77        _cwd: &Path,
 78        _title: Option<SharedString>,
 79        _cx: &mut App,
 80    ) -> Task<Result<Entity<AcpThread>>> {
 81        Task::ready(Err(anyhow::Error::msg(
 82            "Resuming sessions is not supported",
 83        )))
 84    }
 85
 86    /// Whether this agent supports showing session history.
 87    fn supports_session_history(&self) -> bool {
 88        self.supports_load_session() || self.supports_resume_session()
 89    }
 90
 91    fn auth_methods(&self) -> &[acp::AuthMethod];
 92
 93    fn authenticate(&self, method: acp::AuthMethodId, cx: &mut App) -> Task<Result<()>>;
 94
 95    fn prompt(
 96        &self,
 97        user_message_id: Option<UserMessageId>,
 98        params: acp::PromptRequest,
 99        cx: &mut App,
100    ) -> Task<Result<acp::PromptResponse>>;
101
102    fn retry(&self, _session_id: &acp::SessionId, _cx: &App) -> Option<Rc<dyn AgentSessionRetry>> {
103        None
104    }
105
106    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App);
107
108    fn truncate(
109        &self,
110        _session_id: &acp::SessionId,
111        _cx: &App,
112    ) -> Option<Rc<dyn AgentSessionTruncate>> {
113        None
114    }
115
116    fn set_title(
117        &self,
118        _session_id: &acp::SessionId,
119        _cx: &App,
120    ) -> Option<Rc<dyn AgentSessionSetTitle>> {
121        None
122    }
123
124    /// Returns this agent as an [Rc<dyn ModelSelector>] if the model selection capability is supported.
125    ///
126    /// If the agent does not support model selection, returns [None].
127    /// This allows sharing the selector in UI components.
128    fn model_selector(&self, _session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
129        None
130    }
131
132    fn telemetry(&self) -> Option<Rc<dyn AgentTelemetry>> {
133        None
134    }
135
136    fn session_modes(
137        &self,
138        _session_id: &acp::SessionId,
139        _cx: &App,
140    ) -> Option<Rc<dyn AgentSessionModes>> {
141        None
142    }
143
144    fn session_config_options(
145        &self,
146        _session_id: &acp::SessionId,
147        _cx: &App,
148    ) -> Option<Rc<dyn AgentSessionConfigOptions>> {
149        None
150    }
151
152    fn session_list(&self, _cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
153        None
154    }
155
156    fn into_any(self: Rc<Self>) -> Rc<dyn Any>;
157}
158
159impl dyn AgentConnection {
160    pub fn downcast<T: 'static + AgentConnection + Sized>(self: Rc<Self>) -> Option<Rc<T>> {
161        self.into_any().downcast().ok()
162    }
163}
164
165pub trait AgentSessionTruncate {
166    fn run(&self, message_id: UserMessageId, cx: &mut App) -> Task<Result<()>>;
167}
168
169pub trait AgentSessionRetry {
170    fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>>;
171}
172
173pub trait AgentSessionSetTitle {
174    fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>>;
175}
176
177pub trait AgentTelemetry {
178    /// A representation of the current thread state that can be serialized for
179    /// storage with telemetry events.
180    fn thread_data(
181        &self,
182        session_id: &acp::SessionId,
183        cx: &mut App,
184    ) -> Task<Result<serde_json::Value>>;
185}
186
187pub trait AgentSessionModes {
188    fn current_mode(&self) -> acp::SessionModeId;
189
190    fn all_modes(&self) -> Vec<acp::SessionMode>;
191
192    fn set_mode(&self, mode: acp::SessionModeId, cx: &mut App) -> Task<Result<()>>;
193}
194
195pub trait AgentSessionConfigOptions {
196    /// Get all current config options with their state
197    fn config_options(&self) -> Vec<acp::SessionConfigOption>;
198
199    /// Set a config option value
200    /// Returns the full updated list of config options
201    fn set_config_option(
202        &self,
203        config_id: acp::SessionConfigId,
204        value: acp::SessionConfigValueId,
205        cx: &mut App,
206    ) -> Task<Result<Vec<acp::SessionConfigOption>>>;
207
208    /// Whenever the config options are updated the receiver will be notified.
209    /// Optional for agents that don't update their config options dynamically.
210    fn watch(&self, _cx: &mut App) -> Option<watch::Receiver<()>> {
211        None
212    }
213}
214
215#[derive(Debug, Clone, Default)]
216pub struct AgentSessionListRequest {
217    pub cwd: Option<PathBuf>,
218    pub cursor: Option<String>,
219    pub meta: Option<acp::Meta>,
220}
221
222#[derive(Debug, Clone)]
223pub struct AgentSessionListResponse {
224    pub sessions: Vec<AgentSessionInfo>,
225    pub next_cursor: Option<String>,
226    pub meta: Option<acp::Meta>,
227}
228
229impl AgentSessionListResponse {
230    pub fn new(sessions: Vec<AgentSessionInfo>) -> Self {
231        Self {
232            sessions,
233            next_cursor: None,
234            meta: None,
235        }
236    }
237}
238
239#[derive(Debug, Clone, PartialEq)]
240pub struct AgentSessionInfo {
241    pub session_id: acp::SessionId,
242    pub cwd: Option<PathBuf>,
243    pub title: Option<SharedString>,
244    pub updated_at: Option<DateTime<Utc>>,
245    pub meta: Option<acp::Meta>,
246}
247
248impl AgentSessionInfo {
249    pub fn new(session_id: impl Into<acp::SessionId>) -> Self {
250        Self {
251            session_id: session_id.into(),
252            cwd: None,
253            title: None,
254            updated_at: None,
255            meta: None,
256        }
257    }
258}
259
260#[derive(Debug, Clone)]
261pub enum SessionListUpdate {
262    Refresh,
263    SessionInfo {
264        session_id: acp::SessionId,
265        update: acp::SessionInfoUpdate,
266    },
267}
268
269pub trait AgentSessionList {
270    fn list_sessions(
271        &self,
272        request: AgentSessionListRequest,
273        cx: &mut App,
274    ) -> Task<Result<AgentSessionListResponse>>;
275
276    fn supports_delete(&self) -> bool {
277        false
278    }
279
280    fn delete_session(&self, _session_id: &acp::SessionId, _cx: &mut App) -> Task<Result<()>> {
281        Task::ready(Err(anyhow::anyhow!("delete_session not supported")))
282    }
283
284    fn delete_sessions(&self, _cx: &mut App) -> Task<Result<()>> {
285        Task::ready(Err(anyhow::anyhow!("delete_sessions not supported")))
286    }
287
288    fn watch(&self, _cx: &mut App) -> Option<smol::channel::Receiver<SessionListUpdate>> {
289        None
290    }
291
292    fn notify_refresh(&self) {}
293
294    fn into_any(self: Rc<Self>) -> Rc<dyn Any>;
295}
296
297impl dyn AgentSessionList {
298    pub fn downcast<T: 'static + AgentSessionList + Sized>(self: Rc<Self>) -> Option<Rc<T>> {
299        self.into_any().downcast().ok()
300    }
301}
302
303#[derive(Debug)]
304pub struct AuthRequired {
305    pub description: Option<String>,
306    pub provider_id: Option<LanguageModelProviderId>,
307}
308
309impl AuthRequired {
310    pub fn new() -> Self {
311        Self {
312            description: None,
313            provider_id: None,
314        }
315    }
316
317    pub fn with_description(mut self, description: String) -> Self {
318        self.description = Some(description);
319        self
320    }
321
322    pub fn with_language_model_provider(mut self, provider_id: LanguageModelProviderId) -> Self {
323        self.provider_id = Some(provider_id);
324        self
325    }
326}
327
328impl Error for AuthRequired {}
329impl fmt::Display for AuthRequired {
330    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
331        write!(f, "Authentication required")
332    }
333}
334
335/// Trait for agents that support listing, selecting, and querying language models.
336///
337/// This is an optional capability; agents indicate support via [AgentConnection::model_selector].
338pub trait AgentModelSelector: 'static {
339    /// Lists all available language models for this agent.
340    ///
341    /// # Parameters
342    /// - `cx`: The GPUI app context for async operations and global access.
343    ///
344    /// # Returns
345    /// A task resolving to the list of models or an error (e.g., if no models are configured).
346    fn list_models(&self, cx: &mut App) -> Task<Result<AgentModelList>>;
347
348    /// Selects a model for a specific session (thread).
349    ///
350    /// This sets the default model for future interactions in the session.
351    /// If the session doesn't exist or the model is invalid, it returns an error.
352    ///
353    /// # Parameters
354    /// - `model`: The model to select (should be one from [list_models]).
355    /// - `cx`: The GPUI app context.
356    ///
357    /// # Returns
358    /// A task resolving to `Ok(())` on success or an error.
359    fn select_model(&self, model_id: acp::ModelId, cx: &mut App) -> Task<Result<()>>;
360
361    /// Retrieves the currently selected model for a specific session (thread).
362    ///
363    /// # Parameters
364    /// - `cx`: The GPUI app context.
365    ///
366    /// # Returns
367    /// A task resolving to the selected model (always set) or an error (e.g., session not found).
368    fn selected_model(&self, cx: &mut App) -> Task<Result<AgentModelInfo>>;
369
370    /// Whenever the model list is updated the receiver will be notified.
371    /// Optional for agents that don't update their model list.
372    fn watch(&self, _cx: &mut App) -> Option<watch::Receiver<()>> {
373        None
374    }
375
376    /// Returns whether the model picker should render a footer.
377    fn should_render_footer(&self) -> bool {
378        false
379    }
380}
381
382/// Icon for a model in the model selector.
383#[derive(Debug, Clone, PartialEq, Eq)]
384pub enum AgentModelIcon {
385    /// A built-in icon from Zed's icon set.
386    Named(IconName),
387    /// Path to a custom SVG icon file.
388    Path(SharedString),
389}
390
391#[derive(Debug, Clone, PartialEq, Eq)]
392pub struct AgentModelInfo {
393    pub id: acp::ModelId,
394    pub name: SharedString,
395    pub description: Option<SharedString>,
396    pub icon: Option<AgentModelIcon>,
397    pub is_latest: bool,
398    pub cost: Option<SharedString>,
399}
400
401impl From<acp::ModelInfo> for AgentModelInfo {
402    fn from(info: acp::ModelInfo) -> Self {
403        Self {
404            id: info.model_id,
405            name: info.name.into(),
406            description: info.description.map(|desc| desc.into()),
407            icon: None,
408            is_latest: false,
409            cost: None,
410        }
411    }
412}
413
414#[derive(Debug, Clone, PartialEq, Eq, Hash)]
415pub struct AgentModelGroupName(pub SharedString);
416
417#[derive(Debug, Clone)]
418pub enum AgentModelList {
419    Flat(Vec<AgentModelInfo>),
420    Grouped(IndexMap<AgentModelGroupName, Vec<AgentModelInfo>>),
421}
422
423impl AgentModelList {
424    pub fn is_empty(&self) -> bool {
425        match self {
426            AgentModelList::Flat(models) => models.is_empty(),
427            AgentModelList::Grouped(groups) => groups.is_empty(),
428        }
429    }
430
431    pub fn is_flat(&self) -> bool {
432        matches!(self, AgentModelList::Flat(_))
433    }
434}
435
436#[derive(Debug, Clone)]
437pub struct PermissionOptionChoice {
438    pub allow: acp::PermissionOption,
439    pub deny: acp::PermissionOption,
440}
441
442impl PermissionOptionChoice {
443    pub fn label(&self) -> SharedString {
444        self.allow.name.clone().into()
445    }
446}
447
448#[derive(Debug, Clone)]
449pub enum PermissionOptions {
450    Flat(Vec<acp::PermissionOption>),
451    Dropdown(Vec<PermissionOptionChoice>),
452}
453
454impl PermissionOptions {
455    pub fn is_empty(&self) -> bool {
456        match self {
457            PermissionOptions::Flat(options) => options.is_empty(),
458            PermissionOptions::Dropdown(options) => options.is_empty(),
459        }
460    }
461
462    pub fn first_option_of_kind(
463        &self,
464        kind: acp::PermissionOptionKind,
465    ) -> Option<&acp::PermissionOption> {
466        match self {
467            PermissionOptions::Flat(options) => options.iter().find(|option| option.kind == kind),
468            PermissionOptions::Dropdown(options) => options.iter().find_map(|choice| {
469                if choice.allow.kind == kind {
470                    Some(&choice.allow)
471                } else if choice.deny.kind == kind {
472                    Some(&choice.deny)
473                } else {
474                    None
475                }
476            }),
477        }
478    }
479
480    pub fn allow_once_option_id(&self) -> Option<acp::PermissionOptionId> {
481        self.first_option_of_kind(acp::PermissionOptionKind::AllowOnce)
482            .map(|option| option.option_id.clone())
483    }
484
485    pub fn deny_once_option_id(&self) -> Option<acp::PermissionOptionId> {
486        self.first_option_of_kind(acp::PermissionOptionKind::RejectOnce)
487            .map(|option| option.option_id.clone())
488    }
489}
490
491#[cfg(feature = "test-support")]
492mod test_support {
493    //! Test-only stubs and helpers for acp_thread.
494    //!
495    //! This module is gated by the `test-support` feature and is not included
496    //! in production builds. It provides:
497    //! - `StubAgentConnection` for mocking agent connections in tests
498    //! - `create_test_png_base64` for generating test images
499
500    use std::sync::Arc;
501    use std::sync::atomic::{AtomicUsize, Ordering};
502
503    use action_log::ActionLog;
504    use collections::HashMap;
505    use futures::{channel::oneshot, future::try_join_all};
506    use gpui::{AppContext as _, WeakEntity};
507    use parking_lot::Mutex;
508
509    use super::*;
510
511    /// Creates a PNG image encoded as base64 for testing.
512    ///
513    /// Generates a solid-color PNG of the specified dimensions and returns
514    /// it as a base64-encoded string suitable for use in `ImageContent`.
515    pub fn create_test_png_base64(width: u32, height: u32, color: [u8; 4]) -> String {
516        use image::ImageEncoder as _;
517
518        let mut png_data = Vec::new();
519        {
520            let encoder = image::codecs::png::PngEncoder::new(&mut png_data);
521            let mut pixels = Vec::with_capacity((width * height * 4) as usize);
522            for _ in 0..(width * height) {
523                pixels.extend_from_slice(&color);
524            }
525            encoder
526                .write_image(&pixels, width, height, image::ExtendedColorType::Rgba8)
527                .expect("Failed to encode PNG");
528        }
529
530        use image::EncodableLayout as _;
531        base64::Engine::encode(
532            &base64::engine::general_purpose::STANDARD,
533            png_data.as_bytes(),
534        )
535    }
536
537    #[derive(Clone, Default)]
538    pub struct StubAgentConnection {
539        sessions: Arc<Mutex<HashMap<acp::SessionId, Session>>>,
540        permission_requests: HashMap<acp::ToolCallId, PermissionOptions>,
541        next_prompt_updates: Arc<Mutex<Vec<acp::SessionUpdate>>>,
542    }
543
544    struct Session {
545        thread: WeakEntity<AcpThread>,
546        response_tx: Option<oneshot::Sender<acp::StopReason>>,
547    }
548
549    impl StubAgentConnection {
550        pub fn new() -> Self {
551            Self {
552                next_prompt_updates: Default::default(),
553                permission_requests: HashMap::default(),
554                sessions: Arc::default(),
555            }
556        }
557
558        pub fn set_next_prompt_updates(&self, updates: Vec<acp::SessionUpdate>) {
559            *self.next_prompt_updates.lock() = updates;
560        }
561
562        pub fn with_permission_requests(
563            mut self,
564            permission_requests: HashMap<acp::ToolCallId, PermissionOptions>,
565        ) -> Self {
566            self.permission_requests = permission_requests;
567            self
568        }
569
570        pub fn send_update(
571            &self,
572            session_id: acp::SessionId,
573            update: acp::SessionUpdate,
574            cx: &mut App,
575        ) {
576            assert!(
577                self.next_prompt_updates.lock().is_empty(),
578                "Use either send_update or set_next_prompt_updates"
579            );
580
581            self.sessions
582                .lock()
583                .get(&session_id)
584                .unwrap()
585                .thread
586                .update(cx, |thread, cx| {
587                    thread.handle_session_update(update, cx).unwrap();
588                })
589                .unwrap();
590        }
591
592        pub fn end_turn(&self, session_id: acp::SessionId, stop_reason: acp::StopReason) {
593            self.sessions
594                .lock()
595                .get_mut(&session_id)
596                .unwrap()
597                .response_tx
598                .take()
599                .expect("No pending turn")
600                .send(stop_reason)
601                .unwrap();
602        }
603    }
604
605    impl AgentConnection for StubAgentConnection {
606        fn telemetry_id(&self) -> SharedString {
607            "stub".into()
608        }
609
610        fn auth_methods(&self) -> &[acp::AuthMethod] {
611            &[]
612        }
613
614        fn model_selector(
615            &self,
616            _session_id: &acp::SessionId,
617        ) -> Option<Rc<dyn AgentModelSelector>> {
618            Some(self.model_selector_impl())
619        }
620
621        fn new_session(
622            self: Rc<Self>,
623            project: Entity<Project>,
624            cwd: &Path,
625            cx: &mut gpui::App,
626        ) -> Task<gpui::Result<Entity<AcpThread>>> {
627            static NEXT_SESSION_ID: AtomicUsize = AtomicUsize::new(0);
628            let session_id =
629                acp::SessionId::new(NEXT_SESSION_ID.fetch_add(1, Ordering::SeqCst).to_string());
630            let action_log = cx.new(|_| ActionLog::new(project.clone()));
631            let thread = cx.new(|cx| {
632                AcpThread::new(
633                    None,
634                    "Test",
635                    Some(cwd.to_path_buf()),
636                    self.clone(),
637                    project,
638                    action_log,
639                    session_id.clone(),
640                    watch::Receiver::constant(
641                        acp::PromptCapabilities::new()
642                            .image(true)
643                            .audio(true)
644                            .embedded_context(true),
645                    ),
646                    cx,
647                )
648            });
649            self.sessions.lock().insert(
650                session_id,
651                Session {
652                    thread: thread.downgrade(),
653                    response_tx: None,
654                },
655            );
656            Task::ready(Ok(thread))
657        }
658
659        fn authenticate(
660            &self,
661            _method_id: acp::AuthMethodId,
662            _cx: &mut App,
663        ) -> Task<gpui::Result<()>> {
664            unimplemented!()
665        }
666
667        fn prompt(
668            &self,
669            _id: Option<UserMessageId>,
670            params: acp::PromptRequest,
671            cx: &mut App,
672        ) -> Task<gpui::Result<acp::PromptResponse>> {
673            let mut sessions = self.sessions.lock();
674            let Session {
675                thread,
676                response_tx,
677            } = sessions.get_mut(&params.session_id).unwrap();
678            let mut tasks = vec![];
679            if self.next_prompt_updates.lock().is_empty() {
680                let (tx, rx) = oneshot::channel();
681                response_tx.replace(tx);
682                cx.spawn(async move |_| {
683                    let stop_reason = rx.await?;
684                    Ok(acp::PromptResponse::new(stop_reason))
685                })
686            } else {
687                for update in self.next_prompt_updates.lock().drain(..) {
688                    let thread = thread.clone();
689                    let update = update.clone();
690                    let permission_request = if let acp::SessionUpdate::ToolCall(tool_call) =
691                        &update
692                        && let Some(options) = self.permission_requests.get(&tool_call.tool_call_id)
693                    {
694                        Some((tool_call.clone(), options.clone()))
695                    } else {
696                        None
697                    };
698                    let task = cx.spawn(async move |cx| {
699                        if let Some((tool_call, options)) = permission_request {
700                            thread
701                                .update(cx, |thread, cx| {
702                                    thread.request_tool_call_authorization(
703                                        tool_call.clone().into(),
704                                        options.clone(),
705                                        cx,
706                                    )
707                                })??
708                                .await;
709                        }
710                        thread.update(cx, |thread, cx| {
711                            thread.handle_session_update(update.clone(), cx).unwrap();
712                        })?;
713                        anyhow::Ok(())
714                    });
715                    tasks.push(task);
716                }
717
718                cx.spawn(async move |_| {
719                    try_join_all(tasks).await?;
720                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
721                })
722            }
723        }
724
725        fn cancel(&self, session_id: &acp::SessionId, _cx: &mut App) {
726            if let Some(end_turn_tx) = self
727                .sessions
728                .lock()
729                .get_mut(session_id)
730                .unwrap()
731                .response_tx
732                .take()
733            {
734                end_turn_tx.send(acp::StopReason::Cancelled).unwrap();
735            }
736        }
737
738        fn set_title(
739            &self,
740            _session_id: &acp::SessionId,
741            _cx: &App,
742        ) -> Option<Rc<dyn AgentSessionSetTitle>> {
743            Some(Rc::new(StubAgentSessionSetTitle))
744        }
745
746        fn truncate(
747            &self,
748            _session_id: &agent_client_protocol::SessionId,
749            _cx: &App,
750        ) -> Option<Rc<dyn AgentSessionTruncate>> {
751            Some(Rc::new(StubAgentSessionEditor))
752        }
753
754        fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
755            self
756        }
757    }
758
759    struct StubAgentSessionSetTitle;
760
761    impl AgentSessionSetTitle for StubAgentSessionSetTitle {
762        fn run(&self, _title: SharedString, _cx: &mut App) -> Task<Result<()>> {
763            Task::ready(Ok(()))
764        }
765    }
766
767    struct StubAgentSessionEditor;
768
769    impl AgentSessionTruncate for StubAgentSessionEditor {
770        fn run(&self, _: UserMessageId, _: &mut App) -> Task<Result<()>> {
771            Task::ready(Ok(()))
772        }
773    }
774
775    #[derive(Clone)]
776    struct StubModelSelector {
777        selected_model: Arc<Mutex<AgentModelInfo>>,
778    }
779
780    impl StubModelSelector {
781        fn new() -> Self {
782            Self {
783                selected_model: Arc::new(Mutex::new(AgentModelInfo {
784                    id: acp::ModelId::new("visual-test-model"),
785                    name: "Visual Test Model".into(),
786                    description: Some("A stub model for visual testing".into()),
787                    icon: Some(AgentModelIcon::Named(ui::IconName::ZedAssistant)),
788                    is_latest: false,
789                    cost: None,
790                })),
791            }
792        }
793    }
794
795    impl AgentModelSelector for StubModelSelector {
796        fn list_models(&self, _cx: &mut App) -> Task<Result<AgentModelList>> {
797            let model = self.selected_model.lock().clone();
798            Task::ready(Ok(AgentModelList::Flat(vec![model])))
799        }
800
801        fn select_model(&self, model_id: acp::ModelId, _cx: &mut App) -> Task<Result<()>> {
802            self.selected_model.lock().id = model_id;
803            Task::ready(Ok(()))
804        }
805
806        fn selected_model(&self, _cx: &mut App) -> Task<Result<AgentModelInfo>> {
807            Task::ready(Ok(self.selected_model.lock().clone()))
808        }
809    }
810
811    impl StubAgentConnection {
812        /// Returns a model selector for this stub connection.
813        pub fn model_selector_impl(&self) -> Rc<dyn AgentModelSelector> {
814            Rc::new(StubModelSelector::new())
815        }
816    }
817}
818
819#[cfg(feature = "test-support")]
820pub use test_support::*;