1use acp_thread::{AgentConnection, StubAgentConnection};
  2use agent_client_protocol as acp;
  3use agent_servers::{AgentServer, AgentServerDelegate};
  4use gpui::{Entity, Task, TestAppContext, VisualTestContext};
  5use project::AgentId;
  6use project::Project;
  7use settings::SettingsStore;
  8use std::any::Any;
  9use std::rc::Rc;
 10
 11use crate::AgentPanel;
 12use crate::agent_panel;
 13
 14pub struct StubAgentServer<C> {
 15    connection: C,
 16    agent_id: AgentId,
 17}
 18
 19impl<C> StubAgentServer<C>
 20where
 21    C: AgentConnection,
 22{
 23    pub fn new(connection: C) -> Self {
 24        Self {
 25            connection,
 26            agent_id: "Test".into(),
 27        }
 28    }
 29
 30    pub fn with_connection_agent_id(mut self) -> Self {
 31        self.agent_id = self.connection.agent_id();
 32        self
 33    }
 34}
 35
 36impl StubAgentServer<StubAgentConnection> {
 37    pub fn default_response() -> Self {
 38        let conn = StubAgentConnection::new();
 39        conn.set_next_prompt_updates(vec![acp::SessionUpdate::AgentMessageChunk(
 40            acp::ContentChunk::new("Default response".into()),
 41        )]);
 42        Self::new(conn)
 43    }
 44}
 45
 46impl<C> AgentServer for StubAgentServer<C>
 47where
 48    C: 'static + AgentConnection + Send + Clone,
 49{
 50    fn logo(&self) -> ui::IconName {
 51        ui::IconName::ZedAgent
 52    }
 53
 54    fn agent_id(&self) -> AgentId {
 55        self.agent_id.clone()
 56    }
 57
 58    fn connect(
 59        &self,
 60        _delegate: AgentServerDelegate,
 61        _project: Entity<Project>,
 62        _cx: &mut gpui::App,
 63    ) -> Task<gpui::Result<Rc<dyn AgentConnection>>> {
 64        Task::ready(Ok(Rc::new(self.connection.clone())))
 65    }
 66
 67    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
 68        self
 69    }
 70}
 71
 72pub fn init_test(cx: &mut TestAppContext) {
 73    cx.update(|cx| {
 74        let settings_store = SettingsStore::test(cx);
 75        cx.set_global(settings_store);
 76        theme_settings::init(theme::LoadThemes::JustBase, cx);
 77        editor::init(cx);
 78        release_channel::init("0.0.0".parse().unwrap(), cx);
 79        agent_panel::init(cx);
 80    });
 81}
 82
 83pub fn open_thread_with_connection(
 84    panel: &Entity<AgentPanel>,
 85    connection: StubAgentConnection,
 86    cx: &mut VisualTestContext,
 87) {
 88    panel.update_in(cx, |panel, window, cx| {
 89        panel.open_external_thread_with_server(
 90            Rc::new(StubAgentServer::new(connection)),
 91            window,
 92            cx,
 93        );
 94    });
 95    cx.run_until_parked();
 96}
 97
 98pub fn open_thread_with_custom_connection<C>(
 99    panel: &Entity<AgentPanel>,
100    connection: C,
101    cx: &mut VisualTestContext,
102) where
103    C: 'static + AgentConnection + Send + Clone,
104{
105    panel.update_in(cx, |panel, window, cx| {
106        panel.open_external_thread_with_server(
107            Rc::new(StubAgentServer::new(connection).with_connection_agent_id()),
108            window,
109            cx,
110        );
111    });
112    cx.run_until_parked();
113}
114
115pub fn send_message(panel: &Entity<AgentPanel>, cx: &mut VisualTestContext) {
116    let thread_view = panel.read_with(cx, |panel, cx| panel.active_thread_view(cx).unwrap());
117    let message_editor = thread_view.read_with(cx, |view, _cx| view.message_editor.clone());
118    message_editor.update_in(cx, |editor, window, cx| {
119        editor.set_text("Hello", window, cx);
120    });
121    thread_view.update_in(cx, |view, window, cx| view.send(window, cx));
122    cx.run_until_parked();
123}
124
125pub fn active_session_id(panel: &Entity<AgentPanel>, cx: &VisualTestContext) -> acp::SessionId {
126    panel.read_with(cx, |panel, cx| {
127        let thread = panel.active_agent_thread(cx).unwrap();
128        thread.read(cx).session_id().clone()
129    })
130}