acp_utils/client/
session.rs

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

acp_utils/client/
session.rs
