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

aether_cli/acp/
testing.rs
