bamboo_engine/runtime/execution/
event_forwarder.rs

1//! Event forwarding from MPSC to broadcast channels.
2//!
3//! Creates an MPSC channel for agent loop events and spawns a background task
4//! that relays events to the session's broadcast sender while tracking runner
5//! diagnostic state (budget events, tool execution, round progress).
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use chrono::Utc;
11use tokio::sync::{broadcast, mpsc, RwLock};
12
13use bamboo_agent_core::AgentEvent;
14
15use super::runner_state::AgentRunner;
16
17/// Inbox to the account-wide change feed: `(session_id, event)` before the
18/// writer assigns a seq. Threaded as `Option` so engine-internal callers that
19/// have no feed (tests, standalone embeddings) can pass `None`. Defined here so
20/// the engine stays free of any `bamboo-server` dependency.
21pub type AccountFeedInbox = mpsc::Sender<(Option<String>, AgentEvent)>;
22
23/// Forward a durable change event onto the account feed, if an inbox is wired.
24///
25/// Ephemeral events (tokens, heartbeats, …) are filtered out before any clone,
26/// so this is near-free on the hot path. `session_id` is supplied explicitly so
27/// terminal events (which carry no id) still route to the right session.
28fn mirror_to_account_feed(inbox: &Option<AccountFeedInbox>, session_id: &str, event: &AgentEvent) {
29    if let Some(inbox) = inbox {
30        if event.is_durable_change() {
31            let route_session_id = event.session_id().unwrap_or(session_id);
32            let _ = inbox.try_send((Some(route_session_id.to_string()), event.clone()));
33        }
34    }
35}
36
37#[cfg(test)]
38#[allow(clippy::items_after_test_module)]
39mod tests {
40    use super::*;
41
42    #[tokio::test]
43    async fn child_approval_change_routes_to_parent_account_envelope() {
44        let (tx, mut rx) = mpsc::channel(4);
45        let event = AgentEvent::ChildApprovalChanged {
46            parent_session_id: "parent-1".into(),
47            child_session_id: "child-1".into(),
48            child_attempt: 1,
49            request_id: "req-1".into(),
50            version: 2,
51            status: "approved".into(),
52            reason: None,
53            tool_name: "Bash".into(),
54            permission: "execute".into(),
55            resource: "/tmp/x".into(),
56            created_at: "2026-01-01T00:00:00Z".into(),
57            resolved_at: Some("2026-01-01T00:00:01Z".into()),
58        };
59
60        mirror_to_account_feed(&Some(tx), "child-1", &event);
61        let (session_id, mirrored) = rx.recv().await.unwrap();
62        assert_eq!(session_id.as_deref(), Some("parent-1"));
63        assert!(matches!(mirrored, AgentEvent::ChildApprovalChanged { .. }));
64    }
65
66    #[tokio::test]
67    async fn delayed_old_forwarder_cannot_publish_after_successor_reservation() {
68        let session_id = "session-generation";
69        let (broadcast_tx, mut broadcast_rx) = broadcast::channel(16);
70        let mut successor = AgentRunner::new();
71        successor.run_id = "run-new".to_string();
72        successor.status = super::super::runner_state::AgentStatus::Running;
73        successor.event_sender = broadcast_tx.clone();
74        let runners = Arc::new(RwLock::new(HashMap::from([(
75            session_id.to_string(),
76            successor,
77        )])));
78
79        // The successor is already visible on the shared transport before an
80        // old forwarder task finally gets CPU time.
81        broadcast_tx
82            .send(AgentEvent::ExecutionStarted {
83                run_id: "run-new".to_string(),
84                session_id: session_id.to_string(),
85                started_at: Utc::now().to_rfc3339(),
86            })
87            .unwrap();
88        let (old_tx, old_forwarder) = create_event_forwarder(
89            session_id.to_string(),
90            "run-old".to_string(),
91            broadcast_tx.clone(),
92            runners,
93            None,
94        );
95        let _ = old_tx
96            .send(AgentEvent::NeedClarification {
97                question: "stale".to_string(),
98                options: Some(vec!["A".to_string()]),
99                tool_call_id: Some("old-tool".to_string()),
100                tool_name: Some("ConclusionWithOptions".to_string()),
101                allow_custom: false,
102                source: Some(bamboo_agent_core::PendingQuestionSource::PauseTool),
103            })
104            .await;
105        drop(old_tx);
106        old_forwarder.await.unwrap();
107
108        assert!(matches!(
109            broadcast_rx.recv().await.unwrap(),
110            AgentEvent::ExecutionStarted { ref run_id, .. } if run_id == "run-new"
111        ));
112        assert!(
113            tokio::time::timeout(std::time::Duration::from_millis(50), broadcast_rx.recv())
114                .await
115                .is_err(),
116            "old Started/Need/Complete must all be suppressed"
117        );
118    }
119}
120
121/// Create an MPSC channel for agent events and spawn a forwarding task
122/// that relays events to the broadcast sender while tracking runner
123/// diagnostic fields for live visibility.
124///
125/// `account_feed_inbox`, when present, also mirrors durable change events onto
126/// the account-wide feed for resumable multi-client sync.
127///
128/// Returns `(mpsc_tx, forwarder_handle)`.
129pub fn create_event_forwarder(
130    session_id: String,
131    run_id: String,
132    broadcast_tx: broadcast::Sender<AgentEvent>,
133    runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
134    account_feed_inbox: Option<AccountFeedInbox>,
135) -> (mpsc::Sender<AgentEvent>, tokio::task::JoinHandle<()>) {
136    let (mpsc_tx, mut mpsc_rx) = mpsc::channel::<AgentEvent>(100);
137
138    let forwarder = tokio::spawn(async move {
139        // The exact reservation generation is captured synchronously by the
140        // caller. Never re-read the replaceable runner registry here: this
141        // task may be scheduled only after a clarification handoff installs a
142        // successor, which would mis-tag the old terminal as the new run.
143        let started_event = AgentEvent::ExecutionStarted {
144            run_id: run_id.clone(),
145            session_id: session_id.clone(),
146            started_at: Utc::now().to_rfc3339(),
147        };
148        {
149            let runners = runners.read().await;
150            if runners
151                .get(&session_id)
152                .is_none_or(|runner| runner.run_id != run_id)
153            {
154                return;
155            }
156            mirror_to_account_feed(&account_feed_inbox, &session_id, &started_event);
157            let _ = broadcast_tx.send(started_event);
158        }
159
160        while let Some(event) = mpsc_rx.recv().await {
161            let mut runners = runners.write().await;
162            let Some(runner) = runners
163                .get_mut(&session_id)
164                .filter(|runner| runner.run_id == run_id)
165            else {
166                // A clarification handoff installed a successor before this
167                // delayed forwarder/frame ran. Drop the entire stale stream;
168                // broadcasting even its Started/Need would corrupt the shared
169                // session generation state.
170                return;
171            };
172            runner.last_event_at = Some(Utc::now());
173
174            // Cache live state before publication so a subscriber installed
175            // between a clarification pause and its response sees the exact
176            // boundary. This generic forwarder powers Connect, schedules,
177            // SDK spawn, and child-resume paths, so it must preserve the same
178            // replay invariant as the server-owned forwarder.
179            if event.is_replayable_session_state() {
180                runner.push_critical_event(event.clone());
181            }
182
183            match &event {
184                AgentEvent::TokenBudgetUpdated { .. } => {
185                    runner.last_budget_event = Some(event.clone());
186                }
187                AgentEvent::ToolStart { tool_name, .. } => {
188                    runner.last_tool_name = Some(tool_name.clone());
189                    runner.last_tool_phase = Some("begin".to_string());
190                }
191                AgentEvent::ToolLifecycle {
192                    tool_name, phase, ..
193                } => {
194                    runner.last_tool_name = Some(tool_name.clone());
195                    runner.last_tool_phase = Some(phase.clone());
196                }
197                AgentEvent::RunnerProgress { round_count, .. } => {
198                    runner.round_count = *round_count;
199                }
200                _ => {}
201            }
202            mirror_to_account_feed(&account_feed_inbox, &session_id, &event);
203            let _ = broadcast_tx.send(event);
204        }
205    });
206
207    (mpsc_tx, forwarder)
208}
bamboo_engine/runtime/execution/event_forwarder.rs

bamboo_engine/runtime/execution/
event_forwarder.rs
