aether_cli/acp/
testing.rs

1use super::handlers::acp_agent_builder;
2use super::relay::spawn_relay;
3use super::session::Session;
4use super::session_manager::{InitialSessionSelection, SessionManager, SessionManagerConfig};
5use super::session_registry::SessionRegistry;
6use super::session_store::SessionStore;
7use crate::settings_args::SettingsSourceArgs;
8use acp_utils::testing::{TestPeer, duplex_pair};
9use aether_auth::OAuthCredentialStorage;
10use aether_core::core::AgentHandle;
11use aether_core::events::{AgentMessage, UserMessage};
12use agent_client_protocol::schema::SessionId;
13use agent_client_protocol::{Agent, Client, ConnectionTo};
14use llm::ProviderConnectionOverrides;
15use std::sync::Arc;
16use tempfile::TempDir;
17use tokio::sync::{mpsc, oneshot};
18use tokio::task::spawn_local;
19
20fn mock_oauth_store() -> Arc<dyn OAuthCredentialStorage> {
21    Arc::new(aether_auth::FakeOAuthCredentialStore::new())
22}
23
24/// In-memory ACP harness running the real `acp_agent_builder` against a
25/// pre-wired test client. Created via [`AcpTestHarness::start`] inside a
26/// `LocalSet`. The harness owns its own [`SessionRegistry`] and a
27/// temp-dir-backed [`SessionStore`] so tests can register fake-driven
28/// sessions without going through `new_session`.
29pub struct AcpTestHarness {
30    pub client_cx: ConnectionTo<Agent>,
31    pub peer: TestPeer,
32    agent_cx: ConnectionTo<Client>,
33    registry: Arc<SessionRegistry>,
34    session_store: Arc<SessionStore>,
35    _tmp: TempDir,
36}
37
38impl AcpTestHarness {
39    pub async fn start() -> Self {
40        let tmp = tempfile::tempdir().expect("tempdir for session store");
41        let registry = Arc::new(SessionRegistry::new());
42        let session_store = Arc::new(SessionStore::from_path(tmp.path().to_path_buf()));
43        let manager = Arc::new(SessionManager::new(SessionManagerConfig {
44            registry: registry.clone(),
45            session_store: session_store.clone(),
46            oauth_credential_store: mock_oauth_store(),
47            initial_selection: InitialSessionSelection::default(),
48            settings_source: SettingsSourceArgs::default(),
49            provider_connections: ProviderConnectionOverrides::default(),
50        }));
51
52        let (peer, client_builder) = TestPeer::new();
53        let (agent_transport, client_transport) = duplex_pair();
54        let (agent_cx_tx, agent_cx_rx) = oneshot::channel::<ConnectionTo<Client>>();
55        let (client_cx_tx, client_cx_rx) = oneshot::channel::<ConnectionTo<Agent>>();
56
57        spawn_local(async move {
58            let _ = acp_agent_builder(manager)
59                .connect_with(agent_transport, async move |cx: ConnectionTo<Client>| {
60                    let _ = agent_cx_tx.send(cx);
61                    std::future::pending::<()>().await;
62                    Ok(())
63                })
64                .await;
65        });
66
67        spawn_local(async move {
68            let _ = client_builder
69                .connect_with(client_transport, async move |cx: ConnectionTo<Agent>| {
70                    let _ = client_cx_tx.send(cx);
71                    std::future::pending::<()>().await;
72                    Ok(())
73                })
74                .await;
75        });
76
77        let agent_cx = agent_cx_rx.await.expect("agent side connect_with produced a ConnectionTo");
78        let client_cx = client_cx_rx.await.expect("client side connect_with produced a ConnectionTo");
79        Self { client_cx, peer, agent_cx, registry, session_store, _tmp: tmp }
80    }
81
82    /// Register a stub session built from a hand-spawned
83    /// `(agent_tx, agent_rx, agent_handle)` triple — typically from
84    /// `aether_core::core::agent(fake_llm).spawn().await`. MCP channels are
85    /// stubbed: no servers, no events. The session is routable via
86    /// `mgr.prompt(id)` / `mgr.cancel(id)`.
87    pub async fn insert_stub_session(
88        &self,
89        agent_tx: mpsc::Sender<UserMessage>,
90        agent_rx: mpsc::Receiver<AgentMessage>,
91        agent_handle: AgentHandle,
92        id: SessionId,
93        model: &str,
94    ) {
95        let (mcp_tx, _mcp_rx) = mpsc::channel(1);
96        let (_event_tx, event_rx) = mpsc::channel(1);
97        let session = Session {
98            agent_tx,
99            agent_rx,
100            agent_handle,
101            _mcp_handle: tokio::spawn(async {}),
102            mcp_tx,
103            event_rx,
104            initial_server_statuses: vec![],
105            provider_connections: ProviderConnectionOverrides::default(),
106        };
107        let relay =
108            spawn_relay(session, self.agent_cx.clone(), id.clone(), self.session_store.clone(), mock_oauth_store());
109        self.registry.insert(id.0.to_string(), relay, model.to_string(), None, None, vec![]).await;
110    }
111}
aether_cli/acp/testing.rs

aether_cli/acp/
testing.rs
