aether_cli/acp/
testing.rs

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

aether_cli/acp/
testing.rs
