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;
#[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)
}
#[derive(Default)]
pub(crate) struct ApprovalShared {
pub(crate) pending: Mutex<HashMap<ToolCallId, oneshot::Sender<bool>>>,
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,
}
}
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,
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
}
}
}
}
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),
}
}
Err(KaynineError::Store(
kaynine_core::error::StoreError::Internal(
"approval append failed after 8 revision-conflict retries".into(),
),
))
}
async fn broadcast(&self, payload: RealtimeEvent) {
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;
if self
.append(vec![AuthoritativeEvent::ApprovalRequested {
call_id: request.call_id.clone(),
deadline_unix,
}])
.await
.is_err()
{
return ApprovalDecision::Denied {
reason: "审批持久化失败".into(),
};
}
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;
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 == "审批超时");
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
}
}