kaynine_runtime/

approval.rs

//! Interactive approval (SPEC §8.3, §5.1): persist first, then broadcast.
//!
//! The handler is constructed inside the run task so it shares the loop's
//! realtime `mpsc` sender — the run task's forwarder is the only writer on
//! session broadcast feed, so approval realtime events flow through the same
//! channel. The handler sends its envelopes with `run_seq: None`; the
//! forwarder assigns a sequence number at forward time (see `run.rs`).

use kaynine_core::error::KaynineError;
use kaynine_core::event::{EventEnvelope, RealtimeEvent};
use kaynine_core::ids::{BranchId, RunId, SessionId, ToolCallId};
use kaynine_core::policy::{ApprovalDecision, ApprovalHandler, ApprovalRequest};
use kaynine_core::store::{AppendOutcome, AuthoritativeEvent, LeaseOwner, SessionStore};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{mpsc, oneshot, Mutex as AsyncMutex};
use tokio_util::sync::CancellationToken;

/// Outcome of a late `resolve_approval` submission (SPEC §8.3).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ResolutionOutcome {
    Delivered,
    Expired,
    NotFound,
}

impl From<ResolutionOutcome> for crate::service::ApprovalOutcome {
    fn from(outcome: ResolutionOutcome) -> Self {
        match outcome {
            ResolutionOutcome::Delivered => crate::service::ApprovalOutcome::Delivered,
            ResolutionOutcome::Expired => crate::service::ApprovalOutcome::Expired,
            ResolutionOutcome::NotFound => crate::service::ApprovalOutcome::NotFound,
        }
    }
}

fn host_unix_now() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}

/// Shared state between the run task's handler and the actor's
/// `ResolveApproval` command.
#[derive(Default)]
pub(crate) struct ApprovalShared {
    /// Live waiters: call_id → oneshot resolution channel.
    pub(crate) pending: Mutex<HashMap<ToolCallId, oneshot::Sender<bool>>>,
    /// Calls whose wait ended by timeout; call_id → deadline_unix. A resolve
    /// submitted afterwards maps to `Expired` instead of `NotFound`.
    pub(crate) expired: Mutex<HashMap<ToolCallId, i64>>,
}

pub struct InteractiveApprovalHandler {
    store: Arc<dyn SessionStore>,
    session_id: SessionId,
    branch_id: BranchId,
    run_id: RunId,
    revision: Arc<AsyncMutex<u64>>,
    owner: LeaseOwner,
    events: mpsc::Sender<EventEnvelope<RealtimeEvent>>,
    shared: Arc<ApprovalShared>,
    timeout: Duration,
}

impl InteractiveApprovalHandler {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        store: Arc<dyn SessionStore>,
        session_id: SessionId,
        branch_id: BranchId,
        run_id: RunId,
        revision: Arc<AsyncMutex<u64>>,
        owner: LeaseOwner,
        events: mpsc::Sender<EventEnvelope<RealtimeEvent>>,
        shared: Arc<ApprovalShared>,
        timeout: Duration,
    ) -> Self {
        Self {
            store,
            session_id,
            branch_id,
            run_id,
            revision,
            owner,
            events,
            shared,
            timeout,
        }
    }

    /// Resolves a pending approval from outside the run (actor command).
    /// `Delivered` only means the decision reached a live waiter.
    pub(crate) fn resolve(
        shared: &ApprovalShared,
        call_id: &ToolCallId,
        approved: bool,
    ) -> ResolutionOutcome {
        let sender = shared
            .pending
            .lock()
            .expect("pending mutex poisoned")
            .remove(call_id);
        match sender {
            Some(tx) => match tx.send(approved) {
                Ok(()) => ResolutionOutcome::Delivered,
                // The waiter went away (timeout/cancel raced us).
                Err(_) => {
                    if shared
                        .expired
                        .lock()
                        .expect("expired mutex poisoned")
                        .contains_key(call_id)
                    {
                        ResolutionOutcome::Expired
                    } else {
                        ResolutionOutcome::NotFound
                    }
                }
            },
            None => {
                if shared
                    .expired
                    .lock()
                    .expect("expired mutex poisoned")
                    .contains_key(call_id)
                {
                    ResolutionOutcome::Expired
                } else {
                    ResolutionOutcome::NotFound
                }
            }
        }
    }

    /// Append events to the session store, retrying up to 8 attempts on
    /// `RevisionConflict`. Mirrors the steer-append retry loop in `run.rs`.
    ///
    /// On conflict the shared revision is refreshed from the store response
    /// before the next attempt. `NotLeaseHolder` and `LeaseExpired` are
    /// returned immediately as hard failures; all other errors also fail
    /// immediately.
    async fn append(&self, events: Vec<AuthoritativeEvent>) -> Result<u64, KaynineError> {
        for _ in 0..8 {
            let expected = *self.revision.lock().await;
            match self
                .store
                .append_events(
                    &self.session_id,
                    expected,
                    &self.owner,
                    events.clone(),
                    Vec::new(),
                )
                .await
            {
                Ok(AppendOutcome::Appended { new_revision }) => {
                    *self.revision.lock().await = new_revision;
                    return Ok(new_revision);
                }
                Ok(AppendOutcome::RevisionConflict { current_revision }) => {
                    *self.revision.lock().await = current_revision;
                    tokio::time::sleep(Duration::from_millis(10)).await;
                    continue;
                }
                Ok(AppendOutcome::NotLeaseHolder | AppendOutcome::LeaseExpired) => {
                    return Err(KaynineError::Store(
                        kaynine_core::error::StoreError::Internal(
                            "lease lost during approval append".into(),
                        ),
                    ));
                }
                Err(error) => return Err(error),
            }
        }
        // Retries exhausted — fail closed.
        Err(KaynineError::Store(
            kaynine_core::error::StoreError::Internal(
                "approval append failed after 8 revision-conflict retries".into(),
            ),
        ))
    }

    async fn broadcast(&self, payload: RealtimeEvent) {
        // run_seq: None → the run-task forwarder assigns one at forward time.
        let _ = self
            .events
            .send(EventEnvelope {
                session_id: self.session_id.clone(),
                branch_id: self.branch_id.clone(),
                run_id: Some(self.run_id.clone()),
                revision: 0,
                run_seq: None,
                payload,
            })
            .await;
    }

    fn cleanup(&self, call_id: &ToolCallId, expired_at: Option<i64>) {
        if let Some(deadline) = expired_at {
            self.shared
                .expired
                .lock()
                .expect("expired mutex poisoned")
                .insert(call_id.clone(), deadline);
        }
        self.shared
            .pending
            .lock()
            .expect("pending mutex poisoned")
            .remove(call_id);
    }
}

#[async_trait::async_trait]
impl ApprovalHandler for InteractiveApprovalHandler {
    async fn wait(&self, request: ApprovalRequest, cancel: CancellationToken) -> ApprovalDecision {
        let deadline_unix = host_unix_now() + self.timeout.as_secs() as i64;

        // 1. Persist ApprovalRequested first (fail closed on append failure).
        if self
            .append(vec![AuthoritativeEvent::ApprovalRequested {
                call_id: request.call_id.clone(),
                deadline_unix,
            }])
            .await
            .is_err()
        {
            return ApprovalDecision::Denied {
                reason: "审批持久化失败".into(),
            };
        }

        // 2. Broadcast ApprovalPending after the durable write.
        self.broadcast(RealtimeEvent::ApprovalPending {
            call_id: request.call_id.clone(),
            deadline_unix,
        })
        .await;

        let (tx, rx) = oneshot::channel();
        self.shared
            .pending
            .lock()
            .expect("pending mutex poisoned")
            .insert(request.call_id.clone(), tx);
        let mut rx = rx;

        // 3. Await decision / timeout / cancellation.
        let outcome = tokio::select! {
            biased;
            _ = cancel.cancelled() => ApprovalDecision::Denied {
                reason: "工具在执行前被取消".into(),
            },
            _ = tokio::time::sleep(self.timeout) => ApprovalDecision::Denied {
                reason: "审批超时".into(),
            },
            res = &mut rx => match res {
                Ok(true) => ApprovalDecision::Approved,
                Ok(false) => ApprovalDecision::Denied {
                    reason: "审批被拒绝".into(),
                },
                Err(_) => ApprovalDecision::Denied {
                    reason: "审批通道关闭".into(),
                },
            },
        };

        let approved = outcome == ApprovalDecision::Approved;
        let timed_out =
            matches!(&outcome, ApprovalDecision::Denied { reason } if reason == "审批超时");

        // 4. Persist + broadcast ApprovalResolved, then clean up.
        if self
            .append(vec![AuthoritativeEvent::ApprovalResolved {
                call_id: request.call_id.clone(),
                approved,
            }])
            .await
            .is_err()
        {
            tracing::error!(call_id = %request.call_id, "approval resolved append failed");
        }
        self.broadcast(RealtimeEvent::ApprovalResolved {
            call_id: request.call_id.clone(),
            approved,
        })
        .await;
        self.cleanup(&request.call_id, timed_out.then_some(deadline_unix));

        outcome
    }
}