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_infrastructure::LLMProvider;
8use tokio::sync::mpsc;
9use tokio_util::sync::CancellationToken;
10
11use crate::metrics::MetricsCollector;
12use crate::runtime::config::AgentLoopConfig;
13use crate::runtime::managers::lifecycle::LifecycleManager;
14use crate::runtime::runner::state_bridge;
15use crate::runtime::task_context::TaskLoopContext;
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        let mut state = AgentRuntimeState::new(&session.id);
32        state.llm.model_name = config.model_name.clone();
33        state.llm.provider_name = config.provider_name.clone();
34        state.llm.fast_model_name = config.fast_model_name.clone();
35        state.llm.background_model_name = config.background_model_name.clone();
36        state.round.max_rounds = config.max_rounds as u32;
37        state.status = AgentStatusState::Initializing;
38        state
39    }
40
41    #[allow(clippy::too_many_arguments)]
42    async fn prepare_round(
43        &self,
44        session: &mut Session,
45        task_context: &mut Option<TaskLoopContext>,
46        _runtime_state: &mut AgentRuntimeState,
47        round: usize,
48        max_rounds: usize,
49        config: &AgentLoopConfig,
50        cancel_token: &CancellationToken,
51        metrics_collector: Option<&MetricsCollector>,
52        session_id: &str,
53        model_name: &str,
54        tools: &dyn ToolExecutor,
55        _llm: &dyn LLMProvider,
56    ) -> Result<String, AgentError> {
57        crate::runtime::runner::round_prelude::prepare_round(
58            session,
59            task_context,
60            config,
61            self.llm.clone(),
62            tools,
63            &crate::runtime::runner::round_prelude::RoundPreludeFrame {
64                round,
65                max_rounds,
66                debug_enabled: false, // debug logging handled at runner level, not via adapter
67                cancel_token,
68                metrics_collector,
69                session_id,
70                model_name,
71            },
72        )
73        .await
74    }
75
76    async fn handle_round_outcome(
77        &self,
78        session: &mut Session,
79        runtime_state: &mut AgentRuntimeState,
80        _task_context: &mut Option<TaskLoopContext>,
81        round: usize,
82        should_break: bool,
83    ) -> Result<bool, AgentError> {
84        runtime_state.round.current_round = round as u32;
85
86        if should_break {
87            runtime_state.status = AgentStatusState::Finalizing;
88        } else if round as u32 >= runtime_state.round.max_rounds {
89            tracing::info!(
90                "[{}] Reached max rounds ({})",
91                session.id,
92                runtime_state.round.max_rounds
93            );
94            return Ok(true);
95        }
96
97        state_bridge::write_runtime_state(session, runtime_state);
98        Ok(should_break)
99    }
100
101    #[allow(clippy::too_many_arguments)]
102    async fn finalize_run(
103        &self,
104        session: &mut Session,
105        runtime_state: &mut AgentRuntimeState,
106        event_tx: &mpsc::Sender<AgentEvent>,
107        session_id: &str,
108        config: &AgentLoopConfig,
109        metrics_collector: Option<&MetricsCollector>,
110        task_context: Option<TaskLoopContext>,
111    ) {
112        runtime_state.status = AgentStatusState::Completed;
113        state_bridge::write_runtime_state(session, runtime_state);
114
115        crate::runtime::runner::session_finalize::finalize_session(
116            task_context,
117            session,
118            event_tx,
119            session_id,
120            config,
121            metrics_collector,
122            false,
123            runtime_state,
124        )
125        .await;
126    }
127}