aether_cli/acp/
testing.rs1use 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_core::core::AgentHandle;
10use aether_core::events::{AgentMessage, UserMessage};
11use agent_client_protocol::schema::SessionId;
12use agent_client_protocol::{Agent, Client, ConnectionTo};
13use llm::oauth::OAuthCredentialStore;
14use std::sync::Arc;
15use tempfile::TempDir;
16use tokio::sync::{mpsc, oneshot};
17use tokio::task::spawn_local;
18
19fn mock_oauth_store() -> OAuthCredentialStore {
20 OAuthCredentialStore::with_mock_store().unwrap()
21}
22
23pub struct AcpTestHarness {
29 pub client_cx: ConnectionTo<Agent>,
30 pub peer: TestPeer,
31 agent_cx: ConnectionTo<Client>,
32 registry: Arc<SessionRegistry>,
33 session_store: Arc<SessionStore>,
34 _tmp: TempDir,
35}
36
37impl AcpTestHarness {
38 pub async fn start() -> Self {
39 let tmp = tempfile::tempdir().expect("tempdir for session store");
40 let registry = Arc::new(SessionRegistry::new());
41 let session_store = Arc::new(SessionStore::from_path(tmp.path().to_path_buf()));
42 let manager = Arc::new(SessionManager::new(SessionManagerConfig {
43 registry: registry.clone(),
44 session_store: session_store.clone(),
45 oauth_credential_store: mock_oauth_store(),
46 initial_selection: InitialSessionSelection::default(),
47 settings_source: SettingsSourceArgs::default(),
48 }));
49
50 let (peer, client_builder) = TestPeer::new();
51 let (agent_transport, client_transport) = duplex_pair();
52 let (agent_cx_tx, agent_cx_rx) = oneshot::channel::<ConnectionTo<Client>>();
53 let (client_cx_tx, client_cx_rx) = oneshot::channel::<ConnectionTo<Agent>>();
54
55 spawn_local(async move {
56 let _ = acp_agent_builder(manager)
57 .connect_with(agent_transport, async move |cx: ConnectionTo<Client>| {
58 let _ = agent_cx_tx.send(cx);
59 std::future::pending::<()>().await;
60 Ok(())
61 })
62 .await;
63 });
64
65 spawn_local(async move {
66 let _ = client_builder
67 .connect_with(client_transport, async move |cx: ConnectionTo<Agent>| {
68 let _ = client_cx_tx.send(cx);
69 std::future::pending::<()>().await;
70 Ok(())
71 })
72 .await;
73 });
74
75 let agent_cx = agent_cx_rx.await.expect("agent side connect_with produced a ConnectionTo");
76 let client_cx = client_cx_rx.await.expect("client side connect_with produced a ConnectionTo");
77 Self { client_cx, peer, agent_cx, registry, session_store, _tmp: tmp }
78 }
79
80 pub async fn insert_stub_session(
86 &self,
87 agent_tx: mpsc::Sender<UserMessage>,
88 agent_rx: mpsc::Receiver<AgentMessage>,
89 agent_handle: AgentHandle,
90 id: SessionId,
91 model: &str,
92 ) {
93 let (mcp_tx, _mcp_rx) = mpsc::channel(1);
94 let (_event_tx, event_rx) = mpsc::channel(1);
95 let session = Session {
96 agent_tx,
97 agent_rx,
98 agent_handle,
99 _mcp_handle: tokio::spawn(async {}),
100 mcp_tx,
101 event_rx,
102 initial_server_statuses: vec![],
103 };
104 let relay = spawn_relay(session, self.agent_cx.clone(), id.clone(), self.session_store.clone());
105 self.registry.insert(id.0.to_string(), relay, model.to_string(), None, None, vec![]).await;
106 }
107}