aether_cli/acp/
testing.rs

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

aether_cli/acp/
testing.rs
