bamboo_engine/runtime/managers/adapters/
lifecycle.rs1use 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
17pub 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 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, 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}