use acp_thread::{AgentConnection, StubAgentConnection}; use agent_client_protocol as acp; use agent_servers::{AgentServer, AgentServerDelegate}; use gpui::{Entity, Task, TestAppContext, VisualTestContext}; use project::AgentId; use project::Project; use settings::SettingsStore; use std::any::Any; use std::rc::Rc; use crate::AgentPanel; use crate::agent_panel; pub struct StubAgentServer { connection: C, agent_id: AgentId, } impl StubAgentServer where C: AgentConnection, { pub fn new(connection: C) -> Self { Self { connection, agent_id: "Test".into(), } } pub fn with_connection_agent_id(mut self) -> Self { self.agent_id = self.connection.agent_id(); self } } impl StubAgentServer { pub fn default_response() -> Self { let conn = StubAgentConnection::new(); conn.set_next_prompt_updates(vec![acp::SessionUpdate::AgentMessageChunk( acp::ContentChunk::new("Default response".into()), )]); Self::new(conn) } } impl AgentServer for StubAgentServer where C: 'static + AgentConnection + Send + Clone, { fn logo(&self) -> ui::IconName { ui::IconName::ZedAgent } fn agent_id(&self) -> AgentId { self.agent_id.clone() } fn connect( &self, _delegate: AgentServerDelegate, _project: Entity, _cx: &mut gpui::App, ) -> Task>> { Task::ready(Ok(Rc::new(self.connection.clone()))) } fn into_any(self: Rc) -> Rc { self } } pub fn init_test(cx: &mut TestAppContext) { cx.update(|cx| { let settings_store = SettingsStore::test(cx); cx.set_global(settings_store); theme_settings::init(theme::LoadThemes::JustBase, cx); editor::init(cx); release_channel::init("0.0.0".parse().unwrap(), cx); agent_panel::init(cx); }); } pub fn open_thread_with_connection( panel: &Entity, connection: StubAgentConnection, cx: &mut VisualTestContext, ) { panel.update_in(cx, |panel, window, cx| { panel.open_external_thread_with_server( Rc::new(StubAgentServer::new(connection)), window, cx, ); }); cx.run_until_parked(); } pub fn open_thread_with_custom_connection( panel: &Entity, connection: C, cx: &mut VisualTestContext, ) where C: 'static + AgentConnection + Send + Clone, { panel.update_in(cx, |panel, window, cx| { panel.open_external_thread_with_server( Rc::new(StubAgentServer::new(connection).with_connection_agent_id()), window, cx, ); }); cx.run_until_parked(); } pub fn send_message(panel: &Entity, cx: &mut VisualTestContext) { let thread_view = panel.read_with(cx, |panel, cx| panel.active_thread_view(cx).unwrap()); let message_editor = thread_view.read_with(cx, |view, _cx| view.message_editor.clone()); message_editor.update_in(cx, |editor, window, cx| { editor.set_text("Hello", window, cx); }); thread_view.update_in(cx, |view, window, cx| view.send(window, cx)); cx.run_until_parked(); } pub fn active_session_id(panel: &Entity, cx: &VisualTestContext) -> acp::SessionId { panel.read_with(cx, |panel, cx| { let thread = panel.active_agent_thread(cx).unwrap(); thread.read(cx).session_id().clone() }) }