aether_cli/acp/
testing.rs

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

aether_cli/acp/
testing.rs
