aether_cli/acp/
testing.rs

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

aether_cli/acp/
testing.rs
