Skip to main content

bamboo_engine/runtime/managers/adapters/
lifecycle.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use bamboo_agent_core::tools::ToolExecutor;
5use bamboo_agent_core::{AgentError, AgentEvent, Session};
6use bamboo_domain::{AgentRuntimeState, AgentStatusState};
7use bamboo_llm::LLMProvider;
8use tokio::sync::mpsc;
9use tokio_util::sync::CancellationToken;
10
11use crate::runtime::config::AgentLoopConfig;
12use crate::runtime::managers::lifecycle::LifecycleManager;
13use crate::runtime::runner::state_bridge;
14use crate::runtime::task_context::TaskLoopContext;
15use bamboo_metrics::MetricsCollector;
16
17/// Default lifecycle manager that delegates to existing runner functions.
18pub struct DefaultLifecycleManager {
19    llm: Arc<dyn LLMProvider>,
20}
21
22impl DefaultLifecycleManager {
23    pub fn new(llm: Arc<dyn LLMProvider>) -> Self {
24        Self { llm }
25    }
26}
27
28#[async_trait]
29impl LifecycleManager for DefaultLifecycleManager {
30    fn initialize_run(&self, _session: &Session, config: &AgentLoopConfig) -> AgentRuntimeState {
31        // `AgentRuntimeState::run_id` is the existing per-execution identity on
32        // this adapter path. Give every initialized run a fresh value so round
33        // counters can safely restart for the same session.
34        let mut state =
35            AgentRuntimeState::new(crate::runtime::runner::round_prelude::new_execution_id());
36        state.llm.model_name = config.model_name.clone();
37        state.llm.provider_name = config.provider_name.clone();
38        state.llm.fast_model_name = config.fast_model_name.clone();
39        state.llm.background_model_name = config.background_model_name.clone();
40        state.round.max_rounds = config.max_rounds as u32;
41        state.status = AgentStatusState::Initializing;
42        state
43    }
44
45    #[allow(clippy::too_many_arguments)]
46    async fn prepare_round(
47        &self,
48        session: &mut Session,
49        task_context: &mut Option<TaskLoopContext>,
50        runtime_state: &mut AgentRuntimeState,
51        round: usize,
52        max_rounds: usize,
53        config: &AgentLoopConfig,
54        cancel_token: &CancellationToken,
55        metrics_collector: Option<&MetricsCollector>,
56        session_id: &str,
57        model_name: &str,
58        tools: &dyn ToolExecutor,
59        _llm: &dyn LLMProvider,
60    ) -> Result<String, AgentError> {
61        crate::runtime::runner::round_prelude::prepare_round(
62            session,
63            task_context,
64            config,
65            self.llm.clone(),
66            tools,
67            &crate::runtime::runner::round_prelude::RoundPreludeFrame {
68                execution_id: &runtime_state.run_id,
69                round,
70                max_rounds,
71                debug_enabled: false, // debug logging handled at runner level, not via adapter
72                cancel_token,
73                metrics_collector,
74                session_id,
75                model_name,
76            },
77        )
78        .await
79    }
80
81    async fn handle_round_outcome(
82        &self,
83        session: &mut Session,
84        runtime_state: &mut AgentRuntimeState,
85        _task_context: &mut Option<TaskLoopContext>,
86        round: usize,
87        should_break: bool,
88    ) -> Result<bool, AgentError> {
89        runtime_state.round.current_round = round as u32;
90
91        if should_break {
92            runtime_state.status = AgentStatusState::Finalizing;
93        } else if round as u32 >= runtime_state.round.max_rounds {
94            tracing::info!(
95                "[{}] Reached max rounds ({})",
96                session.id,
97                runtime_state.round.max_rounds
98            );
99            return Ok(true);
100        }
101
102        state_bridge::write_runtime_state(session, runtime_state);
103        Ok(should_break)
104    }
105
106    #[allow(clippy::too_many_arguments)]
107    async fn finalize_run(
108        &self,
109        session: &mut Session,
110        runtime_state: &mut AgentRuntimeState,
111        event_tx: &mpsc::Sender<AgentEvent>,
112        session_id: &str,
113        config: &AgentLoopConfig,
114        metrics_collector: Option<&MetricsCollector>,
115        task_context: Option<TaskLoopContext>,
116    ) {
117        runtime_state.status = AgentStatusState::Completed;
118        state_bridge::write_runtime_state(session, runtime_state);
119
120        crate::runtime::runner::session_finalize::finalize_session(
121            task_context,
122            session,
123            event_tx,
124            session_id,
125            config,
126            metrics_collector,
127            false,
128            runtime_state,
129        )
130        .await;
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use bamboo_agent_core::Message;
138    use bamboo_llm::provider::LLMStream;
139    use futures::stream;
140
141    struct UnusedProvider;
142
143    #[async_trait]
144    impl LLMProvider for UnusedProvider {
145        async fn chat_stream(
146            &self,
147            _messages: &[Message],
148            _tools: &[bamboo_agent_core::tools::ToolSchema],
149            _max_output_tokens: Option<u32>,
150            _model: &str,
151        ) -> bamboo_llm::provider::Result<LLMStream> {
152            Ok(Box::pin(stream::iter(vec![Ok(bamboo_llm::LLMChunk::Done)])))
153        }
154    }
155
156    #[test]
157    fn initialize_run_assigns_a_fresh_execution_identity() {
158        let manager = DefaultLifecycleManager::new(Arc::new(UnusedProvider));
159        let session = Session::new("same-session", "model");
160        let config = AgentLoopConfig::default();
161
162        let first = manager.initialize_run(&session, &config);
163        let second = manager.initialize_run(&session, &config);
164
165        assert!(!first.run_id.is_empty());
166        assert_ne!(first.run_id, session.id);
167        assert_ne!(first.run_id, second.run_id);
168    }
169}