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::McpSpawnResult;
24use aether_core::mcp::mcp;
25use aether_core::session::{SessionControlEvent, SessionEvent, SessionMeta, UserEvent, last_agent_from_events};
26use aether_project::AgentCatalog;
27use agent_client_protocol::schema::{SessionId, SessionUpdate};
28use agent_client_protocol::{Agent, Client, ConnectionTo};
29use llm::ProviderConnectionOverrides;
30use llm::testing::FakeLlmProvider;
31use llm::{ChatMessage, Context, LlmResponse, StreamingModelProvider};
32use mcp_utils::client::{McpServer, McpTransport};
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 servers = def.mcp.as_ref().map_or_else(Vec::new, |(server_name, prompt_name)| {
384            vec![McpServer::new(
385                server_name.clone(),
386                McpTransport::InMemory { server: FakePromptMcp::new(prompt_name).into_dyn() },
387                false,
388            )]
389        });
390        let mut spawn = mcp(&self.cwd)
391            .with_servers(servers)
392            .spawn()
393            .await
394            .map_err(|e| SessionError::Build(CliError::McpError(e.to_string())))?;
395        let snapshot = spawn
396            .block_until_ready()
397            .await
398            .ok_or_else(|| SessionError::McpOperation("fake MCP bootstrap aborted".to_string()))?;
399        let McpSpawnResult { command_tx: mcp_tx, event_rx, handle: mcp_handle } = spawn;
400
401        let filtered_tools = spec.tools.apply(snapshot.tool_definitions.clone());
402        let mut builder = AgentBuilder::new(provider).max_auto_continues(0);
403        for prompt in &spec.prompts {
404            builder = builder.system_prompt(prompt.clone());
405        }
406        let (agent_tx, agent_rx, agent_handle) = builder
407            .tools(mcp_tx.clone(), filtered_tools)
408            .messages(initial_messages)
409            .spawn()
410            .await
411            .map_err(|e| SessionError::Build(CliError::AgentError(e.to_string())))?;
412
413        Ok(AgentRuntime::new(
414            agent,
415            spec,
416            agent_tx,
417            agent_rx,
418            Some(agent_handle),
419            mcp_tx,
420            event_rx,
421            mcp_handle,
422            snapshot,
423            runtime_event_tx,
424        ))
425    }
426}
427
428struct StubRuntimeFactory {
429    cwd: PathBuf,
430    agent_parts: Mutex<Option<StubAgentParts>>,
431}
432
433struct StubAgentParts {
434    tx: mpsc::Sender<Command>,
435    rx: mpsc::Receiver<AgentEvent>,
436    handle: AgentHandle,
437}
438
439#[async_trait::async_trait]
440impl RuntimeFactory for StubRuntimeFactory {
441    async fn spawn(
442        &self,
443        agent: AgentKey,
444        spec: &AgentSpec,
445        _initial_messages: Vec<ChatMessage>,
446        runtime_event_tx: mpsc::Sender<RuntimeEvent>,
447    ) -> Result<AgentRuntime, SessionError> {
448        let parts = self
449            .agent_parts
450            .lock()
451            .expect("stub agent parts lock is healthy")
452            .take()
453            .expect("stub runtime spawned more than once");
454
455        let mut spawn =
456            mcp(&self.cwd).spawn().await.map_err(|e| SessionError::Build(CliError::McpError(e.to_string())))?;
457        let snapshot = spawn
458            .block_until_ready()
459            .await
460            .ok_or_else(|| SessionError::McpOperation("stub MCP bootstrap aborted".to_string()))?;
461        let McpSpawnResult { command_tx: mcp_tx, event_rx, handle: mcp_handle } = spawn;
462
463        Ok(AgentRuntime::new(
464            agent,
465            spec,
466            parts.tx,
467            parts.rx,
468            Some(parts.handle),
469            mcp_tx,
470            event_rx,
471            mcp_handle,
472            snapshot,
473            runtime_event_tx,
474        ))
475    }
476}
477
478fn fake_agent(name: &str, server_name: &str, prompt_name: &str, reply: &str) -> (FakeAgentDef, FakeAcpAgent) {
479    let provider =
480        FakeLlmProvider::new(vec![vec![LlmResponse::start("msg"), LlmResponse::text(reply), LlmResponse::done()]])
481            .with_display_name(name);
482    let captured_contexts = provider.captured_contexts();
483    let def = FakeAgentDef {
484        spec: fake_agent_spec(name),
485        provider: Mutex::new(Some(Arc::new(provider))),
486        mcp: Some((server_name.to_string(), prompt_name.to_string())),
487    };
488    let observer = FakeAcpAgent { name: name.to_string(), captured_contexts };
489    (def, observer)
490}
491
492fn fake_oauth_store() -> Arc<dyn OAuthCredentialStorage> {
493    Arc::new(aether_auth::FakeOAuthCredentialStore::new())
494}
495
496fn switching_modes() -> Modes {
497    Modes::new(vec![
498        ValidatedMode {
499            name: "Planner".to_string(),
500            model: "anthropic:claude-sonnet-4-5".to_string(),
501            reasoning_effort: None,
502        },
503        ValidatedMode {
504            name: "Coder".to_string(),
505            model: "deepseek:deepseek-chat".to_string(),
506            reasoning_effort: None,
507        },
508    ])
509}
510
511fn fake_agent_spec(name: &str) -> AgentSpec {
512    let model: llm::catalog::LlmModel = "anthropic:claude-sonnet-4-5".parse().expect("test model parses");
513    let mut spec = AgentSpec::bare(&model, None, vec![Prompt::text(&format!("{name} system prompt"))]);
514    spec.name = name.to_string();
515    spec.description = format!("{name} test agent");
516    spec.exposure = AgentSpecExposure::user_only();
517    spec
518}
519
520fn assert_available_commands(update: SessionUpdate, expected: &[&str], unexpected: &[&str]) {
521    let SessionUpdate::AvailableCommandsUpdate(commands) = update else {
522        panic!("expected available commands update");
523    };
524    let names = commands.available_commands.iter().map(|command| command.name.as_str()).collect::<Vec<_>>();
525    for name in expected {
526        assert!(names.contains(name), "expected command /{name} in {names:?}");
527    }
528    for name in unexpected {
529        assert!(!names.contains(name), "did not expect command /{name} in {names:?}");
530    }
531}
532
533fn assert_server_status(notification: McpNotification, expected: &[&str]) {
534    let McpNotification::ServerStatus { servers } = notification else {
535        panic!("expected server status notification");
536    };
537    let names = servers.iter().map(|server| server.name.as_str()).collect::<Vec<_>>();
538    for server_name in expected {
539        assert!(names.contains(server_name), "expected server {server_name} in {names:?}");
540    }
541}
542
543fn assert_server_status_exact(notification: McpNotification, expected: &[&str]) {
544    let McpNotification::ServerStatus { servers } = notification else {
545        panic!("expected server status notification");
546    };
547    let names = servers.iter().map(|server| server.name.as_str()).collect::<Vec<_>>();
548    assert_eq!(names, expected);
549}
550
551fn conversation_texts(context: &Context) -> Vec<String> {
552    context
553        .messages()
554        .iter()
555        .filter_map(|message| match message {
556            ChatMessage::User { content, .. } => llm::ContentBlock::first_text(content).map(str::to_string),
557            ChatMessage::Assistant { content, .. } if !content.is_empty() => Some(content.clone()),
558            _ => None,
559        })
560        .collect()
561}
aether_cli/acp/testing.rs

aether_cli/acp/
testing.rs
