tenant.openagents/omega

1use crate::{AcpThread, ElicitationStore};
2use agent_client_protocol::schema::v1 as acp;
3use anyhow::Result;
4use chrono::{DateTime, Utc};
5use collections::{HashMap, HashSet, IndexMap};
6use gpui::{Entity, SharedString, Task};
7use language_model::DisabledReason;
8use project::{AgentId, Project};
9use serde::{Deserialize, Serialize};
10use std::{any::Any, error::Error, fmt, path::PathBuf, rc::Rc};
11use task::{HideStrategy, SpawnInTerminal, TaskId};
12use ui::{App, IconName};
13use util::path_list::PathList;
14use uuid::Uuid;
15
16/// A user-message ID generated by Zed and passed to agents that support client IDs.
17#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
18pub struct ClientUserMessageId(SharedString);
19
20impl ClientUserMessageId {
21    pub fn new() -> Self {
22        Self(Uuid::new_v4().to_string().into())
23    }
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]
27pub struct AgentModelId(SharedString);
28
29impl AgentModelId {
30    pub fn new(id: impl Into<Self>) -> Self {
31        id.into()
32    }
33
34    pub fn as_str(&self) -> &str {
35        self.0.as_ref()
36    }
37}
38
39impl AsRef<str> for AgentModelId {
40    fn as_ref(&self) -> &str {
41        self.as_str()
42    }
43}
44
45impl fmt::Display for AgentModelId {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        self.0.fmt(f)
48    }
49}
50
51impl From<SharedString> for AgentModelId {
52    fn from(id: SharedString) -> Self {
53        Self(id)
54    }
55}
56
57impl From<String> for AgentModelId {
58    fn from(id: String) -> Self {
59        Self(SharedString::from(id))
60    }
61}
62
63impl From<&str> for AgentModelId {
64    fn from(id: &str) -> Self {
65        Self(SharedString::from(id.to_owned()))
66    }
67}
68
69pub fn build_terminal_auth_task(
70    id: String,
71    label: String,
72    command: String,
73    args: Vec<String>,
74    env: HashMap<String, String>,
75) -> SpawnInTerminal {
76    SpawnInTerminal {
77        id: TaskId(id),
78        full_label: label.clone(),
79        label: label.clone(),
80        command: Some(command),
81        args,
82        command_label: label,
83        env,
84        use_new_terminal: true,
85        allow_concurrent_runs: true,
86        hide: HideStrategy::Always,
87        ..Default::default()
88    }
89}
90
91pub trait AgentConnection {
92    fn agent_id(&self) -> AgentId;
93
94    fn telemetry_id(&self) -> SharedString;
95
96    fn agent_version(&self) -> Option<SharedString> {
97        None
98    }
99
100    fn new_session(
101        self: Rc<Self>,
102        project: Entity<Project>,
103        _work_dirs: PathList,
104        cx: &mut App,
105    ) -> Task<Result<Entity<AcpThread>>>;
106
107    /// Whether this agent supports loading existing sessions.
108    fn supports_load_session(&self) -> bool {
109        false
110    }
111
112    /// Load an existing session by ID.
113    fn load_session(
114        self: Rc<Self>,
115        _session_id: acp::SessionId,
116        _project: Entity<Project>,
117        _work_dirs: PathList,
118        _title: Option<SharedString>,
119        _cx: &mut App,
120    ) -> Task<Result<Entity<AcpThread>>> {
121        Task::ready(Err(anyhow::Error::msg("Loading sessions is not supported")))
122    }
123
124    /// Whether this agent supports closing existing sessions.
125    fn supports_close_session(&self) -> bool {
126        false
127    }
128
129    /// Close an existing session. Allows the agent to free the session from memory.
130    fn close_session(
131        self: Rc<Self>,
132        _session_id: &acp::SessionId,
133        _cx: &mut App,
134    ) -> Task<Result<()>> {
135        Task::ready(Err(anyhow::Error::msg("Closing sessions is not supported")))
136    }
137
138    /// Whether this agent supports resuming existing sessions without loading history.
139    fn supports_resume_session(&self) -> bool {
140        false
141    }
142
143    /// Resume an existing session by ID without replaying previous messages.
144    fn resume_session(
145        self: Rc<Self>,
146        _session_id: acp::SessionId,
147        _project: Entity<Project>,
148        _work_dirs: PathList,
149        _title: Option<SharedString>,
150        _cx: &mut App,
151    ) -> Task<Result<Entity<AcpThread>>> {
152        Task::ready(Err(anyhow::Error::msg(
153            "Resuming sessions is not supported",
154        )))
155    }
156
157    /// Whether this agent supports showing session history.
158    fn supports_session_history(&self) -> bool {
159        self.supports_load_session() || self.supports_resume_session()
160    }
161
162    /// Whether this agent supports additional session directories.
163    fn supports_session_additional_directories(&self) -> bool {
164        false
165    }
166
167    fn auth_methods(&self) -> &[acp::AuthMethod];
168
169    fn terminal_auth_task(
170        &self,
171        _method: &acp::AuthMethodId,
172        _cx: &App,
173    ) -> Option<Task<Result<SpawnInTerminal>>> {
174        None
175    }
176
177    fn authenticate(&self, method: acp::AuthMethodId, cx: &mut App) -> Task<Result<()>>;
178
179    fn supports_logout(&self) -> bool {
180        false
181    }
182
183    fn logout(&self, _cx: &mut App) -> Task<Result<()>> {
184        Task::ready(Err(anyhow::Error::msg("Logout is not supported")))
185    }
186
187    /// Returns a capability for agents that accept client-generated user message IDs.
188    fn client_user_message_ids(
189        &self,
190        _cx: &App,
191    ) -> Option<Rc<dyn AgentSessionClientUserMessageIds>> {
192        None
193    }
194
195    fn prompt(&self, params: acp::PromptRequest, cx: &mut App)
196    -> Task<Result<acp::PromptResponse>>;
197
198    fn retry(&self, _session_id: &acp::SessionId, _cx: &App) -> Option<Rc<dyn AgentSessionRetry>> {
199        None
200    }
201
202    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App);
203
204    /// Request-scoped elicitations are connection-level because they can arrive before a session
205    /// thread exists. Session-scoped elicitations stay in the thread timeline, but use
206    /// `ElicitationStore` for shared processing.
207    fn request_elicitations(&self) -> Option<Entity<ElicitationStore>> {
208        None
209    }
210
211    fn truncate(
212        &self,
213        _session_id: &acp::SessionId,
214        _cx: &App,
215    ) -> Option<Rc<dyn AgentSessionTruncate>> {
216        None
217    }
218
219    fn set_title(
220        &self,
221        _session_id: &acp::SessionId,
222        _cx: &App,
223    ) -> Option<Rc<dyn AgentSessionSetTitle>> {
224        None
225    }
226
227    /// Returns this agent as an [Rc<dyn ModelSelector>] if the model selection capability is supported.
228    ///
229    /// If the agent does not support model selection, returns [None].
230    /// This allows sharing the selector in UI components.
231    fn model_selector(&self, _session_id: &acp::SessionId) -> Option<Rc<dyn AgentModelSelector>> {
232        None
233    }
234
235    fn telemetry(&self) -> Option<Rc<dyn AgentTelemetry>> {
236        None
237    }
238
239    fn session_modes(
240        &self,
241        _session_id: &acp::SessionId,
242        _cx: &App,
243    ) -> Option<Rc<dyn AgentSessionModes>> {
244        None
245    }
246
247    fn session_config_options(
248        &self,
249        _session_id: &acp::SessionId,
250        _cx: &App,
251    ) -> Option<Rc<dyn AgentSessionConfigOptions>> {
252        None
253    }
254
255    fn session_list(&self, _cx: &mut App) -> Option<Rc<dyn AgentSessionList>> {
256        None
257    }
258
259    fn into_any(self: Rc<Self>) -> Rc<dyn Any>;
260}
261
262impl dyn AgentConnection {
263    pub fn downcast<T: 'static + AgentConnection + Sized>(self: Rc<Self>) -> Option<Rc<T>> {
264        self.into_any().downcast().ok()
265    }
266}
267
268pub trait AgentSessionTruncate {
269    fn run(&self, client_user_message_id: ClientUserMessageId, cx: &mut App) -> Task<Result<()>>;
270}
271
272pub trait AgentSessionClientUserMessageIds {
273    fn new_id(&self) -> ClientUserMessageId {
274        ClientUserMessageId::new()
275    }
276
277    fn prompt(
278        &self,
279        client_user_message_id: ClientUserMessageId,
280        params: acp::PromptRequest,
281        cx: &mut App,
282    ) -> Task<Result<acp::PromptResponse>>;
283}
284
285pub trait AgentSessionRetry {
286    fn run(&self, cx: &mut App) -> Task<Result<acp::PromptResponse>>;
287}
288
289pub trait AgentSessionSetTitle {
290    fn run(&self, title: SharedString, cx: &mut App) -> Task<Result<()>>;
291}
292
293pub trait AgentTelemetry {
294    /// A representation of the current thread state that can be serialized for
295    /// storage with telemetry events.
296    fn thread_data(
297        &self,
298        session_id: &acp::SessionId,
299        cx: &mut App,
300    ) -> Task<Result<serde_json::Value>>;
301}
302
303pub trait AgentSessionModes {
304    fn current_mode(&self) -> acp::SessionModeId;
305
306    fn all_modes(&self) -> Vec<acp::SessionMode>;
307
308    fn set_mode(&self, mode: acp::SessionModeId, cx: &mut App) -> Task<Result<()>>;
309}
310
311pub trait AgentSessionConfigOptions {
312    /// Get all current config options with their state
313    fn config_options(&self) -> Vec<acp::SessionConfigOption>;
314
315    /// Set a config option value
316    /// Returns the full updated list of config options
317    fn set_config_option(
318        &self,
319        config_id: acp::SessionConfigId,
320        value: acp::SessionConfigOptionValue,
321        cx: &mut App,
322    ) -> Task<Result<Vec<acp::SessionConfigOption>>>;
323
324    /// Whenever the config options are updated the receiver will be notified.
325    /// Optional for agents that don't update their config options dynamically.
326    fn watch(&self, _cx: &mut App) -> Option<watch::Receiver<()>> {
327        None
328    }
329}
330
331#[derive(Debug, Clone, Default)]
332pub struct AgentSessionListRequest {
333    pub cwd: Option<PathBuf>,
334    pub cursor: Option<String>,
335    pub meta: Option<acp::Meta>,
336}
337
338#[derive(Debug, Clone)]
339pub struct AgentSessionListResponse {
340    pub sessions: Vec<AgentSessionInfo>,
341    pub next_cursor: Option<String>,
342    pub meta: Option<acp::Meta>,
343}
344
345impl AgentSessionListResponse {
346    pub fn new(sessions: Vec<AgentSessionInfo>) -> Self {
347        Self {
348            sessions,
349            next_cursor: None,
350            meta: None,
351        }
352    }
353}
354
355#[derive(Debug, Clone, PartialEq)]
356pub struct AgentSessionInfo {
357    pub session_id: acp::SessionId,
358    pub work_dirs: Option<PathList>,
359    pub title: Option<SharedString>,
360    pub updated_at: Option<DateTime<Utc>>,
361    pub created_at: Option<DateTime<Utc>>,
362    pub meta: Option<acp::Meta>,
363}
364
365impl AgentSessionInfo {
366    pub fn new(session_id: impl Into<acp::SessionId>) -> Self {
367        Self {
368            session_id: session_id.into(),
369            work_dirs: None,
370            title: None,
371            updated_at: None,
372            created_at: None,
373            meta: None,
374        }
375    }
376}
377
378#[derive(Debug, Clone)]
379pub enum SessionListUpdate {
380    Refresh,
381    SessionInfo {
382        session_id: acp::SessionId,
383        update: acp::SessionInfoUpdate,
384    },
385}
386
387pub trait AgentSessionList {
388    fn list_sessions(
389        &self,
390        request: AgentSessionListRequest,
391        cx: &mut App,
392    ) -> Task<Result<AgentSessionListResponse>>;
393
394    fn supports_delete(&self) -> bool {
395        false
396    }
397
398    fn delete_session(&self, _session_id: &acp::SessionId, _cx: &mut App) -> Task<Result<()>> {
399        Task::ready(Err(anyhow::anyhow!("delete_session not supported")))
400    }
401
402    fn delete_sessions(&self, _cx: &mut App) -> Task<Result<()>> {
403        Task::ready(Err(anyhow::anyhow!("delete_sessions not supported")))
404    }
405
406    fn watch(&self, _cx: &mut App) -> Option<async_channel::Receiver<SessionListUpdate>> {
407        None
408    }
409
410    fn notify_refresh(&self) {}
411
412    fn into_any(self: Rc<Self>) -> Rc<dyn Any>;
413}
414
415impl dyn AgentSessionList {
416    pub fn downcast<T: 'static + AgentSessionList + Sized>(self: Rc<Self>) -> Option<Rc<T>> {
417        self.into_any().downcast().ok()
418    }
419}
420
421#[derive(Debug)]
422pub struct AuthRequired {
423    pub description: Option<String>,
424}
425
426impl AuthRequired {
427    pub fn new() -> Self {
428        Self { description: None }
429    }
430
431    pub fn with_description(mut self, description: String) -> Self {
432        self.description = Some(description);
433        self
434    }
435}
436
437impl Error for AuthRequired {}
438impl fmt::Display for AuthRequired {
439    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
440        write!(f, "Authentication required")
441    }
442}
443
444/// Trait for agents that support listing, selecting, and querying language models.
445///
446/// This is an optional capability; agents indicate support via [AgentConnection::model_selector].
447pub trait AgentModelSelector: 'static {
448    /// Lists all available language models for this agent.
449    ///
450    /// # Parameters
451    /// - `cx`: The GPUI app context for async operations and global access.
452    ///
453    /// # Returns
454    /// A task resolving to the list of models or an error (e.g., if no models are configured).
455    fn list_models(&self, cx: &mut App) -> Task<Result<AgentModelList>>;
456
457    /// Selects a model for a specific session (thread).
458    ///
459    /// # Parameters
460    /// - `model_id`: The model to select (should be one from [list_models]).
461    /// - `cx`: The GPUI app context.
462    ///
463    /// # Returns
464    /// A task resolving to `Ok(())` on success or an error.
465    fn select_model(&self, model_id: AgentModelId, cx: &mut App) -> Task<Result<()>>;
466
467    /// Retrieves the currently selected model for a specific session (thread).
468    ///
469    /// # Parameters
470    /// - `cx`: The GPUI app context.
471    ///
472    /// # Returns
473    /// A task resolving to the selected model (always set) or an error (e.g., session not found).
474    fn selected_model(&self, cx: &mut App) -> Task<Result<AgentModelInfo>>;
475
476    fn favorite_model_ids(&self, _cx: &mut App) -> HashSet<AgentModelId> {
477        HashSet::default()
478    }
479
480    fn toggle_favorite_model(&self, _model_id: AgentModelId, _should_be_favorite: bool, _cx: &App) {
481    }
482
483    /// Whenever the model list is updated the receiver will be notified.
484    /// Optional for agents that don't update their model list.
485    fn watch(&self, _cx: &mut App) -> Option<watch::Receiver<()>> {
486        None
487    }
488
489    /// Returns whether the model picker should render a footer.
490    fn should_render_footer(&self) -> bool {
491        false
492    }
493}
494
495/// Icon for a model in the model selector.
496#[derive(Debug, Clone, PartialEq, Eq)]
497pub enum AgentModelIcon {
498    /// A built-in icon from Zed's icon set.
499    Named(IconName),
500    /// Path to a custom SVG icon file.
501    Path(SharedString),
502}
503
504#[derive(Debug, Clone, PartialEq, Eq)]
505pub struct AgentModelInfo {
506    pub id: AgentModelId,
507    pub name: SharedString,
508    pub description: Option<SharedString>,
509    pub icon: Option<AgentModelIcon>,
510    pub is_latest: bool,
511    pub cost: Option<SharedString>,
512    pub disabled: Option<DisabledReason>,
513}
514
515#[derive(Debug, Clone, PartialEq, Eq, Hash)]
516pub struct AgentModelGroupName(pub SharedString);
517
518#[derive(Debug, Clone)]
519pub enum AgentModelList {
520    Flat(Vec<AgentModelInfo>),
521    Grouped(IndexMap<AgentModelGroupName, Vec<AgentModelInfo>>),
522}
523
524impl AgentModelList {
525    pub fn is_empty(&self) -> bool {
526        match self {
527            AgentModelList::Flat(models) => models.is_empty(),
528            AgentModelList::Grouped(groups) => groups.is_empty(),
529        }
530    }
531
532    pub fn is_flat(&self) -> bool {
533        matches!(self, AgentModelList::Flat(_))
534    }
535}
536
537#[derive(Debug, Clone)]
538pub struct PermissionOptionChoice {
539    pub allow: acp::PermissionOption,
540    pub deny: acp::PermissionOption,
541    pub sub_patterns: Vec<String>,
542}
543
544impl PermissionOptionChoice {
545    pub fn label(&self) -> SharedString {
546        self.allow.name.clone().into()
547    }
548
549    /// Build a `SelectedPermissionOutcome` for this choice.
550    ///
551    /// If the choice carries `sub_patterns`, they are attached as
552    /// `SelectedPermissionParams::Terminal`.
553    pub fn build_outcome(&self, is_allow: bool) -> crate::SelectedPermissionOutcome {
554        let option = if is_allow { &self.allow } else { &self.deny };
555
556        let params = if !self.sub_patterns.is_empty() {
557            Some(crate::SelectedPermissionParams::Terminal {
558                patterns: self.sub_patterns.clone(),
559            })
560        } else {
561            None
562        };
563
564        crate::SelectedPermissionOutcome::new(option.option_id.clone(), option.kind).params(params)
565    }
566}
567
568/// Pairs a tool's permission pattern with its display name
569///
570/// For example, a pattern of `^cargo\\s+build(\\s|$)` would display as `cargo
571/// build`. It's handy to keep these together rather than trying to derive
572/// one from the other.
573#[derive(Debug, Clone, PartialEq)]
574pub struct PermissionPattern {
575    pub pattern: String,
576    pub display_name: String,
577}
578
579#[derive(Debug, Clone)]
580pub enum PermissionOptions {
581    Flat(Vec<acp::PermissionOption>),
582    Dropdown(Vec<PermissionOptionChoice>),
583    DropdownWithPatterns {
584        choices: Vec<PermissionOptionChoice>,
585        patterns: Vec<PermissionPattern>,
586        tool_name: String,
587    },
588}
589
590impl PermissionOptions {
591    pub fn is_empty(&self) -> bool {
592        match self {
593            PermissionOptions::Flat(options) => options.is_empty(),
594            PermissionOptions::Dropdown(options) => options.is_empty(),
595            PermissionOptions::DropdownWithPatterns { choices, .. } => choices.is_empty(),
596        }
597    }
598
599    pub fn first_option_of_kind(
600        &self,
601        kind: acp::PermissionOptionKind,
602    ) -> Option<&acp::PermissionOption> {
603        match self {
604            PermissionOptions::Flat(options) => options.iter().find(|option| option.kind == kind),
605            PermissionOptions::Dropdown(options) => options.iter().find_map(|choice| {
606                if choice.allow.kind == kind {
607                    Some(&choice.allow)
608                } else if choice.deny.kind == kind {
609                    Some(&choice.deny)
610                } else {
611                    None
612                }
613            }),
614            PermissionOptions::DropdownWithPatterns { choices, .. } => {
615                choices.iter().find_map(|choice| {
616                    if choice.allow.kind == kind {
617                        Some(&choice.allow)
618                    } else if choice.deny.kind == kind {
619                        Some(&choice.deny)
620                    } else {
621                        None
622                    }
623                })
624            }
625        }
626    }
627
628    pub fn allow_once_option_id(&self) -> Option<acp::PermissionOptionId> {
629        self.first_option_of_kind(acp::PermissionOptionKind::AllowOnce)
630            .map(|option| option.option_id.clone())
631    }
632
633    pub fn deny_once_option_id(&self) -> Option<acp::PermissionOptionId> {
634        self.first_option_of_kind(acp::PermissionOptionKind::RejectOnce)
635            .map(|option| option.option_id.clone())
636    }
637
638    /// Build a `SelectedPermissionOutcome` for the `DropdownWithPatterns`
639    /// variant when the user has checked specific pattern indices.
640    ///
641    /// Returns `Some` with the always-allow/deny outcome when at least one
642    /// pattern is checked. Returns `None` when zero patterns are checked,
643    /// signaling that the caller should degrade to allow-once / deny-once.
644    ///
645    /// Panics (debug) or returns `None` (release) if called on a non-
646    /// `DropdownWithPatterns` variant.
647    pub fn build_outcome_for_checked_patterns(
648        &self,
649        checked_indices: &[usize],
650        is_allow: bool,
651    ) -> Option<crate::SelectedPermissionOutcome> {
652        let PermissionOptions::DropdownWithPatterns {
653            choices, patterns, ..
654        } = self
655        else {
656            debug_assert!(
657                false,
658                "build_outcome_for_checked_patterns called on non-DropdownWithPatterns"
659            );
660            return None;
661        };
662
663        let checked_patterns: Vec<String> = patterns
664            .iter()
665            .enumerate()
666            .filter(|(index, _)| checked_indices.contains(index))
667            .map(|(_, cp)| cp.pattern.clone())
668            .collect();
669
670        if checked_patterns.is_empty() {
671            return None;
672        }
673
674        // Use the first choice (the "Always" choice) as the base for the outcome.
675        let always_choice = choices.first()?;
676        let option = if is_allow {
677            &always_choice.allow
678        } else {
679            &always_choice.deny
680        };
681
682        let outcome = crate::SelectedPermissionOutcome::new(option.option_id.clone(), option.kind)
683            .params(Some(crate::SelectedPermissionParams::Terminal {
684                patterns: checked_patterns,
685            }));
686        Some(outcome)
687    }
688}
689
690#[cfg(feature = "test-support")]
691mod test_support {
692    //! Test-only stubs and helpers for acp_thread.
693    //!
694    //! This module is gated by the `test-support` feature and is not included
695    //! in production builds. It provides:
696    //! - `StubAgentConnection` for mocking agent connections in tests
697    //! - `create_test_png_base64` for generating test images
698
699    use std::sync::Arc;
700    use std::sync::atomic::{AtomicUsize, Ordering};
701
702    use action_log::ActionLog;
703    use collections::HashMap;
704    use futures::{channel::oneshot, future::try_join_all};
705    use gpui::{AppContext as _, WeakEntity};
706    use parking_lot::Mutex;
707
708    use crate::AuthorizationKind;
709
710    use super::*;
711
712    /// Creates a PNG image encoded as base64 for testing.
713    ///
714    /// Generates a solid-color PNG of the specified dimensions and returns
715    /// it as a base64-encoded string suitable for use in `ImageContent`.
716    pub fn create_test_png_base64(width: u32, height: u32, color: [u8; 4]) -> String {
717        use image::ImageEncoder as _;
718
719        let mut png_data = Vec::new();
720        {
721            let encoder = image::codecs::png::PngEncoder::new(&mut png_data);
722            let mut pixels = Vec::with_capacity((width * height * 4) as usize);
723            for _ in 0..(width * height) {
724                pixels.extend_from_slice(&color);
725            }
726            encoder
727                .write_image(&pixels, width, height, image::ExtendedColorType::Rgba8)
728                .expect("Failed to encode PNG");
729        }
730
731        use image::EncodableLayout as _;
732        base64::Engine::encode(
733            &base64::engine::general_purpose::STANDARD,
734            png_data.as_bytes(),
735        )
736    }
737
738    /// Test-scoped counter for generating unique session IDs across all
739    /// `StubAgentConnection` instances within a test case. Set as a GPUI
740    /// global in test init so each case starts fresh.
741    pub struct StubSessionCounter(pub AtomicUsize);
742    impl gpui::Global for StubSessionCounter {}
743
744    impl StubSessionCounter {
745        pub fn next(cx: &App) -> usize {
746            cx.try_global::<Self>()
747                .map(|g| g.0.fetch_add(1, Ordering::SeqCst))
748                .unwrap_or_else(|| {
749                    static FALLBACK: AtomicUsize = AtomicUsize::new(0);
750                    FALLBACK.fetch_add(1, Ordering::SeqCst)
751                })
752        }
753    }
754
755    #[derive(Clone)]
756    pub struct StubAgentConnection {
757        sessions: Arc<Mutex<HashMap<acp::SessionId, Session>>>,
758        permission_requests: HashMap<acp::ToolCallId, PermissionOptions>,
759        next_prompt_updates: Arc<Mutex<Vec<acp::SessionUpdate>>>,
760        supports_load_session: bool,
761        supports_session_additional_directories: bool,
762        agent_id: AgentId,
763        telemetry_id: SharedString,
764    }
765
766    struct Session {
767        thread: WeakEntity<AcpThread>,
768        response_tx: Option<oneshot::Sender<acp::StopReason>>,
769    }
770
771    impl Default for StubAgentConnection {
772        fn default() -> Self {
773            Self::new()
774        }
775    }
776
777    impl StubAgentConnection {
778        pub fn new() -> Self {
779            Self {
780                next_prompt_updates: Default::default(),
781                permission_requests: HashMap::default(),
782                sessions: Arc::default(),
783                supports_load_session: false,
784                supports_session_additional_directories: false,
785                agent_id: AgentId::new("stub"),
786                telemetry_id: "stub".into(),
787            }
788        }
789
790        pub fn set_next_prompt_updates(&self, updates: Vec<acp::SessionUpdate>) {
791            *self.next_prompt_updates.lock() = updates;
792        }
793
794        pub fn with_permission_requests(
795            mut self,
796            permission_requests: HashMap<acp::ToolCallId, PermissionOptions>,
797        ) -> Self {
798            self.permission_requests = permission_requests;
799            self
800        }
801
802        pub fn with_supports_load_session(mut self, supports_load_session: bool) -> Self {
803            self.supports_load_session = supports_load_session;
804            self
805        }
806
807        pub fn with_supports_session_additional_directories(
808            mut self,
809            supports_session_additional_directories: bool,
810        ) -> Self {
811            self.supports_session_additional_directories = supports_session_additional_directories;
812            self
813        }
814
815        pub fn with_agent_id(mut self, agent_id: AgentId) -> Self {
816            self.agent_id = agent_id;
817            self
818        }
819
820        pub fn with_telemetry_id(mut self, telemetry_id: SharedString) -> Self {
821            self.telemetry_id = telemetry_id;
822            self
823        }
824
825        fn create_session(
826            self: Rc<Self>,
827            session_id: acp::SessionId,
828            project: Entity<Project>,
829            work_dirs: PathList,
830            title: Option<SharedString>,
831            cx: &mut gpui::App,
832        ) -> Entity<AcpThread> {
833            let action_log = cx.new(|_| ActionLog::new(project.clone()));
834            let thread = cx.new(|cx| {
835                AcpThread::new(
836                    None,
837                    title,
838                    Some(work_dirs),
839                    self.clone(),
840                    project,
841                    action_log,
842                    session_id.clone(),
843                    watch::Receiver::constant(
844                        acp::PromptCapabilities::new()
845                            .image(true)
846                            .audio(true)
847                            .embedded_context(true),
848                    ),
849                    cx,
850                )
851            });
852            self.sessions.lock().insert(
853                session_id,
854                Session {
855                    thread: thread.downgrade(),
856                    response_tx: None,
857                },
858            );
859            thread
860        }
861
862        pub fn send_update(
863            &self,
864            session_id: acp::SessionId,
865            update: acp::SessionUpdate,
866            cx: &mut App,
867        ) {
868            assert!(
869                self.next_prompt_updates.lock().is_empty(),
870                "Use either send_update or set_next_prompt_updates"
871            );
872
873            self.sessions
874                .lock()
875                .get(&session_id)
876                .unwrap()
877                .thread
878                .update(cx, |thread, cx| {
879                    thread.handle_session_update(update, cx).unwrap();
880                })
881                .unwrap();
882        }
883
884        pub fn end_turn(&self, session_id: acp::SessionId, stop_reason: acp::StopReason) {
885            self.sessions
886                .lock()
887                .get_mut(&session_id)
888                .unwrap()
889                .response_tx
890                .take()
891                .expect("No pending turn")
892                .send(stop_reason)
893                .unwrap();
894        }
895    }
896
897    impl AgentConnection for StubAgentConnection {
898        fn agent_id(&self) -> AgentId {
899            self.agent_id.clone()
900        }
901
902        fn telemetry_id(&self) -> SharedString {
903            self.telemetry_id.clone()
904        }
905
906        fn auth_methods(&self) -> &[acp::AuthMethod] {
907            &[]
908        }
909
910        fn model_selector(
911            &self,
912            _session_id: &acp::SessionId,
913        ) -> Option<Rc<dyn AgentModelSelector>> {
914            Some(self.model_selector_impl())
915        }
916
917        fn new_session(
918            self: Rc<Self>,
919            project: Entity<Project>,
920            work_dirs: PathList,
921            cx: &mut gpui::App,
922        ) -> Task<gpui::Result<Entity<AcpThread>>> {
923            let session_id = acp::SessionId::new(StubSessionCounter::next(cx).to_string());
924            let thread = self.create_session(session_id, project, work_dirs, None, cx);
925            Task::ready(Ok(thread))
926        }
927
928        fn supports_load_session(&self) -> bool {
929            self.supports_load_session
930        }
931
932        fn supports_session_additional_directories(&self) -> bool {
933            self.supports_session_additional_directories
934        }
935
936        fn load_session(
937            self: Rc<Self>,
938            session_id: acp::SessionId,
939            project: Entity<Project>,
940            work_dirs: PathList,
941            title: Option<SharedString>,
942            cx: &mut App,
943        ) -> Task<Result<Entity<AcpThread>>> {
944            if !self.supports_load_session {
945                return Task::ready(Err(anyhow::Error::msg("Loading sessions is not supported")));
946            }
947
948            let thread = self.create_session(session_id, project, work_dirs, title, cx);
949            Task::ready(Ok(thread))
950        }
951
952        fn authenticate(
953            &self,
954            _method_id: acp::AuthMethodId,
955            _cx: &mut App,
956        ) -> Task<gpui::Result<()>> {
957            unimplemented!()
958        }
959
960        fn prompt(
961            &self,
962            params: acp::PromptRequest,
963            cx: &mut App,
964        ) -> Task<gpui::Result<acp::PromptResponse>> {
965            let mut sessions = self.sessions.lock();
966            let Session {
967                thread,
968                response_tx,
969            } = sessions.get_mut(&params.session_id).unwrap();
970            let mut tasks = vec![];
971            if self.next_prompt_updates.lock().is_empty() {
972                let (tx, rx) = oneshot::channel();
973                response_tx.replace(tx);
974                cx.spawn(async move |_| {
975                    let stop_reason = rx.await?;
976                    Ok(acp::PromptResponse::new(stop_reason))
977                })
978            } else {
979                for update in self.next_prompt_updates.lock().drain(..) {
980                    let thread = thread.clone();
981                    let update = update.clone();
982                    let permission_request = if let acp::SessionUpdate::ToolCall(tool_call) =
983                        &update
984                        && let Some(options) = self.permission_requests.get(&tool_call.tool_call_id)
985                    {
986                        Some((tool_call.clone(), options.clone()))
987                    } else {
988                        None
989                    };
990                    let task = cx.spawn(async move |cx| {
991                        if let Some((tool_call, options)) = permission_request {
992                            thread
993                                .update(cx, |thread, cx| {
994                                    thread.request_tool_call_authorization(
995                                        tool_call.clone().into(),
996                                        options.clone(),
997                                        AuthorizationKind::PermissionGrant,
998                                        cx,
999                                    )
1000                                })??
1001                                .await;
1002                        }
1003                        thread.update(cx, |thread, cx| {
1004                            thread.handle_session_update(update.clone(), cx).unwrap();
1005                        })?;
1006                        anyhow::Ok(())
1007                    });
1008                    tasks.push(task);
1009                }
1010
1011                cx.spawn(async move |_| {
1012                    try_join_all(tasks).await?;
1013                    Ok(acp::PromptResponse::new(acp::StopReason::EndTurn))
1014                })
1015            }
1016        }
1017
1018        fn client_user_message_ids(
1019            &self,
1020            _cx: &App,
1021        ) -> Option<Rc<dyn AgentSessionClientUserMessageIds>> {
1022            Some(Rc::new(StubAgentSessionClientUserMessageIds {
1023                connection: self.clone(),
1024            }))
1025        }
1026
1027        fn cancel(&self, session_id: &acp::SessionId, _cx: &mut App) {
1028            if let Some(end_turn_tx) = self
1029                .sessions
1030                .lock()
1031                .get_mut(session_id)
1032                .unwrap()
1033                .response_tx
1034                .take()
1035            {
1036                end_turn_tx.send(acp::StopReason::Cancelled).unwrap();
1037            }
1038        }
1039
1040        fn set_title(
1041            &self,
1042            _session_id: &acp::SessionId,
1043            _cx: &App,
1044        ) -> Option<Rc<dyn AgentSessionSetTitle>> {
1045            Some(Rc::new(StubAgentSessionSetTitle))
1046        }
1047
1048        fn truncate(
1049            &self,
1050            _session_id: &acp::SessionId,
1051            _cx: &App,
1052        ) -> Option<Rc<dyn AgentSessionTruncate>> {
1053            Some(Rc::new(StubAgentSessionEditor))
1054        }
1055
1056        fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
1057            self
1058        }
1059    }
1060
1061    struct StubAgentSessionSetTitle;
1062
1063    impl AgentSessionSetTitle for StubAgentSessionSetTitle {
1064        fn run(&self, _title: SharedString, _cx: &mut App) -> Task<Result<()>> {
1065            Task::ready(Ok(()))
1066        }
1067    }
1068
1069    struct StubAgentSessionClientUserMessageIds {
1070        connection: StubAgentConnection,
1071    }
1072
1073    impl AgentSessionClientUserMessageIds for StubAgentSessionClientUserMessageIds {
1074        fn prompt(
1075            &self,
1076            _client_user_message_id: ClientUserMessageId,
1077            params: acp::PromptRequest,
1078            cx: &mut App,
1079        ) -> Task<Result<acp::PromptResponse>> {
1080            self.connection.prompt(params, cx)
1081        }
1082    }
1083
1084    struct StubAgentSessionEditor;
1085
1086    impl AgentSessionTruncate for StubAgentSessionEditor {
1087        fn run(&self, _: ClientUserMessageId, _: &mut App) -> Task<Result<()>> {
1088            Task::ready(Ok(()))
1089        }
1090    }
1091
1092    #[derive(Clone)]
1093    struct StubModelSelector {
1094        selected_model: Arc<Mutex<AgentModelInfo>>,
1095    }
1096
1097    impl StubModelSelector {
1098        fn new() -> Self {
1099            Self {
1100                selected_model: Arc::new(Mutex::new(AgentModelInfo {
1101                    id: AgentModelId::new("visual-test-model"),
1102                    name: "Visual Test Model".into(),
1103                    description: Some("A stub model for visual testing".into()),
1104                    icon: Some(AgentModelIcon::Named(ui::IconName::OmegaAssistant)),
1105                    is_latest: false,
1106                    cost: None,
1107                    disabled: None,
1108                })),
1109            }
1110        }
1111    }
1112
1113    impl AgentModelSelector for StubModelSelector {
1114        fn list_models(&self, _cx: &mut App) -> Task<Result<AgentModelList>> {
1115            let model = self.selected_model.lock().clone();
1116            Task::ready(Ok(AgentModelList::Flat(vec![model])))
1117        }
1118
1119        fn select_model(&self, model_id: AgentModelId, _cx: &mut App) -> Task<Result<()>> {
1120            self.selected_model.lock().id = model_id;
1121            Task::ready(Ok(()))
1122        }
1123
1124        fn selected_model(&self, _cx: &mut App) -> Task<Result<AgentModelInfo>> {
1125            Task::ready(Ok(self.selected_model.lock().clone()))
1126        }
1127    }
1128
1129    impl StubAgentConnection {
1130        /// Returns a model selector for this stub connection.
1131        pub fn model_selector_impl(&self) -> Rc<dyn AgentModelSelector> {
1132            Rc::new(StubModelSelector::new())
1133        }
1134    }
1135}
1136
1137#[cfg(feature = "test-support")]
1138pub use test_support::*;
1139