tenant.openagents/omega

1use acp_thread::{AgentConnection, StubAgentConnection};
2use agent_client_protocol::schema::v1 as acp;
3use agent_servers::{AgentServer, AgentServerDelegate};
4use gpui::{
5    App, AppContext as _, Context, Entity, EventEmitter, FocusHandle, Focusable, IntoElement,
6    Pixels, Render, Task, TestAppContext, VisualTestContext, Window, div, px,
7};
8use project::AgentId;
9use project::Project;
10use settings::SettingsStore;
11use std::any::Any;
12use std::cell::RefCell;
13use std::path::Path;
14use std::rc::Rc;
15use std::time::{Duration, SystemTime, UNIX_EPOCH};
16use workspace::{MultiWorkspace, Sidebar as WorkspaceSidebar, SidebarEvent, SidebarSide};
17
18use crate::AgentPanel;
19use crate::agent_panel;
20
21thread_local! {
22    static STUB_AGENT_CONNECTION: RefCell<Option<StubAgentConnection>> = const { RefCell::new(None) };
23}
24
25/// Registers a `StubAgentConnection` that will be used by `Agent::Stub`.
26///
27/// Returns the same connection so callers can hold onto it and control
28/// the stub's behavior (e.g. `connection.set_next_prompt_updates(...)`).
29pub fn set_stub_agent_connection(connection: StubAgentConnection) -> StubAgentConnection {
30    STUB_AGENT_CONNECTION.with(|cell| {
31        *cell.borrow_mut() = Some(connection.clone());
32    });
33    connection
34}
35
36/// Returns the shared `StubAgentConnection` used by `Agent::Stub`,
37/// creating a default one if none was registered.
38pub fn stub_agent_connection() -> StubAgentConnection {
39    STUB_AGENT_CONNECTION.with(|cell| {
40        let mut borrow = cell.borrow_mut();
41        borrow.get_or_insert_with(StubAgentConnection::new).clone()
42    })
43}
44
45pub struct StubAgentServer<C> {
46    connection: C,
47    agent_id: AgentId,
48}
49
50impl<C> StubAgentServer<C>
51where
52    C: AgentConnection,
53{
54    pub fn new(connection: C) -> Self {
55        Self {
56            connection,
57            agent_id: "Test".into(),
58        }
59    }
60
61    pub fn with_connection_agent_id(mut self) -> Self {
62        self.agent_id = self.connection.agent_id();
63        self
64    }
65}
66
67impl StubAgentServer<StubAgentConnection> {
68    pub fn default_response() -> Self {
69        let conn = StubAgentConnection::new();
70        conn.set_next_prompt_updates(vec![acp::SessionUpdate::AgentMessageChunk(
71            acp::ContentChunk::new("Default response".into()),
72        )]);
73        Self::new(conn)
74    }
75}
76
77impl<C> AgentServer for StubAgentServer<C>
78where
79    C: 'static + AgentConnection + Send + Clone,
80{
81    fn logo(&self) -> ui::IconName {
82        ui::IconName::OmegaAgent
83    }
84
85    fn agent_id(&self) -> AgentId {
86        self.agent_id.clone()
87    }
88
89    fn connect(
90        &self,
91        _delegate: AgentServerDelegate,
92        _project: Entity<Project>,
93        _cx: &mut gpui::App,
94    ) -> Task<gpui::Result<Rc<dyn AgentConnection>>> {
95        Task::ready(Ok(Rc::new(self.connection.clone())))
96    }
97
98    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
99        self
100    }
101}
102
103pub fn init_test(cx: &mut TestAppContext) {
104    cx.update(|cx| {
105        let settings_store = SettingsStore::test(cx);
106        cx.set_global(settings_store);
107        // Use an isolated DB so parallel tests can't see each other's
108        // persisted records (e.g. created-worktree records).
109        cx.set_global(db::AppDatabase::test_new());
110        cx.set_global(acp_thread::StubSessionCounter(
111            std::sync::atomic::AtomicUsize::new(0),
112        ));
113        theme_settings::init(theme::LoadThemes::JustBase, cx);
114        editor::init(cx);
115        release_channel::init("0.0.0".parse().unwrap(), cx);
116        agent_panel::init(cx);
117        crate::terminal_thread_metadata_store::TerminalThreadMetadataStore::init_global(cx);
118    });
119}
120
121/// Returns the creation time assigned to a linked worktree's git metadata
122/// directory, mirroring `FakeGitRepository::worktree_created_at` (which uses
123/// the FakeFs directory mtime as a stand-in for the creation time).
124pub async fn fake_worktree_created_at(fs: &dyn fs::Fs, worktree_path: &Path) -> SystemTime {
125    let git_file = fs.load(&worktree_path.join(".git")).await.unwrap();
126    let git_dir = worktree_path.join(git_file.strip_prefix("gitdir:").unwrap().trim());
127    let (seconds, nanos) = fs
128        .metadata(&git_dir)
129        .await
130        .unwrap()
131        .unwrap()
132        .mtime
133        .to_seconds_and_nanos_for_persistence()
134        .unwrap();
135    UNIX_EPOCH + Duration::new(seconds, nanos)
136}
137
138/// Records the worktree in the created-worktrees registry with its actual
139/// (fake) creation time, as the worktree creation flow would. Tests that
140/// expect a worktree to be archivable must call this after setting it up.
141pub async fn record_zed_created_worktree(
142    fs: &dyn fs::Fs,
143    worktree_path: &Path,
144    remote: Option<&remote::RemoteConnectionOptions>,
145    cx: &mut TestAppContext,
146) {
147    let created_at = fake_worktree_created_at(fs, worktree_path).await;
148    cx.update(|cx| {
149        git_ui::created_worktrees::record_created_worktree(worktree_path, remote, created_at, cx)
150    })
151    .await
152    .unwrap();
153}
154
155pub struct TestWorkspaceSidebar {
156    focus_handle: FocusHandle,
157    threads_list_active: bool,
158}
159
160impl TestWorkspaceSidebar {
161    fn new(threads_list_active: bool, cx: &mut Context<Self>) -> Self {
162        Self {
163            focus_handle: cx.focus_handle(),
164            threads_list_active,
165        }
166    }
167}
168
169impl EventEmitter<SidebarEvent> for TestWorkspaceSidebar {}
170
171impl Focusable for TestWorkspaceSidebar {
172    fn focus_handle(&self, _cx: &App) -> FocusHandle {
173        self.focus_handle.clone()
174    }
175}
176
177impl WorkspaceSidebar for TestWorkspaceSidebar {
178    fn width(&self, _cx: &App) -> Pixels {
179        px(300.)
180    }
181
182    fn set_width(&mut self, _width: Option<Pixels>, _cx: &mut Context<Self>) {}
183
184    fn has_notifications(&self, _cx: &App) -> bool {
185        false
186    }
187
188    fn side(&self, _cx: &App) -> SidebarSide {
189        SidebarSide::Left
190    }
191
192    fn is_threads_list_view_active(&self) -> bool {
193        self.threads_list_active
194    }
195}
196
197impl Render for TestWorkspaceSidebar {
198    fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
199        div()
200    }
201}
202
203pub fn register_test_sidebar(
204    threads_list_active: bool,
205    cx: &mut VisualTestContext,
206) -> Entity<TestWorkspaceSidebar> {
207    cx.update(|window, cx| {
208        let multi_workspace = window
209            .root::<MultiWorkspace>()
210            .flatten()
211            .expect("test window should have a MultiWorkspace root");
212        let sidebar = cx.new(|cx| TestWorkspaceSidebar::new(threads_list_active, cx));
213        multi_workspace.update(cx, |multi_workspace, cx| {
214            multi_workspace.register_sidebar(sidebar.clone(), cx);
215        });
216        sidebar
217    })
218}
219
220pub fn open_thread_with_connection(
221    panel: &Entity<AgentPanel>,
222    connection: StubAgentConnection,
223    cx: &mut VisualTestContext,
224) {
225    panel.update_in(cx, |panel, window, cx| {
226        panel.open_external_thread_with_server(
227            Rc::new(StubAgentServer::new(connection)),
228            window,
229            cx,
230        );
231    });
232    cx.run_until_parked();
233}
234
235/// Opens a draft thread against a stub server so the panel's `draft_thread`
236/// pointer is populated for tests that care about draft UX.
237pub fn open_draft_with_connection(
238    panel: &Entity<AgentPanel>,
239    connection: StubAgentConnection,
240    cx: &mut VisualTestContext,
241) {
242    panel.update_in(cx, |panel, window, cx| {
243        panel.open_draft_with_server(Rc::new(StubAgentServer::new(connection)), window, cx);
244    });
245    cx.run_until_parked();
246}
247
248pub fn open_thread_with_custom_connection<C>(
249    panel: &Entity<AgentPanel>,
250    connection: C,
251    cx: &mut VisualTestContext,
252) where
253    C: 'static + AgentConnection + Send + Clone,
254{
255    panel.update_in(cx, |panel, window, cx| {
256        panel.open_external_thread_with_server(
257            Rc::new(StubAgentServer::new(connection).with_connection_agent_id()),
258            window,
259            cx,
260        );
261    });
262    cx.run_until_parked();
263}
264
265pub fn send_message(panel: &Entity<AgentPanel>, cx: &mut VisualTestContext) {
266    let thread_view = panel.read_with(cx, |panel, cx| panel.active_thread_view(cx).unwrap());
267    let message_editor = thread_view.read_with(cx, |view, _cx| view.message_editor.clone());
268    message_editor.update_in(cx, |editor, window, cx| {
269        editor.set_text("Hello", window, cx);
270    });
271    thread_view.update_in(cx, |view, window, cx| view.send(window, cx));
272    cx.run_until_parked();
273}
274
275pub fn type_draft_prompt(panel: &Entity<AgentPanel>, text: &str, cx: &mut VisualTestContext) {
276    let thread_view = panel.read_with(cx, |panel, cx| panel.active_thread_view(cx).unwrap());
277    let message_editor = thread_view.read_with(cx, |view, _cx| view.message_editor.clone());
278    message_editor.update_in(cx, |editor, window, cx| {
279        editor.set_text(text, window, cx);
280    });
281    cx.run_until_parked();
282    // Drain the debounced draft-prompt persist task so the kvp write has
283    // landed by the time we return.
284    cx.executor()
285        .advance_clock(crate::conversation_view::DRAFT_PROMPT_PERSIST_DEBOUNCE * 2);
286    cx.run_until_parked();
287}
288
289pub fn active_session_id(panel: &Entity<AgentPanel>, cx: &VisualTestContext) -> acp::SessionId {
290    panel.read_with(cx, |panel, cx| {
291        let thread = panel.active_agent_thread(cx).unwrap();
292        thread.read(cx).session_id().clone()
293    })
294}
295
296pub fn active_thread_id(
297    panel: &Entity<AgentPanel>,
298    cx: &VisualTestContext,
299) -> crate::thread_metadata_store::ThreadId {
300    panel.read_with(cx, |panel, cx| panel.active_thread_id(cx).unwrap())
301}
302