acp_utils/client/
session.rs

1use super::error::AcpClientError;
2use super::event::AcpEvent;
3use crate::notifications::{AuthMethodsUpdatedParams, ContextClearedParams, McpNotification, SubAgentProgressParams};
4use agent_client_protocol::schema::v2::{
5    AuthMethod, CancelSessionNotification, CreateElicitationRequest, InitializeRequest, InitializeResponse,
6    NewSessionRequest, NewSessionResponse, PermissionOptionId, PermissionOptionKind, PromptCapabilities, PromptRequest,
7    PromptResponse, ReplayFrom, ReplayFromStart, RequestPermissionOutcome, RequestPermissionRequest,
8    RequestPermissionResponse, ResumeSessionRequest, ResumeSessionResponse, SelectedPermissionOutcome,
9    SessionCapabilities, UpdateSessionNotification,
10};
11use agent_client_protocol::util::MatchDispatchFrom;
12use agent_client_protocol::{
13    self as acp, Client, ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, V2ConnectionTo,
14};
15use std::future::Future;
16use std::sync::Arc;
17use tokio::sync::{Mutex, mpsc, oneshot};
18use tokio::task::JoinHandle;
19use tracing::info;
20
21#[derive(Clone)]
22pub struct AcpClientHandle {
23    cx: V2ConnectionTo<acp::Agent>,
24    connection: Arc<ClientConnection>,
25}
26
27pub struct AcpClient {
28    pub initialize_response: InitializeResponse,
29    pub event_rx: mpsc::UnboundedReceiver<AcpEvent>,
30    pub handle: AcpClientHandle,
31}
32
33/// Connect to an ACP agent and complete initialization without creating a session.
34pub async fn connect_acp_client(
35    agent: impl ConnectTo<Client> + 'static,
36    init_request: InitializeRequest,
37) -> Result<AcpClient, AcpClientError> {
38    let (event_tx, event_rx) = mpsc::unbounded_channel();
39    let (init_tx, init_rx) = oneshot::channel();
40    let events = ConnectionEvents(event_tx);
41    let driver = tokio::spawn(run_client_connection(agent, init_request, init_tx, events));
42    let connection = Arc::new(ClientConnection { driver: Mutex::new(Some(driver)) });
43    let (initialize_response, cx) = await_response(init_rx).await?;
44    Ok(AcpClient { initialize_response, event_rx, handle: AcpClientHandle { cx, connection } })
45}
46
47impl AcpClient {
48    /// The agent's display title, falling back to its implementation name.
49    pub fn agent_name(&self) -> String {
50        let info = &self.initialize_response.info;
51        info.title.as_deref().unwrap_or(&info.name).to_string()
52    }
53
54    pub fn prompt_capabilities(&self) -> Option<&PromptCapabilities> {
55        self.session_capabilities().and_then(|session| session.prompt.as_ref())
56    }
57
58    pub fn session_capabilities(&self) -> Option<&SessionCapabilities> {
59        self.initialize_response.capabilities.session.as_ref()
60    }
61
62    pub fn auth_methods(&self) -> &[AuthMethod] {
63        &self.initialize_response.auth_methods
64    }
65}
66
67impl AcpClientHandle {
68    /// Stop this connection and join its driver without sending session/cancel or session/close.
69    pub async fn disconnect(&self) {
70        let mut driver = self.connection.driver.lock().await;
71        if let Some(task) = driver.as_mut() {
72            task.abort();
73            let _ = task.await;
74        }
75        *driver = None;
76    }
77
78    pub fn prompt(
79        &self,
80        request: PromptRequest,
81    ) -> impl Future<Output = Result<PromptResponse, AcpClientError>> + Send + use<> {
82        self.request(request)
83    }
84
85    /// Request history as ordinary session updates, preceding the resume response.
86    pub fn resume_session_with_replay(
87        &self,
88        request: ResumeSessionRequest,
89    ) -> impl Future<Output = Result<ResumeSessionResponse, AcpClientError>> + Send + use<> {
90        self.request(request.replay_from(ReplayFrom::Start(ReplayFromStart::new())))
91    }
92
93    pub fn new_session(
94        &self,
95        request: NewSessionRequest,
96    ) -> impl Future<Output = Result<NewSessionResponse, AcpClientError>> + Send + use<> {
97        self.request(request)
98    }
99
100    /// Send immediately; polling the returned future only waits for the response.
101    pub fn request<R: acp::JsonRpcRequest>(
102        &self,
103        request: R,
104    ) -> impl Future<Output = Result<R::Response, AcpClientError>> + Send + use<R> {
105        let sent = self.cx.send_request(request);
106        async move { sent.block_task().await.map_err(AcpClientError::Protocol) }
107    }
108
109    /// Resume a session without requesting history.
110    pub fn resume_session(
111        &self,
112        request: ResumeSessionRequest,
113    ) -> impl Future<Output = Result<ResumeSessionResponse, AcpClientError>> + Send + use<> {
114        self.request(request.replay_from(None))
115    }
116
117    pub fn cancel(&self, request: CancelSessionNotification) -> Result<(), AcpClientError> {
118        self.notify(request)
119    }
120
121    pub fn notify(&self, request: impl acp::JsonRpcNotification) -> Result<(), AcpClientError> {
122        self.cx.send_notification(request).map_err(AcpClientError::Protocol)
123    }
124}
125
126struct ClientConnection {
127    driver: Mutex<Option<JoinHandle<()>>>,
128}
129
130impl Drop for ClientConnection {
131    fn drop(&mut self) {
132        if let Some(driver) = self.driver.get_mut() {
133            driver.abort();
134        }
135    }
136}
137
138struct ConnectionEvents(mpsc::UnboundedSender<AcpEvent>);
139
140impl Drop for ConnectionEvents {
141    fn drop(&mut self) {
142        let _ = self.0.send(AcpEvent::ConnectionClosed);
143    }
144}
145
146async fn run_client_connection(
147    agent: impl ConnectTo<Client> + 'static,
148    init_request: InitializeRequest,
149    init_tx: oneshot::Sender<Result<(InitializeResponse, V2ConnectionTo<acp::Agent>), AcpClientError>>,
150    events: ConnectionEvents,
151) {
152    let connection_result = Client
153        .v2()
154        .name("wisp")
155        .with_handler(ClientHandlers(events.0.clone()))
156        .connect_with(agent, async move |cx: V2ConnectionTo<acp::Agent>| {
157            let result = cx.send_request(init_request).block_task().await.map_err(AcpClientError::Protocol);
158            let _ = init_tx.send(result.map(|response| {
159                info!("ACP initialized: protocol={:?}, agent_info={:?}", response.protocol_version, response.info);
160                (response, cx.clone())
161            }));
162            cx.incoming_closed().await;
163            Ok(())
164        })
165        .await;
166    if let Err(error) = connection_result {
167        tracing::warn!("ACP connection exited with error: {error:?}");
168    }
169}
170
171struct ClientHandlers(mpsc::UnboundedSender<AcpEvent>);
172
173impl HandleDispatchFrom<acp::Agent> for ClientHandlers {
174    async fn handle_dispatch_from(
175        &mut self,
176        message: Dispatch,
177        cx: ConnectionTo<acp::Agent>,
178    ) -> Result<Handled<Dispatch>, acp::Error> {
179        let emit = |event| {
180            if let Err(error) = self.0.send(event)
181                && let AcpEvent::ElicitationRequest { responder, .. } = error.0
182            {
183                let _ = responder.respond_with_error(acp::Error::internal_error());
184            }
185            Ok::<_, acp::Error>(())
186        };
187        MatchDispatchFrom::new(message, &cx)
188            .if_request(async |request: RequestPermissionRequest, responder| {
189                let outcome = auto_approve_option(&request).map_or(RequestPermissionOutcome::Cancelled, |option| {
190                    RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(option))
191                });
192                let _ = responder.respond(RequestPermissionResponse::new(outcome));
193                Ok(())
194            })
195            .await
196            .if_request(async |params: CreateElicitationRequest, responder| {
197                emit(AcpEvent::ElicitationRequest { params: Box::new(params), responder })
198            })
199            .await
200            .if_notification(async |params: UpdateSessionNotification| emit(params.into()))
201            .await
202            .if_notification(async |params: ContextClearedParams| emit(AcpEvent::ContextCleared(params)))
203            .await
204            .if_notification(async |params: SubAgentProgressParams| emit(AcpEvent::SubAgentProgress(params)))
205            .await
206            .if_notification(async |params: McpNotification| emit(AcpEvent::McpNotification(params)))
207            .await
208            .if_notification(async |params: AuthMethodsUpdatedParams| emit(AcpEvent::AuthMethodsUpdated(params)))
209            .await
210            .done()
211    }
212
213    fn describe_chain(&self) -> impl std::fmt::Debug {
214        "ClientHandlers"
215    }
216}
217
218async fn await_response<T>(receiver: oneshot::Receiver<Result<T, AcpClientError>>) -> Result<T, AcpClientError> {
219    receiver.await.map_err(|_| AcpClientError::AgentCrashed("ACP task ended before responding".to_string()))?
220}
221
222fn auto_approve_option(req: &RequestPermissionRequest) -> Option<PermissionOptionId> {
223    req.options
224        .iter()
225        .find(|option| matches!(option.kind, PermissionOptionKind::AllowOnce | PermissionOptionKind::AllowAlways))
226        .or_else(|| req.options.first())
227        .map(|option| option.option_id.clone())
228}
acp_utils/client/session.rs

acp_utils/client/
session.rs
