bamboo_engine/runtime/execution/
runner_state.rs

1//! Runner state types for background agent execution.
2//!
3//! Provides the `AgentRunner` and `AgentStatus` types that track the lifecycle
4//! of an in-progress agent execution. These are used by the execution
5//! orchestration layer across all background paths (HTTP execute, spawn, schedule).
6
7use chrono::{DateTime, Utc};
8use std::sync::Arc;
9use tokio::sync::broadcast;
10use tokio_util::sync::CancellationToken;
11use uuid::Uuid;
12
13use bamboo_agent_core::AgentEvent;
14
15fn subagent_lifecycle_key(event: &AgentEvent) -> Option<(&str, &str)> {
16    event
17        .legacy_child_lifecycle_observation()
18        .map(|observation| (observation.parent_session_id, observation.child_session_id))
19}
20
21/// Status of an agent execution runner.
22///
23/// Represents the lifecycle state of an agent run from initialization
24/// through completion or error.
25#[derive(Debug, Clone)]
26pub enum AgentStatus {
27    /// Agent is initialized but not yet running.
28    Pending,
29
30    /// Agent is currently executing.
31    Running,
32
33    /// Agent completed successfully.
34    Completed,
35
36    /// Agent execution was cancelled by user.
37    Cancelled,
38
39    /// Agent execution failed with an error message.
40    Error(String),
41}
42
43/// Runner that manages agent execution for a session.
44///
45/// Each active agent run has an associated `AgentRunner` that coordinates
46/// event broadcasting, cancellation, and status tracking.
47///
48/// # Event Broadcasting
49///
50/// Uses a broadcast channel to support multiple subscribers watching
51/// the same agent run simultaneously.
52///
53/// # Cancellation
54///
55/// Provides a cancellation token that can be used to gracefully stop
56/// an in-progress agent execution.
57#[derive(Debug, Clone)]
58pub struct AgentRunner {
59    /// Generation fence and activity clock shared only by this run's producers.
60    pub event_publication: Arc<super::event_publication::EventPublication>,
61
62    /// Message-only visible-text replay and broadcast for this exact run.
63    pub visible_messages: Arc<super::visible_messages::VisibleMessageStream>,
64    /// Broadcast sender for agent events.
65    ///
66    /// Allows multiple clients to subscribe to agent events
67    /// via `event_sender.subscribe()`.
68    pub event_sender: broadcast::Sender<AgentEvent>,
69
70    /// Cancellation token for graceful shutdown.
71    ///
72    /// When triggered, the agent should stop execution at the
73    /// next safe point.
74    pub cancel_token: CancellationToken,
75
76    /// Current status of the agent run.
77    pub status: AgentStatus,
78
79    /// Timestamp when the run was started.
80    pub started_at: DateTime<Utc>,
81
82    /// Timestamp when the run completed (if finished).
83    pub completed_at: Option<DateTime<Utc>>,
84
85    /// Last token budget event to replay for new subscribers.
86    ///
87    /// When a new client subscribes to an ongoing run, this
88    /// allows them to receive the most recent token usage info.
89    pub last_budget_event: Option<AgentEvent>,
90
91    /// Small ring of critical state events (TaskListUpdated, SubAgent*, etc.)
92    /// cached for replay to late/reconnecting subscribers.
93    ///
94    /// Bounded to [`CRITICAL_EVENTS_CAPACITY`] entries; oldest are evicted.
95    pub last_critical_events: Vec<AgentEvent>,
96
97    /// Name of the most recently executed tool (if any).
98    /// Updated live during execution for diagnostic visibility.
99    pub last_tool_name: Option<String>,
100
101    /// Phase of the most recently executed tool: "begin", "finished", or "error".
102    /// Updated live during execution for diagnostic visibility.
103    pub last_tool_phase: Option<String>,
104
105    /// Timestamp of the last event received during this run.
106    /// Updated live during execution for liveness checks.
107    pub last_event_at: Option<DateTime<Utc>>,
108
109    /// Number of completed rounds (turns) so far.
110    /// Updated live during execution for progress tracking.
111    pub round_count: u32,
112
113    /// Unique identifier for this execution run.
114    /// Generated fresh for every `try_reserve_runner` call so that
115    /// frontend SSE events can be matched to the correct run even
116    /// across reconnects.
117    pub run_id: String,
118}
119
120impl Default for AgentRunner {
121    fn default() -> Self {
122        Self::new()
123    }
124}
125
126impl AgentRunner {
127    /// Includes lock-free event traffic as well as legacy diagnostic updates.
128    pub fn last_activity_at(&self) -> Option<DateTime<Utc>> {
129        self.last_event_at
130            .max(self.event_publication.last_event_at())
131    }
132    /// Broadcast channel capacity for agent events.
133    pub const EVENT_CHANNEL_CAPACITY: usize = 1000;
134
135    /// Maximum number of critical events cached for late-subscriber replay.
136    pub const CRITICAL_EVENTS_CAPACITY: usize = 32;
137
138    /// Create a new agent runner with default settings.
139    ///
140    /// Initializes a broadcast channel, a fresh cancellation token,
141    /// and `Pending` status.
142    pub fn new() -> Self {
143        let (event_sender, _) = broadcast::channel(Self::EVENT_CHANNEL_CAPACITY);
144        Self {
145            event_sender,
146            event_publication: Arc::default(),
147            visible_messages: Arc::default(),
148            cancel_token: CancellationToken::new(),
149            status: AgentStatus::Pending,
150            started_at: Utc::now(),
151            completed_at: None,
152            last_budget_event: None,
153            last_critical_events: Vec::new(),
154            last_tool_name: None,
155            last_tool_phase: None,
156            last_event_at: None,
157            round_count: 0,
158            run_id: Uuid::new_v4().to_string(),
159        }
160    }
161
162    /// Push a critical state event into the bounded replay cache.
163    ///
164    /// Sub-agent lifecycle entries are snapshots keyed by stable child id, not
165    /// an event log: a rerunnable/resident child can produce many generations,
166    /// and replaying older Start/Complete pairs makes the current generation
167    /// ambiguous to reconnecting clients. Keep only the newest lifecycle state
168    /// for each child while live subscribers still receive every event.
169    ///
170    /// If the remaining cache is full, the oldest entry is evicted.
171    pub fn push_critical_event(&mut self, event: AgentEvent) {
172        let lifecycle_id = match &event {
173            AgentEvent::WorkflowActivated { event_id, .. }
174            | AgentEvent::WorkflowDeactivated { event_id, .. } => Some(event_id),
175            _ => None,
176        };
177        if lifecycle_id.is_some_and(|event_id| {
178            self.last_critical_events
179                .iter()
180                .any(|existing| match existing {
181                    AgentEvent::WorkflowActivated {
182                        event_id: existing_id,
183                        ..
184                    }
185                    | AgentEvent::WorkflowDeactivated {
186                        event_id: existing_id,
187                        ..
188                    } => existing_id == event_id,
189                    _ => false,
190                })
191        }) {
192            return;
193        }
194        if let Some(key) = subagent_lifecycle_key(&event) {
195            self.last_critical_events
196                .retain(|existing| subagent_lifecycle_key(existing) != Some(key));
197        }
198        if self.last_critical_events.len() >= Self::CRITICAL_EVENTS_CAPACITY {
199            self.last_critical_events.remove(0);
200        }
201        self.last_critical_events.push(event);
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    fn started(child_session_id: &str, title: &str) -> AgentEvent {
210        AgentEvent::SubAgentStarted {
211            parent_session_id: "parent".to_string(),
212            child_session_id: child_session_id.to_string(),
213            title: Some(title.to_string()),
214        }
215    }
216
217    fn completed(child_session_id: &str, status: &str) -> AgentEvent {
218        AgentEvent::SubAgentCompleted {
219            parent_session_id: "parent".to_string(),
220            child_session_id: child_session_id.to_string(),
221            status: status.to_string(),
222            error: None,
223        }
224    }
225
226    fn wrapped(inner: AgentEvent) -> AgentEvent {
227        AgentEvent::SubAgentEvent {
228            parent_session_id: "root".into(),
229            child_session_id: "parent".into(),
230            event: Box::new(inner),
231        }
232    }
233
234    #[test]
235    fn critical_replay_keeps_only_latest_generation_state_per_child() {
236        let mut runner = AgentRunner::new();
237
238        // A stable resident child may be reused repeatedly inside one parent
239        // runner. Reconnect must see S3 only, never the ambiguous historical
240        // sequence S1,C1,S2,C2,S3.
241        for event in [
242            started("resident", "generation-1"),
243            completed("resident", "completed"),
244            started("resident", "generation-2"),
245            completed("resident", "completed"),
246            started("resident", "generation-3"),
247        ] {
248            runner.push_critical_event(event);
249        }
250
251        assert_eq!(runner.last_critical_events.len(), 1);
252        assert!(matches!(
253            &runner.last_critical_events[0],
254            AgentEvent::SubAgentStarted {
255                child_session_id,
256                title: Some(title),
257                ..
258            } if child_session_id == "resident" && title == "generation-3"
259        ));
260
261        runner.push_critical_event(completed("resident", "completed"));
262        assert_eq!(runner.last_critical_events.len(), 1);
263        assert!(matches!(
264            &runner.last_critical_events[0],
265            AgentEvent::SubAgentCompleted {
266                child_session_id,
267                status,
268                ..
269            } if child_session_id == "resident" && status == "completed"
270        ));
271    }
272
273    #[test]
274    fn subagent_coalescing_preserves_other_children_and_recency() {
275        let mut runner = AgentRunner::new();
276        runner.push_critical_event(started("child-a", "a1"));
277        runner.push_critical_event(started("child-b", "b1"));
278        runner.push_critical_event(completed("child-a", "completed"));
279
280        assert_eq!(runner.last_critical_events.len(), 2);
281        assert!(matches!(
282            &runner.last_critical_events[0],
283            AgentEvent::SubAgentStarted {
284                child_session_id, ..
285            } if child_session_id == "child-b"
286        ));
287        assert!(matches!(
288            &runner.last_critical_events[1],
289            AgentEvent::SubAgentCompleted {
290                child_session_id, ..
291            } if child_session_id == "child-a"
292        ));
293    }
294
295    #[test]
296    fn wrapped_legacy_lifecycle_replays_once_and_cannot_double_count_direct_lifecycle() {
297        let mut runner = AgentRunner::new();
298        let legacy_start = wrapped(started("child-a", "initial"));
299        assert!(legacy_start.is_replayable_session_state());
300        runner.push_critical_event(legacy_start);
301        runner.push_critical_event(started("child-a", "direct"));
302        assert_eq!(runner.last_critical_events.len(), 1);
303        assert!(matches!(
304            &runner.last_critical_events[0],
305            AgentEvent::SubAgentStarted { title: Some(title), .. } if title == "direct"
306        ));
307        runner.push_critical_event(wrapped(completed("child-a", "completed")));
308        assert_eq!(runner.last_critical_events.len(), 1);
309        assert!(matches!(
310            &runner.last_critical_events[0],
311            AgentEvent::SubAgentEvent { .. }
312        ));
313        assert!(runner.last_critical_events.len() <= AgentRunner::CRITICAL_EVENTS_CAPACITY);
314
315        let invalid = AgentEvent::SubAgentEvent {
316            parent_session_id: "root".into(),
317            child_session_id: "unrelated".into(),
318            event: Box::new(started("child-b", "forged")),
319        };
320        assert!(!invalid.is_replayable_session_state());
321        assert!(invalid.legacy_child_lifecycle_observation().is_none());
322    }
323}
bamboo_engine/runtime/execution/runner_state.rs

bamboo_engine/runtime/execution/
runner_state.rs
