aether_cli/acp/
testing.rs

1use super::agent::acp_agent_builder;
2use super::fake_prompt_mcp::FakePromptMcp;
3use super::session::actor::{SessionActor, 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, RuntimeEvent, 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::McpNotification;
17use acp_utils::testing::{TestPeer, duplex_pair};
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::v1::{SessionId, SessionUpdate};
27use agent_client_protocol::{Agent, Client, ConnectionTo};
28use futures::FutureExt;
29use llm::ProviderConnectionOverrides;
30use llm::testing::FakeLlmProvider;
31use llm::{ChatMessage, Context, LlmResponse, StreamingModelProvider};
32use mcp_utils::client::{InMemoryServerSpec, McpServer, McpTransport, ToolExposure};
33use std::collections::HashMap;
34use std::path::PathBuf;
35use std::sync::{Arc, Mutex};
36use tokio::sync::{mpsc, oneshot};
37use tokio::task::spawn_local;
38
39const PLANNER_REPLY: &str = "planner reply";
40const CODER_REPLY: &str = "coder reply";
41
42/// In-memory ACP harness running the real `acp_agent_builder` against a
43/// pre-wired test client. Created via [`AcpTestHarness::start`] inside a
44/// `LocalSet`. The harness owns an [`AcpState`] and a temp-dir-backed
45/// [`SessionStore`] so tests can register fake-driven sessions without
46/// going through `new_session`.
47pub struct AcpTestHarness {
48    pub client_cx: ConnectionTo<Agent>,
49    pub peer: TestPeer,
50    resume_agent: FakeAcpAgent,
51    agent_cx: ConnectionTo<Client>,
52    state: Arc<AcpState>,
53    session_store: Arc<SessionStore>,
54    _tmp: tempfile::TempDir,
55}
56
57pub struct FakeAgentSwitchingSession {
58    session_id: SessionId,
59    planner: FakeAcpAgent,
60    coder: FakeAcpAgent,
61}
62
63#[derive(Clone)]
64pub struct FakeAcpAgent {
65    name: String,
66    captured_contexts: Arc<Mutex<Vec<Context>>>,
67}
68
69impl AcpTestHarness {
70    pub async fn start() -> Self {
71        let tmp = tempfile::tempdir().expect("tempdir for session store");
72        let session_store = Arc::new(SessionStore::from_path(tmp.path().to_path_buf()));
73        let workspace_manager = Arc::new(WorkspaceManager::from_registry_path_with_cloner(
74            tmp.path().join("workspaces.json"),
75            Arc::new(StdCopyCloner),
76        ));
77        let (resume_def, resume_agent) = fake_agent("Resume", "resume-mcp", "resume", "resumed reply");
78        let mut resume_agents = HashMap::new();
79        resume_agents.insert(resume_def.spec.name.clone(), resume_def);
80        let runtime_factory = Arc::new(FakeRuntimeFactory { cwd: PathBuf::from("/tmp"), agents: resume_agents });
81        let state = Arc::new(AcpState::new(AcpStateConfig {
82            session_store: session_store.clone(),
83            workspace_manager,
84            oauth_credential_store: fake_oauth_store(),
85            initial_selection: InitialSessionSelection::default(),
86            settings_source: SettingsSourceArgs::default(),
87            provider_connections: ProviderConnectionOverrides::default(),
88            telemetry: None,
89            runtime_factory: Some(runtime_factory),
90        }));
91
92        let (peer, client_builder) = TestPeer::new();
93        let (agent_transport, client_transport) = duplex_pair();
94        let (agent_cx_tx, agent_cx_rx) = oneshot::channel::<ConnectionTo<Client>>();
95        let (client_cx_tx, client_cx_rx) = oneshot::channel::<ConnectionTo<Agent>>();
96
97        let server_state = state.clone();
98        spawn_local(async move {
99            let _ = acp_agent_builder(server_state)
100                .connect_with(agent_transport, async move |cx: ConnectionTo<Client>| {
101                    let _ = agent_cx_tx.send(cx);
102                    std::future::pending::<()>().await;
103                    Ok(())
104                })
105                .await;
106        });
107
108        spawn_local(async move {
109            let _ = client_builder
110                .connect_with(client_transport, async move |cx: ConnectionTo<Agent>| {
111                    let _ = client_cx_tx.send(cx);
112                    std::future::pending::<()>().await;
113                    Ok(())
114                })
115                .await;
116        });
117
118        let agent_cx = agent_cx_rx.await.expect("agent side connect_with produced a ConnectionTo");
119        let client_cx = client_cx_rx.await.expect("client side connect_with produced a ConnectionTo");
120        Self { client_cx, peer, resume_agent, agent_cx, state, session_store, _tmp: tmp }
121    }
122
123    pub fn resume_agent(&self) -> &FakeAcpAgent {
124        &self.resume_agent
125    }
126
127    pub async fn insert_agent_switching_session(&self) -> FakeAgentSwitchingSession {
128        self.insert_switching_session(
129            SessionId::new("agent-switching-session"),
130            Vec::new(),
131            Some("Planner".to_string()),
132            false,
133        )
134        .await
135    }
136
137    pub async fn insert_agent_switching_session_with_serverless_coder(&self) -> FakeAgentSwitchingSession {
138        self.insert_switching_session(
139            SessionId::new("agent-switching-serverless-session"),
140            Vec::new(),
141            Some("Planner".to_string()),
142            true,
143        )
144        .await
145    }
146
147    pub async fn insert_loaded_agent_switching_session(&self, session_id: &str) -> FakeAgentSwitchingSession {
148        let events = self.session_store.load(session_id).map(|(_, events)| events).unwrap_or_default();
149        let selected_mode = last_agent_from_events(Some("Planner".to_string()), &events);
150        self.insert_switching_session(SessionId::new(session_id), events, selected_mode, false).await
151    }
152
153    pub async fn expect_mcp_server_status(&mut self, expected: &[&str]) {
154        assert_server_status(self.peer.next_mcp_notification().await, expected);
155    }
156
157    pub async fn expect_mcp_server_status_exact(&mut self, expected: &[&str]) {
158        assert_server_status_exact(self.peer.next_mcp_notification().await, expected);
159    }
160
161    pub async fn expect_available_commands(&mut self, expected: &[&str], unexpected: &[&str]) {
162        loop {
163            let update = self.peer.next_session_notification().await.update;
164            if matches!(update, SessionUpdate::AvailableCommandsUpdate(_)) {
165                assert_available_commands(update, expected, unexpected);
166                return;
167            }
168        }
169    }
170
171    pub fn append_agent_switch(&self, session_id: &str, from: Option<&str>, to: Option<&str>) {
172        self.append_stored_event(
173            session_id,
174            &SessionEvent::Control(SessionControlEvent::AgentSwitched {
175                from: from.map(str::to_string),
176                to: to.map(str::to_string),
177            }),
178        );
179    }
180
181    /// Register a stub session built from a hand-spawned
182    /// `(agent_tx, agent_rx, agent_handle)` triple — typically from
183    /// `aether_core::core::agent(fake_llm).spawn().await`. Pairs the agent with a
184    /// real but empty in-memory MCP (no servers). The session is routable via
185    /// `state.route_prompt(id)` / `state.cancel(id)`.
186    pub async fn insert_stub_session(
187        &self,
188        agent_tx: mpsc::Sender<Command>,
189        agent_rx: mpsc::Receiver<AgentEvent>,
190        agent_handle: AgentHandle,
191        id: SessionId,
192        model: &str,
193    ) {
194        let model_spec: llm::catalog::LlmModel = "anthropic:claude-sonnet-4-5".parse().expect("test model parses");
195        let mut specs = SessionAgents::new(AgentCatalog::empty(PathBuf::from("/tmp")));
196        specs.set_default(AgentSpec::bare(&model_spec, None, Vec::new()));
197        let factory = Arc::new(StubRuntimeFactory {
198            cwd: PathBuf::from("/tmp"),
199            agent_parts: Mutex::new(Some(StubAgentParts { tx: agent_tx, rx: agent_rx, handle: agent_handle })),
200        });
201
202        let handle = SessionActor::spawn(SessionActorInit {
203            session_id: id.clone(),
204            connection: self.agent_cx.clone(),
205            repository: self.session_store.clone(),
206            oauth_credential_store: fake_oauth_store(),
207            active_agent: AgentKey::Default,
208            specs,
209            runtime_factory: factory,
210            transcript: Vec::new(),
211            modes: Modes::default(),
212            config: SessionConfigState::with_selection(model.to_string(), None, None),
213        })
214        .await
215        .expect("stub session actor spawns");
216        self.state.register_session(&id, handle).await;
217    }
218
219    pub fn append_stored_session(&self, session_id: &str, created_at: &str) {
220        self.append_stored_session_in(session_id, created_at, std::path::Path::new("/tmp"));
221    }
222
223    pub fn append_stored_session_in(&self, session_id: &str, created_at: &str, cwd: &std::path::Path) {
224        let meta = SessionMeta {
225            session_id: session_id.to_string(),
226            cwd: cwd.to_path_buf(),
227            model: "anthropic:claude-sonnet-4-5".to_string(),
228            selected_mode: None,
229            created_at: created_at.to_string(),
230        };
231
232        self.session_store.append_meta(session_id, &meta).expect("stored session meta appends");
233    }
234
235    pub fn append_stored_prompt(&self, session_id: &str, prompt: &str) {
236        self.append_stored_event(
237            session_id,
238            &SessionEvent::User(UserEvent::Message { content: vec![llm::ContentBlock::text(prompt)] }),
239        );
240    }
241
242    pub fn append_stored_user_blocks(&self, session_id: &str, blocks: Vec<llm::ContentBlock>) {
243        self.append_stored_event(session_id, &SessionEvent::User(UserEvent::Message { content: blocks }));
244    }
245
246    pub fn append_stored_agent_turn(&self, session_id: &str, text: &str) {
247        self.append_stored_agent_text(session_id, text);
248        self.append_stored_event(
249            session_id,
250            &SessionEvent::Agent(AgentEvent::Turn(TurnEvent::Ended { outcome: TurnOutcome::Completed })),
251        );
252    }
253
254    pub fn append_stored_agent_text(&self, session_id: &str, text: &str) {
255        self.append_stored_event(
256            session_id,
257            &SessionEvent::Agent(AgentEvent::Message(MessageEvent::Text {
258                message_id: "msg".to_string(),
259                chunk: text.to_string(),
260                is_complete: true,
261            })),
262        );
263    }
264
265    async fn insert_switching_session(
266        &self,
267        acp_session_id: SessionId,
268        events: Vec<SessionEvent>,
269        selected_mode: Option<String>,
270        serverless_coder: bool,
271    ) -> FakeAgentSwitchingSession {
272        let (planner_def, planner) = fake_agent("Planner", "planner-mcp", "plan", PLANNER_REPLY);
273        let (mut coder_def, coder) = fake_agent("Coder", "coder-mcp", "edit", CODER_REPLY);
274        if serverless_coder {
275            coder_def.mcp = None;
276        }
277
278        let mut catalog_specs = Vec::new();
279        let mut agents = HashMap::new();
280        for def in [planner_def, coder_def] {
281            catalog_specs.push(def.spec.clone());
282            agents.insert(def.spec.name.clone(), def);
283        }
284        let specs = SessionAgents::new(AgentCatalog::new(PathBuf::from("/tmp"), catalog_specs, None));
285
286        let factory = Arc::new(FakeRuntimeFactory { cwd: PathBuf::from("/tmp"), agents });
287        let initial_agent = selected_mode.clone().unwrap_or_else(|| "Planner".to_string());
288
289        let handle = SessionActor::spawn(SessionActorInit {
290            session_id: acp_session_id.clone(),
291            connection: self.agent_cx.clone(),
292            repository: self.session_store.clone(),
293            oauth_credential_store: fake_oauth_store(),
294            active_agent: AgentKey::Named(initial_agent),
295            specs,
296            runtime_factory: factory,
297            transcript: events,
298            modes: switching_modes(),
299            config: SessionConfigState::with_selection("anthropic:claude-sonnet-4-5".to_string(), selected_mode, None),
300        })
301        .await
302        .expect("fake agent switching session actor spawns");
303        self.state.register_session(&acp_session_id, handle).await;
304        FakeAgentSwitchingSession { session_id: acp_session_id, planner, coder }
305    }
306
307    fn append_stored_event(&self, session_id: &str, event: &SessionEvent) {
308        self.session_store.append_event(session_id, event).expect("stored session event appends");
309    }
310}
311
312impl FakeAgentSwitchingSession {
313    pub fn session_id(&self) -> &SessionId {
314        &self.session_id
315    }
316
317    pub fn planner(&self) -> &FakeAcpAgent {
318        &self.planner
319    }
320
321    pub fn coder(&self) -> &FakeAcpAgent {
322        &self.coder
323    }
324
325    pub fn agent(&self, name: &str) -> &FakeAcpAgent {
326        match name {
327            "Planner" => &self.planner,
328            "Coder" => &self.coder,
329            other => panic!("unknown fake ACP agent {other:?}"),
330        }
331    }
332}
333
334impl FakeAcpAgent {
335    pub fn name(&self) -> &str {
336        &self.name
337    }
338
339    /// Asserts the agent's most recent turn saw a conversation containing each
340    /// of `expected` (user or assistant text), in addition to anything else.
341    pub fn assert_saw(&self, expected: &[&str]) {
342        let seen = self.latest_conversation();
343        for text in expected {
344            assert!(seen.iter().any(|m| m == text), "{} should have seen {text:?}; saw {seen:?}", self.name);
345        }
346    }
347
348    /// Asserts the agent's most recent turn saw *exactly* `expected` and nothing
349    /// else — used to prove a freshly-activated agent started with no prior
350    /// transcript.
351    pub fn assert_saw_exactly(&self, expected: &[&str]) {
352        let seen = self.latest_conversation();
353        let expected: Vec<String> = expected.iter().map(|t| (*t).to_string()).collect();
354        assert_eq!(seen, expected, "{} conversation mismatch", self.name);
355    }
356
357    /// Asserts the agent never ran a turn (its LLM was never invoked).
358    pub fn assert_never_ran(&self) {
359        let contexts = self.captured_contexts.lock().expect("captured contexts lock is healthy");
360        assert!(contexts.is_empty(), "{} should not have run; captured {} context(s)", self.name, contexts.len());
361    }
362
363    fn latest_conversation(&self) -> Vec<String> {
364        let contexts = self.captured_contexts.lock().expect("captured contexts lock is healthy");
365        let latest = contexts.last().unwrap_or_else(|| panic!("{} should have run a turn", self.name));
366        conversation_texts(latest)
367    }
368}
369
370/// Spawns each agent's runtime through the real [`AgentRuntime`] wiring, but
371/// backed by a [`FakeLlmProvider`] and an in-memory MCP server instead of a
372/// network LLM and external MCP processes.
373struct FakeRuntimeFactory {
374    cwd: PathBuf,
375    agents: HashMap<String, FakeAgentDef>,
376}
377
378struct FakeAgentDef {
379    spec: AgentSpec,
380    provider: Mutex<Option<Arc<dyn StreamingModelProvider>>>,
381    mcp: Option<(String, String)>,
382}
383
384#[async_trait::async_trait]
385impl RuntimeFactory for FakeRuntimeFactory {
386    async fn spawn(
387        &self,
388        agent: AgentKey,
389        spec: &AgentSpec,
390        initial_messages: Vec<ChatMessage>,
391        runtime_event_tx: mpsc::Sender<RuntimeEvent>,
392    ) -> Result<AgentRuntime, SessionError> {
393        let def = self
394            .agents
395            .get(&spec.name)
396            .or_else(|| self.agents.values().next())
397            .ok_or_else(|| SessionError::AgentNotFound(spec.name.clone()))?;
398        let provider = def
399            .provider
400            .lock()
401            .expect("fake provider lock is healthy")
402            .take()
403            .expect("fake agent runtime spawned more than once");
404
405        let mut mcp_builder = mcp(&self.cwd).with_tool_filter(spec.tools.clone());
406        if let Some((server_name, prompt_name)) = &def.mcp {
407            let factory_name = server_name.clone();
408            let prompt_name = prompt_name.clone();
409            let factory: ServerFactory = Box::new(move |_spec, _services| {
410                let prompt_name = prompt_name.clone();
411                async move { FakePromptMcp::new(&prompt_name).into_dyn() }.boxed()
412            });
413            mcp_builder = mcp_builder.register_in_memory_server(factory_name.clone(), factory).with_servers(vec![
414                McpServer::new(
415                    server_name.clone(),
416                    McpTransport::InMemory {
417                        spec: InMemoryServerSpec { factory: factory_name, args: Vec::new(), input: None },
418                    },
419                    ToolExposure::ModelVisible,
420                ),
421            ]);
422        }
423        let mut spawn =
424            mcp_builder.spawn().await.map_err(|e| SessionError::Build(CliError::McpError(e.to_string())))?;
425        spawn
426            .block_until_ready()
427            .await
428            .ok_or_else(|| SessionError::McpOperation("fake MCP bootstrap aborted".to_string()))?;
429        let mcp_handle = spawn.handle().clone();
430        let mut builder = AgentBuilder::new(provider).max_auto_continues(0);
431        for prompt in &spec.prompts {
432            builder = builder.system_prompt(prompt.clone());
433        }
434        let (agent_tx, agent_rx, agent_handle) = builder
435            .tools(mcp_handle, Vec::new())
436            .messages(initial_messages)
437            .spawn()
438            .await
439            .map_err(|e| SessionError::Build(CliError::AgentError(e.to_string())))?;
440        let (mcp_runtime, event_rx) = spawn.connect_agent(agent_tx.clone()).await.split();
441
442        Ok(AgentRuntime::new(agent, agent_tx, agent_rx, Some(agent_handle), event_rx, mcp_runtime, runtime_event_tx))
443    }
444}
445
446struct StubRuntimeFactory {
447    cwd: PathBuf,
448    agent_parts: Mutex<Option<StubAgentParts>>,
449}
450
451struct StubAgentParts {
452    tx: mpsc::Sender<Command>,
453    rx: mpsc::Receiver<AgentEvent>,
454    handle: AgentHandle,
455}
456
457#[async_trait::async_trait]
458impl RuntimeFactory for StubRuntimeFactory {
459    async fn spawn(
460        &self,
461        agent: AgentKey,
462        _spec: &AgentSpec,
463        _initial_messages: Vec<ChatMessage>,
464        runtime_event_tx: mpsc::Sender<RuntimeEvent>,
465    ) -> Result<AgentRuntime, SessionError> {
466        let parts = self
467            .agent_parts
468            .lock()
469            .expect("stub agent parts lock is healthy")
470            .take()
471            .expect("stub runtime spawned more than once");
472
473        let mut spawn =
474            mcp(&self.cwd).spawn().await.map_err(|e| SessionError::Build(CliError::McpError(e.to_string())))?;
475        spawn
476            .block_until_ready()
477            .await
478            .ok_or_else(|| SessionError::McpOperation("stub MCP bootstrap aborted".to_string()))?;
479        let (mcp_runtime, event_rx) = spawn.connect_agent(parts.tx.clone()).await.split();
480
481        Ok(AgentRuntime::new(agent, parts.tx, parts.rx, Some(parts.handle), event_rx, mcp_runtime, runtime_event_tx))
482    }
483}
484
485fn fake_agent(name: &str, server_name: &str, prompt_name: &str, reply: &str) -> (FakeAgentDef, FakeAcpAgent) {
486    let provider =
487        FakeLlmProvider::new(vec![vec![LlmResponse::start("msg"), LlmResponse::text(reply), LlmResponse::done()]])
488            .with_display_name(name);
489    let captured_contexts = provider.captured_contexts();
490    let def = FakeAgentDef {
491        spec: fake_agent_spec(name),
492        provider: Mutex::new(Some(Arc::new(provider))),
493        mcp: Some((server_name.to_string(), prompt_name.to_string())),
494    };
495    let observer = FakeAcpAgent { name: name.to_string(), captured_contexts };
496    (def, observer)
497}
498
499fn fake_oauth_store() -> Arc<dyn OAuthCredentialStorage> {
500    Arc::new(aether_auth::FakeOAuthCredentialStore::new())
501}
502
503fn switching_modes() -> Modes {
504    Modes::new(vec![
505        ValidatedMode {
506            name: "Planner".to_string(),
507            model: "anthropic:claude-sonnet-4-5".to_string(),
508            reasoning_effort: None,
509        },
510        ValidatedMode {
511            name: "Coder".to_string(),
512            model: "deepseek:deepseek-v4-flash".to_string(),
513            reasoning_effort: None,
514        },
515    ])
516}
517
518fn fake_agent_spec(name: &str) -> AgentSpec {
519    let model: llm::catalog::LlmModel = "anthropic:claude-sonnet-4-5".parse().expect("test model parses");
520    let mut spec = AgentSpec::bare(&model, None, vec![Prompt::text(&format!("{name} system prompt"))]);
521    spec.name = name.to_string();
522    spec.description = format!("{name} test agent");
523    spec.exposure = AgentSpecExposure::user_only();
524    spec
525}
526
527fn assert_available_commands(update: SessionUpdate, expected: &[&str], unexpected: &[&str]) {
528    let SessionUpdate::AvailableCommandsUpdate(commands) = update else {
529        panic!("expected available commands update");
530    };
531    let names = commands.available_commands.iter().map(|command| command.name.as_str()).collect::<Vec<_>>();
532    for name in expected {
533        assert!(names.contains(name), "expected command /{name} in {names:?}");
534    }
535    for name in unexpected {
536        assert!(!names.contains(name), "did not expect command /{name} in {names:?}");
537    }
538}
539
540fn assert_server_status(notification: McpNotification, expected: &[&str]) {
541    let McpNotification::ServerStatus { servers } = notification;
542    let names = servers.iter().map(|server| server.name.as_str()).collect::<Vec<_>>();
543    for server_name in expected {
544        assert!(names.contains(server_name), "expected server {server_name} in {names:?}");
545    }
546}
547
548fn assert_server_status_exact(notification: McpNotification, expected: &[&str]) {
549    let McpNotification::ServerStatus { servers } = notification;
550    let names = servers.iter().map(|server| server.name.as_str()).collect::<Vec<_>>();
551    assert_eq!(names, expected);
552}
553
554fn conversation_texts(context: &Context) -> Vec<String> {
555    context
556        .messages()
557        .iter()
558        .filter_map(|message| match message {
559            ChatMessage::User { content, .. } => llm::ContentBlock::first_text(content).map(str::to_string),
560            ChatMessage::Assistant { content, .. } if !content.is_empty() => Some(content.clone()),
561            _ => None,
562        })
563        .collect()
564}
aether_cli/acp/testing.rs

aether_cli/acp/
testing.rs
