aether-acp-utils 0.5.11

use super::error::AcpClientError;
use super::event::AcpEvent;
use crate::notifications::{
    AuthMethodsUpdatedParams, ContextClearedParams, GitDiffEventPayload, McpNotification, SubAgentProgressParams,
};
use agent_client_protocol::schema::ProtocolVersion;
use agent_client_protocol::schema::v2::{
    AuthMethod, CancelSessionNotification, ClientCapabilities, CreateElicitationRequest, ElicitationCapabilities,
    ElicitationFormCapabilities, ElicitationUrlCapabilities, Implementation, InitializeRequest, InitializeResponse,
    NewSessionRequest, NewSessionResponse, PermissionOptionId, PermissionOptionKind, PromptCapabilities, PromptRequest,
    PromptResponse, RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse,
    ResumeSessionRequest, ResumeSessionResponse, SelectedPermissionOutcome, SessionCapabilities,
    UpdateSessionNotification,
};
use agent_client_protocol::util::MatchDispatchFrom;
use agent_client_protocol::{
    self as acp, Client, ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, V2ConnectionTo,
};
use futures::future::{AbortHandle, abortable};
use std::future::Future;
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use tracing::info;

#[derive(Clone)]
pub struct AcpClientHandle {
    cx: V2ConnectionTo<acp::Agent>,
    connection: Arc<ClientConnection>,
}

pub struct AcpClient {
    pub initialize_response: InitializeResponse,
    pub event_rx: mpsc::UnboundedReceiver<AcpEvent>,
    pub handle: AcpClientHandle,
}

/// Connect to an ACP agent and complete initialization without creating a session.
pub async fn connect_acp_client(
    agent: impl ConnectTo<Client> + 'static,
    init_request: InitializeRequest,
) -> Result<AcpClient, AcpClientError> {
    let (event_tx, event_rx) = mpsc::unbounded_channel();
    let (init_tx, init_rx) = oneshot::channel();
    let closed = CancellationToken::new();
    let events = ConnectionEvents { event_tx, closed: closed.clone() };
    let (driver, abort) = abortable(run_client_connection(agent, init_request, init_tx, events));
    spawn(async move {
        let _ = driver.await;
    });
    let connection = Arc::new(ClientConnection { abort, closed });
    let (initialize_response, cx) = await_response(init_rx).await?;
    Ok(AcpClient { initialize_response, event_rx, handle: AcpClientHandle { cx, connection } })
}

pub fn initialize_request(info: Implementation) -> InitializeRequest {
    InitializeRequest::new(ProtocolVersion::V2, info).capabilities(ClientCapabilities::new().elicitation(
        ElicitationCapabilities::new().form(ElicitationFormCapabilities::new()).url(ElicitationUrlCapabilities::new()),
    ))
}

impl AcpClient {
    /// The agent's display title, falling back to its implementation name.
    pub fn agent_name(&self) -> String {
        let info = &self.initialize_response.info;
        info.title.as_deref().unwrap_or(&info.name).to_string()
    }

    pub fn prompt_capabilities(&self) -> Option<&PromptCapabilities> {
        self.session_capabilities().and_then(|session| session.prompt.as_ref())
    }

    pub fn session_capabilities(&self) -> Option<&SessionCapabilities> {
        self.initialize_response.capabilities.session.as_ref()
    }

    pub fn auth_methods(&self) -> &[AuthMethod] {
        &self.initialize_response.auth_methods
    }
}

impl AcpClientHandle {
    /// Stop this connection and wait for it to close without sending session/cancel or session/close.
    pub async fn disconnect(&self) {
        self.connection.abort.abort();
        self.connection.closed.cancelled().await;
    }

    pub fn prompt(
        &self,
        request: PromptRequest,
    ) -> impl Future<Output = Result<PromptResponse, AcpClientError>> + Send + use<> {
        self.request(request)
    }

    pub fn new_session(
        &self,
        request: NewSessionRequest,
    ) -> impl Future<Output = Result<NewSessionResponse, AcpClientError>> + Send + use<> {
        self.request(request)
    }

    /// Send immediately; polling the returned future only waits for the response.
    pub fn request<R: acp::JsonRpcRequest>(
        &self,
        request: R,
    ) -> impl Future<Output = Result<R::Response, AcpClientError>> + Send + use<R> {
        let sent = self.cx.send_request(request);
        async move { sent.block_task().await.map_err(AcpClientError::Protocol) }
    }

    pub fn resume_session(
        &self,
        request: ResumeSessionRequest,
    ) -> impl Future<Output = Result<ResumeSessionResponse, AcpClientError>> + Send + use<> {
        self.request(request)
    }

    pub fn cancel(&self, request: CancelSessionNotification) -> Result<(), AcpClientError> {
        self.notify(request)
    }

    pub fn notify(&self, request: impl acp::JsonRpcNotification) -> Result<(), AcpClientError> {
        self.cx.send_notification(request).map_err(AcpClientError::Protocol)
    }
}

struct ClientConnection {
    abort: AbortHandle,
    closed: CancellationToken,
}

impl Drop for ClientConnection {
    fn drop(&mut self) {
        self.abort.abort();
    }
}

struct ConnectionEvents {
    event_tx: mpsc::UnboundedSender<AcpEvent>,
    closed: CancellationToken,
}

impl Drop for ConnectionEvents {
    fn drop(&mut self) {
        let _ = self.event_tx.send(AcpEvent::ConnectionClosed);
        self.closed.cancel();
    }
}

async fn run_client_connection(
    agent: impl ConnectTo<Client> + 'static,
    init_request: InitializeRequest,
    init_tx: oneshot::Sender<Result<(InitializeResponse, V2ConnectionTo<acp::Agent>), AcpClientError>>,
    events: ConnectionEvents,
) {
    let connection_result = Client
        .v2()
        .name(&init_request.info.name)
        .with_handler(ClientHandlers(events.event_tx.clone()))
        .connect_with(agent, async move |cx: V2ConnectionTo<acp::Agent>| {
            let result = cx.send_request(init_request).block_task().await.map_err(AcpClientError::Protocol);
            let _ = init_tx.send(result.map(|response| {
                info!("ACP initialized: protocol={:?}, agent_info={:?}", response.protocol_version, response.info);
                (response, cx.clone())
            }));
            cx.incoming_closed().await;
            Ok(())
        })
        .await;
    if let Err(error) = connection_result {
        tracing::warn!("ACP connection exited with error: {error:?}");
    }
}

struct ClientHandlers(mpsc::UnboundedSender<AcpEvent>);

impl HandleDispatchFrom<acp::Agent> for ClientHandlers {
    async fn handle_dispatch_from(
        &mut self,
        message: Dispatch,
        cx: ConnectionTo<acp::Agent>,
    ) -> Result<Handled<Dispatch>, acp::Error> {
        let emit = |event| {
            if let Err(error) = self.0.send(event)
                && let AcpEvent::ElicitationRequest { responder, .. } = error.0
            {
                let _ = responder.respond_with_error(acp::Error::internal_error());
            }
            Ok::<_, acp::Error>(())
        };
        MatchDispatchFrom::new(message, &cx)
            .if_request(async |request: RequestPermissionRequest, responder| {
                let outcome = auto_approve_option(&request).map_or(RequestPermissionOutcome::Cancelled, |option| {
                    RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(option))
                });
                let _ = responder.respond(RequestPermissionResponse::new(outcome));
                Ok(())
            })
            .await
            .if_request(async |params: CreateElicitationRequest, responder| {
                emit(AcpEvent::ElicitationRequest { params: Box::new(params), responder })
            })
            .await
            .if_notification(async |params: UpdateSessionNotification| emit(params.into()))
            .await
            .if_notification(async |params: ContextClearedParams| emit(AcpEvent::ContextCleared(params)))
            .await
            .if_notification(async |params: SubAgentProgressParams| emit(AcpEvent::SubAgentProgress(params)))
            .await
            .if_notification(async |params: GitDiffEventPayload| emit(AcpEvent::GitDiffEvent(params)))
            .await
            .if_notification(async |params: McpNotification| emit(AcpEvent::McpNotification(params)))
            .await
            .if_notification(async |params: AuthMethodsUpdatedParams| emit(AcpEvent::AuthMethodsUpdated(params)))
            .await
            .done()
    }

    fn describe_chain(&self) -> impl std::fmt::Debug {
        "ClientHandlers"
    }
}

#[cfg(not(target_family = "wasm"))]
fn spawn(future: impl Future<Output = ()> + Send + 'static) {
    tokio::spawn(future);
}

#[cfg(target_family = "wasm")]
fn spawn(future: impl Future<Output = ()> + 'static) {
    wasm_bindgen_futures::spawn_local(future);
}

async fn await_response<T>(receiver: oneshot::Receiver<Result<T, AcpClientError>>) -> Result<T, AcpClientError> {
    receiver.await.map_err(|_| AcpClientError::AgentCrashed("ACP task ended before responding".to_string()))?
}

fn auto_approve_option(req: &RequestPermissionRequest) -> Option<PermissionOptionId> {
    req.options
        .iter()
        .find(|option| matches!(option.kind, PermissionOptionKind::AllowOnce | PermissionOptionKind::AllowAlways))
        .or_else(|| req.options.first())
        .map(|option| option.option_id.clone())
}