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

aether_cli/acp/
testing.rs
