use std::collections::{BTreeMap, BTreeSet};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use chrono::Utc;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use crate::runtime::config::AgentLoopConfig;
use crate::runtime::guardian_state::{
ensure_guardian_state, guardian_read_only_disabled_tools, write_guardian_config,
write_guardian_state, GuardianPhase, GUARDIAN_REVIEW_RUBRIC,
};
use crate::runtime::runner::loop_execution::startup::{
resolve_auxiliary_models, InFlightTaskEvaluation, LoopRunState,
};
use crate::runtime::runner::prompt_context::PromptMemoryRuntimeContext;
use crate::runtime::runner::session_setup::tool_schemas::{
resolve_available_tool_schemas_for_session, resolve_tool_schemas_for_round,
};
use crate::runtime::stream::handler::{StreamHandlingOutput, VisibleMessageIdentity};
use crate::runtime::task_context::TaskLoopContext;
use bamboo_agent_core::tools::ToolExecutor;
use bamboo_agent_core::{AgentError, AgentEvent, Message, Role, Session};
use bamboo_domain::session::runtime_state::{
AgentRuntimeState, AgentStatusState, ChildWaitPolicy, SuspensionState, WaitingForBashState,
WaitingForChildrenState,
};
use bamboo_domain::{
AgentHookPoint, CapabilityInvocationTarget, CapabilityLoadingClass, CapabilityLoadingMode,
CapabilityMatch, CapabilitySource, ClassifiedToolSchema, DiscoverCapabilitiesRequest,
EffectiveCallableSet, HookPayload, HookResult, ProviderFamily, ProviderProtocol,
ProviderTranscriptAuthor, ProviderTranscriptItem, ProviderTranscriptItemKind,
ProviderTranscriptOrigin,
};
use bamboo_llm::LLMProvider;
use bamboo_metrics::{
MetricsCollector, RoundStatus as MetricsRoundStatus, SessionStatus as MetricsSessionStatus,
TokenUsage as MetricsTokenUsage,
};
use super::super::to_event_token_usage;
use super::gold::{
apply_completed_gold_evaluation, evaluate_gold_terminal, poll_completed_gold_evaluation,
spawn_gold_evaluation_if_needed, start_queued_gold_evaluation_if_idle, GoldTerminalDecision,
};
use crate::runtime::runner::state_bridge;
const MAX_LLM_TURN_ATTEMPTS: usize = 3;
const LLM_RETRY_BASE_DELAY_MS: u64 = 400;
const STICKY_DISCOVERY_RUNTIME_KIND: &str = "sticky_capability_discovery";
const STICKY_DISCOVERY_RUNTIME_VERSION: u64 = 1;
const STICKY_DISCOVERY_RESULT_START: &str = "<loaded_tools>";
const STICKY_DISCOVERY_RESULT_END: &str = "</loaded_tools>";
#[cfg(test)]
const TEST_POST_LLM_RETRY_FAILURES_KEY: &str = "test.pipeline.post_llm_retry_failures";
#[cfg(test)]
fn take_test_post_llm_retry_failure(session: &mut Session) -> Option<AgentError> {
let remaining = session
.metadata
.get(TEST_POST_LLM_RETRY_FAILURES_KEY)?
.parse::<usize>()
.ok()?;
if remaining == 0 {
return None;
}
if remaining == 1 {
session.metadata.remove(TEST_POST_LLM_RETRY_FAILURES_KEY);
} else {
session.metadata.insert(
TEST_POST_LLM_RETRY_FAILURES_KEY.to_string(),
(remaining - 1).to_string(),
);
}
Some(AgentError::LLM(
"transient test-injected post-LLM handler failure".to_string(),
))
}
pub(in crate::runtime::runner) fn effective_callable_set_for_round(
session: &Session,
tool_schemas: &[bamboo_agent_core::tools::ToolSchema],
capability_loading_mode: CapabilityLoadingMode,
) -> EffectiveCallableSet {
if capability_loading_mode == CapabilityLoadingMode::LegacyFullCatalog {
if legacy_browser_needs_discovery(
session,
tool_schemas,
crate::runtime::runner::round_lifecycle::required_tool_for_session(session),
) {
let without_browser = tool_schemas
.iter()
.filter(|schema| schema.function.name != "browser")
.cloned()
.collect::<Vec<_>>();
return crate::runtime::runner::tool_execution::legacy_effective_callable_set(
&without_browser,
);
}
return crate::runtime::runner::tool_execution::legacy_effective_callable_set(tool_schemas);
}
let loaded_names = if capability_loading_mode == CapabilityLoadingMode::StickyFallback {
validated_sticky_fallback_loaded_tool_names(session)
} else {
let transcript = &session.provider_transcript;
let family = transcript.active_family();
let protocol = transcript.active_protocol();
let groups = match (
family,
protocol,
transcript.active_provider_boundary_sha256(),
) {
(Some(family), Some(protocol), Some(boundary)) => {
transcript.replayable_groups(family, protocol, boundary)
}
_ => Vec::new(),
};
match (family, protocol) {
(
Some(ProviderFamily::Anthropic),
Some(ProviderProtocol::AnthropicMessages2023_06_01),
) => bamboo_llm::providers::anthropic::validated_anthropic_loaded_tool_names(
groups.iter().copied(),
tool_schemas,
),
(Some(family @ ProviderFamily::OpenAi), Some(ProviderProtocol::OpenAiResponsesV1)) => {
bamboo_domain::validated_openai_loaded_tool_names(groups.iter().copied(), family)
}
_ => Vec::new(),
}
};
let catalog = tool_schemas
.iter()
.cloned()
.filter_map(ClassifiedToolSchema::new)
.collect::<Vec<_>>();
EffectiveCallableSet::from_catalog(
&catalog,
capability_loading_mode,
loaded_names.iter().map(String::as_str),
)
}
fn openai_client_tool_search_requests(
items: &[ProviderTranscriptItem],
) -> Result<Vec<(String, DiscoverCapabilitiesRequest)>, AgentError> {
items
.iter()
.filter(|item| {
item.family() == ProviderFamily::OpenAi
&& item.protocol() == ProviderProtocol::OpenAiResponsesV1
&& item.kind() == ProviderTranscriptItemKind::OpenAiToolSearchCall
&& item.payload()["execution"].as_str() == Some("client")
})
.map(|item| {
let call_id = item.payload()["call_id"]
.as_str()
.expect("validated client tool_search_call has call_id")
.to_string();
let request = serde_json::from_value::<DiscoverCapabilitiesRequest>(
item.payload()["arguments"].clone(),
)
.map_err(|error| {
AgentError::LLM(format!(
"OpenAI client tool-search arguments are invalid: {error}"
))
})?;
Ok((call_id, request))
})
.collect()
}
fn capability_source_name(source: CapabilitySource) -> &'static str {
match source {
CapabilitySource::Builtin => "builtin",
CapabilitySource::Server => "server",
CapabilitySource::Mcp => "mcp",
CapabilitySource::Custom => "custom",
CapabilitySource::Project => "project",
CapabilitySource::Workspace => "workspace",
CapabilitySource::User => "user",
CapabilitySource::Plugin => "plugin",
}
}
fn scope_discovered_gateway_schema(
entry: &ClassifiedToolSchema,
matches: &[&CapabilityMatch],
) -> bamboo_agent_core::tools::ToolSchema {
let mut schema = entry.schema().clone();
let (catalog_kind, id_property) = match entry.execution_name() {
"load_skill" => ("skill", "skill_id"),
"workflow_run" => ("workflow", "workflow_id"),
_ => return schema,
};
let mut ids = Vec::new();
let mut workflow_revisions = Vec::new();
let mut metadata = Vec::new();
for matched in matches {
let scoped = match &matched.invocation_target {
CapabilityInvocationTarget::Skill {
skill_id,
source,
revision,
..
} if catalog_kind == "skill" => Some((skill_id, *source, *revision)),
CapabilityInvocationTarget::Workflow {
workflow_id,
source,
revision,
..
} if catalog_kind == "workflow" => {
if !workflow_revisions.contains(revision) {
workflow_revisions.push(*revision);
}
Some((workflow_id, *source, *revision))
}
_ => None,
};
let Some((id, source, revision)) = scoped else {
continue;
};
if !ids.contains(id) {
ids.push(id.clone());
}
let detail = format!(
"{id} — {} — {} [revision={revision}, source={}]",
matched.display_name,
matched.summary,
capability_source_name(source)
);
if !metadata.contains(&detail) {
metadata.push(detail);
}
}
if ids.is_empty() {
return schema;
}
if !schema.function.parameters.is_object() {
schema.function.parameters = serde_json::json!({"type": "object"});
}
let parameters = schema
.function
.parameters
.as_object_mut()
.expect("gateway parameters were normalized to an object");
let properties = parameters
.entry("properties")
.or_insert_with(|| serde_json::json!({}));
if !properties.is_object() {
*properties = serde_json::json!({});
}
let properties = properties
.as_object_mut()
.expect("scoped gateway properties are an object");
let id_schema = properties
.entry(id_property)
.or_insert_with(|| serde_json::json!({"type": "string"}));
if !id_schema.is_object() {
*id_schema = serde_json::json!({"type": "string"});
}
id_schema["enum"] = serde_json::json!(ids);
id_schema["description"] = serde_json::json!(format!(
"Use only the {id_property} values advertised by this discovery output."
));
if catalog_kind == "workflow" && !workflow_revisions.is_empty() {
let revision_schema = properties
.entry("revision")
.or_insert_with(|| serde_json::json!({"type": "integer", "minimum": 1}));
if !revision_schema.is_object() {
*revision_schema = serde_json::json!({"type": "integer", "minimum": 1});
}
revision_schema["enum"] = serde_json::json!(workflow_revisions);
}
schema.function.description = format!(
"{}. Scoped {catalog_kind} matches in discovery relevance order: {}.",
entry.schema().function.description.trim_end_matches('.'),
metadata.join("; ")
);
schema
}
struct CompleteCapabilityDiscovery {
catalog: Vec<ClassifiedToolSchema>,
tools: bamboo_tools::tool_search::ToolSearchIndex,
commands: Option<crate::capability_discovery::CapabilityDiscoveryIndex>,
}
impl CompleteCapabilityDiscovery {
async fn new(
session: &Session,
config: &AgentLoopConfig,
tool_schemas: &[bamboo_agent_core::tools::ToolSchema],
browser_only: bool,
include_commands: bool,
) -> Result<Self, AgentError> {
let catalog = tool_schemas
.iter()
.cloned()
.filter_map(ClassifiedToolSchema::new)
.filter(|entry| session.allows_model_tool_execution(entry.execution_name()))
.collect::<Vec<_>>();
let searchable_tool_catalog = catalog.iter().filter(|entry| {
entry.loading_class() == CapabilityLoadingClass::Deferred
&& (!browser_only || entry.execution_name() == "browser")
});
let tools = bamboo_tools::tool_search::ToolSearchIndex::from_resolved_catalog(
searchable_tool_catalog,
);
if browser_only || !include_commands || session.model_orchestration_only_enabled() {
return Ok(Self {
catalog,
tools,
commands: None,
});
}
let (_, disabled_skill_ids) = config.resolve_disabled_filters();
let catalog_names = catalog
.iter()
.map(|entry| entry.execution_name())
.collect::<BTreeSet<_>>();
let eligibility = crate::capability_discovery::CapabilityDiscoveryEligibility {
disabled_skill_ids: disabled_skill_ids.into_owned(),
allowed_skill_ids: config
.selected_skill_ids
.as_ref()
.map(|ids| ids.iter().cloned().collect()),
skill_gateway_available: catalog_names.contains("load_skill"),
workflow_gateway_available: catalog_names.contains("workflow_run"),
..Default::default()
};
let commands = match crate::runtime::runner::session_setup::skill_context::resolve_skill_store_for_session(
config, session,
)
.await
.map_err(AgentError::Tool)?
{
Some(store) => {
crate::capability_discovery::CapabilityDiscoveryIndex::from_resolved_classified_store(
&[],
store.as_ref(),
&eligibility,
)
.await
}
None => {
let empty_skills = bamboo_skills::WorkflowCatalogSnapshot::default();
let empty_workflows = bamboo_skills::WorkflowCatalogSnapshot::default();
crate::capability_discovery::CapabilityDiscoveryIndex::from_snapshots(
Vec::new(),
&empty_skills,
&empty_workflows,
&eligibility,
)
}
};
Ok(Self {
catalog,
tools,
commands: Some(commands),
})
}
fn discover_complete_schemas(
&self,
request: &DiscoverCapabilitiesRequest,
) -> Result<Vec<bamboo_agent_core::tools::ToolSchema>, AgentError> {
if request.kinds.as_ref().is_some_and(|kinds| kinds.len() > 3) {
return Err(AgentError::LLM(
"capability discovery accepts at most three kinds".to_string(),
));
}
let search_tools = request
.kinds
.as_ref()
.is_none_or(|kinds| kinds.contains(&bamboo_domain::CapabilityKind::Tool));
let mut tools = if search_tools {
self.tools
.search(&request.query, request.limit)
.map_err(|error| AgentError::LLM(format!("tool search failed: {error}")))?
.into_iter()
.filter_map(|name| {
self.catalog
.iter()
.find(|entry| entry.execution_name() == name)
.map(|entry| entry.schema().clone())
})
.collect::<Vec<_>>()
} else {
self.tools
.search(&request.query, request.limit)
.map_err(|error| AgentError::LLM(format!("tool search failed: {error}")))?;
Vec::new()
};
let limit = request
.limit
.unwrap_or(bamboo_domain::MAX_DISCOVERY_RESULTS);
let Some(commands) = self.commands.as_ref() else {
return Ok(tools);
};
let mut command_request = request.clone();
command_request.kinds = request.kinds.as_ref().map(|kinds| {
kinds
.iter()
.copied()
.filter(|kind| *kind != bamboo_domain::CapabilityKind::Tool)
.collect()
});
command_request.limit = Some(bamboo_domain::MAX_DISCOVERY_RESULTS);
let result = commands
.discover(&command_request)
.map_err(|error| AgentError::LLM(format!("command discovery failed: {error}")))?;
let mut matches_by_function = Vec::<(String, Vec<_>)>::new();
for matched in &result.matches {
let name = match &matched.invocation_target {
CapabilityInvocationTarget::Tool { name }
| CapabilityInvocationTarget::Skill { name, .. }
| CapabilityInvocationTarget::Workflow { name, .. } => name,
};
if let Some((_, matches)) = matches_by_function
.iter_mut()
.find(|(function, _)| function == name)
{
matches.push(matched);
} else {
matches_by_function.push((name.clone(), vec![matched]));
}
}
let command_tools = matches_by_function
.into_iter()
.filter_map(|(name, matches)| {
let entry = self
.catalog
.iter()
.find(|entry| entry.execution_name() == name)?;
if entry.loading_class() != CapabilityLoadingClass::Deferred {
return None;
}
Some(scope_discovered_gateway_schema(entry, &matches))
})
.collect::<Vec<_>>();
for command in command_tools {
if let Some(existing) = tools
.iter_mut()
.find(|tool| tool.function.name == command.function.name)
{
*existing = command;
} else if tools.len() < limit {
tools.push(command);
}
}
Ok(tools)
}
}
fn sticky_discovery_call_ids(session: &Session) -> BTreeSet<&str> {
session
.messages
.iter()
.filter(|message| matches!(message.role, Role::Assistant))
.filter_map(|message| message.tool_calls.as_ref())
.flatten()
.filter(|call| call.function.name == bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME)
.map(|call| call.id.as_str())
.collect()
}
fn sticky_result_definition_values(message: &Message) -> Option<Vec<serde_json::Value>> {
let body = message
.content
.strip_prefix(STICKY_DISCOVERY_RESULT_START)?
.strip_suffix(STICKY_DISCOVERY_RESULT_END)?;
let payload = serde_json::from_str::<serde_json::Value>(body).ok()?;
payload["tools"].as_array().cloned()
}
fn sticky_definition_name(definition: &serde_json::Value) -> Option<&str> {
(definition["type"].as_str() == Some("function"))
.then(|| definition["function"]["name"].as_str())
.flatten()
}
fn validated_sticky_fallback_results(session: &Session) -> Vec<&Message> {
let call_ids = sticky_discovery_call_ids(session);
session
.messages
.iter()
.filter(|message| {
matches!(message.role, Role::Tool)
&& message.tool_success == Some(true)
&& message
.tool_call_id
.as_deref()
.is_some_and(|call_id| call_ids.contains(call_id))
&& message.metadata.as_ref().is_some_and(|metadata| {
metadata["runtime_kind"].as_str() == Some(STICKY_DISCOVERY_RUNTIME_KIND)
&& metadata["version"].as_u64() == Some(STICKY_DISCOVERY_RUNTIME_VERSION)
&& metadata["canonical_new_names"].is_array()
})
&& sticky_result_definition_values(message).is_some()
})
.collect()
}
fn validated_sticky_fallback_loaded_tool_names(session: &Session) -> Vec<String> {
let mut loaded = Vec::new();
for message in validated_sticky_fallback_results(session) {
let definition_names = sticky_result_definition_values(message)
.unwrap_or_default()
.into_iter()
.filter_map(|definition| sticky_definition_name(&definition).map(str::to_string))
.collect::<BTreeSet<_>>();
let Some(names) = message
.metadata
.as_ref()
.and_then(|metadata| metadata["canonical_new_names"].as_array())
else {
continue;
};
for name in names.iter().filter_map(serde_json::Value::as_str) {
if definition_names.contains(name) && !loaded.iter().any(|seen| seen == name) {
loaded.push(name.to_string());
}
}
}
loaded
}
pub(in crate::runtime::runner) fn legacy_browser_loaded_result_content(
session: &Session,
) -> BTreeMap<String, String> {
validated_sticky_fallback_results(session)
.into_iter()
.filter(|message| {
let canonical_browser_only = message
.metadata
.as_ref()
.and_then(|metadata| metadata["canonical_new_names"].as_array())
.is_some_and(|names| names.len() == 1 && names[0].as_str() == Some("browser"));
let browser_definition_only =
sticky_result_definition_values(message).is_some_and(|definitions| {
definitions.len() == 1
&& sticky_definition_name(&definitions[0]) == Some("browser")
});
canonical_browser_only && browser_definition_only
})
.map(|message| (message.id.clone(), message.content.clone()))
.collect()
}
pub(in crate::runtime::runner) fn legacy_browser_needs_discovery(
session: &Session,
tool_schemas: &[bamboo_agent_core::tools::ToolSchema],
required_tool: Option<&str>,
) -> bool {
required_tool != Some("browser")
&& tool_schemas
.iter()
.any(|schema| schema.function.name == "browser")
&& !validated_sticky_fallback_loaded_tool_names(session)
.iter()
.any(|name| name == "browser")
}
fn prior_sticky_fallback_definitions(session: &Session) -> Vec<serde_json::Value> {
validated_sticky_fallback_results(session)
.into_iter()
.filter_map(sticky_result_definition_values)
.flatten()
.collect()
}
fn sticky_fallback_definition_delta(
session: &Session,
schemas: &[bamboo_agent_core::tools::ToolSchema],
) -> Result<Vec<serde_json::Value>, serde_json::Error> {
let previous = prior_sticky_fallback_definitions(session);
let mut definitions = schemas
.iter()
.map(serde_json::to_value)
.collect::<Result<Vec<_>, _>>()?;
definitions.retain(|definition| !previous.contains(definition));
Ok(definitions)
}
fn sticky_fallback_tool_result(
tool_call_id: &str,
content: String,
success: bool,
canonical_new_names: &[String],
) -> Message {
let mut message = Message::tool_result_with_status(tool_call_id, content, success);
message.never_compress = true;
message.metadata = Some(serde_json::json!({
"runtime_kind": STICKY_DISCOVERY_RUNTIME_KIND,
"version": STICKY_DISCOVERY_RUNTIME_VERSION,
"canonical_new_names": canonical_new_names,
}));
message
}
async fn ensure_discovery_allowed(
session: &mut Session,
config: &AgentLoopConfig,
) -> Result<(), AgentError> {
state_bridge::refresh_round_root_tool_authority(session, config.storage.as_ref()).await?;
if session.root_orchestration_only_enabled() && session.supervisor_incarnation_id().is_none() {
return Err(AgentError::Tool(
"capability discovery is outside this Root's tool authority".into(),
));
}
Ok(())
}
async fn commit_sticky_fallback_discovery_round(
stream_output: StreamHandlingOutput,
session: &mut Session,
config: &AgentLoopConfig,
tool_schemas: &[bamboo_agent_core::tools::ToolSchema],
browser_only: bool,
) -> Result<(), AgentError> {
ensure_discovery_allowed(session, config).await?;
let reasoning = (!stream_output.reasoning_content.trim().is_empty())
.then_some(stream_output.reasoning_content);
let reasoning_signature = reasoning
.as_ref()
.and_then(|_| stream_output.reasoning_signature.clone());
let tool_calls = stream_output.tool_calls;
let mut assistant = assistant_with_visible_identity(
Message::assistant_with_reasoning(
stream_output.content,
Some(tool_calls.clone()),
reasoning,
)
.with_reasoning_signature(reasoning_signature),
stream_output.visible_message,
);
assistant.never_compress = true;
assistant.metadata = Some(serde_json::json!({
"runtime_kind": STICKY_DISCOVERY_RUNTIME_KIND,
"version": STICKY_DISCOVERY_RUNTIME_VERSION,
}));
let mut native_items = Some(stream_output.provider_transcript_items);
commit_assistant_message(session, assistant, &mut native_items)?;
if let Some(persistence) = config.persistence.as_ref() {
persistence
.save_runtime_session(session)
.await
.map_err(|error| {
AgentError::Tool(format!(
"sticky discovery assistant checkpoint could not be persisted: {error}"
))
})?;
}
let discovery_is_alone = tool_calls.len() == 1
&& tool_calls[0].function.name == bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME;
if !discovery_is_alone {
const ERROR: &str = "discovery must be called alone";
for call in &tool_calls {
session.add_message(sticky_fallback_tool_result(
&call.id,
ERROR.to_string(),
false,
&[],
));
}
} else {
let call = &tool_calls[0];
let discovery_result =
serde_json::from_str::<DiscoverCapabilitiesRequest>(&call.function.arguments)
.map_err(|error| format!("invalid discovery arguments: {error}"));
let definitions = match discovery_result {
Ok(request) => {
let include_commands = request.kinds.as_ref().is_some_and(|kinds| {
kinds
.iter()
.any(|kind| *kind != bamboo_domain::CapabilityKind::Tool)
});
match CompleteCapabilityDiscovery::new(
session,
config,
tool_schemas,
browser_only,
include_commands,
)
.await
.and_then(|discovery| discovery.discover_complete_schemas(&request))
{
Ok(mut schemas) => {
if browser_only {
schemas.retain(|schema| schema.function.name == "browser");
}
sticky_fallback_definition_delta(session, &schemas).map_err(|error| {
format!("discovered tool definitions could not be serialized: {error}")
})
}
Err(error) => Err(error.to_string()),
}
}
Err(error) => Err(error),
};
match definitions {
Ok(definitions) => {
let already_loaded = validated_sticky_fallback_loaded_tool_names(session)
.into_iter()
.collect::<BTreeSet<_>>();
let mut canonical_new_names = Vec::new();
for name in definitions.iter().filter_map(sticky_definition_name) {
if !already_loaded.contains(name)
&& !canonical_new_names.iter().any(|seen| seen == name)
{
canonical_new_names.push(name.to_string());
}
}
let payload = serde_json::to_string(&serde_json::json!({
"tools": definitions
}))
.map_err(|error| {
AgentError::LLM(format!(
"sticky discovery result could not be serialized: {error}"
))
})?;
session.add_message(sticky_fallback_tool_result(
&call.id,
format!(
"{STICKY_DISCOVERY_RESULT_START}{payload}{STICKY_DISCOVERY_RESULT_END}"
),
true,
&canonical_new_names,
));
}
Err(error) => {
session.add_message(sticky_fallback_tool_result(
&call.id,
format!("capability discovery failed: {error}"),
false,
&[],
));
}
}
}
if let Some(persistence) = config.persistence.as_ref() {
persistence
.save_runtime_session(session)
.await
.map_err(|error| {
AgentError::Tool(format!(
"sticky discovery result checkpoint could not be persisted: {error}"
))
})?;
}
Ok(())
}
async fn build_openai_client_tool_search_outputs(
session: &Session,
config: &AgentLoopConfig,
tool_schemas: &[bamboo_agent_core::tools::ToolSchema],
provider_items: &[ProviderTranscriptItem],
) -> Result<Vec<ProviderTranscriptItem>, AgentError> {
let requests = openai_client_tool_search_requests(provider_items)?;
let discovery =
CompleteCapabilityDiscovery::new(session, config, tool_schemas, false, false).await?;
let mut outputs = Vec::with_capacity(requests.len());
for (call_id, request) in requests {
if request.kinds.as_ref().is_some_and(|kinds| {
kinds
.iter()
.any(|kind| *kind != bamboo_domain::CapabilityKind::Tool)
}) {
return Err(AgentError::LLM(
"native tool search accepts Tools only".to_string(),
));
}
let tools = discovery
.discover_complete_schemas(&request)?
.iter()
.map(bamboo_llm::providers::common::openai_responses::loaded_tool_to_responses_json)
.collect::<Vec<_>>();
let item = ProviderTranscriptItem::try_from_payload(
ProviderFamily::OpenAi,
ProviderProtocol::OpenAiResponsesV1,
ProviderTranscriptOrigin::HostToolSearch,
ProviderTranscriptAuthor::ToolResult,
serde_json::json!({
"type": "tool_search_output",
"execution": "client",
"call_id": call_id,
"status": "completed",
"tools": tools,
}),
)
.map_err(|error| {
AgentError::LLM(format!(
"OpenAI client tool-search output could not be constructed: {error}"
))
})?;
outputs.push(item);
}
Ok(outputs)
}
async fn commit_openai_client_tool_search_round(
stream_output: StreamHandlingOutput,
session: &mut Session,
config: &AgentLoopConfig,
tool_schemas: &[bamboo_agent_core::tools::ToolSchema],
) -> Result<Message, AgentError> {
ensure_discovery_allowed(session, config).await?;
let host_outputs = build_openai_client_tool_search_outputs(
session,
config,
tool_schemas,
&stream_output.provider_transcript_items,
)
.await?;
let reasoning = (!stream_output.reasoning_content.trim().is_empty())
.then_some(stream_output.reasoning_content);
let reasoning_signature = reasoning
.as_ref()
.and_then(|_| stream_output.reasoning_signature.clone());
let message = assistant_with_visible_identity(
Message::assistant_with_reasoning(stream_output.content, None, reasoning)
.with_reasoning_signature(reasoning_signature),
stream_output.visible_message,
);
let anchor = message.id.clone();
let mut provider_items = Some(stream_output.provider_transcript_items);
commit_assistant_message(session, message.clone(), &mut provider_items)?;
for output in host_outputs {
session
.append_provider_transcript_group(&anchor, None, vec![output])
.map_err(|error| {
AgentError::LLM(format!(
"OpenAI client tool-search result could not be committed: {error}"
))
})?;
}
if let Some(persistence) = config.persistence.as_ref() {
persistence
.save_runtime_session(session)
.await
.map_err(|error| {
AgentError::Tool(format!(
"OpenAI client tool-search checkpoint could not be persisted: {error}"
))
})?;
}
Ok(message)
}
fn should_retry_turn_error(error: &AgentError) -> bool {
let message = match error {
AgentError::StreamTimeout(timeout) => return timeout.retry_safe(),
AgentError::LLM(message) => message,
_ => return false,
};
let message = message.trim().to_ascii_lowercase();
if message.is_empty() {
return false;
}
let non_retryable_patterns = [
"authentication error",
"invalid api key",
"invalid_request_error",
"unsupported model",
"model_name is required",
"http 400",
"http 401",
"http 403",
"http 404",
];
!non_retryable_patterns
.iter()
.any(|pattern| message.contains(pattern))
}
fn is_overflow_recoverable(error: &AgentError) -> bool {
matches!(error, AgentError::LLMOverflow(_))
}
struct TurnOutcome {
should_break: bool,
sent_complete: bool,
}
const SUBAGENT_TOOL_NAME: &str = "SubAgent";
const PLAN_TOOL_NAME: &str = "Plan";
fn is_child_spawn_call(call: &bamboo_agent_core::tools::ToolCall) -> bool {
match call.function.name.as_str() {
PLAN_TOOL_NAME => return true,
SUBAGENT_TOOL_NAME => {}
_ => return false,
}
serde_json::from_str::<serde_json::Value>(&call.function.arguments)
.ok()
.and_then(|value| {
value
.get("action")
.and_then(|a| a.as_str())
.map(str::to_string)
})
.is_none_or(|action| action == "create")
}
struct RunBudgetExceeded {
kind: &'static str,
limit: u64,
actual: u64,
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
struct RoundActivity {
prompt_tokens: u64,
completion_tokens: u64,
tool_call_count: u32,
subagent_spawn_count: u32,
}
impl RoundActivity {
fn absorb_attempt(
&mut self,
stream_output: &StreamHandlingOutput,
attempt_usage: MetricsTokenUsage,
) {
let mut accumulated = self.token_usage();
accumulated.add_assign_durable(attempt_usage);
self.prompt_tokens = accumulated.prompt_tokens;
self.completion_tokens = accumulated.completion_tokens;
self.tool_call_count = self
.tool_call_count
.saturating_add(stream_output.tool_calls.len() as u32);
self.subagent_spawn_count = self.subagent_spawn_count.saturating_add(
stream_output
.tool_calls
.iter()
.filter(|call| is_child_spawn_call(call))
.count() as u32,
);
}
fn token_usage(&self) -> MetricsTokenUsage {
MetricsTokenUsage {
prompt_tokens: self.prompt_tokens,
completion_tokens: self.completion_tokens,
total_tokens: self.prompt_tokens.saturating_add(self.completion_tokens),
}
.clamped_for_durable_metrics()
}
fn commit_to_runtime(&self, runtime_state: &mut AgentRuntimeState) {
let mut accumulated = MetricsTokenUsage {
prompt_tokens: runtime_state.round.total_prompt_tokens,
completion_tokens: runtime_state.round.total_completion_tokens,
total_tokens: 0,
}
.clamped_for_durable_metrics();
accumulated.add_assign_durable(self.token_usage());
runtime_state.round.total_prompt_tokens = accumulated.prompt_tokens;
runtime_state.round.total_completion_tokens = accumulated.completion_tokens;
runtime_state.round.total_tool_calls = runtime_state
.round
.total_tool_calls
.saturating_add(self.tool_call_count);
runtime_state.round.total_subagents_spawned = runtime_state
.round
.total_subagents_spawned
.saturating_add(self.subagent_spawn_count);
}
}
fn check_run_budget_exceeded(
round: &bamboo_domain::session::runtime_state::RoundRuntimeState,
budget: &bamboo_config::RunBudgetConfig,
) -> Option<RunBudgetExceeded> {
let total_tokens = round
.total_prompt_tokens
.saturating_add(round.total_completion_tokens)
.min(bamboo_domain::MAX_DURABLE_TOKEN_COUNT);
if let Some(limit) = budget.max_total_tokens {
if total_tokens >= limit {
return Some(RunBudgetExceeded {
kind: "max_total_tokens",
limit,
actual: total_tokens,
});
}
}
if let Some(limit) = budget.max_tool_calls {
if round.total_tool_calls >= limit {
return Some(RunBudgetExceeded {
kind: "max_tool_calls",
limit: limit as u64,
actual: round.total_tool_calls as u64,
});
}
}
if let Some(limit) = budget.max_subagents {
if round.total_subagents_spawned >= limit {
return Some(RunBudgetExceeded {
kind: "max_subagents",
limit: limit as u64,
actual: round.total_subagents_spawned as u64,
});
}
}
None
}
fn is_terminal_child_status(status: &str) -> bool {
matches!(
status,
"completed" | "error" | "timeout" | "cancelled" | "skipped"
)
}
fn reconcile_child_wait_at_suspend(
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
durable_wait: Option<WaitingForChildrenState>,
) {
runtime_state.waiting_for_children = durable_wait;
if runtime_state.waiting_for_children.is_none() {
session.metadata.remove("runtime.suspend_reason");
runtime_state.suspension = None;
if runtime_state.status == AgentStatusState::Suspended {
runtime_state.status = AgentStatusState::Idle;
}
}
}
async fn suspend_to_wait_for_children(
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
persistence: Option<&Arc<dyn bamboo_domain::RuntimeSessionPersistence>>,
child_session_ids: Vec<String>,
wait_for: ChildWaitPolicy,
) -> TurnOutcome {
let now = Utc::now();
let count = child_session_ids.len();
runtime_state.waiting_for_children = Some(WaitingForChildrenState::for_children(
child_session_ids,
wait_for,
now,
));
state_bridge::write_runtime_state(session, runtime_state);
session.metadata.insert(
"runtime.suspend_reason".to_string(),
"waiting_for_children".to_string(),
);
session.updated_at = now;
if let Some(persistence) = persistence {
if let Err(error) = persistence.save_runtime_session(session).await {
tracing::warn!(
"[{}] suspend-to-wait failed to persist parent wait on {} child(ren): {}",
session.id,
count,
error
);
}
}
TurnOutcome {
should_break: true,
sent_complete: false,
}
}
async fn maybe_suspend_for_orphaned_children(
session: &mut Session,
config: &AgentLoopConfig,
runtime_state: &mut AgentRuntimeState,
) -> Result<Option<TurnOutcome>, AgentError> {
if session
.metadata
.get("runtime.canonical_subagent_host")
.is_some_and(|value| value == "true")
{
return Ok(None);
}
let inherited_wait = runtime_state.waiting_for_children.is_some();
let Some(storage) = config.storage.as_ref() else {
return Ok(None);
};
let parent_lock =
crate::session_app::child_completion_coordinator::session_resume_lock(&session.id);
let _parent_guard = parent_lock.lock().await;
let mut active: Vec<String> = storage
.list_child_run_statuses(&session.id)
.await
.unwrap_or_default()
.into_iter()
.filter(|(_, status)| !status.as_deref().is_some_and(is_terminal_child_status))
.map(|(id, _)| id)
.collect();
if !inherited_wait {
let mut ordinary = Vec::with_capacity(active.len());
for child_id in active {
if !crate::ticket_runtime::is_independent_ticket_child(
storage.as_ref(),
session,
&child_id,
)
.await
{
ordinary.push(child_id);
}
}
active = ordinary;
}
if active.is_empty() && !inherited_wait {
return Ok(None);
}
if crate::runtime::loop_control::suspension::active(config) {
return Err(AgentError::Tool(
"managed loop cannot establish an asynchronous child wait".into(),
));
}
active.sort();
active.dedup();
let durable = storage.load_session(&session.id).await.map_err(|_| {
AgentError::Tool("parent wait observation failed; refusing to replace its wait".into())
})?;
if inherited_wait && durable.is_none() {
return Err(AgentError::Tool(
"parent wait observation failed; current session is missing".into(),
));
}
if let Some(durable) = durable {
if durable.id != session.id || durable.created_at != session.created_at {
return Err(AgentError::Tool(
"parent wait observation belongs to a different session lifetime".into(),
));
}
if let Some(wait) = durable
.agent_runtime_state
.and_then(|state| state.waiting_for_children)
{
runtime_state.waiting_for_children = Some(wait);
state_bridge::write_runtime_state(session, runtime_state);
session.metadata.insert(
"runtime.suspend_reason".into(),
"waiting_for_children".into(),
);
return Ok(Some(TurnOutcome {
should_break: true,
sent_complete: false,
}));
}
}
if inherited_wait {
runtime_state.waiting_for_children = None;
if session
.metadata
.get("runtime.suspend_reason")
.map(String::as_str)
== Some("waiting_for_children")
{
session.metadata.remove("runtime.suspend_reason");
}
if runtime_state
.suspension
.as_ref()
.is_some_and(|s| s.reason == "waiting_for_children")
{
runtime_state.suspension = None;
if runtime_state.status == AgentStatusState::Suspended {
runtime_state.status = AgentStatusState::Idle;
}
}
state_bridge::write_runtime_state(session, runtime_state);
return Ok(None);
}
tracing::info!(
"[{}] end-of-turn safety net: suspending to wait for {} orphaned child session(s) the model did not explicitly wait on",
session.id,
active.len(),
);
Ok(Some(
suspend_to_wait_for_children(
session,
runtime_state,
config.persistence.as_ref(),
active,
ChildWaitPolicy::All,
)
.await,
))
}
async fn suspend_to_wait_for_bash(
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
persistence: Option<&Arc<dyn bamboo_domain::RuntimeSessionPersistence>>,
bash_ids: Vec<String>,
) -> TurnOutcome {
let now = Utc::now();
let count = bash_ids.len();
runtime_state.waiting_for_bash = Some(WaitingForBashState::for_bash(bash_ids, now));
state_bridge::write_runtime_state(session, runtime_state);
session.metadata.insert(
"runtime.suspend_reason".to_string(),
"waiting_for_bash".to_string(),
);
session.updated_at = now;
if let Some(persistence) = persistence {
if let Err(error) = persistence.save_runtime_session(session).await {
tracing::warn!(
"[{}] suspend-to-wait-bash failed to persist bash wait on {} shell(s): {}",
session.id,
count,
error
);
}
}
TurnOutcome {
should_break: true,
sent_complete: false,
}
}
async fn maybe_suspend_for_outstanding_bash(
session: &mut Session,
config: &AgentLoopConfig,
runtime_state: &mut AgentRuntimeState,
) -> Option<TurnOutcome> {
if runtime_state.waiting_for_bash.is_some() {
return None;
}
config.persistence.as_ref()?;
let hook = config.bash_resume_hook.as_ref()?;
let mut bash_ids = bamboo_tools::tools::bash_runtime::running_shells_for_session(&session.id);
if bash_ids.is_empty() {
return None;
}
bash_ids.sort();
bash_ids.dedup();
if bamboo_tools::tools::bash_runtime::running_shells_for_session(&session.id).is_empty() {
tracing::info!(
"[{}] end-of-turn bash gate: all {} shell(s) finished during the snapshot window; not suspending",
session.id,
bash_ids.len(),
);
return None;
}
tracing::info!(
"[{}] end-of-turn safety net: suspending to wait for {} background bash shell(s) still running",
session.id,
bash_ids.len(),
);
let hook_ids = bash_ids.clone();
let outcome = suspend_to_wait_for_bash(
session,
runtime_state,
config.persistence.as_ref(),
bash_ids,
)
.await;
hook.arrange_bash_self_resume(session.id.clone(), hook_ids);
Some(outcome)
}
fn build_guardian_review_prompt(
task_context: &Option<TaskLoopContext>,
config: &AgentLoopConfig,
final_assistant_content: Option<&str>,
) -> String {
let mut prompt = String::from(GUARDIAN_REVIEW_RUBRIC);
let criteria: Vec<String> = task_context
.as_ref()
.and_then(|ctx| {
ctx.items
.iter()
.find(|item| Some(&item.id) == ctx.active_item_id.as_ref())
})
.map(|item| item.completion_criteria.clone())
.unwrap_or_default();
if !criteria.is_empty() {
prompt.push_str("\n\n## Completion criteria (verify EACH against real evidence)\n");
for (idx, criterion) in criteria.iter().enumerate() {
prompt.push_str(&format!("{}. {}\n", idx + 1, criterion));
}
}
let goal = config.active_goal();
if let Some(goal) = goal {
prompt.push_str("\n\n## Session goal\n");
prompt.push_str(goal);
prompt.push('\n');
}
if criteria.is_empty() && goal.is_none() {
prompt.push_str(
"\n\n(No explicit completion criteria or goal were provided; review the diff for correctness, completeness, and obvious bugs.)\n",
);
}
if let Some(content) = final_assistant_content {
let trimmed = content.trim();
if !trimmed.is_empty() {
prompt.push_str(
"\n\n## Agent's final message (context only — not yet part of the session transcript)\n",
);
prompt.push_str(trimmed);
prompt.push('\n');
}
}
prompt
}
async fn maybe_spawn_guardian_review(
session: &mut Session,
config: &AgentLoopConfig,
task_context: &Option<TaskLoopContext>,
runtime_state: &mut AgentRuntimeState,
iteration: u32,
final_assistant_content: Option<&str>,
) -> Option<TurnOutcome> {
if runtime_state.waiting_for_children.is_some() {
return None;
}
if !config.guardian_active() {
return None;
}
let spawner = config.guardian_spawner.as_ref()?;
let max_reviews = config.guardian_max_reviews();
let mut guardian_state = ensure_guardian_state(session);
match guardian_state.phase {
GuardianPhase::Pending => return None,
GuardianPhase::Reviewed => {
if guardian_state.last_approved() {
return None;
}
if guardian_state.budget_exhausted(max_reviews) {
tracing::warn!(
"[{}] guardian: review budget ({}) exhausted with unresolved findings; allowing completion",
session.id,
max_reviews
);
return None;
}
}
GuardianPhase::None => {
if guardian_state.budget_exhausted(max_reviews) {
return None;
}
}
}
if let Some(guardian_config) = config.guardian_config.as_ref() {
write_guardian_config(session, guardian_config);
}
let review_prompt = build_guardian_review_prompt(task_context, config, final_assistant_content);
let Some(model) = config
.guardian_model()
.map(str::to_string)
.or_else(|| config.model_name.clone())
.map(|model| model.trim().to_string())
.filter(|model| !model.is_empty())
else {
tracing::warn!(
"[{}] guardian: no reviewer model resolved; skipping review at this terminal",
session.id
);
return None;
};
let disabled_tools = Some(guardian_read_only_disabled_tools());
match spawner
.spawn_guardian_review(session, review_prompt, model, disabled_tools)
.await
{
Ok(child_id) => {
guardian_state.record_spawn(&child_id);
guardian_state.last_reviewed_at_round = iteration;
let pass = guardian_state.review_count;
write_guardian_state(session, guardian_state);
tracing::info!(
"[{}] guardian: spawned read-only review child {} (pass {}/{}); suspending until verdict",
session.id,
child_id,
pass,
max_reviews
);
Some(
suspend_to_wait_for_children(
session,
runtime_state,
config.persistence.as_ref(),
vec![child_id],
ChildWaitPolicy::All,
)
.await,
)
}
Err(error) => {
tracing::warn!(
"[{}] guardian: failed to spawn review child: {}; allowing completion",
session.id,
error
);
None
}
}
}
fn map_turn_error_status(error: &AgentError) -> (MetricsRoundStatus, MetricsSessionStatus) {
if matches!(error, AgentError::Cancelled) {
(
MetricsRoundStatus::Cancelled,
MetricsSessionStatus::Cancelled,
)
} else {
(MetricsRoundStatus::Error, MetricsSessionStatus::Error)
}
}
fn record_turn_failure(
metrics_collector: Option<&MetricsCollector>,
round_id: &str,
session_id: &str,
message_count: u32,
round_usage: MetricsTokenUsage,
error: &AgentError,
) {
let (round_status, session_status) = map_turn_error_status(error);
crate::runtime::runner::metrics_lifecycle::record_round_and_session_error(
metrics_collector,
round_id,
session_id,
message_count,
round_status,
round_usage,
Some(error.to_string()),
session_status,
);
}
#[allow(clippy::too_many_arguments)]
fn record_task_evaluation_terminal_metrics(
metrics_collector: Option<&MetricsCollector>,
metrics_round_id: &str,
metrics_started: &AtomicBool,
metrics_terminal: &AtomicBool,
completed_at: chrono::DateTime<Utc>,
status: MetricsRoundStatus,
usage: MetricsTokenUsage,
error: Option<String>,
) {
if !metrics_started.load(Ordering::Acquire)
|| metrics_terminal
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
return;
}
if let Some(metrics) = metrics_collector {
metrics.round_completed(
metrics_round_id.to_string(),
completed_at,
status,
usage,
0,
0,
error,
);
}
}
fn task_evaluation_result_usage(
result: &crate::runtime::runner::task_lifecycle::AsyncTaskEvaluationResult,
) -> MetricsTokenUsage {
let prompt_tokens = result.evaluation_result.prompt_tokens;
let completion_tokens = result.evaluation_result.completion_tokens;
MetricsTokenUsage {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens.saturating_add(completion_tokens),
}
}
fn record_harvested_task_evaluation_metrics(
metrics_collector: Option<&MetricsCollector>,
result: &crate::runtime::runner::task_lifecycle::AsyncTaskEvaluationResult,
) {
let (status, error) = match result.error.clone() {
Some(error) => (MetricsRoundStatus::Error, Some(error)),
None => (MetricsRoundStatus::Success, None),
};
record_task_evaluation_terminal_metrics(
metrics_collector,
&result.metrics_round_id,
result.metrics_started.as_ref(),
result.metrics_terminal.as_ref(),
result.finished_at,
status,
task_evaluation_result_usage(result),
error,
);
}
async fn poll_completed_task_evaluation(state: &mut LoopRunState) {
let finished = state
.task_evaluation
.in_flight
.as_ref()
.is_some_and(|in_flight| in_flight.join_handle.is_finished());
if !finished {
return;
}
let Some(in_flight) = state.task_evaluation.in_flight.take() else {
return;
};
match in_flight.join_handle.await {
Ok(Some(result)) => {
state.task_evaluation.completed = Some(result);
}
Ok(None) => {
tracing::debug!(
"[{}] Async task evaluation cancelled for round {}",
state.session_id,
in_flight.request.round_number
);
record_task_evaluation_terminal_metrics(
state.metrics_collector.as_ref(),
&in_flight.request.metrics_round_id,
in_flight.metrics_started.as_ref(),
in_flight.metrics_terminal.as_ref(),
Utc::now(),
MetricsRoundStatus::Cancelled,
MetricsTokenUsage::default(),
Some("task evaluation cancelled".to_string()),
);
}
Err(error) => {
tracing::warn!(
"[{}] Async task evaluation join failed for round {}: {}",
state.session_id,
in_flight.request.round_number,
error
);
record_task_evaluation_terminal_metrics(
state.metrics_collector.as_ref(),
&in_flight.request.metrics_round_id,
in_flight.metrics_started.as_ref(),
in_flight.metrics_terminal.as_ref(),
Utc::now(),
MetricsRoundStatus::Error,
MetricsTokenUsage::default(),
Some(format!("task evaluation join failed: {error}")),
);
}
}
}
async fn apply_completed_task_evaluation(
session: &mut Session,
event_tx: &mpsc::Sender<AgentEvent>,
config: &AgentLoopConfig,
state: &mut LoopRunState,
) {
let Some(result) = state.task_evaluation.completed.take() else {
return;
};
let mut staged_task_context = state.task_context.clone();
let mut staged_session = session.clone();
let mut apply_outcome = crate::runtime::runner::task_lifecycle::apply_task_evaluation_result(
&mut staged_task_context,
&mut staged_session,
&state.session_id,
result.clone(),
);
let mut persistence_error = None;
if !apply_outcome.stale && apply_outcome.applied_updates > 0 {
if let Some(ref ctx) = staged_task_context {
let task_list_title = result
.task_list_title
.clone()
.or_else(|| {
staged_session
.task_list
.as_ref()
.map(|task_list| task_list.title.clone())
})
.unwrap_or_else(|| "Agent Tasks".to_string());
let task_list = ctx.to_task_list_with_title(task_list_title);
let new_version = ctx.version.to_string();
let expected_version = result.based_on_task_context_version.to_string();
let persisted = if let Some(persistence) = config.persistence.as_ref() {
match persistence
.update_task_list_control_planes_if_version(
&state.session_id,
&result.shared_session_id,
&expected_version,
&result.based_on_task_list,
&task_list,
&new_version,
)
.await
{
Ok(updated) => updated,
Err(error) => {
tracing::warn!(
"[{}] Failed to compare-and-patch evaluated Task state on {}: {}",
state.session_id,
result.shared_session_id,
error
);
persistence_error = Some(error.to_string());
false
}
}
} else {
true
};
if persisted {
state.task_context = staged_task_context;
let event_version = new_version.parse::<u64>().ok();
session.set_task_list_version_meta(new_version);
session.set_task_list(task_list.clone());
let _ = event_tx
.send(AgentEvent::TaskListUpdated {
task_list,
version: event_version,
})
.await;
} else {
apply_outcome.stale = true;
tracing::debug!(
"[{}] Dropping evaluated Task update because durable generation {} no longer matches",
state.session_id,
expected_version
);
}
}
}
let (metrics_status, completed_at, metrics_error) = if let Some(error) = result.error.clone() {
(MetricsRoundStatus::Error, result.finished_at, Some(error))
} else if let Some(error) = persistence_error {
(MetricsRoundStatus::Error, Utc::now(), Some(error))
} else if apply_outcome.stale {
(
MetricsRoundStatus::Cancelled,
Utc::now(),
Some("stale task evaluation result".to_string()),
)
} else {
(MetricsRoundStatus::Success, result.finished_at, None)
};
record_task_evaluation_terminal_metrics(
state.metrics_collector.as_ref(),
&result.metrics_round_id,
result.metrics_started.as_ref(),
result.metrics_terminal.as_ref(),
completed_at,
metrics_status,
apply_outcome.usage,
metrics_error,
);
}
fn spawn_task_evaluation_request(
state: &mut LoopRunState,
event_tx: &mpsc::Sender<AgentEvent>,
request: crate::runtime::runner::task_lifecycle::AsyncTaskEvaluationRequest,
llm: Arc<dyn LLMProvider>,
cancel_token: CancellationToken,
configured_limit: usize,
) {
let task_round = request.round_number;
let session_id = state.session_id.clone();
let event_tx = event_tx.clone();
let request_for_spawn = request.clone();
let metrics_collector = state.metrics_collector.clone();
let metrics_started = Arc::new(AtomicBool::new(false));
let metrics_terminal = Arc::new(AtomicBool::new(false));
let metrics_started_for_spawn = metrics_started.clone();
let metrics_terminal_for_spawn = metrics_terminal.clone();
let join_handle = bamboo_tools::local_task_resources::spawn_cancellable(async move {
tokio::select! {
biased;
_ = cancel_token.cancelled() => None,
result = crate::runtime::runner::task_lifecycle::execute_async_task_evaluation(
request_for_spawn,
llm,
event_tx,
configured_limit,
metrics_collector,
metrics_started_for_spawn,
metrics_terminal_for_spawn,
) => Some(result),
}
});
tracing::debug!(
evaluation_id = %request.evaluation_id,
"[{}] Spawned async task evaluation for round {}",
session_id,
task_round
);
state.task_evaluation.in_flight = Some(InFlightTaskEvaluation {
request,
metrics_started,
metrics_terminal,
join_handle,
});
}
async fn abort_in_flight_evaluations(
state: &mut LoopRunState,
event_tx: &mpsc::Sender<AgentEvent>,
reason: &'static str,
) {
let mut task_was_cancelled = false;
let mut task_generation = None;
let mut gold_was_cancelled = false;
if let Some(in_flight) = state.task_evaluation.in_flight.take() {
task_generation = Some(in_flight.request.based_on_task_context_version);
let was_finished = in_flight.join_handle.is_finished();
if !was_finished {
in_flight.join_handle.abort();
}
let metrics_round_id = in_flight.request.metrics_round_id.clone();
let metrics_started = in_flight.metrics_started.clone();
let metrics_terminal = in_flight.metrics_terminal.clone();
match in_flight.join_handle.await {
Ok(Some(result)) => {
record_harvested_task_evaluation_metrics(state.metrics_collector.as_ref(), &result);
}
Ok(None) => {
task_was_cancelled = true;
record_task_evaluation_terminal_metrics(
state.metrics_collector.as_ref(),
&metrics_round_id,
metrics_started.as_ref(),
metrics_terminal.as_ref(),
Utc::now(),
MetricsRoundStatus::Cancelled,
MetricsTokenUsage::default(),
Some(reason.to_string()),
);
}
Err(error) if error.is_cancelled() => {
task_was_cancelled = true;
record_task_evaluation_terminal_metrics(
state.metrics_collector.as_ref(),
&metrics_round_id,
metrics_started.as_ref(),
metrics_terminal.as_ref(),
Utc::now(),
MetricsRoundStatus::Cancelled,
MetricsTokenUsage::default(),
Some(reason.to_string()),
);
}
Err(error) => {
tracing::warn!(
"[{}] Async task evaluation join failed during cleanup for round {}: {}",
state.session_id,
in_flight.request.round_number,
error
);
record_task_evaluation_terminal_metrics(
state.metrics_collector.as_ref(),
&metrics_round_id,
metrics_started.as_ref(),
metrics_terminal.as_ref(),
Utc::now(),
MetricsRoundStatus::Error,
MetricsTokenUsage::default(),
Some(format!("task evaluation join failed: {error}")),
);
}
}
}
if let Some(in_flight) = state.gold_evaluation.in_flight.take() {
let was_finished = in_flight.join_handle.is_finished();
if !was_finished {
in_flight.join_handle.abort();
}
match in_flight.join_handle.await {
Ok(Some(_result)) => {}
Ok(None) => gold_was_cancelled = true,
Err(error) if error.is_cancelled() => gold_was_cancelled = true,
Err(error) => {
tracing::warn!(
"[{}] Async Gold evaluation join failed during cleanup for round {}: {}",
state.session_id,
in_flight.request.round_number,
error
);
}
}
}
state.task_evaluation.queued_request = None;
state.gold_evaluation.queued_request = None;
if task_was_cancelled {
let _ = event_tx
.send(AgentEvent::TaskEvaluationCancelled {
session_id: state.session_id.clone(),
reason: reason.to_string(),
generation: task_generation,
})
.await;
}
if gold_was_cancelled {
let _ = event_tx
.send(AgentEvent::GoldEvaluationCancelled {
session_id: state.session_id.clone(),
reason: reason.to_string(),
})
.await;
}
}
fn spawn_task_evaluation_if_needed(
turn: usize,
session: &Session,
event_tx: &mpsc::Sender<AgentEvent>,
config: &AgentLoopConfig,
state: &mut LoopRunState,
llm: Arc<dyn LLMProvider>,
cancel_token: CancellationToken,
) -> Result<(), AgentError> {
if config.ticket_worker_plan.is_some() {
if let Some(ctx) = state.task_context.as_mut() {
ctx.task_list_dirty = false;
}
return Ok(());
}
let task_list_dirty = state
.task_context
.as_ref()
.is_some_and(|ctx| ctx.task_list_dirty);
if !task_list_dirty {
return Ok(());
}
if let Some(ctx) = state.task_context.as_mut() {
ctx.task_list_dirty = false;
}
let eval_model = state
.auxiliary_models
.fast_model_name
.as_deref()
.or(Some(state.model_name.as_str()));
let request = crate::runtime::runner::task_lifecycle::build_async_task_evaluation_request(
&state.task_context,
session,
&state.session_id,
turn + 1,
eval_model,
config.reasoning_effort,
crate::runtime::stream::handler::StreamTimeoutContext::new(
config.stream_timeout,
config.provider_name.as_deref(),
eval_model,
),
)?;
let Some(request) = request else {
return Ok(());
};
if state.task_evaluation.in_flight.is_some() {
state.task_evaluation.queued_request = Some(request);
tracing::debug!(
"[{}] Queued latest async task evaluation snapshot for round {} while another evaluation is still in flight",
state.session_id,
turn + 1
);
return Ok(());
}
spawn_task_evaluation_request(
state,
event_tx,
request,
llm,
cancel_token,
config.auxiliary_evaluation_max_concurrency,
);
Ok(())
}
fn refresh_auxiliary_models_for_round(state: &mut LoopRunState, config: &AgentLoopConfig) {
state.auxiliary_models = resolve_auxiliary_models(config);
state.runtime_state.llm.fast_model_name = state.auxiliary_models.fast_model_name.clone();
state.runtime_state.llm.background_model_name =
state.auxiliary_models.background_model_name.clone();
}
fn assistant_with_visible_identity(
message: Message,
visible_message: Option<VisibleMessageIdentity>,
) -> Message {
match visible_message {
Some(identity) => identity.apply_to(message),
None => message,
}
}
fn commit_assistant_message(
session: &mut Session,
message: Message,
native_items: &mut Option<Vec<ProviderTranscriptItem>>,
) -> Result<(), AgentError> {
let anchor = message.id.clone();
let previous_updated_at = session.updated_at;
session.add_message(message);
let Some(items) = native_items.take().filter(|items| !items.is_empty()) else {
return Ok(());
};
if let Err(error) = session.append_provider_transcript_group(&anchor, None, items) {
if session
.messages
.last()
.is_some_and(|message| message.id == anchor)
{
session.messages.pop();
session.updated_at = previous_updated_at;
}
return Err(AgentError::LLM(format!(
"provider-native transcript group rejected: {error}"
)));
}
Ok(())
}
fn record_no_tool_calls_round_completed(
metrics_collector: Option<&MetricsCollector>,
round_id: &str,
session_id: &str,
session: &Session,
round_usage: MetricsTokenUsage,
) {
crate::runtime::runner::metrics_lifecycle::record_round_completed(
metrics_collector,
round_id,
session_id,
session.messages.len() as u32,
MetricsRoundStatus::Success,
round_usage,
session
.token_usage
.as_ref()
.map(|usage| usage.prompt_cached_tool_outputs)
.unwrap_or(0)
.min(u32::MAX as usize) as u32,
session
.token_usage
.as_ref()
.map(|usage| usage.prompt_cached_tool_tokens_saved)
.unwrap_or(0),
None,
);
}
#[allow(clippy::too_many_arguments)]
async fn handle_no_tool_calls_with_native(
content: String,
visible_message: Option<VisibleMessageIdentity>,
reasoning: Option<String>,
reasoning_signature: Option<String>,
prompt_tokens: u64,
completion_tokens: u64,
round_usage: MetricsTokenUsage,
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
event_tx: &mpsc::Sender<AgentEvent>,
metrics_collector: Option<&MetricsCollector>,
round_id: &str,
session_id: &str,
config: &AgentLoopConfig,
task_context: &Option<TaskLoopContext>,
eval_model: &str,
iteration: u32,
llm: Arc<dyn LLMProvider>,
provider_transcript_items: Vec<ProviderTranscriptItem>,
evidence_llm: Arc<dyn LLMProvider>,
cancel_token: &CancellationToken,
) -> Result<TurnOutcome, AgentError> {
let add_message_before_gold =
config.goal_loop_active() || config.effective_counted_loop().is_some();
let mut deferred_assistant_message = Some(assistant_with_visible_identity(
Message::assistant_with_reasoning(content, None, reasoning)
.with_reasoning_signature(reasoning_signature),
visible_message,
));
let buffered_candidate = config.features_final_evidence_check.then(|| {
deferred_assistant_message
.as_ref()
.expect("candidate exists")
.clone()
});
let mut native_items = Some(provider_transcript_items);
if add_message_before_gold {
if let Some(message) = deferred_assistant_message.take() {
commit_assistant_message(session, message, &mut native_items)?;
}
}
if config.effective_counted_loop().is_some() && cancel_token.is_cancelled() {
return Err(AgentError::Cancelled);
}
let counted_decision = crate::runtime::counted_loop::terminal(session, config).await?;
if config.effective_counted_loop().is_some() && cancel_token.is_cancelled() {
return Err(AgentError::Cancelled);
}
match counted_decision {
crate::runtime::counted_loop::TerminalDecision::Continue => {
if let Some(message) = buffered_candidate.as_ref() {
publish_buffered_message(event_tx, message).await;
}
record_no_tool_calls_round_completed(
metrics_collector,
round_id,
session_id,
session,
round_usage,
);
return Ok(TurnOutcome {
should_break: false,
sent_complete: false,
});
}
crate::runtime::counted_loop::TerminalDecision::Stop => {
if let Some(message) = buffered_candidate.as_ref() {
publish_buffered_message(event_tx, message).await;
}
record_no_tool_calls_round_completed(
metrics_collector,
round_id,
session_id,
session,
round_usage,
);
return Ok(TurnOutcome {
should_break: true,
sent_complete: false,
});
}
crate::runtime::counted_loop::TerminalDecision::Finalize => {}
}
let decision = evaluate_gold_terminal(
session,
task_context,
config,
eval_model,
config.reasoning_effort,
session_id,
iteration,
llm,
event_tx,
)
.await;
if let GoldTerminalDecision::Continue { continuation_count } = decision {
if let Some(message) = buffered_candidate.as_ref() {
publish_buffered_message(event_tx, message).await;
}
tracing::info!(
"[{}] Goal terminal gate: continuing toward goal (continuation {})",
session_id,
continuation_count
);
record_no_tool_calls_round_completed(
metrics_collector,
round_id,
session_id,
session,
round_usage,
);
return Ok(TurnOutcome {
should_break: false,
sent_complete: false,
});
}
let final_assistant_content_for_guardian = deferred_assistant_message
.as_ref()
.map(|message| message.content.as_str());
if let Some(review) = maybe_spawn_guardian_review(
session,
config,
task_context,
runtime_state,
iteration,
final_assistant_content_for_guardian,
)
.await
{
if let Some(message) = buffered_candidate.as_ref() {
publish_buffered_message(event_tx, message).await;
}
return Ok(review);
}
const MAX_STOP_HOOK_CONTINUATIONS: u8 = 5;
if config
.hook_runner
.has_hooks_for(AgentHookPoint::BeforeFinalize)
{
let outcome = config
.hook_runner
.run_hooks_with_inputs(
AgentHookPoint::BeforeFinalize,
&HookPayload::Finalize {
stop_hook_active: runtime_state.stop_hook_forced_continuations > 0,
},
session,
runtime_state,
Some(event_tx),
bamboo_hooks::portable::PortableInputs {
final_assistant_content: final_assistant_content_for_guardian,
..Default::default()
},
)
.await;
if let HookResult::Deny { reason } = &outcome.decision {
if runtime_state.stop_hook_forced_continuations < MAX_STOP_HOOK_CONTINUATIONS {
runtime_state.stop_hook_forced_continuations += 1;
if let Some(message) = deferred_assistant_message.take() {
commit_assistant_message(session, message, &mut native_items)?;
}
if let Some(message) = buffered_candidate.as_ref() {
publish_buffered_message(event_tx, message).await;
}
let extra_context = outcome
.injected_contexts
.iter()
.map(|context| context.trim())
.filter(|context| !context.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
let context_suffix = if extra_context.is_empty() {
String::new()
} else {
format!("\n\nAdditional hook context:\n{extra_context}")
};
let mut continuation = Message::user(format!(
"A Stop lifecycle hook requires another work round ({}/{}): {}{}\n\nContinue working and address this feedback before attempting to finish again.",
runtime_state.stop_hook_forced_continuations,
MAX_STOP_HOOK_CONTINUATIONS,
reason.trim(),
context_suffix,
));
continuation.metadata = Some(serde_json::json!({
"runtime_kind": "stop_hook_continuation"
}));
session.add_message(continuation);
state_bridge::write_runtime_state(session, runtime_state);
record_no_tool_calls_round_completed(
metrics_collector,
round_id,
session_id,
session,
round_usage,
);
return Ok(TurnOutcome {
should_break: false,
sent_complete: false,
});
}
tracing::warn!(
"[{}] Stop hook continuation cap ({}) reached; ignoring further block",
session_id,
MAX_STOP_HOOK_CONTINUATIONS
);
session.metadata.insert(
"runtime.completion_reason".to_string(),
"stop_hook_continuation_cap".to_string(),
);
} else {
let hook_result = crate::runtime::hooks::apply_hook_outcome(
AgentHookPoint::BeforeFinalize,
outcome,
session,
runtime_state,
);
state_bridge::write_runtime_state(session, runtime_state);
hook_result?;
}
}
let mut complete_usage = MetricsTokenUsage {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens.saturating_add(completion_tokens),
};
if let Some(mut answer) = buffered_candidate {
let check = super::final_answer::check(
session,
runtime_state,
&mut answer,
config,
evidence_llm,
eval_model,
cancel_token,
metrics_collector,
round_id,
round_usage,
)
.await;
let check = match check {
Ok(check) => check,
Err(error) => {
if add_message_before_gold {
session.messages.retain(|message| message.id != answer.id);
session.updated_at = Utc::now();
}
reset_final_answer_context(session);
return Err(error);
}
};
complete_usage.add_assign_durable(check.usage);
if check.revised {
if let Some(message) = deferred_assistant_message.as_mut() {
*message = answer.clone();
} else {
let stored = session
.messages
.iter_mut()
.find(|message| message.id == answer.id)
.ok_or_else(|| {
AgentError::LLM("final candidate is no longer present".to_string())
})?;
*stored = answer.clone();
session.updated_at = Utc::now();
}
native_items = None;
reset_final_answer_context(session);
}
publish_buffered_message(event_tx, &answer).await;
}
if let Some(message) = deferred_assistant_message.take() {
commit_assistant_message(session, message, &mut native_items)?;
}
let usage = to_event_token_usage(
complete_usage.prompt_tokens,
complete_usage.completion_tokens,
);
if !super::defers_terminal_complete(config) {
let _ = event_tx.send(AgentEvent::Complete { usage }).await;
}
record_no_tool_calls_round_completed(
metrics_collector,
round_id,
session_id,
session,
round_usage,
);
Ok(TurnOutcome {
should_break: true,
sent_complete: !super::defers_terminal_complete(config),
})
}
fn reset_final_answer_context(session: &mut Session) {
session.metadata.remove("responses.previous_response_id");
session.reset_model_context_epoch(
bamboo_domain::session::model_context::ModelContextResetReason::ExplicitHistoryRewrite,
);
}
async fn publish_buffered_message(event_tx: &mpsc::Sender<AgentEvent>, message: &Message) {
crate::runtime::stream::handler::publish_buffered_response(
event_tx,
&VisibleMessageIdentity {
message_id: message.id.clone(),
created_at: message.created_at,
},
&message.content,
message.reasoning.as_deref(),
)
.await;
}
#[cfg(test)]
#[allow(clippy::too_many_arguments)]
async fn handle_no_tool_calls(
content: String,
reasoning: Option<String>,
reasoning_signature: Option<String>,
prompt_tokens: u64,
completion_tokens: u64,
round_usage: MetricsTokenUsage,
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
event_tx: &mpsc::Sender<AgentEvent>,
metrics_collector: Option<&MetricsCollector>,
round_id: &str,
session_id: &str,
config: &AgentLoopConfig,
task_context: &Option<TaskLoopContext>,
eval_model: &str,
iteration: u32,
llm: Arc<dyn LLMProvider>,
) -> Result<TurnOutcome, AgentError> {
handle_no_tool_calls_with_native(
content,
None,
reasoning,
reasoning_signature,
prompt_tokens,
completion_tokens,
round_usage,
session,
runtime_state,
event_tx,
metrics_collector,
round_id,
session_id,
config,
task_context,
eval_model,
iteration,
llm.clone(),
Vec::new(),
llm,
&CancellationToken::new(),
)
.await
}
#[allow(clippy::too_many_arguments)]
async fn handle_tool_calls_path(
frame: &crate::runtime::runner::round_frame::RoundFrame<'_>,
stream_output: StreamHandlingOutput,
mut round_usage: MetricsTokenUsage,
session: &mut Session,
runtime_state: &mut AgentRuntimeState,
policy_guard: &mut crate::runtime::runner::tool_execution::ToolPolicyGuard,
auxiliary_models: &crate::runtime::config::AuxiliaryModelConfig,
model_name: &str,
task_context: &mut Option<TaskLoopContext>,
cancel_token: &CancellationToken,
) -> Result<TurnOutcome, AgentError> {
let reasoning = (!stream_output.reasoning_content.trim().is_empty())
.then_some(stream_output.reasoning_content);
let reasoning_signature = reasoning
.as_ref()
.and_then(|_| stream_output.reasoning_signature.clone());
let mut native_items = Some(stream_output.provider_transcript_items.clone());
commit_assistant_message(
session,
assistant_with_visible_identity(
Message::assistant_with_reasoning(
stream_output.content,
Some(stream_output.tool_calls.clone()),
reasoning,
)
.with_reasoning_signature(reasoning_signature),
stream_output.visible_message,
),
&mut native_items,
)?;
if let Some(persistence) = frame.config.persistence.as_ref() {
persistence
.save_runtime_session(session)
.await
.map_err(|error| {
AgentError::Tool(format!(
"assistant tool-call checkpoint could not be persisted: {error}"
))
})?;
}
let compression_model = Some(model_name.to_string())
.or_else(|| (!session.model.trim().is_empty()).then_some(session.model.trim().to_string()));
if compression_model.is_none() {
tracing::warn!(
"[{}] Skipping mid-turn context compression after tool execution: missing model name",
frame.session_id
);
}
let eligible_tool_schemas =
resolve_available_tool_schemas_for_session(frame.config, frame.tools.as_ref(), session);
let required_tool = crate::runtime::runner::round_lifecycle::required_tool_for_session(session);
let request_tool_schemas = crate::runtime::runner::round_lifecycle::effective_tool_schemas(
session,
&eligible_tool_schemas,
);
let tool_schemas = request_tool_schemas.as_ref();
let capability_loading_mode = frame
.llm
.capability_loading_mode(model_name, required_tool)
.await;
let effective_callable_set =
effective_callable_set_for_round(session, tool_schemas, capability_loading_mode);
let tool_execution = tokio::select! {
biased;
_ = cancel_token.cancelled() => return Err(AgentError::Cancelled),
result = crate::runtime::runner::tool_execution::execute_round_tool_calls(
crate::runtime::runner::tool_execution::RoundToolExecution {
tool_calls: &stream_output.tool_calls,
frame,
session,
runtime_state,
policy_guard,
task_context,
compression_model_name: compression_model
.as_deref()
.or(auxiliary_models.background_model_name.as_deref()),
compression_model_provider: auxiliary_models
.summarization_model_provider
.as_ref()
.or(auxiliary_models.background_model_provider.as_ref()),
tool_schemas,
effective_callable_set: &effective_callable_set,
},
) => result?,
};
let mut awaiting_clarification = false;
let mut waiting_for_children = false;
let mut round_status = MetricsRoundStatus::Success;
let mut round_error: Option<String> = None;
if tool_execution.round_status != MetricsRoundStatus::Success {
round_status = tool_execution.round_status;
}
if let Some(e) = tool_execution.round_error {
round_error = Some(e);
}
if tool_execution.awaiting_clarification {
awaiting_clarification = true;
}
if tool_execution.waiting_for_children {
waiting_for_children = true;
}
let ticket_question_yield = frame.config.ticket_worker_plan.is_some()
&& session
.metadata
.contains_key(crate::ticket_worker_plan::TICKET_QUESTION_YIELD_KEY);
if ticket_question_yield || awaiting_clarification || waiting_for_children {
crate::runtime::runner::metrics_lifecycle::record_round_completed(
frame.metrics_collector,
frame.round_id,
frame.session_id,
session.messages.len() as u32,
round_status,
round_usage,
session
.token_usage
.as_ref()
.map(|usage| usage.prompt_cached_tool_outputs)
.unwrap_or(0)
.min(u32::MAX as usize) as u32,
session
.token_usage
.as_ref()
.map(|usage| usage.prompt_cached_tool_tokens_saved)
.unwrap_or(0),
round_error,
);
if ticket_question_yield {
session.metadata.insert(
"runtime.completion_reason".into(),
"ticket_question_yield".into(),
);
if !super::defers_terminal_complete(frame.config) {
let _ = frame
.event_tx
.send(AgentEvent::Complete {
usage: to_event_token_usage(
round_usage.prompt_tokens,
round_usage.completion_tokens,
),
})
.await;
}
}
return Ok(TurnOutcome {
should_break: true,
sent_complete: ticket_question_yield && !super::defers_terminal_complete(frame.config),
});
}
if frame.debug_enabled {
tracing::debug!(
"[{}] round_complete: {}",
frame.session_id,
serde_json::json!({
"round": frame.turn + 1,
"message_count": session.messages.len(),
})
);
}
let _complexity = if frame.config.features_dynamic_model_routing {
let round_tool_calls = &stream_output.tool_calls;
let classifier_model = auxiliary_models
.fast_model_name
.as_deref()
.or(Some(model_name));
let _classifier_provider = auxiliary_models
.fast_model_provider
.clone()
.unwrap_or_else(|| frame.llm.clone());
if let Some(_model) = classifier_model {
let complexity = heuristic_complexity(round_tool_calls);
tracing::info!(
"[{}] Dynamic model routing: round {} complexity={:?}",
frame.session_id,
frame.turn + 1,
complexity
);
session.metadata.insert(
"last_round_complexity".to_string(),
format!("{:?}", complexity),
);
Some(complexity)
} else {
None
}
} else {
None
};
round_usage.recompute_total();
crate::runtime::runner::metrics_lifecycle::record_round_completed(
frame.metrics_collector,
frame.round_id,
frame.session_id,
session.messages.len() as u32,
round_status,
round_usage,
session
.token_usage
.as_ref()
.map(|usage| usage.prompt_cached_tool_outputs)
.unwrap_or(0)
.min(u32::MAX as usize) as u32,
session
.token_usage
.as_ref()
.map(|usage| usage.prompt_cached_tool_tokens_saved)
.unwrap_or(0),
round_error,
);
Ok(TurnOutcome {
should_break: false,
sent_complete: false,
})
}
use crate::runtime::runner::session_setup::legacy_instruction::{
apply_successful_attempt as apply_successful_explicit_activation,
validate_first_step as validate_explicit_activation_first_step,
};
pub(super) async fn run_pipeline(
session: &mut Session,
event_tx: &mpsc::Sender<AgentEvent>,
llm: Arc<dyn LLMProvider>,
tools: Arc<dyn ToolExecutor>,
cancel_token: &CancellationToken,
config: &AgentLoopConfig,
state: &mut LoopRunState,
) -> super::super::Result<bool> {
let result =
boxed_pipeline_inner(session, event_tx, llm, tools, cancel_token, config, state).await;
let result = if config.effective_counted_loop().is_some()
&& cancel_token.is_cancelled()
&& matches!(result, Ok(false))
{
Err(AgentError::Cancelled)
} else {
result
};
let reason = if cancel_token.is_cancelled() {
"run_cancelled"
} else if result.is_err() {
"run_failed"
} else if session.metadata.contains_key("runtime.suspend_reason") {
"run_suspended"
} else {
"run_completed"
};
abort_in_flight_evaluations(state, event_tx, reason).await;
state.current_inputs = None;
result
}
fn boxed_pipeline_inner<'a>(
session: &'a mut Session,
event_tx: &'a mpsc::Sender<AgentEvent>,
llm: Arc<dyn LLMProvider>,
tools: Arc<dyn ToolExecutor>,
cancel_token: &'a CancellationToken,
config: &'a AgentLoopConfig,
state: &'a mut LoopRunState,
) -> std::pin::Pin<Box<impl std::future::Future<Output = super::super::Result<bool>> + 'a>> {
Box::pin(run_pipeline_inner(
session,
event_tx,
llm,
tools,
cancel_token,
config,
state,
))
}
async fn run_pipeline_inner(
session: &mut Session,
event_tx: &mpsc::Sender<AgentEvent>,
llm: Arc<dyn LLMProvider>,
tools: Arc<dyn ToolExecutor>,
cancel_token: &CancellationToken,
config: &AgentLoopConfig,
state: &mut LoopRunState,
) -> super::super::Result<bool> {
if config.run_budget.max_rounds == Some(0) {
return Err(AgentError::Budget(
"run_budget.max_rounds must be at least 1".to_string(),
));
}
let mut sent_complete = false;
let mut turn_counter: u32 = 0;
let mut max_rounds_summary_used = false;
let mut budget_summary_used = false;
session.metadata.remove("runtime.completion_reason");
session.metadata.remove("runtime.final_evidence_check");
session.metadata.remove("runtime.budget_exceeded_kind");
let mut tool_policy_guard = crate::runtime::runner::tool_execution::ToolPolicyGuard::new(
config.max_tool_calls_per_round,
config.max_consecutive_failures_per_tool,
);
loop {
if config.effective_counted_loop().is_some() && cancel_token.is_cancelled() {
return Err(AgentError::Cancelled);
}
if crate::session_app::no_progress::resume_after_user_message(
session,
&mut state.runtime_state,
) {
tool_policy_guard.reset_observation_progress();
}
if crate::session_app::no_progress::retain_pending_pause(session, &mut state.runtime_state)
{
break;
}
crate::runtime::loop_control::scoped::drain(session, config, turn_counter).await?;
if !crate::runtime::counted_loop::before_round(session, config, &state.execution_id).await?
{
break;
}
if let Some(message) = tool_policy_guard.delegation_failure_message() {
let error = if cancel_token.is_cancelled() {
AgentError::Cancelled
} else {
AgentError::Tool(message)
};
if let Some(metrics) = state.metrics_collector.as_ref() {
metrics.session_completed(
state.session_id.clone(),
map_turn_error_status(&error).1,
Utc::now(),
);
}
state_bridge::write_runtime_state(session, &state.runtime_state);
abort_in_flight_evaluations(state, event_tx, "delegation_failure_limit").await;
return Err(error);
}
refresh_auxiliary_models_for_round(state, config);
poll_completed_task_evaluation(state).await;
apply_completed_task_evaluation(session, event_tx, config, state).await;
if state.task_evaluation.in_flight.is_none() {
if let Some(request) = state.task_evaluation.queued_request.take() {
let eval_provider = state
.auxiliary_models
.fast_model_provider
.clone()
.unwrap_or_else(|| llm.clone());
spawn_task_evaluation_request(
state,
event_tx,
request,
eval_provider,
cancel_token.clone(),
config.auxiliary_evaluation_max_concurrency,
);
}
}
poll_completed_gold_evaluation(state).await;
apply_completed_gold_evaluation(session, config, state).await;
start_queued_gold_evaluation_if_idle(
state,
event_tx,
state
.auxiliary_models
.fast_model_provider
.clone()
.unwrap_or_else(|| llm.clone()),
cancel_token.clone(),
config.auxiliary_evaluation_max_concurrency,
);
state.runtime_state.round.current_round = turn_counter;
let round_id = crate::runtime::runner::round_prelude::build_round_id(
&state.session_id,
&state.execution_id,
turn_counter as usize,
);
state.runtime_state.round.last_round_id = Some(round_id.clone());
if config.effective_counted_loop().is_some() {
state_bridge::write_runtime_state(session, &state.runtime_state);
}
if config
.hook_runner
.has_hooks_for(AgentHookPoint::BeforeRound)
{
let outcome = config
.hook_runner
.run_hooks(
AgentHookPoint::BeforeRound,
&HookPayload::Round {
round: turn_counter + 1,
},
session,
&mut state.runtime_state,
Some(event_tx),
)
.await;
let hook_result = crate::runtime::hooks::apply_hook_outcome(
AgentHookPoint::BeforeRound,
outcome,
session,
&mut state.runtime_state,
);
state_bridge::write_runtime_state(session, &state.runtime_state);
hook_result?;
}
let runtime_context = PromptMemoryRuntimeContext {
llm: state
.auxiliary_models
.background_model_provider
.clone()
.unwrap_or_else(|| llm.clone()),
background_model_name: state.auxiliary_models.background_model_name.clone(),
};
let (prompt_memory_exposure, observation) =
crate::runtime::runner::round_prelude::refresh_round_boundary_with_observation(
session,
&mut state.runtime_state,
config,
Some(event_tx),
cancel_token,
state.metrics_collector.as_ref(),
Some(&runtime_context),
&state.execution_id,
)
.await?;
let fresh_human_input = observation.new_inputs().is_some_and(|batch| {
batch.records().iter().any(|input| {
(input.source == bamboo_domain::SessionMessageSource::User
&& input.kind == bamboo_domain::SessionMessageKind::UserInput)
|| input.wrapper.as_deref() == Some("root_chat_turn_v1")
})
});
observation.update_current(
&mut state.current_inputs,
&state.session_id,
&state.execution_id,
);
let resumed_progress_question = crate::session_app::no_progress::resume_after_user_message(
session,
&mut state.runtime_state,
);
if fresh_human_input || resumed_progress_question {
tool_policy_guard.reset_observation_progress();
}
let feedback_active = crate::runtime::runner::tool_execution::progress_feedback_active(
session,
&state.runtime_state,
config,
);
let observation_progress_hint =
match tool_policy_guard.feedback_for_next_round(feedback_active) {
crate::runtime::runner::tool_execution::FeedbackDecision::Continue(hint) => hint,
crate::runtime::runner::tool_execution::FeedbackDecision::Exhausted {
feedback,
summary,
} => {
crate::runtime::runner::tool_execution::finish_exhausted_progress(
session,
&state.runtime_state,
event_tx,
config,
&feedback,
&summary,
)
.await?;
break;
}
};
if let Some(ctx) = state.task_context.as_mut() {
ctx.current_round = turn_counter;
ctx.max_rounds = config.run_budget.max_rounds;
}
if state.debug_logger.enabled {
tracing::debug!(
"[{}] round_start: {}",
state.session_id,
serde_json::json!({
"round": turn_counter + 1,
"round_cap": config.run_budget.max_rounds,
"message_count": session.messages.len(),
})
);
}
let _ = event_tx
.send(AgentEvent::RunnerProgress {
session_id: state.session_id.clone(),
round_count: turn_counter,
})
.await;
crate::runtime::runner::metrics_lifecycle::record_round_started(
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
&state.model_name,
);
let mut overflow_recovery_attempted = false;
let mut turn_outcome: Option<TurnOutcome> = None;
let mut terminal_error: Option<AgentError> = None;
let mut hook_suspension: Option<AgentError> = None;
let mut round_activity = RoundActivity::default();
if let Some(host) = config.sdk_skill_execution_host.as_ref() {
host.observe_current_inputs(
&state.session_id,
&state.execution_id,
state.current_inputs.as_ref(),
)?;
}
let current_inputs = state.current_inputs.take();
let (round_result, current_inputs) =
crate::runtime::managers::lifecycle::scope_input_request_data(current_inputs, async {
if let Some(host) = config.sdk_skill_execution_host.as_ref() {
let context = host.render_skill_prompt(session, &state.execution_id).await?;
if context.is_empty() {
session.metadata.remove("skill.context");
} else {
session.metadata.insert("skill.context".into(), context);
}
super::super::session_setup::refresh_prompt_snapshot(session);
}
let tool_schemas = resolve_tool_schemas_for_round(config, tools.as_ref(), session);
if config.goal_loop_active() {
let goal = crate::runtime::goal_state::ensure_goal_state(
session,
config.active_goal().expect("active goal"),
);
crate::runtime::goal_state::write_goal_state(session, goal);
}
let extra_attempts = if config.goal_loop_active() {
config
.gold_config
.as_ref()
.map_or(0, |gold| gold.recovery.max_attempts.min(10)) as usize
} else {
0
};
for attempt in 1..=MAX_LLM_TURN_ATTEMPTS + extra_attempts {
if config.goal_loop_active() && extra_attempts > 0 {
if let Some(delay_ms) = crate::runtime::goal_recovery::pending_delay(
session,
Utc::now().timestamp_millis(),
) {
tokio::select! {
_ = cancel_token.cancelled() => {
terminal_error = Some(AgentError::Cancelled);
break;
},
_ = tokio::time::sleep(Duration::from_millis(delay_ms)) => {}
}
}
}
let attempt_tail_message_id = session.messages.last().map(|message| message.id.clone());
let llm_output = match crate::runtime::runner::round_lifecycle::execute_llm_round(
session,
config,
&llm,
event_tx,
cancel_token,
&state.session_id,
&state.model_name,
&tool_schemas,
observation_progress_hint.as_deref(),
Some(
crate::runtime::runner::round_lifecycle::PromptMemoryExposureFrame {
round_id: &round_id,
provenance: &prompt_memory_exposure,
},
),
)
.await
{
Ok(output) => output,
Err(error) => {
if is_overflow_recoverable(&error) && !overflow_recovery_attempted {
overflow_recovery_attempted = true;
if !state.overflow_recovery.can_attempt_recovery() {
let breaker_error = AgentError::LLMOverflow(format!(
"overflow recovery circuit breaker opened after {} consecutive recoveries",
state.overflow_recovery.consecutive_recoveries
));
tracing::error!(
"[{}] Turn {} overflow recovery skipped by circuit breaker: {}",
state.session_id,
turn_counter + 1,
breaker_error,
);
terminal_error = Some(breaker_error);
break;
}
tracing::warn!(
"[{}] Turn {} detected overflow error (attempt {}/{}): {}. Trying forced overflow recovery.",
state.session_id,
turn_counter + 1,
attempt,
MAX_LLM_TURN_ATTEMPTS,
error,
);
let overflow_tool_schemas =
crate::runtime::runner::round_lifecycle::context_tool_schemas_for_session(
session,
&llm,
&state.model_name,
&tool_schemas,
)
.await;
let recovered =
match crate::runtime::runner::round_lifecycle::force_overflow_context_recovery(
session,
config,
&state.model_name,
&state.session_id,
overflow_tool_schemas.as_ref(),
&llm,
Some(event_tx),
)
.await
{
Ok(recovered) => recovered,
Err(error) => {
terminal_error = Some(error);
break;
}
};
if recovered {
state
.overflow_recovery
.record_recovery(turn_counter as usize);
tracing::info!(
"[{}] Overflow recovery applied: total_recoveries={}, consecutive_recoveries={}, turn={}",
state.session_id,
state.overflow_recovery.total_recoveries,
state.overflow_recovery.consecutive_recoveries,
turn_counter + 1,
);
let tool_schemas_after_recovery =
resolve_available_tool_schemas_for_session(
config,
tools.as_ref(),
session,
);
match crate::runtime::runner::round_lifecycle::execute_llm_round(
session,
config,
&llm,
event_tx,
cancel_token,
&state.session_id,
&state.model_name,
&tool_schemas_after_recovery,
observation_progress_hint.as_deref(),
Some(
crate::runtime::runner::round_lifecycle::PromptMemoryExposureFrame {
round_id: &round_id,
provenance: &prompt_memory_exposure,
},
),
)
.await
{
Ok(output) => output,
Err(recovery_error) => {
tracing::error!(
"[{}] Turn {} overflow recovery retry failed: {}",
state.session_id,
turn_counter + 1,
recovery_error,
);
terminal_error = Some(recovery_error);
break;
}
}
} else {
tracing::error!(
"[{}] Turn {} overflow recovery was attempted but no compression was applied.",
state.session_id,
turn_counter + 1,
);
terminal_error = Some(error);
break;
}
} else if should_retry_turn_error(&error) && attempt < MAX_LLM_TURN_ATTEMPTS {
if let Some(message_id) = crate::runtime::runner::round_lifecycle::discard_latest_interrupted_assistant_output(
session,
attempt_tail_message_id.as_deref(),
) {
if event_tx
.send(AgentEvent::VisibleMessageDiscard { message_id })
.await
.is_err()
{
tracing::warn!(
"[{}] event channel closed; visible retry rollback was not broadcast",
state.session_id,
);
}
}
let delay_ms = LLM_RETRY_BASE_DELAY_MS * (1u64 << (attempt - 1));
tracing::warn!(
"[{}] Turn {} LLM call failed (attempt {}/{}): {}. Retrying in {}ms",
state.session_id,
turn_counter + 1,
attempt,
MAX_LLM_TURN_ATTEMPTS,
error,
delay_ms
);
tokio::select! {
_ = cancel_token.cancelled() => {
terminal_error = Some(AgentError::Cancelled);
break;
},
_ = tokio::time::sleep(Duration::from_millis(delay_ms)) => {}
}
continue;
} else {
let mut recovery_budget_state = state.runtime_state.clone();
round_activity.commit_to_runtime(&mut recovery_budget_state);
let recovery_delay = if attempt >= MAX_LLM_TURN_ATTEMPTS
&& config.goal_loop_active()
&& state.runtime_state.suspension.is_none()
&& state.runtime_state.waiting_for_children.is_none()
&& state.runtime_state.waiting_for_bash.is_none()
&& check_run_budget_exceeded(
&recovery_budget_state.round,
&config.run_budget,
)
.is_none()
&& config.persistence.is_some()
{
crate::runtime::goal_recovery::reserve_retry(
session,
&config
.gold_config
.as_ref()
.expect("active goal config")
.recovery,
&error,
Utc::now().timestamp_millis(),
)
} else {
None
};
if let Some(delay_ms) = recovery_delay {
state_bridge::write_runtime_state(session, &recovery_budget_state);
config
.persistence
.as_ref()
.expect("checked persistence")
.checkpoint_runtime_session(session)
.await
.map_err(|error| {
AgentError::LLM(format!(
"Goal recovery checkpoint failed: {error}"
))
})?;
tokio::select! {
_ = cancel_token.cancelled() => {
terminal_error = Some(AgentError::Cancelled);
break;
},
_ = tokio::time::sleep(Duration::from_millis(delay_ms)) => {}
}
continue;
}
tracing::error!(
"[{}] Turn {} LLM call failed terminally (attempt {}/{}): {}",
state.session_id,
turn_counter + 1,
attempt,
MAX_LLM_TURN_ATTEMPTS,
error,
);
terminal_error = Some(error);
break;
}
}
};
let stream_output = llm_output.stream_output;
round_activity.absorb_attempt(&stream_output, llm_output.attempt_usage);
if let Some(error) = llm_output.terminal_validation_error {
terminal_error = Some(error);
break;
}
let activation_attempt =
match validate_explicit_activation_first_step(session, &stream_output.tool_calls) {
Ok(attempt) => attempt,
Err(error) => {
terminal_error = Some(error);
break;
}
};
if stream_output.tool_calls.is_empty() {
if crate::runtime::runner::round_lifecycle::is_openai_client_tool_search_boundary(
&stream_output.provider_transcript_items,
) {
match commit_openai_client_tool_search_round(
stream_output,
session,
config,
&tool_schemas,
)
.await
{
Ok(message) => {
if config.features_final_evidence_check {
publish_buffered_message(event_tx, &message).await;
}
record_no_tool_calls_round_completed(
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
session,
round_activity.token_usage(),
);
turn_outcome = Some(TurnOutcome {
should_break: false,
sent_complete: false,
});
}
Err(error) => terminal_error = Some(error),
}
break;
}
let inherited_untagged_wait = state
.runtime_state
.waiting_for_children
.as_ref()
.is_some_and(|wait| wait.registered_by_tool_call_id.is_none());
match maybe_suspend_for_orphaned_children(session, config, &mut state.runtime_state)
.await
{
Ok(Some(suspend)) => {
if inherited_untagged_wait {
let reasoning = (!stream_output.reasoning_content.trim().is_empty())
.then_some(stream_output.reasoning_content);
let reasoning_signature = reasoning
.as_ref()
.and_then(|_| stream_output.reasoning_signature.clone());
let message = assistant_with_visible_identity(
Message::assistant_with_reasoning(
stream_output.content,
None,
reasoning,
)
.with_reasoning_signature(reasoning_signature),
stream_output.visible_message,
);
let buffered_message = config
.features_final_evidence_check
.then(|| message.clone());
let mut native_items = Some(stream_output.provider_transcript_items);
if let Err(error) =
commit_assistant_message(session, message, &mut native_items)
{
terminal_error = Some(error);
break;
}
if let Some(message) = buffered_message.as_ref() {
publish_buffered_message(event_tx, message).await;
}
record_no_tool_calls_round_completed(
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
session,
round_activity.token_usage(),
);
}
turn_outcome = Some(suspend);
break;
}
Ok(None) => {}
Err(error) => {
terminal_error = Some(error);
break;
}
}
if crate::runtime::loop_control::suspension::active(config)
&& !bamboo_tools::tools::bash_runtime::running_shells_for_session(&session.id).is_empty()
{
terminal_error = Some(AgentError::Tool(
"managed loop cannot establish an asynchronous shell wait".into(),
));
break;
}
if let Some(suspend) =
maybe_suspend_for_outstanding_bash(session, config, &mut state.runtime_state)
.await
{
turn_outcome = Some(suspend);
break;
}
let reasoning = (!stream_output.reasoning_content.trim().is_empty())
.then_some(stream_output.reasoning_content);
let reasoning_signature = reasoning
.as_ref()
.and_then(|_| stream_output.reasoning_signature.clone());
let eval_model = state
.auxiliary_models
.fast_model_name
.clone()
.unwrap_or_else(|| state.model_name.clone());
match handle_no_tool_calls_with_native(
stream_output.content,
stream_output.visible_message,
reasoning,
reasoning_signature,
llm_output.prompt_tokens,
llm_output.completion_tokens,
round_activity.token_usage(),
session,
&mut state.runtime_state,
event_tx,
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
config,
&state.task_context,
&eval_model,
turn_counter + 1,
llm.clone(),
stream_output.provider_transcript_items,
state
.auxiliary_models
.fast_model_provider
.clone()
.unwrap_or_else(|| llm.clone()),
cancel_token,
)
.await
{
Ok(outcome) => turn_outcome = Some(outcome),
Err(error) if error.is_hook_suspended() => {
record_no_tool_calls_round_completed(
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
session,
round_activity.token_usage(),
);
hook_suspension = Some(error);
}
Err(error) => terminal_error = Some(error),
}
break;
}
let required_tool =
crate::runtime::runner::round_lifecycle::required_tool_for_session(session);
let capability_loading_mode = llm
.capability_loading_mode(&state.model_name, required_tool)
.await;
let legacy_browser_discovery = capability_loading_mode
== CapabilityLoadingMode::LegacyFullCatalog
&& legacy_browser_needs_discovery(session, &tool_schemas, required_tool);
if (capability_loading_mode == CapabilityLoadingMode::StickyFallback
|| legacy_browser_discovery)
&& stream_output.tool_calls.iter().any(|call| {
call.function.name == bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME
})
{
match commit_sticky_fallback_discovery_round(
stream_output,
session,
config,
&tool_schemas,
legacy_browser_discovery,
)
.await
{
Ok(()) => {
record_no_tool_calls_round_completed(
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
session,
round_activity.token_usage(),
);
turn_outcome = Some(TurnOutcome {
should_break: false,
sent_complete: false,
});
}
Err(error) => terminal_error = Some(error),
}
break;
}
let frame = crate::runtime::runner::round_frame::RoundFrame {
vision_support: llm.supports_vision(&state.model_name).await,
session_id: &state.session_id,
round_id: &round_id,
turn: turn_counter as usize,
debug_enabled: state.debug_logger.enabled,
event_tx,
metrics_collector: state.metrics_collector.as_ref(),
config,
llm: &llm,
tools: &tools,
};
#[cfg(test)]
let injected_handler_error = take_test_post_llm_retry_failure(session);
#[cfg(not(test))]
let injected_handler_error: Option<AgentError> = None;
let handler_result = match injected_handler_error {
Some(error) => Err(error),
None => {
handle_tool_calls_path(
&frame,
stream_output,
round_activity.token_usage(),
session,
&mut state.runtime_state,
&mut tool_policy_guard,
&state.auxiliary_models,
&state.model_name,
&mut state.task_context,
cancel_token,
)
.await
}
};
match handler_result {
Ok(outcome) => {
if let Some(attempt) = activation_attempt.as_ref() {
if !outcome.should_break {
if let Err(error) =
apply_successful_explicit_activation(session, attempt)
{
terminal_error = Some(error);
break;
}
}
}
turn_outcome = Some(outcome);
break;
}
Err(error) => {
if should_retry_turn_error(&error) && attempt < MAX_LLM_TURN_ATTEMPTS {
if let Some(message_id) = crate::runtime::runner::round_lifecycle::discard_latest_interrupted_assistant_output(
session,
attempt_tail_message_id.as_deref(),
) {
if event_tx
.send(AgentEvent::VisibleMessageDiscard { message_id })
.await
.is_err()
{
tracing::warn!(
"[{}] event channel closed; visible post-LLM retry rollback was not broadcast",
state.session_id,
);
}
}
let delay_ms = LLM_RETRY_BASE_DELAY_MS * (1u64 << (attempt - 1));
tracing::warn!(
"[{}] Turn {} post-LLM handling failed (attempt {}/{}): {}. Retrying in {}ms",
state.session_id,
turn_counter + 1,
attempt,
MAX_LLM_TURN_ATTEMPTS,
error,
delay_ms
);
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
continue;
}
tracing::error!(
"[{}] Turn {} post-LLM handling failed terminally (attempt {}/{}): {}",
state.session_id,
turn_counter + 1,
attempt,
MAX_LLM_TURN_ATTEMPTS,
error,
);
terminal_error = Some(error);
break;
}
}
}
Ok::<(), AgentError>(())
}).await;
state.current_inputs = current_inputs;
round_result?;
round_activity.commit_to_runtime(&mut state.runtime_state);
if let Some(error) = hook_suspension {
return Err(error);
}
if let Some(error) = terminal_error {
state_bridge::write_runtime_state(session, &state.runtime_state);
record_turn_failure(
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
session.messages.len() as u32,
round_activity.token_usage(),
&error,
);
abort_in_flight_evaluations(state, event_tx, "terminal_error").await;
return Err(error);
}
let Some(outcome) = turn_outcome else {
let error = AgentError::LLM(format!(
"[{}] turn {} completed without outcome",
state.session_id,
turn_counter + 1
));
state_bridge::write_runtime_state(session, &state.runtime_state);
record_turn_failure(
state.metrics_collector.as_ref(),
&round_id,
&state.session_id,
session.messages.len() as u32,
round_activity.token_usage(),
&error,
);
abort_in_flight_evaluations(state, event_tx, "run_stopped").await;
return Err(error);
};
if !overflow_recovery_attempted {
state.overflow_recovery.reset_after_stable_round();
}
state.runtime_state.memory.overflow_recovery_total =
state.overflow_recovery.total_recoveries as u32;
state.runtime_state.memory.overflow_recovery_consecutive =
state.overflow_recovery.consecutive_recoveries as u32;
match session
.metadata
.get("runtime.suspend_reason")
.map(String::as_str)
{
Some("awaiting_clarification") => {
state.runtime_state.status = AgentStatusState::Suspended;
state.runtime_state.suspension = Some(SuspensionState {
reason: "awaiting_clarification".to_string(),
suspended_at: Utc::now(),
resumable: true,
hook_point: Some("AfterToolExecution".to_string()),
});
}
Some("awaiting_parent_approval") => {
state.runtime_state.status = AgentStatusState::Suspended;
state.runtime_state.suspension = Some(SuspensionState {
reason: "awaiting_parent_approval".to_string(),
suspended_at: Utc::now(),
resumable: true,
hook_point: Some("AfterToolExecution".to_string()),
});
}
Some("waiting_for_children") => {
state.runtime_state.status = AgentStatusState::Suspended;
state.runtime_state.suspension = Some(SuspensionState {
reason: "waiting_for_children".to_string(),
suspended_at: Utc::now(),
resumable: true,
hook_point: Some("AfterToolExecution".to_string()),
});
if let Some(storage) = config.storage.as_ref() {
if let Ok(Some(persisted)) = storage.load_session(&state.session_id).await {
if let Some(runtime_state) = persisted.agent_runtime_state {
reconcile_child_wait_at_suspend(
session,
&mut state.runtime_state,
runtime_state.waiting_for_children,
);
}
let existing_ids: std::collections::HashSet<String> = session
.messages
.iter()
.map(|message| message.id.clone())
.collect();
let mut appended = 0usize;
for message in persisted.messages {
let hidden_runtime_resume = message
.metadata
.as_ref()
.and_then(|metadata| metadata.get("runtime_kind"))
.and_then(|value| value.as_str())
.is_some_and(|kind| {
matches!(
kind,
"child_completion_resume" | "guardian_review_resume"
)
});
if hidden_runtime_resume && !existing_ids.contains(message.id.as_str())
{
session.messages.push(message);
appended += 1;
}
}
if appended > 0 {
tracing::info!(
"[{}] Preserved {} hidden child-completion resume message(s) during parent suspension save",
state.session_id,
appended
);
}
}
}
}
Some("waiting_for_bash") => {
state.runtime_state.status = AgentStatusState::Suspended;
state.runtime_state.suspension = Some(SuspensionState {
reason: "waiting_for_bash".to_string(),
suspended_at: Utc::now(),
resumable: true,
hook_point: Some("AfterToolExecution".to_string()),
});
if let Some(storage) = config.storage.as_ref() {
if let Ok(Some(persisted)) = storage.load_session(&state.session_id).await {
if let Some(runtime_state) = persisted.agent_runtime_state {
if runtime_state.waiting_for_bash.is_some() {
state.runtime_state.waiting_for_bash =
runtime_state.waiting_for_bash;
}
}
let existing_ids: std::collections::HashSet<String> = session
.messages
.iter()
.map(|message| message.id.clone())
.collect();
let mut appended = 0usize;
for message in persisted.messages {
let hidden_runtime_resume = message
.metadata
.as_ref()
.and_then(|metadata| metadata.get("runtime_kind"))
.and_then(|value| value.as_str())
.is_some_and(|kind| {
kind == crate::runtime::config::BASH_COMPLETION_RESUME_KIND
});
if hidden_runtime_resume && !existing_ids.contains(message.id.as_str())
{
session.messages.push(message);
appended += 1;
}
}
if appended > 0 {
tracing::info!(
"[{}] Preserved {} hidden bash-completion resume message(s) during suspension save",
state.session_id,
appended
);
}
}
}
}
_ => {}
}
if !config.hook_runner.is_empty() {
crate::runtime::hooks::merge_session_hook_checkpoints(
session,
&mut state.runtime_state,
);
}
if config.hook_runner.has_hooks_for(AgentHookPoint::AfterRound) {
let hook_outcome = config
.hook_runner
.run_hooks(
AgentHookPoint::AfterRound,
&HookPayload::Round {
round: turn_counter + 1,
},
session,
&mut state.runtime_state,
Some(event_tx),
)
.await;
let hook_result = crate::runtime::hooks::apply_hook_outcome(
AgentHookPoint::AfterRound,
hook_outcome,
session,
&mut state.runtime_state,
);
state_bridge::write_runtime_state(session, &state.runtime_state);
hook_result?;
}
state_bridge::write_runtime_state(session, &state.runtime_state);
sent_complete = sent_complete || outcome.sent_complete;
if outcome.should_break {
break;
}
if let Err(error) = spawn_task_evaluation_if_needed(
turn_counter as usize,
session,
event_tx,
config,
state,
state
.auxiliary_models
.fast_model_provider
.clone()
.unwrap_or_else(|| llm.clone()),
cancel_token.clone(),
) {
tracing::warn!(
"[{}] Failed to spawn async task evaluation after round {}: {}",
state.session_id,
turn_counter + 1,
error
);
}
if let Err(error) = spawn_gold_evaluation_if_needed(
turn_counter as usize,
session,
event_tx,
config,
state,
state
.auxiliary_models
.fast_model_provider
.clone()
.unwrap_or_else(|| llm.clone()),
cancel_token.clone(),
) {
tracing::warn!(
"[{}] Failed to spawn async Gold evaluation after round {}: {}",
state.session_id,
turn_counter + 1,
error
);
}
turn_counter += 1;
if (config.effective_counted_loop().is_some() || config.loop_control.is_some())
&& check_run_budget_exceeded(&state.runtime_state.round, &config.run_budget).is_some()
{
crate::runtime::counted_loop::exhaust_budget(session, config).await?;
break;
}
if let Some(exceeded) =
check_run_budget_exceeded(&state.runtime_state.round, &config.run_budget)
{
if !budget_summary_used {
tracing::warn!(
"[{}] Run budget exceeded ({} limit={} actual={}) — granting one summary turn before stopping.",
state.session_id,
exceeded.kind,
exceeded.limit,
exceeded.actual,
);
session.metadata.insert(
"runtime.completion_reason".to_string(),
"budget_exceeded".to_string(),
);
session.metadata.insert(
"runtime.budget_exceeded_kind".to_string(),
exceeded.kind.to_string(),
);
let mut summary = Message::user(format!(
"The run's resource budget ({}, limit={}, reached={}) was exceeded; the \
task was stopped before completion. Stop working now and summarize your \
progress so far and what remains.",
exceeded.kind, exceeded.limit, exceeded.actual
));
summary.metadata = Some(serde_json::json!({"runtime_kind": "run_budget_summary"}));
session.add_message(summary);
let _ = event_tx
.send(AgentEvent::BudgetExceeded {
session_id: state.session_id.clone(),
kind: exceeded.kind.to_string(),
limit: exceeded.limit,
actual: exceeded.actual,
})
.await;
budget_summary_used = true;
continue;
}
tracing::warn!(
"[{}] Run budget exceeded ({} limit={} actual={}) — stopping the run before completion.",
state.session_id,
exceeded.kind,
exceeded.limit,
exceeded.actual,
);
break;
}
if let Some(max_rounds) = config.run_budget.max_rounds {
if (config.effective_counted_loop().is_some() || config.loop_control.is_some())
&& turn_counter >= max_rounds
{
crate::runtime::counted_loop::exhaust_budget(session, config).await?;
break;
}
if turn_counter >= max_rounds {
if !max_rounds_summary_used {
tracing::warn!(
"[{}] Reached max rounds ({}) — granting one summary turn before stopping.",
state.session_id,
max_rounds
);
session.metadata.insert(
"runtime.completion_reason".to_string(),
"max_rounds_reached".to_string(),
);
let mut summary = Message::user(format!(
"Reached the maximum of {0} rounds; the task was stopped before \
completion. Stop working now and summarize your progress so far \
and what remains.",
max_rounds
));
summary.metadata =
Some(serde_json::json!({"runtime_kind": "max_rounds_summary"}));
session.add_message(summary);
max_rounds_summary_used = true;
continue;
}
tracing::warn!(
"[{}] Reached max rounds ({}) — stopping the run before completion.",
state.session_id,
max_rounds
);
break;
}
}
}
poll_completed_task_evaluation(state).await;
apply_completed_task_evaluation(session, event_tx, config, state).await;
poll_completed_gold_evaluation(state).await;
apply_completed_gold_evaluation(session, config, state).await;
if !session.metadata.contains_key("runtime.suspend_reason") {
if let Some(storage) = config.storage.as_ref() {
if let Ok(Some(persisted)) = storage.load_session(&state.session_id).await {
if let Some(persisted_state) = persisted.agent_runtime_state {
state.runtime_state.waiting_for_children = persisted_state.waiting_for_children;
state.runtime_state.waiting_for_bash = persisted_state.waiting_for_bash;
}
}
}
let reason = if state.runtime_state.waiting_for_children.is_some() {
Some(("waiting_for_children", "ChildCompletion"))
} else if state.runtime_state.waiting_for_bash.is_some() {
Some(("waiting_for_bash", "BashCompletion"))
} else {
None
};
if let Some((reason, hook_point)) = reason {
state.runtime_state.status = AgentStatusState::Suspended;
state.runtime_state.suspension = Some(SuspensionState {
reason: reason.to_string(),
suspended_at: Utc::now(),
resumable: true,
hook_point: Some(hook_point.to_string()),
});
session
.metadata
.insert("runtime.suspend_reason".to_string(), reason.to_string());
state_bridge::write_runtime_state(session, &state.runtime_state);
}
}
let evaluation_stop_reason = if session.metadata.contains_key("runtime.suspend_reason") {
"run_suspended"
} else {
"run_completed"
};
abort_in_flight_evaluations(state, event_tx, evaluation_stop_reason).await;
Ok(sent_complete)
}
fn heuristic_complexity(
tool_calls: &[bamboo_agent_core::tools::ToolCall],
) -> crate::runtime::complexity_classifier::TaskComplexity {
use crate::runtime::complexity_classifier::TaskComplexity;
let simple_tools = [
"Read",
"Glob",
"Grep",
"Bash",
"exec_command",
"write_stdin",
];
let complex_tools = ["Agent", "Plan", "SubAgent", "TodoWrite"];
let names: Vec<&str> = tool_calls
.iter()
.map(|tc| tc.function.name.as_str())
.collect();
if names.iter().any(|n| complex_tools.contains(n)) {
return TaskComplexity::Complex;
}
if names.iter().all(|n| simple_tools.contains(n)) && !names.is_empty() {
return TaskComplexity::Simple;
}
TaskComplexity::Standard
}
#[cfg(test)]
#[tokio::test]
async fn supervisor_delegation_first_discovery_is_readonly_bounded_and_keeps_ordinary_ultra_denial()
{
use bamboo_domain::{SessionAuthorityIdentity, DEFAULT_SUPERVISOR_SESSION_ID};
let mut supervisor = Session::new(DEFAULT_SUPERVISOR_SESSION_ID, "model");
supervisor.authority_identity = SessionAuthorityIdentity::Supervisor {
incarnation_id: uuid::Uuid::new_v4(),
};
let schemas = [
"Read",
"SubAgent",
"exec_command",
"load_skill",
"workflow_run",
"vendor::Read",
]
.into_iter()
.map(|name| bamboo_agent_core::tools::ToolSchema {
schema_type: "function".into(),
function: bamboo_agent_core::tools::FunctionSchema {
name: name.into(),
description: name.into(),
parameters: serde_json::json!({"type":"object"}),
},
})
.collect::<Vec<_>>();
let config = AgentLoopConfig::default();
for ultra in [false, true] {
supervisor.set_root_orchestration_only(ultra).unwrap();
let before = serde_json::to_value(&supervisor).unwrap();
ensure_discovery_allowed(&mut supervisor, &config)
.await
.unwrap();
let discovery =
CompleteCapabilityDiscovery::new(&supervisor, &config, &schemas, false, true)
.await
.unwrap();
assert!(
discovery.commands.is_none(),
"Supervisor must not open Skill/Workflow discovery"
);
assert_eq!(
discovery
.catalog
.iter()
.map(|entry| entry.execution_name())
.collect::<Vec<_>>(),
["Read", "SubAgent"]
);
let request = serde_json::from_value(
serde_json::json!({"query":"SubAgent","kinds":["tool"],"limit":1}),
)
.unwrap();
let found = discovery.discover_complete_schemas(&request).unwrap();
assert_eq!(
found
.iter()
.map(|schema| schema.function.name.as_str())
.collect::<Vec<_>>(),
["SubAgent"]
);
assert_eq!(serde_json::to_value(&supervisor).unwrap(), before);
}
let mut ordinary = Session::new("ordinary-ultra", "model");
ordinary.metadata.insert("role".into(), "supervisor".into());
ordinary.set_root_orchestration_only(true).unwrap();
assert!(ensure_discovery_allowed(&mut ordinary, &config)
.await
.is_err());
}
#[cfg(test)]
mod tests {
#[tokio::test]
async fn ql_main_pipeline_observes_real_new_then_terminal_or_cancel_drops_owner() {
for cancelled in [false, true] {
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 mut session = Session::new_child("ql-main", "parent", "model", "Child");
session.add_message(Message::system("base system"));
storage.save_session(&session).await.unwrap();
let mut expected = Vec::new();
for text in [
"first current input",
"second current input without selection",
] {
let envelope = SessionMessageEnvelope::user_input(&session.id, text);
expected.push(envelope.id.to_string());
let receipt = inbox.deliver(&envelope).await.unwrap();
inbox
.mark_activation_eligible(
&session.id,
receipt.generation,
SessionActivationPolicy::InterruptSpecificWait,
)
.await
.unwrap();
}
let config = AgentLoopConfig {
storage: Some(storage),
persistence: Some(persistence),
session_inbox: Some(inbox.clone()),
skip_initial_user_message: true,
model_name: Some("model".into()),
run_budget: bamboo_config::RunBudgetConfig {
max_rounds: Some(2),
..Default::default()
},
..Default::default()
};
let provider = Arc::new(ContextProbeProvider::default());
let tools = Arc::new(bamboo_tools::BuiltinToolExecutorBuilder::new().build());
let cancel = tokio_util::sync::CancellationToken::new();
let mut state = e2e_loop_state(&session.id);
let (tx, _rx) = tokio::sync::mpsc::channel(256);
assert!(super::run_pipeline_inner(
&mut session,
&tx,
provider.clone(),
tools.clone(),
&cancel,
&config,
&mut state
)
.await
.unwrap());
let batch = state
.current_inputs
.as_ref()
.expect("production shared prelude moved real New");
assert_eq!(
batch
.records()
.iter()
.map(|r| r.input_id.clone())
.collect::<Vec<_>>(),
expected
);
assert!(batch.records().iter().all(|r| r.request.is_none()));
assert_eq!(batch.execution_id(), state.execution_id);
assert_eq!(provider.calls.load(std::sync::atomic::Ordering::SeqCst), 1);
if cancelled {
cancel.cancel();
}
let result = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
tools,
&cancel,
&config,
&mut state,
)
.await;
assert!(
state.current_inputs.is_none(),
"outer terminal fence drops prior N for every return"
);
if cancelled {
assert!(matches!(result, Err(AgentError::Cancelled)));
} else {
assert!(result.unwrap());
}
assert_eq!(
provider.calls.load(std::sync::atomic::Ordering::SeqCst),
if cancelled { 1 } else { 2 }
);
assert_eq!(inbox.inspect(&session.id).await.unwrap().claimed, 0);
}
}
use super::super::startup::{InFlightTaskEvaluation, OverflowRecoveryState};
use super::{
apply_successful_explicit_activation, assistant_with_visible_identity,
build_guardian_review_prompt, build_openai_client_tool_search_outputs,
check_run_budget_exceeded, commit_assistant_message,
commit_openai_client_tool_search_round, commit_sticky_fallback_discovery_round,
effective_callable_set_for_round, is_child_spawn_call, is_overflow_recoverable,
is_terminal_child_status, map_turn_error_status, maybe_spawn_guardian_review,
maybe_suspend_for_orphaned_children, maybe_suspend_for_outstanding_bash,
scope_discovered_gateway_schema, should_retry_turn_error, sticky_fallback_definition_delta,
sticky_fallback_tool_result, sticky_result_definition_values, suspend_to_wait_for_bash,
validate_explicit_activation_first_step, validated_sticky_fallback_loaded_tool_names,
};
use crate::project_context::{
ProjectContextError, ProjectContextResolver, ProjectContextSource, ProjectDescriptor,
};
use crate::runtime::config::{AgentLoopConfig, GuardianConfig, GuardianSpawner};
use crate::runtime::goal_state::{
ensure_goal_state, read_goal_state, write_goal_state, GoalDeclaredStatus, GoalRuntimeStatus,
};
use crate::runtime::guardian_state::{
ensure_guardian_state, read_guardian_state, write_guardian_state, GuardianPhase,
GuardianVerdict,
};
use crate::runtime::runner::state_bridge;
use bamboo_agent_core::storage::Storage;
use bamboo_agent_core::{
AgentError, AgentEvent, AgentHook, Message, Session, StreamTimeoutError, StreamTimeoutPhase,
};
use bamboo_domain::session::runtime_state::{
AgentStatusState, ChildWaitPolicy, WaitingForChildrenState,
};
use bamboo_domain::{
AgentHookPoint, AgentRuntimeState, CapabilityLoadingMode, HookPayload, HookResult,
ProjectId, ProjectResourceSummary, SessionActivationPolicy, SessionInboxLimits,
SessionInboxPort, SessionMessageEnvelope,
};
use bamboo_llm::{LLMChunk, LLMError, LLMProvider, LLMStream};
use bamboo_metrics::{
RoundStatus as MetricsRoundStatus, SessionStatus as MetricsSessionStatus,
TokenUsage as MetricsTokenUsage,
};
use bamboo_skills::{SkillManager, SkillStoreConfig};
use chrono::Utc;
use futures::stream;
use std::collections::HashMap;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use tokio::sync::RwLock;
#[test]
fn visible_identity_is_reused_by_persisted_assistant_message() {
let created_at = chrono::DateTime::parse_from_rfc3339("2026-09-22T19:00:00Z")
.unwrap()
.with_timezone(&Utc);
let identity = crate::runtime::stream::handler::VisibleMessageIdentity {
message_id: "stable-visible-message".to_string(),
created_at,
};
let message = assistant_with_visible_identity(
Message::assistant("visible text", None),
Some(identity),
);
assert_eq!(message.id, "stable-visible-message");
assert_eq!(message.created_at, created_at);
assert_eq!(message.content, "visible text");
assert_eq!(message.role, bamboo_agent_core::Role::Assistant);
}
fn pending_explicit_session() -> Session {
let mut session = Session::new("explicit-gate", "model");
session.metadata.insert(
bamboo_skills::runtime_metadata::SKILL_RUNTIME_SELECTION_SOURCE_KEY.to_string(),
"explicit".to_string(),
);
session.metadata.insert(
bamboo_skills::runtime_metadata::SKILL_RUNTIME_SELECTED_SKILL_IDS_KEY.to_string(),
"[\"review\"]".to_string(),
);
session
}
fn activation_call(id: &str, name: &str, arguments: &str) -> ToolCall {
ToolCall {
id: id.to_string(),
tool_type: "function".to_string(),
function: FunctionCall {
name: name.to_string(),
arguments: arguments.to_string(),
},
}
}
fn native_client_search_item_with_arguments(
call_id: &str,
arguments: serde_json::Value,
) -> bamboo_domain::ProviderTranscriptItem {
bamboo_domain::ProviderTranscriptItem::try_from_payload(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
bamboo_domain::ProviderTranscriptOrigin::Provider,
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"tool_search_call","id":format!("tsc_pipeline_{call_id}"),
"execution":"client","call_id":call_id,
"status":"completed","arguments":arguments
}),
)
.unwrap()
}
fn native_client_search_item_for(
call_id: &str,
query: &str,
) -> bamboo_domain::ProviderTranscriptItem {
native_client_search_item_with_arguments(call_id, serde_json::json!({"query":query}))
}
fn native_client_search_item() -> bamboo_domain::ProviderTranscriptItem {
native_client_search_item_for("search_1", "orders")
}
fn loading_test_schema(name: &str) -> bamboo_agent_core::tools::ToolSchema {
loading_test_schema_with_description(name, "")
}
fn loading_test_schema_with_description(
name: &str,
description: &str,
) -> bamboo_agent_core::tools::ToolSchema {
bamboo_agent_core::tools::ToolSchema {
schema_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionSchema {
name: name.to_string(),
description: description.to_string(),
parameters: serde_json::json!({"type":"object"}),
},
}
}
#[test]
fn discovered_catalog_gateways_are_scoped_to_matching_ids_and_metadata() {
let skill_gateway =
bamboo_domain::ClassifiedToolSchema::new(bamboo_agent_core::tools::ToolSchema {
schema_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionSchema {
name: "load_skill".to_string(),
description: "Load one instruction Skill".to_string(),
parameters: serde_json::json!({
"type":"object",
"properties":{
"skill_id":{"type":"string"},
"detail":{"type":"string"}
},
"required":["skill_id"]
}),
},
})
.unwrap();
let skill_match = bamboo_domain::CapabilityMatch {
capability_ref: "skill:review-helper".to_string(),
kind: bamboo_domain::CapabilityKind::Skill,
display_name: "Review Helper".to_string(),
summary: "Review a change".to_string(),
source: bamboo_domain::CapabilitySource::User,
revision: Some(7),
status: bamboo_domain::CapabilityStatus::Valid,
invocation_policy: None,
invocation_target: bamboo_domain::CapabilityInvocationTarget::Skill {
name: "load_skill".to_string(),
skill_id: "review-helper".to_string(),
source: bamboo_domain::CapabilitySource::User,
revision: 7,
},
};
let second_skill_match = bamboo_domain::CapabilityMatch {
capability_ref: "skill:lint-helper".to_string(),
kind: bamboo_domain::CapabilityKind::Skill,
display_name: "Lint Helper".to_string(),
summary: "Run focused lint checks".to_string(),
source: bamboo_domain::CapabilitySource::Project,
revision: Some(3),
status: bamboo_domain::CapabilityStatus::Valid,
invocation_policy: None,
invocation_target: bamboo_domain::CapabilityInvocationTarget::Skill {
name: "load_skill".to_string(),
skill_id: "lint-helper".to_string(),
source: bamboo_domain::CapabilitySource::Project,
revision: 3,
},
};
let skill_schema =
scope_discovered_gateway_schema(&skill_gateway, &[&skill_match, &second_skill_match]);
let skill = bamboo_llm::providers::common::openai_responses::loaded_tool_to_responses_json(
&skill_schema,
);
assert_eq!(
skill["parameters"]["properties"]["skill_id"]["enum"],
serde_json::json!(["review-helper", "lint-helper"])
);
assert!(!skill["parameters"]["properties"]["skill_id"]["enum"]
.as_array()
.unwrap()
.contains(&serde_json::json!("unmatched-skill")));
assert_eq!(
skill["parameters"]["properties"]["detail"]["type"], "string",
"the rest of the real gateway schema remains complete"
);
assert!(skill["description"]
.as_str()
.unwrap()
.contains("revision=7"));
assert!(skill["description"]
.as_str()
.unwrap()
.contains("source=user"));
let skill_description = skill["description"].as_str().unwrap();
let first_skill = skill_description
.find("review-helper — Review Helper — Review a change [revision=7, source=user]")
.unwrap();
let second_skill = skill_description
.find(
"lint-helper — Lint Helper — Run focused lint checks [revision=3, source=project]",
)
.unwrap();
assert!(
first_skill < second_skill,
"discovery relevance order is kept"
);
let workflow_gateway =
bamboo_domain::ClassifiedToolSchema::new(bamboo_agent_core::tools::ToolSchema {
schema_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionSchema {
name: "workflow_run".to_string(),
description: "Run a catalog Workflow".to_string(),
parameters: serde_json::json!({
"type":"object",
"properties":{
"action":{"type":"string","enum":["start","list"]},
"workflow_id":{"type":"string"},
"revision":{"type":"integer","minimum":1}
},
"required":["action"],
"additionalProperties":false
}),
},
})
.unwrap();
let workflow_match = bamboo_domain::CapabilityMatch {
capability_ref: "workflow:review-pipeline".to_string(),
kind: bamboo_domain::CapabilityKind::Workflow,
display_name: "Review Pipeline".to_string(),
summary: "Review a repository".to_string(),
source: bamboo_domain::CapabilitySource::Workspace,
revision: Some(9),
status: bamboo_domain::CapabilityStatus::Valid,
invocation_policy: None,
invocation_target: bamboo_domain::CapabilityInvocationTarget::Workflow {
name: "workflow_run".to_string(),
workflow_id: "review-pipeline".to_string(),
source: bamboo_domain::CapabilitySource::Workspace,
revision: 9,
},
};
let second_workflow_match = bamboo_domain::CapabilityMatch {
capability_ref: "workflow:lint-pipeline".to_string(),
kind: bamboo_domain::CapabilityKind::Workflow,
display_name: "Lint Pipeline".to_string(),
summary: "Lint the selected package".to_string(),
source: bamboo_domain::CapabilitySource::Project,
revision: Some(4),
status: bamboo_domain::CapabilityStatus::Valid,
invocation_policy: None,
invocation_target: bamboo_domain::CapabilityInvocationTarget::Workflow {
name: "workflow_run".to_string(),
workflow_id: "lint-pipeline".to_string(),
source: bamboo_domain::CapabilitySource::Project,
revision: 4,
},
};
let workflow_schema = scope_discovered_gateway_schema(
&workflow_gateway,
&[&workflow_match, &second_workflow_match],
);
let workflow =
bamboo_llm::providers::common::openai_responses::loaded_tool_to_responses_json(
&workflow_schema,
);
assert_eq!(
workflow["parameters"]["properties"]["workflow_id"]["enum"],
serde_json::json!(["review-pipeline", "lint-pipeline"])
);
assert_eq!(
workflow["parameters"]["properties"]["revision"]["enum"],
serde_json::json!([9, 4])
);
assert!(!workflow["parameters"]["properties"]["workflow_id"]["enum"]
.as_array()
.unwrap()
.contains(&serde_json::json!("unmatched-workflow")));
assert_eq!(
workflow["parameters"]["properties"]["action"]["enum"],
serde_json::json!(["start", "list"])
);
assert!(workflow["description"]
.as_str()
.unwrap()
.contains("revision=9"));
assert!(workflow["description"]
.as_str()
.unwrap()
.contains("source=workspace"));
let workflow_description = workflow["description"].as_str().unwrap();
let first_workflow = workflow_description
.find(
"review-pipeline — Review Pipeline — Review a repository [revision=9, source=workspace]",
)
.unwrap();
let second_workflow = workflow_description
.find(
"lint-pipeline — Lint Pipeline — Lint the selected package [revision=4, source=project]",
)
.unwrap();
assert!(
first_workflow < second_workflow,
"workflow relevance order is kept"
);
}
#[test]
fn sticky_skill_a_then_b_keeps_each_complete_definition_and_repeats_b_as_empty() {
fn scoped_skill(
skill_id: &str,
display_name: &str,
summary: &str,
revision: u64,
source: bamboo_domain::CapabilitySource,
) -> bamboo_agent_core::tools::ToolSchema {
let gateway =
bamboo_domain::ClassifiedToolSchema::new(bamboo_agent_core::tools::ToolSchema {
schema_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionSchema {
name: "load_skill".to_string(),
description: "Load one instruction Skill".to_string(),
parameters: serde_json::json!({
"type":"object",
"properties":{
"skill_id":{"type":"string"},
"detail":{"type":"string"}
},
"required":["skill_id"],
"additionalProperties":false
}),
},
})
.unwrap();
let matched = bamboo_domain::CapabilityMatch {
capability_ref: format!("skill:{skill_id}"),
kind: bamboo_domain::CapabilityKind::Skill,
display_name: display_name.to_string(),
summary: summary.to_string(),
source,
revision: Some(revision),
status: bamboo_domain::CapabilityStatus::Valid,
invocation_policy: None,
invocation_target: bamboo_domain::CapabilityInvocationTarget::Skill {
name: "load_skill".to_string(),
skill_id: skill_id.to_string(),
source,
revision,
},
};
scope_discovered_gateway_schema(&gateway, &[&matched])
}
fn append_sticky_definition(
session: &mut Session,
call_id: &str,
schema: &bamboo_agent_core::tools::ToolSchema,
canonical_new_names: &[String],
) {
let mut assistant = Message::assistant(
"",
Some(vec![activation_call(
call_id,
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
r#"{"query":"skill"}"#,
)]),
);
assistant.never_compress = true;
session.add_message(assistant);
let payload = serde_json::to_string(&serde_json::json!({
"tools":[serde_json::to_value(schema).unwrap()]
}))
.unwrap();
session.add_message(sticky_fallback_tool_result(
call_id,
format!("<loaded_tools>{payload}</loaded_tools>"),
true,
canonical_new_names,
));
}
let skill_a = scoped_skill(
"skill-a",
"Skill A",
"Review alpha changes",
3,
bamboo_domain::CapabilitySource::Project,
);
let skill_b = scoped_skill(
"skill-b",
"Skill B",
"Review beta changes",
5,
bamboo_domain::CapabilitySource::User,
);
let mut session = Session::new("sticky-skill-delta", "chat-model");
append_sticky_definition(
&mut session,
"skill-search-a",
&skill_a,
&["load_skill".to_string()],
);
let delta_b =
sticky_fallback_definition_delta(&session, std::slice::from_ref(&skill_b)).unwrap();
assert_eq!(delta_b, vec![serde_json::to_value(&skill_b).unwrap()]);
assert_eq!(
delta_b[0]["function"]["parameters"]["properties"]["skill_id"]["enum"],
serde_json::json!(["skill-b"])
);
assert!(delta_b[0]["function"]["description"]
.as_str()
.unwrap()
.contains("skill-b — Skill B — Review beta changes [revision=5, source=user]"));
append_sticky_definition(&mut session, "skill-search-b", &skill_b, &[]);
assert!(
sticky_fallback_definition_delta(&session, std::slice::from_ref(&skill_b))
.unwrap()
.is_empty(),
"repeating the same scoped Skill B definition emits an empty delta"
);
assert_eq!(
super::prior_sticky_fallback_definitions(&session),
vec![
serde_json::to_value(&skill_a).unwrap(),
serde_json::to_value(&skill_b).unwrap()
],
"different scoped definitions sharing load_skill remain independently visible"
);
}
#[tokio::test]
async fn client_search_filters_core_before_applying_the_result_limit() {
let tools = vec![
loading_test_schema_with_description("Read", "Read repository files"),
loading_test_schema_with_description("ReadArchive", "Read archived repository files"),
];
let call = native_client_search_item_with_arguments(
"search_deferred_limit",
serde_json::json!({"query":"read","kinds":["tool"],"limit":1}),
);
let outputs = build_openai_client_tool_search_outputs(
&Session::new("deferred-before-limit", "gpt-5.6"),
&AgentLoopConfig::default(),
&tools,
&[call],
)
.await
.unwrap();
assert_eq!(outputs.len(), 1);
assert_eq!(
outputs[0].payload()["tools"],
serde_json::json!([{
"type":"function",
"name":"ReadArchive",
"description":"Read archived repository files",
"parameters":{"type":"object","properties":{}},
"strict":false,
"defer_loading":true
}]),
"the initially visible Core candidate must not consume limit=1"
);
}
#[tokio::test]
async fn legacy_browser_discovery_loads_only_this_chats_browser() {
let mut browser = loading_test_schema_with_description(
"browser",
"Operate the browser page shared with this chat",
);
browser.function.parameters = serde_json::json!({
"type":"object",
"properties":{
"action":{"type":"string","enum":["navigate","snapshot","click","fill","press","scroll","screenshot"]},
"url":{"type":"string"},
"selector":{"type":"string"},
"text":{"type":"string"},
"expected_epoch":{"type":"integer"}
},
"required":["action"],
"additionalProperties":false
});
let tools = vec![
loading_test_schema_with_description("Read", "Read files"),
loading_test_schema_with_description("Glob", "Find paths"),
browser.clone(),
];
let provider: Arc<dyn LLMProvider> = Arc::new(StubProvider);
let mut first_chat = Session::new("browser-chat-a", "chat-model");
let mut second_chat = Session::new("browser-chat-b", "chat-model");
let initial = crate::runtime::runner::round_lifecycle::request_tool_schemas_for_session(
&first_chat,
&provider,
"chat-model",
&tools,
)
.await;
assert!(initial.iter().all(|tool| tool.function.name != "browser"));
assert!(initial
.iter()
.any(|tool| tool.function.name == "discover_capabilities"));
assert_eq!(initial[1].function.name, "Glob");
let initial_bytes = serde_json::to_vec(initial.as_ref()).unwrap().len();
let full_bytes = serde_json::to_vec(&tools).unwrap().len();
assert!(
initial_bytes < full_bytes,
"initial projection should be smaller: {initial_bytes} vs {full_bytes} bytes"
);
let first_before = effective_callable_set_for_round(
&first_chat,
&tools,
bamboo_domain::CapabilityLoadingMode::LegacyFullCatalog,
);
assert!(first_before.contains_execution_name("Glob"));
assert!(!first_before.contains_execution_name("browser"));
commit_sticky_fallback_discovery_round(
stream_output_with_tool_call(activation_call(
"browser-search-a",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
r#"{"query":"browser","kinds":["tool"],"limit":1}"#,
)),
&mut first_chat,
&AgentLoopConfig::default(),
&tools,
true,
)
.await
.unwrap();
let loaded = first_chat.messages.last().unwrap();
let definitions = sticky_result_definition_values(loaded).unwrap();
assert_eq!(definitions.len(), 1);
assert_eq!(definitions[0]["function"]["name"], "browser");
assert_eq!(
definitions[0]["function"]["parameters"],
browser.function.parameters
);
assert_eq!(
validated_sticky_fallback_loaded_tool_names(&first_chat),
vec!["browser"]
);
let first_after = effective_callable_set_for_round(
&first_chat,
&tools,
bamboo_domain::CapabilityLoadingMode::LegacyFullCatalog,
);
assert!(first_after.contains_execution_name("browser"));
let resumed: Session =
serde_json::from_value(serde_json::to_value(&first_chat).unwrap()).unwrap();
assert!(effective_callable_set_for_round(
&resumed,
&tools,
bamboo_domain::CapabilityLoadingMode::LegacyFullCatalog,
)
.contains_execution_name("browser"));
let loaded_projection =
crate::runtime::runner::round_lifecycle::request_tool_schemas_for_session(
&first_chat,
&provider,
"chat-model",
&tools,
)
.await;
assert_eq!(
serde_json::to_value(loaded_projection.as_ref()).unwrap(),
serde_json::to_value(&tools).unwrap()
);
let second_before = effective_callable_set_for_round(
&second_chat,
&tools,
bamboo_domain::CapabilityLoadingMode::LegacyFullCatalog,
);
assert!(!second_before.contains_execution_name("browser"));
let second_projection =
crate::runtime::runner::round_lifecycle::request_tool_schemas_for_session(
&second_chat,
&provider,
"chat-model",
&tools,
)
.await;
assert!(second_projection
.iter()
.all(|tool| tool.function.name != "browser"));
commit_sticky_fallback_discovery_round(
stream_output_with_tool_call(activation_call(
"unrelated-search-b",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
r#"{"query":"paths","kinds":["tool"],"limit":1}"#,
)),
&mut second_chat,
&AgentLoopConfig::default(),
&tools,
true,
)
.await
.unwrap();
assert!(
sticky_result_definition_values(second_chat.messages.last().unwrap())
.unwrap()
.is_empty()
);
second_chat.add_message(Message::assistant(
"",
Some(vec![activation_call("ordinary-browser", "browser", "{}")]),
));
second_chat.add_message(Message::tool_result_with_status(
"ordinary-browser",
format!(
"<loaded_tools>{}</loaded_tools>",
serde_json::json!({"tools":[serde_json::to_value(&browser).unwrap()]})
),
true,
));
assert!(validated_sticky_fallback_loaded_tool_names(&second_chat).is_empty());
assert!(!effective_callable_set_for_round(
&second_chat,
&tools,
bamboo_domain::CapabilityLoadingMode::LegacyFullCatalog,
)
.contains_execution_name("browser"));
}
struct FailingBrowserDiscoveryProjectSource {
lookups: AtomicUsize,
}
#[async_trait::async_trait]
impl ProjectContextSource for FailingBrowserDiscoveryProjectSource {
async fn find_project(
&self,
_project_id: &ProjectId,
) -> Result<Option<ProjectDescriptor>, ProjectContextError> {
self.lookups.fetch_add(1, Ordering::SeqCst);
Err(ProjectContextError::Source(
"project catalog is unavailable".to_string(),
))
}
}
#[tokio::test]
async fn legacy_browser_discovery_ignores_failing_project_skill_resolver() {
let directory = tempfile::tempdir().unwrap();
let source = Arc::new(FailingBrowserDiscoveryProjectSource {
lookups: AtomicUsize::new(0),
});
let config = AgentLoopConfig {
skill_manager: Some(Arc::new(SkillManager::with_config(SkillStoreConfig {
skills_dir: directory.path().join("skills"),
..Default::default()
}))),
project_context_resolver: Some(Arc::new(ProjectContextResolver::new(source.clone()))),
..AgentLoopConfig::default()
};
let mut session = Session::new("browser-with-broken-skill-store", "chat-model");
session.set_project_id_meta("browser-discovery-project".to_string());
let tools = vec![loading_test_schema_with_description(
"browser",
"Operate the browser page shared with this chat",
)];
commit_sticky_fallback_discovery_round(
stream_output_with_tool_call(activation_call(
"browser-search-broken-store",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
r#"{"query":"browser","kinds":["tool"],"limit":1}"#,
)),
&mut session,
&config,
&tools,
true,
)
.await
.unwrap();
assert_eq!(
validated_sticky_fallback_loaded_tool_names(&session),
vec!["browser"]
);
assert_eq!(source.lookups.load(Ordering::SeqCst), 0);
assert!(
crate::runtime::runner::session_setup::skill_context::resolve_skill_store_for_session(
&config, &session
)
.await
.is_err()
);
assert_eq!(source.lookups.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn codex_native_tool_search_reads_parameters_without_resolving_skill_store() {
let directory = tempfile::tempdir().unwrap();
let source = Arc::new(FailingBrowserDiscoveryProjectSource {
lookups: AtomicUsize::new(0),
});
let config = AgentLoopConfig {
skill_manager: Some(Arc::new(SkillManager::with_config(SkillStoreConfig {
skills_dir: directory.path().join("skills"),
..Default::default()
}))),
project_context_resolver: Some(Arc::new(ProjectContextResolver::new(source.clone()))),
..Default::default()
};
let mut session = Session::new("native-search-with-broken-skills", "gpt-5.6");
session.set_project_id_meta("browser-discovery-project".to_string());
let mut tool = loading_test_schema_with_description("lookup", "Search records");
tool.function.parameters = serde_json::json!({
"type":"object", "properties":{
"entries":{"type":"array", "items":{"anyOf":[
{"type":"string", "description":"Meteorological forecast"},
{"type":"object", "properties":{"timezone":{"type":"string"}}}
]}}
}, "required":["entries"], "additionalProperties":false
});
let outputs = build_openai_client_tool_search_outputs(
&session,
&config,
std::slice::from_ref(&tool),
&[native_client_search_item_for(
"search_parameters",
"meteorological!",
)],
)
.await
.unwrap();
assert_eq!(
outputs[0].payload()["tools"],
serde_json::json!([
bamboo_llm::providers::common::openai_responses::loaded_tool_to_responses_json(
&tool
)
])
);
assert_eq!(source.lookups.load(Ordering::SeqCst), 0);
tool.function.parameters = serde_json::json!({
"type":"object", "properties":{"postal_code":{"type":"string", "description":"District address"}},
"required":["postal_code"]
});
let refreshed = build_openai_client_tool_search_outputs(
&session,
&config,
std::slice::from_ref(&tool),
&[native_client_search_item_for(
"refreshed_parameters",
"district",
)],
)
.await
.unwrap();
assert_eq!(
refreshed[0].payload()["tools"],
serde_json::json!([
bamboo_llm::providers::common::openai_responses::loaded_tool_to_responses_json(
&tool
)
])
);
let missing = build_openai_client_tool_search_outputs(
&session,
&config,
std::slice::from_ref(&tool),
&[native_client_search_item_for(
"old_parameters",
"meteorological",
)],
)
.await
.unwrap();
assert_eq!(missing[0].payload()["tools"], serde_json::json!([]));
let invalid_kind = native_client_search_item_with_arguments(
"search_skill",
serde_json::json!({
"query":"review", "kinds":["skill"]
}),
);
assert!(build_openai_client_tool_search_outputs(
&session,
&config,
&[tool],
&[invalid_kind]
)
.await
.unwrap_err()
.to_string()
.contains("Tools only"));
assert_eq!(source.lookups.load(Ordering::SeqCst), 0);
}
#[test]
fn codex_search_explicit_legacy_commands_follow_tools_without_score_mixing() {
use bamboo_skills::{
WorkflowCatalogEntry, WorkflowCatalogSnapshot, WorkflowKind, WorkflowSource,
WorkflowStatus,
};
let entry = WorkflowCatalogEntry {
id: "calendar-review".into(),
name: "Calendar review".into(),
description: "Review calendar events".into(),
kind: WorkflowKind::Instruction,
source: WorkflowSource::User,
revision: 1,
content_digest: "digest".into(),
version: "1".into(),
invocation_policy: serde_json::json!({"explicit":true,"automatic":true}),
argument_schema: serde_json::json!({"type":"object"}),
status: WorkflowStatus::Valid,
legacy: false,
migration_status: None,
last_error: None,
winner: true,
shadowed_candidates: Vec::new(),
};
let mut skills = WorkflowCatalogSnapshot::default();
skills.entries.push(entry);
let tool =
loading_test_schema_with_description("calendar_lookup", "Look up calendar events");
let mut gateway = loading_test_schema_with_description("load_skill", "Load instructions");
gateway.function.parameters = serde_json::json!({"type":"object","properties":{"skill_id":{"type":"string"}},"required":["skill_id"]});
let catalog = [tool.clone(), gateway]
.into_iter()
.filter_map(bamboo_domain::ClassifiedToolSchema::new)
.collect::<Vec<_>>();
let discovery = super::CompleteCapabilityDiscovery {
tools: bamboo_tools::tool_search::ToolSearchIndex::from_resolved_catalog(&catalog),
commands: Some(
crate::capability_discovery::CapabilityDiscoveryIndex::from_snapshots(
Vec::new(),
&skills,
&WorkflowCatalogSnapshot::default(),
&Default::default(),
),
),
catalog,
};
let request = bamboo_domain::DiscoverCapabilitiesRequest {
query: "calendar".into(),
kinds: Some(vec![
bamboo_domain::CapabilityKind::Tool,
bamboo_domain::CapabilityKind::Skill,
]),
limit: Some(2),
};
let schemas = discovery.discover_complete_schemas(&request).unwrap();
assert_eq!(schemas.len(), 2);
assert_eq!(
serde_json::to_value(&schemas[0]).unwrap(),
serde_json::to_value(&tool).unwrap()
);
assert_eq!(schemas[1].function.name, "load_skill");
assert_eq!(
schemas[1].function.parameters["properties"]["skill_id"]["enum"],
serde_json::json!(["calendar-review"])
);
let bounded = discovery
.discover_complete_schemas(&bamboo_domain::DiscoverCapabilitiesRequest {
limit: Some(1),
..request.clone()
})
.unwrap();
assert_eq!(
serde_json::to_value(bounded).unwrap(),
serde_json::json!([tool])
);
let mut overlapping_catalog = discovery.catalog.clone();
let mut overlapping_gateway = overlapping_catalog[1].schema().clone();
overlapping_gateway.function.description = "Calendar instructions".into();
overlapping_catalog[1] =
bamboo_domain::ClassifiedToolSchema::new(overlapping_gateway).unwrap();
let overlapping = super::CompleteCapabilityDiscovery {
tools: bamboo_tools::tool_search::ToolSearchIndex::from_resolved_catalog(
&overlapping_catalog,
),
commands: discovery.commands.clone(),
catalog: overlapping_catalog,
};
for limit in [1, 2, 3] {
let tool_request = bamboo_domain::DiscoverCapabilitiesRequest {
limit: Some(limit),
kinds: Some(vec![bamboo_domain::CapabilityKind::Tool]),
..request.clone()
};
let complete_tools = overlapping
.discover_complete_schemas(&tool_request)
.unwrap();
let expected_names = if limit == 1 {
vec!["calendar_lookup"]
} else {
vec!["calendar_lookup", "load_skill"]
};
assert_eq!(
complete_tools
.iter()
.map(|schema| schema.function.name.as_str())
.collect::<Vec<_>>(),
expected_names,
"the BM25 winner/order must stay explicit at limit={limit}"
);
for schema in &complete_tools {
let original = overlapping
.catalog
.iter()
.find(|entry| entry.execution_name() == schema.function.name)
.unwrap();
assert_eq!(
serde_json::to_value(schema).unwrap(),
serde_json::to_value(original.schema()).unwrap(),
"Tool-only {name} must keep its complete original definition",
name = schema.function.name
);
}
let mixed = overlapping
.discover_complete_schemas(&bamboo_domain::DiscoverCapabilitiesRequest {
limit: Some(limit),
..request.clone()
})
.unwrap();
assert!(mixed.len() <= limit);
assert_eq!(
mixed
.iter()
.map(|schema| schema.function.name.as_str())
.collect::<Vec<_>>(),
expected_names,
"command replacement must preserve Tools positions at limit={limit}"
);
assert_eq!(
serde_json::to_value(&mixed[0]).unwrap(),
serde_json::to_value(&complete_tools[0]).unwrap(),
"limit=1 admits calendar_lookup and cannot add the unrelated gateway"
);
if limit >= 2 {
let gateway = &mixed[1];
assert_eq!(
gateway.function.parameters["properties"]["skill_id"]["enum"],
serde_json::json!(["calendar-review"])
);
assert!(gateway.function.description.contains("revision=1"));
assert!(gateway.function.description.contains("source=user"));
assert_ne!(serde_json::to_value(gateway).unwrap(), serde_json::to_value(&complete_tools[1]).unwrap(),
"explicit command must narrow an existing generic gateway even at the full limit=2");
}
}
}
struct SupervisorLoadingProvider {
mode: CapabilityLoadingMode,
calls: AtomicUsize,
failure: Option<&'static str>,
}
#[async_trait::async_trait]
impl LLMProvider for SupervisorLoadingProvider {
async fn capability_loading_mode(
&self,
_: &str,
required: Option<&str>,
) -> CapabilityLoadingMode {
assert!(
required.is_none(),
"delegation must not become a forced named-tool call"
);
self.mode
}
async fn provider_visible_tool_footprint(
&self,
ir: &bamboo_llm::PromptIR,
tools: &[bamboo_agent_core::tools::ToolSchema],
_: &str,
_: Option<&str>,
) -> Result<bamboo_llm::ProviderVisibleToolFootprint, LLMError> {
if self.mode == CapabilityLoadingMode::Progressive {
bamboo_llm::providers::openai::OpenAIProvider::new("test")
.with_responses_only_models(vec!["gpt-5*".into()])
.with_tool_search_execution(bamboo_llm::providers::common::openai_responses::ResponsesToolSearchExecution::Client)
.provider_visible_tool_footprint(ir, tools, "gpt-5.6", None)
.await
} else {
let projected =
bamboo_llm::providers::common::openai_compat::tools_to_openai_compat_json(
tools,
);
Ok(bamboo_llm::ProviderVisibleToolFootprint {
segments: vec![bamboo_llm::ProviderVisibleToolSegment {
kind: bamboo_llm::ProviderVisibleToolSegmentKind::InitialFullDefinition,
serialized: serde_json::to_string(&projected)?,
}],
})
}
}
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
panic!("the controlled provider must receive the real PromptIR");
}
async fn chat_stream_ir(
&self,
ir: &bamboo_llm::PromptIR,
tools: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
options: Option<&bamboo_llm::LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
assert!(options
.and_then(|options| options.required_tool.as_ref())
.is_none());
assert!(tools.iter().all(|schema| !matches!(
schema.function.name.as_str(),
"exec_command" | "load_skill" | "vendor::SubAgent"
)));
let messages = ir.flatten();
let count = self.calls.fetch_add(1, Ordering::SeqCst);
let marker = "Current delegation route observation: ";
let latest = messages
.iter()
.rev()
.find_map(|message| {
message
.content
.rfind(marker)
.map(|index| &message.content[index + marker.len()..])
})
.expect("Supervisor request must report current delegation state");
let observation = serde_json::Deserializer::from_str(latest)
.into_iter::<serde_json::Value>()
.next()
.unwrap()
.unwrap();
if count == 0 {
assert_eq!(
observation["status"], "discovery_required",
"deferred definitions are not callable yet"
);
} else {
assert_eq!(
observation["status"], "routes_exposed",
"validated discovery must be reflected in the next request"
);
assert_eq!(observation["routes"], serde_json::json!(["SubAgent"]));
}
if self.mode == CapabilityLoadingMode::StickyFallback {
assert!(tools.iter().any(|schema| schema.function.name
== bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME));
assert!(!tools
.iter()
.any(|schema| schema.function.name == "SubAgent"));
} else {
let response_options = ir.responses_request_options(None);
let body = bamboo_llm::providers::common::openai_responses::build_responses_body_with_capability_loading(
"gpt-5.6", &messages, tools, None, None, Some(&response_options), None, None,
CapabilityLoadingMode::Progressive,
bamboo_llm::providers::common::openai_responses::ResponsesToolSearchExecution::Client);
let initial = body["tools"].as_array().unwrap();
assert!(initial
.iter()
.any(|tool| tool["type"] == "tool_search" && tool["execution"] == "client"));
assert!(initial
.iter()
.all(|tool| tool["name"] != "SubAgent" && tool["name"] != "exec_command"));
}
let chunks = match count {
0 if self.mode == CapabilityLoadingMode::StickyFallback => vec![
Ok(LLMChunk::ToolCalls(vec![activation_call(
"supervisor-search",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
r#"{"query":"SubAgent","kinds":["tool"],"limit":1}"#,
)])),
Ok(LLMChunk::Done),
],
0 => vec![
Ok(LLMChunk::ProviderTranscriptItem(
native_client_search_item_for("supervisor-search", "SubAgent"),
)),
Ok(LLMChunk::Done),
],
1 => {
if self.mode == CapabilityLoadingMode::StickyFallback {
let result = messages
.iter()
.find(|message| {
message.tool_call_id.as_deref() == Some("supervisor-search")
})
.unwrap();
let definitions = sticky_result_definition_values(result).unwrap();
assert_eq!(definitions.len(), 1);
assert_eq!(definitions[0]["function"]["name"], "SubAgent");
} else {
let result = ir
.provider_transcript_groups
.iter()
.flat_map(|group| group.items())
.find(|item| item.payload()["type"] == "tool_search_output")
.unwrap();
assert_eq!(result.payload()["tools"].as_array().unwrap().len(), 1);
assert_eq!(result.payload()["tools"][0]["name"], "SubAgent");
}
vec![
Ok(LLMChunk::ToolCalls(vec![activation_call(
"supervisor-delegate",
"SubAgent",
r#"{"message":"Implement only the assigned change and return verification evidence."}"#,
)])),
Ok(LLMChunk::Done),
]
}
2 => {
let result = messages
.iter()
.find(|message| {
message.tool_call_id.as_deref() == Some("supervisor-delegate")
})
.unwrap();
assert_eq!(result.tool_success, Some(self.failure.is_none()));
if let Some(reason) = self.failure {
assert!(result.content.contains(reason));
}
let text = if self.failure.is_some() {
"Delegation blocked; owner action is required before a retry."
} else {
"Delegation admitted; the worker result is still pending."
};
vec![Ok(LLMChunk::Token(text.into())), Ok(LLMChunk::Done)]
}
_ => panic!("unexpected Supervisor request {count}"),
};
Ok(Box::pin(stream::iter(chunks)))
}
}
struct SupervisorLoadingExecutor {
calls: AtomicUsize,
failure: Option<&'static str>,
}
#[async_trait::async_trait]
impl bamboo_agent_core::tools::ToolExecutor for SupervisorLoadingExecutor {
async fn execute(
&self,
call: &ToolCall,
) -> bamboo_agent_core::tools::executor::Result<bamboo_agent_core::tools::ToolResult>
{
assert_eq!(
call.function.name, "SubAgent",
"only delegated work may enter the executor"
);
self.calls.fetch_add(1, Ordering::SeqCst);
match self.failure {
Some(reason) => Err(bamboo_agent_core::tools::ToolError::Execution(
reason.into(),
)),
None => Ok(bamboo_agent_core::tools::ToolResult::text(
true,
"child assignment admitted",
)),
}
}
fn list_tools(&self) -> Vec<bamboo_agent_core::tools::ToolSchema> {
[
"Read",
"SubAgent",
"exec_command",
"load_skill",
"vendor::SubAgent",
]
.into_iter()
.map(|name| {
let mut schema = loading_test_schema_with_description(name, name);
if name == "SubAgent" {
schema.function.parameters = serde_json::json!({
"type": "object",
"properties": {
"message": {"type": "string"},
"intent": {"type": "string"},
"role": {"type": "string"},
"target": {"type": "string"},
"reply_to": {"type": "string"}
},
"required": ["message"],
"additionalProperties": false
});
}
schema
})
.collect()
}
}
#[tokio::test]
async fn supervisor_delegation_first_real_runner_loads_only_scoped_delegation_and_reports_worker_blockers(
) {
use bamboo_domain::{SessionAuthorityIdentity, DEFAULT_SUPERVISOR_SESSION_ID};
for ultra in [false, true] {
for mode in [
CapabilityLoadingMode::StickyFallback,
CapabilityLoadingMode::Progressive,
] {
for failure in [
None,
Some("worker_unavailable: configure an authorized worker"),
Some("scope_denied: target is not attached"),
] {
let mut session = Session::new(DEFAULT_SUPERVISOR_SESSION_ID, "gpt-5.6");
session.authority_identity = SessionAuthorityIdentity::Supervisor {
incarnation_id: uuid::Uuid::new_v4(),
};
session.set_root_orchestration_only(ultra).unwrap();
session.add_message(Message::user(
"Delegate the scoped implementation and report its admission or blocker.",
));
let provider = Arc::new(SupervisorLoadingProvider {
mode,
calls: AtomicUsize::new(0),
failure,
});
let executor = Arc::new(SupervisorLoadingExecutor {
calls: AtomicUsize::new(0),
failure,
});
let config = AgentLoopConfig {
provider_name: Some("supervisor-loading-test".into()),
provider_type: Some("openai".into()),
model_name: Some("gpt-5.6".into()),
..delegation_loop_config()
};
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let mut state = e2e_loop_state(&session.id);
state.model_name = "gpt-5.6".into();
assert!(super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state
)
.await
.unwrap());
assert_eq!(provider.calls.load(Ordering::SeqCst), 3);
assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
let mut restored: Session =
serde_json::from_value(serde_json::to_value(&session).unwrap()).unwrap();
let current = super::resolve_tool_schemas_for_round(
&config,
executor.as_ref(),
&mut restored,
);
let callable = effective_callable_set_for_round(&restored, ¤t, mode);
assert!(callable.contains_execution_name("SubAgent"));
assert!(!callable.contains_execution_name("exec_command"));
assert!(!callable.contains_execution_name("load_skill"));
let removed = effective_callable_set_for_round(&restored, &[], mode);
assert!(
!removed.contains_execution_name("SubAgent"),
"a saved loading receipt cannot restore a removed tool"
);
let observation = crate::runtime::runner::session_setup::prompt_envelope::build_root_orchestration_context_block_for_tools(
&restored, &[], mode,
).unwrap();
assert!(observation.content.contains("delegation_tools_unavailable"));
assert!(!observation.content.contains("routes_exposed"));
}
}
}
}
#[tokio::test]
async fn supervisor_delegation_first_repeated_worker_failures_keep_existing_bounded_circuit() {
use bamboo_domain::{SessionAuthorityIdentity, DEFAULT_SUPERVISOR_SESSION_ID};
for ultra in [false, true] {
let mut session = Session::new(DEFAULT_SUPERVISOR_SESSION_ID, "model");
session.authority_identity = SessionAuthorityIdentity::Supervisor {
incarnation_id: uuid::Uuid::new_v4(),
};
session.set_root_orchestration_only(ultra).unwrap();
session.add_message(Message::user("Delegate scoped work"));
let provider = Arc::new(DelegationLoopProvider {
calls: AtomicUsize::new(0),
finish_after: None,
calls_per_round: 1,
});
let executor = Arc::new(SupervisorLoadingExecutor {
calls: AtomicUsize::new(0),
failure: Some("worker_unavailable: owner must configure an authorized worker"),
});
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let mut state = e2e_loop_state(&session.id);
let error = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&tokio_util::sync::CancellationToken::new(),
&delegation_loop_config(),
&mut state,
)
.await
.unwrap_err();
assert!(
matches!(error,AgentError::Tool(ref reason) if reason.contains("3 consecutive failures"))
);
assert_eq!(executor.calls.load(Ordering::SeqCst), 3);
assert_eq!(provider.calls.load(Ordering::SeqCst), 3);
assert!(session
.messages
.iter()
.any(|message| message.tool_success == Some(false)
&& message.content.contains("worker_unavailable")));
}
}
struct CodexSearchLoopProvider {
calls: AtomicUsize,
}
#[async_trait::async_trait]
impl LLMProvider for CodexSearchLoopProvider {
async fn capability_loading_mode(
&self,
_: &str,
_: Option<&str>,
) -> bamboo_domain::CapabilityLoadingMode {
bamboo_domain::CapabilityLoadingMode::StickyFallback
}
async fn chat_stream(
&self,
messages: &[Message],
tools: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
assert!(tools.iter().all(
|tool| tool.function.name != "lookup" && tool.function.name != "hidden_lookup"
));
assert!(tools
.iter()
.any(|tool| tool.function.name
== bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME));
let count = self.calls.fetch_add(1, Ordering::SeqCst);
let chunks = match count {
0 => vec![
Ok(LLMChunk::ToolCalls(vec![activation_call(
"codex-search",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
r#"{"query":"timezone!","limit":1}"#,
)])),
Ok(LLMChunk::Done),
],
1 => {
let result = messages
.iter()
.find(|message| message.tool_call_id.as_deref() == Some("codex-search"))
.unwrap();
let definitions = sticky_result_definition_values(result).unwrap();
assert_eq!(definitions.len(), 1);
assert_eq!(definitions[0]["function"]["name"], "lookup");
assert_eq!(
definitions[0]["function"]["parameters"]["required"],
serde_json::json!(["timezone"])
);
vec![
Ok(LLMChunk::ToolCalls(vec![activation_call(
"codex-lookup",
"lookup",
r#"{"timezone":"UTC"}"#,
)])),
Ok(LLMChunk::Done),
]
}
2 => {
assert!(messages
.iter()
.any(
|message| message.tool_call_id.as_deref() == Some("codex-lookup")
&& message.tool_success == Some(true)
));
vec![Ok(LLMChunk::Token("done".to_string())), Ok(LLMChunk::Done)]
}
_ => panic!("unexpected request {count}"),
};
Ok(Box::pin(stream::iter(chunks)))
}
}
struct CodexSearchLoopExecutor {
calls: AtomicUsize,
}
#[async_trait::async_trait]
impl bamboo_agent_core::tools::ToolExecutor for CodexSearchLoopExecutor {
async fn execute(
&self,
call: &ToolCall,
) -> bamboo_agent_core::tools::executor::Result<bamboo_agent_core::tools::ToolResult>
{
assert_eq!(call.function.name, "lookup");
assert_eq!(call.function.arguments, r#"{"timezone":"UTC"}"#);
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(bamboo_agent_core::tools::ToolResult::text(
true,
"verified UTC record",
))
}
fn list_tools(&self) -> Vec<bamboo_agent_core::tools::ToolSchema> {
["lookup", "hidden_lookup"].into_iter().map(|name| {
let mut schema = loading_test_schema_with_description(name, "Search records");
schema.function.parameters = serde_json::json!({
"type":"object", "properties":{"timezone":{"type":"string", "description":"Regional clock offset"}},
"required":["timezone"], "additionalProperties":false
});
schema
}).collect()
}
}
#[tokio::test]
async fn codex_search_real_runner_hides_then_loads_schema_and_accepts_next_call() {
let mut session = Session::new("codex-search-runner", "model");
session.add_message(Message::user("Read the UTC record"));
let provider = Arc::new(CodexSearchLoopProvider {
calls: AtomicUsize::new(0),
});
let executor = Arc::new(CodexSearchLoopExecutor {
calls: AtomicUsize::new(0),
});
let config = AgentLoopConfig {
disabled_tools: std::collections::BTreeSet::from(["hidden_lookup".to_string()]),
..delegation_loop_config()
};
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let mut state = e2e_loop_state(&session.id);
assert!(super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state
)
.await
.unwrap());
assert_eq!(provider.calls.load(Ordering::SeqCst), 3);
assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
let mut resumed: Session =
serde_json::from_value(serde_json::to_value(&session).unwrap()).unwrap();
let current =
super::resolve_tool_schemas_for_round(&config, executor.as_ref(), &mut resumed);
let callable = effective_callable_set_for_round(
&resumed,
¤t,
bamboo_domain::CapabilityLoadingMode::StickyFallback,
);
assert!(callable.contains_execution_name("lookup"));
assert!(!callable.contains_execution_name("hidden_lookup"));
let removed = effective_callable_set_for_round(
&resumed,
&[],
bamboo_domain::CapabilityLoadingMode::StickyFallback,
);
assert!(!removed.contains_execution_name("lookup"));
}
struct CodexNativeSearchLoopProvider {
calls: AtomicUsize,
}
#[async_trait::async_trait]
impl LLMProvider for CodexNativeSearchLoopProvider {
async fn capability_loading_mode(
&self,
_: &str,
_: Option<&str>,
) -> bamboo_domain::CapabilityLoadingMode {
bamboo_domain::CapabilityLoadingMode::Progressive
}
async fn provider_visible_tool_footprint(
&self,
ir: &bamboo_llm::PromptIR,
tools: &[bamboo_agent_core::tools::ToolSchema],
_: &str,
_: Option<&str>,
) -> Result<bamboo_llm::ProviderVisibleToolFootprint, LLMError> {
bamboo_llm::providers::openai::OpenAIProvider::new("test")
.with_responses_only_models(vec!["gpt-5*".into()])
.with_tool_search_execution(bamboo_llm::providers::common::openai_responses::ResponsesToolSearchExecution::Client)
.provider_visible_tool_footprint(ir, tools, "gpt-5.6", None)
.await
}
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
panic!("native test must consume the actual PromptIR")
}
async fn chat_stream_ir(
&self,
ir: &bamboo_llm::PromptIR,
tools: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
_: Option<&bamboo_llm::LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
let response_options = ir.responses_request_options(None);
let body = bamboo_llm::providers::common::openai_responses::build_responses_body_with_capability_loading(
"gpt-5.6", &ir.flatten(), tools, None, None, Some(&response_options), None, None,
bamboo_domain::CapabilityLoadingMode::Progressive,
bamboo_llm::providers::common::openai_responses::ResponsesToolSearchExecution::Client);
let initial = body["tools"].as_array().unwrap();
assert!(initial
.iter()
.all(|tool| tool["name"] != "lookup" && tool["name"] != "hidden_lookup"));
assert!(initial
.iter()
.any(|tool| tool["type"] == "tool_search" && tool["execution"] == "client"));
let count = self.calls.fetch_add(1, Ordering::SeqCst);
let chunks = match count {
0 => vec![
Ok(LLMChunk::ProviderTranscriptItem(
native_client_search_item_for("codex-native-search", "timezone!"),
)),
Ok(LLMChunk::Done),
],
1 => {
let output = ir
.provider_transcript_groups
.iter()
.flat_map(|group| group.items())
.find(|item| item.payload()["type"] == "tool_search_output")
.unwrap();
assert_eq!(output.payload()["tools"].as_array().unwrap().len(), 1);
assert_eq!(output.payload()["tools"][0]["name"], "lookup");
assert_eq!(
output.payload()["tools"][0]["parameters"]["required"],
serde_json::json!(["timezone"])
);
let call = activation_call("codex-lookup", "lookup", r#"{"timezone":"UTC"}"#);
vec![Ok(LLMChunk::ToolCalls(vec![call])), Ok(LLMChunk::Done)]
}
2 => {
assert!(body["input"]
.as_array()
.unwrap()
.iter()
.any(|item| item["type"] == "function_call_output"
&& item["call_id"] == "codex-lookup"
&& item["output"].to_string().contains("verified UTC record")));
vec![Ok(LLMChunk::Token("done".to_string())), Ok(LLMChunk::Done)]
}
_ => panic!("unexpected native request {count}"),
};
Ok(Box::pin(stream::iter(chunks)))
}
}
#[tokio::test]
async fn codex_native_search_real_runner_reveals_complete_schema_on_next_request() {
let mut session = Session::new("codex-native-search-runner", "gpt-5.6");
session.add_message(Message::user("Read the UTC record"));
let provider = Arc::new(CodexNativeSearchLoopProvider {
calls: AtomicUsize::new(0),
});
let executor = Arc::new(CodexSearchLoopExecutor {
calls: AtomicUsize::new(0),
});
let config = AgentLoopConfig {
disabled_tools: std::collections::BTreeSet::from(["hidden_lookup".to_string()]),
provider_name: Some("codex-search-test".to_string()),
provider_type: Some("openai".to_string()),
model_name: Some("gpt-5.6".to_string()),
..delegation_loop_config()
};
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let mut state = e2e_loop_state(&session.id);
state.model_name = "gpt-5.6".to_string();
assert!(super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state
)
.await
.unwrap());
assert_eq!(provider.calls.load(Ordering::SeqCst), 3);
assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
}
#[test]
fn codex_search_keeps_completed_legacy_command_transcripts_replayable() {
let mut session = Session::new("old-command-search", "gpt-5.6");
session
.activate_provider_transcript_route(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
&"a".repeat(64),
)
.unwrap();
let assistant = Message::assistant("", None);
let anchor = assistant.id.clone();
session.add_message(assistant);
let call = native_client_search_item_with_arguments(
"old-skill-search",
serde_json::json!({
"query":"review", "kinds":["skill"]
}),
);
session
.append_provider_transcript_group(&anchor, None, vec![call])
.unwrap();
let gateway = loading_test_schema("load_skill");
let output = bamboo_domain::ProviderTranscriptItem::try_from_payload(
bamboo_domain::ProviderFamily::OpenAi, bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
bamboo_domain::ProviderTranscriptOrigin::HostToolSearch, bamboo_domain::ProviderTranscriptAuthor::ToolResult,
serde_json::json!({"type":"tool_search_output","execution":"client","call_id":"old-skill-search",
"status":"completed","tools":[bamboo_llm::providers::common::openai_responses::loaded_tool_to_responses_json(&gateway)]})
).unwrap();
session
.append_provider_transcript_group(&anchor, None, vec![output])
.unwrap();
let resumed: Session =
serde_json::from_value(serde_json::to_value(session).unwrap()).unwrap();
assert!(effective_callable_set_for_round(
&resumed,
&[gateway],
bamboo_domain::CapabilityLoadingMode::Progressive
)
.contains_execution_name("load_skill"));
}
#[tokio::test]
async fn mid_round_root_tightening_rejects_both_discovery_paths() {
let storage: Arc<dyn Storage> = Arc::new(TestStorage::default());
let mut running = Session::new("mid-round-discovery", "model");
let mut durable = running.clone();
storage.save_session(&durable).await.unwrap();
let config = AgentLoopConfig {
storage: Some(storage.clone()),
..Default::default()
};
let tools = vec![loading_test_schema("Read"), loading_test_schema("Bash")];
durable.set_root_orchestration_only(true).unwrap();
storage.save_session(&durable).await.unwrap();
let sticky = stream_output_with_tool_call(activation_call(
"stale-discovery",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
r#"{"query":"bash"}"#,
));
let error =
commit_sticky_fallback_discovery_round(sticky, &mut running, &config, &tools, false)
.await
.expect_err("StickyFallback discovery must stop before transcript mutation");
assert!(error
.to_string()
.contains("outside this Root's tool authority"));
assert!(running.messages.is_empty());
let mut native = stream_output_with_tool_call(activation_call("unused", "Read", "{}"));
native.tool_calls.clear();
native.provider_transcript_items = vec![native_client_search_item()];
let error = commit_openai_client_tool_search_round(native, &mut running, &config, &tools)
.await
.expect_err("OpenAI client search must stop before returning definitions");
assert!(error
.to_string()
.contains("outside this Root's tool authority"));
assert!(running.messages.is_empty());
}
#[tokio::test]
async fn sticky_discovery_persists_canonical_delta_and_resumes_callable_membership() {
let mut deferred =
loading_test_schema_with_description("ReadArchive", "Read archived repository files");
deferred.function.parameters = serde_json::json!({
"type":"object",
"properties":{"path":{"type":"string"}},
"required":["path"],
"additionalProperties":false
});
let tools = vec![
loading_test_schema_with_description("Read", "Read repository files"),
deferred,
loading_test_schema_with_description("Glob", "Match repository paths"),
];
let config = AgentLoopConfig::default();
let mut session = Session::new("sticky-discovery", "chat-model");
let discovery_arguments = r#"{"query":"read","kinds":["tool"],"limit":1}"#;
commit_sticky_fallback_discovery_round(
stream_output_with_tool_call(activation_call(
"sticky-search-1",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
discovery_arguments,
)),
&mut session,
&config,
&tools,
false,
)
.await
.unwrap();
assert_eq!(session.messages.len(), 2);
let assistant = &session.messages[0];
let result = &session.messages[1];
assert!(matches!(assistant.role, bamboo_agent_core::Role::Assistant));
assert!(matches!(result.role, bamboo_agent_core::Role::Tool));
assert_eq!(result.tool_call_id.as_deref(), Some("sticky-search-1"));
assert_eq!(result.tool_success, Some(true));
assert!(assistant.never_compress && result.never_compress);
let definitions = sticky_result_definition_values(result).unwrap();
assert_eq!(definitions.len(), 1);
let definition = &definitions[0];
assert_eq!(definition["type"], "function");
assert_eq!(definition["function"]["name"], "ReadArchive");
assert_eq!(
definition["function"]["parameters"],
tools[1].function.parameters
);
assert!(
definition.get("defer_loading").is_none(),
"fallback history uses the provider-neutral ToolSchema/Chat shape"
);
assert_eq!(
result.metadata.as_ref().unwrap()["canonical_new_names"],
serde_json::json!(["ReadArchive"])
);
assert_eq!(
validated_sticky_fallback_loaded_tool_names(&session),
vec!["ReadArchive"]
);
let mut resumed: Session =
serde_json::from_value(serde_json::to_value(&session).unwrap()).unwrap();
assert!(resumed.messages[0].never_compress && resumed.messages[1].never_compress);
let resumed_effective = effective_callable_set_for_round(
&resumed,
&tools,
bamboo_domain::CapabilityLoadingMode::StickyFallback,
);
assert_eq!(
resumed_effective.resolve_callable_reference("ReadArchive"),
Some("ReadArchive".to_string())
);
assert_eq!(resumed_effective.resolve_callable_reference("Glob"), None);
assert_eq!(
resumed_effective.resolve_callable_reference("invented_tool"),
None
);
let forged_definition = serde_json::to_value(&tools[2]).unwrap();
resumed.add_message(Message::assistant(
"",
Some(vec![activation_call("ordinary-glob", "Glob", "{}")]),
));
resumed.add_message(Message::tool_result_with_status(
"ordinary-glob",
format!(
"<loaded_tools>{}</loaded_tools>",
serde_json::json!({"tools":[forged_definition]})
),
true,
));
assert_eq!(
validated_sticky_fallback_loaded_tool_names(&resumed),
vec!["ReadArchive"],
"an ordinary function call/result cannot manufacture loaded state"
);
commit_sticky_fallback_discovery_round(
stream_output_with_tool_call(activation_call(
"sticky-search-2",
bamboo_domain::DISCOVERY_CONTROL_FALLBACK_TOOL_NAME,
discovery_arguments,
)),
&mut resumed,
&config,
&tools,
false,
)
.await
.unwrap();
let repeated = resumed.messages.last().unwrap();
assert_eq!(repeated.tool_call_id.as_deref(), Some("sticky-search-2"));
assert_eq!(repeated.tool_success, Some(true));
assert_eq!(
sticky_result_definition_values(repeated).unwrap(),
Vec::<serde_json::Value>::new(),
"repeated discovery closes the call with an empty definition delta"
);
assert!(resumed.messages[4].never_compress && repeated.never_compress);
assert_eq!(
validated_sticky_fallback_loaded_tool_names(&resumed),
vec!["ReadArchive"]
);
}
#[test]
fn anthropic_history_drives_progressive_round_membership_while_legacy_stays_full() {
const BOUNDARY: &str = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb";
let tools = vec![
loading_test_schema("exec_command"),
loading_test_schema("get_weather"),
loading_test_schema("Glob"),
];
let mut first_round = Session::new("anthropic-first-round", "model");
first_round
.activate_provider_transcript_route(
bamboo_domain::ProviderFamily::Anthropic,
bamboo_domain::ProviderProtocol::AnthropicMessages2023_06_01,
BOUNDARY,
)
.unwrap();
let first_effective = effective_callable_set_for_round(
&first_round,
&tools,
bamboo_domain::CapabilityLoadingMode::Progressive,
);
assert!(first_effective.contains_execution_name("exec_command"));
assert!(!first_effective.contains_execution_name("get_weather"));
let assistant = Message::assistant("normalized", None);
let anchor = assistant.id.clone();
first_round.add_message(assistant);
let item = |author, payload| {
bamboo_domain::ProviderTranscriptItem::try_from_payload(
bamboo_domain::ProviderFamily::Anthropic,
bamboo_domain::ProviderProtocol::AnthropicMessages2023_06_01,
bamboo_domain::ProviderTranscriptOrigin::Provider,
author,
payload,
)
.unwrap()
};
first_round
.append_provider_transcript_group(
&anchor,
None,
vec![
item(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"server_tool_use","id":"srv_1",
"name":"tool_search_tool_regex","input":{"pattern":"weather"}
}),
),
item(
bamboo_domain::ProviderTranscriptAuthor::ToolResult,
serde_json::json!({
"type":"tool_search_tool_result","tool_use_id":"srv_1",
"content":{"type":"tool_search_tool_search_result","tool_references":[
{"type":"tool_reference","tool_name":"get_weather"}
]}
}),
),
item(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"tool_use","id":"tool_1","name":"get_weather","input":{}
}),
),
item(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"tool_use","id":"ordinary_tool_use",
"name":"Glob","input":{}
}),
),
],
)
.unwrap();
let loaded_effective = effective_callable_set_for_round(
&first_round,
&tools,
bamboo_domain::CapabilityLoadingMode::Progressive,
);
assert!(loaded_effective.contains_execution_name("exec_command"));
assert!(loaded_effective.contains_execution_name("get_weather"));
assert!(!loaded_effective.contains_execution_name("Glob"));
let legacy = Session::new("legacy-round", "model");
let legacy_effective = effective_callable_set_for_round(
&legacy,
&tools,
bamboo_domain::CapabilityLoadingMode::LegacyFullCatalog,
);
assert!(legacy_effective.contains_execution_name("exec_command"));
assert!(legacy_effective.contains_execution_name("get_weather"));
assert!(legacy_effective.contains_execution_name("Glob"));
}
#[test]
fn openai_search_output_drives_progressive_membership_across_resume() {
const BOUNDARY: &str = "cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc";
let tools = vec![
loading_test_schema("exec_command"),
loading_test_schema("search_orders"),
loading_test_schema("Glob"),
];
let mut session = Session::new("openai-loaded-round", "gpt-5.6");
session
.activate_provider_transcript_route(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
BOUNDARY,
)
.unwrap();
let first = effective_callable_set_for_round(
&session,
&tools,
bamboo_domain::CapabilityLoadingMode::Progressive,
);
assert!(first.contains_execution_name("exec_command"));
assert!(!first.contains_execution_name("search_orders"));
let assistant = Message::assistant("", None);
let anchor = assistant.id.clone();
session.add_message(assistant);
session
.append_provider_transcript_group(&anchor, None, vec![native_client_search_item()])
.unwrap();
let output = bamboo_domain::ProviderTranscriptItem::try_from_payload(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
bamboo_domain::ProviderTranscriptOrigin::HostToolSearch,
bamboo_domain::ProviderTranscriptAuthor::ToolResult,
serde_json::json!({
"type":"tool_search_output","execution":"client","call_id":"search_1",
"status":"completed","tools":[{
"type":"function","name":"search_orders","description":"Search orders",
"parameters":{"type":"object"},"strict":false,"defer_loading":true
}]
}),
)
.unwrap();
session
.append_provider_transcript_group(&anchor, None, vec![output])
.unwrap();
let loaded = effective_callable_set_for_round(
&session,
&tools,
bamboo_domain::CapabilityLoadingMode::Progressive,
);
assert!(loaded.contains_execution_name("exec_command"));
assert!(loaded.contains_execution_name("search_orders"));
assert!(!loaded.contains_execution_name("Glob"));
let resumed: Session =
serde_json::from_value(serde_json::to_value(&session).unwrap()).unwrap();
let resumed_loaded = effective_callable_set_for_round(
&resumed,
&tools,
bamboo_domain::CapabilityLoadingMode::Progressive,
);
assert!(resumed_loaded.contains_execution_name("search_orders"));
assert!(!resumed_loaded.contains_execution_name("Glob"));
}
#[test]
fn hosted_search_output_enables_its_same_response_function_but_not_an_ordinary_call() {
const BOUNDARY: &str = "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee";
let tools = vec![
loading_test_schema("exec_command"),
loading_test_schema("search_orders"),
loading_test_schema("Glob"),
];
let mut session = Session::new("openai-hosted-search", "gpt-5.6");
session
.activate_provider_transcript_route(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
BOUNDARY,
)
.unwrap();
let assistant = Message::assistant("", None);
let anchor = assistant.id.clone();
session.add_message(assistant);
let item = |author, payload| {
bamboo_domain::ProviderTranscriptItem::try_from_payload(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
bamboo_domain::ProviderTranscriptOrigin::Provider,
author,
payload,
)
.unwrap()
};
session
.append_provider_transcript_group(
&anchor,
None,
vec![
item(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"tool_search_call","id":"tsc_hosted","execution":"server",
"call_id":"search_hosted","status":"completed",
"arguments":{"query":"orders"}
}),
),
item(
bamboo_domain::ProviderTranscriptAuthor::ToolResult,
serde_json::json!({
"type":"tool_search_output","id":"tso_hosted","execution":"server",
"call_id":"search_hosted","status":"completed",
"tools":[{"type":"function","name":"search_orders"}]
}),
),
item(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"function_call","id":"fc_loaded","call_id":"call_loaded",
"name":"search_orders","arguments":"{}","status":"completed"
}),
),
item(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"function_call","id":"fc_ordinary","call_id":"call_ordinary",
"name":"Glob","arguments":"{}","status":"completed"
}),
),
],
)
.unwrap();
let effective = effective_callable_set_for_round(
&session,
&tools,
bamboo_domain::CapabilityLoadingMode::Progressive,
);
assert!(effective.contains_execution_name("exec_command"));
assert!(effective.contains_execution_name("search_orders"));
assert!(
!effective.contains_execution_name("Glob"),
"an ordinary function_call cannot manufacture loaded state"
);
}
#[tokio::test]
async fn client_search_builds_host_output_and_commits_an_internal_next_round_boundary() {
const BOUNDARY: &str = "dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd";
let tools = vec![
loading_test_schema("Read"),
loading_test_schema("search_orders"),
loading_test_schema("Glob"),
];
let config = AgentLoopConfig::default();
let mut session = Session::new("client-search-next-round", "gpt-5.6");
session
.activate_provider_transcript_route(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
BOUNDARY,
)
.unwrap();
let call = native_client_search_item();
let outputs = build_openai_client_tool_search_outputs(
&session,
&config,
&tools,
std::slice::from_ref(&call),
)
.await
.unwrap();
assert_eq!(outputs.len(), 1);
assert_eq!(outputs[0].payload()["type"], "tool_search_output");
assert_eq!(outputs[0].payload()["execution"], "client");
assert_eq!(outputs[0].payload()["call_id"], "search_1");
assert_eq!(outputs[0].payload()["status"], "completed");
let discovered = outputs[0].payload()["tools"].as_array().unwrap();
assert_eq!(discovered.len(), 1);
assert_eq!(discovered[0]["name"], "search_orders");
assert_eq!(discovered[0]["strict"], false);
assert_eq!(discovered[0]["defer_loading"], true);
let empty_call = native_client_search_item_for("search_empty", "zzqxvplmn");
let empty_outputs = build_openai_client_tool_search_outputs(
&session,
&config,
&tools,
std::slice::from_ref(&empty_call),
)
.await
.unwrap();
assert_eq!(empty_outputs.len(), 1);
assert_eq!(empty_outputs[0].payload()["call_id"], "search_empty");
assert_eq!(
empty_outputs[0].payload()["tools"],
serde_json::json!([]),
"an empty discovery result remains an explicit completed tools array"
);
let stream_output = crate::runtime::stream::handler::StreamHandlingOutput {
response_id: Some("resp_client_search".to_string()),
visible_message: None,
content: String::new(),
reasoning_content: String::new(),
reasoning_signature: None,
token_count: 0,
tool_calls: Vec::new(),
output_tokens: 0,
thinking_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
provider_usage: None,
input_tokens: 0,
provider_transcript_items: vec![call],
};
commit_openai_client_tool_search_round(stream_output, &mut session, &config, &tools)
.await
.unwrap();
assert_eq!(session.messages.len(), 1);
let groups = session.provider_transcript.replayable_groups(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
BOUNDARY,
);
assert_eq!(groups.len(), 2, "provider call then host output");
assert_eq!(groups[0].anchor_message_id(), groups[1].anchor_message_id());
assert_eq!(groups[0].items()[0].payload()["type"], "tool_search_call");
assert_eq!(groups[1].items()[0].payload()["type"], "tool_search_output");
let effective = effective_callable_set_for_round(
&session,
&tools,
bamboo_domain::CapabilityLoadingMode::Progressive,
);
assert!(effective.contains_execution_name("Read"));
assert!(effective.contains_execution_name("search_orders"));
assert!(!effective.contains_execution_name("Glob"));
}
#[tokio::test]
async fn explicit_activation_uses_one_legacy_request_slice_for_wire_and_admission() {
let tools = vec![
loading_test_schema("load_skill"),
loading_test_schema("Read"),
];
let mut session = Session::new("anthropic-explicit-activation", "claude-sonnet-4-6");
session.metadata.insert(
bamboo_skills::runtime_metadata::SKILL_RUNTIME_SELECTION_SOURCE_KEY.to_string(),
"explicit".to_string(),
);
session.metadata.insert(
bamboo_skills::runtime_metadata::SKILL_RUNTIME_SELECTED_SKILL_IDS_KEY.to_string(),
"[\"review\"]".to_string(),
);
let required_tool =
crate::runtime::runner::round_lifecycle::required_tool_for_session(&session);
assert_eq!(required_tool, Some("load_skill"));
let request_tools =
crate::runtime::runner::round_lifecycle::effective_tool_schemas(&session, &tools);
assert_eq!(
request_tools
.iter()
.map(|tool| tool.function.name.as_str())
.collect::<Vec<_>>(),
vec!["load_skill"]
);
let provider = bamboo_llm::providers::anthropic::AnthropicProvider::new("test-key");
let mode = provider
.capability_loading_mode("claude-sonnet-4-6", required_tool)
.await;
assert_eq!(
mode,
bamboo_domain::CapabilityLoadingMode::LegacyFullCatalog
);
let effective = effective_callable_set_for_round(&session, request_tools.as_ref(), mode);
assert!(effective.contains_execution_name("load_skill"));
assert!(!effective.contains_execution_name("Read"));
}
#[test]
fn assistant_and_native_group_commit_or_rollback_as_one_unit() {
let mut session = Session::new("native-commit", "model");
let message = Message::assistant("normalized", None);
let anchor = message.id.clone();
let mut items = Some(vec![native_client_search_item()]);
commit_assistant_message(&mut session, message, &mut items).unwrap();
assert_eq!(session.messages.len(), 1);
let boundary = session
.provider_transcript
.active_provider_boundary_sha256()
.unwrap();
let groups = session.provider_transcript.replayable_groups(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
boundary,
);
assert_eq!(groups.len(), 1);
assert_eq!(groups[0].anchor_message_id(), anchor);
let mut rejected = Session::new("native-rejected", "model");
let rejected_boundary = bamboo_domain::provider_transcript_boundary_sha256(
Some("anthropic-rejected"),
Some("anthropic"),
)
.unwrap();
rejected
.activate_provider_transcript_route(
bamboo_domain::ProviderFamily::Anthropic,
bamboo_domain::ProviderProtocol::AnthropicMessages2023_06_01,
&rejected_boundary,
)
.unwrap();
let before = rejected.clone();
let mut items = Some(vec![native_client_search_item()]);
let error = commit_assistant_message(
&mut rejected,
Message::assistant("must roll back", None),
&mut items,
)
.unwrap_err();
assert!(error
.to_string()
.contains("provider-native transcript group rejected"));
assert!(rejected.messages.is_empty());
assert!(before.messages.is_empty());
assert_eq!(rejected.provider_transcript, before.provider_transcript);
assert_eq!(rejected.updated_at, before.updated_at);
}
#[test]
fn explicit_activation_first_step_gate_rejects_missing_wrong_and_multiple_calls() {
let session = pending_explicit_session();
assert!(validate_explicit_activation_first_step(&session, &[]).is_err());
assert!(validate_explicit_activation_first_step(
&session,
&[activation_call("wrong", "Read", r#"{"file_path":"x"}"#)],
)
.is_err());
assert!(validate_explicit_activation_first_step(
&session,
&[
activation_call("load", "load_skill", r#"{"skill_id":"review"}"#),
activation_call("other", "Read", r#"{"file_path":"x"}"#),
],
)
.is_err());
}
#[test]
fn explicit_activation_first_step_gate_accepts_only_matching_load_skill() {
let session = pending_explicit_session();
assert!(validate_explicit_activation_first_step(
&session,
&[activation_call(
"wrong-skill",
"load_skill",
r#"{"skill_id":"plan"}"#,
)],
)
.is_err());
let attempt = validate_explicit_activation_first_step(
&session,
&[activation_call(
"load-review",
"load_skill",
r#"{"skill_id":"review"}"#,
)],
)
.expect("matching load_skill should pass")
.expect("pending activation attempt");
assert_eq!(attempt.call_id(), "load-review");
assert_eq!(attempt.skill_id(), "review");
}
#[test]
fn explicit_activation_clears_pending_only_after_successful_tool_result() {
let call = activation_call("load-review", "load_skill", r#"{"skill_id":"review"}"#);
let mut failed = pending_explicit_session();
let attempt = validate_explicit_activation_first_step(&failed, std::slice::from_ref(&call))
.expect("valid first step")
.expect("activation attempt");
failed.add_message(Message::tool_result_with_status(
"load-review",
"durable save failed",
false,
));
assert!(apply_successful_explicit_activation(&mut failed, &attempt).is_err());
assert!(
crate::runtime::runner::session_setup::skill_context::explicit_activation_pending(
&failed
)
);
let mut succeeded = pending_explicit_session();
let attempt = validate_explicit_activation_first_step(&succeeded, &[call])
.expect("valid first step")
.expect("activation attempt");
succeeded.add_message(Message::tool_result_with_status(
"load-review",
"loaded",
true,
));
succeeded.metadata.insert(
bamboo_skills::runtime_metadata::LOADED_SKILL_IDS_METADATA_KEY.to_string(),
"[\"review\"]".to_string(),
);
succeeded.metadata.insert(
bamboo_skills::ACTIVE_WORKFLOW_METADATA_KEY.to_string(),
serde_json::json!({
"id": "review",
"source": "builtin",
"revision": 1,
"kind": "instruction",
"args": {},
"invoked_by": "user",
"activated_at": "2026-07-21T00:00:00Z",
"status": "active"
})
.to_string(),
);
succeeded.metadata.insert(
bamboo_skills::ACTIVE_WORKFLOW_SNAPSHOT_METADATA_KEY.to_string(),
"{}".to_string(),
);
apply_successful_explicit_activation(&mut succeeded, &attempt)
.expect("successful tool result activates workflow");
assert!(
!crate::runtime::runner::session_setup::skill_context::explicit_activation_pending(
&succeeded
)
);
let mut degraded = pending_explicit_session();
let call = activation_call("load-review", "load_skill", r#"{"skill_id":"review"}"#);
let attempt = validate_explicit_activation_first_step(°raded, &[call])
.expect("valid first step")
.expect("activation attempt");
degraded.add_message(Message::tool_result_with_status(
"load-review",
r#"{"activation_status":"degraded"}"#,
true,
));
degraded.metadata.insert(
bamboo_skills::runtime_metadata::SKILL_RUNTIME_ACTIVATION_ERROR_KEY.to_string(),
r#"{"code":"provider_failed"}"#.to_string(),
);
apply_successful_explicit_activation(&mut degraded, &attempt)
.expect("typed degraded activation lets the main session continue fail-closed");
assert!(!degraded
.metadata
.contains_key(bamboo_skills::ACTIVE_WORKFLOW_METADATA_KEY));
assert!(
!crate::runtime::runner::session_setup::skill_context::explicit_activation_pending(
°raded
)
);
}
struct MockGuardianSpawner {
child_id: String,
}
#[async_trait::async_trait]
impl GuardianSpawner for MockGuardianSpawner {
async fn spawn_guardian_review(
&self,
_parent_session: &Session,
_review_prompt: String,
_model: String,
_disabled_tools: Option<std::collections::BTreeSet<String>>,
) -> Result<String, String> {
Ok(self.child_id.clone())
}
}
fn guardian_enabled_config(max_reviews: u32) -> AgentLoopConfig {
let spawner: Arc<dyn GuardianSpawner> = Arc::new(MockGuardianSpawner {
child_id: "guardian-child".to_string(),
});
AgentLoopConfig {
guardian_config: Some(GuardianConfig {
enabled: true,
model_name: Some("guardian-test-model".to_string()),
max_reviews,
}),
guardian_spawner: Some(spawner),
..Default::default()
}
}
#[tokio::test]
async fn guardian_gate_spawns_and_suspends_on_first_terminal() {
let mut session = Session::new("s1", "model");
let config = guardian_enabled_config(2);
let mut runtime_state = AgentRuntimeState::new("s1".to_string());
let outcome =
maybe_spawn_guardian_review(&mut session, &config, &None, &mut runtime_state, 1, None)
.await
.expect("guardian should engage a review and suspend");
assert!(outcome.should_break && !outcome.sent_complete);
assert!(runtime_state.waiting_for_children.is_some());
let guardian_state = read_guardian_state(&session).expect("guardian state persisted");
assert_eq!(guardian_state.phase, GuardianPhase::Pending);
assert_eq!(
guardian_state.guardian_child_id.as_deref(),
Some("guardian-child")
);
assert_eq!(guardian_state.review_count, 1);
}
#[tokio::test]
async fn guardian_gate_inert_without_config() {
let mut session = Session::new("s1", "model");
let config = AgentLoopConfig::default(); let mut runtime_state = AgentRuntimeState::new("s1".to_string());
assert!(maybe_spawn_guardian_review(
&mut session,
&config,
&None,
&mut runtime_state,
1,
None
)
.await
.is_none());
assert!(runtime_state.waiting_for_children.is_none());
}
#[tokio::test]
async fn guardian_gate_skips_when_no_model_resolves() {
let spawner: Arc<dyn GuardianSpawner> = Arc::new(MockGuardianSpawner {
child_id: "guardian-child".to_string(),
});
let config = AgentLoopConfig {
guardian_config: Some(GuardianConfig {
enabled: true,
model_name: None,
max_reviews: 2,
}),
guardian_spawner: Some(spawner),
..Default::default()
};
let mut session = Session::new("s1", "model");
let mut runtime_state = AgentRuntimeState::new("s1".to_string());
assert!(maybe_spawn_guardian_review(
&mut session,
&config,
&None,
&mut runtime_state,
1,
None
)
.await
.is_none());
assert!(runtime_state.waiting_for_children.is_none());
assert!(
read_guardian_state(&session).is_none(),
"no guardian review budget should be charged when skipped"
);
}
#[tokio::test]
async fn guardian_gate_completes_after_approval() {
let mut session = Session::new("s1", "model");
let mut guardian_state = ensure_guardian_state(&session);
guardian_state.record_spawn("guardian-child");
guardian_state.record_verdict(GuardianVerdict::approved(), 1);
write_guardian_state(&mut session, guardian_state);
let config = guardian_enabled_config(2);
let mut runtime_state = AgentRuntimeState::new("s1".to_string());
assert!(maybe_spawn_guardian_review(
&mut session,
&config,
&None,
&mut runtime_state,
2,
None
)
.await
.is_none());
assert!(runtime_state.waiting_for_children.is_none());
}
#[tokio::test]
async fn guardian_gate_re_reviews_after_reject_then_completes_on_budget() {
let mut session = Session::new("s1", "model");
let mut guardian_state = ensure_guardian_state(&session);
guardian_state.record_spawn("guardian-child");
guardian_state.record_verdict(GuardianVerdict::rejected(vec!["bug".to_string()]), 1);
write_guardian_state(&mut session, guardian_state);
let config = guardian_enabled_config(2);
let mut runtime_state = AgentRuntimeState::new("s1".to_string());
let outcome =
maybe_spawn_guardian_review(&mut session, &config, &None, &mut runtime_state, 2, None)
.await
.expect("rejected within budget → re-review (suspend)");
assert!(outcome.should_break && !outcome.sent_complete);
let after = read_guardian_state(&session).expect("state persisted");
assert_eq!(after.review_count, 2, "second review spawned");
assert_eq!(after.phase, GuardianPhase::Pending);
let mut exhausted = ensure_guardian_state(&session);
exhausted.record_verdict(GuardianVerdict::rejected(vec!["still".to_string()]), 3);
write_guardian_state(&mut session, exhausted);
let mut runtime_state2 = AgentRuntimeState::new("s1".to_string());
assert!(
maybe_spawn_guardian_review(&mut session, &config, &None, &mut runtime_state2, 4, None)
.await
.is_none(),
"budget exhausted → allow completion despite unresolved findings"
);
}
struct StubProvider;
#[async_trait::async_trait]
impl LLMProvider for StubProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
}
#[derive(Default)]
struct ContextProbeProvider {
calls: AtomicUsize,
}
#[async_trait::async_trait]
impl LLMProvider for ContextProbeProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("done".to_string())),
Ok(LLMChunk::Done),
])))
}
}
struct CancelOnSecondProjectLookup {
descriptor: ProjectDescriptor,
cancel_token: tokio_util::sync::CancellationToken,
lookups: AtomicUsize,
}
#[async_trait::async_trait]
impl ProjectContextSource for CancelOnSecondProjectLookup {
async fn find_project(
&self,
project_id: &ProjectId,
) -> Result<Option<ProjectDescriptor>, ProjectContextError> {
assert_eq!(project_id, &self.descriptor.id);
if self.lookups.fetch_add(1, Ordering::SeqCst) == 1 {
self.cancel_token.cancel();
}
Ok(Some(self.descriptor.clone()))
}
}
struct ScriptedGoldProvider {
decision: &'static str,
confidence: &'static str,
}
#[async_trait::async_trait]
impl LLMProvider for ScriptedGoldProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
let arguments = format!(
r#"{{"decision":"{}","confidence":"{}","reasoning":"gate test"}}"#,
self.decision, self.confidence
);
let call = bamboo_agent_core::tools::ToolCall {
id: "gold-call-1".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "report_gold_evaluation".to_string(),
arguments,
},
};
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::Done),
])))
}
}
fn gold_continue_config() -> crate::runtime::config::AgentLoopConfig {
crate::runtime::config::AgentLoopConfig {
gold_config: Some(crate::runtime::config::GoldConfig {
enabled: true,
auto_continue_enabled: true,
goal: Some("finish the task".to_string()),
max_auto_continuations: 3,
..crate::runtime::config::GoldConfig::default()
}),
..crate::runtime::config::AgentLoopConfig::default()
}
}
fn round_usage() -> MetricsTokenUsage {
MetricsTokenUsage {
prompt_tokens: 1,
completion_tokens: 1,
total_tokens: 2,
}
}
struct BlockFirstStopHook;
#[async_trait::async_trait]
impl AgentHook for BlockFirstStopHook {
fn point(&self) -> AgentHookPoint {
AgentHookPoint::BeforeFinalize
}
async fn run(
&self,
_point: AgentHookPoint,
payload: &HookPayload,
_session: &Session,
) -> HookResult {
match payload {
HookPayload::Finalize {
stop_hook_active: false,
} => HookResult::WithContext {
result: Box::new(HookResult::Deny {
reason: "verify the result".to_string(),
}),
text: "run the focused check".to_string(),
},
HookPayload::Finalize {
stop_hook_active: true,
} => HookResult::Continue,
payload => panic!("unexpected stop payload: {payload:?}"),
}
}
}
struct AlwaysBlockStopHook;
#[async_trait::async_trait]
impl AgentHook for AlwaysBlockStopHook {
fn point(&self) -> AgentHookPoint {
AgentHookPoint::BeforeFinalize
}
async fn run(
&self,
_point: AgentHookPoint,
_payload: &HookPayload,
_session: &Session,
) -> HookResult {
HookResult::Deny {
reason: "keep going".to_string(),
}
}
}
fn stop_hook_config(hook: Arc<dyn AgentHook>) -> AgentLoopConfig {
let mut runner = crate::runtime::hooks::HookRunner::new();
runner.register(hook);
AgentLoopConfig {
hook_runner: Arc::new(runner),
..Default::default()
}
}
#[tokio::test]
async fn stop_hook_block_continues_then_allows_completion_with_active_payload() {
let mut session = Session::new("stop-hook", "model");
let mut runtime_state = AgentRuntimeState::new("stop-hook");
let (tx, mut rx) = tokio::sync::mpsc::channel(16);
let config = stop_hook_config(Arc::new(BlockFirstStopHook));
let first = super::handle_no_tool_calls(
"first final answer".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"stop-hook",
&config,
&None,
"model",
1,
Arc::new(StubProvider),
)
.await
.unwrap();
assert!(!first.should_break);
assert!(!first.sent_complete);
assert_eq!(runtime_state.stop_hook_forced_continuations, 1);
assert!(session.messages.last().is_some_and(|message| {
message.content.contains("verify the result")
&& message.content.contains("run the focused check")
}));
while let Ok(event) = rx.try_recv() {
assert!(
!matches!(event, AgentEvent::Complete { .. }),
"blocked stop must not emit Complete"
);
}
let second = super::handle_no_tool_calls(
"verified final answer".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-2",
"stop-hook",
&config,
&None,
"model",
2,
Arc::new(StubProvider),
)
.await
.unwrap();
assert!(second.should_break);
assert!(second.sent_complete);
let mut saw_complete = false;
while let Ok(event) = rx.try_recv() {
saw_complete |= matches!(event, AgentEvent::Complete { .. });
}
assert!(saw_complete, "allowed stop must emit Complete");
}
#[tokio::test]
async fn stop_hook_continuation_cap_completes_despite_further_blocks() {
let mut session = Session::new("stop-hook-cap", "model");
let mut runtime_state = AgentRuntimeState::new("stop-hook-cap");
runtime_state.stop_hook_forced_continuations = 5;
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let config = stop_hook_config(Arc::new(AlwaysBlockStopHook));
let outcome = super::handle_no_tool_calls(
"forced final answer".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-cap",
"stop-hook-cap",
&config,
&None,
"model",
6,
Arc::new(StubProvider),
)
.await
.unwrap();
assert!(outcome.should_break);
assert!(outcome.sent_complete);
assert_eq!(
session
.metadata
.get("runtime.completion_reason")
.map(String::as_str),
Some("stop_hook_continuation_cap")
);
}
#[tokio::test]
async fn no_tool_calls_does_not_complete_when_gold_continues() {
let mut session = Session::new("session-1", "model");
let mut runtime_state = AgentRuntimeState::new("session-1".to_string());
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
let outcome = super::handle_no_tool_calls(
"tentative answer".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"session-1",
&gold_continue_config(),
&None,
"model",
1,
Arc::new(ScriptedGoldProvider {
decision: "continue",
confidence: "high",
}),
)
.await
.unwrap();
assert!(!outcome.should_break);
assert!(!outcome.sent_complete);
assert_eq!(session.messages.len(), 2);
let last = session.messages.last().unwrap();
assert!(matches!(last.role, bamboo_agent_core::Role::User));
let metadata = last.metadata.as_ref().expect("runtime metadata");
assert_eq!(
metadata.get("runtime_kind").and_then(|v| v.as_str()),
Some("goal_continue")
);
drop(tx);
let mut saw_complete = false;
while let Some(event) = rx.recv().await {
if matches!(event, AgentEvent::Complete { .. }) {
saw_complete = true;
}
}
assert!(
!saw_complete,
"Complete must not be emitted on gold continue"
);
}
#[tokio::test]
async fn no_tool_calls_completes_when_gold_achieved() {
let mut session = Session::new("session-1", "model");
let mut runtime_state = AgentRuntimeState::new("session-1".to_string());
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
let outcome = super::handle_no_tool_calls(
"final answer".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"session-1",
&gold_continue_config(),
&None,
"model",
1,
Arc::new(ScriptedGoldProvider {
decision: "achieved",
confidence: "high",
}),
)
.await
.unwrap();
assert!(outcome.should_break);
assert!(outcome.sent_complete);
assert_eq!(session.messages.len(), 1);
drop(tx);
let mut saw_complete = false;
while let Some(event) = rx.recv().await {
if matches!(event, AgentEvent::Complete { .. }) {
saw_complete = true;
}
}
assert!(
saw_complete,
"Complete must be emitted when gold is achieved"
);
}
#[tokio::test]
async fn e2e_goal_loop_continue_then_declare_then_complete() {
let mut session = Session::new("session-e2e", "model");
let config = gold_continue_config();
let mut runtime_state = AgentRuntimeState::new("session-e2e".to_string());
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let r1 = super::handle_no_tool_calls(
"I think that's everything.".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"session-e2e",
&config,
&None,
"model",
1,
Arc::new(ScriptedGoldProvider {
decision: "continue",
confidence: "high",
}),
)
.await
.unwrap();
assert!(!r1.should_break, "undeclared + continue → keep working");
assert!(!r1.sent_complete);
let st = read_goal_state(&session).expect("goal state persisted after round 1");
assert_eq!(st.continuation_count, 1);
assert_eq!(st.status, GoalRuntimeStatus::Active);
assert_eq!(st.eval_history.len(), 1);
assert!(session
.messages
.last()
.unwrap()
.content
.contains("update_goal"));
let mut st = ensure_goal_state(&session, "finish the task");
st.declare(GoalDeclaredStatus::Complete, 2);
write_goal_state(&mut session, st);
let r2 = super::handle_no_tool_calls(
"Done — shipped and verified.".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-2",
"session-e2e",
&config,
&None,
"model",
2,
Arc::new(ScriptedGoldProvider {
decision: "achieved",
confidence: "high",
}),
)
.await
.unwrap();
assert!(r2.should_break, "declared complete + achieved → stop");
assert!(r2.sent_complete);
let st = read_goal_state(&session).expect("goal state persisted after round 2");
assert_eq!(st.status, GoalRuntimeStatus::Complete);
assert_eq!(st.declared_status, None, "declaration cleared after acting");
assert_eq!(st.eval_history.len(), 2, "both double-checks persisted");
drop(tx);
let mut completes = 0;
while let Some(event) = rx.recv().await {
if matches!(event, AgentEvent::Complete { .. }) {
completes += 1;
}
}
assert_eq!(
completes, 1,
"exactly one terminal Complete across the whole loop"
);
}
#[tokio::test]
async fn e2e_goal_loop_double_check_vetoes_premature_complete() {
let mut session = Session::new("session-e2e2", "model");
let config = gold_continue_config();
let mut runtime_state = AgentRuntimeState::new("session-e2e2".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let mut st = ensure_goal_state(&session, "finish the task");
st.declare(GoalDeclaredStatus::Complete, 1);
write_goal_state(&mut session, st);
let outcome = super::handle_no_tool_calls(
"All done!".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"session-e2e2",
&config,
&None,
"model",
1,
Arc::new(ScriptedGoldProvider {
decision: "continue",
confidence: "high",
}),
)
.await
.unwrap();
assert!(!outcome.should_break, "premature completion vetoed");
assert!(!outcome.sent_complete);
let st = read_goal_state(&session).expect("goal state persisted");
assert_eq!(st.status, GoalRuntimeStatus::Active);
assert_eq!(
st.declared_status, None,
"stale declaration cleared on veto"
);
assert_eq!(st.continuation_count, 1);
}
fn guardian_and_gold_config(max_reviews: u32) -> crate::runtime::config::AgentLoopConfig {
let spawner: Arc<dyn GuardianSpawner> = Arc::new(MockGuardianSpawner {
child_id: "guardian-child".to_string(),
});
crate::runtime::config::AgentLoopConfig {
gold_config: Some(crate::runtime::config::GoldConfig {
enabled: true,
auto_continue_enabled: true,
goal: Some("finish the task".to_string()),
max_auto_continuations: 3,
..crate::runtime::config::GoldConfig::default()
}),
guardian_config: Some(GuardianConfig {
enabled: true,
model_name: Some("guardian-test-model".to_string()),
max_reviews,
}),
guardian_spawner: Some(spawner),
..crate::runtime::config::AgentLoopConfig::default()
}
}
#[tokio::test]
async fn gold_continue_skips_guardian_review() {
let mut session = Session::new("s343-continue", "model");
let config = guardian_and_gold_config(2);
let mut runtime_state = AgentRuntimeState::new("s343-continue".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let outcome = super::handle_no_tool_calls(
"tentative — I think that's everything".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"s343-continue",
&config,
&None,
"model",
1,
Arc::new(ScriptedGoldProvider {
decision: "continue",
confidence: "high",
}),
)
.await
.unwrap();
assert!(!outcome.should_break);
assert!(!outcome.sent_complete);
assert!(
runtime_state.waiting_for_children.is_none(),
"a premature terminal must NOT suspend on a guardian review",
);
assert!(
read_guardian_state(&session).is_none(),
"no guardian review budget may be spent before the goal is met",
);
assert_eq!(session.messages.len(), 2);
let last = session.messages.last().unwrap();
assert_eq!(
last.metadata
.as_ref()
.and_then(|m| m.get("runtime_kind"))
.and_then(|v| v.as_str()),
Some("goal_continue"),
);
}
#[tokio::test]
async fn gold_stop_reaches_guardian_review_on_final_state() {
let mut session = Session::new("s343-stop", "model");
let config = guardian_and_gold_config(2);
let mut goal = ensure_goal_state(&session, "finish the task");
goal.declare(GoalDeclaredStatus::Complete, 1);
write_goal_state(&mut session, goal);
let mut runtime_state = AgentRuntimeState::new("s343-stop".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let outcome = super::handle_no_tool_calls(
"Done — shipped and verified.".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"s343-stop",
&config,
&None,
"model",
1,
Arc::new(ScriptedGoldProvider {
decision: "achieved",
confidence: "high",
}),
)
.await
.unwrap();
assert!(outcome.should_break);
assert!(
!outcome.sent_complete,
"Gold STOP must reach the guardian and suspend, not complete outright",
);
assert!(
runtime_state.waiting_for_children.is_some(),
"the guardian must review the final state and suspend on its verdict",
);
let guardian = read_guardian_state(&session).expect("guardian state persisted");
assert_eq!(guardian.phase, GuardianPhase::Pending);
assert_eq!(guardian.review_count, 1);
}
struct GoalLoopE2eProvider {
main_calls: std::sync::atomic::AtomicUsize,
}
#[async_trait::async_trait]
impl LLMProvider for GoalLoopE2eProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
async fn chat_stream_with_options(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
options: Option<&bamboo_llm::LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
let purpose = options
.and_then(|o| o.request_purpose.as_deref())
.unwrap_or("agent_loop");
if purpose == "gold_evaluation" {
let call = bamboo_agent_core::tools::ToolCall {
id: "gold-1".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "report_gold_evaluation".to_string(),
arguments: r#"{"decision":"achieved","confidence":"high","reasoning":"objective verified"}"#.to_string(),
},
};
return Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::Done),
])));
}
let n = self
.main_calls
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if n == 0 {
let call = bamboo_agent_core::tools::ToolCall {
id: "ug-1".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "update_goal".to_string(),
arguments: r#"{"status":"complete"}"#.to_string(),
},
};
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::Done),
])))
} else {
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("Done — shipped and verified.".to_string())),
Ok(LLMChunk::Done),
])))
}
}
}
fn e2e_loop_state(
session_id: &str,
) -> crate::runtime::runner::loop_execution::startup::LoopRunState {
use crate::runtime::runner::loop_execution::startup::{
GoldEvaluationState, LoopRunState, OverflowRecoveryState, TaskEvaluationState,
};
LoopRunState {
session_id: session_id.to_string(),
execution_id: "test-execution".to_string(),
current_inputs: None,
model_name: "model".to_string(),
metrics_collector: None,
debug_logger: crate::runtime::runner::logging::DebugLogger::new(false),
task_context: None,
overflow_recovery: OverflowRecoveryState::default(),
task_evaluation: TaskEvaluationState::default(),
gold_evaluation: GoldEvaluationState {
in_flight: None,
completed: None,
queued_request: None,
},
auxiliary_models: crate::runtime::config::AuxiliaryModelConfig::default(),
runtime_state: AgentRuntimeState::new(session_id),
}
}
#[tokio::test]
async fn e2e_full_loop_update_goal_tool_then_double_check_completes() {
use crate::runtime::config::PromptMemoryFlags;
let mut session = Session::new("session-full-e2e", "model");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let llm: Arc<dyn LLMProvider> = Arc::new(GoalLoopE2eProvider {
main_calls: std::sync::atomic::AtomicUsize::new(0),
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> =
Arc::new(bamboo_tools::BuiltinToolExecutor::new());
let config = AgentLoopConfig {
gold_config: Some(crate::runtime::config::GoldConfig {
enabled: true,
auto_continue_enabled: true,
goal: Some("ship it".to_string()),
max_auto_continuations: 3,
..crate::runtime::config::GoldConfig::default()
}),
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
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(5),
..Default::default()
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-full-e2e");
let cancel = tokio_util::sync::CancellationToken::new();
let sent_complete =
super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state)
.await
.expect("pipeline runs to completion");
assert!(sent_complete, "the run emits a terminal Complete");
let goal_state = read_goal_state(&session).expect("goal state persisted");
assert_eq!(goal_state.status, GoalRuntimeStatus::Complete);
assert_eq!(
goal_state.declared_status, None,
"declaration cleared after the terminal gate acted"
);
assert!(
!goal_state.eval_history.is_empty(),
"the double-check verdict was persisted into the goal's eval trail"
);
drop(tx);
let mut completes = 0;
while let Some(event) = rx.recv().await {
if matches!(event, AgentEvent::Complete { .. }) {
completes += 1;
}
}
assert_eq!(completes, 1, "exactly one terminal Complete");
}
struct DelegationLoopProvider {
calls: AtomicUsize,
finish_after: Option<usize>,
calls_per_round: usize,
}
#[async_trait::async_trait]
impl LLMProvider for DelegationLoopProvider {
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
let round = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if self.finish_after == Some(round) {
return Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("Delegation corrected.".into())),
Ok(LLMChunk::Done),
])));
}
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls((0..self.calls_per_round).map(|offset| bamboo_agent_core::tools::ToolCall {
id: format!("delegation-{}", round * self.calls_per_round + offset),
tool_type: "function".into(),
function: bamboo_agent_core::tools::FunctionCall {
name: "SubAgent".into(),
arguments: r#"{"intent":"chat","role":"explorer","target":"","reply_to":"","message":"Audit worktrees"}"#.into(),
},
}).collect())),
Ok(LLMChunk::Done),
])))
}
}
struct DelegationLoopExecutor {
calls: AtomicUsize,
success_on: Option<usize>,
}
#[async_trait::async_trait]
impl bamboo_agent_core::tools::ToolExecutor for DelegationLoopExecutor {
async fn execute(
&self,
_: &bamboo_agent_core::tools::ToolCall,
) -> bamboo_agent_core::tools::executor::Result<bamboo_agent_core::tools::ToolResult>
{
let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if self.success_on == Some(call) {
Ok(bamboo_agent_core::tools::ToolResult::text(
true,
"child created",
))
} else {
Err(bamboo_agent_core::tools::ToolError::InvalidArguments(
"ParentRequest reply requires chat intent and omits target and role".into(),
))
}
}
fn list_tools(&self) -> Vec<bamboo_agent_core::tools::ToolSchema> {
vec![bamboo_agent_core::tools::ToolSchema {
schema_type: "function".into(),
function: bamboo_agent_core::tools::FunctionSchema {
name: "SubAgent".into(),
description: "Delegate a task".into(),
parameters: serde_json::json!({"type":"object", "properties":{}}),
},
}]
}
}
fn delegation_loop_config() -> AgentLoopConfig {
AgentLoopConfig {
prompt_memory_flags: crate::runtime::config::PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
},
model_name: Some("model".into()),
run_budget: bamboo_config::RunBudgetConfig {
max_rounds: Some(10),
..Default::default()
},
..AgentLoopConfig::default()
}
}
#[tokio::test]
async fn repeated_subagent_failures_stop_native_run_and_a_fresh_run_resets_the_guard() {
use std::sync::atomic::Ordering;
let mut session = Session::new("delegation-failure-loop", "model");
session.add_message(Message::user("Audit worktrees"));
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let provider = Arc::new(DelegationLoopProvider {
calls: AtomicUsize::new(0),
finish_after: None,
calls_per_round: 1,
});
let executor = Arc::new(DelegationLoopExecutor {
calls: AtomicUsize::new(0),
success_on: None,
});
let config = delegation_loop_config();
let cancel = tokio_util::sync::CancellationToken::new();
for run in 1..=2 {
let mut state = e2e_loop_state(&session.id);
let error = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&cancel,
&config,
&mut state,
)
.await
.expect_err("delegation circuit must stop before the round budget");
assert!(
matches!(error, AgentError::Tool(ref message) if message.contains("3 consecutive failures") && message.contains("omit unused target and reply_to"))
);
assert_eq!(provider.calls.load(Ordering::SeqCst), 3 * run);
assert_eq!(executor.calls.load(Ordering::SeqCst), 3 * run);
for call in 0..3 * run {
assert_eq!(
session
.messages
.iter()
.filter(|message| message.tool_call_id.as_deref()
== Some(format!("delegation-{call}").as_str()))
.count(),
1,
"every failed tool call needs one durable response"
);
}
assert!(!session.metadata.contains_key("runtime.completion_reason"));
session.add_message(Message::user("Retry after correcting the call"));
}
}
#[tokio::test]
async fn corrected_subagent_call_clears_failure_streak_in_the_real_loop() {
use std::sync::atomic::Ordering;
let mut session = Session::new("delegation-correction-loop", "model");
session.add_message(Message::user("Audit worktrees"));
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let provider = Arc::new(DelegationLoopProvider {
calls: AtomicUsize::new(0),
finish_after: Some(5),
calls_per_round: 1,
});
let executor = Arc::new(DelegationLoopExecutor {
calls: AtomicUsize::new(0),
success_on: Some(2),
});
let config = delegation_loop_config();
let cancel = tokio_util::sync::CancellationToken::new();
let mut state = e2e_loop_state(&session.id);
let complete = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&cancel,
&config,
&mut state,
)
.await
.expect("a corrected call must reset the streak");
assert!(complete);
assert_eq!(provider.calls.load(Ordering::SeqCst), 6);
assert_eq!(executor.calls.load(Ordering::SeqCst), 5);
}
struct CancelAfterThirdDelegationRound(tokio_util::sync::CancellationToken);
#[async_trait::async_trait]
impl AgentHook for CancelAfterThirdDelegationRound {
fn point(&self) -> AgentHookPoint {
AgentHookPoint::AfterRound
}
async fn run(&self, _: AgentHookPoint, payload: &HookPayload, _: &Session) -> HookResult {
if matches!(payload, HookPayload::Round { round: 3 }) {
self.0.cancel();
}
HookResult::Continue
}
}
#[tokio::test]
async fn delegation_circuit_records_terminal_session_metrics_and_prioritizes_cancel() {
use bamboo_metrics::storage::MetricsStorage;
use std::sync::atomic::Ordering;
for cancelled in [false, true] {
let mut session = Session::new("delegation-terminal-metrics", "model");
session.add_message(Message::user("Audit worktrees"));
let (_dir, collector, storage) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
&session.id,
"model",
session.created_at,
session.messages.len() as u32,
);
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let provider = Arc::new(DelegationLoopProvider {
calls: AtomicUsize::new(0),
finish_after: None,
calls_per_round: 1,
});
let executor = Arc::new(DelegationLoopExecutor {
calls: AtomicUsize::new(0),
success_on: None,
});
let cancel = tokio_util::sync::CancellationToken::new();
let mut config = delegation_loop_config();
if cancelled {
let mut hooks = crate::runtime::hooks::HookRunner::new();
hooks.register(Arc::new(CancelAfterThirdDelegationRound(cancel.clone())));
config.hook_runner = Arc::new(hooks);
}
let mut state = e2e_loop_state(&session.id);
state.metrics_collector = Some(collector);
let error = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&cancel,
&config,
&mut state,
)
.await
.expect_err("terminal delegation failure");
assert_eq!(matches!(error, AgentError::Cancelled), cancelled);
assert_eq!(provider.calls.load(Ordering::SeqCst), 3);
assert_eq!(executor.calls.load(Ordering::SeqCst), 3);
let expected = if cancelled {
MetricsSessionStatus::Cancelled
} else {
MetricsSessionStatus::Error
};
let detail = tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
if let Some(detail) = storage.session_detail(&session.id).await.unwrap() {
if detail.session.status == expected
&& detail.session.completed_at.is_some()
{
break detail;
}
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
})
.await
.expect("terminal metric persisted");
assert_eq!(detail.rounds.len(), 3, "no phantom fourth round");
assert!(detail
.rounds
.iter()
.all(|round| round.status == MetricsRoundStatus::Error));
assert_eq!(
detail.session.total_token_usage.prompt_tokens,
state.runtime_state.round.total_prompt_tokens
);
assert_eq!(
detail.session.total_token_usage.completion_tokens,
state.runtime_state.round.total_completion_tokens
);
}
}
#[tokio::test]
async fn failing_subagent_batch_executes_at_most_three_calls_and_records_every_response() {
use std::sync::atomic::Ordering;
let mut session = Session::new("delegation-batch", "model");
session.add_message(Message::user("Audit worktrees"));
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let provider = Arc::new(DelegationLoopProvider {
calls: AtomicUsize::new(0),
finish_after: None,
calls_per_round: 4,
});
let executor = Arc::new(DelegationLoopExecutor {
calls: AtomicUsize::new(0),
success_on: None,
});
let mut state = e2e_loop_state(&session.id);
let error = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
executor.clone(),
&tokio_util::sync::CancellationToken::new(),
&delegation_loop_config(),
&mut state,
)
.await
.expect_err("bounded failing batch");
assert!(matches!(error, AgentError::Tool(_)));
assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
assert_eq!(executor.calls.load(Ordering::SeqCst), 3);
for call in 0..4 {
assert_eq!(
session
.messages
.iter()
.filter(|message| message.tool_call_id.as_deref()
== Some(format!("delegation-{call}").as_str()))
.count(),
1
);
}
assert!(session
.messages
.iter()
.any(
|message| message.tool_call_id.as_deref() == Some("delegation-3")
&& message.content.contains("circuit limit")
));
}
struct MaxRoundsProvider {
main_calls: std::sync::atomic::AtomicUsize,
}
#[async_trait::async_trait]
impl LLMProvider for MaxRoundsProvider {
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
async fn chat_stream_with_options(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
_: Option<&bamboo_llm::LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
let n = self
.main_calls
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let call = bamboo_agent_core::tools::ToolCall {
id: format!("tool-{n}"),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "noop".to_string(),
arguments: "{}".to_string(),
},
};
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::Done),
])))
}
}
struct AlwaysOkExecutor;
#[async_trait::async_trait]
impl bamboo_agent_core::tools::ToolExecutor for AlwaysOkExecutor {
async fn execute(
&self,
_call: &bamboo_agent_core::tools::ToolCall,
) -> std::result::Result<
bamboo_agent_core::tools::ToolResult,
bamboo_agent_core::tools::ToolError,
> {
Ok(bamboo_agent_core::tools::ToolResult {
success: true,
result: "ok".to_string(),
display_preference: None,
images: Vec::new(),
})
}
fn list_tools(&self) -> Vec<bamboo_agent_core::tools::ToolSchema> {
Vec::new()
}
}
#[tokio::test]
async fn max_rounds_exhaustion_is_distinguishable_and_runs_one_summary_turn() {
use crate::runtime::config::PromptMemoryFlags;
use std::sync::atomic::Ordering;
const MAX_ROUNDS: usize = 3;
let mut session = Session::new("session-max-rounds", "model");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider = Arc::new(MaxRoundsProvider {
main_calls: std::sync::atomic::AtomicUsize::new(0),
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = AgentLoopConfig {
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
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(MAX_ROUNDS as u32),
..Default::default()
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-max-rounds");
let cancel = tokio_util::sync::CancellationToken::new();
let sent_complete =
super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state)
.await
.expect("pipeline runs to completion");
assert_eq!(
session
.metadata
.get("runtime.completion_reason")
.map(String::as_str),
Some("max_rounds_reached"),
"exhaustion must be stamped in session metadata"
);
assert!(
session.messages.iter().any(|m| m.content.contains(
"Reached the maximum of 3 rounds; the task was stopped before completion."
)),
"a visible max_rounds notification message must be appended"
);
assert!(
!session
.messages
.windows(2)
.any(|w| w[0].role == bamboo_domain::Role::User
&& w[1].role == bamboo_domain::Role::User),
"max_rounds injection must not produce consecutive user messages"
);
let main_calls = provider.main_calls.load(Ordering::SeqCst);
assert_eq!(
main_calls,
MAX_ROUNDS + 1,
"exactly one extra summary turn after {MAX_ROUNDS} normal rounds (got {main_calls})"
);
assert!(
!sent_complete,
"worst-case summary turn (tool calls) leaves sent_complete false"
);
drop(tx);
let mut completes = 0;
while let Some(event) = rx.recv().await {
if matches!(event, AgentEvent::Complete { .. }) {
completes += 1;
}
}
assert_eq!(
completes, 0,
"no Complete emitted during this worst-case run"
);
}
struct UsageProvider {
calls: std::sync::atomic::AtomicUsize,
prompt_tokens_per_round: u64,
completion_tokens_per_round: u64,
subagent_calls: bool,
}
#[async_trait::async_trait]
impl LLMProvider for UsageProvider {
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
async fn chat_stream_with_options(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
_: Option<&bamboo_llm::LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
let n = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let (name, arguments) = if self.subagent_calls {
(
"SubAgent".to_string(),
r#"{"action":"create","prompt":"do work"}"#.to_string(),
)
} else {
("noop".to_string(), "{}".to_string())
};
let call = bamboo_agent_core::tools::ToolCall {
id: format!("tool-{n}"),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall { name, arguments },
};
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::CacheUsage {
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
input_tokens: self.prompt_tokens_per_round,
}),
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::UsageSummary {
output_tokens: self.completion_tokens_per_round,
thinking_tokens: 0,
}),
Ok(LLMChunk::Done),
])))
}
}
struct CanonicalUsageProvider {
prompt_tokens: u64,
completion_tokens: u64,
content: Option<&'static str>,
}
#[async_trait::async_trait]
impl LLMProvider for CanonicalUsageProvider {
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
let mut chunks = vec![
Ok(LLMChunk::ProviderUsage {
input_tokens: Some(self.prompt_tokens),
output_tokens: Some(self.completion_tokens),
total_tokens: Some(9_999),
reasoning_tokens: None,
cache_creation_input_tokens: None,
cache_read_input_tokens: None,
cache_write_input_tokens: None,
}),
Ok(LLMChunk::ResponseId("canonical-usage-response".to_string())),
];
if let Some(content) = self.content {
chunks.push(Ok(LLMChunk::Token(content.to_string())));
}
chunks.push(Ok(LLMChunk::Done));
Ok(Box::pin(stream::iter(chunks)))
}
}
struct IntentionalFinalizeSuspendHook;
#[async_trait::async_trait]
impl AgentHook for IntentionalFinalizeSuspendHook {
fn point(&self) -> AgentHookPoint {
AgentHookPoint::BeforeFinalize
}
async fn run(
&self,
_point: AgentHookPoint,
_payload: &HookPayload,
_session: &Session,
) -> HookResult {
HookResult::Suspend {
reason: "wait for external approval".to_string(),
}
}
}
struct AbortFinalizeHook;
#[async_trait::async_trait]
impl AgentHook for AbortFinalizeHook {
fn point(&self) -> AgentHookPoint {
AgentHookPoint::BeforeFinalize
}
async fn run(
&self,
_point: AgentHookPoint,
_payload: &HookPayload,
_session: &Session,
) -> HookResult {
HookResult::Abort {
reason: "injected terminal policy failure".to_string(),
}
}
}
struct BilledRetryProvider {
calls: std::sync::atomic::AtomicUsize,
}
#[async_trait::async_trait]
impl LLMProvider for BilledRetryProvider {
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
let attempt = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let usage = match attempt {
0 => LLMChunk::ProviderUsage {
input_tokens: Some(10),
output_tokens: Some(3),
total_tokens: Some(13),
reasoning_tokens: None,
cache_creation_input_tokens: None,
cache_read_input_tokens: None,
cache_write_input_tokens: None,
},
1 => LLMChunk::ProviderUsage {
input_tokens: Some(20),
output_tokens: Some(5),
total_tokens: Some(25),
reasoning_tokens: None,
cache_creation_input_tokens: None,
cache_read_input_tokens: None,
cache_write_input_tokens: None,
},
other => panic!("unexpected provider attempt {other}"),
};
let mut chunks = vec![Ok(usage)];
if attempt == 0 {
chunks.push(Ok(LLMChunk::ToolCalls(vec![
bamboo_agent_core::tools::ToolCall {
id: "retry-tool-attempt-1".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "noop".to_string(),
arguments: "{}".to_string(),
},
},
])));
} else {
chunks.push(Ok(LLMChunk::Token(
"second billed attempt completed".to_string(),
)));
}
chunks.push(Ok(LLMChunk::Done));
Ok(Box::pin(stream::iter(chunks)))
}
}
struct BilledThenTimeoutProvider {
first: BilledRetryProvider,
calls: std::sync::atomic::AtomicUsize,
failed: Arc<tokio::sync::Notify>,
silent_stream: bool,
}
#[async_trait::async_trait]
impl LLMProvider for BilledThenTimeoutProvider {
async fn chat_stream(
&self,
messages: &[Message],
schemas: &[bamboo_agent_core::tools::ToolSchema],
max_output_tokens: Option<u32>,
model: &str,
) -> Result<LLMStream, LLMError> {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
return self
.first
.chat_stream(messages, schemas, max_output_tokens, model)
.await;
}
self.failed.notify_one();
if self.silent_stream {
Ok(Box::pin(stream::pending()))
} else {
Err(LLMError::Api("retryable transport failure".into()))
}
}
}
struct FailBeforeUsageProvider;
#[async_trait::async_trait]
impl LLMProvider for FailBeforeUsageProvider {
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
Err(LLMError::Api(
"authentication error before a stream was created".to_string(),
))
}
}
async fn create_pipeline_metrics() -> (
tempfile::TempDir,
bamboo_metrics::MetricsCollector,
Arc<bamboo_metrics::SqliteMetricsStorage>,
) {
use bamboo_metrics::storage::MetricsStorage;
let dir = tempfile::tempdir().expect("temp metrics dir");
let storage = Arc::new(bamboo_metrics::SqliteMetricsStorage::new(
dir.path().join("metrics.db"),
));
storage.init().await.expect("init metrics storage");
let collector = bamboo_metrics::MetricsCollector::spawn(storage.clone(), 7);
(dir, collector, storage)
}
struct DelayedTaskMetricsProvider {
calls: std::sync::atomic::AtomicUsize,
windows: std::sync::Mutex<Vec<(chrono::DateTime<Utc>, chrono::DateTime<Utc>)>>,
}
#[async_trait::async_trait]
impl LLMProvider for DelayedTaskMetricsProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let started = Utc::now();
tokio::time::sleep(if call == 0 {
Duration::from_millis(60)
} else {
Duration::from_millis(150)
})
.await;
let completed = Utc::now();
self.windows
.lock()
.expect("metrics windows lock")
.push((started, completed));
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
}
#[tokio::test]
async fn zero_max_rounds_is_rejected_before_the_first_provider_call() {
use std::sync::atomic::Ordering;
let mut session = Session::new("session-zero-max-rounds", "model");
let (tx, _rx) = tokio::sync::mpsc::channel(8);
let provider = Arc::new(MaxRoundsProvider {
main_calls: std::sync::atomic::AtomicUsize::new(0),
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = AgentLoopConfig {
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_rounds: Some(0),
..Default::default()
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-zero-max-rounds");
let cancel = tokio_util::sync::CancellationToken::new();
let error =
super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state)
.await
.expect_err("a zero round cap must be rejected before execution");
assert!(matches!(error, AgentError::Budget(_)));
assert_eq!(provider.main_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn repeated_task_evaluation_rounds_use_distinct_rows_and_dispatch_durations() {
use bamboo_metrics::storage::MetricsStorage;
let session_id = "task-evaluation-metrics-repeated-activation";
let (_dir, collector, storage) = create_pipeline_metrics().await;
let base = chrono::Utc::now();
collector.session_started(session_id, "model", base);
let mut session = Session::new(session_id, "model");
session.set_task_list(bamboo_domain::TaskList {
session_id: session_id.to_string(),
title: "Metrics task".to_string(),
items: vec![bamboo_domain::TaskItem {
id: "task-1".to_string(),
description: "measure evaluator".to_string(),
status: bamboo_domain::TaskItemStatus::InProgress,
..bamboo_domain::TaskItem::default()
}],
created_at: base,
updated_at: base,
});
session.set_task_list_version_meta("1");
let task_context = crate::runtime::task_context::TaskLoopContext::from_session(&session);
let timeout_context = crate::runtime::stream::handler::StreamTimeoutContext::new(
bamboo_config::StreamTimeoutConfig::default(),
Some("provider"),
Some("fast-model"),
);
let provider = Arc::new(DelayedTaskMetricsProvider {
calls: std::sync::atomic::AtomicUsize::new(0),
windows: std::sync::Mutex::new(Vec::new()),
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let (event_tx, _event_rx) = tokio::sync::mpsc::channel(16);
let mut metric_ids = Vec::new();
for activation in 0..2 {
let request =
crate::runtime::runner::task_lifecycle::build_async_task_evaluation_request(
&task_context,
&session,
session_id,
1,
Some("fast-model"),
None,
timeout_context.clone(),
)
.expect("activation request")
.expect("task context exists");
let request_id = request.evaluation_id.clone();
let result = crate::runtime::runner::task_lifecycle::execute_async_task_evaluation(
request,
llm.clone(),
event_tx.clone(),
1,
Some(collector.clone()),
Arc::new(std::sync::atomic::AtomicBool::new(false)),
Arc::new(std::sync::atomic::AtomicBool::new(false)),
)
.await;
assert!(result.metrics_started.load(Ordering::Acquire));
metric_ids.push((request_id, result.metrics_round_id.clone()));
let mut activation_state = e2e_loop_state(session_id);
activation_state.execution_id = format!("activation-{activation}");
activation_state.task_context = task_context.clone();
activation_state.metrics_collector = Some(collector.clone());
activation_state.task_evaluation.completed = Some(result);
super::apply_completed_task_evaluation(
&mut session,
&event_tx,
&AgentLoopConfig::default(),
&mut activation_state,
)
.await;
}
assert_ne!(metric_ids[0].0, metric_ids[1].0);
assert_ne!(metric_ids[0].1, metric_ids[1].1);
let detail = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let Some(detail) = storage
.session_detail(session_id)
.await
.expect("session metrics query")
{
if detail.rounds.len() == 2
&& detail
.rounds
.iter()
.all(|round| round.completed_at.is_some())
{
break detail;
}
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("two evaluator metric rows are persisted");
let first_row = detail
.rounds
.iter()
.find(|row| row.round_id == metric_ids[0].1)
.expect("first activation row");
let second_row = detail
.rounds
.iter()
.find(|row| row.round_id == metric_ids[1].1)
.expect("second activation row");
let windows = provider.windows.lock().expect("metrics windows lock");
assert_eq!(windows.len(), 2);
for (row, (provider_started, provider_completed)) in
[(first_row, windows[0]), (second_row, windows[1])]
{
assert!(row.started_at <= provider_started);
assert!(
provider_started - row.started_at < chrono::Duration::milliseconds(50),
"metric start must be adjacent to real provider dispatch"
);
let completed_at = row.completed_at.expect("terminal timestamp");
assert!(completed_at >= provider_completed);
assert!(
completed_at - provider_completed < chrono::Duration::milliseconds(50),
"metric completion must use the evaluator's real finish"
);
}
assert!(first_row.duration_ms.is_some_and(|duration| duration >= 60));
assert!(second_row
.duration_ms
.is_some_and(|duration| duration >= 150));
assert!(second_row.duration_ms.unwrap() > first_row.duration_ms.unwrap() + 50);
}
async fn wait_for_pipeline_metrics(
storage: &bamboo_metrics::SqliteMetricsStorage,
session_id: &str,
expected_status: MetricsRoundStatus,
expected_usage: MetricsTokenUsage,
) -> bamboo_metrics::types::SessionDetail {
use bamboo_metrics::storage::MetricsStorage;
for _ in 0..100 {
if let Some(detail) = storage
.session_detail(session_id)
.await
.expect("session detail query")
{
if detail.rounds.first().is_some_and(|round| {
round.status == expected_status && round.token_usage == expected_usage
}) && detail.session.total_token_usage == expected_usage
{
return detail;
}
}
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
panic!(
"session {session_id} did not persist a {expected_status:?} round with usage {expected_usage:?}"
);
}
fn canonical_usage_pipeline_config() -> AgentLoopConfig {
use crate::runtime::config::PromptMemoryFlags;
AgentLoopConfig {
model_name: Some("model".to_string()),
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
},
..AgentLoopConfig::default()
}
}
#[tokio::test]
async fn pipeline_success_clamps_overflow_and_reconciles_runtime_round_and_session() {
let session_id = "canonical-success";
let (_dir, collector, storage) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
session_id,
"model",
chrono::Utc::now(),
1,
);
let mut session = Session::new(session_id, "model");
session.add_message(Message::user(
"provider reports values beyond the durable signed range",
));
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let llm: Arc<dyn LLMProvider> = Arc::new(CanonicalUsageProvider {
prompt_tokens: u64::MAX,
completion_tokens: u64::MAX,
content: Some("provider-backed answer"),
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let mut state = e2e_loop_state(session_id);
state.metrics_collector = Some(collector);
super::run_pipeline(
&mut session,
&tx,
llm,
tools,
&tokio_util::sync::CancellationToken::new(),
&canonical_usage_pipeline_config(),
&mut state,
)
.await
.expect("successful pipeline");
drop(tx);
drain(&mut rx).await;
let max = bamboo_domain::MAX_DURABLE_TOKEN_COUNT;
let expected = MetricsTokenUsage {
prompt_tokens: max,
completion_tokens: max,
total_tokens: max,
};
let detail = wait_for_pipeline_metrics(
storage.as_ref(),
session_id,
MetricsRoundStatus::Success,
expected,
)
.await;
assert_eq!(detail.rounds.len(), 1);
assert_eq!(detail.rounds[0].token_usage, expected);
assert_eq!(detail.session.total_token_usage, expected);
assert_eq!(state.runtime_state.round.total_prompt_tokens, max);
assert_eq!(state.runtime_state.round.total_completion_tokens, max);
}
#[tokio::test]
async fn pipeline_billed_retry_attempts_persist_exactly_once_across_all_consumers() {
let session_id = "canonical-billed-retry";
let (_dir, collector, storage) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
session_id,
"model",
chrono::Utc::now(),
1,
);
let mut session = Session::new(session_id, "model");
session.add_message(Message::user("retry after post-LLM handler failure"));
session.metadata.insert(
super::TEST_POST_LLM_RETRY_FAILURES_KEY.to_string(),
"1".to_string(),
);
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider = Arc::new(BilledRetryProvider {
calls: std::sync::atomic::AtomicUsize::new(0),
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let mut state = e2e_loop_state(session_id);
state.metrics_collector = Some(collector.clone());
let mut config = canonical_usage_pipeline_config();
config.metrics_collector = Some(collector);
super::run_pipeline(
&mut session,
&tx,
provider.clone(),
tools,
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state,
)
.await
.expect("second billed attempt completes the same round");
drop(tx);
drain(&mut rx).await;
assert_eq!(
provider.calls.load(std::sync::atomic::Ordering::SeqCst),
2,
"one completed attempt fails in post-LLM handling and exactly one retry is billed"
);
let expected = MetricsTokenUsage {
prompt_tokens: 30,
completion_tokens: 8,
total_tokens: 38,
};
let detail = wait_for_pipeline_metrics(
storage.as_ref(),
session_id,
MetricsRoundStatus::Success,
expected,
)
.await;
assert_eq!(detail.rounds.len(), 1, "both attempts belong to one round");
assert_eq!(detail.rounds[0].token_usage, expected);
assert_eq!(detail.session.total_token_usage, expected);
let exposure = storage
.prompt_memory_exposure(&detail.rounds[0].round_id)
.await
.expect("query retry exposure")
.expect("both successful bootstraps share one first-wins observation");
assert_eq!(
exposure.recall_outcome,
bamboo_metrics::types::PromptMemoryRecallOutcome::Disabled
);
assert_eq!(exposure.all_compact_exposed_count, 0);
assert_eq!(state.runtime_state.round.total_prompt_tokens, 30);
assert_eq!(state.runtime_state.round.total_completion_tokens, 8);
}
async fn assert_cancelled_backoff_preserves_billed_usage(goal_recovery: bool) {
let session_id = if goal_recovery {
"goal-backoff-billed"
} else {
"retry-backoff-billed"
};
let (directory, collector, metrics) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
session_id,
"model",
chrono::Utc::now(),
1,
);
let storage: Arc<dyn Storage> = Arc::new(
bamboo_storage::SessionStoreV2::new(directory.path().join("sessions"))
.await
.unwrap(),
);
let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
Arc::new(bamboo_storage::LockedSessionStore::new(storage.clone()));
let mut session = Session::new(session_id, "model");
session.add_message(Message::user("retain billed retries when stopped"));
session
.metadata
.insert(super::TEST_POST_LLM_RETRY_FAILURES_KEY.into(), "1".into());
storage.save_session(&session).await.unwrap();
let failed = Arc::new(tokio::sync::Notify::new());
let provider = Arc::new(BilledThenTimeoutProvider {
first: BilledRetryProvider {
calls: std::sync::atomic::AtomicUsize::new(0),
},
calls: std::sync::atomic::AtomicUsize::new(0),
failed: failed.clone(),
silent_stream: goal_recovery,
});
let cancel = tokio_util::sync::CancellationToken::new();
let trigger = cancel.clone();
let watched_storage = storage.clone();
let stopper = tokio::spawn(async move {
if goal_recovery {
tokio::time::timeout(Duration::from_secs(10), async {
loop {
if watched_storage
.load_session(session_id)
.await
.unwrap()
.is_some_and(|saved| saved.metadata.contains_key("goal.recovery"))
{
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("goal recovery reservation persisted");
} else {
failed.notified().await;
tokio::time::sleep(Duration::from_millis(20)).await;
}
trigger.cancel();
});
let mut config = canonical_usage_pipeline_config();
config.persistence = Some(persistence.clone());
config.metrics_collector = Some(collector.clone());
if goal_recovery {
config.stream_timeout.transport_idle_timeout_secs = 1;
config.gold_config = Some(crate::runtime::config::GoldConfig {
enabled: true,
auto_continue_enabled: true,
goal: Some("finish".into()),
recovery: crate::runtime::goal_recovery::GoalRecoveryPolicy {
max_attempts: 1,
max_elapsed_seconds: 120,
},
..Default::default()
});
}
let mut state = e2e_loop_state(session_id);
state.metrics_collector = Some(collector);
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let result = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
Arc::new(AlwaysOkExecutor),
&cancel,
&config,
&mut state,
)
.await;
stopper.await.unwrap();
assert!(matches!(result, Err(AgentError::Cancelled)));
assert_eq!(
provider.calls.load(Ordering::SeqCst),
if goal_recovery { 3 } else { 2 }
);
assert_eq!(state.runtime_state.round.total_prompt_tokens, 10);
assert_eq!(state.runtime_state.round.total_completion_tokens, 3);
assert_eq!(state.runtime_state.round.total_tool_calls, 1);
state_bridge::write_runtime_state(&mut session, &state.runtime_state);
persistence
.checkpoint_runtime_session(&mut session)
.await
.unwrap();
let saved = storage.load_session(session_id).await.unwrap().unwrap();
let round = saved.agent_runtime_state.unwrap().round;
assert_eq!(round.total_prompt_tokens, 10);
assert_eq!(round.total_completion_tokens, 3);
assert_eq!(round.total_tool_calls, 1);
drop(tx);
drain(&mut rx).await;
let expected = MetricsTokenUsage {
prompt_tokens: 10,
completion_tokens: 3,
total_tokens: 13,
};
let detail = wait_for_pipeline_metrics(
metrics.as_ref(),
session_id,
MetricsRoundStatus::Cancelled,
expected,
)
.await;
assert_eq!(detail.rounds.len(), 1);
}
#[tokio::test]
async fn cancelled_ordinary_backoff_preserves_billed_usage() {
assert_cancelled_backoff_preserves_billed_usage(false).await;
}
#[tokio::test]
async fn cancelled_goal_recovery_backoff_preserves_billed_usage() {
assert_cancelled_backoff_preserves_billed_usage(true).await;
}
#[tokio::test]
async fn pipeline_terminal_validation_retains_billed_usage_in_runtime_and_metrics() {
let session_id = "canonical-terminal";
let (_dir, collector, storage) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
session_id,
"model",
chrono::Utc::now(),
1,
);
let mut session = Session::new(session_id, "model");
session.add_message(Message::user("returning no answer is terminal"));
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let llm: Arc<dyn LLMProvider> = Arc::new(CanonicalUsageProvider {
prompt_tokens: 31,
completion_tokens: 7,
content: None,
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let mut state = e2e_loop_state(session_id);
state.metrics_collector = Some(collector);
let error = super::run_pipeline(
&mut session,
&tx,
llm,
tools,
&tokio_util::sync::CancellationToken::new(),
&canonical_usage_pipeline_config(),
&mut state,
)
.await
.expect_err("empty assistant response is terminal");
assert!(matches!(error, AgentError::EmptyAssistantResponse { .. }));
drop(tx);
drain(&mut rx).await;
let expected = MetricsTokenUsage {
prompt_tokens: 31,
completion_tokens: 7,
total_tokens: 38,
};
let detail = wait_for_pipeline_metrics(
storage.as_ref(),
session_id,
MetricsRoundStatus::Error,
expected,
)
.await;
assert_eq!(detail.rounds.len(), 1);
assert_eq!(detail.rounds[0].token_usage, expected);
assert_eq!(detail.session.total_token_usage, expected);
assert_eq!(state.runtime_state.round.total_prompt_tokens, 31);
assert_eq!(state.runtime_state.round.total_completion_tokens, 7);
let persisted_runtime = session
.agent_runtime_state
.as_ref()
.expect("terminal runtime state mirrored to session");
assert_eq!(persisted_runtime.round.total_prompt_tokens, 31);
assert_eq!(persisted_runtime.round.total_completion_tokens, 7);
}
#[tokio::test]
async fn pipeline_post_llm_hook_abort_retains_usage_in_terminal_metrics_and_runtime() {
let session_id = "canonical-post-llm-abort";
let (_dir, collector, storage) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
session_id,
"model",
chrono::Utc::now(),
1,
);
let mut session = Session::new(session_id, "model");
session.add_message(Message::user("abort only after the stream completes"));
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let llm: Arc<dyn LLMProvider> = Arc::new(CanonicalUsageProvider {
prompt_tokens: 23,
completion_tokens: 6,
content: Some("completed provider response"),
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let mut state = e2e_loop_state(session_id);
state.metrics_collector = Some(collector);
let mut config = canonical_usage_pipeline_config();
let mut hook_runner = crate::runtime::hooks::HookRunner::new();
hook_runner.register(Arc::new(AbortFinalizeHook));
config.hook_runner = Arc::new(hook_runner);
let error = super::run_pipeline(
&mut session,
&tx,
llm,
tools,
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state,
)
.await
.expect_err("finalize hook aborts after the completed stream");
assert!(matches!(error, AgentError::Tool(message) if message.contains("hook aborted")));
drop(tx);
drain(&mut rx).await;
let expected = MetricsTokenUsage {
prompt_tokens: 23,
completion_tokens: 6,
total_tokens: 29,
};
let detail = wait_for_pipeline_metrics(
storage.as_ref(),
session_id,
MetricsRoundStatus::Error,
expected,
)
.await;
assert_eq!(detail.rounds[0].token_usage, expected);
assert_eq!(detail.session.total_token_usage, expected);
assert_eq!(state.runtime_state.round.total_prompt_tokens, 23);
assert_eq!(state.runtime_state.round.total_completion_tokens, 6);
}
#[tokio::test]
async fn pipeline_hook_suspension_keeps_success_usage_without_error_status() {
let session_id = "canonical-hook-suspend";
let (_dir, collector, storage) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
session_id,
"model",
chrono::Utc::now(),
1,
);
let mut session = Session::new(session_id, "model");
session.add_message(Message::user("pause before finalizing"));
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let llm: Arc<dyn LLMProvider> = Arc::new(CanonicalUsageProvider {
prompt_tokens: 19,
completion_tokens: 4,
content: Some("ready, pending approval"),
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let mut state = e2e_loop_state(session_id);
state.metrics_collector = Some(collector);
let mut config = canonical_usage_pipeline_config();
let mut hook_runner = crate::runtime::hooks::HookRunner::new();
hook_runner.register(Arc::new(IntentionalFinalizeSuspendHook));
config.hook_runner = Arc::new(hook_runner);
let error = super::run_pipeline(
&mut session,
&tx,
llm,
tools,
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state,
)
.await
.expect_err("hook suspension remains a control-flow signal");
assert!(error.is_hook_suspended());
drop(tx);
drain(&mut rx).await;
let expected = MetricsTokenUsage {
prompt_tokens: 19,
completion_tokens: 4,
total_tokens: 23,
};
let detail = wait_for_pipeline_metrics(
storage.as_ref(),
session_id,
MetricsRoundStatus::Success,
expected,
)
.await;
assert_eq!(detail.rounds[0].token_usage, expected);
assert_eq!(detail.session.total_token_usage, expected);
assert_eq!(
detail.session.status,
MetricsSessionStatus::Running,
"the outer finalizer, not terminal-error metrics, owns suspension status"
);
assert_eq!(state.runtime_state.round.total_prompt_tokens, 19);
assert_eq!(state.runtime_state.round.total_completion_tokens, 4);
}
#[tokio::test]
async fn pipeline_failure_before_provider_usage_remains_zero_everywhere() {
let session_id = "canonical-pre-stream-error";
let (_dir, collector, storage) = create_pipeline_metrics().await;
crate::runtime::runner::metrics_lifecycle::record_session_started(
Some(&collector),
session_id,
"model",
chrono::Utc::now(),
1,
);
let mut session = Session::new(session_id, "model");
session.add_message(Message::user("this request fails before streaming"));
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let mut state = e2e_loop_state(session_id);
state.metrics_collector = Some(collector);
super::run_pipeline(
&mut session,
&tx,
Arc::new(FailBeforeUsageProvider),
tools,
&tokio_util::sync::CancellationToken::new(),
&canonical_usage_pipeline_config(),
&mut state,
)
.await
.expect_err("provider authentication error is terminal");
drop(tx);
drain(&mut rx).await;
let detail = wait_for_pipeline_metrics(
storage.as_ref(),
session_id,
MetricsRoundStatus::Error,
MetricsTokenUsage::default(),
)
.await;
assert_eq!(detail.rounds[0].token_usage, MetricsTokenUsage::default());
assert_eq!(
detail.session.total_token_usage,
MetricsTokenUsage::default()
);
assert_eq!(state.runtime_state.round.total_prompt_tokens, 0);
assert_eq!(state.runtime_state.round.total_completion_tokens, 0);
}
#[tokio::test]
async fn run_budget_token_limit_stops_run_gracefully() {
use crate::runtime::config::PromptMemoryFlags;
let mut session = Session::new("session-token-budget", "model");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider = Arc::new(UsageProvider {
calls: std::sync::atomic::AtomicUsize::new(0),
prompt_tokens_per_round: 10,
completion_tokens_per_round: 5,
subagent_calls: false,
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = AgentLoopConfig {
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
},
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_total_tokens: Some(20),
max_tool_calls: None,
max_subagents: None,
max_rounds: Some(50),
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-token-budget");
let cancel = tokio_util::sync::CancellationToken::new();
let sent_complete =
super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state)
.await
.expect("pipeline runs to completion");
assert_eq!(
session
.metadata
.get("runtime.completion_reason")
.map(String::as_str),
Some("budget_exceeded"),
"budget trip must be stamped in session metadata"
);
assert_eq!(
session
.metadata
.get("runtime.budget_exceeded_kind")
.map(String::as_str),
Some("max_total_tokens"),
);
assert!(
session
.messages
.iter()
.any(|m| m.content.contains("max_total_tokens")),
"a visible budget-exceeded notification message must be appended"
);
assert!(
!sent_complete,
"budget trip does not send a normal complete"
);
assert_eq!(provider.calls.load(std::sync::atomic::Ordering::SeqCst), 3);
drop(tx);
let mut budget_events = Vec::new();
while let Some(event) = rx.recv().await {
if let AgentEvent::BudgetExceeded {
kind,
limit,
actual,
..
} = event
{
budget_events.push((kind, limit, actual));
}
}
assert_eq!(
budget_events.len(),
1,
"exactly one structured BudgetExceeded event must be emitted"
);
let (kind, limit, actual) = &budget_events[0];
assert_eq!(kind, "max_total_tokens");
assert_eq!(*limit, 20);
assert_eq!(*actual, 30, "trips on round 2's cumulative total (15+15)");
}
#[tokio::test]
async fn run_budget_tool_call_limit_stops_run_gracefully() {
use crate::runtime::config::PromptMemoryFlags;
let mut session = Session::new("session-tool-call-budget", "model");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider = Arc::new(UsageProvider {
calls: std::sync::atomic::AtomicUsize::new(0),
prompt_tokens_per_round: 0,
completion_tokens_per_round: 0,
subagent_calls: false,
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = AgentLoopConfig {
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
},
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_total_tokens: None,
max_tool_calls: Some(2),
max_subagents: None,
max_rounds: Some(50),
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-tool-call-budget");
let cancel = tokio_util::sync::CancellationToken::new();
let _ = super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state)
.await
.expect("pipeline runs to completion");
assert_eq!(
session
.metadata
.get("runtime.budget_exceeded_kind")
.map(String::as_str),
Some("max_tool_calls"),
);
assert_eq!(provider.calls.load(std::sync::atomic::Ordering::SeqCst), 3);
drop(tx);
drain(&mut rx).await;
}
#[tokio::test]
async fn run_budget_subagent_limit_counts_only_create_calls() {
use crate::runtime::config::PromptMemoryFlags;
let mut session = Session::new("session-subagent-budget", "model");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider = Arc::new(UsageProvider {
calls: std::sync::atomic::AtomicUsize::new(0),
prompt_tokens_per_round: 0,
completion_tokens_per_round: 0,
subagent_calls: true,
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = AgentLoopConfig {
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
},
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_total_tokens: None,
max_tool_calls: None,
max_subagents: Some(1),
max_rounds: Some(50),
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-subagent-budget");
let cancel = tokio_util::sync::CancellationToken::new();
let _ = super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state)
.await
.expect("pipeline runs to completion");
assert_eq!(
session
.metadata
.get("runtime.budget_exceeded_kind")
.map(String::as_str),
Some("max_subagents"),
);
assert_eq!(provider.calls.load(std::sync::atomic::Ordering::SeqCst), 2);
drop(tx);
drain(&mut rx).await;
}
#[tokio::test]
async fn run_under_budget_is_unaffected() {
use crate::runtime::config::PromptMemoryFlags;
let mut session = Session::new("session-under-budget", "model");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider = Arc::new(MaxRoundsProvider {
main_calls: std::sync::atomic::AtomicUsize::new(0),
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = AgentLoopConfig {
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
},
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_total_tokens: Some(1_000_000),
max_tool_calls: Some(1_000_000),
max_subagents: Some(1_000_000),
max_rounds: Some(2),
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-under-budget");
let cancel = tokio_util::sync::CancellationToken::new();
let _ = super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state)
.await
.expect("pipeline runs to completion");
assert_eq!(
session
.metadata
.get("runtime.completion_reason")
.map(String::as_str),
Some("max_rounds_reached"),
"an under-budget run still hits its real stop reason (max_rounds here), not budget_exceeded"
);
assert!(!session
.metadata
.contains_key("runtime.budget_exceeded_kind"));
drop(tx);
let mut budget_events = 0;
while let Some(event) = rx.recv().await {
if matches!(event, AgentEvent::BudgetExceeded { .. }) {
budget_events += 1;
}
}
assert_eq!(budget_events, 0, "no budget event for an under-budget run");
}
async fn drain(rx: &mut tokio::sync::mpsc::Receiver<AgentEvent>) {
while rx.recv().await.is_some() {}
}
#[test]
fn round_activity_accumulates_across_retry_attempts_instead_of_overwriting() {
use crate::runtime::stream::handler::StreamHandlingOutput;
fn attempt(input: u64, output: u64, tool_calls: Vec<&str>) -> StreamHandlingOutput {
StreamHandlingOutput {
response_id: None,
visible_message: None,
content: "x".to_string(),
reasoning_content: String::new(),
reasoning_signature: None,
token_count: 0,
tool_calls: tool_calls
.into_iter()
.enumerate()
.map(|(i, name)| bamboo_agent_core::tools::ToolCall {
id: format!("t{i}"),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: name.to_string(),
arguments: "{}".to_string(),
},
})
.collect(),
output_tokens: output,
thinking_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
provider_usage: None,
input_tokens: input,
provider_transcript_items: Vec::new(),
}
}
let mut activity = super::RoundActivity::default();
let attempt_1 = attempt(100, 50, vec!["Bash", "SubAgent"]);
let attempt_1_usage =
crate::runtime::runner::round_lifecycle::canonical_attempt_usage(&attempt_1, 999, 999);
activity.absorb_attempt(&attempt_1, attempt_1_usage);
assert_eq!(activity.prompt_tokens, 100);
assert_eq!(activity.completion_tokens, 50);
assert_eq!(activity.tool_call_count, 2);
assert_eq!(activity.subagent_spawn_count, 1);
let attempt_2 = attempt(120, 30, vec!["Bash"]);
let attempt_2_usage =
crate::runtime::runner::round_lifecycle::canonical_attempt_usage(&attempt_2, 999, 999);
activity.absorb_attempt(&attempt_2, attempt_2_usage);
assert_eq!(
activity.prompt_tokens, 220,
"attempt 1's billed prompt tokens must not be dropped on retry"
);
assert_eq!(activity.completion_tokens, 80);
assert_eq!(
activity.token_usage(),
MetricsTokenUsage {
prompt_tokens: 220,
completion_tokens: 80,
total_tokens: 300,
},
"durable usage must contain every billed retry attempt exactly once"
);
assert_eq!(activity.tool_call_count, 3);
assert_eq!(activity.subagent_spawn_count, 1);
let saturating_attempt = attempt(u64::MAX, u64::MAX, vec![]);
let saturating_usage = crate::runtime::runner::round_lifecycle::canonical_attempt_usage(
&saturating_attempt,
1,
1,
);
activity.absorb_attempt(&saturating_attempt, saturating_usage);
assert_eq!(
activity.prompt_tokens,
bamboo_domain::MAX_DURABLE_TOKEN_COUNT
);
assert_eq!(
activity.completion_tokens,
bamboo_domain::MAX_DURABLE_TOKEN_COUNT
);
}
#[test]
fn round_activity_prefers_provider_prompt_total_without_adding_reasoning_twice() {
use crate::runtime::stream::handler::{ProviderUsageSnapshot, StreamHandlingOutput};
let mut output = StreamHandlingOutput {
response_id: None,
visible_message: None,
content: "answer".to_string(),
reasoning_content: "thought".to_string(),
reasoning_signature: None,
token_count: 6,
tool_calls: Vec::new(),
output_tokens: 56,
thinking_tokens: 78,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 768,
provider_usage: Some(ProviderUsageSnapshot {
input_tokens: Some(1000),
output_tokens: Some(120),
total_tokens: Some(1120),
reasoning_tokens: Some(20),
cache_creation_input_tokens: None,
cache_read_input_tokens: Some(768),
cache_write_input_tokens: None,
}),
input_tokens: 232,
provider_transcript_items: Vec::new(),
};
let mut activity = super::RoundActivity::default();
let usage = crate::runtime::runner::round_lifecycle::canonical_attempt_usage(&output, 7, 9);
activity.absorb_attempt(&output, usage);
assert_eq!(activity.prompt_tokens, 1000);
assert_eq!(
activity.completion_tokens, 120,
"reasoning is a subset of provider output, not additional output"
);
output
.provider_usage
.as_mut()
.expect("provider usage")
.output_tokens = Some(0);
let mut zero_activity = super::RoundActivity::default();
let zero_usage =
crate::runtime::runner::round_lifecycle::canonical_attempt_usage(&output, 7, 9);
zero_activity.absorb_attempt(&output, zero_usage);
assert_eq!(
zero_activity.completion_tokens, 0,
"explicit provider zero must beat a nonzero legacy flat value"
);
}
#[test]
fn canonical_attempt_usage_falls_back_per_component_only_when_unreported() {
use crate::runtime::stream::handler::{ProviderUsageSnapshot, StreamHandlingOutput};
let output = StreamHandlingOutput {
response_id: None,
visible_message: None,
content: "answer".to_string(),
reasoning_content: String::new(),
reasoning_signature: None,
token_count: 6,
tool_calls: Vec::new(),
output_tokens: 0,
thinking_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
provider_usage: Some(ProviderUsageSnapshot {
input_tokens: Some(41),
output_tokens: None,
total_tokens: Some(999),
reasoning_tokens: None,
cache_creation_input_tokens: None,
cache_read_input_tokens: None,
cache_write_input_tokens: None,
}),
input_tokens: 0,
provider_transcript_items: Vec::new(),
};
assert_eq!(
crate::runtime::runner::round_lifecycle::canonical_attempt_usage(&output, 3, 5),
MetricsTokenUsage {
prompt_tokens: 41,
completion_tokens: 5,
total_tokens: 46,
},
"reported prompt wins, missing completion estimates, and provider total is not trusted"
);
}
#[tokio::test]
async fn budget_exceeded_kind_metadata_is_cleared_on_the_next_run() {
use crate::runtime::config::PromptMemoryFlags;
let flags = PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
relevant_recall_rerank: false,
project_first_dream: false,
ledger_agenda: false,
};
let mut session = Session::new("session-budget-metadata-hygiene", "model");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider = Arc::new(UsageProvider {
calls: std::sync::atomic::AtomicUsize::new(0),
prompt_tokens_per_round: 0,
completion_tokens_per_round: 0,
subagent_calls: false,
});
let llm: Arc<dyn LLMProvider> = provider.clone();
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let tripping_config = AgentLoopConfig {
prompt_memory_flags: flags,
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_total_tokens: None,
max_tool_calls: Some(1),
max_subagents: None,
max_rounds: Some(50),
},
..AgentLoopConfig::default()
};
let mut state = e2e_loop_state("session-budget-metadata-hygiene");
let cancel = tokio_util::sync::CancellationToken::new();
let _ = super::run_pipeline(
&mut session,
&tx,
llm,
tools.clone(),
&cancel,
&tripping_config,
&mut state,
)
.await
.expect("run 1 completes");
drop(tx);
drain(&mut rx).await;
assert_eq!(
session
.metadata
.get("runtime.budget_exceeded_kind")
.map(String::as_str),
Some("max_tool_calls"),
"run 1 must stamp the trip detail"
);
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let provider2 = Arc::new(MaxRoundsProvider {
main_calls: std::sync::atomic::AtomicUsize::new(0),
});
let llm2: Arc<dyn LLMProvider> = provider2.clone();
let unlimited_config = AgentLoopConfig {
prompt_memory_flags: flags,
model_name: Some("model".to_string()),
run_budget: bamboo_config::RunBudgetConfig {
max_rounds: Some(2),
..Default::default()
},
..AgentLoopConfig::default()
};
let mut state2 = e2e_loop_state("session-budget-metadata-hygiene");
let _ = super::run_pipeline(
&mut session,
&tx,
llm2,
tools,
&cancel,
&unlimited_config,
&mut state2,
)
.await
.expect("run 2 completes");
drop(tx);
drain(&mut rx).await;
assert!(
!session
.metadata
.contains_key("runtime.budget_exceeded_kind"),
"stale budget trip detail must be cleared by the next run"
);
assert_eq!(
session
.metadata
.get("runtime.completion_reason")
.map(String::as_str),
Some("max_rounds_reached"),
"run 2's own stop reason replaces run 1's budget_exceeded"
);
}
#[test]
fn is_child_spawn_call_counts_plan_and_subagent_create_only() {
let call = |name: &str, arguments: &str| bamboo_agent_core::tools::ToolCall {
id: "id".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: name.to_string(),
arguments: arguments.to_string(),
},
};
assert!(
is_child_spawn_call(&call("Plan", r#"{"task":"design it"}"#)),
"every Plan call creates one planner child"
);
assert!(
is_child_spawn_call(&call("SubAgent", r#"{"action":"create","prompt":"x"}"#)),
"explicit action=create counts"
);
assert!(
is_child_spawn_call(&call("SubAgent", r#"{"prompt":"x"}"#)),
"missing action defaults to the tool's legacy create behavior"
);
assert!(
!is_child_spawn_call(&call("SubAgent", r#"{"action":"wait"}"#)),
"action=wait manages an existing child, not a spawn"
);
assert!(
!is_child_spawn_call(&call("SubAgent", r#"{"action":"list"}"#)),
"action=list is read-only, not a spawn"
);
assert!(
!is_child_spawn_call(&call("Bash", r#"{"action":"create"}"#)),
"a differently named tool is never counted, regardless of args"
);
}
#[test]
fn check_run_budget_exceeded_reports_first_tripped_kind_in_priority_order() {
use bamboo_domain::session::runtime_state::RoundRuntimeState;
let unlimited = bamboo_config::RunBudgetConfig::default();
let round = RoundRuntimeState {
total_prompt_tokens: 5,
total_completion_tokens: 5,
total_tool_calls: 3,
total_subagents_spawned: 1,
..Default::default()
};
assert!(
check_run_budget_exceeded(&round, &unlimited).is_none(),
"unlimited config never trips"
);
let all_exceeded = bamboo_config::RunBudgetConfig {
max_total_tokens: Some(5),
max_tool_calls: Some(1),
max_subagents: Some(1),
max_rounds: None,
};
let exceeded =
check_run_budget_exceeded(&round, &all_exceeded).expect("some guardrail trips");
assert_eq!(exceeded.kind, "max_total_tokens");
assert_eq!(exceeded.actual, 10);
let tool_calls_only = bamboo_config::RunBudgetConfig {
max_total_tokens: None,
max_tool_calls: Some(3),
max_subagents: None,
max_rounds: None,
};
let exceeded =
check_run_budget_exceeded(&round, &tool_calls_only).expect("tool-call guardrail trips");
assert_eq!(exceeded.kind, "max_tool_calls");
assert_eq!(exceeded.actual, 3);
}
#[derive(Default)]
struct TestStorage {
sessions: RwLock<HashMap<String, Session>>,
}
#[async_trait::async_trait]
impl Storage for TestStorage {
async fn save_session(&self, session: &Session) -> std::io::Result<()> {
self.sessions
.write()
.await
.insert(session.id.clone(), session.clone());
Ok(())
}
async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
Ok(self.sessions.read().await.get(session_id).cloned())
}
async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
Ok(self.sessions.write().await.remove(session_id).is_some())
}
}
struct TestPersistence(Arc<dyn Storage>);
#[async_trait::async_trait]
impl bamboo_domain::RuntimeSessionPersistence for TestPersistence {
async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
self.0.save_session(session).await
}
}
#[derive(Default)]
struct EvaluationCasPersistence {
full_saves: std::sync::atomic::AtomicUsize,
control_plane_saves: std::sync::atomic::AtomicUsize,
paired_task_patches: std::sync::atomic::AtomicUsize,
}
#[async_trait::async_trait]
impl bamboo_domain::RuntimeSessionPersistence for EvaluationCasPersistence {
async fn save_runtime_session(&self, _session: &mut Session) -> std::io::Result<()> {
self.full_saves.fetch_add(1, Ordering::SeqCst);
Ok(())
}
async fn save_runtime_control_plane(&self, _session: &mut Session) -> std::io::Result<()> {
self.control_plane_saves.fetch_add(1, Ordering::SeqCst);
Ok(())
}
async fn update_task_list_control_planes_if_version(
&self,
session_id: &str,
shared_session_id: &str,
expected_version: &str,
_expected_task_list: &bamboo_domain::TaskList,
_task_list: &bamboo_domain::TaskList,
version: &str,
) -> std::io::Result<bool> {
assert_eq!(session_id, shared_session_id);
assert_eq!(expected_version, "1");
assert!(version.parse::<u64>().expect("numeric version") > 1);
self.paired_task_patches.fetch_add(1, Ordering::SeqCst);
Ok(true)
}
}
#[tokio::test]
async fn pending_injected_messages_are_merged_once_and_cleared_from_storage() {
let storage: Arc<dyn Storage> = Arc::new(TestStorage::default());
let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
Arc::new(TestPersistence(storage.clone()));
let mut persisted = Session::new_child("child-merge", "parent", "model", "Child");
persisted.add_message(Message::system("system"));
persisted.add_message(Message::user("original task"));
persisted.metadata.insert(
"pending_injected_messages".to_string(),
serde_json::json!([
{
"content": "queued correction",
"created_at": chrono::Utc::now(),
}
])
.to_string(),
);
storage
.save_session(&persisted)
.await
.expect("persisted child should be saved");
let mut running = persisted.clone();
running.metadata.remove("pending_injected_messages");
state_bridge::merge_pending_injected_messages(
&mut running,
Some(&storage),
Some(&persistence),
)
.await;
assert_eq!(
running
.messages
.last()
.map(|message| message.content.as_str()),
Some("queued correction")
);
assert!(!running.metadata.contains_key("pending_injected_messages"));
let saved = storage
.load_session("child-merge")
.await
.expect("load should succeed")
.expect("session should exist");
assert!(!saved.metadata.contains_key("pending_injected_messages"));
let count_after_first_merge = running.messages.len();
state_bridge::merge_pending_injected_messages(
&mut running,
Some(&storage),
Some(&persistence),
)
.await;
assert_eq!(running.messages.len(), count_after_first_merge);
}
#[tokio::test]
async fn pipeline_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,
"Pipeline cobalt orchid rule",
"The cobalt orchid request must use the pipeline memory boundary.",
&["cobalt".to_string(), "orchid".to_string()],
Some("pipeline-typed-inbox-recall"),
"model",
false,
None,
)
.await
.unwrap();
let mut persisted = Session::new("pipeline-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 pipeline cobalt orchid 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 (event_tx, mut event_rx) = tokio::sync::mpsc::channel(64);
let mut state = e2e_loop_state(&running.id);
let cancel = tokio_util::sync::CancellationToken::new();
super::run_pipeline(
&mut running,
&event_tx,
Arc::new(ContextProbeProvider::default()),
Arc::new(AlwaysOkExecutor),
&cancel,
&config,
&mut state,
)
.await
.expect("the pipeline should complete after refreshing the boundary");
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("Pipeline cobalt orchid rule"));
assert!(
rendered.contains("The cobalt orchid request must use the pipeline 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);
assert!(
std::iter::from_fn(|| event_rx.try_recv().ok()).any(|event| {
matches!(
event,
AgentEvent::MessageAppended {
ref message_id,
ref content,
..
} if message_id == envelope.id.as_str() && content == query
)
})
);
}
#[tokio::test]
async fn pipeline_cancelled_during_context_refresh_skips_primary_provider() {
let directory = tempfile::tempdir().unwrap();
let workspace = directory.path().join("workspace");
std::fs::create_dir_all(&workspace).unwrap();
let workspace = workspace.canonicalize().unwrap();
let project_id = ProjectId::parse("pipeline-context-cancel-project").unwrap();
let cancel = tokio_util::sync::CancellationToken::new();
let source = Arc::new(CancelOnSecondProjectLookup {
descriptor: ProjectDescriptor {
id: project_id.clone(),
name: "Pipeline cancellation fixture".to_string(),
project_path: Some(workspace),
home: directory
.path()
.join("projects/pipeline-context-cancel-project"),
workspace_bindings: Vec::new(),
resources: ProjectResourceSummary {
project_id: project_id.clone(),
resource_revision: 1,
resources: Vec::new(),
},
},
cancel_token: cancel.clone(),
lookups: AtomicUsize::new(0),
});
let provider = Arc::new(ContextProbeProvider::default());
let mut session = Session::new("pipeline-context-cancel", "model");
session.add_message(Message::system("base prompt"));
session.add_message(Message::user("run the cancellation fixture"));
session.set_project_id_meta(project_id.to_string());
let config = AgentLoopConfig {
project_context_resolver: Some(Arc::new(ProjectContextResolver::new(source.clone()))),
app_data_dir: Some(directory.path().to_path_buf()),
prompt_memory_flags: crate::runtime::config::PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
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 (event_tx, _event_rx) = tokio::sync::mpsc::channel(32);
let mut state = e2e_loop_state(&session.id);
let error = super::run_pipeline(
&mut session,
&event_tx,
provider.clone(),
Arc::new(AlwaysOkExecutor),
&cancel,
&config,
&mut state,
)
.await
.expect_err("cancellation observed after context refresh must stop the round");
assert!(matches!(error, AgentError::Cancelled));
assert_eq!(source.lookups.load(Ordering::SeqCst), 2);
assert!(session
.metadata
.contains_key(crate::runtime::runner::prompt_context::PROMPT_MEMORY_OBSERVABILITY_KEY));
assert_eq!(provider.calls.load(Ordering::SeqCst), 0);
}
#[test]
fn retries_transient_llm_errors() {
assert!(should_retry_turn_error(&AgentError::LLM(
"HTTP error: timeout while connecting".to_string(),
)));
assert!(should_retry_turn_error(&AgentError::LLM(
"API error: HTTP 503: Service Unavailable".to_string(),
)));
}
#[test]
fn empty_assistant_response_is_terminal_by_variant_not_message_text() {
assert!(!should_retry_turn_error(
&AgentError::EmptyAssistantResponse {
response_id: Some("resp_740".to_string()),
}
));
assert!(should_retry_turn_error(&AgentError::LLM(
"empty assistant response".to_string(),
)));
}
#[test]
fn retries_reqwest_transport_errors() {
assert!(should_retry_turn_error(&AgentError::LLM(
"HTTP error: error sending request for url (https://api.githubcopilot.com/chat/completions)".to_string(),
)));
}
#[test]
fn retries_stream_decode_transport_errors() {
assert!(should_retry_turn_error(&AgentError::LLM(
"Stream error: Transport error: error decoding response body".to_string(),
)));
}
#[test]
fn retries_unknown_llm_errors_by_default() {
assert!(should_retry_turn_error(&AgentError::LLM(
"some completely unknown error".to_string(),
)));
}
#[test]
fn does_not_retry_non_retryable_llm_errors() {
assert!(!should_retry_turn_error(&AgentError::LLM(
"Authentication error: Invalid API key".to_string(),
)));
assert!(!should_retry_turn_error(&AgentError::LLM(
"API error: HTTP 400: invalid request".to_string(),
)));
}
#[test]
fn does_not_retry_non_llm_errors() {
assert!(!should_retry_turn_error(&AgentError::Cancelled));
assert!(!should_retry_turn_error(&AgentError::Tool(
"tool failed".to_string(),
)));
assert!(!should_retry_turn_error(&AgentError::Budget(
"budget exceeded".to_string(),
)));
}
#[test]
fn retries_only_structurally_safe_stream_timeouts() {
let retry_safe = AgentError::StreamTimeout(StreamTimeoutError::new(
StreamTimeoutPhase::Bootstrap,
Duration::from_secs(120),
Some("provider".to_string()),
Some("model".to_string()),
Duration::from_secs(120),
None,
true,
));
assert!(should_retry_turn_error(&retry_safe));
let partial_output = AgentError::StreamTimeout(StreamTimeoutError::new(
StreamTimeoutPhase::TransportIdle,
Duration::from_secs(120),
Some("provider".to_string()),
Some("model".to_string()),
Duration::from_secs(120),
Some(Duration::from_secs(120)),
true,
));
assert!(!should_retry_turn_error(&partial_output));
let auxiliary_timeout = AgentError::StreamTimeout(StreamTimeoutError::new(
StreamTimeoutPhase::Bootstrap,
Duration::from_secs(120),
Some("provider".to_string()),
Some("model".to_string()),
Duration::from_secs(120),
None,
false,
));
assert!(!should_retry_turn_error(&auxiliary_timeout));
}
#[test]
fn does_not_retry_empty_llm_error() {
assert!(!should_retry_turn_error(&AgentError::LLM("".to_string())));
assert!(!should_retry_turn_error(&AgentError::LLM(
" ".to_string()
)));
}
#[test]
fn overflow_errors_use_dedicated_recovery_path() {
assert!(is_overflow_recoverable(&AgentError::LLMOverflow(
"prompt too long".to_string(),
)));
assert!(!is_overflow_recoverable(&AgentError::LLM(
"timeout while connecting".to_string(),
)));
assert!(!should_retry_turn_error(&AgentError::LLMOverflow(
"maximum context length exceeded".to_string(),
)));
}
#[test]
fn overflow_recovery_state_opens_circuit_breaker_after_threshold() {
let mut state = OverflowRecoveryState::default();
assert!(state.can_attempt_recovery());
state.record_recovery(0);
state.record_recovery(1);
state.record_recovery(2);
assert!(!state.can_attempt_recovery());
}
#[test]
fn test_map_turn_error_status_cancelled() {
let error = AgentError::Cancelled;
let (round_status, session_status) = map_turn_error_status(&error);
assert_eq!(round_status, MetricsRoundStatus::Cancelled);
assert_eq!(session_status, MetricsSessionStatus::Cancelled);
}
#[test]
fn test_map_turn_error_status_tool_error() {
let error = AgentError::Tool("Tool failed".to_string());
let (round_status, session_status) = map_turn_error_status(&error);
assert_eq!(round_status, MetricsRoundStatus::Error);
assert_eq!(session_status, MetricsSessionStatus::Error);
}
#[test]
fn test_map_turn_error_status_llm_error() {
let error = AgentError::LLM("LLM provider error".to_string());
let (round_status, session_status) = map_turn_error_status(&error);
assert_eq!(round_status, MetricsRoundStatus::Error);
assert_eq!(session_status, MetricsSessionStatus::Error);
}
#[test]
fn test_map_turn_error_status_session_not_found() {
let error = AgentError::SessionNotFound("session-123".to_string());
let (round_status, session_status) = map_turn_error_status(&error);
assert_eq!(round_status, MetricsRoundStatus::Error);
assert_eq!(session_status, MetricsSessionStatus::Error);
}
#[test]
fn test_map_turn_error_status_budget_error() {
let error = AgentError::Budget("Budget exceeded".to_string());
let (round_status, session_status) = map_turn_error_status(&error);
assert_eq!(round_status, MetricsRoundStatus::Error);
assert_eq!(session_status, MetricsSessionStatus::Error);
}
#[test]
fn test_map_turn_error_status_cancelled_is_distinct() {
let cancelled_error = AgentError::Cancelled;
let other_error = AgentError::Tool("Tool error".to_string());
let (cancelled_round, cancelled_session) = map_turn_error_status(&cancelled_error);
let (other_round, other_session) = map_turn_error_status(&other_error);
assert_ne!(cancelled_round, other_round);
assert_ne!(cancelled_session, other_session);
}
#[test]
fn test_map_turn_error_only_cancelled_gets_cancelled_status() {
let errors = vec![
AgentError::LLM("error".to_string()),
AgentError::EmptyAssistantResponse { response_id: None },
AgentError::Tool("error".to_string()),
AgentError::SessionNotFound("id".to_string()),
AgentError::Budget("error".to_string()),
];
for error in errors {
let (round_status, session_status) = map_turn_error_status(&error);
assert_eq!(round_status, MetricsRoundStatus::Error);
assert_eq!(session_status, MetricsSessionStatus::Error);
}
let (round_status, session_status) = map_turn_error_status(&AgentError::Cancelled);
assert_eq!(round_status, MetricsRoundStatus::Cancelled);
assert_eq!(session_status, MetricsSessionStatus::Cancelled);
}
#[tokio::test]
async fn handle_no_tool_calls_emits_complete_and_appends_assistant_message() {
let mut session = Session::new("session-1", "model");
let mut runtime_state = AgentRuntimeState::new("session-1".to_string());
let (tx, mut rx) = tokio::sync::mpsc::channel(4);
let outcome = super::handle_no_tool_calls(
"final answer".to_string(),
Some("reasoning trace".to_string()),
None,
11,
7,
MetricsTokenUsage {
prompt_tokens: 11,
completion_tokens: 7,
total_tokens: 18,
},
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"session-1",
&crate::runtime::config::AgentLoopConfig::default(),
&None,
"model",
1,
Arc::new(StubProvider),
)
.await
.unwrap();
assert!(outcome.should_break);
assert!(outcome.sent_complete);
assert_eq!(session.messages.len(), 1);
assert!(matches!(
session.messages[0].role,
bamboo_agent_core::Role::Assistant
));
assert_eq!(session.messages[0].content, "final answer");
assert_eq!(
session.messages[0].reasoning.as_deref(),
Some("reasoning trace")
);
let event = rx.recv().await.expect("complete event should be sent");
match event {
AgentEvent::Complete { usage } => {
assert_eq!(usage.prompt_tokens, 11);
assert_eq!(usage.completion_tokens, 7);
assert_eq!(usage.total_tokens, 18);
}
other => panic!("unexpected event: {other:?}"),
}
}
#[tokio::test]
async fn apply_completed_task_evaluation_updates_task_list_and_emits_event() {
let persistence = Arc::new(EvaluationCasPersistence::default());
let mut session = Session::new("session-task-eval", "model");
session.set_task_list(bamboo_domain::TaskList {
session_id: "session-task-eval".to_string(),
title: "Eval Tasks".to_string(),
items: vec![bamboo_domain::TaskItem {
id: "task-1".to_string(),
description: "Do work".to_string(),
status: bamboo_domain::TaskItemStatus::InProgress,
..bamboo_domain::TaskItem::default()
}],
created_at: chrono::Utc::now(),
updated_at: chrono::Utc::now(),
});
session
.metadata
.insert("task_list_version".to_string(), "1".to_string());
let based_on_task_list = session.task_list.clone().expect("task list");
let mut state = super::super::startup::LoopRunState {
session_id: "session-task-eval".to_string(),
execution_id: "task-eval-execution".to_string(),
current_inputs: None,
model_name: "model".to_string(),
metrics_collector: None,
debug_logger: crate::runtime::runner::logging::DebugLogger::new(false),
task_context: crate::runtime::task_context::TaskLoopContext::from_session(&session),
overflow_recovery: super::super::startup::OverflowRecoveryState::default(),
task_evaluation: super::super::startup::TaskEvaluationState {
in_flight: None,
completed: Some(
crate::runtime::runner::task_lifecycle::AsyncTaskEvaluationResult {
metrics_round_id: "session-task-eval-task-evaluation-test".to_string(),
shared_session_id: "session-task-eval".to_string(),
round_number: 1,
based_on_task_context_version: 1,
based_on_task_list,
task_list_title: Some("Eval Tasks".to_string()),
evaluation_result: crate::runtime::task_evaluation::TaskEvaluationResult {
needs_evaluation: true,
updates: vec![crate::runtime::task_evaluation::TaskItemUpdate {
item_id: "task-1".to_string(),
status: bamboo_domain::TaskItemStatus::Completed,
notes: Some("done".to_string()),
evidence: Some("verified".to_string()),
blocker: None,
criteria_met: None,
}],
reasoning: "complete".to_string(),
prompt_tokens: 4,
completion_tokens: 2,
},
finished_at: chrono::Utc::now(),
error: None,
metrics_started: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(
true,
)),
metrics_terminal: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(
false,
)),
},
),
queued_request: None,
},
gold_evaluation: super::super::startup::GoldEvaluationState::default(),
auxiliary_models: crate::runtime::config::AuxiliaryModelConfig::default(),
runtime_state: AgentRuntimeState::new("session-task-eval"),
};
let config = crate::runtime::config::AgentLoopConfig {
persistence: Some(persistence.clone()),
..Default::default()
};
let (tx, mut rx) = tokio::sync::mpsc::channel(8);
super::apply_completed_task_evaluation(&mut session, &tx, &config, &mut state).await;
assert_eq!(
session.task_list.as_ref().unwrap().items[0].status,
bamboo_domain::TaskItemStatus::Completed
);
let event = rx
.recv()
.await
.expect("task update event should be emitted");
match event {
AgentEvent::TaskListUpdated { task_list, .. } => {
assert_eq!(
task_list.items[0].status,
bamboo_domain::TaskItemStatus::Completed
);
}
other => panic!("unexpected event: {other:?}"),
}
assert_eq!(
persistence.paired_task_patches.load(Ordering::SeqCst),
1,
"evaluation must persist exactly one targeted Task CAS"
);
assert_eq!(persistence.full_saves.load(Ordering::SeqCst), 0);
assert_eq!(
persistence.control_plane_saves.load(Ordering::SeqCst),
0,
"evaluation must not add a child/root whole control-plane save"
);
}
#[test]
fn test_build_round_id() {
let id =
crate::runtime::runner::round_prelude::build_round_id("session-123", "execution-a", 0);
assert_eq!(id, "session-123-run-execution-a-round-1");
let id = crate::runtime::runner::round_prelude::build_round_id("test", "execution-b", 4);
assert_eq!(id, "test-run-execution-b-round-5");
assert_ne!(
crate::runtime::runner::round_prelude::build_round_id("session-123", "execution-a", 0,),
crate::runtime::runner::round_prelude::build_round_id("session-123", "execution-b", 0,),
"the per-run round counter may reset, but the execution namespace must not"
);
let task_a = crate::runtime::runner::round_prelude::build_auxiliary_round_id(
"session-123",
"execution-a",
"task-evaluation",
1,
);
let task_b = crate::runtime::runner::round_prelude::build_auxiliary_round_id(
"session-123",
"execution-b",
"task-evaluation",
1,
);
let gold_a = crate::runtime::runner::round_prelude::build_auxiliary_round_id(
"session-123",
"execution-a",
"gold-evaluation",
1,
);
assert_eq!(
task_a,
"session-123-run-execution-a-task-evaluation-round-1"
);
assert_ne!(task_a, task_b);
assert_ne!(task_a, gold_a);
}
#[tokio::test]
async fn ensure_not_cancelled_returns_ok_when_not_cancelled() {
let token = tokio_util::sync::CancellationToken::new();
assert!(!token.is_cancelled());
}
#[tokio::test]
async fn ensure_not_cancelled_returns_error_when_cancelled() {
let token = tokio_util::sync::CancellationToken::new();
token.cancel();
assert!(token.is_cancelled());
}
#[test]
fn accumulate_round_usage_saturates_components_and_recomputes_total() {
let mut usage = MetricsTokenUsage {
prompt_tokens: u64::MAX - 5,
completion_tokens: u64::MAX - 9,
total_tokens: 0,
};
let delta = MetricsTokenUsage {
prompt_tokens: 10,
completion_tokens: 20,
total_tokens: 30,
};
usage.prompt_tokens = usage.prompt_tokens.saturating_add(delta.prompt_tokens);
usage.completion_tokens = usage
.completion_tokens
.saturating_add(delta.completion_tokens);
usage.recompute_total();
assert_eq!(usage.prompt_tokens, u64::MAX);
assert_eq!(usage.completion_tokens, u64::MAX);
assert_eq!(usage.total_tokens, u64::MAX);
}
#[test]
fn is_terminal_child_status_classifies_correctly() {
for s in ["completed", "error", "timeout", "cancelled", "skipped"] {
assert!(is_terminal_child_status(s), "{s} should be terminal");
}
for s in ["running", "pending", "queued", ""] {
assert!(!is_terminal_child_status(s), "{s} should be active");
}
}
struct FastChildWaitProvider;
#[async_trait::async_trait]
impl LLMProvider for FastChildWaitProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![activation_call(
"fast-wait-call",
"WaitControl",
"{}",
)])),
Ok(LLMChunk::Done),
])))
}
}
struct FastChildWaitExecutor(Arc<dyn Storage>);
struct ChildOutcomeConsumerProvider(Arc<AtomicBool>);
#[async_trait::async_trait]
impl LLMProvider for ChildOutcomeConsumerProvider {
async fn chat_stream(
&self,
messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
self.0.store(
messages
.iter()
.any(|message| message.content == "fast child outcome"),
Ordering::SeqCst,
);
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("continued after child".to_string())),
Ok(LLMChunk::Done),
])))
}
}
#[async_trait::async_trait]
impl bamboo_agent_core::tools::ToolExecutor for FastChildWaitExecutor {
async fn execute(
&self,
_call: &bamboo_agent_core::tools::ToolCall,
) -> std::result::Result<
bamboo_agent_core::tools::ToolResult,
bamboo_agent_core::tools::ToolError,
> {
let mut durable = self.0.load_session("fast-parent").await.unwrap().unwrap();
let mut runtime = AgentRuntimeState::new("fast-parent");
runtime.waiting_for_children = Some(WaitingForChildrenState::for_children(
vec!["fast-child".to_string()],
ChildWaitPolicy::All,
chrono::Utc::now(),
));
durable.agent_runtime_state = Some(runtime);
durable.metadata.insert(
"runtime.suspend_reason".to_string(),
"waiting_for_children".to_string(),
);
self.0.save_session(&durable).await.unwrap();
durable
.agent_runtime_state
.as_mut()
.unwrap()
.waiting_for_children = None;
durable.metadata.remove("runtime.suspend_reason");
let mut resume = Message::user("fast child outcome");
resume.metadata = Some(serde_json::json!({
"runtime_kind": "child_completion_resume"
}));
durable.add_message(resume);
self.0.save_session(&durable).await.unwrap();
Ok(bamboo_agent_core::tools::ToolResult {
success: true,
result: serde_json::json!({
"runtime_control": "waiting_for_children"
})
.to_string(),
display_preference: Some("runtime_control:waiting_for_children".to_string()),
images: Vec::new(),
})
}
fn list_tools(&self) -> Vec<bamboo_agent_core::tools::ToolSchema> {
vec![loading_test_schema("WaitControl")]
}
}
#[tokio::test]
async fn fast_child_completion_reconciles_wait_and_preserves_outcome_for_successor() {
let storage: Arc<dyn Storage> = Arc::new(TestStorage::default());
let mut session = Session::new("fast-parent", "model");
session.agent_runtime_state = Some(AgentRuntimeState::new(&session.id));
storage.save_session(&session).await.unwrap();
let mut config = config_with_storage(storage.clone());
config.model_name = Some("model".to_string());
let mut state = e2e_loop_state(&session.id);
let (tx, _rx) = tokio::sync::mpsc::channel(64);
let sent_complete = super::run_pipeline(
&mut session,
&tx,
Arc::new(FastChildWaitProvider),
Arc::new(FastChildWaitExecutor(storage.clone())),
&CancellationToken::new(),
&config,
&mut state,
)
.await
.expect("waiting-control round finalizes");
assert!(!sent_complete, "tool round hands the resume to a successor");
assert!(!session.metadata.contains_key("runtime.suspend_reason"));
assert!(state.runtime_state.waiting_for_children.is_none());
assert!(state.runtime_state.suspension.is_none());
assert_eq!(state.runtime_state.status, AgentStatusState::Idle);
assert!(session.messages.iter().any(|message| {
message.content == "fast child outcome"
&& message
.metadata
.as_ref()
.and_then(|metadata| metadata.get("runtime_kind"))
.and_then(|value| value.as_str())
== Some("child_completion_resume")
}));
let outcome_seen = Arc::new(AtomicBool::new(false));
let continued = super::run_pipeline(
&mut session,
&tx,
Arc::new(ChildOutcomeConsumerProvider(outcome_seen.clone())),
Arc::new(FastChildWaitExecutor(storage)),
&CancellationToken::new(),
&config,
&mut state,
)
.await
.expect("successor round consumes child completion");
assert!(continued, "successor round completes");
assert!(
outcome_seen.load(Ordering::SeqCst),
"successor provider must receive the child outcome"
);
}
struct ChildIndexStorage {
inner: Arc<TestStorage>,
children: Vec<(String, Option<String>)>,
}
#[async_trait::async_trait]
impl Storage for ChildIndexStorage {
async fn save_session(&self, session: &Session) -> std::io::Result<()> {
self.inner.save_session(session).await
}
async fn load_session(&self, id: &str) -> std::io::Result<Option<Session>> {
self.inner.load_session(id).await
}
async fn delete_session(&self, id: &str) -> std::io::Result<bool> {
self.inner.delete_session(id).await
}
async fn list_child_run_statuses(
&self,
_parent: &str,
) -> std::io::Result<Vec<(String, Option<String>)>> {
Ok(self.children.clone())
}
}
fn config_with_storage(storage: Arc<dyn Storage>) -> AgentLoopConfig {
let persistence: Arc<dyn bamboo_domain::RuntimeSessionPersistence> =
Arc::new(TestPersistence(storage.clone()));
AgentLoopConfig {
storage: Some(storage),
persistence: Some(persistence),
..AgentLoopConfig::default()
}
}
#[tokio::test]
async fn safety_net_suspends_on_orphaned_active_children() {
let inner = Arc::new(TestStorage::default());
let storage: Arc<dyn Storage> = Arc::new(ChildIndexStorage {
inner: inner.clone(),
children: vec![
("c-run".into(), Some("running".into())),
("c-pend".into(), None),
("c-done".into(), Some("completed".into())),
],
});
let config = config_with_storage(storage.clone());
let mut session = Session::new("parent-orphan", "model");
let mut runtime_state = AgentRuntimeState::new("parent-orphan");
let outcome =
maybe_suspend_for_orphaned_children(&mut session, &config, &mut runtime_state)
.await
.expect("parent wait observation succeeds")
.expect("must suspend when active children remain");
assert!(outcome.should_break && !outcome.sent_complete);
let wait = runtime_state
.waiting_for_children
.expect("durable wait registered");
assert_eq!(
wait.child_session_ids,
vec!["c-pend".to_string(), "c-run".to_string()]
);
assert_eq!(
session
.metadata
.get("runtime.suspend_reason")
.map(String::as_str),
Some("waiting_for_children")
);
let persisted = storage
.load_session("parent-orphan")
.await
.unwrap()
.unwrap();
assert!(persisted
.agent_runtime_state
.and_then(|s| s.waiting_for_children)
.is_some());
}
#[tokio::test]
async fn safety_net_noop_when_all_children_terminal() {
let inner = Arc::new(TestStorage::default());
let storage: Arc<dyn Storage> = Arc::new(ChildIndexStorage {
inner,
children: vec![
("a".into(), Some("completed".into())),
("b".into(), Some("error".into())),
],
});
let config = config_with_storage(storage);
let mut session = Session::new("parent-done", "model");
let mut runtime_state = AgentRuntimeState::new("parent-done");
assert!(
maybe_suspend_for_orphaned_children(&mut session, &config, &mut runtime_state)
.await
.unwrap()
.is_none(),
"no active children → must not suspend"
);
assert!(runtime_state.waiting_for_children.is_none());
}
struct FinishingChildIndexStorage {
inner: Arc<TestStorage>,
terminal: AtomicBool,
reads: AtomicUsize,
}
#[async_trait::async_trait]
impl Storage for FinishingChildIndexStorage {
async fn save_session(&self, session: &Session) -> std::io::Result<()> {
self.inner.save_session(session).await
}
async fn load_session(&self, id: &str) -> std::io::Result<Option<Session>> {
self.inner.load_session(id).await
}
async fn delete_session(&self, id: &str) -> std::io::Result<bool> {
self.inner.delete_session(id).await
}
async fn list_child_run_statuses(
&self,
_parent: &str,
) -> std::io::Result<Vec<(String, Option<String>)>> {
self.reads.fetch_add(1, Ordering::SeqCst);
Ok(vec![(
"fast-child".into(),
Some(
if self.terminal.load(Ordering::SeqCst) {
"completed"
} else {
"running"
}
.into(),
),
)])
}
}
#[tokio::test]
async fn orphan_wait_does_not_rearm_after_completion_wins_parent_lock() {
let storage = Arc::new(FinishingChildIndexStorage {
inner: Arc::new(TestStorage::default()),
terminal: AtomicBool::new(false),
reads: AtomicUsize::new(0),
});
let config = config_with_storage(storage.clone());
let parent_lock =
crate::session_app::child_completion_coordinator::session_resume_lock("fast-parent");
let parent_guard = parent_lock.lock().await;
let (started, ready) = tokio::sync::oneshot::channel();
let candidate = tokio::spawn(async move {
let mut session = Session::new("fast-parent", "model");
let mut runtime = AgentRuntimeState::new("fast-parent");
let _ = started.send(());
let outcome = maybe_suspend_for_orphaned_children(&mut session, &config, &mut runtime)
.await
.expect("orphan gate observation succeeds");
(outcome, runtime)
});
ready.await.expect("orphan gate task started");
tokio::task::yield_now().await;
assert_eq!(storage.reads.load(Ordering::SeqCst), 0);
storage.terminal.store(true, Ordering::SeqCst);
drop(parent_guard);
let (outcome, runtime) = candidate.await.expect("orphan gate task joins");
assert!(outcome.is_none(), "completed child must not arm a new wait");
assert!(runtime.waiting_for_children.is_none());
assert_eq!(storage.reads.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn safety_net_resuspends_existing_untagged_wait_before_terminal_gates() {
let storage: Arc<dyn Storage> = Arc::new(ChildIndexStorage {
inner: Arc::new(TestStorage::default()),
children: vec![("x".into(), Some("running".into()))],
});
let config = config_with_storage(storage.clone());
let mut session = Session::new("parent-waiting", "model");
let mut runtime_state = AgentRuntimeState::new("parent-waiting");
let wait = super::WaitingForChildrenState {
child_session_ids: vec!["x".into()],
wait_for: super::ChildWaitPolicy::All,
registered_at: chrono::Utc::now() - chrono::Duration::minutes(5),
timeout_at: None,
registered_by_tool_call_id: None,
};
runtime_state.waiting_for_children = Some(wait.clone());
session.agent_runtime_state = Some(runtime_state.clone());
storage.save_session(&session).await.unwrap();
let outcome =
maybe_suspend_for_orphaned_children(&mut session, &config, &mut runtime_state)
.await
.unwrap()
.expect("remaining inherited child wait suspends the reasoning turn");
assert!(outcome.should_break && !outcome.sent_complete);
assert_eq!(runtime_state.waiting_for_children.as_ref(), Some(&wait));
assert_eq!(
session
.metadata
.get("runtime.suspend_reason")
.map(String::as_str),
Some("waiting_for_children")
);
}
struct PartialChildWorkProvider {
calls: AtomicUsize,
outcome_seen: AtomicBool,
}
#[async_trait::async_trait]
impl LLMProvider for PartialChildWorkProvider {
async fn chat_stream(
&self,
messages: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
if self.calls.fetch_add(1, Ordering::SeqCst) == 0 {
self.outcome_seen.store(
messages
.iter()
.any(|message| message.content == "fast child outcome"),
Ordering::SeqCst,
);
return Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![
activation_call(
"partial-message",
"SubAgent",
r#"{"action":"send_message","child_session_id":"slow-child","message":"Use the fast child finding"}"#,
),
activation_call(
"partial-cancel",
"SubAgent",
r#"{"action":"cancel","child_session_id":"obsolete-child"}"#,
),
activation_call(
"partial-create",
"SubAgent",
r#"{"action":"create","prompt":"Verify the fast child finding","wait":false}"#,
),
])),
Ok(LLMChunk::Done),
])));
}
for (id, action) in [
("partial-message", "send_message"),
("partial-cancel", "cancel"),
("partial-create", "create"),
] {
assert!(
messages.iter().any(|message| {
message.tool_call_id.as_deref() == Some(id)
&& message.content.contains(&format!("handled {action}"))
}),
"tool feedback for {action} reaches the next reasoning round"
);
}
let native = |author, payload| {
Ok(LLMChunk::ProviderTranscriptItem(
bamboo_domain::ProviderTranscriptItem::try_from_payload(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
bamboo_domain::ProviderTranscriptOrigin::Provider,
author,
payload,
)
.unwrap(),
))
};
Ok(Box::pin(stream::iter(vec![
native(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"tool_search_call","id":"partial-native-search","execution":"server",
"call_id":"partial-search","status":"completed","arguments":{"paths":["children"]}
}),
),
native(
bamboo_domain::ProviderTranscriptAuthor::ToolResult,
serde_json::json!({
"type":"tool_search_output","id":"partial-native-search-result","execution":"server",
"call_id":"partial-search","status":"completed","tools":[]
}),
),
native(
bamboo_domain::ProviderTranscriptAuthor::Model,
serde_json::json!({
"type":"message","id":"partial-native-reply","role":"assistant","status":"completed",
"content":[{"type":"output_text","text":"continued after child","annotations":[]}]
}),
),
Ok(LLMChunk::ReasoningToken("remaining child work".into())),
Ok(LLMChunk::ReasoningSignature(
"partial-reasoning-signature".into(),
)),
Ok(LLMChunk::Token("continued after child".into())),
Ok(LLMChunk::Done),
])))
}
}
#[derive(Default)]
struct PartialChildWorkExecutor(std::sync::Mutex<Vec<String>>);
#[async_trait::async_trait]
impl bamboo_agent_core::tools::ToolExecutor for PartialChildWorkExecutor {
async fn execute(
&self,
call: &bamboo_agent_core::tools::ToolCall,
) -> bamboo_agent_core::tools::executor::Result<bamboo_agent_core::tools::ToolResult>
{
let arguments: serde_json::Value =
serde_json::from_str(&call.function.arguments).unwrap();
let action = arguments["action"].as_str().unwrap();
self.0.lock().unwrap().push(action.into());
Ok(bamboo_agent_core::tools::ToolResult::text(
true,
format!("handled {action}"),
))
}
fn list_tools(&self) -> Vec<bamboo_agent_core::tools::ToolSchema> {
vec![loading_test_schema("SubAgent")]
}
}
#[tokio::test]
async fn partial_child_outcome_can_drive_tool_work_before_resuspending_exact_inherited_wait() {
for (tagged, buffered) in [true, false].into_iter().flat_map(|tagged| {
[false, true]
.into_iter()
.map(move |buffered| (tagged, buffered))
}) {
let storage: Arc<dyn Storage> = Arc::new(ChildIndexStorage {
inner: Arc::new(TestStorage::default()),
children: vec![("slow-child".into(), Some("running".into()))],
});
let mut config = config_with_storage(storage.clone());
config.model_name = Some("model".into());
config.prompt_memory_flags = delegation_loop_config().prompt_memory_flags;
config.features_final_evidence_check = buffered;
let mut session = Session::new("partial-parent", "model");
let mut outcome = Message::user("fast child outcome");
outcome.metadata = Some(serde_json::json!({
"runtime_kind": "child_completion_resume"
}));
session.add_message(outcome);
let mut state = e2e_loop_state(&session.id);
let mut wait = WaitingForChildrenState::for_children(
vec!["fast-child".into(), "slow-child".into()],
ChildWaitPolicy::All,
Utc::now() - chrono::Duration::minutes(5),
);
wait.registered_by_tool_call_id = tagged.then(|| "original-wait-call".into());
state.runtime_state.waiting_for_children = Some(wait.clone());
session.agent_runtime_state = Some(state.runtime_state.clone());
storage.save_session(&session).await.unwrap();
let provider = Arc::new(PartialChildWorkProvider {
calls: AtomicUsize::new(0),
outcome_seen: AtomicBool::new(false),
});
let tools = Arc::new(PartialChildWorkExecutor::default());
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let sent_complete = super::run_pipeline(
&mut session,
&tx,
provider.clone(),
tools.clone(),
&CancellationToken::new(),
&config,
&mut state,
)
.await
.expect("partial child reasoning turn runs without waiting for the slow child");
assert!(provider.outcome_seen.load(Ordering::SeqCst));
assert_eq!(provider.calls.load(Ordering::SeqCst), 2);
let mut actions = tools.0.lock().unwrap().clone();
actions.sort();
assert_eq!(actions, vec!["cancel", "create", "send_message"]);
assert!(
!sent_complete,
"remaining child work bypasses terminal gates"
);
assert_eq!(state.runtime_state.status, AgentStatusState::Suspended);
assert_eq!(
state.runtime_state.waiting_for_children.as_ref(),
Some(&wait)
);
let replies: Vec<_> = session
.messages
.iter()
.filter(|message| message.content == "continued after child")
.collect();
assert_eq!(replies.len(), usize::from(!tagged));
let mut events = Vec::new();
while let Ok(event) = rx.try_recv() {
assert!(
!matches!(event, AgentEvent::Complete { .. }),
"an inherited wait must not publish Complete"
);
events.push(event);
}
assert_eq!(
events
.iter()
.filter(|event| matches!(event, AgentEvent::Token { content } if content == "continued after child"))
.count(),
usize::from(!buffered || !tagged),
"the untagged reply is published once in either streaming mode"
);
if let Some(reply) = replies.first() {
assert_eq!(reply.reasoning.as_deref(), Some("remaining child work"));
assert_eq!(
reply.reasoning_signature.as_deref(),
Some("partial-reasoning-signature")
);
assert_eq!(
events
.iter()
.filter(|event| matches!(event, AgentEvent::VisibleMessageStart { message_id, created_at } if message_id == &reply.id && *created_at == reply.created_at))
.count(),
1,
"persistence retains the published reply identity"
);
let boundary = session
.provider_transcript
.active_provider_boundary_sha256()
.unwrap();
let groups = session.provider_transcript.replayable_groups(
bamboo_domain::ProviderFamily::OpenAi,
bamboo_domain::ProviderProtocol::OpenAiResponsesV1,
boundary,
);
assert_eq!(groups.len(), 1);
assert_eq!(groups[0].anchor_message_id(), reply.id);
assert_eq!(groups[0].items().len(), 3);
assert_eq!(groups[0].items()[2].payload()["id"], "partial-native-reply");
}
}
}
struct InterruptedWaitStorage {
inner: Arc<TestStorage>,
read_error: bool,
saves: AtomicUsize,
}
#[async_trait::async_trait]
impl Storage for InterruptedWaitStorage {
async fn save_session(&self, session: &Session) -> std::io::Result<()> {
self.saves.fetch_add(1, Ordering::SeqCst);
self.inner.save_session(session).await
}
async fn load_session(&self, id: &str) -> std::io::Result<Option<Session>> {
if self.read_error {
return Err(std::io::Error::other("private fixture source failure"));
}
self.inner.load_session(id).await
}
async fn delete_session(&self, id: &str) -> std::io::Result<bool> {
self.inner.delete_session(id).await
}
async fn list_child_run_statuses(
&self,
_parent: &str,
) -> std::io::Result<Vec<(String, Option<String>)>> {
Ok(vec![("child-running".into(), Some("running".into()))])
}
}
#[tokio::test]
async fn interrupted_wait_restores_exact_value_without_save_and_completion_wins() {
for (tagged, policy) in [true, false].into_iter().flat_map(|tagged| {
[
ChildWaitPolicy::All,
ChildWaitPolicy::Any,
ChildWaitPolicy::FirstError,
]
.into_iter()
.map(move |policy| (tagged, policy))
}) {
let inner = Arc::new(TestStorage::default());
let mut parent = Session::new("interrupted-parent", "model");
let registered_at = Utc::now() - chrono::Duration::minutes(5);
let wait = WaitingForChildrenState {
child_session_ids: vec!["child-other".into(), "child-running".into()],
wait_for: policy,
registered_at,
timeout_at: match policy {
ChildWaitPolicy::All => None,
ChildWaitPolicy::Any => Some(registered_at),
ChildWaitPolicy::FirstError => {
Some(registered_at + chrono::Duration::minutes(10))
}
},
registered_by_tool_call_id: tagged.then(|| "original-subagent-call".into()),
};
let mut durable_runtime = AgentRuntimeState::new("original-parent-run");
durable_runtime.status = AgentStatusState::Suspended;
durable_runtime.waiting_for_children = Some(wait.clone());
parent.agent_runtime_state = Some(durable_runtime);
inner.save_session(&parent).await.unwrap();
let storage = Arc::new(InterruptedWaitStorage {
inner: inner.clone(),
read_error: false,
saves: AtomicUsize::new(0),
});
let config = config_with_storage(storage.clone());
let mut runtime = AgentRuntimeState::new("interrupt-reasoning-run");
runtime.waiting_for_children = Some(wait.clone());
parent.agent_runtime_state = Some(runtime.clone());
let outcome = maybe_suspend_for_orphaned_children(&mut parent, &config, &mut runtime)
.await
.unwrap()
.unwrap();
assert!(outcome.should_break && !outcome.sent_complete);
assert_eq!(
runtime.waiting_for_children.as_ref(),
Some(&wait),
"tagged={tagged}, policy={policy:?}"
);
assert_eq!(
parent
.agent_runtime_state
.as_ref()
.unwrap()
.waiting_for_children
.as_ref(),
Some(&wait)
);
assert_eq!(storage.saves.load(Ordering::SeqCst), 0);
assert!(
maybe_suspend_for_orphaned_children(&mut parent, &config, &mut runtime)
.await
.unwrap()
.unwrap()
.should_break
);
assert_eq!(runtime.waiting_for_children.as_ref(), Some(&wait));
assert_eq!(storage.saves.load(Ordering::SeqCst), 0);
let mut durable = inner.load_session(&parent.id).await.unwrap().unwrap();
assert_eq!(
durable
.agent_runtime_state
.as_ref()
.unwrap()
.waiting_for_children
.as_ref(),
Some(&wait)
);
durable
.agent_runtime_state
.as_mut()
.unwrap()
.waiting_for_children = None;
inner.save_session(&durable).await.unwrap();
let cleared = storage.load_session(&parent.id).await.unwrap().unwrap();
runtime.status = AgentStatusState::Suspended;
super::reconcile_child_wait_at_suspend(
&mut parent,
&mut runtime,
cleared.agent_runtime_state.unwrap().waiting_for_children,
);
assert!(runtime.waiting_for_children.is_none());
assert_eq!(runtime.status, AgentStatusState::Idle);
assert!(!parent.metadata.contains_key("runtime.suspend_reason"));
assert_eq!(storage.saves.load(Ordering::SeqCst), 0);
runtime.status = AgentStatusState::Running;
runtime.waiting_for_children = Some(wait);
assert!(
maybe_suspend_for_orphaned_children(&mut parent, &config, &mut runtime)
.await
.unwrap()
.is_none()
);
assert!(runtime.waiting_for_children.is_none());
assert_eq!(runtime.status, AgentStatusState::Running);
assert!(!parent.metadata.contains_key("runtime.suspend_reason"));
assert_eq!(storage.saves.load(Ordering::SeqCst), 0);
}
}
#[tokio::test]
async fn interrupted_wait_absent_preserves_new_orphan_wait_behavior() {
for persisted_parent in [false, true] {
let inner = Arc::new(TestStorage::default());
let mut parent = Session::new("no-durable-wait", "model");
if persisted_parent {
inner.save_session(&parent).await.unwrap();
}
let storage = Arc::new(InterruptedWaitStorage {
inner: inner.clone(),
read_error: false,
saves: AtomicUsize::new(0),
});
let config = config_with_storage(storage.clone());
let mut runtime = AgentRuntimeState::new("orphan-run");
let before = Utc::now();
let outcome = maybe_suspend_for_orphaned_children(&mut parent, &config, &mut runtime)
.await
.unwrap()
.unwrap();
assert!(outcome.should_break && !outcome.sent_complete);
assert_eq!(storage.saves.load(Ordering::SeqCst), 1);
let persisted = inner.load_session(&parent.id).await.unwrap().unwrap();
let wait = persisted
.agent_runtime_state
.unwrap()
.waiting_for_children
.unwrap();
assert_eq!(wait.child_session_ids, vec!["child-running"]);
assert_eq!(wait.wait_for, ChildWaitPolicy::All);
assert!(wait.registered_at >= before);
assert_eq!(
wait.timeout_at,
Some(wait.registered_at + chrono::Duration::hours(6))
);
assert!(wait.registered_by_tool_call_id.is_none());
}
}
#[tokio::test]
async fn interrupted_wait_observation_fails_closed_on_read_error_or_replaced_parent() {
for read_error in [true, false] {
let inner = Arc::new(TestStorage::default());
let mut parent = Session::new("observation-parent", "model");
let mut replacement = parent.clone();
replacement.created_at += chrono::Duration::nanoseconds(1);
inner.save_session(&replacement).await.unwrap();
let storage = Arc::new(InterruptedWaitStorage {
inner,
read_error,
saves: AtomicUsize::new(0),
});
let config = config_with_storage(storage.clone());
let original = serde_json::to_value(&parent).unwrap();
let mut runtime = AgentRuntimeState::new("parent-run");
let original_runtime = runtime.clone();
let error = maybe_suspend_for_orphaned_children(&mut parent, &config, &mut runtime)
.await
.err()
.expect("observation must not renew on uncertainty");
assert!(matches!(error, AgentError::Tool(_)));
assert!(!error.to_string().contains("private fixture source failure"));
assert_eq!(serde_json::to_value(&parent).unwrap(), original);
assert_eq!(runtime, original_runtime);
assert_eq!(storage.saves.load(Ordering::SeqCst), 0);
for tagged in [true, false] {
let mut live_wait = WaitingForChildrenState::for_children(
vec!["child-running".into()],
ChildWaitPolicy::All,
Utc::now(),
);
live_wait.registered_by_tool_call_id = tagged.then(|| "wait-call".into());
runtime.waiting_for_children = Some(live_wait.clone());
let error = maybe_suspend_for_orphaned_children(&mut parent, &config, &mut runtime)
.await
.err()
.expect("inherited wait observation must fail closed");
assert!(matches!(error, AgentError::Tool(_)));
assert_eq!(runtime.waiting_for_children.as_ref(), Some(&live_wait));
assert_eq!(storage.saves.load(Ordering::SeqCst), 0);
}
}
}
#[tokio::test]
async fn suspend_to_wait_for_bash_sets_reason_and_state() {
let mut session = Session::new("s-bash", "model");
let mut runtime_state = AgentRuntimeState::new("s-bash");
let outcome = suspend_to_wait_for_bash(
&mut session,
&mut runtime_state,
None,
vec!["bg-1".to_string(), "bg-2".to_string()],
)
.await;
assert!(outcome.should_break, "must break the turn");
assert!(!outcome.sent_complete, "must not send complete");
let wait = runtime_state
.waiting_for_bash
.expect("durable bash wait should be registered");
assert_eq!(wait.bash_ids, vec!["bg-1".to_string(), "bg-2".to_string()]);
assert_eq!(
session
.metadata
.get("runtime.suspend_reason")
.map(String::as_str),
Some("waiting_for_bash"),
"metadata reason must match the discriminant arm"
);
}
#[tokio::test]
async fn bash_safety_net_noop_when_already_waiting() {
let config = AgentLoopConfig::default();
let mut session = Session::new("s-bash-waiting", "model");
let mut runtime_state = AgentRuntimeState::new("s-bash-waiting");
runtime_state.waiting_for_bash = Some(super::WaitingForBashState {
bash_ids: vec!["bg-1".to_string()],
registered_at: chrono::Utc::now(),
timeout_at: None,
});
assert!(
maybe_suspend_for_outstanding_bash(&mut session, &config, &mut runtime_state)
.await
.is_none(),
"must not re-suspend when a bash wait is already registered"
);
}
struct StubBashPersistence;
#[async_trait::async_trait]
impl bamboo_domain::RuntimeSessionPersistence for StubBashPersistence {
async fn save_runtime_session(&self, _session: &mut Session) -> std::io::Result<()> {
Ok(())
}
}
#[derive(Clone)]
struct RecordingBashResumeHook {
calls: Arc<std::sync::Mutex<Vec<(String, Vec<String>)>>>,
}
impl crate::runtime::config::BashResumeHook for RecordingBashResumeHook {
fn arrange_bash_self_resume(&self, session_id: String, bash_ids: Vec<String>) {
self.calls
.lock()
.expect("hook mutex")
.push((session_id, bash_ids));
}
}
struct NoopBashResumeHook;
impl crate::runtime::config::BashResumeHook for NoopBashResumeHook {
fn arrange_bash_self_resume(&self, _: String, _: Vec<String>) {}
}
#[tokio::test]
async fn bash_gate_arranges_self_resume_hook_on_suspend() {
let session_id = "s-bash-liveness";
let mut config = AgentLoopConfig::default();
config.persistence = Some(Arc::new(StubBashPersistence));
let hook = RecordingBashResumeHook {
calls: Arc::new(std::sync::Mutex::new(Vec::new())),
};
config.bash_resume_hook = Some(Arc::new(hook.clone()));
let shell = bamboo_tools::tools::bash_runtime::spawn_background(
"sleep 5",
None,
None,
Some(session_id.to_string()),
false,
None,
)
.await
.expect("spawn");
let mut session = Session::new(session_id, "model");
let mut runtime_state = AgentRuntimeState::new(session_id);
let outcome =
maybe_suspend_for_outstanding_bash(&mut session, &config, &mut runtime_state).await;
let _ = shell.kill().await;
assert!(
outcome.is_some(),
"gate should suspend with a running shell"
);
assert!(
runtime_state.waiting_for_bash.is_some(),
"durable wait registered"
);
let calls = hook.calls.lock().expect("hook calls");
assert_eq!(calls.len(), 1, "hook called exactly once");
assert_eq!(calls[0].0, session_id);
assert!(!calls[0].1.is_empty(), "hook received bash ids");
}
#[tokio::test]
async fn bash_gate_no_suspend_when_all_shells_finished() {
let session_id = "s-bash-toctou";
let mut config = AgentLoopConfig::default();
config.persistence = Some(Arc::new(StubBashPersistence));
config.bash_resume_hook = Some(Arc::new(NoopBashResumeHook));
let shell = bamboo_tools::tools::bash_runtime::spawn_background(
"true",
None,
None,
Some(session_id.to_string()),
false,
None,
)
.await
.expect("spawn");
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
loop {
if shell.status() != "running" {
break;
}
if tokio::time::Instant::now() >= deadline {
panic!("test shell did not finish in 5s");
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
let mut session = Session::new(session_id, "model");
let mut runtime_state = AgentRuntimeState::new(session_id);
let outcome =
maybe_suspend_for_outstanding_bash(&mut session, &config, &mut runtime_state).await;
assert!(
outcome.is_none(),
"must not suspend when no shells are running"
);
assert!(
runtime_state.waiting_for_bash.is_none(),
"no bash wait registered"
);
}
#[tokio::test]
async fn bash_suspend_reason_matches_suspended_discriminant() {
let mut session = Session::new("s-discriminant", "model");
let mut runtime_state = AgentRuntimeState::new("s-discriminant");
suspend_to_wait_for_bash(
&mut session,
&mut runtime_state,
None,
vec!["bg-1".to_string()],
)
.await;
let reason = session
.metadata
.get("runtime.suspend_reason")
.map(String::as_str);
assert_eq!(reason, Some("waiting_for_bash"));
let produces_suspended = matches!(
reason,
Some("awaiting_clarification")
| Some("awaiting_parent_approval")
| Some("waiting_for_children")
| Some("waiting_for_bash")
);
assert!(
produces_suspended,
"waiting_for_bash must be Suspended-producing"
);
}
use super::handle_tool_calls_path;
use crate::runtime::runner::round_frame::RoundFrame;
use crate::runtime::runner::tool_execution::execute_round_tool_calls;
use crate::runtime::stream::handler::StreamHandlingOutput;
use crate::runtime::task_context::TaskLoopContext;
use bamboo_agent_core::tools::{
FunctionCall, FunctionSchema, ToolCall, ToolExecutor, ToolResult, ToolSchema,
};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
struct CancelProbeToolExecutor {
block: bool,
started: Arc<AtomicBool>,
}
#[async_trait::async_trait]
impl ToolExecutor for CancelProbeToolExecutor {
async fn execute(
&self,
_call: &ToolCall,
) -> bamboo_agent_core::tools::executor::Result<ToolResult> {
self.started.store(true, Ordering::SeqCst);
if self.block {
tokio::time::sleep(Duration::from_secs(120)).await;
}
Ok(ToolResult {
success: true,
result: "tool-result-123".to_string(),
display_preference: None,
images: Vec::new(),
})
}
fn list_tools(&self) -> Vec<ToolSchema> {
vec![ToolSchema {
schema_type: "function".to_string(),
function: FunctionSchema {
name: "Read".to_string(),
description: "read tool".to_string(),
parameters: serde_json::json!({ "type": "object", "properties": {} }),
},
}]
}
}
fn single_read_call() -> ToolCall {
ToolCall {
id: "call-read".to_string(),
tool_type: "function".to_string(),
function: FunctionCall {
name: "Read".to_string(),
arguments: "{}".to_string(),
},
}
}
fn stream_output_with_tool_call(call: ToolCall) -> StreamHandlingOutput {
StreamHandlingOutput {
response_id: None,
visible_message: None,
content: String::new(),
reasoning_content: String::new(),
reasoning_signature: None,
token_count: 0,
tool_calls: vec![call],
output_tokens: 0,
thinking_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
provider_usage: None,
input_tokens: 0,
provider_transcript_items: Vec::new(),
}
}
#[tokio::test]
async fn tool_execution_cancel_returns_promptly() {
let started = Arc::new(AtomicBool::new(false));
let tools: Arc<dyn ToolExecutor> = Arc::new(CancelProbeToolExecutor {
block: true,
started: started.clone(),
});
let (event_tx, _event_rx) = mpsc::channel::<AgentEvent>(128);
let llm: Arc<dyn LLMProvider> = Arc::new(StubProvider);
let config = AgentLoopConfig::default();
let mut session = Session::new("s-cancel", "model");
let frame = RoundFrame {
vision_support: true,
session_id: "s-cancel",
round_id: "r1",
turn: 0,
debug_enabled: false,
event_tx: &event_tx,
metrics_collector: None,
config: &config,
llm: &llm,
tools: &tools,
};
let auxiliary_models = crate::runtime::config::AuxiliaryModelConfig::default();
let mut runtime_state = AgentRuntimeState::new("s-cancel");
let mut task_context: Option<TaskLoopContext> = None;
let cancel_token = CancellationToken::new();
let driver_started = started.clone();
let driver_token = cancel_token.clone();
let driver = tokio::spawn(async move {
for _ in 0..500 {
if driver_started.load(Ordering::SeqCst) {
break;
}
tokio::time::sleep(Duration::from_millis(2)).await;
}
assert!(
driver_started.load(Ordering::SeqCst),
"tool never started executing"
);
driver_token.cancel();
});
let t0 = std::time::Instant::now();
let result = tokio::time::timeout(
Duration::from_secs(5),
handle_tool_calls_path(
&frame,
stream_output_with_tool_call(single_read_call()),
MetricsTokenUsage {
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
},
&mut session,
&mut runtime_state,
&mut crate::runtime::runner::tool_execution::ToolPolicyGuard::default(),
&auxiliary_models,
"model",
&mut task_context,
&cancel_token,
),
)
.await;
let elapsed = t0.elapsed();
let _ = driver.await;
let inner = result.expect(
"handle_tool_calls_path did not return within 5s — cancel not honored during tool execution",
);
assert!(
matches!(inner, Err(AgentError::Cancelled)),
"expected Err(AgentError::Cancelled), got {:?}",
inner.as_ref().err()
);
assert!(
elapsed < Duration::from_secs(2),
"cancel was not prompt (tool would otherwise block for ~120s): {:?}",
elapsed
);
}
#[tokio::test]
async fn normal_tool_batch_completes_unchanged() {
let tools: Arc<dyn ToolExecutor> = Arc::new(CancelProbeToolExecutor {
block: false,
started: Arc::new(AtomicBool::new(false)),
});
let (event_tx, _event_rx) = mpsc::channel::<AgentEvent>(128);
let llm: Arc<dyn LLMProvider> = Arc::new(StubProvider);
let config = AgentLoopConfig::default();
let mut session = Session::new("s-normal", "model");
let frame = RoundFrame {
vision_support: true,
session_id: "s-normal",
round_id: "r1",
turn: 0,
debug_enabled: false,
event_tx: &event_tx,
metrics_collector: None,
config: &config,
llm: &llm,
tools: &tools,
};
let tool_schemas = tools.list_tools();
let effective_callable_set =
crate::runtime::runner::tool_execution::legacy_effective_callable_set(&tool_schemas);
let mut runtime_state = AgentRuntimeState::new("s-normal");
let mut task_context: Option<TaskLoopContext> = None;
let result = tokio::time::timeout(
Duration::from_secs(10),
execute_round_tool_calls(crate::runtime::runner::tool_execution::RoundToolExecution {
tool_calls: std::slice::from_ref(&single_read_call()),
frame: &frame,
session: &mut session,
runtime_state: &mut runtime_state,
policy_guard: &mut crate::runtime::runner::tool_execution::ToolPolicyGuard::default(
),
task_context: &mut task_context,
compression_model_name: None,
compression_model_provider: None,
tool_schemas: &tool_schemas,
effective_callable_set: &effective_callable_set,
}),
)
.await
.expect("normal tool batch did not complete within 10s");
let round_result = result.expect("normal batch should return Ok");
assert!(!round_result.awaiting_clarification);
assert!(!round_result.waiting_for_children);
assert!(
session
.messages
.iter()
.any(|m| m.role == bamboo_agent_core::Role::Tool
&& m.content.contains("tool-result-123")),
"expected a tool-result message, got {} message(s)",
session.messages.len()
);
}
struct RecordingSequentialExecutor {
executed: Arc<std::sync::Mutex<Vec<String>>>,
}
#[async_trait::async_trait]
impl ToolExecutor for RecordingSequentialExecutor {
async fn execute(
&self,
call: &ToolCall,
) -> bamboo_agent_core::tools::executor::Result<ToolResult> {
self.executed
.lock()
.unwrap()
.push(call.function.name.clone());
Ok(ToolResult {
success: true,
result: format!("result-of-{}", call.function.name),
display_preference: None,
images: Vec::new(),
})
}
fn list_tools(&self) -> Vec<ToolSchema> {
["compact_context", "tool_b", "tool_c"]
.iter()
.map(|name| ToolSchema {
schema_type: "function".to_string(),
function: FunctionSchema {
name: name.to_string(),
description: "test tool".to_string(),
parameters: serde_json::json!({ "type": "object", "properties": {} }),
},
})
.collect()
}
fn call_parallel_classification(
&self,
_call: &ToolCall,
) -> (bamboo_agent_core::tools::ToolMutability, bool) {
(bamboo_agent_core::tools::ToolMutability::Mutating, false)
}
}
struct FailingCompressionProvider {
compression_calls: Arc<std::sync::atomic::AtomicUsize>,
}
#[async_trait::async_trait]
impl LLMProvider for FailingCompressionProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
async fn chat_stream_with_options(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
options: Option<&bamboo_llm::LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
let purpose = options
.and_then(|o| o.request_purpose.as_deref())
.unwrap_or("");
if purpose == "compression" {
self.compression_calls.fetch_add(1, Ordering::SeqCst);
return Err(LLMError::Api(
"http 500 transient upstream failure (compression summarization)".to_string(),
));
}
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
}
fn tool_call(id: &str, name: &str) -> ToolCall {
ToolCall {
id: id.to_string(),
tool_type: "function".to_string(),
function: FunctionCall {
name: name.to_string(),
arguments: "{}".to_string(),
},
}
}
#[tokio::test]
async fn mid_turn_compression_failure_is_best_effort_and_does_not_retry_turn() {
let compression_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let executed = Arc::new(std::sync::Mutex::new(Vec::<String>::new()));
let llm: Arc<dyn LLMProvider> = Arc::new(FailingCompressionProvider {
compression_calls: compression_calls.clone(),
});
let tools: Arc<dyn ToolExecutor> = Arc::new(RecordingSequentialExecutor {
executed: executed.clone(),
});
let (event_tx, _event_rx) = mpsc::channel::<AgentEvent>(128);
let config = AgentLoopConfig {
model_name: Some("model".to_string()),
background_model_name: Some("summarizer".to_string()),
..AgentLoopConfig::default()
};
let mut session = Session::new("s-compress-fail", "model");
session.add_message(Message::system("system"));
session.add_message(Message::user("do the work"));
session.add_message(Message::assistant("prior assistant turn".to_string(), None));
session.add_message(Message::user("keep going"));
let frame = RoundFrame {
vision_support: true,
session_id: "s-compress-fail",
round_id: "r1",
turn: 0,
debug_enabled: false,
event_tx: &event_tx,
metrics_collector: None,
config: &config,
llm: &llm,
tools: &tools,
};
let auxiliary_models = crate::runtime::config::AuxiliaryModelConfig::default();
let mut runtime_state = AgentRuntimeState::new("s-compress-fail");
let mut task_context: Option<TaskLoopContext> = None;
let cancel_token = CancellationToken::new();
let stream_output = StreamHandlingOutput {
response_id: None,
visible_message: None,
content: String::new(),
reasoning_content: String::new(),
reasoning_signature: None,
token_count: 0,
tool_calls: vec![
tool_call("call-compact", "compact_context"),
tool_call("call-b", "tool_b"),
tool_call("call-c", "tool_c"),
],
output_tokens: 0,
thinking_tokens: 0,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 0,
provider_usage: None,
input_tokens: 0,
provider_transcript_items: Vec::new(),
};
let result = tokio::time::timeout(
Duration::from_secs(10),
handle_tool_calls_path(
&frame,
stream_output,
MetricsTokenUsage {
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
},
&mut session,
&mut runtime_state,
&mut crate::runtime::runner::tool_execution::ToolPolicyGuard::default(),
&auxiliary_models,
"model",
&mut task_context,
&cancel_token,
),
)
.await
.expect("handle_tool_calls_path did not return within 10s");
let _outcome = result.expect(
"mid-turn compression failure must be best-effort (Ok), not a whole-turn error/retry",
);
assert!(
compression_calls.load(Ordering::SeqCst) >= 1,
"mid-turn compression summarization must have been attempted (and failed)"
);
let ran = executed.lock().unwrap().clone();
assert_eq!(
ran,
vec![
"compact_context".to_string(),
"tool_b".to_string(),
"tool_c".to_string(),
],
"all tools must execute in order despite the mid-turn compression failure"
);
let assistant_turns = session
.messages
.iter()
.filter(|m| {
m.role == bamboo_agent_core::Role::Assistant
&& m.tool_calls.as_ref().is_some_and(|calls| {
calls.iter().any(|c| c.function.name == "compact_context")
})
})
.count();
assert_eq!(
assistant_turns, 1,
"exactly one assistant message must exist for the turn (no duplicate)"
);
for (id, name) in [
("call-compact", "compact_context"),
("call-b", "tool_b"),
("call-c", "tool_c"),
] {
let count = session
.messages
.iter()
.filter(|m| {
m.role == bamboo_agent_core::Role::Tool && m.tool_call_id.as_deref() == Some(id)
})
.count();
assert_eq!(count, 1, "tool {name} must have exactly one result message");
}
}
struct AbortBeforeRoundHook;
#[async_trait::async_trait]
impl AgentHook for AbortBeforeRoundHook {
fn point(&self) -> AgentHookPoint {
AgentHookPoint::BeforeRound
}
async fn run(
&self,
_point: AgentHookPoint,
_payload: &HookPayload,
_session: &Session,
) -> HookResult {
HookResult::Abort {
reason: "injected before-round failure".to_string(),
}
}
}
struct EvalDropFlag(Arc<AtomicBool>);
impl Drop for EvalDropFlag {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
fn seed_blocked_task_evaluation(
session: &mut Session,
state: &mut super::super::startup::LoopRunState,
) -> Arc<AtomicBool> {
let now = chrono::Utc::now();
session.set_task_list(bamboo_domain::TaskList {
session_id: session.id.clone(),
title: "Pending evaluator".to_string(),
items: vec![bamboo_domain::TaskItem {
id: "pending".to_string(),
description: "wait".to_string(),
status: bamboo_domain::TaskItemStatus::InProgress,
..bamboo_domain::TaskItem::default()
}],
created_at: now,
updated_at: now,
});
session.set_task_list_version_meta("1");
state.task_context = crate::runtime::task_context::TaskLoopContext::from_session(session);
let request = crate::runtime::runner::task_lifecycle::build_async_task_evaluation_request(
&state.task_context,
session,
&state.session_id,
1,
Some("fast-model"),
None,
crate::runtime::stream::handler::StreamTimeoutContext::new(
bamboo_config::StreamTimeoutConfig::default(),
Some("provider"),
Some("fast-model"),
),
)
.expect("evaluation request")
.expect("task context");
state.task_evaluation.queued_request = Some(request.clone());
let dropped = Arc::new(AtomicBool::new(false));
let drop_flag = EvalDropFlag(dropped.clone());
let join_handle = tokio::spawn(async move {
let _drop_flag = drop_flag;
futures::future::pending::<
Option<crate::runtime::runner::task_lifecycle::AsyncTaskEvaluationResult>,
>()
.await
});
state.task_evaluation.in_flight = Some(InFlightTaskEvaluation {
request,
metrics_started: Arc::new(AtomicBool::new(false)),
metrics_terminal: Arc::new(AtomicBool::new(false)),
join_handle,
});
dropped
}
#[tokio::test]
async fn before_round_hook_failure_aborts_task_eval_and_clears_latest_snapshot() {
let mut session = Session::new("task-eval-before-round-hook", "model");
let mut state = e2e_loop_state(&session.id);
let dropped = seed_blocked_task_evaluation(&mut session, &mut state);
let mut config = stop_hook_config(Arc::new(AbortBeforeRoundHook));
config.model_name = Some("model".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let result = super::run_pipeline(
&mut session,
&tx,
Arc::new(StubProvider),
Arc::new(AlwaysOkExecutor),
&CancellationToken::new(),
&config,
&mut state,
)
.await;
assert!(result.is_err(), "BeforeRound abort must fail the run");
assert!(state.task_evaluation.in_flight.is_none());
assert!(state.task_evaluation.queued_request.is_none());
assert!(
dropped.load(Ordering::SeqCst),
"evaluator future was aborted"
);
}
#[tokio::test]
async fn no_tool_finalize_hook_failure_aborts_task_eval_and_clears_latest_snapshot() {
let mut session = Session::new("task-eval-no-tool-hook", "model");
let mut state = e2e_loop_state(&session.id);
let dropped = seed_blocked_task_evaluation(&mut session, &mut state);
let mut config = stop_hook_config(Arc::new(AbortFinalizeHook));
config.model_name = Some("model".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let result = super::run_pipeline(
&mut session,
&tx,
Arc::new(CanonicalUsageProvider {
prompt_tokens: 1,
completion_tokens: 1,
content: Some("final answer"),
}),
Arc::new(AlwaysOkExecutor),
&CancellationToken::new(),
&config,
&mut state,
)
.await;
assert!(
result.is_err(),
"BeforeFinalize abort must fail no-tool path"
);
assert!(state.task_evaluation.in_flight.is_none());
assert!(state.task_evaluation.queued_request.is_none());
assert!(
dropped.load(Ordering::SeqCst),
"evaluator future was aborted"
);
}
#[tokio::test]
async fn abort_while_waiting_for_auxiliary_permit_does_not_start_metric_row() {
use bamboo_metrics::storage::MetricsStorage;
let session_id = "task-eval-waiting-budget-cancel";
let (_dir, collector, storage) = create_pipeline_metrics().await;
collector.session_started(session_id, "model", Utc::now());
let provider: Arc<dyn LLMProvider> = Arc::new(StubProvider);
let held =
crate::runtime::runner::auxiliary_budget::acquire(&provider, "fast-model", 1).await;
let mut session = Session::new(session_id, "model");
let mut state = e2e_loop_state(session_id);
seed_blocked_task_evaluation(&mut session, &mut state);
let (reset_tx, _reset_rx) = mpsc::channel(4);
super::abort_in_flight_evaluations(&mut state, &reset_tx, "test_reset").await;
let request = crate::runtime::runner::task_lifecycle::build_async_task_evaluation_request(
&state.task_context,
&session,
session_id,
1,
Some("fast-model"),
None,
crate::runtime::stream::handler::StreamTimeoutContext::default(),
)
.expect("request builds")
.expect("task context exists");
state.metrics_collector = Some(collector);
let (event_tx, _event_rx) = mpsc::channel(8);
super::spawn_task_evaluation_request(
&mut state,
&event_tx,
request,
provider,
CancellationToken::new(),
1,
);
tokio::task::yield_now().await;
let started = state
.task_evaluation
.in_flight
.as_ref()
.expect("evaluation waits")
.metrics_started
.clone();
assert!(!started.load(Ordering::Acquire));
super::abort_in_flight_evaluations(&mut state, &event_tx, "run_cancelled").await;
drop(held);
assert!(!started.load(Ordering::Acquire));
assert!(state.task_evaluation.in_flight.is_none());
assert!(state.task_evaluation.queued_request.is_none());
let detail = tokio::time::timeout(Duration::from_secs(1), async {
loop {
if let Some(detail) = storage
.session_detail(session_id)
.await
.expect("session metrics query")
{
break detail;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("session start metric persists");
assert!(
detail.rounds.is_empty(),
"permit-wait cancellation must not create a started evaluator row"
);
}
struct CompletedTaskEvaluationProvider;
#[async_trait::async_trait]
impl LLMProvider for CompletedTaskEvaluationProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
let call = bamboo_agent_core::tools::ToolCall {
id: "completed-unharvested-evaluation".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "update_task_item".to_string(),
arguments: r#"{"item_id":"pending","status":"completed","notes":"verified"}"#
.to_string(),
},
};
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::Done),
])))
}
}
#[tokio::test]
async fn cleanup_harvests_finished_task_evaluation_without_false_cancellation() {
use bamboo_metrics::storage::MetricsStorage;
let session_id = "task-eval-finished-before-cleanup";
let (_dir, collector, storage) = create_pipeline_metrics().await;
collector.session_started(session_id, "model", Utc::now());
let mut session = Session::new(session_id, "model");
let mut state = e2e_loop_state(session_id);
seed_blocked_task_evaluation(&mut session, &mut state);
let (reset_tx, _reset_rx) = mpsc::channel(4);
super::abort_in_flight_evaluations(&mut state, &reset_tx, "test_reset").await;
let request = crate::runtime::runner::task_lifecycle::build_async_task_evaluation_request(
&state.task_context,
&session,
session_id,
1,
Some("fast-model"),
None,
crate::runtime::stream::handler::StreamTimeoutContext::default(),
)
.expect("request builds")
.expect("task context exists");
let metrics_round_id = request.metrics_round_id.clone();
state.metrics_collector = Some(collector);
let (event_tx, mut event_rx) = mpsc::channel(16);
super::spawn_task_evaluation_request(
&mut state,
&event_tx,
request,
Arc::new(CompletedTaskEvaluationProvider),
CancellationToken::new(),
1,
);
let metrics_started = state
.task_evaluation
.in_flight
.as_ref()
.expect("evaluation spawned")
.metrics_started
.clone();
let metrics_terminal = state
.task_evaluation
.in_flight
.as_ref()
.expect("evaluation spawned")
.metrics_terminal
.clone();
tokio::time::timeout(Duration::from_secs(1), async {
while !state
.task_evaluation
.in_flight
.as_ref()
.expect("unharvested handle remains installed")
.join_handle
.is_finished()
{
tokio::task::yield_now().await;
}
})
.await
.expect("evaluator should finish before cleanup");
assert!(metrics_started.load(Ordering::Acquire));
assert!(!metrics_terminal.load(Ordering::Acquire));
super::abort_in_flight_evaluations(&mut state, &event_tx, "run_completed").await;
assert!(state.task_evaluation.in_flight.is_none());
assert!(metrics_terminal.load(Ordering::Acquire));
let mut started_events = 0;
let mut completed_events = 0;
let mut cancelled_events = 0;
while let Ok(event) = event_rx.try_recv() {
match event {
AgentEvent::TaskEvaluationStarted { .. } => started_events += 1,
AgentEvent::TaskEvaluationCompleted { .. } => completed_events += 1,
AgentEvent::TaskEvaluationCancelled { .. } => cancelled_events += 1,
_ => {}
}
}
assert_eq!(started_events, 1);
assert_eq!(completed_events, 1);
assert_eq!(
cancelled_events, 0,
"a completed evaluator must not receive a duplicate cancellation event"
);
let detail = tokio::time::timeout(Duration::from_secs(2), async {
loop {
if let Some(detail) = storage
.session_detail(session_id)
.await
.expect("session metrics query")
{
if detail.rounds.len() == 1 && detail.rounds[0].completed_at.is_some() {
break detail;
}
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("harvested evaluator terminal metric persists");
let row = &detail.rounds[0];
assert_eq!(row.round_id, metrics_round_id);
assert_eq!(row.status, MetricsRoundStatus::Success);
assert!(row.error.is_none());
assert!(row.token_usage.prompt_tokens > 0);
assert!(row.token_usage.completion_tokens > 0);
assert_eq!(
row.token_usage.total_tokens,
row.token_usage
.prompt_tokens
.saturating_add(row.token_usage.completion_tokens)
);
}
#[derive(Clone, Copy)]
enum SecondRoundBehavior {
BlockForever,
TerminalError,
}
struct EvalAbortProbeProvider {
main_calls: std::sync::atomic::AtomicUsize,
gold_started: Arc<AtomicBool>,
gold_completed: Arc<AtomicBool>,
release: Arc<tokio::sync::Notify>,
finished: Arc<tokio::sync::Notify>,
second_round: SecondRoundBehavior,
}
#[async_trait::async_trait]
impl LLMProvider for EvalAbortProbeProvider {
async fn chat_stream(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
) -> Result<LLMStream, LLMError> {
Ok(Box::pin(stream::iter(vec![Ok(LLMChunk::Done)])))
}
async fn chat_stream_with_options(
&self,
_messages: &[Message],
_tools: &[bamboo_agent_core::tools::ToolSchema],
_max_output_tokens: Option<u32>,
_model: &str,
options: Option<&bamboo_llm::LLMRequestOptions>,
) -> Result<LLMStream, LLMError> {
let purpose = options
.and_then(|o| o.request_purpose.as_deref())
.unwrap_or("agent_loop");
if purpose == "gold_evaluation" {
self.gold_started.store(true, Ordering::SeqCst);
self.release.notified().await;
self.gold_completed.store(true, Ordering::SeqCst);
self.finished.notify_one();
let call = bamboo_agent_core::tools::ToolCall {
id: "gold-eval-async".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "report_gold_evaluation".to_string(),
arguments:
r#"{"decision":"achieved","confidence":"high","reasoning":"done"}"#
.to_string(),
},
};
return Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::Done),
])));
}
let n = self
.main_calls
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if n == 0 {
let call = bamboo_agent_core::tools::ToolCall {
id: "noop-1".to_string(),
tool_type: "function".to_string(),
function: bamboo_agent_core::tools::FunctionCall {
name: "noop".to_string(),
arguments: "{}".to_string(),
},
};
return Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![call])),
Ok(LLMChunk::Done),
])));
}
for _ in 0..2000 {
if self.gold_started.load(Ordering::SeqCst) {
break;
}
tokio::time::sleep(Duration::from_millis(2)).await;
}
match self.second_round {
SecondRoundBehavior::BlockForever => Ok(Box::pin(stream::pending())),
SecondRoundBehavior::TerminalError => Err(LLMError::Auth(
"terminal error injected to exercise #347 abort".to_string(),
)),
}
}
}
fn eval_abort_config() -> AgentLoopConfig {
use crate::runtime::config::PromptMemoryFlags;
AgentLoopConfig {
gold_config: Some(crate::runtime::config::GoldConfig {
enabled: true,
auto_continue_enabled: true,
goal: Some("ship it".to_string()),
max_auto_continuations: 3,
..crate::runtime::config::GoldConfig::default()
}),
prompt_memory_flags: PromptMemoryFlags {
project_prompt_injection: false,
relevant_recall: false,
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(5),
..Default::default()
},
..AgentLoopConfig::default()
}
}
#[tokio::test]
async fn cleanup_harvests_finished_gold_evaluation_without_false_cancellation() {
let session_id = "gold-eval-finished-before-cleanup";
let mut state = e2e_loop_state(session_id);
let checkpoint = bamboo_agent_core::GoldCheckpoint::PostRound;
let request = crate::runtime::gold_evaluation::AsyncGoldEvaluationRequest {
session_id: session_id.to_string(),
round_number: 1,
model_name: "fast".to_string(),
reasoning_effort: None,
checkpoint,
timeout_context: crate::runtime::stream::handler::StreamTimeoutContext::default(),
session_snapshot: Session::new(session_id, "model"),
task_context_snapshot: None,
gold_config: crate::runtime::config::GoldConfig::default(),
};
let result = crate::runtime::gold_evaluation::AsyncGoldEvaluationResult {
round_number: 1,
model_name: "fast".to_string(),
evaluation_result: crate::runtime::gold_evaluation::GoldEvaluationResult {
checkpoint,
iteration: 1,
decision: bamboo_agent_core::GoldDecision::Continue,
confidence: bamboo_agent_core::GoldConfidence::High,
reasoning: "complete".to_string(),
missing_information: Vec::new(),
next_action: None,
prompt_tokens: 1,
completion_tokens: 1,
},
};
let (tx, mut rx) = mpsc::channel(8);
tx.send(AgentEvent::GoldEvaluationStarted {
session_id: session_id.to_string(),
checkpoint,
iteration: 1,
})
.await
.expect("started event");
tx.send(AgentEvent::GoldEvaluationCompleted {
session_id: session_id.to_string(),
checkpoint,
iteration: 1,
decision: bamboo_agent_core::GoldDecision::Continue,
confidence: bamboo_agent_core::GoldConfidence::High,
reasoning: "complete".to_string(),
})
.await
.expect("completed event");
state.gold_evaluation.in_flight = Some(super::super::startup::InFlightGoldEvaluation {
request,
join_handle: tokio::spawn(async move { Some(result) }),
});
tokio::time::timeout(Duration::from_secs(1), async {
while !state
.gold_evaluation
.in_flight
.as_ref()
.expect("unharvested Gold handle")
.join_handle
.is_finished()
{
tokio::task::yield_now().await;
}
})
.await
.expect("Gold evaluator should finish before cleanup");
super::abort_in_flight_evaluations(&mut state, &tx, "run_completed").await;
assert!(state.gold_evaluation.in_flight.is_none());
let mut started = 0;
let mut completed = 0;
let mut cancelled = 0;
while let Ok(event) = rx.try_recv() {
match event {
AgentEvent::GoldEvaluationStarted { .. } => started += 1,
AgentEvent::GoldEvaluationCompleted { .. } => completed += 1,
AgentEvent::GoldEvaluationCancelled { .. } => cancelled += 1,
_ => {}
}
}
assert_eq!(started, 1);
assert_eq!(completed, 1);
assert_eq!(
cancelled, 0,
"a completed Gold evaluator must not receive a contradictory cancellation"
);
}
#[tokio::test]
async fn cancelled_run_does_not_complete_in_flight_gold_eval() {
let gold_started = Arc::new(AtomicBool::new(false));
let gold_completed = Arc::new(AtomicBool::new(false));
let release = Arc::new(tokio::sync::Notify::new());
let finished = Arc::new(tokio::sync::Notify::new());
let llm: Arc<dyn LLMProvider> = Arc::new(EvalAbortProbeProvider {
main_calls: std::sync::atomic::AtomicUsize::new(0),
gold_started: gold_started.clone(),
gold_completed: gold_completed.clone(),
release: release.clone(),
finished: finished.clone(),
second_round: SecondRoundBehavior::BlockForever,
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = eval_abort_config();
let mut session = Session::new("session-eval-cancel", "model");
let mut state = e2e_loop_state("session-eval-cancel");
let (tx, _rx) = tokio::sync::mpsc::channel(64);
let cancel = CancellationToken::new();
let driver_started = gold_started.clone();
let driver_token = cancel.clone();
let driver = tokio::spawn(async move {
for _ in 0..2000 {
if driver_started.load(Ordering::SeqCst) {
break;
}
tokio::time::sleep(Duration::from_millis(2)).await;
}
driver_token.cancel();
});
let result = tokio::time::timeout(
Duration::from_secs(5),
super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state),
)
.await
.expect("run_pipeline did not return within 5s after cancel");
let _ = driver.await;
assert!(
matches!(result, Err(AgentError::Cancelled)),
"cancelled run must return Cancelled, got {result:?}"
);
assert!(
gold_started.load(Ordering::SeqCst),
"the Gold eval must have been genuinely in flight (else nothing was tested)"
);
assert!(
state.gold_evaluation.in_flight.is_none(),
"the in-flight Gold eval slot must be cleared on the cancel early-exit"
);
release.notify_one();
let finished_within =
tokio::time::timeout(Duration::from_millis(500), finished.notified()).await;
assert!(
finished_within.is_err(),
"cancelled Gold eval kept running to completion (spend not stopped)"
);
assert!(
!gold_completed.load(Ordering::SeqCst),
"cancelled Gold eval must not complete its LLM request"
);
}
#[tokio::test]
async fn terminal_error_aborts_in_flight_gold_eval() {
let gold_started = Arc::new(AtomicBool::new(false));
let gold_completed = Arc::new(AtomicBool::new(false));
let release = Arc::new(tokio::sync::Notify::new());
let finished = Arc::new(tokio::sync::Notify::new());
let llm: Arc<dyn LLMProvider> = Arc::new(EvalAbortProbeProvider {
main_calls: std::sync::atomic::AtomicUsize::new(0),
gold_started: gold_started.clone(),
gold_completed: gold_completed.clone(),
release: release.clone(),
finished: finished.clone(),
second_round: SecondRoundBehavior::TerminalError,
});
let tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(AlwaysOkExecutor);
let config = eval_abort_config();
let mut session = Session::new("session-eval-terminal", "model");
let mut state = e2e_loop_state("session-eval-terminal");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let cancel = CancellationToken::new();
let result = tokio::time::timeout(
Duration::from_secs(5),
super::run_pipeline(&mut session, &tx, llm, tools, &cancel, &config, &mut state),
)
.await
.expect("run_pipeline did not return within 5s");
assert!(
matches!(result, Err(AgentError::LLM(_))),
"the injected terminal error must surface as Err(LLM), got {result:?}"
);
assert!(
!cancel.is_cancelled(),
"this test must NOT rely on cancellation — it isolates the abort path"
);
assert!(
gold_started.load(Ordering::SeqCst),
"the Gold eval must have been genuinely in flight (else nothing was tested)"
);
assert!(
state.gold_evaluation.in_flight.is_none(),
"the in-flight Gold eval slot must be aborted+cleared on the terminal early-exit"
);
let mut saw_cancelled = false;
while let Ok(event) = rx.try_recv() {
if matches!(
event,
AgentEvent::GoldEvaluationCancelled { ref reason, .. }
if reason == "terminal_error"
) {
saw_cancelled = true;
}
}
assert!(
saw_cancelled,
"an observed evaluation start must receive an explicit terminal cancellation event"
);
release.notify_one();
let finished_within =
tokio::time::timeout(Duration::from_millis(500), finished.notified()).await;
assert!(
finished_within.is_err(),
"in-flight Gold eval was detached, not aborted, on the terminal early-exit (#347)"
);
assert!(
!gold_completed.load(Ordering::SeqCst),
"aborted Gold eval must not complete its LLM request"
);
}
#[tokio::test]
async fn abort_helper_is_nonblocking_for_completion_and_suspension() {
let mut session = Session::new("session-eval-complete", "model");
let mut state = e2e_loop_state("session-eval-complete");
let (tx, mut rx) = tokio::sync::mpsc::channel(64);
let request = crate::runtime::gold_evaluation::AsyncGoldEvaluationRequest {
session_id: session.id.clone(),
round_number: 1,
model_name: "fast".to_string(),
reasoning_effort: None,
checkpoint: bamboo_agent_core::GoldCheckpoint::PostRound,
timeout_context: crate::runtime::stream::handler::StreamTimeoutContext::default(),
session_snapshot: session.clone(),
task_context_snapshot: None,
gold_config: crate::runtime::config::GoldConfig::default(),
};
state.gold_evaluation.in_flight = Some(super::super::startup::InFlightGoldEvaluation {
request: request.clone(),
join_handle: tokio::spawn(std::future::pending()),
});
state.gold_evaluation.queued_request = Some(request.clone());
tokio::time::timeout(
Duration::from_millis(100),
super::abort_in_flight_evaluations(&mut state, &tx, "run_completed"),
)
.await
.expect("normal finalization waited for the blocked Gold evaluation");
assert!(state.gold_evaluation.in_flight.is_none());
assert!(state.gold_evaluation.queued_request.is_none());
let mut saw_cancelled = false;
while let Ok(event) = rx.try_recv() {
if matches!(
event,
AgentEvent::GoldEvaluationCancelled { ref reason, .. }
if reason == "run_completed"
) {
saw_cancelled = true;
}
}
assert!(saw_cancelled);
state.gold_evaluation.in_flight = Some(super::super::startup::InFlightGoldEvaluation {
request: request.clone(),
join_handle: tokio::spawn(std::future::pending()),
});
tokio::time::timeout(
Duration::from_millis(100),
super::abort_in_flight_evaluations(&mut state, &tx, "run_suspended"),
)
.await
.expect("suspension finalization waited for the blocked Gold evaluation");
assert!(matches!(
rx.try_recv(),
Ok(AgentEvent::GoldEvaluationCancelled { reason, .. })
if reason == "run_suspended"
));
session.metadata.insert("verified".into(), "true".into());
}
struct RecordingGuardianSpawner {
child_id: String,
prompts: Arc<std::sync::Mutex<Vec<String>>>,
}
#[async_trait::async_trait]
impl GuardianSpawner for RecordingGuardianSpawner {
async fn spawn_guardian_review(
&self,
_parent_session: &Session,
review_prompt: String,
_model: String,
_disabled_tools: Option<std::collections::BTreeSet<String>>,
) -> Result<String, String> {
self.prompts.lock().unwrap().push(review_prompt);
Ok(self.child_id.clone())
}
}
fn guardian_only_config_with_recorder(
max_reviews: u32,
) -> (AgentLoopConfig, Arc<std::sync::Mutex<Vec<String>>>) {
let prompts = Arc::new(std::sync::Mutex::new(Vec::new()));
let spawner: Arc<dyn GuardianSpawner> = Arc::new(RecordingGuardianSpawner {
child_id: "guardian-child".to_string(),
prompts: prompts.clone(),
});
let config = AgentLoopConfig {
guardian_config: Some(GuardianConfig {
enabled: true,
model_name: Some("guardian-test-model".to_string()),
max_reviews,
}),
guardian_spawner: Some(spawner),
..Default::default()
};
(config, prompts)
}
fn guardian_and_gold_config_with_recorder(
max_reviews: u32,
) -> (
crate::runtime::config::AgentLoopConfig,
Arc<std::sync::Mutex<Vec<String>>>,
) {
let prompts = Arc::new(std::sync::Mutex::new(Vec::new()));
let spawner: Arc<dyn GuardianSpawner> = Arc::new(RecordingGuardianSpawner {
child_id: "guardian-child".to_string(),
prompts: prompts.clone(),
});
let config = crate::runtime::config::AgentLoopConfig {
gold_config: Some(crate::runtime::config::GoldConfig {
enabled: true,
auto_continue_enabled: true,
goal: Some("finish the task".to_string()),
max_auto_continuations: 3,
..crate::runtime::config::GoldConfig::default()
}),
guardian_config: Some(GuardianConfig {
enabled: true,
model_name: Some("guardian-test-model".to_string()),
max_reviews,
}),
guardian_spawner: Some(spawner),
..crate::runtime::config::AgentLoopConfig::default()
};
(config, prompts)
}
const GUARDIAN_FINAL_MESSAGE_HEADER: &str = "## Agent's final message";
#[test]
fn guardian_review_prompt_includes_final_assistant_content() {
let config = AgentLoopConfig::default();
let prompt = build_guardian_review_prompt(
&None,
&config,
Some("Final handoff: shipped the fix and ran the tests."),
);
assert!(prompt.contains(GUARDIAN_FINAL_MESSAGE_HEADER));
assert!(prompt.contains("Final handoff: shipped the fix and ran the tests."));
}
#[test]
fn guardian_review_prompt_omits_section_when_content_is_none() {
let config = AgentLoopConfig::default();
let prompt = build_guardian_review_prompt(&None, &config, None);
assert!(!prompt.contains(GUARDIAN_FINAL_MESSAGE_HEADER));
}
#[test]
fn guardian_review_prompt_omits_section_when_content_is_blank() {
let config = AgentLoopConfig::default();
let prompt = build_guardian_review_prompt(&None, &config, Some(" \n\t "));
assert!(!prompt.contains(GUARDIAN_FINAL_MESSAGE_HEADER));
}
#[tokio::test]
async fn guardian_only_review_context_includes_final_content_without_persisting_it() {
let mut session = Session::new("s400-guardian-only", "model");
let (config, prompts) = guardian_only_config_with_recorder(2);
let mut runtime_state = AgentRuntimeState::new("s400-guardian-only".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let final_text = "Final handoff: implemented the feature and verified with cargo test.";
let outcome = super::handle_no_tool_calls(
final_text.to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"s400-guardian-only",
&config,
&None,
"model",
1,
Arc::new(StubProvider),
)
.await
.unwrap();
assert!(outcome.should_break);
assert!(!outcome.sent_complete);
assert!(runtime_state.waiting_for_children.is_some());
assert!(
session.messages.is_empty(),
"the deferred final message must not be persisted into the transcript \
at the guardian suspend point, got {:?}",
session.messages
);
let recorded = prompts.lock().unwrap();
assert_eq!(recorded.len(), 1, "exactly one review was spawned");
assert!(
recorded[0].contains(GUARDIAN_FINAL_MESSAGE_HEADER),
"guardian review prompt must include the final-message section:\n{}",
recorded[0]
);
assert!(
recorded[0].contains(final_text),
"guardian review prompt must include the agent's actual final content:\n{}",
recorded[0]
);
}
#[tokio::test]
async fn goal_loop_active_final_content_not_duplicated_in_guardian_prompt() {
let mut session = Session::new("s400-goal-loop", "model");
let (config, prompts) = guardian_and_gold_config_with_recorder(2);
let mut goal = ensure_goal_state(&session, "finish the task");
goal.declare(GoalDeclaredStatus::Complete, 1);
write_goal_state(&mut session, goal);
let mut runtime_state = AgentRuntimeState::new("s400-goal-loop".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let final_text = "Done — shipped and verified.";
let outcome = super::handle_no_tool_calls(
final_text.to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"s400-goal-loop",
&config,
&None,
"model",
1,
Arc::new(ScriptedGoldProvider {
decision: "achieved",
confidence: "high",
}),
)
.await
.unwrap();
assert!(outcome.should_break);
assert!(!outcome.sent_complete);
assert!(runtime_state.waiting_for_children.is_some());
assert!(
session
.messages
.iter()
.any(|message| message.content == final_text),
"goal-loop path must add the final assistant message to the transcript"
);
let recorded = prompts.lock().unwrap();
assert_eq!(recorded.len(), 1, "exactly one review was spawned");
assert!(
!recorded[0].contains(GUARDIAN_FINAL_MESSAGE_HEADER),
"goal-loop case must not duplicate the final message into the guardian prompt \
(it is already in the forked transcript):\n{}",
recorded[0]
);
}
#[tokio::test]
async fn guardian_only_blank_final_content_adds_no_stray_context_block() {
let mut session = Session::new("s400-blank", "model");
let (config, prompts) = guardian_only_config_with_recorder(2);
let mut runtime_state = AgentRuntimeState::new("s400-blank".to_string());
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let outcome = super::handle_no_tool_calls(
" \n ".to_string(),
None,
None,
5,
5,
round_usage(),
&mut session,
&mut runtime_state,
&tx,
None,
"round-1",
"s400-blank",
&config,
&None,
"model",
1,
Arc::new(StubProvider),
)
.await
.unwrap();
assert!(outcome.should_break);
assert!(!outcome.sent_complete);
assert!(session.messages.is_empty());
let recorded = prompts.lock().unwrap();
assert_eq!(recorded.len(), 1);
assert!(
!recorded[0].contains(GUARDIAN_FINAL_MESSAGE_HEADER),
"blank final content must not add a stray context block:\n{}",
recorded[0]
);
}
#[tokio::test]
async fn goal_recovery_backoff_stops_without_waiting_for_the_deadline() {
let mut session = Session::new("goal-backoff-stop", "model");
session.add_message(Message::user("finish"));
let goal = crate::runtime::goal_state::ensure_goal_state(&session, "finish");
crate::runtime::goal_state::write_goal_state(&mut session, goal);
let policy = crate::runtime::goal_recovery::GoalRecoveryPolicy {
max_attempts: 3,
max_elapsed_seconds: 900,
};
let timeout = AgentError::StreamTimeout(bamboo_agent_core::StreamTimeoutError::new(
bamboo_agent_core::StreamTimeoutPhase::FirstSemantic,
std::time::Duration::from_secs(30),
None,
None,
std::time::Duration::ZERO,
None,
true,
));
assert_eq!(
crate::runtime::goal_recovery::reserve_retry(
&mut session,
&policy,
&timeout,
chrono::Utc::now().timestamp_millis()
),
Some(5000)
);
let config = crate::runtime::config::AgentLoopConfig {
gold_config: Some(crate::runtime::config::GoldConfig {
enabled: true,
auto_continue_enabled: true,
goal: Some("finish".into()),
recovery: policy,
..Default::default()
}),
..Default::default()
};
let cancel = tokio_util::sync::CancellationToken::new();
let trigger = cancel.clone();
tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
trigger.cancel();
});
let (tx, _rx) = tokio::sync::mpsc::channel(64);
let mut state = e2e_loop_state("goal-backoff-stop");
let result = tokio::time::timeout(
std::time::Duration::from_secs(1),
super::run_pipeline(
&mut session,
&tx,
Arc::new(StubProvider),
Arc::new(AlwaysOkExecutor),
&cancel,
&config,
&mut state,
),
)
.await
.expect("stop interrupts the persisted five-second backoff");
assert!(matches!(result, Err(AgentError::Cancelled)));
}
struct VisionLoopProvider {
support: Option<bool>,
path: String,
expect_description: bool,
alias: &'static str,
calls: AtomicUsize,
received_images: std::sync::Mutex<Vec<bool>>,
}
#[async_trait::async_trait]
impl LLMProvider for VisionLoopProvider {
async fn vision_support_override(&self, model: &str) -> Option<bool> {
assert_eq!(model, "actual-model");
self.support
}
async fn chat_stream(
&self,
messages: &[Message],
_: &[bamboo_agent_core::tools::ToolSchema],
_: Option<u32>,
model: &str,
) -> Result<LLMStream, LLMError> {
assert_eq!(model, "actual-model");
if self.support == Some(false) {
assert!(
messages.iter().all(|message| {
message.content_parts.as_ref().is_none_or(|parts| {
parts.iter().all(|part| {
!matches!(part, bamboo_domain::MessagePart::ImageUrl { .. })
})
})
}),
"non-Vision main never receives pixels from upload or tool output"
);
if self.path.starts_with("bamboo-") {
assert!(messages
.iter()
.any(|message| message.content.contains(&self.path)));
}
}
let n = self.calls.fetch_add(1, Ordering::SeqCst);
if n == 0 {
return Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::ToolCalls(vec![activation_call(
"vision-call",
self.alias,
&serde_json::json!({"path":self.path}).to_string(),
)])),
Ok(LLMChunk::Done),
])));
}
let result = messages
.iter()
.find(|m| m.tool_call_id.as_deref() == Some("vision-call"))
.expect("paired tool feedback reaches next model request");
let images = result.content_parts.as_ref().is_some_and(|parts| {
parts
.iter()
.any(|p| matches!(p, bamboo_domain::MessagePart::ImageUrl { .. }))
});
self.received_images.lock().unwrap().push(images);
if self.support == Some(false) {
assert!(!images);
if self.expect_description {
assert!(result.content.contains("Vision model output"));
assert!(result.content.contains("untrusted data"));
assert!(result.content.contains("visual fixture"));
} else {
assert!(result.content.contains("configured Vision model"));
assert!(result.content.contains("No image was sent"));
}
} else {
assert!(
images,
"real ViewImage result is delivered in the following request"
);
}
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("done".into())),
Ok(LLMChunk::Done),
])))
}
}
#[tokio::test]
async fn vision_real_viewimage_loop_delivers_direct_images_or_configuration_feedback() {
let image = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
std::fs::write(image.path(), b"\x89PNG\r\n\x1a\nvision-test").unwrap();
for support in [Some(true), Some(false), None] {
let provider = Arc::new(VisionLoopProvider {
support,
path: if support == Some(false) {
"/does-not-exist-vision-rejected-before-read.png".into()
} else {
image.path().to_string_lossy().into_owned()
},
expect_description: false,
alias: "default::ViewImage",
calls: AtomicUsize::new(0),
received_images: Default::default(),
});
let mut session =
Session::new(format!("vision-loop-{support:?}"), "stale-session-model");
session.add_message(Message::user("View the image"));
let mut state = e2e_loop_state(&session.id);
state.model_name = "actual-model".into();
let config = AgentLoopConfig {
model_name: Some("actual-model".into()),
image_fallback: (support == Some(true)).then_some(crate::ImageFallbackConfig {
mode: crate::ImageFallbackMode::Error,
vision_model: None,
}),
vision_model_resolver: if support == Some(false) {
None
} else {
Some(Arc::new(|| {
panic!("Vision-capable main must not invoke an auxiliary model")
}))
},
..delegation_loop_config()
};
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let tools = Arc::new(bamboo_tools::BuiltinToolExecutor::new());
assert!(super::run_pipeline(
&mut session,
&tx,
provider.clone(),
tools,
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state
)
.await
.unwrap());
assert_eq!(provider.calls.load(Ordering::SeqCst), 2);
assert_eq!(
*provider.received_images.lock().unwrap(),
vec![support != Some(false)]
);
}
}
struct AuxiliaryVisionMock(Arc<AtomicUsize>);
#[async_trait::async_trait]
impl LLMProvider for AuxiliaryVisionMock {
async fn chat_stream(
&self,
messages: &[Message],
tools: &[bamboo_domain::ToolSchema],
max: Option<u32>,
model: &str,
) -> Result<LLMStream, LLMError> {
assert_eq!(model, "configured-vision");
assert_eq!(max, Some(1024));
assert!(tools.is_empty());
assert!(messages[0].content.contains("untrusted"));
assert!(messages[1]
.content_parts
.as_ref()
.unwrap()
.iter()
.any(|part| matches!(part, bamboo_domain::MessagePart::ImageUrl { .. })));
self.0.fetch_add(1, Ordering::SeqCst);
Ok(Box::pin(stream::iter(vec![
Ok(LLMChunk::Token("visual fixture".into())),
Ok(LLMChunk::Done),
])))
}
}
#[tokio::test]
async fn vision_nonvision_upload_and_history_followup_use_only_requested_reference_fallback() {
use bamboo_agent_core::storage::Storage;
use base64::Engine;
for (inline, followup) in [(false, false), (false, true), (true, false), (true, true)] {
let home = tempfile::tempdir().unwrap();
let store = Arc::new(
bamboo_storage::SessionStoreV2::new(home.path().to_path_buf())
.await
.unwrap(),
);
let mut session = Session::new("vision-reference-loop", "stale-session-model");
store.save_session(&session).await.unwrap();
let png = b"\x89PNG\r\n\x1a\nvision-test";
let encoded = base64::engine::general_purpose::STANDARD.encode(png);
let (attachment_id, url) = if inline {
(None, format!("data:image/png;base64,{encoded}"))
} else {
let (id, url) = store
.write_image_attachment(&session, &encoded, Some("image/png"))
.await
.unwrap();
(Some(id), url)
};
let image = Message::user_with_parts(
"inspect this image",
vec![bamboo_domain::MessagePart::ImageUrl {
image_url: bamboo_domain::ImageUrlRef {
url: url.clone(),
detail: None,
},
}],
);
let image_id = image.id.clone();
let original = serde_json::to_vec(&image).unwrap();
let reference = crate::runtime::runner::image_fallback::conversation_image_reference(
&session.id,
&image.id,
0,
&url,
);
session.add_message(image);
if followup {
session.add_message(Message::assistant("earlier turn", None));
session.add_message(Message::user("inspect the earlier retained image"));
}
store.save_session(&session).await.unwrap();
let session_file = store
.sessions_root_dir()
.join(&session.id)
.join("session.json");
let original_record_bytes = tokio::fs::read(&session_file).await.unwrap();
let disk_before = serde_json::to_vec(
&store
.load_session(&session.id)
.await
.unwrap()
.unwrap()
.messages,
)
.unwrap();
let provider = Arc::new(VisionLoopProvider {
support: Some(false),
path: reference,
expect_description: true,
alias: "view_image",
calls: AtomicUsize::new(0),
received_images: Default::default(),
});
let vision_calls = Arc::new(AtomicUsize::new(0));
let vision: Arc<dyn LLMProvider> = Arc::new(AuxiliaryVisionMock(vision_calls.clone()));
let target_resolutions = Arc::new(AtomicUsize::new(0));
let resolutions = target_resolutions.clone();
let config = AgentLoopConfig {
model_name: Some("actual-model".into()),
attachment_reader: Some(store.clone()),
image_fallback: Some(crate::ImageFallbackConfig {
mode: crate::ImageFallbackMode::Vision,
vision_model: Some("configured-vision".into()),
}),
vision_model_resolver: Some(Arc::new(move || {
resolutions.fetch_add(1, Ordering::SeqCst);
let provider = vision.clone();
Box::pin(async move {
Ok(bamboo_llm::ResolvedModel::new(
provider,
"configured-vision",
))
})
})),
..delegation_loop_config()
};
crate::session_app::execute::validate_image_fallback_for_session(
&session,
config.image_fallback.as_ref(),
"actual-model",
Some(false),
)
.unwrap();
assert_eq!(provider.calls.load(Ordering::SeqCst), 0);
assert_eq!(target_resolutions.load(Ordering::SeqCst), 0);
assert_eq!(vision_calls.load(Ordering::SeqCst), 0);
let mut state = e2e_loop_state(&session.id);
state.model_name = "actual-model".into();
let (tx, _rx) = tokio::sync::mpsc::channel(256);
assert!(super::run_pipeline(
&mut session,
&tx,
provider.clone(),
Arc::new(bamboo_tools::BuiltinToolExecutor::new()),
&tokio_util::sync::CancellationToken::new(),
&config,
&mut state,
)
.await
.unwrap());
assert_eq!(provider.calls.load(Ordering::SeqCst), 2);
assert_eq!(target_resolutions.load(Ordering::SeqCst), 1);
assert_eq!(vision_calls.load(Ordering::SeqCst), 1);
assert_eq!(*provider.received_images.lock().unwrap(), vec![false]);
assert_eq!(session.model, "stale-session-model");
assert_eq!(
serde_json::to_vec(
session
.messages
.iter()
.find(|message| message.id == image_id)
.unwrap(),
)
.unwrap(),
original
);
assert_eq!(
serde_json::to_vec(
&store
.load_session(&session.id)
.await
.unwrap()
.unwrap()
.messages,
)
.unwrap(),
disk_before
);
assert_eq!(
tokio::fs::read(&session_file).await.unwrap(),
original_record_bytes,
"the durable session record bytes are unchanged",
);
if let Some(id) = attachment_id {
assert_eq!(
store
.read_attachment(&session.id, &id)
.await
.unwrap()
.unwrap()
.0
.as_slice(),
&png[..],
);
}
let result = session
.messages
.iter()
.find(|message| message.tool_call_id.as_deref() == Some("vision-call"))
.unwrap();
assert!(result.content_parts.is_none());
assert!(result.content.contains("Vision model output"));
assert!(result.content.contains("untrusted data"));
}
}
struct WaitingVisionMock {
started: Arc<tokio::sync::Notify>,
dropped: Arc<AtomicBool>,
}
struct WaitingVisionDrop(Arc<AtomicBool>);
impl Drop for WaitingVisionDrop {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
#[async_trait::async_trait]
impl LLMProvider for WaitingVisionMock {
async fn chat_stream(
&self,
_: &[Message],
_: &[bamboo_domain::ToolSchema],
_: Option<u32>,
_: &str,
) -> Result<LLMStream, LLMError> {
let guard = WaitingVisionDrop(self.dropped.clone());
self.started.notify_one();
futures::future::pending::<()>().await;
drop(guard);
unreachable!()
}
}
#[tokio::test]
async fn vision_pipeline_cancellation_drops_only_the_viewimage_fallback_request() {
let image = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
std::fs::write(image.path(), b"\x89PNG\r\n\x1a\nvision-test").unwrap();
let main = Arc::new(VisionLoopProvider {
support: Some(false),
path: image.path().to_string_lossy().into_owned(),
expect_description: true,
alias: "viewImage",
calls: AtomicUsize::new(0),
received_images: Default::default(),
});
let started = Arc::new(tokio::sync::Notify::new());
let dropped = Arc::new(AtomicBool::new(false));
let vision: Arc<dyn LLMProvider> = Arc::new(WaitingVisionMock {
started: started.clone(),
dropped: dropped.clone(),
});
let mut session = Session::new("cancel-viewimage-fallback", "actual-model");
session.add_message(Message::user("View the image"));
let mut state = e2e_loop_state(&session.id);
state.model_name = "actual-model".into();
let config = AgentLoopConfig {
model_name: Some("actual-model".into()),
vision_model_resolver: Some(Arc::new(move || {
let provider = vision.clone();
Box::pin(async move {
Ok(bamboo_llm::ResolvedModel::new(
provider,
"configured-vision",
))
})
})),
..delegation_loop_config()
};
let cancel = tokio_util::sync::CancellationToken::new();
let (tx, _rx) = tokio::sync::mpsc::channel(256);
let work = super::run_pipeline(
&mut session,
&tx,
main.clone(),
Arc::new(bamboo_tools::BuiltinToolExecutor::new()),
&cancel,
&config,
&mut state,
);
let (result, ()) = tokio::join!(work, async {
started.notified().await;
cancel.cancel();
});
assert!(matches!(result, Err(AgentError::Cancelled)));
assert!(dropped.load(Ordering::SeqCst));
assert_eq!(main.calls.load(Ordering::SeqCst), 1);
assert_eq!(session.model, "actual-model");
}
}
#[cfg(test)]
#[path = "pipeline_final_evidence_tests.rs"]
mod final_evidence_tests;
#[cfg(all(test, unix))]
#[path = "pipeline_exec_sessions_tests.rs"]
mod exec_sessions_tests;