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
19pub struct AcpTestHarness {
25 pub client_cx: ConnectionTo<Agent>,
26 pub peer: TestPeer,
27 agent_cx: ConnectionTo<Client>,
28 registry: Arc<SessionRegistry>,
29 session_store: Arc<SessionStore>,
30 _tmp: TempDir,
31}
32
33impl AcpTestHarness {
34 pub async fn start() -> Self {
35 let tmp = tempfile::tempdir().expect("tempdir for session store");
36 let registry = Arc::new(SessionRegistry::new());
37 let session_store = Arc::new(SessionStore::from_path(tmp.path().to_path_buf()));
38 let manager = Arc::new(SessionManager::new(SessionManagerConfig {
39 registry: registry.clone(),
40 session_store: session_store.clone(),
41 has_oauth_credential: OAuthCredentialStore::has_credential,
42 initial_selection: InitialSessionSelection::default(),
43 settings_source: SettingsSourceArgs::default(),
44 }));
45
46 let (peer, client_builder) = TestPeer::new();
47 let (agent_transport, client_transport) = duplex_pair();
48 let (agent_cx_tx, agent_cx_rx) = oneshot::channel::<ConnectionTo<Client>>();
49 let (client_cx_tx, client_cx_rx) = oneshot::channel::<ConnectionTo<Agent>>();
50
51 spawn_local(async move {
52 let _ = acp_agent_builder(manager)
53 .connect_with(agent_transport, async move |cx: ConnectionTo<Client>| {
54 let _ = agent_cx_tx.send(cx);
55 std::future::pending::<()>().await;
56 Ok(())
57 })
58 .await;
59 });
60
61 spawn_local(async move {
62 let _ = client_builder
63 .connect_with(client_transport, async move |cx: ConnectionTo<Agent>| {
64 let _ = client_cx_tx.send(cx);
65 std::future::pending::<()>().await;
66 Ok(())
67 })
68 .await;
69 });
70
71 let agent_cx = agent_cx_rx.await.expect("agent side connect_with produced a ConnectionTo");
72 let client_cx = client_cx_rx.await.expect("client side connect_with produced a ConnectionTo");
73 Self { client_cx, peer, agent_cx, registry, session_store, _tmp: tmp }
74 }
75
76 pub async fn insert_stub_session(
82 &self,
83 agent_tx: mpsc::Sender<UserMessage>,
84 agent_rx: mpsc::Receiver<AgentMessage>,
85 agent_handle: AgentHandle,
86 id: SessionId,
87 model: &str,
88 ) {
89 let (mcp_tx, _mcp_rx) = mpsc::channel(1);
90 let (_event_tx, event_rx) = mpsc::channel(1);
91 let session = Session {
92 agent_tx,
93 agent_rx,
94 agent_handle,
95 _mcp_handle: tokio::spawn(async {}),
96 mcp_tx,
97 event_rx,
98 initial_server_statuses: vec![],
99 };
100 let relay = spawn_relay(session, self.agent_cx.clone(), id.clone(), self.session_store.clone());
101 self.registry.insert(id.0.to_string(), relay, model.to_string(), None, None, vec![]).await;
102 }
103}