acp_utils/client/
session.rs

1use super::error::AcpClientError;
2use super::event::AcpEvent;
3use crate::notifications::{
4    AuthMethodsUpdatedParams, ContextClearedParams, ContextCompactionParams, McpNotification, McpRequest,
5    PromptSearchParams, PromptSearchResponse, SessionPreviewParams, SessionPreviewResponse, SessionUsageParams,
6    SubAgentProgressParams, WorkspaceListParams, WorkspaceListResponse, WorkspaceMoveParams, WorkspaceMoveResponse,
7};
8use agent_client_protocol::schema::v1::{
9    AuthMethod, AuthenticateRequest, AuthenticateResponse, CancelNotification, CloseSessionRequest,
10    CloseSessionResponse, CreateElicitationRequest, InitializeRequest, InitializeResponse, ListSessionsRequest,
11    ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, NewSessionRequest, NewSessionResponse,
12    PermissionOptionId, PermissionOptionKind, PromptCapabilities, PromptRequest, PromptResponse,
13    RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse, ResumeSessionRequest,
14    ResumeSessionResponse, SelectedPermissionOutcome, SessionCapabilities, SessionId, SessionNotification,
15    SetSessionConfigOptionRequest, SetSessionConfigOptionResponse,
16};
17use agent_client_protocol::{self as acp, Client, ConnectTo, ConnectionTo, JsonRpcNotification, JsonRpcRequest};
18use std::sync::{Arc, Mutex};
19use tokio::sync::{mpsc, oneshot};
20use tracing::info;
21
22/// A cloneable handle for issuing typed lifecycle requests and prompt commands.
23#[derive(Clone)]
24pub struct AcpClientHandle {
25    cmd_tx: mpsc::UnboundedSender<ClientCommand>,
26}
27
28/// An initialized ACP connection that can create and manage multiple sessions.
29pub struct AcpClient {
30    pub initialize_response: InitializeResponse,
31    pub event_rx: mpsc::UnboundedReceiver<AcpEvent>,
32    pub handle: AcpClientHandle,
33}
34
35/// The result of loading an ACP session, including notifications sent before the response.
36pub struct LoadedSession {
37    pub session_id: SessionId,
38    pub response: LoadSessionResponse,
39    pub replay: Vec<SessionNotification>,
40}
41
42/// Connect to an ACP agent and complete initialization without creating a session.
43pub async fn connect_acp_client(
44    agent: impl ConnectTo<Client> + 'static,
45    init_request: InitializeRequest,
46) -> Result<AcpClient, AcpClientError> {
47    let (event_tx, event_rx) = mpsc::unbounded_channel::<AcpEvent>();
48    let (cmd_tx, cmd_rx) = mpsc::unbounded_channel::<ClientCommand>();
49    let (init_tx, init_rx) = oneshot::channel::<InitializeResult>();
50    let init_tx = Arc::new(Mutex::new(Some(init_tx)));
51    let replay_state = Arc::new(Mutex::new(None));
52
53    tokio::spawn(run_client_connection(
54        agent,
55        event_tx,
56        cmd_rx,
57        Arc::clone(&init_tx),
58        init_request,
59        Arc::clone(&replay_state),
60    ));
61
62    let initialize_response = init_rx
63        .await
64        .map_err(|_| AcpClientError::AgentCrashed("ACP task died during initialization".to_string()))??;
65
66    Ok(AcpClient { initialize_response, event_rx, handle: AcpClientHandle { cmd_tx } })
67}
68
69impl AcpClient {
70    /// The agent's display title, falling back to its implementation name.
71    pub fn agent_name(&self) -> String {
72        self.initialize_response
73            .agent_info
74            .as_ref()
75            .map_or_else(|| "agent".to_string(), |info| info.title.as_deref().unwrap_or(&info.name).to_string())
76    }
77
78    pub fn prompt_capabilities(&self) -> &PromptCapabilities {
79        &self.initialize_response.agent_capabilities.prompt_capabilities
80    }
81
82    pub fn session_capabilities(&self) -> &SessionCapabilities {
83        &self.initialize_response.agent_capabilities.session_capabilities
84    }
85
86    pub fn auth_methods(&self) -> &[AuthMethod] {
87        &self.initialize_response.auth_methods
88    }
89}
90
91impl AcpClientHandle {
92    #[cfg(feature = "testing")]
93    pub fn detached() -> Self {
94        let (cmd_tx, _) = mpsc::unbounded_channel();
95        Self { cmd_tx }
96    }
97
98    pub async fn prompt(&self, request: PromptRequest) -> Result<PromptResponse, AcpClientError> {
99        let (response, receiver) = oneshot::channel();
100        self.send(ClientCommand::Prompt { request, response })?;
101        await_response(receiver).await
102    }
103
104    /// Load a session and collect its replay notifications in wire order.
105    pub async fn load_session(&self, request: LoadSessionRequest) -> Result<LoadedSession, AcpClientError> {
106        let (response, receiver) = oneshot::channel();
107        self.send(ClientCommand::LoadSession { request, response })?;
108        await_response(receiver).await
109    }
110
111    pub async fn new_session(&self, request: NewSessionRequest) -> Result<NewSessionResponse, AcpClientError> {
112        self.request(request, false).await
113    }
114
115    pub async fn list_sessions(&self, request: ListSessionsRequest) -> Result<ListSessionsResponse, AcpClientError> {
116        self.request(request, false).await
117    }
118
119    /// Resume a session without collecting or replaying its prior notifications.
120    pub async fn resume_session(&self, request: ResumeSessionRequest) -> Result<ResumeSessionResponse, AcpClientError> {
121        self.request(request, false).await
122    }
123
124    pub async fn close_session(&self, request: CloseSessionRequest) -> Result<CloseSessionResponse, AcpClientError> {
125        self.request(request, false).await
126    }
127
128    /// Search the agent's prompt history through Aether's ACP extension.
129    pub async fn search_prompts(&self, params: PromptSearchParams) -> Result<PromptSearchResponse, AcpClientError> {
130        self.request(params, false).await
131    }
132
133    /// Load a session preview through Aether's ACP extension.
134    pub async fn preview_session(
135        &self,
136        params: SessionPreviewParams,
137    ) -> Result<SessionPreviewResponse, AcpClientError> {
138        self.request(params, false).await
139    }
140
141    /// List workspaces through Aether's ACP extension.
142    pub async fn list_workspaces(&self, params: WorkspaceListParams) -> Result<WorkspaceListResponse, AcpClientError> {
143        self.request(params, false).await
144    }
145
146    /// Move a session through Aether's ACP extension.
147    pub async fn move_workspace(&self, params: WorkspaceMoveParams) -> Result<WorkspaceMoveResponse, AcpClientError> {
148        self.request(params, false).await
149    }
150
151    pub async fn set_config_option(
152        &self,
153        request: SetSessionConfigOptionRequest,
154    ) -> Result<SetSessionConfigOptionResponse, AcpClientError> {
155        self.request(request, true).await
156    }
157
158    pub async fn authenticate(&self, request: AuthenticateRequest) -> Result<AuthenticateResponse, AcpClientError> {
159        self.request(request, true).await
160    }
161
162    pub async fn cancel(&self, request: CancelNotification) -> Result<(), AcpClientError> {
163        self.notify(request).await
164    }
165
166    pub async fn authenticate_mcp_server(&self, request: McpRequest) -> Result<(), AcpClientError> {
167        self.notify(request).await
168    }
169
170    async fn request<T>(&self, request: T, allow_during_prompt: bool) -> Result<T::Response, AcpClientError>
171    where
172        T: JsonRpcRequest + Send + 'static,
173        T::Response: Send,
174    {
175        let (response, receiver) = oneshot::channel();
176        self.send(ClientCommand::Request {
177            allow_during_prompt,
178            run: Box::new(move |cx| match cx {
179                Ok(cx) => send_typed_response(cx, request, response),
180                Err(error) => {
181                    let _ = response.send(Err(error));
182                }
183            }),
184        })?;
185        await_response(receiver).await
186    }
187
188    async fn notify<T>(&self, notification: T) -> Result<(), AcpClientError>
189    where
190        T: JsonRpcNotification + Send + 'static,
191    {
192        let (response, receiver) = oneshot::channel();
193        self.send(ClientCommand::Request {
194            allow_during_prompt: true,
195            run: Box::new(move |cx| {
196                let result = cx.and_then(|cx| cx.send_notification(notification).map_err(AcpClientError::Protocol));
197                let _ = response.send(result);
198            }),
199        })?;
200        await_response(receiver).await
201    }
202
203    fn send(&self, command: ClientCommand) -> Result<(), AcpClientError> {
204        self.cmd_tx.send(command).map_err(|_| AcpClientError::AgentCrashed("ACP task is no longer running".to_string()))
205    }
206}
207
208type InitializeResult = Result<InitializeResponse, AcpClientError>;
209type InitializeSender = Arc<Mutex<Option<oneshot::Sender<InitializeResult>>>>;
210type Response<T> = oneshot::Sender<Result<T, AcpClientError>>;
211type RequestFn = Box<dyn FnOnce(Result<&ConnectionTo<acp::Agent>, AcpClientError>) + Send>;
212
213enum ClientCommand {
214    Prompt { request: PromptRequest, response: Response<PromptResponse> },
215    LoadSession { request: LoadSessionRequest, response: Response<LoadedSession> },
216    Request { allow_during_prompt: bool, run: RequestFn },
217}
218
219struct ReplayState {
220    session_id: SessionId,
221    notifications: Vec<SessionNotification>,
222}
223
224async fn await_response<T>(receiver: oneshot::Receiver<Result<T, AcpClientError>>) -> Result<T, AcpClientError> {
225    receiver.await.map_err(|_| AcpClientError::AgentCrashed("ACP task ended before responding".to_string()))?
226}
227
228#[allow(clippy::too_many_lines)]
229async fn run_client_connection(
230    agent: impl ConnectTo<Client> + 'static,
231    event_tx: mpsc::UnboundedSender<AcpEvent>,
232    mut cmd_rx: mpsc::UnboundedReceiver<ClientCommand>,
233    init_tx: InitializeSender,
234    init_request: InitializeRequest,
235    replay_state: Arc<Mutex<Option<ReplayState>>>,
236) {
237    let connection_result = Client
238        .builder()
239        .on_receive_request(
240            async move |req: RequestPermissionRequest, responder, _cx| {
241                responder.respond(RequestPermissionResponse::new(RequestPermissionOutcome::Selected(
242                    SelectedPermissionOutcome::new(auto_approve_option(&req)),
243                )))
244            },
245            acp::on_receive_request!(),
246        )
247        .on_receive_request(
248            {
249                let event_tx = event_tx.clone();
250                async move |params: CreateElicitationRequest, responder, _cx| {
251                    if let Err(send_err) =
252                        event_tx.send(AcpEvent::ElicitationRequest { params: Box::new(params), responder })
253                        && let AcpEvent::ElicitationRequest { responder, .. } = send_err.0
254                    {
255                        return responder.respond_with_error(acp::Error::internal_error());
256                    }
257                    Ok(())
258                }
259            },
260            acp::on_receive_request!(),
261        )
262        .on_receive_notification(
263            {
264                let event_tx = event_tx.clone();
265                let replay_state = Arc::clone(&replay_state);
266                async move |notification: SessionNotification, _cx| {
267                    let passthrough = {
268                        let mut replay = replay_state.lock().expect("replay state lock poisoned");
269                        match replay.as_mut() {
270                            Some(state) if state.session_id == notification.session_id => {
271                                state.notifications.push(notification);
272                                None
273                            }
274                            _ => Some(notification),
275                        }
276                    };
277                    if let Some(SessionNotification { session_id, update, .. }) = passthrough {
278                        let _ = event_tx.send(AcpEvent::SessionUpdate { session_id, update: Box::new(update) });
279                    }
280                    Ok(())
281                }
282            },
283            acp::on_receive_notification!(),
284        )
285        .on_receive_notification(
286            {
287                let event_tx = event_tx.clone();
288                async move |params: ContextCompactionParams, _cx| {
289                    let _ = event_tx.send(AcpEvent::ContextCompaction(params));
290                    Ok(())
291                }
292            },
293            acp::on_receive_notification!(),
294        )
295        .on_receive_notification(
296            {
297                let event_tx = event_tx.clone();
298                async move |params: ContextClearedParams, _cx| {
299                    let _ = event_tx.send(AcpEvent::ContextCleared(params));
300                    Ok(())
301                }
302            },
303            acp::on_receive_notification!(),
304        )
305        .on_receive_notification(
306            {
307                let event_tx = event_tx.clone();
308                async move |params: SubAgentProgressParams, _cx| {
309                    let _ = event_tx.send(AcpEvent::SubAgentProgress(params));
310                    Ok(())
311                }
312            },
313            acp::on_receive_notification!(),
314        )
315        .on_receive_notification(
316            {
317                let event_tx = event_tx.clone();
318                async move |params: SessionUsageParams, _cx| {
319                    let _ = event_tx.send(AcpEvent::SessionUsage(Box::new(params)));
320                    Ok(())
321                }
322            },
323            acp::on_receive_notification!(),
324        )
325        .on_receive_notification(
326            {
327                let event_tx = event_tx.clone();
328                async move |params: AuthMethodsUpdatedParams, _cx| {
329                    let _ = event_tx.send(AcpEvent::AuthMethodsUpdated(params));
330                    Ok(())
331                }
332            },
333            acp::on_receive_notification!(),
334        )
335        .on_receive_notification(
336            {
337                let event_tx = event_tx.clone();
338                async move |params: McpNotification, _cx| {
339                    let _ = event_tx.send(AcpEvent::McpNotification(params));
340                    Ok(())
341                }
342            },
343            acp::on_receive_notification!(),
344        )
345        .connect_with(agent, {
346            let event_tx = event_tx.clone();
347            let init_tx = Arc::clone(&init_tx);
348            async move |cx: ConnectionTo<acp::Agent>| {
349                run_main(cx, event_tx, &mut cmd_rx, Arc::clone(&init_tx), init_request, replay_state).await;
350                Ok(())
351            }
352        })
353        .await;
354
355    if let Err(e) = connection_result {
356        tracing::warn!("ACP connection exited with error: {e:?}");
357        send_initialization(&init_tx, Err(AcpClientError::ConnectFailed(e)));
358    }
359    let _ = event_tx.send(AcpEvent::ConnectionClosed);
360}
361
362async fn run_main(
363    cx: ConnectionTo<acp::Agent>,
364    event_tx: mpsc::UnboundedSender<AcpEvent>,
365    cmd_rx: &mut mpsc::UnboundedReceiver<ClientCommand>,
366    init_tx: InitializeSender,
367    init_request: InitializeRequest,
368    replay_state: Arc<Mutex<Option<ReplayState>>>,
369) {
370    let init_resp = match cx.send_request(init_request).block_task().await {
371        Ok(response) => response,
372        Err(error) => {
373            send_initialization(&init_tx, Err(AcpClientError::Protocol(error)));
374            return;
375        }
376    };
377    info!("ACP initialized: protocol={:?}, agent_info={:?}", init_resp.protocol_version, init_resp.agent_info);
378    if !send_initialization(&init_tx, Ok(init_resp)) {
379        return;
380    }
381
382    while let Some(command) = cmd_rx.recv().await {
383        handle_command(&cx, &event_tx, command, ClientState::Idle, &replay_state, cmd_rx).await;
384    }
385}
386
387async fn run_prompt(
388    cx: &ConnectionTo<acp::Agent>,
389    event_tx: &mpsc::UnboundedSender<AcpEvent>,
390    cmd_rx: &mut mpsc::UnboundedReceiver<ClientCommand>,
391    replay_state: &Arc<Mutex<Option<ReplayState>>>,
392    request: PromptRequest,
393    response: Response<PromptResponse>,
394) {
395    let prompt_fut = cx.send_request(request).block_task();
396    tokio::pin!(prompt_fut);
397
398    loop {
399        tokio::select! {
400            result = &mut prompt_fut => {
401                match result {
402                    Ok(prompt_response) => {
403                        let _ = event_tx.send(AcpEvent::PromptCompleted(prompt_response.stop_reason));
404                        let _ = response.send(Ok(prompt_response));
405                    }
406                    Err(error) => {
407                        let _ = response.send(Err(AcpClientError::Protocol(error)));
408                    }
409                }
410                break;
411            }
412            Some(command) = cmd_rx.recv() => {
413                Box::pin(handle_command(cx, event_tx, command, ClientState::Prompting, replay_state, cmd_rx)).await;
414            }
415            else => break,
416        }
417    }
418}
419
420fn send_initialization(sender: &InitializeSender, result: InitializeResult) -> bool {
421    sender.lock().expect("initialization lock poisoned").take().is_some_and(|sender| sender.send(result).is_ok())
422}
423
424#[derive(Clone, Copy, PartialEq, Eq)]
425enum ClientState {
426    Idle,
427    Prompting,
428}
429
430async fn handle_command(
431    cx: &ConnectionTo<acp::Agent>,
432    event_tx: &mpsc::UnboundedSender<AcpEvent>,
433    command: ClientCommand,
434    state: ClientState,
435    replay_state: &Arc<Mutex<Option<ReplayState>>>,
436    cmd_rx: &mut mpsc::UnboundedReceiver<ClientCommand>,
437) {
438    match command {
439        ClientCommand::Prompt { request, response } => {
440            if state == ClientState::Prompting {
441                let _ = response.send(Err(AcpClientError::Busy));
442            } else {
443                Box::pin(run_prompt(cx, event_tx, cmd_rx, replay_state, request, response)).await;
444            }
445        }
446        ClientCommand::LoadSession { request, response } => {
447            if state == ClientState::Prompting {
448                let _ = response.send(Err(AcpClientError::Busy));
449                return;
450            }
451            let session_id = request.session_id.clone();
452            *replay_state.lock().expect("replay state lock poisoned") =
453                Some(ReplayState { session_id: session_id.clone(), notifications: vec![] });
454            let result = cx.send_request(request).block_task().await.map_err(AcpClientError::Protocol);
455            let replay = replay_state
456                .lock()
457                .expect("replay state lock poisoned")
458                .take()
459                .map_or_else(Vec::new, |state| state.notifications);
460            let _ = response.send(result.map(|response| LoadedSession { session_id, response, replay }));
461        }
462        ClientCommand::Request { allow_during_prompt, run } => {
463            if state == ClientState::Prompting && !allow_during_prompt {
464                run(Err(AcpClientError::Busy));
465            } else {
466                run(Ok(cx));
467            }
468        }
469    }
470}
471
472fn send_typed_response<T: JsonRpcRequest + 'static>(
473    cx: &ConnectionTo<acp::Agent>,
474    request: T,
475    response: Response<T::Response>,
476) {
477    let request = cx.send_request(request).block_task();
478    if let Err(error) = cx.spawn(async move {
479        let result = request.await.map_err(AcpClientError::Protocol);
480        let _ = response.send(result);
481        Ok(())
482    }) {
483        tracing::warn!("failed to spawn ACP request: {error:?}");
484    }
485}
486
487fn auto_approve_option(req: &RequestPermissionRequest) -> PermissionOptionId {
488    debug_assert!(!req.options.is_empty(), "ACP guarantees at least one permission option");
489    req.options
490        .iter()
491        .find(|option| matches!(option.kind, PermissionOptionKind::AllowOnce | PermissionOptionKind::AllowAlways))
492        .map_or_else(|| req.options[0].option_id.clone(), |option| option.option_id.clone())
493}
acp_utils/client/session.rs

acp_utils/client/
session.rs
