1use std::collections::HashMap;
12use std::sync::{Mutex, OnceLock};
13
14use bamboo_agent_core::AgentEvent;
15use bamboo_domain::poison::PoisonRecover;
16use bamboo_subagent::proto::ParentFrame;
17use tokio::sync::mpsc;
18
19use super::approval_registry::{
20 ApprovalRegistry, ApprovalState, DurableApproval, SharedApprovalRegistry,
21};
22
23type ScopeId = usize;
24type LiveKey = (ScopeId, String, u32);
25type PendingKey = (ScopeId, String, String, u32, String);
26
27fn scope_id(registry: Option<&SharedApprovalRegistry>) -> ScopeId {
28 registry.map_or(0, |registry| registry.lock().recover_poison().scope_id())
29}
30
31fn map() -> &'static Mutex<HashMap<LiveKey, mpsc::UnboundedSender<ParentFrame>>> {
32 static MAP: OnceLock<Mutex<HashMap<LiveKey, mpsc::UnboundedSender<ParentFrame>>>> =
33 OnceLock::new();
34 MAP.get_or_init(|| Mutex::new(HashMap::new()))
35}
36
37#[derive(Clone)]
44struct PendingApproval {
45 parent_session_id: String,
46 tool_name: String,
47 permission: String,
48 resource: String,
49 created_at: String,
50 version: u64,
51 child_attempt: u32,
52 event_tx: mpsc::Sender<AgentEvent>,
53}
54
55fn pending() -> &'static Mutex<HashMap<PendingKey, PendingApproval>> {
56 static PENDING: OnceLock<Mutex<HashMap<PendingKey, PendingApproval>>> = OnceLock::new();
57 PENDING.get_or_init(|| Mutex::new(HashMap::new()))
58}
59
60pub fn initialize_durable_approvals(
63 path: std::path::PathBuf,
64) -> std::io::Result<(SharedApprovalRegistry, Vec<AgentEvent>)> {
65 let mut registry = ApprovalRegistry::open(path)?;
66 let reconciled = registry.reconcile_restart()?;
67 let events = reconciled.into_iter().map(record_event).collect();
68 Ok((std::sync::Arc::new(Mutex::new(registry)), events))
69}
70
71pub struct PendingApprovalObservation<'a> {
73 pub registry: Option<&'a SharedApprovalRegistry>,
74 pub parent_session_id: &'a str,
75 pub child_id: &'a str,
76 pub child_attempt: u32,
77 pub request_id: &'a str,
78 pub tool_name: &'a str,
79 pub permission: &'a str,
80 pub resource: &'a str,
81 pub event_tx: mpsc::Sender<AgentEvent>,
82}
83
84pub fn observe_pending_approval(observation: PendingApprovalObservation<'_>) -> (u64, String) {
88 let PendingApprovalObservation {
89 registry,
90 parent_session_id,
91 child_id,
92 child_attempt,
93 request_id,
94 tool_name,
95 permission,
96 resource,
97 event_tx,
98 } = observation;
99 let now = chrono::Utc::now();
100 let version = now.timestamp_micros().max(0) as u64;
101 let created_at = now.to_rfc3339();
102 let durable_record = DurableApproval {
103 parent_session_id: parent_session_id.to_string(),
104 child_session_id: child_id.to_string(),
105 child_attempt,
106 request_id: request_id.to_string(),
107 tool_name: tool_name.to_string(),
108 permission: permission.to_string(),
109 resource: resource.to_string(),
110 created_at: created_at.clone(),
111 updated_at: created_at.clone(),
112 version,
113 state: ApprovalState::Pending,
114 approved: None,
115 reason: None,
116 };
117 if let Some(registry) = registry {
118 if let Err(error) = registry.lock().recover_poison().register(durable_record) {
119 tracing::error!("failed to persist pending child approval: {error}");
120 return (0, created_at);
121 }
122 }
123 pending().lock().recover_poison().insert(
124 (
125 scope_id(registry),
126 parent_session_id.to_string(),
127 child_id.to_string(),
128 child_attempt,
129 request_id.to_string(),
130 ),
131 PendingApproval {
132 parent_session_id: parent_session_id.to_string(),
133 tool_name: tool_name.to_string(),
134 permission: permission.to_string(),
135 resource: resource.to_string(),
136 created_at: created_at.clone(),
137 version,
138 child_attempt,
139 event_tx,
140 },
141 );
142 (version, created_at)
143}
144
145#[allow(
150 clippy::too_many_arguments,
151 reason = "public compatibility wrapper; the typed observation is the canonical API"
152)]
153pub fn register_pending_approval_observed(
154 registry: Option<&SharedApprovalRegistry>,
155 parent_session_id: &str,
156 child_id: &str,
157 child_attempt: u32,
158 request_id: &str,
159 tool_name: &str,
160 permission: &str,
161 resource: &str,
162 event_tx: mpsc::Sender<AgentEvent>,
163) -> (u64, String) {
164 observe_pending_approval(PendingApprovalObservation {
165 registry,
166 parent_session_id,
167 child_id,
168 child_attempt,
169 request_id,
170 tool_name,
171 permission,
172 resource,
173 event_tx,
174 })
175}
176
177#[cfg(test)]
178fn register_pending_approval(child_id: &str, request_id: &str) {
179 let (event_tx, _rx) = mpsc::channel(1);
180 let _ = observe_pending_approval(PendingApprovalObservation {
181 registry: None,
182 parent_session_id: "test-parent",
183 child_id,
184 child_attempt: 0,
185 request_id,
186 tool_name: "test-tool",
187 permission: "test-permission",
188 resource: "test-resource",
189 event_tx,
190 });
191}
192
193pub fn take_pending_approval(child_id: &str, request_id: &str) -> bool {
197 remove_unique_pending(None, child_id, request_id).is_some()
198}
199
200pub fn clear_pending_approvals_for(
202 registry: Option<&SharedApprovalRegistry>,
203 child_id: &str,
204 child_attempt: u32,
205) {
206 let records: Vec<_> = {
207 let mut guard = pending().lock().recover_poison();
208 let keys: Vec<_> = guard
209 .keys()
210 .filter(|(scope, _, child, attempt, _)| {
211 *scope == scope_id(registry) && child == child_id && *attempt == child_attempt
212 })
213 .cloned()
214 .collect();
215 keys.into_iter()
216 .filter_map(|key| guard.remove(&key).map(|record| (key.4, record)))
217 .collect()
218 };
219 for (request_id, record) in records {
220 let durable = finish_durable(
221 registry,
222 &record.parent_session_id,
223 child_id,
224 record.child_attempt,
225 &request_id,
226 false,
227 Some("child_disconnected"),
228 );
229 if registry.is_none() || durable.is_some() {
230 emit_resolution(
231 child_id,
232 &request_id,
233 record,
234 "delivery_failed",
235 Some("child_disconnected"),
236 durable.as_ref(),
237 );
238 }
239 }
240}
241
242pub fn expire_pending_approval(
243 registry: Option<&SharedApprovalRegistry>,
244 child_id: &str,
245 request_id: &str,
246) -> bool {
247 let record = remove_unique_pending(registry, child_id, request_id);
248 let Some(record) = record else {
249 return false;
250 };
251 let durable = finish_durable(
252 registry,
253 &record.parent_session_id,
254 child_id,
255 record.child_attempt,
256 request_id,
257 false,
258 Some("approval_timeout"),
259 );
260 if registry.is_none() || durable.is_some() {
261 emit_resolution(
262 child_id,
263 request_id,
264 record,
265 "expired",
266 Some("approval_timeout"),
267 durable.as_ref(),
268 );
269 }
270 true
271}
272
273fn emit_resolution(
274 child_id: &str,
275 request_id: &str,
276 record: PendingApproval,
277 status: &str,
278 reason: Option<&str>,
279 durable: Option<&DurableApproval>,
280) {
281 let now = chrono::Utc::now();
282 let event = AgentEvent::ChildApprovalChanged {
283 parent_session_id: record.parent_session_id,
284 child_session_id: child_id.to_string(),
285 child_attempt: record.child_attempt,
286 request_id: request_id.to_string(),
287 version: durable.map_or_else(
288 || (now.timestamp_micros().max(0) as u64).max(record.version.saturating_add(1)),
289 |record| record.version,
290 ),
291 status: status.to_string(),
292 reason: reason.map(str::to_string),
293 tool_name: record.tool_name,
294 permission: record.permission,
295 resource: record.resource,
296 created_at: record.created_at,
297 resolved_at: Some(
298 durable
299 .map(|record| record.updated_at.clone())
300 .unwrap_or_else(|| now.to_rfc3339()),
301 ),
302 };
303 match record.event_tx.try_send(event) {
304 Ok(()) | Err(mpsc::error::TrySendError::Closed(_)) => {}
305 Err(mpsc::error::TrySendError::Full(event)) => {
306 let tx = record.event_tx;
307 tokio::spawn(async move {
308 let _ = tx.send(event).await;
309 });
310 }
311 }
312}
313
314pub fn deliver_approval_checked(
322 registry: Option<&SharedApprovalRegistry>,
323 child_id: &str,
324 request_id: &str,
325 approved: bool,
326) -> bool {
327 let Some((key, record)) = find_unique_pending(registry, child_id, request_id) else {
328 return false;
329 };
330 if let Some(registry) = registry {
333 match registry.lock().recover_poison().record_decision(
334 &record.parent_session_id,
335 child_id,
336 record.child_attempt,
337 request_id,
338 approved,
339 ) {
340 Ok(Some(_)) => {}
341 Ok(None) => return false,
342 Err(error) => {
343 tracing::error!("failed to persist child approval decision: {error}");
344 return false;
345 }
346 }
347 }
348 let Some(record) = pending().lock().recover_poison().remove(&key) else {
349 let _ = finish_durable(
350 registry,
351 &record.parent_session_id,
352 child_id,
353 record.child_attempt,
354 request_id,
355 false,
356 Some("pending_state_lost"),
357 );
358 return false;
359 };
360 let delivered = deliver_approval_scoped(
361 registry,
362 child_id,
363 record.child_attempt,
364 request_id,
365 approved,
366 );
367 let status = if delivered {
368 if approved {
369 "approved"
370 } else {
371 "denied"
372 }
373 } else {
374 "delivery_failed"
375 };
376 let durable = finish_durable(
377 registry,
378 &record.parent_session_id,
379 child_id,
380 record.child_attempt,
381 request_id,
382 delivered,
383 (!delivered).then_some("child_not_live"),
384 );
385 if registry.is_none() || durable.is_some() {
386 emit_resolution(
387 child_id,
388 request_id,
389 record,
390 status,
391 (!delivered).then_some("child_not_live"),
392 durable.as_ref(),
393 );
394 }
395 delivered
396}
397
398fn finish_durable(
399 registry: Option<&SharedApprovalRegistry>,
400 parent_id: &str,
401 child_id: &str,
402 child_attempt: u32,
403 request_id: &str,
404 delivered: bool,
405 reason: Option<&str>,
406) -> Option<DurableApproval> {
407 if let Some(registry) = registry {
408 match registry.lock().recover_poison().finish(
409 parent_id,
410 child_id,
411 child_attempt,
412 request_id,
413 delivered,
414 reason,
415 ) {
416 Ok(record) => return record,
417 Err(error) => {
418 tracing::error!("failed to persist child approval resolution: {error}");
419 }
420 }
421 }
422 None
423}
424
425fn record_event(record: DurableApproval) -> AgentEvent {
426 AgentEvent::ChildApprovalChanged {
427 parent_session_id: record.parent_session_id,
428 child_session_id: record.child_session_id,
429 child_attempt: record.child_attempt,
430 request_id: record.request_id,
431 version: record.version,
432 status: match record.state {
433 ApprovalState::Pending => "pending",
434 ApprovalState::DecisionRecorded => "decision_recorded",
435 ApprovalState::Delivered if record.approved == Some(true) => "approved",
436 ApprovalState::Delivered => "denied",
437 ApprovalState::DeliveryFailed => "delivery_failed",
438 ApprovalState::Expired => "expired",
439 }
440 .to_string(),
441 reason: record.reason,
442 tool_name: record.tool_name,
443 permission: record.permission,
444 resource: record.resource,
445 created_at: record.created_at,
446 resolved_at: Some(record.updated_at),
447 }
448}
449
450pub struct LiveActorGuard {
453 scope_id: ScopeId,
454 child_id: String,
455 child_attempt: u32,
456 approval_registry: Option<SharedApprovalRegistry>,
457}
458
459impl Drop for LiveActorGuard {
460 fn drop(&mut self) {
461 map().lock().recover_poison().remove(&(
462 self.scope_id,
463 self.child_id.clone(),
464 self.child_attempt,
465 ));
466 clear_pending_approvals_for(
469 self.approval_registry.as_ref(),
470 &self.child_id,
471 self.child_attempt,
472 );
473 }
474}
475
476pub fn register(
478 child_id: &str,
479 tx: mpsc::UnboundedSender<ParentFrame>,
480 child_attempt: u32,
481 approval_registry: Option<SharedApprovalRegistry>,
482) -> LiveActorGuard {
483 let scope_id = scope_id(approval_registry.as_ref());
484 map()
485 .lock()
486 .recover_poison()
487 .insert((scope_id, child_id.to_string(), child_attempt), tx);
488 LiveActorGuard {
489 scope_id,
490 child_id: child_id.to_string(),
491 child_attempt,
492 approval_registry,
493 }
494}
495
496fn find_unique_pending(
497 registry: Option<&SharedApprovalRegistry>,
498 child_id: &str,
499 request_id: &str,
500) -> Option<(PendingKey, PendingApproval)> {
501 let guard = pending().lock().recover_poison();
502 let scope = scope_id(registry);
503 let mut matches = guard
504 .iter()
505 .filter(|(key, _)| key.0 == scope && key.2 == child_id && key.4 == request_id);
506 let (key, record) = matches.next()?;
507 if matches.next().is_some() {
508 return None;
509 }
510 Some((key.clone(), record.clone()))
511}
512
513fn remove_unique_pending(
514 registry: Option<&SharedApprovalRegistry>,
515 child_id: &str,
516 request_id: &str,
517) -> Option<PendingApproval> {
518 let mut guard = pending().lock().recover_poison();
519 let scope = scope_id(registry);
520 let mut keys = guard
521 .keys()
522 .filter(|(key_scope, _, child, _, request)| {
523 *key_scope == scope && child == child_id && request == request_id
524 })
525 .cloned();
526 let key = keys.next()?;
527 if keys.next().is_some() {
528 return None;
529 }
530 guard.remove(&key)
531}
532
533pub fn deliver_message(child_id: &str, text: &str) -> bool {
536 let guard = map().lock().recover_poison();
537 let mut senders = guard
538 .iter()
539 .filter(|((_, child, _), _)| child == child_id)
540 .map(|(_, sender)| sender);
541 let Some(tx) = senders.next() else {
542 return false;
543 };
544 if senders.next().is_some() {
545 return false;
546 }
547 tx.send(ParentFrame::Message {
548 text: text.to_string(),
549 })
550 .is_ok()
551}
552
553pub fn deliver_approval(child_id: &str, request_id: &str, approved: bool) -> bool {
564 deliver_approval_scoped(None, child_id, 0, request_id, approved)
565}
566
567pub fn deliver_approval_scoped(
568 registry: Option<&SharedApprovalRegistry>,
569 child_id: &str,
570 child_attempt: u32,
571 request_id: &str,
572 approved: bool,
573) -> bool {
574 let guard = map().lock().recover_poison();
575 match guard.get(&(scope_id(registry), child_id.to_string(), child_attempt)) {
576 Some(tx) => tx
577 .send(ParentFrame::ApprovalReply {
578 id: request_id.to_string(),
579 approved,
580 })
581 .is_ok(),
582 None => false,
583 }
584}
585
586pub fn is_live(child_id: &str) -> bool {
588 map()
589 .lock()
590 .recover_poison()
591 .keys()
592 .any(|(_, child, _)| child == child_id)
593}
594
595#[cfg(test)]
596mod tests {
597 use super::*;
598
599 fn registry() -> SharedApprovalRegistry {
600 std::sync::Arc::new(Mutex::new(
601 ApprovalRegistry::open(tempfile::tempdir().unwrap().keep().join("registry.json"))
602 .unwrap(),
603 ))
604 }
605
606 #[test]
607 fn register_deliver_unregister() {
608 let (tx, mut rx) = mpsc::unbounded_channel();
609 let guard = register("c-live", tx, 0, None);
610 assert!(is_live("c-live"));
611 assert!(deliver_message("c-live", "hi"));
612 match rx.try_recv() {
613 Ok(ParentFrame::Message { text }) => assert_eq!(text, "hi"),
614 other => panic!("expected message frame, got {other:?}"),
615 }
616
617 drop(guard);
618 assert!(!is_live("c-live"));
619 assert!(!deliver_message("c-live", "gone"));
620 }
621
622 #[test]
623 fn deliver_fails_when_receiver_dropped() {
624 let (tx, rx) = mpsc::unbounded_channel();
625 let _guard = register("c-dead", tx, 0, None);
626 drop(rx);
627 assert!(!deliver_message("c-dead", "hi"));
628 }
629
630 #[test]
631 fn deliver_approval_routes_reply_frame() {
632 let (tx, mut rx) = mpsc::unbounded_channel();
633 let guard = register("c-appr", tx, 0, None);
634 assert!(deliver_approval("c-appr", "req-7", true));
635 match rx.try_recv() {
636 Ok(ParentFrame::ApprovalReply { id, approved }) => {
637 assert_eq!(id, "req-7");
638 assert!(approved);
639 }
640 other => panic!("expected approval reply, got {other:?}"),
641 }
642 drop(guard);
643 assert!(!deliver_approval("c-appr", "req-8", false));
645 }
646
647 #[test]
648 fn app_scopes_and_attempt_guards_do_not_cross_talk() {
649 let first_registry = registry();
650 let second_registry = registry();
651 let (first_tx, mut first_rx) = mpsc::unbounded_channel();
652 let (retry_tx, mut retry_rx) = mpsc::unbounded_channel();
653 let (second_tx, mut second_rx) = mpsc::unbounded_channel();
654 let old_guard = register("shared-child", first_tx, 1, Some(first_registry.clone()));
655 let retry_guard = register("shared-child", retry_tx, 2, Some(first_registry.clone()));
656 let second_guard = register("shared-child", second_tx, 1, Some(second_registry.clone()));
657
658 drop(old_guard);
659 assert!(deliver_approval_scoped(
660 Some(&first_registry),
661 "shared-child",
662 2,
663 "retry-request",
664 true,
665 ));
666 assert!(matches!(
667 retry_rx.try_recv(),
668 Ok(ParentFrame::ApprovalReply { id, .. }) if id == "retry-request"
669 ));
670 assert!(deliver_approval_scoped(
671 Some(&second_registry),
672 "shared-child",
673 1,
674 "other-app-request",
675 false,
676 ));
677 assert!(matches!(
678 second_rx.try_recv(),
679 Ok(ParentFrame::ApprovalReply { id, .. }) if id == "other-app-request"
680 ));
681 assert!(first_rx.try_recv().is_err());
682 drop(retry_guard);
683 drop(second_guard);
684 }
685
686 #[test]
687 fn pending_approval_is_one_shot() {
688 register_pending_approval("c-pend", "req-1");
689 assert!(take_pending_approval("c-pend", "req-1"));
691 assert!(!take_pending_approval("c-pend", "req-1"));
692 }
693
694 #[test]
695 fn take_of_unregistered_pair_is_false() {
696 assert!(!take_pending_approval("c-unknown", "req-x"));
698 register_pending_approval("c-known", "req-real");
700 assert!(!take_pending_approval("c-known", "req-bogus"));
701 assert!(take_pending_approval("c-known", "req-real"));
703 }
704
705 #[test]
706 fn deliver_approval_checked_only_delivers_for_registered_pair() {
707 let (tx, mut rx) = mpsc::unbounded_channel();
708 let _guard = register("c-checked", tx, 0, None);
709
710 assert!(!deliver_approval_checked(
712 None,
713 "c-checked",
714 "req-stray",
715 true
716 ));
717 assert!(rx.try_recv().is_err());
718
719 register_pending_approval("c-checked", "req-ok");
721 assert!(deliver_approval_checked(None, "c-checked", "req-ok", true));
722 match rx.try_recv() {
723 Ok(ParentFrame::ApprovalReply { id, approved }) => {
724 assert_eq!(id, "req-ok");
725 assert!(approved);
726 }
727 other => panic!("expected approval reply, got {other:?}"),
728 }
729 assert!(!deliver_approval_checked(None, "c-checked", "req-ok", true));
731 assert!(rx.try_recv().is_err());
732 }
733
734 #[test]
735 fn clear_pending_approvals_for_drops_them() {
736 register_pending_approval("c-clear", "req-a");
737 register_pending_approval("c-clear", "req-b");
738 clear_pending_approvals_for(None, "c-clear", 0);
739 assert!(!take_pending_approval("c-clear", "req-a"));
740 assert!(!take_pending_approval("c-clear", "req-b"));
741 }
742
743 #[tokio::test]
744 async fn observed_approval_emits_exactly_one_terminal_outcome() {
745 let (wire_tx, _wire_rx) = mpsc::unbounded_channel();
746 let _guard = register("c-audit", wire_tx, 0, None);
747 let (event_tx, mut event_rx) = mpsc::channel(8);
748 observe_pending_approval(PendingApprovalObservation {
749 registry: None,
750 parent_session_id: "parent-audit",
751 child_id: "c-audit",
752 child_attempt: 0,
753 request_id: "req-audit",
754 tool_name: "Bash",
755 permission: "execute",
756 resource: "/tmp/x",
757 event_tx,
758 });
759
760 assert!(deliver_approval_checked(
761 None,
762 "c-audit",
763 "req-audit",
764 false
765 ));
766 assert!(!deliver_approval_checked(
767 None,
768 "c-audit",
769 "req-audit",
770 true
771 ));
772 assert!(matches!(
773 event_rx.recv().await,
774 Some(AgentEvent::ChildApprovalChanged { status, .. }) if status == "denied"
775 ));
776 assert!(event_rx.try_recv().is_err());
777 }
778
779 #[tokio::test]
780 async fn durable_resolution_uses_registry_version_and_attempt() {
781 let registry = registry();
782 let (wire_tx, _wire_rx) = mpsc::unbounded_channel();
783 let _guard = register("c-versioned", wire_tx, 7, Some(registry.clone()));
784 let (event_tx, mut event_rx) = mpsc::channel(4);
785 let (pending_version, _) = observe_pending_approval(PendingApprovalObservation {
786 registry: Some(®istry),
787 parent_session_id: "parent-versioned",
788 child_id: "c-versioned",
789 child_attempt: 7,
790 request_id: "req-versioned",
791 tool_name: "Bash",
792 permission: "execute",
793 resource: "/tmp/versioned",
794 event_tx,
795 });
796
797 assert!(deliver_approval_checked(
798 Some(®istry),
799 "c-versioned",
800 "req-versioned",
801 true,
802 ));
803 assert!(matches!(
804 event_rx.recv().await,
805 Some(AgentEvent::ChildApprovalChanged {
806 child_attempt: 7,
807 version,
808 status,
809 ..
810 }) if version == pending_version + 2 && status == "approved"
811 ));
812 }
813
814 #[tokio::test]
815 async fn timeout_and_disconnect_emit_terminal_outcomes() {
816 let (event_tx, mut event_rx) = mpsc::channel(8);
817 observe_pending_approval(PendingApprovalObservation {
818 registry: None,
819 parent_session_id: "parent-audit",
820 child_id: "c-expire",
821 child_attempt: 0,
822 request_id: "req-expire",
823 tool_name: "Bash",
824 permission: "execute",
825 resource: "/tmp/x",
826 event_tx: event_tx.clone(),
827 });
828 assert!(expire_pending_approval(None, "c-expire", "req-expire"));
829 assert!(!expire_pending_approval(None, "c-expire", "req-expire"));
830 assert!(matches!(
831 event_rx.recv().await,
832 Some(AgentEvent::ChildApprovalChanged { status, .. }) if status == "expired"
833 ));
834
835 observe_pending_approval(PendingApprovalObservation {
836 registry: None,
837 parent_session_id: "parent-audit",
838 child_id: "c-disconnect",
839 child_attempt: 0,
840 request_id: "req-disconnect",
841 tool_name: "Write",
842 permission: "write",
843 resource: "/tmp/y",
844 event_tx,
845 });
846 clear_pending_approvals_for(None, "c-disconnect", 0);
847 assert!(matches!(
848 event_rx.recv().await,
849 Some(AgentEvent::ChildApprovalChanged { status, .. }) if status == "delivery_failed"
850 ));
851 }
852}