use std::sync::{mpsc, Arc, Mutex};
use async_trait::async_trait;
use crate::mcp::{ElicitationRequest, ElicitationResponse, McpElicitationHandler};
use crate::permissions::{ApprovalOutcome, ApprovalRequest, PermissionsApprovalHandler};
use crate::subagents::QueuedApproval;
use super::bridge::{
PendingApprovalRequest, PendingChildApproval, PendingElicitation, PendingOAuthDisplay,
};
pub struct TuiApprovalHandler {
tx: mpsc::Sender<PendingApprovalRequest>,
}
impl TuiApprovalHandler {
pub fn new(tx: mpsc::Sender<PendingApprovalRequest>) -> Self {
TuiApprovalHandler { tx }
}
}
impl PermissionsApprovalHandler for TuiApprovalHandler {
fn ask(&self, req: &ApprovalRequest) -> ApprovalOutcome {
let (reply_tx, reply_rx) = mpsc::channel();
let pending = PendingApprovalRequest {
tool: req.tool.to_string(),
subject: req.subject.map(String::from),
raw_args: req.raw_args.clone(),
reply_tx,
};
if self.tx.send(pending).is_err() {
return ApprovalOutcome::Deny;
}
reply_rx.recv().unwrap_or(ApprovalOutcome::Deny)
}
}
pub struct TuiChildApprovalHandler {
child_agent_id: String,
queue: Arc<Mutex<Vec<QueuedApproval>>>,
tx: mpsc::Sender<PendingChildApproval>,
}
impl TuiChildApprovalHandler {
pub fn new(
child_agent_id: String,
queue: Arc<Mutex<Vec<QueuedApproval>>>,
tx: mpsc::Sender<PendingChildApproval>,
) -> Self {
TuiChildApprovalHandler {
child_agent_id,
queue,
tx,
}
}
}
impl PermissionsApprovalHandler for TuiChildApprovalHandler {
fn ask(&self, req: &ApprovalRequest) -> ApprovalOutcome {
let queued = crate::subagents::queue_approval(
&self.queue,
QueuedApproval {
child_agent_id: self.child_agent_id.clone(),
tool: req.tool.to_string(),
subject: req.subject.map(String::from),
queued_at_ms: now_ms(),
outcome: None,
},
);
let (reply_tx, reply_rx) = mpsc::channel();
let pending = PendingChildApproval {
child_agent_id: self.child_agent_id.clone(),
tool: req.tool.to_string(),
subject: req.subject.map(String::from),
raw_args: req.raw_args.clone(),
reply_tx,
};
let outcome = if self.tx.send(pending).is_err() {
ApprovalOutcome::Deny
} else {
reply_rx.recv().unwrap_or(ApprovalOutcome::Deny)
};
if let Some(index) = queued {
crate::subagents::record_queued_outcome(&self.queue, index, outcome.into());
}
outcome
}
}
fn now_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0)
}
pub struct TuiElicitationHandler {
tx: mpsc::Sender<PendingElicitation>,
}
impl TuiElicitationHandler {
pub fn new(tx: mpsc::Sender<PendingElicitation>) -> Self {
TuiElicitationHandler { tx }
}
}
#[async_trait]
impl McpElicitationHandler for TuiElicitationHandler {
async fn handle(&self, request: &ElicitationRequest) -> ElicitationResponse {
let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();
let pending = PendingElicitation {
message: request.message.clone(),
requested_schema: request.requested_schema.clone(),
reply_tx,
};
if self.tx.send(pending).is_err() {
return ElicitationResponse {
action: crate::mcp::ElicitationAction::Cancel,
content: None,
};
}
reply_rx.await.unwrap_or(ElicitationResponse {
action: crate::mcp::ElicitationAction::Cancel,
content: None,
})
}
}
pub struct TuiBridge {
approval_tx: mpsc::Sender<PendingApprovalRequest>,
pub approval_rx: mpsc::Receiver<PendingApprovalRequest>,
child_approval_tx: mpsc::Sender<PendingChildApproval>,
pub child_approval_rx: mpsc::Receiver<PendingChildApproval>,
elicitation_tx: mpsc::Sender<PendingElicitation>,
pub elicitation_rx: mpsc::Receiver<PendingElicitation>,
oauth_tx: mpsc::Sender<PendingOAuthDisplay>,
pub oauth_rx: mpsc::Receiver<PendingOAuthDisplay>,
}
impl Default for TuiBridge {
fn default() -> Self {
Self::new()
}
}
impl TuiBridge {
pub fn new() -> Self {
let (approval_tx, approval_rx) = mpsc::channel();
let (child_approval_tx, child_approval_rx) = mpsc::channel();
let (elicitation_tx, elicitation_rx) = mpsc::channel();
let (oauth_tx, oauth_rx) = mpsc::channel();
TuiBridge {
approval_tx,
approval_rx,
child_approval_tx,
child_approval_rx,
elicitation_tx,
elicitation_rx,
oauth_tx,
oauth_rx,
}
}
pub fn install_on(&self, agent: &mut crate::agent::Agent) {
agent.set_permissions_approval_handler(TuiApprovalHandler::new(self.approval_tx.clone()));
agent.set_user_question_handler(self.elicitation_handler());
agent.set_child_approval_handler_factory({
let tx = self.child_approval_tx.clone();
move |child_id, queue| {
Arc::new(TuiChildApprovalHandler::new(child_id, queue, tx.clone()))
as Arc<dyn PermissionsApprovalHandler>
}
});
}
pub fn elicitation_handler(&self) -> Arc<dyn McpElicitationHandler> {
Arc::new(TuiElicitationHandler::new(self.elicitation_tx.clone()))
}
pub fn oauth_sender(&self) -> mpsc::Sender<PendingOAuthDisplay> {
self.oauth_tx.clone()
}
}