1use std::collections::HashMap;
8use std::sync::Arc;
9
10use chrono::Utc;
11use tokio::sync::{broadcast, mpsc, watch, RwLock};
12
13use bamboo_agent_core::AgentEvent;
14
15use super::runner_state::AgentRunner;
16
17fn visible_terminal_reason(event: &AgentEvent) -> Option<&'static str> {
18 match event {
19 AgentEvent::Complete { .. } => Some("complete"),
20 AgentEvent::Cancelled { .. } => Some("cancelled"),
21 AgentEvent::Error { .. } => Some("error"),
22 _ => None,
23 }
24}
25
26pub type AccountFeedInbox = mpsc::Sender<(Option<String>, AgentEvent)>;
31
32#[cfg(test)]
33#[path = "root_actor_event_forwarder_tests.rs"]
34mod root_actor_tests;
35
36#[derive(Clone, Debug)]
41pub struct HistoryCommitBarrier {
42 revision: watch::Receiver<u64>,
43}
44
45#[derive(Clone, Debug)]
47pub struct HistoryCommitAcknowledger {
48 revision: watch::Sender<u64>,
49}
50
51pub fn history_commit_barrier() -> (HistoryCommitAcknowledger, HistoryCommitBarrier) {
53 let (revision_tx, revision_rx) = watch::channel(0);
54 (
55 HistoryCommitAcknowledger {
56 revision: revision_tx,
57 },
58 HistoryCommitBarrier {
59 revision: revision_rx,
60 },
61 )
62}
63
64impl HistoryCommitAcknowledger {
65 pub fn acknowledge(&self) {
68 self.revision
69 .send_modify(|revision| *revision = revision.saturating_add(1));
70 }
71}
72
73impl HistoryCommitBarrier {
74 pub async fn send_and_wait(
77 &mut self,
78 event_tx: &mpsc::Sender<AgentEvent>,
79 session_id: String,
80 ) -> bool {
81 let observed = *self.revision.borrow_and_update();
82 if event_tx
83 .send(AgentEvent::SessionHistoryCommitted { session_id })
84 .await
85 .is_err()
86 {
87 return false;
88 }
89
90 loop {
91 if *self.revision.borrow() > observed {
92 return true;
93 }
94 if self.revision.changed().await.is_err() {
95 return false;
96 }
97 }
98 }
99}
100
101fn mirror_to_account_feed(inbox: &Option<AccountFeedInbox>, session_id: &str, event: &AgentEvent) {
107 if let Some(inbox) = inbox {
108 if event.is_durable_change() {
109 let route_session_id = event.session_id().unwrap_or(session_id);
110 let _ = inbox.try_send((Some(route_session_id.to_string()), event.clone()));
111 }
112 }
113}
114
115#[cfg(test)]
116#[allow(clippy::items_after_test_module)]
117mod tests {
118 use super::*;
119
120 use bamboo_agent_core::ToolResult;
121 use serde_json::json;
122
123 #[tokio::test]
124 async fn generic_forwarder_projects_native_download_without_changing_ordinary_results() {
125 let session_id = "generic-download-display";
126 let (broadcast_tx, mut broadcast_rx) = broadcast::channel(16);
127 let mut runner = AgentRunner::new();
128 runner.status = super::super::runner_state::AgentStatus::Running;
129 runner.event_sender = broadcast_tx.clone();
130 let run_id = runner.run_id.clone();
131 let runners = Arc::new(RwLock::new(HashMap::from([(
132 session_id.to_string(),
133 runner,
134 )])));
135 let (input, forwarder) =
136 create_event_forwarder(session_id.to_string(), run_id, broadcast_tx, runners, None);
137 assert!(matches!(
138 broadcast_rx.recv().await.unwrap(),
139 AgentEvent::ExecutionStarted { .. }
140 ));
141
142 let private = "private-file-url-and-bytes";
143 input
144 .send(AgentEvent::ToolStart {
145 tool_call_id: "download".into(),
146 tool_name: "browser".into(),
147 arguments: json!({"action":"download","selector":private,"expected_epoch":9}),
148 })
149 .await
150 .unwrap();
151 input
152 .send(AgentEvent::ToolError {
153 tool_call_id: "download".into(),
154 error: private.into(),
155 })
156 .await
157 .unwrap();
158 input
159 .send(AgentEvent::ToolComplete {
160 tool_call_id: "download".into(),
161 result: ToolResult::text(false, private),
162 })
163 .await
164 .unwrap();
165 input
166 .send(AgentEvent::ToolStart {
167 tool_call_id: "read".into(),
168 tool_name: "Read".into(),
169 arguments: json!({"path":"readme.md"}),
170 })
171 .await
172 .unwrap();
173 input
174 .send(AgentEvent::ToolComplete {
175 tool_call_id: "read".into(),
176 result: ToolResult::text(true, "ordinary Read result"),
177 })
178 .await
179 .unwrap();
180 input
181 .send(AgentEvent::HookLifecycle {
182 hook_name: private.into(),
183 point: bamboo_domain::AgentHookPoint::AfterToolExecution,
184 phase: private.into(),
185 duration_ms: 7,
186 decision: bamboo_domain::HookResult::WithContext {
187 result: Box::new(bamboo_domain::HookResult::Deny {
188 reason: private.into(),
189 }),
190 text: private.into(),
191 },
192 })
193 .await
194 .unwrap();
195 drop(input);
196
197 let mut visible = Vec::new();
198 for _ in 0..6 {
199 visible.push(
200 tokio::time::timeout(std::time::Duration::from_secs(5), broadcast_rx.recv())
201 .await
202 .unwrap()
203 .unwrap(),
204 );
205 }
206 for event in &visible[..3] {
207 assert!(!serde_json::to_string(event).unwrap().contains(private));
208 }
209 assert!(
210 matches!(&visible[2], AgentEvent::ToolComplete { result, .. }
211 if result.result == "Tool result hidden" && result.images.is_empty())
212 );
213 assert!(
214 matches!(&visible[4], AgentEvent::ToolComplete { result, .. }
215 if result.result == "ordinary Read result")
216 );
217 assert!(!serde_json::to_string(&visible[5])
218 .unwrap()
219 .contains(private));
220 assert!(matches!(
221 &visible[5],
222 AgentEvent::HookLifecycle {
223 point: bamboo_domain::AgentHookPoint::AfterToolExecution,
224 duration_ms: 7,
225 decision: bamboo_domain::HookResult::WithContext { .. },
226 ..
227 }
228 ));
229 forwarder.await.unwrap();
230 }
231
232 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
233 async fn hundreds_of_child_streams_progress_while_global_registry_is_locked() {
234 let runners = Arc::new(RwLock::new(HashMap::new()));
235 let mut streams = Vec::new();
236 for index in 0..512 {
237 let id = format!("child-{index}");
238 let mut runner = AgentRunner::new();
239 runner.status = super::super::runner_state::AgentStatus::Running;
240 let mut receiver = runner.event_sender.subscribe();
241 let sender = runner.event_sender.clone();
242 let run_id = runner.run_id.clone();
243 runners.write().await.insert(id.clone(), runner);
244 let (input, task) = create_event_forwarder(id, run_id, sender, runners.clone(), None);
245 assert!(matches!(
246 receiver.recv().await.unwrap(),
247 AgentEvent::ExecutionStarted { .. }
248 ));
249 streams.push((input, receiver, task));
250 }
251 let held_registry = runners.write().await;
252 let started = std::time::Instant::now();
253 tokio::time::timeout(std::time::Duration::from_secs(10), async {
254 futures::future::join_all(streams.iter_mut().map(|(input, receiver, _)| async move {
255 for _ in 0..64 {
256 input
257 .send(AgentEvent::Token {
258 content: "delta".into(),
259 })
260 .await
261 .unwrap();
262 }
263 for _ in 0..64 {
264 assert!(matches!(
265 receiver.recv().await.unwrap(),
266 AgentEvent::Token { .. }
267 ));
268 }
269 }))
270 .await;
271 })
272 .await
273 .expect("independent token streams must not wait for registry ownership");
274 eprintln!(
275 "512 child streams / 32768 tokens with held registry: {:?}",
276 started.elapsed()
277 );
278 assert!(held_registry
279 .values()
280 .all(|runner| runner.last_activity_at().is_some()));
281 drop(held_registry);
282 for (input, _, task) in streams {
283 drop(input);
284 task.await.unwrap();
285 }
286 }
287
288 #[tokio::test]
289 async fn child_approval_change_routes_to_parent_account_envelope() {
290 let (tx, mut rx) = mpsc::channel(4);
291 let event = AgentEvent::ChildApprovalChanged {
292 parent_session_id: "parent-1".into(),
293 child_session_id: "child-1".into(),
294 child_attempt: 1,
295 request_id: "req-1".into(),
296 version: 2,
297 status: "approved".into(),
298 reason: None,
299 tool_name: "Bash".into(),
300 permission: "execute".into(),
301 resource: "/tmp/x".into(),
302 created_at: "2026-01-01T00:00:00Z".into(),
303 resolved_at: Some("2026-01-01T00:00:01Z".into()),
304 };
305
306 mirror_to_account_feed(&Some(tx), "child-1", &event);
307 let (session_id, mirrored) = rx.recv().await.unwrap();
308 assert_eq!(session_id.as_deref(), Some("parent-1"));
309 assert!(matches!(mirrored, AgentEvent::ChildApprovalChanged { .. }));
310 }
311
312 #[tokio::test]
313 async fn child_browser_approval_is_redacted_before_broadcast_replay_and_account_feed() {
314 let session_id = "child-browser-private";
315 let private = "private-download-selector";
316 let (broadcast_tx, mut broadcast_rx) = broadcast::channel(8);
317 let mut runner = AgentRunner::new();
318 runner.status = super::super::runner_state::AgentStatus::Running;
319 runner.event_sender = broadcast_tx.clone();
320 let run_id = runner.run_id.clone();
321 let runners = Arc::new(RwLock::new(HashMap::from([(
322 session_id.to_string(),
323 runner,
324 )])));
325 let (account_tx, mut account_rx) = mpsc::channel(8);
326 let (event_tx, forwarder) = create_event_forwarder(
327 session_id.into(),
328 run_id,
329 broadcast_tx,
330 runners.clone(),
331 Some(account_tx),
332 );
333 assert!(matches!(
334 broadcast_rx.recv().await.unwrap(),
335 AgentEvent::ExecutionStarted { .. }
336 ));
337 event_tx
338 .send(AgentEvent::ChildApprovalChanged {
339 parent_session_id: "parent-private".into(),
340 child_session_id: session_id.into(),
341 child_attempt: 4,
342 request_id: "request-private".into(),
343 version: 7,
344 status: "denied".into(),
345 reason: Some(private.into()),
346 tool_name: "browser".into(),
347 permission: private.into(),
348 resource: format!("browser:17:download:css:{private}"),
349 created_at: "2026-09-24T00:00:00Z".into(),
350 resolved_at: Some("2026-09-24T00:00:01Z".into()),
351 })
352 .await
353 .unwrap();
354 let live = broadcast_rx.recv().await.unwrap();
355 assert!(!serde_json::to_string(&live).unwrap().contains(private));
356 let _started = account_rx.recv().await.unwrap();
357 let (route, feed) = account_rx.recv().await.unwrap();
358 assert_eq!(route.as_deref(), Some("parent-private"));
359 assert!(!serde_json::to_string(&feed).unwrap().contains(private));
360 assert!(
361 !serde_json::to_string(&runners.read().await[session_id].last_critical_events)
362 .unwrap()
363 .contains(private)
364 );
365 drop(event_tx);
366 forwarder.await.unwrap();
367 }
368
369 #[tokio::test]
370 async fn history_commit_barrier_acknowledges_actual_publication() {
371 let session_id = "session-history-barrier";
372 let (broadcast_tx, mut broadcast_rx) = broadcast::channel(16);
373 let mut runner = AgentRunner::new();
374 runner.status = super::super::runner_state::AgentStatus::Running;
375 runner.event_sender = broadcast_tx.clone();
376 let run_id = runner.run_id.clone();
377 let runners = Arc::new(RwLock::new(HashMap::from([(
378 session_id.to_string(),
379 runner,
380 )])));
381 let (event_tx, forwarder, mut barrier) = create_event_forwarder_with_history_commit_barrier(
382 session_id.to_string(),
383 run_id,
384 broadcast_tx,
385 runners,
386 None,
387 );
388
389 assert!(matches!(
390 broadcast_rx.recv().await.unwrap(),
391 AgentEvent::ExecutionStarted { .. }
392 ));
393 assert!(
394 barrier
395 .send_and_wait(&event_tx, session_id.to_string())
396 .await,
397 "producer must not continue until the forwarder publishes the barrier"
398 );
399 assert!(matches!(
400 broadcast_rx.recv().await.unwrap(),
401 AgentEvent::SessionHistoryCommitted { session_id: id } if id == session_id
402 ));
403
404 drop(event_tx);
405 forwarder.await.unwrap();
406 }
407
408 #[tokio::test]
409 async fn delayed_old_forwarder_cannot_publish_after_successor_reservation() {
410 let session_id = "session-generation";
411 let (broadcast_tx, mut broadcast_rx) = broadcast::channel(16);
412 let mut successor = AgentRunner::new();
413 successor.run_id = "run-new".to_string();
414 successor.status = super::super::runner_state::AgentStatus::Running;
415 successor.event_sender = broadcast_tx.clone();
416 let runners = Arc::new(RwLock::new(HashMap::from([(
417 session_id.to_string(),
418 successor,
419 )])));
420
421 broadcast_tx
424 .send(AgentEvent::ExecutionStarted {
425 run_id: "run-new".to_string(),
426 session_id: session_id.to_string(),
427 started_at: Utc::now().to_rfc3339(),
428 })
429 .unwrap();
430 let (old_tx, old_forwarder) = create_event_forwarder(
431 session_id.to_string(),
432 "run-old".to_string(),
433 broadcast_tx.clone(),
434 runners,
435 None,
436 );
437 let _ = old_tx
438 .send(AgentEvent::NeedClarification {
439 question: "stale".to_string(),
440 options: Some(vec!["A".to_string()]),
441 tool_call_id: Some("old-tool".to_string()),
442 tool_name: Some("ConclusionWithOptions".to_string()),
443 allow_custom: false,
444 source: Some(bamboo_agent_core::PendingQuestionSource::PauseTool),
445 })
446 .await;
447 drop(old_tx);
448 old_forwarder.await.unwrap();
449
450 assert!(matches!(
451 broadcast_rx.recv().await.unwrap(),
452 AgentEvent::ExecutionStarted { ref run_id, .. } if run_id == "run-new"
453 ));
454 assert!(
455 tokio::time::timeout(std::time::Duration::from_millis(50), broadcast_rx.recv())
456 .await
457 .is_err(),
458 "old Started/Need/Complete must all be suppressed"
459 );
460 }
461
462 #[tokio::test]
463 async fn wrapped_child_lifecycle_reaches_bounded_legacy_replay_consumer() {
464 let session_id = "root";
465 let (broadcast_tx, mut broadcast_rx) = broadcast::channel(8);
466 let mut runner = AgentRunner::new();
467 runner.status = super::super::runner_state::AgentStatus::Running;
468 runner.event_sender = broadcast_tx.clone();
469 let run_id = runner.run_id.clone();
470 let runners = Arc::new(RwLock::new(HashMap::from([(session_id.into(), runner)])));
471 let (input, forwarder) = create_event_forwarder(
472 session_id.into(),
473 run_id,
474 broadcast_tx,
475 runners.clone(),
476 None,
477 );
478 assert!(matches!(
479 broadcast_rx.recv().await.unwrap(),
480 AgentEvent::ExecutionStarted { .. }
481 ));
482 let legacy = AgentEvent::SubAgentEvent {
483 parent_session_id: "root".into(),
484 child_session_id: "parent".into(),
485 event: Box::new(AgentEvent::SubAgentStarted {
486 parent_session_id: "parent".into(),
487 child_session_id: "grandchild".into(),
488 title: Some("work".into()),
489 }),
490 };
491 input.send(legacy).await.unwrap();
492 assert!(matches!(
493 broadcast_rx.recv().await.unwrap(),
494 AgentEvent::SubAgentEvent { .. }
495 ));
496 let guard = runners.read().await;
497 let cached = &guard[session_id].last_critical_events;
498 assert_eq!(cached.len(), 1);
499 assert!(matches!(&cached[0], AgentEvent::SubAgentEvent { .. }));
500 drop(guard);
501 drop(input);
502 forwarder.await.unwrap();
503 }
504}
505
506pub fn create_event_forwarder(
515 session_id: String,
516 run_id: String,
517 broadcast_tx: broadcast::Sender<AgentEvent>,
518 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
519 account_feed_inbox: Option<AccountFeedInbox>,
520) -> (mpsc::Sender<AgentEvent>, tokio::task::JoinHandle<()>) {
521 let (event_tx, forwarder, _history_commit_barrier) =
522 create_event_forwarder_with_history_commit_barrier(
523 session_id,
524 run_id,
525 broadcast_tx,
526 runners,
527 account_feed_inbox,
528 );
529 (event_tx, forwarder)
530}
531
532pub fn create_event_forwarder_with_history_commit_barrier(
536 session_id: String,
537 run_id: String,
538 broadcast_tx: broadcast::Sender<AgentEvent>,
539 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
540 account_feed_inbox: Option<AccountFeedInbox>,
541) -> (
542 mpsc::Sender<AgentEvent>,
543 tokio::task::JoinHandle<()>,
544 HistoryCommitBarrier,
545) {
546 create_event_forwarder_with_root_actor(
547 session_id,
548 run_id,
549 broadcast_tx,
550 runners,
551 account_feed_inbox,
552 None,
553 )
554}
555
556async fn publish_forwarded_event(
557 root_actor: Option<&crate::events::RootActorEventPublication>,
558 session_id: &str,
559 event: &AgentEvent,
560 publish: Box<dyn FnOnce() -> bool + Send>,
561) -> std::io::Result<()> {
562 if let Some(owner) = root_actor {
563 owner
564 .publish(event.session_id().unwrap_or(session_id), event, publish)
565 .await
566 } else if publish() {
567 Ok(())
568 } else {
569 Err(std::io::Error::other("event publication was retired"))
570 }
571}
572
573pub fn create_event_forwarder_with_root_actor(
576 session_id: String,
577 run_id: String,
578 broadcast_tx: broadcast::Sender<AgentEvent>,
579 runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
580 account_feed_inbox: Option<AccountFeedInbox>,
581 root_actor: Option<crate::events::RootActorEventPublication>,
582) -> (
583 mpsc::Sender<AgentEvent>,
584 tokio::task::JoinHandle<()>,
585 HistoryCommitBarrier,
586) {
587 let (mpsc_tx, mut mpsc_rx) = mpsc::channel::<AgentEvent>(100);
588 let (history_commit_acknowledger, history_commit_barrier) = history_commit_barrier();
589 let forward = async move {
590 if root_actor
591 .as_ref()
592 .is_some_and(|owner| !owner.matches_execution(&session_id, &run_id))
593 {
594 tracing::warn!(%session_id, %run_id, "Root forwarder handoff does not match its execution");
595 return;
596 }
597 let legacy_feed = if root_actor.is_some() {
599 None
600 } else {
601 account_feed_inbox
602 };
603 let started_event = AgentEvent::ExecutionStarted {
604 run_id: run_id.clone(),
605 session_id: session_id.clone(),
606 started_at: Utc::now().to_rfc3339(),
607 };
608 let (publication, visible_messages) = {
609 let guard = runners.clone().read_owned().await;
610 let Some(runner) = guard
611 .get(&session_id)
612 .filter(|runner| runner.run_id == run_id)
613 else {
614 return;
615 };
616 let publication = runner.event_publication.clone();
617 let visible_messages = runner.visible_messages.clone();
618 let feed = legacy_feed.clone();
619 let sid = session_id.clone();
620 let tx = broadcast_tx.clone();
621 let frame = started_event.clone();
622 if let Err(error) = publish_forwarded_event(
623 root_actor.as_ref(),
624 &session_id,
625 &frame,
626 Box::new(move || {
627 let _guard = guard;
628 mirror_to_account_feed(&feed, &sid, &started_event);
629 let _ = tx.send(started_event);
630 true
631 }),
632 )
633 .await
634 {
635 tracing::warn!(%session_id, %run_id, %error, "Root Started publication rejected");
636 return;
637 }
638 (publication, visible_messages)
639 };
640 let mut tool_event_display = bamboo_agent_core::NativeToolEventDisplay::default();
641 while let Some(event) = mpsc_rx.recv().await {
642 match &event {
643 AgentEvent::VisibleMessageStart {
644 message_id,
645 created_at,
646 } => {
647 if !publication.publish(|| {
648 visible_messages.start(message_id.clone(), created_at.to_owned())
649 }) {
650 return;
651 }
652 continue;
653 }
654 AgentEvent::VisibleMessageDiscard { message_id } => {
655 if !publication.publish(|| visible_messages.discard(message_id)) {
656 return;
657 }
658 continue;
659 }
660 _ => {}
661 }
662 let event = tool_event_display.project(event);
663 let event = crate::external_agents::live::approval_event_for_display(event);
664 let needs_runner_update = event.is_replayable_session_state()
665 || matches!(
666 &event,
667 AgentEvent::TokenBudgetUpdated { .. }
668 | AgentEvent::ToolStart { .. }
669 | AgentEvent::ToolLifecycle { .. }
670 | AgentEvent::RunnerProgress { .. }
671 );
672 let is_history_commit = matches!(&event, AgentEvent::SessionHistoryCommitted { .. });
673 let frame = event.clone();
674 let feed = legacy_feed.clone();
675 let sid = session_id.clone();
676 let tx = broadcast_tx.clone();
677 let result = if !needs_runner_update {
678 let visible_token = match &event {
679 AgentEvent::Token { content } => Some(content.clone()),
680 _ => None,
681 };
682 let terminal_reason = visible_terminal_reason(&event);
683 let owner = root_actor
685 .as_ref()
686 .filter(|_| event.is_durable_change() || terminal_reason.is_some());
687 let publication = publication.clone();
688 let visible = visible_messages.clone();
689 publish_forwarded_event(
690 owner,
691 &session_id,
692 &frame,
693 Box::new(move || {
694 publication.publish(|| {
695 if let Some(content) = visible_token {
696 visible.append(content);
697 }
698 mirror_to_account_feed(&feed, &sid, &event);
699 let _ = tx.send(event);
700 if let Some(reason) = terminal_reason {
701 visible.mark_terminal(reason);
702 }
703 if is_history_commit {
704 visible.history_committed();
705 }
706 })
707 }),
708 )
709 .await
710 } else {
711 let mut guard = runners.clone().write_owned().await;
712 if !guard
713 .get(&session_id)
714 .is_some_and(|runner| runner.run_id == run_id)
715 {
716 return;
717 }
718 let publication = publication.clone();
719 publish_forwarded_event(
720 root_actor.as_ref(),
721 &session_id,
722 &frame,
723 Box::new(move || {
724 let runner = guard.get_mut(&sid).expect("retained exact runner guard");
725 runner.last_event_at = Some(Utc::now());
726 publication.touch();
727 if event.is_replayable_session_state() {
728 runner.push_critical_event(event.clone());
729 }
730 match &event {
731 AgentEvent::TokenBudgetUpdated { .. } => {
732 runner.last_budget_event = Some(event.clone())
733 }
734 AgentEvent::ToolStart { tool_name, .. } => {
735 runner.last_tool_name = Some(tool_name.clone());
736 runner.last_tool_phase = Some("begin".into());
737 }
738 AgentEvent::ToolLifecycle {
739 tool_name, phase, ..
740 } => {
741 runner.last_tool_name = Some(tool_name.clone());
742 runner.last_tool_phase = Some(phase.clone());
743 }
744 AgentEvent::RunnerProgress { round_count, .. } => {
745 runner.round_count = *round_count;
746 runner.visible_messages.begin_round(*round_count);
747 }
748 _ => {}
749 }
750 mirror_to_account_feed(&feed, &sid, &event);
751 let _ = tx.send(event);
752 true
753 }),
754 )
755 .await
756 };
757 if let Err(error) = result {
758 tracing::warn!(%session_id, %run_id, %error, "Root runtime publication rejected");
759 return;
760 }
761 if is_history_commit {
762 history_commit_acknowledger.acknowledge();
763 }
764 }
765 };
766 (mpsc_tx, tokio::spawn(forward), history_commit_barrier)
767}