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

aether_cli/acp/
testing.rs
