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.run_budget.max_rounds;
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        // Round cap; `None` = unlimited.
53        max_rounds: Option<usize>,
54        config: &AgentLoopConfig,
55        cancel_token: &CancellationToken,
56        metrics_collector: Option<&MetricsCollector>,
57        session_id: &str,
58        model_name: &str,
59        tools: &dyn ToolExecutor,
60        _llm: &dyn LLMProvider,
61    ) -> Result<String, AgentError> {
62        let execution_id = runtime_state.run_id.clone();
63        crate::runtime::runner::round_prelude::prepare_round(
64            session,
65            task_context,
66            runtime_state,
67            config,
68            self.llm.clone(),
69            tools,
70            &crate::runtime::runner::round_prelude::RoundPreludeFrame {
71                execution_id: &execution_id,
72                round,
73                max_rounds,
74                debug_enabled: false, // debug logging handled at runner level, not via adapter
75                cancel_token,
76                metrics_collector,
77                session_id,
78                model_name,
79            },
80        )
81        .await
82    }
83
84    async fn prepare_round_with_observation(
85        &self,
86        session: &mut Session,
87        task_context: &mut Option<TaskLoopContext>,
88        runtime_state: &mut AgentRuntimeState,
89        round: usize,
90        // Round cap; `None` = unlimited.
91        max_rounds: Option<usize>,
92        config: &AgentLoopConfig,
93        cancel_token: &CancellationToken,
94        metrics_collector: Option<&MetricsCollector>,
95        session_id: &str,
96        model_name: &str,
97        tools: &dyn ToolExecutor,
98        _llm: &dyn LLMProvider,
99    ) -> Result<crate::runtime::managers::lifecycle::ObservedRoundPreparation, AgentError> {
100        let execution_id = runtime_state.run_id.clone();
101        crate::runtime::runner::round_prelude::prepare_round_with_observation(
102            session,
103            task_context,
104            runtime_state,
105            config,
106            self.llm.clone(),
107            tools,
108            &crate::runtime::runner::round_prelude::RoundPreludeFrame {
109                execution_id: &execution_id,
110                round,
111                max_rounds,
112                debug_enabled: false, // debug logging handled at runner level, not via adapter
113                cancel_token,
114                metrics_collector,
115                session_id,
116                model_name,
117            },
118        )
119        .await
120    }
121
122    async fn handle_round_outcome(
123        &self,
124        session: &mut Session,
125        runtime_state: &mut AgentRuntimeState,
126        _task_context: &mut Option<TaskLoopContext>,
127        round: usize,
128        should_break: bool,
129    ) -> Result<bool, AgentError> {
130        runtime_state.round.current_round = round as u32;
131
132        if should_break {
133            runtime_state.status = AgentStatusState::Finalizing;
134        } else if runtime_state
135            .round
136            .max_rounds
137            .is_some_and(|max_rounds| round as u32 >= max_rounds)
138        {
139            tracing::info!(
140                "[{}] Reached max rounds ({})",
141                session.id,
142                runtime_state.round.max_rounds.unwrap_or_default()
143            );
144            return Ok(true);
145        }
146
147        state_bridge::write_runtime_state(session, runtime_state);
148        Ok(should_break)
149    }
150
151    #[allow(clippy::too_many_arguments)]
152    async fn finalize_run(
153        &self,
154        session: &mut Session,
155        runtime_state: &mut AgentRuntimeState,
156        event_tx: &mpsc::Sender<AgentEvent>,
157        session_id: &str,
158        config: &AgentLoopConfig,
159        metrics_collector: Option<&MetricsCollector>,
160        task_context: Option<TaskLoopContext>,
161    ) {
162        runtime_state.status = AgentStatusState::Completed;
163        state_bridge::write_runtime_state(session, runtime_state);
164
165        crate::runtime::runner::session_finalize::finalize_session(
166            task_context,
167            session,
168            event_tx,
169            session_id,
170            config,
171            metrics_collector,
172            false,
173            runtime_state,
174        )
175        .await;
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    struct QlOldCustom {
182        calls: std::sync::atomic::AtomicUsize,
183        fail: bool,
184        builtin: DefaultLifecycleManager,
185    }
186
187    #[async_trait]
188    impl LifecycleManager for QlOldCustom {
189        fn initialize_run(&self, s: &Session, c: &AgentLoopConfig) -> AgentRuntimeState {
190            self.builtin.initialize_run(s, c)
191        }
192        async fn prepare_round(
193            &self,
194            s: &mut Session,
195            _: &mut Option<TaskLoopContext>,
196            _: &mut AgentRuntimeState,
197            _: usize,
198            _: Option<usize>,
199            _: &AgentLoopConfig,
200            _: &CancellationToken,
201            _: Option<&MetricsCollector>,
202            _: &str,
203            _: &str,
204            _: &dyn ToolExecutor,
205            _: &dyn LLMProvider,
206        ) -> Result<String, AgentError> {
207            self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
208            s.metadata
209                .insert("old-custom-effect".into(), "exact".into());
210            if self.fail {
211                Err(AgentError::Tool("old-custom-error".into()))
212            } else {
213                Ok("old-custom-round".into())
214            }
215        }
216        async fn handle_round_outcome(
217            &self,
218            s: &mut Session,
219            r: &mut AgentRuntimeState,
220            t: &mut Option<TaskLoopContext>,
221            n: usize,
222            b: bool,
223        ) -> Result<bool, AgentError> {
224            self.builtin.handle_round_outcome(s, r, t, n, b).await
225        }
226        async fn finalize_run(
227            &self,
228            s: &mut Session,
229            r: &mut AgentRuntimeState,
230            e: &mpsc::Sender<AgentEvent>,
231            id: &str,
232            c: &AgentLoopConfig,
233            m: Option<&MetricsCollector>,
234            t: Option<TaskLoopContext>,
235        ) {
236            self.builtin.finalize_run(s, r, e, id, c, m, t).await;
237        }
238    }
239
240    #[tokio::test]
241    async fn ql_old_custom_default_companion_calls_once_and_preserves_return_error_effect() {
242        for fail in [false, true] {
243            let manager = QlOldCustom {
244                calls: Default::default(),
245                fail,
246                builtin: DefaultLifecycleManager::new(Arc::new(UnusedProvider)),
247            };
248            let mut session = Session::new("old-custom", "model");
249            let config = AgentLoopConfig::default();
250            let mut runtime = manager.initialize_run(&session, &config);
251            let result = manager
252                .prepare_round_with_observation(
253                    &mut session,
254                    &mut None,
255                    &mut runtime,
256                    0,
257                    None,
258                    &config,
259                    &CancellationToken::new(),
260                    None,
261                    "old-custom",
262                    "model",
263                    &EmptyTools,
264                    &UnusedProvider,
265                )
266                .await;
267            assert_eq!(manager.calls.load(std::sync::atomic::Ordering::SeqCst), 1);
268            assert_eq!(session.metadata["old-custom-effect"], "exact");
269            match result {
270                Ok(prepared) => {
271                    assert!(!fail);
272                    assert_eq!(prepared.round_id, "old-custom-round");
273                    assert!(prepared.observation.is_unavailable());
274                }
275                Err(AgentError::Tool(error)) => {
276                    assert!(fail);
277                    assert_eq!(error, "old-custom-error");
278                }
279                other => panic!("changed old result: {other:?}"),
280            }
281        }
282    }
283
284    #[tokio::test]
285    async fn ql_same_builtin_manager_returns_separate_data_to_real_run_owners() {
286        let home = tempfile::tempdir().unwrap();
287        let store = Arc::new(
288            bamboo_storage::SessionStoreV2::new(home.path().into())
289                .await
290                .unwrap(),
291        );
292        let storage: Arc<dyn Storage> = store.clone();
293        let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
294            Arc::new(bamboo_storage::LockedSessionStore::new(storage.clone()));
295        let inbox: Arc<dyn SessionInboxPort> = Arc::new(bamboo_storage::FileSessionInbox::new(
296            store,
297            SessionInboxLimits::default(),
298        ));
299        let config = AgentLoopConfig {
300            storage: Some(storage.clone()),
301            persistence: Some(persistence),
302            session_inbox: Some(inbox.clone()),
303            ..Default::default()
304        };
305        let manager = DefaultLifecycleManager::new(Arc::new(UnusedProvider));
306        let mut owners = Vec::new();
307        for id in ["ql-manager-a", "ql-manager-b"] {
308            let mut session = Session::new(id, "model");
309            session.add_message(Message::system("system"));
310            storage.save_session(&session).await.unwrap();
311            let envelope = SessionMessageEnvelope::user_input(id, format!("input {id}"));
312            let receipt = inbox.deliver(&envelope).await.unwrap();
313            inbox
314                .mark_activation_eligible(
315                    id,
316                    receipt.generation,
317                    SessionActivationPolicy::InterruptSpecificWait,
318                )
319                .await
320                .unwrap();
321            let mut runtime = manager.initialize_run(&session, &config);
322            let execution = runtime.run_id.clone();
323            let prepared = manager
324                .prepare_round_with_observation(
325                    &mut session,
326                    &mut None,
327                    &mut runtime,
328                    0,
329                    None,
330                    &config,
331                    &CancellationToken::new(),
332                    None,
333                    id,
334                    "model",
335                    &EmptyTools,
336                    &UnusedProvider,
337                )
338                .await
339                .unwrap();
340            assert_eq!(prepared.round_id, format!("{id}-run-{execution}-round-1"));
341            let mut current = None;
342            prepared
343                .observation
344                .update_current(&mut current, id, &execution);
345            let batch = current.as_ref().unwrap();
346            assert_eq!(batch.records()[0].input_id, envelope.id.as_str());
347            assert_eq!(batch.session_id(), id);
348            let prepared = manager
349                .prepare_round_with_observation(
350                    &mut session,
351                    &mut None,
352                    &mut runtime,
353                    1,
354                    None,
355                    &config,
356                    &CancellationToken::new(),
357                    None,
358                    id,
359                    "model",
360                    &EmptyTools,
361                    &UnusedProvider,
362                )
363                .await
364                .unwrap();
365            assert!(prepared.observation.is_successful_no_new_input());
366            prepared
367                .observation
368                .update_current(&mut current, id, &execution);
369            owners.push(current.unwrap());
370        }
371        assert_ne!(owners[0].execution_id(), owners[1].execution_id());
372        assert_ne!(owners[0].records().as_ptr(), owners[1].records().as_ptr());
373        assert_ne!(
374            owners[0].records()[0].input_id,
375            owners[1].records()[0].input_id
376        );
377    }
378
379    use super::*;
380    use bamboo_agent_core::storage::Storage;
381    use bamboo_agent_core::tools::{ToolCall, ToolError, ToolResult, ToolSchema};
382    use bamboo_agent_core::Message;
383    use bamboo_domain::{
384        SessionActivationPolicy, SessionInboxLimits, SessionInboxPort, SessionMessageEnvelope,
385    };
386    use bamboo_llm::provider::LLMStream;
387    use futures::stream;
388
389    struct UnusedProvider;
390
391    #[async_trait]
392    impl LLMProvider for UnusedProvider {
393        async fn chat_stream(
394            &self,
395            _messages: &[Message],
396            _tools: &[bamboo_agent_core::tools::ToolSchema],
397            _max_output_tokens: Option<u32>,
398            _model: &str,
399        ) -> bamboo_llm::provider::Result<LLMStream> {
400            Ok(Box::pin(stream::iter(vec![Ok(bamboo_llm::LLMChunk::Done)])))
401        }
402    }
403
404    struct EmptyTools;
405
406    #[async_trait]
407    impl ToolExecutor for EmptyTools {
408        async fn execute(&self, call: &ToolCall) -> std::result::Result<ToolResult, ToolError> {
409            Err(ToolError::NotFound(call.function.name.clone()))
410        }
411
412        fn list_tools(&self) -> Vec<ToolSchema> {
413            Vec::new()
414        }
415    }
416
417    #[test]
418    fn initialize_run_assigns_a_fresh_execution_identity() {
419        let manager = DefaultLifecycleManager::new(Arc::new(UnusedProvider));
420        let session = Session::new("same-session", "model");
421        let config = AgentLoopConfig::default();
422
423        let first = manager.initialize_run(&session, &config);
424        let second = manager.initialize_run(&session, &config);
425
426        assert!(!first.run_id.is_empty());
427        assert_ne!(first.run_id, session.id);
428        assert_ne!(first.run_id, second.run_id);
429    }
430
431    #[tokio::test]
432    async fn adapter_admits_typed_inbox_input_then_cancels_before_prompt_context() {
433        let directory = tempfile::tempdir().unwrap();
434        let store = Arc::new(
435            bamboo_storage::SessionStoreV2::new(directory.path().to_path_buf())
436                .await
437                .unwrap(),
438        );
439        let storage: Arc<dyn Storage> = store.clone();
440        let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
441            Arc::new(bamboo_storage::LockedSessionStore::new(storage.clone()));
442        let inbox: Arc<dyn SessionInboxPort> = Arc::new(bamboo_storage::FileSessionInbox::new(
443            store,
444            SessionInboxLimits::default(),
445        ));
446        let mut persisted = Session::new("adapter-boundary-before-recall", "model");
447        persisted.add_message(Message::system("base prompt"));
448        storage.save_session(&persisted).await.unwrap();
449
450        let mut running = persisted;
451        let content = "admit this before observing cancellation";
452        let envelope = SessionMessageEnvelope::user_input(&running.id, content);
453        let receipt = inbox.deliver(&envelope).await.unwrap();
454        inbox
455            .mark_activation_eligible(
456                &running.id,
457                receipt.generation,
458                SessionActivationPolicy::InterruptSpecificWait,
459            )
460            .await
461            .unwrap();
462        let config = AgentLoopConfig {
463            storage: Some(storage),
464            persistence: Some(persistence),
465            session_inbox: Some(inbox),
466            app_data_dir: Some(directory.path().to_path_buf()),
467            ..AgentLoopConfig::default()
468        };
469        let provider: Arc<dyn LLMProvider> = Arc::new(UnusedProvider);
470        let manager = DefaultLifecycleManager::new(provider.clone());
471        let mut runtime_state = manager.initialize_run(&running, &config);
472        let mut task_context = None;
473        let cancel = CancellationToken::new();
474        cancel.cancel();
475        let session_id = running.id.clone();
476
477        let error = manager
478            .prepare_round(
479                &mut running,
480                &mut task_context,
481                &mut runtime_state,
482                0,
483                Some(1),
484                &config,
485                &cancel,
486                None,
487                &session_id,
488                "model",
489                &EmptyTools,
490                provider.as_ref(),
491            )
492            .await
493            .expect_err("a cancelled adapter round must stop before prompt context");
494
495        assert!(matches!(error, AgentError::Cancelled));
496        assert!(running.messages.iter().any(|message| {
497            message.id == envelope.id.as_str()
498                && message.role == bamboo_domain::Role::User
499                && message.content == content
500        }));
501        assert!(!running
502            .metadata
503            .contains_key(crate::runtime::runner::prompt_context::PROMPT_MEMORY_OBSERVABILITY_KEY));
504    }
505
506    #[tokio::test]
507    async fn adapter_typed_inbox_input_drives_current_round_memory_recall() {
508        let directory = tempfile::tempdir().unwrap();
509        let store = Arc::new(
510            bamboo_storage::SessionStoreV2::new(directory.path().to_path_buf())
511                .await
512                .unwrap(),
513        );
514        let storage: Arc<dyn Storage> = store.clone();
515        let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
516            Arc::new(bamboo_storage::LockedSessionStore::new(storage.clone()));
517        let inbox: Arc<dyn SessionInboxPort> = Arc::new(bamboo_storage::FileSessionInbox::new(
518            store,
519            SessionInboxLimits::default(),
520        ));
521
522        let memory = bamboo_memory::memory_store::MemoryStore::new(directory.path().join("jiandu"));
523        memory
524            .write_memory(
525                bamboo_memory::memory_store::MemoryScope::Global,
526                None,
527                bamboo_memory::memory_store::DurableMemoryType::Reference,
528                "Adapter silver heron rule",
529                "The silver heron request must use the lifecycle memory boundary.",
530                &["silver".to_string(), "heron".to_string()],
531                Some("adapter-typed-inbox-recall"),
532                "model",
533                false,
534                None,
535            )
536            .await
537            .unwrap();
538
539        let mut persisted = Session::new("adapter-typed-inbox-recall", "model");
540        persisted.add_message(Message::system("base prompt"));
541        persisted.add_message(Message::user("unrelated earlier request"));
542        storage.save_session(&persisted).await.unwrap();
543
544        let mut running = persisted;
545        let query = "what is the adapter silver heron rule?";
546        let envelope = SessionMessageEnvelope::user_input(&running.id, query);
547        let receipt = inbox.deliver(&envelope).await.unwrap();
548        inbox
549            .mark_activation_eligible(
550                &running.id,
551                receipt.generation,
552                SessionActivationPolicy::InterruptSpecificWait,
553            )
554            .await
555            .unwrap();
556        let config = AgentLoopConfig {
557            storage: Some(storage),
558            persistence: Some(persistence),
559            session_inbox: Some(inbox),
560            app_data_dir: Some(directory.path().join("bamboo")),
561            memory_store: memory,
562            prompt_memory_flags: crate::runtime::config::PromptMemoryFlags {
563                project_prompt_injection: false,
564                relevant_recall: true,
565                relevant_recall_rerank: false,
566                project_first_dream: false,
567                ledger_agenda: false,
568            },
569            model_name: Some("model".to_string()),
570            run_budget: bamboo_config::RunBudgetConfig {
571                max_rounds: Some(1),
572                ..Default::default()
573            },
574            ..AgentLoopConfig::default()
575        };
576        let provider: Arc<dyn LLMProvider> = Arc::new(UnusedProvider);
577        let manager = DefaultLifecycleManager::new(provider.clone());
578        let mut runtime_state = manager.initialize_run(&running, &config);
579        let mut task_context = None;
580        let cancel = CancellationToken::new();
581        let session_id = running.id.clone();
582
583        manager
584            .prepare_round(
585                &mut running,
586                &mut task_context,
587                &mut runtime_state,
588                0,
589                Some(1),
590                &config,
591                &cancel,
592                None,
593                &session_id,
594                "model",
595                &EmptyTools,
596                provider.as_ref(),
597            )
598            .await
599            .expect("typed inbox recall should prepare the current round");
600
601        assert!(running.messages.iter().any(|message| {
602            message.id == envelope.id.as_str()
603                && message.role == bamboo_domain::Role::User
604                && message.content == query
605        }));
606        let rendered =
607            crate::runtime::runner::prompt_context::render_external_memory_section(&running)
608                .expect("current-round recall should render external memory");
609        assert!(rendered.contains("Adapter silver heron rule"));
610        assert!(
611            rendered.contains("The silver heron request must use the lifecycle memory boundary.")
612        );
613
614        let observability: bamboo_agent_core::PromptMemoryObservability = serde_json::from_str(
615            running
616                .metadata
617                .get(crate::runtime::runner::prompt_context::PROMPT_MEMORY_OBSERVABILITY_KEY)
618                .expect("prompt memory refresh should persist observability"),
619        )
620        .unwrap();
621        assert!(observability.latest_user_query_present);
622        assert_eq!(observability.relevant_memory_status, "lexical");
623        assert_eq!(observability.relevant_memory_count, 1);
624    }
625}
bamboo_engine/runtime/managers/adapters/lifecycle.rs

bamboo_engine/runtime/managers/adapters/
lifecycle.rs
