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::acp::session_store::SessionMeta;
8use crate::settings_args::SettingsSourceArgs;
9use acp_utils::testing::{TestPeer, duplex_pair};
10use aether_auth::OAuthCredentialStorage;
11use aether_core::context::ext::{SessionEvent, UserEvent};
12use aether_core::core::AgentHandle;
13use aether_core::events::{AgentMessage, UserMessage};
14use agent_client_protocol::schema::SessionId;
15use agent_client_protocol::{Agent, Client, ConnectionTo};
16use llm::ProviderConnectionOverrides;
17use std::path::PathBuf;
18use std::sync::Arc;
19use tempfile::TempDir;
20use tokio::sync::{mpsc, oneshot};
21use tokio::task::spawn_local;
22
23fn mock_oauth_store() -> Arc<dyn OAuthCredentialStorage> {
24    Arc::new(aether_auth::FakeOAuthCredentialStore::new())
25}
26
27/// In-memory ACP harness running the real `acp_agent_builder` against a
28/// pre-wired test client. Created via [`AcpTestHarness::start`] inside a
29/// `LocalSet`. The harness owns its own [`SessionRegistry`] and a
30/// temp-dir-backed [`SessionStore`] so tests can register fake-driven
31/// sessions without going through `new_session`.
32pub struct AcpTestHarness {
33    pub client_cx: ConnectionTo<Agent>,
34    pub peer: TestPeer,
35    agent_cx: ConnectionTo<Client>,
36    registry: Arc<SessionRegistry>,
37    session_store: Arc<SessionStore>,
38    _tmp: TempDir,
39}
40
41impl AcpTestHarness {
42    pub async fn start() -> Self {
43        let tmp = tempfile::tempdir().expect("tempdir for session store");
44        let registry = Arc::new(SessionRegistry::new());
45        let session_store = Arc::new(SessionStore::from_path(tmp.path().to_path_buf()));
46        let manager = Arc::new(SessionManager::new(SessionManagerConfig {
47            registry: registry.clone(),
48            session_store: session_store.clone(),
49            oauth_credential_store: mock_oauth_store(),
50            initial_selection: InitialSessionSelection::default(),
51            settings_source: SettingsSourceArgs::default(),
52            provider_connections: ProviderConnectionOverrides::default(),
53        }));
54
55        let (peer, client_builder) = TestPeer::new();
56        let (agent_transport, client_transport) = duplex_pair();
57        let (agent_cx_tx, agent_cx_rx) = oneshot::channel::<ConnectionTo<Client>>();
58        let (client_cx_tx, client_cx_rx) = oneshot::channel::<ConnectionTo<Agent>>();
59
60        spawn_local(async move {
61            let _ = acp_agent_builder(manager)
62                .connect_with(agent_transport, async move |cx: ConnectionTo<Client>| {
63                    let _ = agent_cx_tx.send(cx);
64                    std::future::pending::<()>().await;
65                    Ok(())
66                })
67                .await;
68        });
69
70        spawn_local(async move {
71            let _ = client_builder
72                .connect_with(client_transport, async move |cx: ConnectionTo<Agent>| {
73                    let _ = client_cx_tx.send(cx);
74                    std::future::pending::<()>().await;
75                    Ok(())
76                })
77                .await;
78        });
79
80        let agent_cx = agent_cx_rx.await.expect("agent side connect_with produced a ConnectionTo");
81        let client_cx = client_cx_rx.await.expect("client side connect_with produced a ConnectionTo");
82        Self { client_cx, peer, agent_cx, registry, session_store, _tmp: tmp }
83    }
84
85    /// Register a stub session built from a hand-spawned
86    /// `(agent_tx, agent_rx, agent_handle)` triple — typically from
87    /// `aether_core::core::agent(fake_llm).spawn().await`. MCP channels are
88    /// stubbed: no servers, no events. The session is routable via
89    /// `mgr.prompt(id)` / `mgr.cancel(id)`.
90    pub async fn insert_stub_session(
91        &self,
92        agent_tx: mpsc::Sender<UserMessage>,
93        agent_rx: mpsc::Receiver<AgentMessage>,
94        agent_handle: AgentHandle,
95        id: SessionId,
96        model: &str,
97    ) {
98        let (mcp_tx, _mcp_rx) = mpsc::channel(1);
99        let (_event_tx, event_rx) = mpsc::channel(1);
100        let session = Session {
101            agent_tx,
102            agent_rx,
103            agent_handle,
104            _mcp_handle: tokio::spawn(async {}),
105            mcp_tx,
106            event_rx,
107            initial_server_statuses: vec![],
108            provider_connections: ProviderConnectionOverrides::default(),
109        };
110        let relay =
111            spawn_relay(session, self.agent_cx.clone(), id.clone(), self.session_store.clone(), mock_oauth_store());
112        self.registry.insert(id.0.to_string(), relay, model.to_string(), None, None, vec![]).await;
113    }
114
115    pub fn append_stored_session(&self, session_id: &str, created_at: &str) {
116        let meta = SessionMeta {
117            session_id: session_id.to_string(),
118            cwd: PathBuf::from("/tmp"),
119            model: "test-model".to_string(),
120            selected_mode: None,
121            created_at: created_at.to_string(),
122        };
123
124        self.session_store.append_meta(session_id, &meta).expect("stored session meta appends");
125    }
126
127    pub fn append_stored_prompt(&self, session_id: &str, prompt: &str) {
128        self.append_stored_event(
129            session_id,
130            &SessionEvent::User(UserEvent::Message { content: vec![llm::ContentBlock::text(prompt)] }),
131        );
132    }
133
134    pub fn append_stored_user_blocks(&self, session_id: &str, blocks: Vec<llm::ContentBlock>) {
135        self.append_stored_event(session_id, &SessionEvent::User(UserEvent::Message { content: blocks }));
136    }
137
138    pub fn append_stored_agent_text(&self, session_id: &str, text: &str) {
139        self.append_stored_event(
140            session_id,
141            &SessionEvent::Agent(AgentMessage::Text {
142                message_id: "msg".to_string(),
143                chunk: text.to_string(),
144                is_complete: true,
145                model_name: "test".to_string(),
146            }),
147        );
148    }
149
150    fn append_stored_event(&self, session_id: &str, event: &SessionEvent) {
151        self.session_store.append_event(session_id, event).expect("stored session event appends");
152    }
153}
aether_cli/acp/testing.rs

aether_cli/acp/
testing.rs
