use std::sync::Arc;
use async_trait::async_trait;
use bamboo_agent_core::tools::ToolExecutor;
use bamboo_agent_core::{AgentError, AgentEvent, Session};
use bamboo_domain::{AgentRuntimeState, AgentStatusState};
use bamboo_llm::LLMProvider;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use crate::runtime::config::AgentLoopConfig;
use crate::runtime::managers::lifecycle::LifecycleManager;
use crate::runtime::runner::state_bridge;
use crate::runtime::task_context::TaskLoopContext;
use bamboo_metrics::MetricsCollector;
pub struct DefaultLifecycleManager {
llm: Arc<dyn LLMProvider>,
}
impl DefaultLifecycleManager {
pub fn new(llm: Arc<dyn LLMProvider>) -> Self {
Self { llm }
}
}
#[async_trait]
impl LifecycleManager for DefaultLifecycleManager {
fn initialize_run(&self, _session: &Session, config: &AgentLoopConfig) -> AgentRuntimeState {
let mut state =
AgentRuntimeState::new(crate::runtime::runner::round_prelude::new_execution_id());
state.llm.model_name = config.model_name.clone();
state.llm.provider_name = config.provider_name.clone();
state.llm.fast_model_name = config.fast_model_name.clone();
state.llm.background_model_name = config.background_model_name.clone();
state.round.max_rounds = config.run_budget.max_rounds;
state.status = AgentStatusState::Initializing;
state
}
#[allow(clippy::too_many_arguments)]
async fn prepare_round(
&self,
session: &mut Session,
task_context: &mut Option<TaskLoopContext>,
runtime_state: &mut AgentRuntimeState,
round: usize,
max_rounds: Option<usize>,
config: &AgentLoopConfig,
cancel_token: &CancellationToken,
metrics_collector: Option<&MetricsCollector>,
session_id: &str,
model_name: &str,
tools: &dyn ToolExecutor,
_llm: &dyn LLMProvider,
) -> Result<String, AgentError> {
let execution_id = runtime_state.run_id.clone();
crate::runtime::runner::round_prelude::prepare_round(
session,
task_context,
runtime_state,
config,
self.llm.clone(),
tools,
&crate::runtime::runner::round_prelude::RoundPreludeFrame {
execution_id: &execution_id,
round,
max_rounds,
debug_enabled: false, cancel_token,
metrics_collector,
session_id,
model_name,
},
)
.await
}
async fn prepare_round_with_observation(
&self,
session: &mut Session,
task_context: &mut Option<TaskLoopContext>,
runtime_state: &mut AgentRuntimeState,
round: usize,
max_rounds: Option<usize>,
config: &AgentLoopConfig,
cancel_token: &CancellationToken,
metrics_collector: Option<&MetricsCollector>,
session_id: &str,
model_name: &str,
tools: &dyn ToolExecutor,
_llm: &dyn LLMProvider,
) -> Result<crate::runtime::managers::lifecycle::ObservedRoundPreparation, AgentError> {
let execution_id = runtime_state.run_id.clone();
crate::runtime::runner::round_prelude::prepare_round_with_observation(
session,
task_context,
runtime_state,
config,
self.llm.clone(),
tools,
&crate::runtime::runner::round_prelude::RoundPreludeFrame {
execution_id: &execution_id,
round,
max_rounds,
debug_enabled: false, cancel_token,
metrics_collector,
session_id,
model_name,
},
)
.await
}
async fn handle_round_outcome(
&self,
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
_task_context: &mut Option<TaskLoopContext>,
round: usize,
should_break: bool,
) -> Result<bool, AgentError> {
runtime_state.round.current_round = round as u32;
if should_break {
runtime_state.status = AgentStatusState::Finalizing;
} else if runtime_state
.round
.max_rounds
.is_some_and(|max_rounds| round as u32 >= max_rounds)
{
tracing::info!(
"[{}] Reached max rounds ({})",
session.id,
runtime_state.round.max_rounds.unwrap_or_default()
);
return Ok(true);
}
state_bridge::write_runtime_state(session, runtime_state);
Ok(should_break)
}
#[allow(clippy::too_many_arguments)]
async fn finalize_run(
&self,
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
event_tx: &mpsc::Sender<AgentEvent>,
session_id: &str,
config: &AgentLoopConfig,
metrics_collector: Option<&MetricsCollector>,
task_context: Option<TaskLoopContext>,
) {
runtime_state.status = AgentStatusState::Completed;
state_bridge::write_runtime_state(session, runtime_state);
crate::runtime::runner::session_finalize::finalize_session(
task_context,
session,
event_tx,
session_id,
config,
metrics_collector,
false,
runtime_state,
)
.await;
}
}
#[cfg(test)]
mod tests {
struct QlOldCustom {
calls: std::sync::atomic::AtomicUsize,
fail: bool,
builtin: DefaultLifecycleManager,
}
#[async_trait]
impl LifecycleManager for QlOldCustom {
fn initialize_run(&self, s: &Session, c: &AgentLoopConfig) -> AgentRuntimeState {
self.builtin.initialize_run(s, c)
}
async fn prepare_round(
&self,
s: &mut Session,
_: &mut Option<TaskLoopContext>,
_: &mut AgentRuntimeState,
_: usize,
_: Option<usize>,
_: &AgentLoopConfig,
_: &CancellationToken,
_: Option<&MetricsCollector>,
_: &str,
_: &str,
_: &dyn ToolExecutor,
_: &dyn LLMProvider,
) -> Result<String, AgentError> {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
s.metadata
.insert("old-custom-effect".into(), "exact".into());
if self.fail {
Err(AgentError::Tool("old-custom-error".into()))
} else {
Ok("old-custom-round".into())
}
}
async fn handle_round_outcome(
&self,
s: &mut Session,
r: &mut AgentRuntimeState,
t: &mut Option<TaskLoopContext>,
n: usize,
b: bool,
) -> Result<bool, AgentError> {
self.builtin.handle_round_outcome(s, r, t, n, b).await
}
async fn finalize_run(
&self,
s: &mut Session,
r: &mut AgentRuntimeState,
e: &mpsc::Sender<AgentEvent>,
id: &str,
c: &AgentLoopConfig,
m: Option<&MetricsCollector>,
t: Option<TaskLoopContext>,
) {
self.builtin.finalize_run(s, r, e, id, c, m, t).await;
}
}
#[tokio::test]
async fn ql_old_custom_default_companion_calls_once_and_preserves_return_error_effect() {
for fail in [false, true] {
let manager = QlOldCustom {
calls: Default::default(),
fail,
builtin: DefaultLifecycleManager::new(Arc::new(UnusedProvider)),
};
let mut session = Session::new("old-custom", "model");
let config = AgentLoopConfig::default();
let mut runtime = manager.initialize_run(&session, &config);
let result = manager
.prepare_round_with_observation(
&mut session,
&mut None,
&mut runtime,
0,
None,
&config,
&CancellationToken::new(),
None,
"old-custom",
"model",
&EmptyTools,
&UnusedProvider,
)
.await;
assert_eq!(manager.calls.load(std::sync::atomic::Ordering::SeqCst), 1);
assert_eq!(session.metadata["old-custom-effect"], "exact");
match result {
Ok(prepared) => {
assert!(!fail);
assert_eq!(prepared.round_id, "old-custom-round");
assert!(prepared.observation.is_unavailable());
}
Err(AgentError::Tool(error)) => {
assert!(fail);
assert_eq!(error, "old-custom-error");
}
other => panic!("changed old result: {other:?}"),
}
}
}
#[tokio::test]
async fn ql_same_builtin_manager_returns_separate_data_to_real_run_owners() {
let home = tempfile::tempdir().unwrap();
let store = Arc::new(
bamboo_storage::SessionStoreV2::new(home.path().into())
.await
.unwrap(),
);
let storage: Arc<dyn Storage> = store.clone();
let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
Arc::new(bamboo_storage::LockedSessionStore::new(storage.clone()));
let inbox: Arc<dyn SessionInboxPort> = Arc::new(bamboo_storage::FileSessionInbox::new(
store,
SessionInboxLimits::default(),
));
let config = AgentLoopConfig {
storage: Some(storage.clone()),
persistence: Some(persistence),
session_inbox: Some(inbox.clone()),
..Default::default()
};
let manager = DefaultLifecycleManager::new(Arc::new(UnusedProvider));
let mut owners = Vec::new();
for id in ["ql-manager-a", "ql-manager-b"] {
let mut session = Session::new(id, "model");
session.add_message(Message::system("system"));
storage.save_session(&session).await.unwrap();
let envelope = SessionMessageEnvelope::user_input(id, format!("input {id}"));
let receipt = inbox.deliver(&envelope).await.unwrap();
inbox
.mark_activation_eligible(
id,
receipt.generation,
SessionActivationPolicy::InterruptSpecificWait,
)
.await
.unwrap();
let mut runtime = manager.initialize_run(&session, &config);
let execution = runtime.run_id.clone();
let prepared = manager
.prepare_round_with_observation(
&mut session,
&mut None,
&mut runtime,
0,
None,
&config,
&CancellationToken::new(),
None,
id,
"model",
&EmptyTools,
&UnusedProvider,
)
.await
.unwrap();
assert_eq!(prepared.round_id, format!("{id}-run-{execution}-round-1"));
let mut current = None;
prepared
.observation
.update_current(&mut current, id, &execution);
let batch = current.as_ref().unwrap();
assert_eq!(batch.records()[0].input_id, envelope.id.as_str());
assert_eq!(batch.session_id(), id);
let prepared = manager
.prepare_round_with_observation(
&mut session,
&mut None,
&mut runtime,
1,
None,
&config,
&CancellationToken::new(),
None,
id,
"model",
&EmptyTools,
&UnusedProvider,
)
.await
.unwrap();
assert!(prepared.observation.is_successful_no_new_input());
prepared
.observation
.update_current(&mut current, id, &execution);
owners.push(current.unwrap());
}
assert_ne!(owners[0].execution_id(), owners[1].execution_id());
assert_ne!(owners[0].records().as_ptr(), owners[1].records().as_ptr());
assert_ne!(
owners[0].records()[0].input_id,
owners[1].records()[0].input_id
);
}
use super::*;
use bamboo_agent_core::storage::Storage;
use bamboo_agent_core::tools::{ToolCall, ToolError, ToolResult, ToolSchema};
use bamboo_agent_core::Message;
use bamboo_domain::{
SessionActivationPolicy, SessionInboxLimits, SessionInboxPort, SessionMessageEnvelope,
};
use bamboo_llm::provider::LLMStream;
use futures::stream;
struct UnusedProvider;
#[async_trait]
impl LLMProvider for UnusedProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> bamboo_llm::provider::Result<LLMStream> {
Ok(Box::pin(stream::iter(vec![Ok(bamboo_llm::LLMChunk::Done)])))
}
}
struct EmptyTools;
#[async_trait]
impl ToolExecutor for EmptyTools {
async fn execute(&self, call: &ToolCall) -> std::result::Result<ToolResult, ToolError> {
Err(ToolError::NotFound(call.function.name.clone()))
}
fn list_tools(&self) -> Vec<ToolSchema> {
Vec::new()
}
}
#[test]
fn initialize_run_assigns_a_fresh_execution_identity() {
let manager = DefaultLifecycleManager::new(Arc::new(UnusedProvider));
let session = Session::new("same-session", "model");
let config = AgentLoopConfig::default();
let first = manager.initialize_run(&session, &config);
let second = manager.initialize_run(&session, &config);
assert!(!first.run_id.is_empty());
assert_ne!(first.run_id, session.id);
assert_ne!(first.run_id, second.run_id);
}
#[tokio::test]
async fn adapter_admits_typed_inbox_input_then_cancels_before_prompt_context() {
let directory = tempfile::tempdir().unwrap();
let store = Arc::new(
bamboo_storage::SessionStoreV2::new(directory.path().to_path_buf())
.await
.unwrap(),
);
let storage: Arc<dyn Storage> = store.clone();
let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
Arc::new(bamboo_storage::LockedSessionStore::new(storage.clone()));
let inbox: Arc<dyn SessionInboxPort> = Arc::new(bamboo_storage::FileSessionInbox::new(
store,
SessionInboxLimits::default(),
));
let mut persisted = Session::new("adapter-boundary-before-recall", "model");
persisted.add_message(Message::system("base prompt"));
storage.save_session(&persisted).await.unwrap();
let mut running = persisted;
let content = "admit this before observing cancellation";
let envelope = SessionMessageEnvelope::user_input(&running.id, content);
let receipt = inbox.deliver(&envelope).await.unwrap();
inbox
.mark_activation_eligible(
&running.id,
receipt.generation,
SessionActivationPolicy::InterruptSpecificWait,
)
.await
.unwrap();
let config = AgentLoopConfig {
storage: Some(storage),
persistence: Some(persistence),
session_inbox: Some(inbox),
app_data_dir: Some(directory.path().to_path_buf()),
..AgentLoopConfig::default()
};
let provider: Arc<dyn LLMProvider> = Arc::new(UnusedProvider);
let manager = DefaultLifecycleManager::new(provider.clone());
let mut runtime_state = manager.initialize_run(&running, &config);
let mut task_context = None;
let cancel = CancellationToken::new();
cancel.cancel();
let session_id = running.id.clone();
let error = manager
.prepare_round(
&mut running,
&mut task_context,
&mut runtime_state,
0,
Some(1),
&config,
&cancel,
None,
&session_id,
"model",
&EmptyTools,
provider.as_ref(),
)
.await
.expect_err("a cancelled adapter round must stop before prompt context");
assert!(matches!(error, AgentError::Cancelled));
assert!(running.messages.iter().any(|message| {
message.id == envelope.id.as_str()
&& message.role == bamboo_domain::Role::User
&& message.content == content
}));
assert!(!running
.metadata
.contains_key(crate::runtime::runner::prompt_context::PROMPT_MEMORY_OBSERVABILITY_KEY));
}
#[tokio::test]
async fn adapter_typed_inbox_input_drives_current_round_memory_recall() {
let directory = tempfile::tempdir().unwrap();
let store = Arc::new(
bamboo_storage::SessionStoreV2::new(directory.path().to_path_buf())
.await
.unwrap(),
);
let storage: Arc<dyn Storage> = store.clone();
let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
Arc::new(bamboo_storage::LockedSessionStore::new(storage.clone()));
let inbox: Arc<dyn SessionInboxPort> = Arc::new(bamboo_storage::FileSessionInbox::new(
store,
SessionInboxLimits::default(),
));
let memory = bamboo_memory::memory_store::MemoryStore::new(directory.path().join("jiandu"));
memory
.write_memory(
bamboo_memory::memory_store::MemoryScope::Global,
None,
bamboo_memory::memory_store::DurableMemoryType::Reference,
"Adapter silver heron rule",
"The silver heron request must use the lifecycle memory boundary.",
&["silver".to_string(), "heron".to_string()],
Some("adapter-typed-inbox-recall"),
"model",
false,
None,
)
.await
.unwrap();
let mut persisted = Session::new("adapter-typed-inbox-recall", "model");
persisted.add_message(Message::system("base prompt"));
persisted.add_message(Message::user("unrelated earlier request"));
storage.save_session(&persisted).await.unwrap();
let mut running = persisted;
let query = "what is the adapter silver heron rule?";
let envelope = SessionMessageEnvelope::user_input(&running.id, query);
let receipt = inbox.deliver(&envelope).await.unwrap();
inbox
.mark_activation_eligible(
&running.id,
receipt.generation,
SessionActivationPolicy::InterruptSpecificWait,
)
.await
.unwrap();
let config = AgentLoopConfig {
storage: Some(storage),
persistence: Some(persistence),
session_inbox: Some(inbox),
app_data_dir: Some(directory.path().join("bamboo")),
memory_store: memory,
prompt_memory_flags: crate::runtime::config::PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: true,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
},
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_rounds: Some(1),
..Default::default()
},
..AgentLoopConfig::default()
};
let provider: Arc<dyn LLMProvider> = Arc::new(UnusedProvider);
let manager = DefaultLifecycleManager::new(provider.clone());
let mut runtime_state = manager.initialize_run(&running, &config);
let mut task_context = None;
let cancel = CancellationToken::new();
let session_id = running.id.clone();
manager
.prepare_round(
&mut running,
&mut task_context,
&mut runtime_state,
0,
Some(1),
&config,
&cancel,
None,
&session_id,
"model",
&EmptyTools,
provider.as_ref(),
)
.await
.expect("typed inbox recall should prepare the current round");
assert!(running.messages.iter().any(|message| {
message.id == envelope.id.as_str()
&& message.role == bamboo_domain::Role::User
&& message.content == query
}));
let rendered =
crate::runtime::runner::prompt_context::render_external_memory_section(&running)
.expect("current-round recall should render external memory");
assert!(rendered.contains("Adapter silver heron rule"));
assert!(
rendered.contains("The silver heron request must use the lifecycle memory boundary.")
);
let observability: bamboo_agent_core::PromptMemoryObservability = serde_json::from_str(
running
.metadata
.get(crate::runtime::runner::prompt_context::PROMPT_MEMORY_OBSERVABILITY_KEY)
.expect("prompt memory refresh should persist observability"),
)
.unwrap();
assert!(observability.latest_user_query_present);
assert_eq!(observability.relevant_memory_status, "lexical");
assert_eq!(observability.relevant_memory_count, 1);
}
}