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::v2::{
7    AuthMethod, CancelSessionNotification, CreateElicitationRequest, InitializeRequest, InitializeResponse,
8    NewSessionRequest, NewSessionResponse, PermissionOptionId, PermissionOptionKind, PromptCapabilities, PromptRequest,
9    PromptResponse, ReplayFrom, ReplayFromStart, RequestPermissionOutcome, RequestPermissionRequest,
10    RequestPermissionResponse, ResumeSessionRequest, ResumeSessionResponse, SelectedPermissionOutcome,
11    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 std::future::Future;
18use std::sync::Arc;
19use tokio::sync::{mpsc, oneshot};
20use tokio::task::JoinHandle;
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 = tokio::spawn(run_client_connection(agent, init_request, init_tx, events));
46    let connection = Arc::new(ClientConnection { driver, closed });
47    let (initialize_response, cx) = await_response(init_rx).await?;
48    Ok(AcpClient { initialize_response, event_rx, handle: AcpClientHandle { cx, connection } })
49}
50
51impl AcpClient {
52    /// The agent's display title, falling back to its implementation name.
53    pub fn agent_name(&self) -> String {
54        let info = &self.initialize_response.info;
55        info.title.as_deref().unwrap_or(&info.name).to_string()
56    }
57
58    pub fn prompt_capabilities(&self) -> Option<&PromptCapabilities> {
59        self.session_capabilities().and_then(|session| session.prompt.as_ref())
60    }
61
62    pub fn session_capabilities(&self) -> Option<&SessionCapabilities> {
63        self.initialize_response.capabilities.session.as_ref()
64    }
65
66    pub fn auth_methods(&self) -> &[AuthMethod] {
67        &self.initialize_response.auth_methods
68    }
69}
70
71impl AcpClientHandle {
72    /// Stop this connection and wait for it to close without sending session/cancel or session/close.
73    pub async fn disconnect(&self) {
74        self.connection.driver.abort();
75        self.connection.closed.cancelled().await;
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: JoinHandle<()>,
128    closed: CancellationToken,
129}
130
131impl Drop for ClientConnection {
132    fn drop(&mut self) {
133        self.driver.abort();
134    }
135}
136
137struct ConnectionEvents {
138    event_tx: mpsc::UnboundedSender<AcpEvent>,
139    closed: CancellationToken,
140}
141
142impl Drop for ConnectionEvents {
143    fn drop(&mut self) {
144        let _ = self.event_tx.send(AcpEvent::ConnectionClosed);
145        self.closed.cancel();
146    }
147}
148
149async fn run_client_connection(
150    agent: impl ConnectTo<Client> + 'static,
151    init_request: InitializeRequest,
152    init_tx: oneshot::Sender<Result<(InitializeResponse, V2ConnectionTo<acp::Agent>), AcpClientError>>,
153    events: ConnectionEvents,
154) {
155    let connection_result = Client
156        .v2()
157        .name("wisp")
158        .with_handler(ClientHandlers(events.event_tx.clone()))
159        .connect_with(agent, async move |cx: V2ConnectionTo<acp::Agent>| {
160            let result = cx.send_request(init_request).block_task().await.map_err(AcpClientError::Protocol);
161            let _ = init_tx.send(result.map(|response| {
162                info!("ACP initialized: protocol={:?}, agent_info={:?}", response.protocol_version, response.info);
163                (response, cx.clone())
164            }));
165            cx.incoming_closed().await;
166            Ok(())
167        })
168        .await;
169    if let Err(error) = connection_result {
170        tracing::warn!("ACP connection exited with error: {error:?}");
171    }
172}
173
174struct ClientHandlers(mpsc::UnboundedSender<AcpEvent>);
175
176impl HandleDispatchFrom<acp::Agent> for ClientHandlers {
177    async fn handle_dispatch_from(
178        &mut self,
179        message: Dispatch,
180        cx: ConnectionTo<acp::Agent>,
181    ) -> Result<Handled<Dispatch>, acp::Error> {
182        let emit = |event| {
183            if let Err(error) = self.0.send(event)
184                && let AcpEvent::ElicitationRequest { responder, .. } = error.0
185            {
186                let _ = responder.respond_with_error(acp::Error::internal_error());
187            }
188            Ok::<_, acp::Error>(())
189        };
190        MatchDispatchFrom::new(message, &cx)
191            .if_request(async |request: RequestPermissionRequest, responder| {
192                let outcome = auto_approve_option(&request).map_or(RequestPermissionOutcome::Cancelled, |option| {
193                    RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(option))
194                });
195                let _ = responder.respond(RequestPermissionResponse::new(outcome));
196                Ok(())
197            })
198            .await
199            .if_request(async |params: CreateElicitationRequest, responder| {
200                emit(AcpEvent::ElicitationRequest { params: Box::new(params), responder })
201            })
202            .await
203            .if_notification(async |params: UpdateSessionNotification| emit(params.into()))
204            .await
205            .if_notification(async |params: ContextClearedParams| emit(AcpEvent::ContextCleared(params)))
206            .await
207            .if_notification(async |params: SubAgentProgressParams| emit(AcpEvent::SubAgentProgress(params)))
208            .await
209            .if_notification(async |params: GitDiffEventPayload| emit(AcpEvent::GitDiffEvent(params)))
210            .await
211            .if_notification(async |params: McpNotification| emit(AcpEvent::McpNotification(params)))
212            .await
213            .if_notification(async |params: AuthMethodsUpdatedParams| emit(AcpEvent::AuthMethodsUpdated(params)))
214            .await
215            .done()
216    }
217
218    fn describe_chain(&self) -> impl std::fmt::Debug {
219        "ClientHandlers"
220    }
221}
222
223async fn await_response<T>(receiver: oneshot::Receiver<Result<T, AcpClientError>>) -> Result<T, AcpClientError> {
224    receiver.await.map_err(|_| AcpClientError::AgentCrashed("ACP task ended before responding".to_string()))?
225}
226
227fn auto_approve_option(req: &RequestPermissionRequest) -> Option<PermissionOptionId> {
228    req.options
229        .iter()
230        .find(|option| matches!(option.kind, PermissionOptionKind::AllowOnce | PermissionOptionKind::AllowAlways))
231        .or_else(|| req.options.first())
232        .map(|option| option.option_id.clone())
233}
acp_utils/client/session.rs

acp_utils/client/
session.rs
