@corbet-labs/ccht 0.1.0

package/source/src/native/session.rs

@@ -0,0 +1,541 @@
+use std::collections::BTreeMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+
+use agent_client_protocol::util::MatchDispatch;
+use agent_client_protocol::{ActiveSession, Agent, Dispatch, Responder, SessionMessage};
+use tokio::sync::{mpsc, oneshot, watch};
+use tokio::task::JoinHandle;
+use tokio::time::{Instant, interval, timeout};
+
+use crate::acp::{
+    CancelNotification, CloseSessionRequest, ContentBlock, PromptRequest, PromptResponse,
+    RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse, SessionId,
+    SessionNotification, StopReason,
+};
+use crate::{Event, Prompt, SessionEvent};
+
+use super::client::{ClientLease, ClientState};
+use super::{NativeError, PermissionDecision, PermissionPolicy, Result, lock};
+
+static NEXT_PERMISSION: AtomicU64 = AtomicU64::new(1);
+
+/// A native session and its single ordered event stream.
+///
+/// Take the receiver once, then retain a handle while consuming it. A cloned
+/// handle keeps the session alive independently of this value.
+pub struct NativeSession {
+    handle: SessionHandle,
+    events: Option<mpsc::Receiver<SessionEvent>>,
+}
+
+/// Cloneable, thread-safe control of a native conversation.
+#[derive(Clone)]
+pub struct SessionHandle {
+    lease: Arc<SessionLease>,
+}
+
+struct SessionLease {
+    state: Arc<SessionState>,
+    commands: mpsc::Sender<PromptCommand>,
+    client: Arc<ClientLease>,
+}
+
+struct SessionState {
+    id: SessionId,
+    client: Arc<ClientState>,
+    connection: agent_client_protocol::ConnectionTo<Agent>,
+    turn: Mutex<TurnState>,
+    activity: watch::Sender<bool>,
+    closed: watch::Sender<bool>,
+    failure: Mutex<Option<NativeError>>,
+    events: Mutex<Option<mpsc::Sender<SessionEvent>>>,
+    permissions: Mutex<BTreeMap<String, PendingPermission>>,
+    task: Mutex<Option<JoinHandle<()>>>,
+}
+
+#[derive(Default)]
+struct TurnState {
+    reserved: Option<String>,
+    current: Option<String>,
+    cancelled: bool,
+}
+
+struct PendingPermission {
+    request: RequestPermissionRequest,
+    responder: Responder<RequestPermissionResponse>,
+    expires: Instant,
+}
+
+struct PromptCommand {
+    prompt: Prompt,
+    result: oneshot::Sender<Result<StopReason>>,
+}
+
+struct PromptGuard(Option<(Arc<SessionState>, String)>);
+
+impl Drop for PromptGuard {
+    fn drop(&mut self) {
+        let Some((state, request_id)) = &self.0 else {
+            return;
+        };
+        // Release the turn lock before cleanup takes it again.
+        let active = lock(&state.turn).reserved.as_ref() == Some(request_id);
+        if active {
+            // Abandoning a prompt cannot leave unattended agent work running.
+            state.close(NativeError::Closed, true);
+        }
+    }
+}
+
+impl Drop for SessionLease {
+    fn drop(&mut self) {
+        let active = lock(&self.state.turn).reserved.is_some();
+        self.state.close(NativeError::Closed, active);
+    }
+}
+
+impl SessionState {
+    fn is_closed(&self) -> bool {
+        *self.closed.borrow() || self.client.is_closed()
+    }
+
+    fn error(&self) -> NativeError {
+        lock(&self.failure)
+            .clone()
+            .unwrap_or_else(|| self.client.error())
+    }
+
+    fn emit(&self, event: Event) -> Result<()> {
+        let request_id = lock(&self.turn).current.clone();
+        let result = lock(&self.events)
+            .as_ref()
+            .ok_or(NativeError::Closed)?
+            .try_send(SessionEvent { request_id, event });
+        match result {
+            Ok(()) => Ok(()),
+            Err(mpsc::error::TrySendError::Full(_)) => {
+                self.close(NativeError::EventBufferFull, true);
+                Err(NativeError::EventBufferFull)
+            }
+            Err(mpsc::error::TrySendError::Closed(_)) => {
+                self.close(NativeError::Closed, true);
+                Err(NativeError::Closed)
+            }
+        }
+    }
+
+    fn deny_permissions(&self) {
+        for (_, pending) in std::mem::take(&mut *lock(&self.permissions)) {
+            let _ = pending.responder.respond(RequestPermissionResponse::new(
+                RequestPermissionOutcome::Cancelled,
+            ));
+        }
+    }
+
+    fn close(&self, error: NativeError, stop_connection: bool) {
+        if !*self.closed.borrow() {
+            *lock(&self.failure) = Some(error.clone());
+            let _ = self
+                .connection
+                .send_notification(CancelNotification::new(self.id.clone()));
+            self.deny_permissions();
+            let request_id = lock(&self.turn).current.clone();
+            if let Some(events) = lock(&self.events).take()
+                && (request_id.is_some() || error != NativeError::Closed)
+            {
+                let _ = events.try_send(SessionEvent {
+                    request_id,
+                    event: Event::Error {
+                        code: error.code().into(),
+                        message: error.to_string(),
+                    },
+                });
+            }
+            self.closed.send_replace(true);
+            self.activity.send_replace(false);
+            if let Some(task) = lock(&self.task).as_ref() {
+                task.abort();
+            }
+        }
+        if stop_connection {
+            self.client.abort();
+        }
+    }
+
+    fn finish(&self, outcome: &Result<StopReason>) -> Result<()> {
+        self.deny_permissions();
+        let event = match outcome {
+            Ok(stop_reason) => Event::Completed {
+                stop_reason: *stop_reason,
+            },
+            Err(error) => Event::Error {
+                code: error.code().into(),
+                message: error.to_string(),
+            },
+        };
+        self.emit(event)?;
+        *lock(&self.turn) = TurnState::default();
+        self.activity.send_replace(false);
+        Ok(())
+    }
+}
+
+impl NativeSession {
+    pub(super) fn start(
+        client: Arc<ClientLease>,
+        session: ActiveSession<'static, Agent>,
+    ) -> Result<Self> {
+        let (events_tx, events_rx) = mpsc::channel(client.options.event_capacity);
+        let (commands_tx, commands_rx) = mpsc::channel(1);
+        let (activity, _) = watch::channel(false);
+        let (closed, _) = watch::channel(false);
+        let state = Arc::new(SessionState {
+            id: session.session_id().clone(),
+            client: client.state.clone(),
+            connection: session.connection().clone(),
+            turn: Mutex::new(TurnState::default()),
+            activity,
+            closed,
+            failure: Mutex::new(None),
+            events: Mutex::new(Some(events_tx)),
+            permissions: Mutex::new(BTreeMap::new()),
+            task: Mutex::new(None),
+        });
+        let worker_state = state.clone();
+        let options = client.options.clone();
+        let task = tokio::spawn(async move {
+            if let Err(error) =
+                run_session(session, commands_rx, worker_state.clone(), options).await
+            {
+                worker_state.close(error, true);
+            }
+        });
+        *lock(&state.task) = Some(task);
+        Ok(Self {
+            handle: SessionHandle {
+                lease: Arc::new(SessionLease {
+                    state,
+                    commands: commands_tx,
+                    client,
+                }),
+            },
+            events: Some(events_rx),
+        })
+    }
+
+    /// Clone the control handle; it can be used while the receiver is polled elsewhere.
+    pub fn handle(&self) -> SessionHandle {
+        self.handle.clone()
+    }
+
+    /// Take the sole stream, including any history queued during session restoration.
+    pub fn take_events(&mut self) -> Result<mpsc::Receiver<SessionEvent>> {
+        self.events.take().ok_or(NativeError::EventsTaken)
+    }
+}
+
+impl SessionHandle {
+    /// The agent-owned session identifier, suitable for explicit load or resume.
+    pub fn id(&self) -> &str {
+        self.lease.state.id.0.as_ref()
+    }
+
+    /// Send one prompt and await its terminal result while events stream independently.
+    ///
+    /// A concurrent prompt fails with `Busy`. Dropping this future after submission
+    /// closes the connection, ensuring the abandoned turn cannot keep running.
+    pub async fn prompt(&self, prompt: Prompt) -> Result<StopReason> {
+        let state = &self.lease.state;
+        if state.is_closed() {
+            return Err(state.error());
+        }
+        if prompt.request_id.is_empty() || prompt.content.is_empty() {
+            return Err(NativeError::InvalidOptions(
+                "a prompt needs a request id and content",
+            ));
+        }
+        let capabilities = &self
+            .lease
+            .client
+            .info
+            .agent_capabilities
+            .prompt_capabilities;
+        for content in &prompt.content {
+            match content {
+                ContentBlock::Text(_) | ContentBlock::ResourceLink(_) => {}
+                ContentBlock::Image(_) if capabilities.image => {}
+                ContentBlock::Audio(_) if capabilities.audio => {}
+                ContentBlock::Resource(_) if capabilities.embedded_context => {}
+                _ => return Err(NativeError::Unsupported("the requested prompt content")),
+            }
+        }
+        {
+            let mut turn = lock(&state.turn);
+            if turn.reserved.is_some() {
+                return Err(NativeError::Busy);
+            }
+            turn.reserved = Some(prompt.request_id.clone());
+            turn.cancelled = false;
+            state.activity.send_replace(true);
+        }
+        let mut guard = PromptGuard(Some((state.clone(), prompt.request_id.clone())));
+        let (result_tx, result_rx) = oneshot::channel();
+        self.lease
+            .commands
+            .send(PromptCommand {
+                prompt,
+                result: result_tx,
+            })
+            .await
+            .map_err(|_| state.error())?;
+        let mut closed = state.closed.subscribe();
+        let mut connection_closed = state.client.closed.subscribe();
+        let result = timeout(self.lease.client.options.prompt_timeout, async {
+            if state.is_closed() {
+                return Err(state.error());
+            }
+            tokio::select! {
+                result = result_rx => result.unwrap_or_else(|_| Err(state.error())),
+                _ = closed.changed() => Err(state.error()),
+                _ = connection_closed.changed() => Err(state.error()),
+            }
+        })
+        .await;
+        match result {
+            Ok(outcome) => {
+                guard.0 = None;
+                outcome
+            }
+            Err(_) => {
+                state.close(NativeError::Timeout, true);
+                guard.0 = None;
+                Err(NativeError::Timeout)
+            }
+        }
+    }
+
+    /// Request cancellation and wait a bounded time for the active turn to finish.
+    ///
+    /// If the agent ignores cancellation, the SDK connection is closed. This also
+    /// closes other sessions sharing that agent process.
+    pub async fn cancel(&self) -> Result<()> {
+        let state = &self.lease.state;
+        if state.is_closed() {
+            return Err(state.error());
+        }
+        {
+            let mut turn = lock(&state.turn);
+            if turn.reserved.is_none() {
+                return Ok(());
+            }
+            turn.cancelled = true;
+        }
+        state.deny_permissions();
+        state
+            .connection
+            .send_notification(CancelNotification::new(state.id.clone()))?;
+        let mut activity = state.activity.subscribe();
+        if timeout(self.lease.client.options.shutdown_timeout, async {
+            while *activity.borrow_and_update() {
+                if activity.changed().await.is_err() {
+                    break;
+                }
+            }
+        })
+        .await
+        .is_err()
+        {
+            state.close(NativeError::Timeout, true);
+            return Err(NativeError::Timeout);
+        }
+        if state.is_closed() {
+            Err(state.error())
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Answer one pending permission request using an option the agent offered.
+    ///
+    /// Invalid selections leave the request pending so the application can correct
+    /// its response. Expired or previously answered requests fail explicitly.
+    pub async fn respond_permission(
+        &self,
+        request_id: &str,
+        decision: PermissionDecision,
+    ) -> Result<()> {
+        let state = &self.lease.state;
+        if state.is_closed() {
+            return Err(state.error());
+        }
+        let pending = {
+            let mut permissions = lock(&state.permissions);
+            let pending = permissions
+                .get(request_id)
+                .ok_or(NativeError::UnknownPermission)?;
+            if pending.expires <= Instant::now() || pending.responder.cancellation().is_cancelled()
+            {
+                return Err(NativeError::UnknownPermission);
+            }
+            match &decision {
+                RequestPermissionOutcome::Cancelled => {}
+                RequestPermissionOutcome::Selected(selected)
+                    if pending
+                        .request
+                        .options
+                        .iter()
+                        .any(|option| option.option_id == selected.option_id) => {}
+                _ => return Err(NativeError::InvalidPermissionOption),
+            }
+            permissions
+                .remove(request_id)
+                .ok_or(NativeError::UnknownPermission)?
+        };
+        pending
+            .responder
+            .respond(RequestPermissionResponse::new(decision))?;
+        Ok(())
+    }
+
+    /// Cancel work, detach local routing, and close the remote session when supported.
+    ///
+    /// Without `session/close`, persisted agent history remains available for an
+    /// explicit future load; this method never claims to delete remote history.
+    pub async fn close(&self) -> Result<()> {
+        let state = &self.lease.state;
+        if state.is_closed() {
+            return Ok(());
+        }
+        self.cancel().await?;
+        if self
+            .lease
+            .client
+            .info
+            .agent_capabilities
+            .session_capabilities
+            .close
+            .is_some()
+        {
+            let response = timeout(
+                self.lease.client.options.shutdown_timeout,
+                state
+                    .connection
+                    .send_request(CloseSessionRequest::new(state.id.clone()))
+                    .block_task(),
+            )
+            .await;
+            match response {
+                Ok(Ok(_)) => {}
+                Ok(Err(error)) => {
+                    state.close(error.clone().into(), false);
+                    return Err(error.into());
+                }
+                Err(_) => {
+                    state.close(NativeError::Timeout, true);
+                    return Err(NativeError::Timeout);
+                }
+            }
+        }
+        state.close(NativeError::Closed, false);
+        Ok(())
+    }
+}
+
+async fn run_session(
+    mut session: ActiveSession<'static, Agent>,
+    mut commands: mpsc::Receiver<PromptCommand>,
+    state: Arc<SessionState>,
+    options: super::NativeOptions,
+) -> Result<()> {
+    let (completion_tx, mut completion_rx) = mpsc::unbounded_channel::<Result<StopReason>>();
+    let mut active: Option<oneshot::Sender<Result<StopReason>>> = None;
+    let mut closed = state.closed.subscribe();
+    let mut connection_closed = state.client.closed.subscribe();
+    let mut expiry = interval(
+        options
+            .permission_timeout
+            .min(std::time::Duration::from_secs(1)),
+    );
+    loop {
+        if state.is_closed() {
+            return Err(state.error());
+        }
+        tokio::select! {
+            biased;
+            _ = closed.changed() => return Ok(()),
+            _ = connection_closed.changed() => return Err(state.client.error()),
+            message = session.read_update() => {
+                match message? {
+                    SessionMessage::SessionMessage(dispatch) => {
+                        MatchDispatch::new(dispatch)
+                            .if_notification(async |notification: SessionNotification| {
+                                state.emit(Event::Update { update: notification.update })
+                                    .map_err(|_| agent_client_protocol::Error::internal_error())
+                            }).await
+                            .if_request(async |request: RequestPermissionRequest, responder| {
+                                if options.permissions == PermissionPolicy::Deny
+                                    || lock(&state.turn).cancelled
+                                    || lock(&state.permissions).len() >= options.event_capacity
+                                {
+                                    return responder.respond(RequestPermissionResponse::new(RequestPermissionOutcome::Cancelled));
+                                }
+                                let request_id = format!("permission-{}", NEXT_PERMISSION.fetch_add(1, Ordering::Relaxed));
+                                lock(&state.permissions).insert(request_id.clone(), PendingPermission {
+                                    request: request.clone(), responder,
+                                    expires: Instant::now() + options.permission_timeout,
+                                });
+                                state.emit(Event::Permission { request_id, request })
+                                    .map_err(|_| agent_client_protocol::Error::internal_error())
+                            }).await
+                            .otherwise(async |dispatch| {
+                                match dispatch {
+                                    Dispatch::Request(_, responder) => responder.respond_with_result(Err(agent_client_protocol::Error::method_not_found())),
+                                    _ => Ok(()),
+                                }
+                            }).await?;
+                    }
+                    SessionMessage::StopReason(stop_reason) => {
+                        let _ = completion_tx.send(Ok(stop_reason));
+                    }
+                    _ => return Err(NativeError::Unsupported("the received SDK session event")),
+                }
+            }
+            Some(outcome) = completion_rx.recv(), if active.is_some() => {
+                state.finish(&outcome)?;
+                if let Some(result) = active.take() { let _ = result.send(outcome); }
+            }
+            Some(command) = commands.recv(), if active.is_none() => {
+                {
+                    let mut turn = lock(&state.turn);
+                    turn.current = Some(command.prompt.request_id.clone());
+                }
+                if lock(&state.turn).cancelled || command.result.is_closed() {
+                    let outcome = Ok(StopReason::Cancelled);
+                    state.finish(&outcome)?;
+                    let _ = command.result.send(outcome);
+                    continue;
+                }
+                active = Some(command.result);
+                let tx = completion_tx.clone();
+                session.connection().send_request(PromptRequest::new(state.id.clone(), command.prompt.content))
+                    .on_receiving_result(async move |result: std::result::Result<PromptResponse, agent_client_protocol::Error>| {
+                        // Earlier notifications are already on ActiveSession's queue.
+                        // The biased reader drains them before publishing completion.
+                        let _ = tx.send(result.map(|response| response.stop_reason).map_err(NativeError::from));
+                        Ok(())
+                    })?;
+            }
+            _ = expiry.tick() => {
+                let expired: Vec<_> = lock(&state.permissions).iter()
+                    .filter(|(_, pending)| pending.expires <= Instant::now() || pending.responder.cancellation().is_cancelled())
+                    .map(|(id, _)| id.clone()).collect();
+                for id in expired {
+                    if let Some(pending) = lock(&state.permissions).remove(&id) {
+                        let _ = pending.responder.respond(RequestPermissionResponse::new(RequestPermissionOutcome::Cancelled));
+                    }
+                }
+            }
+        }
+    }
+}