aether_cli/acp/
testing.rs

1use super::fake_prompt_mcp::FakePromptMcp;
2use super::server::{AcpServer, ServerRunError};
3use super::session::actor::SessionActorInit;
4use super::session::agent_key::AgentKey;
5use super::session::agents::SessionAgents;
6use super::session::config::SessionConfigState;
7use super::session::error::SessionError;
8use super::session::hooks::SessionHooks;
9use super::session::model::{Modes, ValidatedMode};
10use super::session::runtime::{AgentRuntime, RuntimeFactory};
11use super::state::{AcpState, AcpStateConfig};
12use crate::error::CliError;
13use crate::resolve::InitialSessionSelection;
14use crate::settings_args::SettingsSourceArgs;
15use crate::workspace::WorkspaceManager;
16use crate::workspace::testing::StdCopyCloner;
17use acp_utils::notifications::{AuthMethodsUpdatedParams, McpNotification};
18use acp_utils::testing::{TestPeer, initialize_request};
19use aether_auth::OAuthCredentialStorage;
20use aether_core::agent_spec::{AgentSpec, AgentSpecExposure};
21use aether_core::core::{AgentBuilder, AgentHandle, Prompt};
22use aether_core::events::{AgentEvent, Command, MessageEvent, ToolEvent, TurnEvent, TurnOutcome, UserCommand};
23use aether_core::mcp::{McpBuilder, McpSession, ServerFactory, mcp};
24use aether_project::AgentCatalog;
25use aether_sessions::SessionStore;
26use aether_sessions::{SessionControlEvent, SessionEvent, SessionMeta, UserEvent, last_agent_from_events};
27use agent_client_protocol::schema::v2::{
28    AbsolutePath, InitializeResponse, ReplayFrom, ReplayFromStart, ResumeSessionRequest, SessionId, SessionUpdate,
29    StateUpdate, StopReason,
30};
31use agent_client_protocol::{Agent, Channel, Client, ConnectionTo, on_receive_notification};
32use futures::FutureExt;
33use llm::testing::FakeLlmProvider;
34use llm::{ChatMessage, Context, LlmModel, LlmResponse, SessionUsageEvent, StreamingModelProvider};
35use llm::{MessageId, ProviderConnectionOverrides, ToolCallRequest, ToolCallResult};
36use mcp_utils::client::{InMemoryServerSpec, McpServer, McpTransport, ToolExposure};
37use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse, completed_task_payload};
38use rmcp::RoleServer;
39use rmcp::model::{CallToolResult, CreateTaskResult, DetailedTask, Task, TaskPayload, TaskStatus};
40use rmcp::service::DynService;
41use std::collections::HashMap;
42use std::future::Future;
43use std::net::SocketAddr;
44use std::path::PathBuf;
45use std::sync::{Arc, Mutex};
46use tokio::sync::{mpsc, oneshot, watch};
47use tokio::task::{JoinHandle, JoinSet, LocalSet};
48use tokio_util::sync::CancellationToken;
49
50const PLANNER_REPLY: &str = "planner reply";
51const CODER_REPLY: &str = "coder reply";
52
53pub struct AcpTestHarness {
54    pub client_cx: ConnectionTo<Agent>,
55    pub peer: TestPeer,
56    pub initialize_response: InitializeResponse,
57    connection: Option<HarnessConnectionTasks>,
58    pub auth_updates: mpsc::UnboundedReceiver<acp_utils::notifications::AuthMethodsUpdatedParams>,
59    resume_agent: FakeAcpAgent,
60    runtime_control: Arc<Mutex<FakeRuntimeControl>>,
61    pub oauth_store: Arc<aether_auth::FakeOAuthCredentialStore>,
62    agent_cx: ConnectionTo<Client>,
63    state: Arc<AcpState>,
64    session_store: Arc<SessionStore>,
65    _tmp: tempfile::TempDir,
66}
67
68/// A real loopback listener backed by the harness's existing state and fake runtimes.
69pub struct AcpWebSocketTestServer {
70    pub address: SocketAddr,
71    detached: watch::Receiver<u64>,
72    shutdown: Option<oneshot::Sender<()>>,
73    task: JoinHandle<Result<(), ServerRunError>>,
74}
75
76impl AcpWebSocketTestServer {
77    pub fn url(&self) -> String {
78        format!("ws://{}", self.address)
79    }
80
81    pub fn subscribe(&self) -> watch::Receiver<u64> {
82        self.detached.clone()
83    }
84
85    /// Join networking cleanup without shutting down the harness-owned host.
86    pub async fn shutdown(mut self) {
87        let _ = self.shutdown.take().expect("server is running").send(());
88        (&mut self.task).await.expect("listener task joins").expect("listener shuts down");
89    }
90
91    /// Drop the running server without executing the normal async shutdown path.
92    pub async fn abort(mut self) {
93        self.task.abort();
94        assert!((&mut self.task).await.expect_err("listener task is aborted").is_cancelled());
95    }
96}
97
98impl Drop for AcpWebSocketTestServer {
99    fn drop(&mut self) {
100        self.task.abort();
101    }
102}
103
104pub struct FakeAgentSwitchingSession {
105    session_id: SessionId,
106    planner: FakeAcpAgent,
107    coder: FakeAcpAgent,
108}
109
110#[derive(Clone)]
111pub struct FakeAcpAgent {
112    name: String,
113    captured_contexts: Arc<Mutex<Vec<Context>>>,
114}
115
116pub struct FakeBackgroundTask {
117    server: FakeMcpServer,
118}
119
120pub struct PendingRuntime {
121    started: oneshot::Receiver<()>,
122    release: oneshot::Sender<bool>,
123}
124
125impl PendingRuntime {
126    pub async fn wait_until_started(&mut self) {
127        (&mut self.started).await.expect("runtime startup reached");
128    }
129
130    pub fn finish(self, succeed: bool) {
131        self.release.send(succeed).expect("runtime startup is waiting");
132    }
133}
134
135impl AcpTestHarness {
136    pub fn pause_next_runtime(&self) -> PendingRuntime {
137        let (started, observed) = oneshot::channel();
138        let (release, proceed) = oneshot::channel();
139        self.runtime_control.lock().unwrap().pending = Some((started, proceed));
140        PendingRuntime { started: observed, release }
141    }
142
143    pub fn pause_prompt_expansion(&self) -> (Arc<tokio::sync::Notify>, Arc<tokio::sync::Notify>) {
144        let gate = (Arc::new(tokio::sync::Notify::new()), Arc::new(tokio::sync::Notify::new()));
145        self.runtime_control.lock().unwrap().prompt_gate = Some(gate.clone());
146        gate
147    }
148
149    pub fn watch_agent_prompts(&self) -> watch::Receiver<usize> {
150        self.runtime_control.lock().unwrap().agent_prompts.subscribe()
151    }
152
153    pub fn elicit_during_prompt_expansion(&self) -> mpsc::UnboundedReceiver<rmcp::model::ElicitResult> {
154        let (sender, receiver) = mpsc::unbounded_channel();
155        self.runtime_control.lock().unwrap().elicitation_results = Some(sender);
156        receiver
157    }
158
159    pub fn live_runtime_count(&self) -> usize {
160        self.runtime_control.lock().unwrap().agents.iter().filter(|sender| !sender.is_closed()).count()
161    }
162
163    pub async fn disconnect(&mut self) {
164        if let Some(connection) = self.connection.take() {
165            connection.shutdown().await;
166        }
167    }
168
169    pub async fn start_unattached_session(&mut self, prompt: impl Into<String>) -> SessionId {
170        self.disconnect().await;
171        self.state.start_session(prompt.into()).await.expect("unattached session starts")
172    }
173
174    /// Connect a fresh initialized client to the same host and session store.
175    pub async fn reconnect(&mut self) {
176        assert!(self.connection.is_none(), "disconnect the current client before reconnecting");
177        assert!(!self.state.stop_token().is_cancelled(), "host is running");
178        let connection = connect_client(self.state.clone()).await;
179        self.client_cx = connection.client_cx;
180        self.agent_cx = connection.agent_cx;
181        self.peer = connection.peer;
182        self.initialize_response = connection.initialize_response;
183        self.auth_updates = connection.auth_updates;
184        self.connection = Some(connection.tasks);
185    }
186
187    pub async fn resume(&self, session_id: &SessionId) {
188        self.client_cx
189            .send_request(ResumeSessionRequest::new(session_id.clone(), AbsolutePath::new("/tmp")))
190            .block_task()
191            .await
192            .expect("session resumes");
193    }
194
195    pub async fn resume_with_replay(&self, session_id: &SessionId) {
196        self.client_cx
197            .send_request(
198                ResumeSessionRequest::new(session_id.clone(), AbsolutePath::new("/tmp"))
199                    .replay_from(ReplayFrom::Start(ReplayFromStart::new())),
200            )
201            .block_task()
202            .await
203            .expect("session resumes with replay");
204    }
205
206    /// Move this persistent host from its in-memory connection to a real port-zero listener.
207    pub async fn serve_websocket(&mut self) -> AcpWebSocketTestServer {
208        self.disconnect().await;
209        self.listen_websocket().await
210    }
211
212    /// Listen without disconnecting any existing client.
213    pub async fn listen_websocket(&self) -> AcpWebSocketTestServer {
214        assert!(!self.state.stop_token().is_cancelled(), "host is running");
215        let server =
216            AcpServer::bind("127.0.0.1:0".parse().unwrap(), self.state.clone()).await.expect("bind test server");
217        let address = server.local_addr().expect("listener address");
218        let detached = self.state.client_slot().subscribe();
219        let (shutdown, stopped) = oneshot::channel();
220        let task = tokio::spawn(server.run_until(async move {
221            let _ = stopped.await;
222            Ok(())
223        }));
224        AcpWebSocketTestServer { address, detached, shutdown: Some(shutdown), task }
225    }
226
227    /// Stop the connection and all host-owned session runtimes.
228    pub async fn shutdown(&mut self) {
229        self.disconnect().await;
230        self.state.shutdown_all().await;
231    }
232
233    pub fn stored_events(&self, session_id: &SessionId) -> Vec<SessionEvent> {
234        self.session_store.load(session_id.0.as_ref()).expect("stored session loads").1
235    }
236
237    /// Runs `body` with a fresh harness on the current-thread `LocalSet` the
238    /// host needs for its `spawn_local` tasks; tests must be annotated with
239    /// `#[tokio::test(flavor = "current_thread")]`.
240    pub async fn run<F, Fut, T>(body: F) -> T
241    where
242        F: FnOnce(Self) -> Fut,
243        Fut: Future<Output = T>,
244    {
245        LocalSet::new().run_until(Box::pin(async move { body(Self::start().await).await })).await
246    }
247
248    pub async fn run_with_hooks<F, Fut>(hooks: SessionHooks, body: F)
249    where
250        F: FnOnce(Self) -> Fut,
251        Fut: Future<Output = ()>,
252    {
253        LocalSet::new().run_until(Box::pin(async move { body(Self::start_with(hooks).await).await })).await;
254    }
255
256    pub async fn start() -> Self {
257        Self::start_with(SessionHooks::default()).await
258    }
259
260    async fn start_with(hooks: SessionHooks) -> Self {
261        let tmp = tempfile::tempdir().expect("tempdir for session store");
262        let session_store = Arc::new(SessionStore::from_path(tmp.path().to_path_buf()));
263        let workspace_manager = Arc::new(WorkspaceManager::from_registry_path_with_cloner(
264            tmp.path().join("workspaces.json"),
265            Arc::new(StdCopyCloner),
266        ));
267        let (resume_def, resume_agent) = fake_agent("Resume", "resume-mcp", "resume", "resumed reply");
268        let mut resume_agents = HashMap::new();
269        resume_agents.insert(resume_def.spec.name.clone(), resume_def);
270        let runtime_control = Arc::new(Mutex::new(FakeRuntimeControl::default()));
271        let runtime_factory = Arc::new(FakeRuntimeFactory {
272            cwd: PathBuf::from("/tmp"),
273            agents: resume_agents,
274            control: runtime_control.clone(),
275        });
276        let oauth_store = Arc::new(aether_auth::FakeOAuthCredentialStore::new());
277        let state = Arc::new(AcpState::with_login(
278            AcpStateConfig {
279                session_store: session_store.clone(),
280                workspace_manager,
281                oauth_credential_store: oauth_store.clone(),
282                initial_selection: InitialSessionSelection::default(),
283                settings_source: SettingsSourceArgs::default(),
284                provider_connections: ProviderConnectionOverrides::default(),
285                telemetry: None,
286                runtime_factory: Some(runtime_factory),
287                cwd: PathBuf::from("/tmp"),
288                hooks,
289            },
290            Arc::new(FakeProviderLogin),
291        ));
292
293        let connection = connect_client(state.clone()).await;
294        Self {
295            client_cx: connection.client_cx,
296            peer: connection.peer,
297            initialize_response: connection.initialize_response,
298            connection: Some(connection.tasks),
299            auth_updates: connection.auth_updates,
300            resume_agent,
301            runtime_control,
302            oauth_store,
303            agent_cx: connection.agent_cx,
304            state,
305            session_store,
306            _tmp: tmp,
307        }
308    }
309
310    pub fn resume_agent(&self) -> &FakeAcpAgent {
311        &self.resume_agent
312    }
313
314    pub async fn insert_agent_switching_session(&self) -> FakeAgentSwitchingSession {
315        self.insert_switching_session(
316            SessionId::new("agent-switching-session"),
317            Vec::new(),
318            Some("Planner".to_string()),
319            false,
320        )
321        .await
322    }
323
324    pub async fn insert_agent_switching_session_with_serverless_coder(&self) -> FakeAgentSwitchingSession {
325        self.insert_switching_session(
326            SessionId::new("agent-switching-serverless-session"),
327            Vec::new(),
328            Some("Planner".to_string()),
329            true,
330        )
331        .await
332    }
333
334    pub async fn insert_loaded_agent_switching_session(&self, session_id: &str) -> FakeAgentSwitchingSession {
335        let events = self.session_store.load(session_id).map(|(_, events)| events).unwrap_or_default();
336        let selected_mode = last_agent_from_events(Some("Planner".to_string()), &events);
337        self.insert_switching_session(SessionId::new(session_id), events, selected_mode, false).await
338    }
339
340    pub async fn expect_idle(&mut self, session_id: &SessionId, expected: StopReason) {
341        loop {
342            let notification = self.peer.next_session_notification().await;
343            if notification.session_id == *session_id
344                && let SessionUpdate::StateUpdate(StateUpdate::Idle(idle)) = notification.update
345            {
346                assert_eq!(idle.stop_reason, Some(expected));
347                return;
348            }
349        }
350    }
351
352    pub async fn updates_until(
353        &mut self,
354        session_id: &SessionId,
355        last: impl Fn(&SessionUpdate) -> bool,
356    ) -> Vec<SessionUpdate> {
357        let mut updates = Vec::new();
358        loop {
359            let notification = self.peer.next_session_notification().await;
360            if notification.session_id != *session_id {
361                continue;
362            }
363            let done = last(&notification.update);
364            updates.push(notification.update);
365            if done {
366                return updates;
367            }
368        }
369    }
370
371    pub async fn expect_mcp_server_status(&mut self, expected: &[&str]) {
372        assert_server_status(self.peer.next_mcp_notification().await, expected);
373    }
374
375    pub async fn expect_mcp_server_status_exact(&mut self, expected: &[&str]) {
376        assert_server_status_exact(self.peer.next_mcp_notification().await, expected);
377    }
378
379    pub async fn expect_available_commands(&mut self, expected: &[&str], unexpected: &[&str]) {
380        loop {
381            let update = self.peer.next_session_notification().await.update;
382            if matches!(update, SessionUpdate::AvailableCommandsUpdate(_)) {
383                assert_available_commands(update, expected, unexpected);
384                return;
385            }
386        }
387    }
388
389    pub fn append_agent_switch(&self, session_id: &str, from: Option<&str>, to: Option<&str>) {
390        self.append_stored_event(
391            session_id,
392            &SessionEvent::Control(SessionControlEvent::AgentSwitched {
393                from: from.map(str::to_string),
394                to: to.map(str::to_string),
395            }),
396        );
397    }
398
399    /// Register a stub session built from a hand-spawned
400    /// `(agent_tx, agent_rx, agent_handle)` triple — typically from
401    /// `aether_core::core::agent(fake_llm).spawn().await`. Pairs the agent with a
402    /// real but empty in-memory MCP (no servers). The session is routable via
403    /// `state.route_prompt(id)` / `state.cancel(id)`.
404    pub async fn insert_stub_session(
405        &self,
406        agent_tx: mpsc::Sender<Command>,
407        agent_rx: mpsc::Receiver<AgentEvent>,
408        agent_handle: AgentHandle,
409        id: SessionId,
410        model: &str,
411    ) {
412        let mut mcp = mcp("/tmp").spawn().await.expect("stub MCP spawns");
413        mcp.block_until_ready().await.expect("stub MCP is ready");
414        self.register_stub_session(StubAgentParts { tx: agent_tx, rx: agent_rx, handle: agent_handle, mcp }, id, model)
415            .await;
416    }
417
418    pub async fn insert_background_task_session(
419        &self,
420        provider: FakeLlmProvider,
421        task: &FakeBackgroundTask,
422        id: SessionId,
423    ) {
424        let server = task.server.clone();
425        let gate = self.runtime_control.lock().unwrap().prompt_gate.clone();
426        let mcp = with_in_memory_server(mcp("/tmp"), "tasks", move || server.clone().into_dyn());
427        let mcp = with_in_memory_server(mcp, "prompts", move || {
428            FakePromptMcp::new("plan").with_gate(gate.clone()).into_dyn()
429        });
430        let mut mcp = mcp.spawn().await.expect("background task MCP spawns");
431        mcp.block_until_ready().await.expect("background task MCP is ready");
432        let (tx, rx, handle) = AgentBuilder::new(Arc::new(provider))
433            .max_auto_continues(0)
434            .tools(mcp.handle().clone(), Vec::new())
435            .spawn()
436            .await
437            .expect("background task agent spawns");
438        self.register_stub_session(StubAgentParts { tx, rx, handle, mcp }, id, "fake:fake").await;
439    }
440
441    pub fn append_stored_session(&self, session_id: &str, created_at: &str) {
442        self.append_stored_session_in(session_id, created_at, std::path::Path::new("/tmp"));
443    }
444
445    pub fn append_stored_session_in(&self, session_id: &str, created_at: &str, cwd: &std::path::Path) {
446        let meta = SessionMeta {
447            session_id: session_id.to_string(),
448            cwd: cwd.to_path_buf(),
449            model: "anthropic:claude-sonnet-4-5".to_string(),
450            selected_mode: None,
451            created_at: created_at.to_string(),
452        };
453
454        self.session_store.append_meta(session_id, &meta).expect("stored session meta appends");
455    }
456
457    pub fn append_stored_prompt(&self, session_id: &str, prompt: &str) {
458        self.append_stored_event(
459            session_id,
460            &SessionEvent::User(UserEvent::Message {
461                message_id: llm::MessageId::new(),
462                content: vec![llm::ContentBlock::text(prompt)],
463                display_content: None,
464            }),
465        );
466    }
467
468    pub fn append_stored_user_blocks(&self, session_id: &str, blocks: Vec<llm::ContentBlock>) {
469        self.append_stored_event(
470            session_id,
471            &SessionEvent::User(UserEvent::Message {
472                message_id: llm::MessageId::new(),
473                content: blocks,
474                display_content: None,
475            }),
476        );
477    }
478
479    pub fn append_stored_agent_turn(&self, session_id: &str, text: &str) {
480        self.append_stored_agent_text(session_id, text);
481        self.append_stored_event(
482            session_id,
483            &SessionEvent::Agent(AgentEvent::Turn(TurnEvent::Ended { outcome: TurnOutcome::Completed })),
484        );
485    }
486
487    pub fn append_stored_agent_text(&self, session_id: &str, text: &str) {
488        self.append_stored_event(
489            session_id,
490            &SessionEvent::Agent(AgentEvent::Message(MessageEvent::Text {
491                message_id: MessageId::new(),
492                chunk: text.to_string(),
493                is_complete: true,
494            })),
495        );
496    }
497
498    pub fn append_stored_tool_call(&self, session_id: &str, id: &str, name: &str, arguments: &str) {
499        self.append_stored_event(
500            session_id,
501            &SessionEvent::Agent(AgentEvent::Tool(ToolEvent::Call {
502                request: ToolCallRequest {
503                    id: id.to_string(),
504                    name: name.to_string(),
505                    arguments: arguments.to_string(),
506                },
507            })),
508        );
509    }
510
511    pub fn append_stored_tool_result(&self, session_id: &str, id: &str, name: &str, result: &str) {
512        self.append_stored_event(
513            session_id,
514            &SessionEvent::Agent(AgentEvent::Tool(ToolEvent::Result {
515                result: ToolCallResult {
516                    id: id.to_string(),
517                    name: name.to_string(),
518                    arguments: "{}".to_string(),
519                    result: result.to_string(),
520                },
521                result_meta: None,
522            })),
523        );
524    }
525
526    async fn insert_switching_session(
527        &self,
528        acp_session_id: SessionId,
529        events: Vec<SessionEvent>,
530        selected_mode: Option<String>,
531        serverless_coder: bool,
532    ) -> FakeAgentSwitchingSession {
533        let (planner_def, planner) = fake_agent("Planner", "planner-mcp", "plan", PLANNER_REPLY);
534        let (mut coder_def, coder) = fake_agent("Coder", "coder-mcp", "edit", CODER_REPLY);
535        if serverless_coder {
536            coder_def.mcp = None;
537        }
538
539        let mut catalog_specs = Vec::new();
540        let mut agents = HashMap::new();
541        for def in [planner_def, coder_def] {
542            catalog_specs.push(def.spec.clone());
543            agents.insert(def.spec.name.clone(), def);
544        }
545        let specs = SessionAgents::new(AgentCatalog::new(PathBuf::from("/tmp"), catalog_specs, None));
546
547        let factory =
548            Arc::new(FakeRuntimeFactory { cwd: PathBuf::from("/tmp"), agents, control: self.runtime_control.clone() });
549        let initial_agent = selected_mode.clone().unwrap_or_else(|| "Planner".to_string());
550
551        self.state
552            .register_session(SessionActorInit {
553                cwd: PathBuf::from("/tmp"),
554                mcp_servers: Vec::new(),
555                session_id: acp_session_id.clone(),
556                connection: Some(self.agent_cx.clone()),
557                repository: self.session_store.clone(),
558                oauth_credential_store: self.oauth_store.clone(),
559                active_agent: AgentKey::Named(initial_agent),
560                specs,
561                runtime_factory: factory,
562                transcript: events,
563                replay: false,
564                modes: switching_modes(),
565                config: SessionConfigState::with_selection(
566                    "anthropic:claude-sonnet-4-5".to_string(),
567                    selected_mode,
568                    None,
569                ),
570                hooks: SessionHooks::default(),
571            })
572            .await;
573        FakeAgentSwitchingSession { session_id: acp_session_id, planner, coder }
574    }
575
576    pub fn append_stored_event(&self, session_id: &str, event: &SessionEvent) {
577        self.session_store.append_event(session_id, event).expect("stored session event appends");
578    }
579
580    async fn register_stub_session(&self, mut parts: StubAgentParts, id: SessionId, model: &str) {
581        let model_spec: llm::catalog::LlmModel = "anthropic:claude-sonnet-4-5".parse().expect("test model parses");
582        let mut specs = SessionAgents::new(AgentCatalog::empty(PathBuf::from("/tmp")));
583        specs.set_default(AgentSpec::bare(&model_spec, None, Vec::new()));
584        let agent_prompts = self.runtime_control.lock().unwrap().agent_prompts.clone();
585        self.runtime_control.lock().unwrap().agents.push(parts.tx.clone());
586        parts.tx = count_prompts(parts.tx, agent_prompts);
587        let factory = Arc::new(StubRuntimeFactory { agent_parts: Mutex::new(Some(parts)) });
588
589        self.state
590            .register_session(SessionActorInit {
591                cwd: PathBuf::from("/tmp"),
592                mcp_servers: Vec::new(),
593                session_id: id.clone(),
594                connection: Some(self.agent_cx.clone()),
595                repository: self.session_store.clone(),
596                oauth_credential_store: self.oauth_store.clone(),
597                active_agent: AgentKey::Default,
598                specs,
599                runtime_factory: factory,
600                transcript: Vec::new(),
601                replay: false,
602                modes: Modes::default(),
603                config: SessionConfigState::with_selection(model.to_string(), None, None),
604                hooks: SessionHooks::default(),
605            })
606            .await;
607    }
608}
609
610impl FakeAgentSwitchingSession {
611    pub fn session_id(&self) -> &SessionId {
612        &self.session_id
613    }
614
615    pub fn planner(&self) -> &FakeAcpAgent {
616        &self.planner
617    }
618
619    pub fn coder(&self) -> &FakeAcpAgent {
620        &self.coder
621    }
622}
623
624impl FakeAcpAgent {
625    /// Asserts the agent's most recent turn saw a conversation containing each
626    /// of `expected` (user or assistant text), in addition to anything else.
627    pub fn assert_saw(&self, expected: &[&str]) {
628        let seen = self.latest_conversation();
629        for text in expected {
630            assert!(seen.iter().any(|m| m == text), "{} should have seen {text:?}; saw {seen:?}", self.name);
631        }
632    }
633
634    /// Asserts the agent's most recent turn saw *exactly* `expected` and nothing
635    /// else — used to prove a freshly-activated agent started with no prior
636    /// transcript.
637    pub fn assert_saw_exactly(&self, expected: &[&str]) {
638        let seen = self.latest_conversation();
639        let expected: Vec<String> = expected.iter().map(|t| (*t).to_string()).collect();
640        assert_eq!(seen, expected, "{} conversation mismatch", self.name);
641    }
642
643    pub fn assert_saw_user_content(&self, expected: &[llm::ContentBlock]) {
644        let contexts = self.captured_contexts.lock().expect("captured contexts lock is healthy");
645        let latest = contexts.last().expect("agent should have run a turn");
646        assert!(
647            latest
648                .messages()
649                .iter()
650                .any(|message| { matches!(message, ChatMessage::User { content, .. } if content == expected) }),
651            "{} should have seen user content {expected:?}",
652            self.name
653        );
654    }
655
656    /// Asserts the agent never ran a turn (its LLM was never invoked).
657    pub fn assert_never_ran(&self) {
658        let contexts = self.captured_contexts.lock().expect("captured contexts lock is healthy");
659        assert!(contexts.is_empty(), "{} should not have run; captured {} context(s)", self.name, contexts.len());
660    }
661
662    fn latest_conversation(&self) -> Vec<String> {
663        let contexts = self.captured_contexts.lock().expect("captured contexts lock is healthy");
664        let latest = contexts.last().unwrap_or_else(|| panic!("{} should have run a turn", self.name));
665        conversation_texts(latest)
666    }
667}
668
669impl FakeBackgroundTask {
670    pub const TOOL: &str = "tasks__start";
671    const ID: &str = "background-task";
672
673    pub fn new() -> Self {
674        let working = background_task(TaskStatus::Working);
675        let server = FakeMcpServer::new()
676            .with_tool(FakeTool::new("start").responds(FakeToolResponse::task(CreateTaskResult::new(working.clone()))))
677            .with_task(Self::ID, [DetailedTask::new(working, TaskPayload::Working)]);
678        Self { server }
679    }
680
681    pub fn complete(&self, output: &str) {
682        let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text(output)]);
683        let completed = DetailedTask::new(background_task(TaskStatus::Completed), completed_task_payload(result));
684        self.server.state().script_task(Self::ID, [completed]);
685    }
686}
687
688impl Default for FakeBackgroundTask {
689    fn default() -> Self {
690        Self::new()
691    }
692}
693
694struct HarnessConnection {
695    client_cx: ConnectionTo<Agent>,
696    agent_cx: ConnectionTo<Client>,
697    peer: TestPeer,
698    initialize_response: InitializeResponse,
699    auth_updates: mpsc::UnboundedReceiver<acp_utils::notifications::AuthMethodsUpdatedParams>,
700    tasks: HarnessConnectionTasks,
701}
702
703struct HarnessConnectionTasks {
704    stop: CancellationToken,
705    tasks: JoinSet<Result<(), agent_client_protocol::Error>>,
706}
707
708impl HarnessConnectionTasks {
709    async fn shutdown(mut self) {
710        self.stop.cancel();
711        while let Some(result) = self.tasks.join_next().await {
712            result.expect("connection task joins").expect("connection cleanup completes");
713        }
714    }
715}
716
717async fn connect_client(state: Arc<AcpState>) -> HarnessConnection {
718    let (peer, client_builder) = TestPeer::new();
719    let (auth_tx, auth_updates) = mpsc::unbounded_channel();
720    let client_builder = client_builder.on_receive_notification(
721        async move |notification: AuthMethodsUpdatedParams, _cx| {
722            let _ = auth_tx.send(notification);
723            Ok(())
724        },
725        on_receive_notification!(),
726    );
727    let (agent_transport, client_transport) = Channel::duplex();
728    let (send_agent_connection, agent_ready) = oneshot::channel();
729    let (send_client_connection, client_ready) = oneshot::channel();
730    let stop = state.stop_token().child_token();
731    let mut tasks = JoinSet::new();
732    let agent_stop = stop.clone();
733    tasks.spawn_local(async move { state.serve_with(agent_transport, agent_stop, send_agent_connection).await });
734    tasks.spawn_local(async move {
735        client_builder
736            .with_runner(acp_utils::testing::CaptureConnection(send_client_connection))
737            .connect_to(client_transport)
738            .await
739    });
740    let agent_cx = agent_ready.await.expect("agent connection");
741    let client_cx = client_ready.await.expect("client connection");
742    let initialize_response =
743        client_cx.send_request(initialize_request()).block_task().await.expect("initialize harness");
744    HarnessConnection {
745        client_cx,
746        agent_cx,
747        peer,
748        initialize_response,
749        auth_updates,
750        tasks: HarnessConnectionTasks { stop, tasks },
751    }
752}
753
754/// Spawns each agent's runtime through the real [`AgentRuntime`] wiring, but
755/// backed by a [`FakeLlmProvider`] and an in-memory MCP server instead of a
756/// network LLM and external MCP processes.
757struct FakeRuntimeFactory {
758    cwd: PathBuf,
759    agents: HashMap<String, FakeAgentDef>,
760    control: Arc<Mutex<FakeRuntimeControl>>,
761}
762
763#[derive(Default)]
764struct FakeRuntimeControl {
765    agents: Vec<mpsc::Sender<Command>>,
766    pending: Option<(oneshot::Sender<()>, oneshot::Receiver<bool>)>,
767    prompt_gate: Option<(Arc<tokio::sync::Notify>, Arc<tokio::sync::Notify>)>,
768    elicitation_results: Option<mpsc::UnboundedSender<rmcp::model::ElicitResult>>,
769    agent_prompts: watch::Sender<usize>,
770}
771
772struct FakeAgentDef {
773    spec: AgentSpec,
774    provider: Arc<dyn StreamingModelProvider>,
775    mcp: Option<(String, String)>,
776}
777
778#[async_trait::async_trait]
779impl RuntimeFactory for FakeRuntimeFactory {
780    async fn spawn(
781        &self,
782        _agent: AgentKey,
783        spec: &AgentSpec,
784        initial_messages: Vec<ChatMessage>,
785        usage_seed: Option<SessionUsageEvent>,
786    ) -> Result<AgentRuntime, SessionError> {
787        let pending = self.control.lock().unwrap().pending.take();
788        if let Some((started, proceed)) = pending {
789            let _ = started.send(());
790            if !proceed.await.unwrap_or(false) {
791                return Err(SessionError::AgentNotFound("injected startup failure".to_string()));
792            }
793        }
794        let def = self
795            .agents
796            .get(&spec.name)
797            .or_else(|| self.agents.values().next())
798            .ok_or_else(|| SessionError::AgentNotFound(spec.name.clone()))?;
799        let provider = def.provider.clone();
800
801        let mut mcp_builder = mcp(&self.cwd).with_tool_filter(spec.tools.clone());
802        if let Some((server_name, prompt_name)) = &def.mcp {
803            let prompt_name = prompt_name.clone();
804            let gate = self.control.lock().unwrap().prompt_gate.clone();
805            let elicitation_results = self.control.lock().unwrap().elicitation_results.clone();
806            mcp_builder = with_in_memory_server(mcp_builder, server_name, move || {
807                FakePromptMcp::new(&prompt_name)
808                    .with_gate(gate.clone())
809                    .with_elicitation(elicitation_results.clone())
810                    .into_dyn()
811            });
812        }
813        let mut spawn =
814            mcp_builder.spawn().await.map_err(|e| SessionError::Build(CliError::McpError(e.to_string())))?;
815        spawn.block_until_ready().await.ok_or(SessionError::McpStartupStopped)?;
816        let mcp_handle = spawn.handle().clone();
817        let mut builder = AgentBuilder::new(provider).max_auto_continues(0);
818        if let Some(last) = &usage_seed {
819            builder = builder.resume_usage(last);
820        }
821        for prompt in &spec.prompts {
822            builder = builder.system_prompt(prompt.clone());
823        }
824        let (agent_tx, agent_rx, agent_handle) = builder
825            .tools(mcp_handle, Vec::new())
826            .messages(initial_messages)
827            .spawn()
828            .await
829            .map_err(|e| SessionError::Build(CliError::AgentError(e.to_string())))?;
830        self.control.lock().unwrap().agents.push(agent_tx.clone());
831        let (mcp_runtime, event_rx) = spawn.connect_agent(agent_tx.clone()).await.split();
832
833        Ok(AgentRuntime::new(agent_tx, agent_rx, Some(agent_handle), event_rx, mcp_runtime))
834    }
835}
836
837struct StubRuntimeFactory {
838    agent_parts: Mutex<Option<StubAgentParts>>,
839}
840
841struct StubAgentParts {
842    tx: mpsc::Sender<Command>,
843    rx: mpsc::Receiver<AgentEvent>,
844    handle: AgentHandle,
845    mcp: McpSession,
846}
847
848#[async_trait::async_trait]
849impl RuntimeFactory for StubRuntimeFactory {
850    async fn spawn(
851        &self,
852        _agent: AgentKey,
853        _spec: &AgentSpec,
854        _initial_messages: Vec<ChatMessage>,
855        _usage_seed: Option<SessionUsageEvent>,
856    ) -> Result<AgentRuntime, SessionError> {
857        let parts = self
858            .agent_parts
859            .lock()
860            .expect("stub agent parts lock is healthy")
861            .take()
862            .expect("stub runtime spawned more than once");
863        let (mcp_runtime, event_rx) = parts.mcp.connect_agent(parts.tx.clone()).await.split();
864
865        Ok(AgentRuntime::new(parts.tx, parts.rx, Some(parts.handle), event_rx, mcp_runtime))
866    }
867}
868
869fn with_in_memory_server(
870    builder: McpBuilder,
871    name: &str,
872    server: impl Fn() -> Box<dyn DynService<RoleServer>> + Send + Sync + 'static,
873) -> McpBuilder {
874    let factory: ServerFactory = Box::new(move |_spec, _services| std::future::ready(server()).boxed());
875    let transport = McpTransport::InMemory {
876        spec: InMemoryServerSpec { factory: name.to_string(), args: Vec::new(), input: None },
877    };
878    builder.register_in_memory_server(name, factory).with_servers(vec![McpServer::new(
879        name,
880        transport,
881        ToolExposure::ModelVisible,
882    )])
883}
884
885fn count_prompts(agent_tx: mpsc::Sender<Command>, prompts: watch::Sender<usize>) -> mpsc::Sender<Command> {
886    let (tx, mut rx) = mpsc::channel(32);
887    tokio::spawn(async move {
888        while let Some(command) = rx.recv().await {
889            let is_prompt = matches!(command, Command::UserCommand(UserCommand::Text { .. }));
890            if agent_tx.send(command).await.is_err() {
891                return;
892            }
893            if is_prompt {
894                prompts.send_modify(|count| *count += 1);
895            }
896        }
897    });
898    tx
899}
900
901fn background_task(status: TaskStatus) -> Task {
902    const CREATED_AT: &str = "2026-01-01T00:00:00Z";
903    Task::new(FakeBackgroundTask::ID, status, CREATED_AT.to_string(), CREATED_AT.to_string()).with_poll_interval_ms(10)
904}
905
906fn fake_agent(name: &str, server_name: &str, prompt_name: &str, reply: &str) -> (FakeAgentDef, FakeAcpAgent) {
907    const TURNS_BEFORE_AND_AFTER_REATTACH: usize = 2;
908    let provider = FakeLlmProvider::new(vec![
909        vec![LlmResponse::Start, LlmResponse::text(reply), LlmResponse::done()];
910        TURNS_BEFORE_AND_AFTER_REATTACH
911    ])
912    .with_display_name(name);
913    let captured_contexts = provider.captured_contexts();
914    let def = FakeAgentDef {
915        spec: fake_agent_spec(name),
916        provider: Arc::new(provider),
917        mcp: Some((server_name.to_string(), prompt_name.to_string())),
918    };
919    let observer = FakeAcpAgent { name: name.to_string(), captured_contexts };
920    (def, observer)
921}
922
923struct FakeProviderLogin;
924
925#[async_trait::async_trait]
926impl super::state::ProviderLogin for FakeProviderLogin {
927    async fn login(&self, store: &dyn OAuthCredentialStorage) -> Result<(), llm::LlmError> {
928        store
929            .save("codex", serde_json::json!({"access_token": "fake-access", "refresh_token": "fake-refresh"}))
930            .await?;
931        Ok(())
932    }
933}
934
935fn switching_modes() -> Modes {
936    Modes::new(vec![
937        ValidatedMode {
938            name: "Planner".to_string(),
939            model: "anthropic:claude-sonnet-4-5".to_string(),
940            reasoning_effort: None,
941        },
942        ValidatedMode { name: "Coder".to_string(), model: text_only_model(), reasoning_effort: None },
943    ])
944}
945
946fn text_only_model() -> String {
947    LlmModel::all()
948        .iter()
949        .find(|model| !model.supports_image())
950        .expect("catalog contains a text-only model")
951        .to_string()
952}
953
954fn fake_agent_spec(name: &str) -> AgentSpec {
955    let model: llm::catalog::LlmModel = "anthropic:claude-sonnet-4-5".parse().expect("test model parses");
956    let mut spec = AgentSpec::bare(&model, None, vec![Prompt::text(&format!("{name} system prompt"))]);
957    spec.name = name.to_string();
958    spec.description = format!("{name} test agent");
959    spec.exposure = AgentSpecExposure::user_only();
960    spec
961}
962
963fn assert_available_commands(update: SessionUpdate, expected: &[&str], unexpected: &[&str]) {
964    let SessionUpdate::AvailableCommandsUpdate(commands) = update else {
965        panic!("expected available commands update");
966    };
967    let names = commands.available_commands.iter().map(|command| command.name.as_str()).collect::<Vec<_>>();
968    for name in expected {
969        assert!(names.contains(name), "expected command /{name} in {names:?}");
970    }
971    for name in unexpected {
972        assert!(!names.contains(name), "did not expect command /{name} in {names:?}");
973    }
974}
975
976fn assert_server_status(notification: McpNotification, expected: &[&str]) {
977    let McpNotification::ServerStatus { servers } = notification;
978    let names = servers.iter().map(|server| server.name.as_str()).collect::<Vec<_>>();
979    for server_name in expected {
980        assert!(names.contains(server_name), "expected server {server_name} in {names:?}");
981    }
982}
983
984fn assert_server_status_exact(notification: McpNotification, expected: &[&str]) {
985    let McpNotification::ServerStatus { servers } = notification;
986    let names = servers.iter().map(|server| server.name.as_str()).collect::<Vec<_>>();
987    assert_eq!(names, expected);
988}
989
990fn conversation_texts(context: &Context) -> Vec<String> {
991    context
992        .messages()
993        .iter()
994        .filter_map(|message| match message {
995            ChatMessage::User { content, .. } => llm::ContentBlock::first_text(content).map(str::to_string),
996            ChatMessage::Assistant { content, .. } if !content.is_empty() => Some(content.clone()),
997            _ => None,
998        })
999        .collect()
1000}
aether_cli/acp/testing.rs

aether_cli/acp/
testing.rs
