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.max_rounds as u32;
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: 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 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 round as u32 >= runtime_state.round.max_rounds {
tracing::info!(
"[{}] Reached max rounds ({})",
session.id,
runtime_state.round.max_rounds
);
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 {
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,
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()),
max_rounds: 1,
..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,
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);
}
}