use async_trait::async_trait;
use futures::stream::{Stream, StreamExt};
use parking_lot::RwLock;
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Weak};
use std::time::{Duration, Instant};
use tracing::{debug, error, info, instrument, warn};
const DISAMBIGUATION_STATE_GENERATION_KEY: &str = "_runtime.disambiguation_state_generation";
const MAX_TOOL_FALLBACK_HOPS: usize = 16;
pub(crate) type RootTurnGate = Arc<tokio::sync::Mutex<()>>;
pub(crate) type RootTurnGateIdentityStack = Arc<[RootTurnGate]>;
tokio::task_local! {
static RUNTIME_GATE_IDENTITY_STACK: RootTurnGateIdentityStack;
}
pub(crate) fn current_runtime_gate_identity_stack() -> RootTurnGateIdentityStack {
RUNTIME_GATE_IDENTITY_STACK
.try_with(Arc::clone)
.unwrap_or_default()
}
pub(crate) async fn scope_runtime_gate_identity_stack<F, T>(
identity_stack: &RootTurnGateIdentityStack,
future: F,
) -> T
where
F: Future<Output = T>,
{
RUNTIME_GATE_IDENTITY_STACK
.scope(Arc::clone(identity_stack), future)
.await
}
pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Weak<tokio::sync::Mutex<()>>>>>;
struct ToolResourceGuards {
guards: Vec<tokio::sync::OwnedMutexGuard<()>>,
locks: ToolResourceLocks,
}
struct RootTurnAdmission {
guard: tokio::sync::OwnedMutexGuard<()>,
identity_stack: RootTurnGateIdentityStack,
}
#[derive(Clone)]
struct StoredSessionRestore {
snapshot: AgentSnapshot,
metadata: Option<ai_agents_core::SessionMetadata>,
}
struct RuntimeSessionRestorePoint {
snapshot: AgentSnapshot,
metadata: ai_agents_core::SessionMetadata,
actor_id: Option<String>,
session_id: Option<String>,
}
impl Drop for ToolResourceGuards {
fn drop(&mut self) {
self.guards.clear();
self.locks.write().retain(|_, lock| lock.strong_count() > 0);
}
}
#[derive(Clone)]
struct RuntimeSafetySnapshot {
version: u64,
emergency_deny: bool,
tool_security: ToolSecurityEngine,
tool_scope_override: Option<Vec<String>>,
}
#[derive(Clone, Copy)]
struct ToolDecisionVersions {
policy: u64,
registry: u64,
runtime_control: u64,
state: Option<u64>,
}
#[derive(Clone, Debug, Default)]
struct ToolFallbackState {
visited_canonical_ids: Vec<String>,
}
impl ToolFallbackState {
fn rejection_reason(&self, canonical_id: &str) -> Option<String> {
if self
.visited_canonical_ids
.iter()
.any(|visited| visited == canonical_id)
{
return Some(format!(
"Tool fallback cycle detected at '{canonical_id}' after [{}]",
self.visited_canonical_ids.join(" -> ")
));
}
if self.visited_canonical_ids.len() > MAX_TOOL_FALLBACK_HOPS {
return Some(format!(
"Tool fallback chain exceeds the maximum of {MAX_TOOL_FALLBACK_HOPS} hops"
));
}
None
}
fn with_current(mut self, canonical_id: String) -> Self {
self.visited_canonical_ids.push(canonical_id);
self
}
fn final_rejection_reason(
&self,
admitted_canonical_id: &str,
final_canonical_id: &str,
) -> Option<String> {
if admitted_canonical_id == final_canonical_id {
return None;
}
if self
.visited_canonical_ids
.iter()
.any(|visited| visited == final_canonical_id)
{
return Some(format!(
"Tool fallback cycle detected after final resolution changed '{admitted_canonical_id}' to '{final_canonical_id}'"
));
}
Some(format!(
"Tool canonical target changed after initial admission from '{admitted_canonical_id}' to '{final_canonical_id}'"
))
}
}
#[derive(Clone, Copy, Debug)]
struct ValidatedToolTimeout {
timer: Duration,
deadline_delta: chrono::Duration,
}
struct AvailableToolIdsSnapshot {
tool_ids: Vec<String>,
state_generation: Option<u64>,
}
#[derive(Clone)]
struct ToolApprovalBinding {
canonical_id: String,
arguments: Value,
confirmation_required: bool,
policy_version: u64,
runtime_control_version: u64,
state_generation: Option<u64>,
reviewed_tool: Arc<dyn ai_agents_core::Tool>,
}
fn merge_approved_record(record: &mut Option<ToolApprovalRecord>) {
if record
.as_ref()
.is_some_and(|record| matches!(record.status, ToolApprovalStatus::Modified))
{
return;
}
*record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Approved,
reason: None,
modified_arguments: None,
});
}
impl ToolApprovalBinding {
fn is_stale(
&self,
canonical_id: &str,
arguments: &Value,
confirmation_required: bool,
versions: ToolDecisionVersions,
resolved_tool: &Arc<dyn ai_agents_core::Tool>,
) -> bool {
self.canonical_id != canonical_id
|| self.arguments != *arguments
|| self.confirmation_required != confirmation_required
|| self.policy_version != versions.policy
|| self.runtime_control_version != versions.runtime_control
|| self.state_generation != versions.state
|| !Arc::ptr_eq(&self.reviewed_tool, resolved_tool)
}
}
use crate::turn_context::{current_turn_actor_context, scope_actor_context};
use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
use ai_agents_core::traits::storage::StorageCapability;
use ai_agents_core::{
AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMError, LLMProvider,
LLMResponse, LLMToolDefinition, LLMToolRequest, PermissionOutcome, Result, ToolActorContext,
ToolApprovalRecord, ToolApprovalStatus, ToolCallClassification, ToolCallSource,
ToolCancellationToken, ToolChoice, ToolExecutionContext, ToolExecutionLimits,
ToolExecutionRecord, ToolExecutionRequest, ToolInvoker, ToolPolicyDecisionRecord, ToolResult,
ToolSafetyMetadata,
};
use ai_agents_disambiguation::{
ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
ConfirmationParseFuture, DisambiguationConfig, DisambiguationContext, DisambiguationManager,
DisambiguationResult,
};
use ai_agents_hitl::{
ApprovalHandler, ApprovalResolvedOutcome, ApprovalResult, ApprovalTrigger, HITLCheckResult,
HITLEngine, RejectAllHandler, TimeoutAction,
};
use ai_agents_hooks::{AgentHooks, NoopHooks};
use ai_agents_llm::LLMRegistry;
use ai_agents_memory::{
CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
};
use ai_agents_observability::{
EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
current_observation_context, new_session_id as new_observation_session_id,
resolve_language_from_context, with_observation_context, with_observation_purpose,
};
use ai_agents_process::{
ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
};
use ai_agents_reasoning::{
CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
ReflectionMetadata, StepFailureAction,
};
use ai_agents_recovery::{
ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
};
use ai_agents_relationships::RelationshipManager;
use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
use ai_agents_state::{
PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, Transition,
TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
};
use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
use ai_agents_tools::{
CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
MAX_TOOL_TIMEOUT_MS, QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord,
ToolRegistry, ToolSecurityConfig, ToolSecurityEngine,
};
use super::{
Agent, AgentInfo, AgentResponse, AgentStreamEvent, ParallelToolsConfig, StreamChunk,
StreamingConfig, ToolCall,
};
use crate::optimization::{
AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
};
use crate::spec::StorageConfig;
enum ToolCallOutcome {
Continue,
TransitionFired,
Rejected(AgentResponse),
}
#[derive(Clone)]
struct MainToolProtocol {
choice: Option<ToolChoice>,
tool_ids: Vec<String>,
definitions: Vec<LLMToolDefinition>,
}
struct MainProviderResponse {
response: LLMResponse,
used_native_tools: bool,
}
struct CommittedTextResponse<'a> {
processed_input: &'a str,
input_context: &'a HashMap<String, Value>,
answer: String,
reasoning_mode: ReasoningMode,
auto_detected: bool,
iterations: u32,
thinking_content: Option<String>,
all_tool_calls: Vec<ToolCall>,
}
struct AgentResponseParts {
content: String,
all_tool_calls: Vec<ToolCall>,
reasoning_mode: ReasoningMode,
auto_detected: bool,
iterations: u32,
thinking: Option<String>,
reflection_metadata: Option<ReflectionMetadata>,
}
type RuntimeStreamTerminalSlot = Arc<RwLock<Option<AgentResponse>>>;
fn new_runtime_stream_terminal_slot() -> RuntimeStreamTerminalSlot {
Arc::new(RwLock::new(None))
}
fn record_runtime_stream_final(slot: &RuntimeStreamTerminalSlot, response: AgentResponse) {
*slot.write() = Some(response);
}
#[derive(Clone, Copy)]
struct DisambiguationOwnership {
epoch: u64,
state_generation: Option<u64>,
}
enum SkillRouteResult {
NoMatch,
Response { skill_id: String, content: String },
NeedsClarification {
response: AgentResponse,
ownership: Option<DisambiguationOwnership>,
},
}
enum ParallelTransitionSelection {
Candidate(TransitionCandidate),
NoMatch,
ReservationExhausted,
}
enum PostLoopResult {
NoTransition(String),
Transitioned(String),
NeedsRedispatch,
}
struct StateTransitionReservation<'a> {
reserved: &'a AtomicBool,
}
impl Drop for StateTransitionReservation<'_> {
fn drop(&mut self) {
self.reserved.store(false, Ordering::SeqCst);
}
}
struct RootTurnCleanup<'a> {
agent: &'a RuntimeAgent,
}
impl<'a> RootTurnCleanup<'a> {
fn new(agent: &'a RuntimeAgent) -> Self {
Self { agent }
}
}
impl Drop for RootTurnCleanup<'_> {
fn drop(&mut self) {
self.agent.end_root_turn();
}
}
#[derive(Debug)]
struct RuntimeControlState {
snapshot_guard: RwLock<()>,
version: AtomicU64,
emergency_deny: Arc<AtomicBool>,
tool_security_override: RwLock<Option<ToolSecurityEngine>>,
tool_scope_override: RwLock<Option<Vec<String>>>,
}
impl Default for RuntimeControlState {
fn default() -> Self {
Self {
snapshot_guard: RwLock::new(()),
version: AtomicU64::new(1),
emergency_deny: Arc::new(AtomicBool::new(false)),
tool_security_override: RwLock::new(None),
tool_scope_override: RwLock::new(None),
}
}
}
#[derive(Clone)]
pub struct RuntimeControlHandle {
state: Arc<RuntimeControlState>,
}
impl RuntimeControlHandle {
pub fn version(&self) -> u64 {
self.state.version.load(Ordering::SeqCst)
}
fn bump(&self) -> u64 {
self.state.version.fetch_add(1, Ordering::SeqCst) + 1
}
pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
self.try_set_tool_security(config)
.expect("invalid tool security configuration")
}
pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
config.validate()?;
let _guard = self.state.snapshot_guard.write();
let generation = self.bump();
*self.state.tool_security_override.write() = Some(
ToolSecurityEngine::new_with_policy_version(config, generation),
);
Ok(generation)
}
pub fn clear_tool_security_override(&self) -> u64 {
let _guard = self.state.snapshot_guard.write();
*self.state.tool_security_override.write() = None;
self.bump()
}
pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
let _guard = self.state.snapshot_guard.write();
*self.state.tool_scope_override.write() = Some(tool_ids);
self.bump()
}
pub fn clear_tool_scope_override(&self) -> u64 {
let _guard = self.state.snapshot_guard.write();
*self.state.tool_scope_override.write() = None;
self.bump()
}
pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
let _guard = self.state.snapshot_guard.write();
self.state.emergency_deny.store(enabled, Ordering::SeqCst);
self.bump()
}
pub fn cancel_all(&self) -> u64 {
self.set_emergency_deny(true)
}
}
pub struct RuntimeAgent {
info: AgentInfo,
llm_registry: Arc<LLMRegistry>,
memory: Arc<dyn Memory>,
tools: Arc<ToolRegistry>,
skills: Vec<SkillDefinition>,
skill_router: Option<SkillRouter>,
skill_executor: Option<SkillExecutor>,
base_system_prompt: String,
max_iterations: u32,
iteration_count: RwLock<u32>,
max_context_tokens: u32,
memory_token_budget: Option<MemoryTokenBudget>,
recovery_manager: RecoveryManager,
tool_security: ToolSecurityEngine,
process_processor: Option<ProcessProcessor>,
message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
state_machine: Option<Arc<StateMachine>>,
transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
context_manager: Arc<ContextManager>,
template_renderer: TemplateRenderer,
tool_call_history: RwLock<Vec<ToolCallRecord>>,
parallel_tools: ParallelToolsConfig,
streaming: StreamingConfig,
hooks: Arc<dyn AgentHooks>,
hitl_engine: Option<HITLEngine>,
approval_handler: Arc<dyn ApprovalHandler>,
storage_config: StorageConfig,
storage: RwLock<Option<Arc<dyn AgentStorage>>>,
storage_init: tokio::sync::Mutex<()>,
reasoning_config: ReasoningConfig,
reflection_config: ReflectionConfig,
disambiguation_manager: Option<DisambiguationManager>,
disambiguation_epoch: AtomicU64,
disambiguation_admission: tokio::sync::RwLock<()>,
state_transition_reserved: AtomicBool,
persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
pending_skill_id: RwLock<Option<String>>,
current_plan: RwLock<Option<Plan>>,
declared_tool_ids: Option<Vec<String>>,
context_initialized: AtomicBool,
spawner: Option<Arc<crate::spawner::AgentSpawner>>,
spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
redispatch_depth: RwLock<u32>,
active_turn_context: RwLock<Option<TurnOptimizationContext>>,
root_user_message_committed: AtomicBool,
actor_id: RwLock<Option<String>>,
fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
messages_since_extraction: Arc<RwLock<usize>>,
actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
facts_config: Option<ai_agents_facts::FactsConfig>,
session_metadata: RwLock<ai_agents_core::SessionMetadata>,
current_session_id: RwLock<Option<String>>,
relationship_manager: Option<Arc<RelationshipManager>>,
observability_manager: Option<Arc<ObservabilityManager>>,
runtime_config: RuntimeConfig,
background_maintenance: Arc<BackgroundMaintenanceQueue>,
resource_locks: ToolResourceLocks,
runtime_control: Arc<RuntimeControlState>,
root_turn_gate: RootTurnGate,
}
impl std::fmt::Debug for RuntimeAgent {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RuntimeAgent")
.field("info", &self.info)
.field("base_system_prompt", &self.base_system_prompt)
.field("max_iterations", &self.max_iterations)
.field("skills_count", &self.skills.len())
.field("max_context_tokens", &self.max_context_tokens)
.field("has_state_machine", &self.state_machine.is_some())
.field("parallel_tools", &self.parallel_tools)
.field("streaming", &self.streaming)
.field("has_hooks", &true)
.field("has_hitl", &self.hitl_engine.is_some())
.field("storage_type", &self.storage_config.storage_type())
.field("reasoning_mode", &self.reasoning_config.mode)
.field("reflection_enabled", &self.reflection_config.enabled)
.field("declared_tool_ids", &self.declared_tool_ids)
.field("has_persona", &self.persona_manager.is_some())
.field("has_observability", &self.observability_manager.is_some())
.finish_non_exhaustive()
}
}
struct ObservabilityClarificationObserver;
impl ClarificationObserver for ObservabilityClarificationObserver {
fn observe_question<'a>(
&'a self,
future: ClarificationQuestionFuture<'a>,
) -> ClarificationQuestionFuture<'a> {
Box::pin(async move {
with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
})
}
fn observe_parse<'a>(
&'a self,
future: ClarificationParseFuture<'a>,
) -> ClarificationParseFuture<'a> {
Box::pin(async move {
with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
})
}
fn observe_confirmation_parse<'a>(
&'a self,
future: ConfirmationParseFuture<'a>,
) -> ConfirmationParseFuture<'a> {
Box::pin(async move {
with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
})
}
}
struct ObservabilityProcessStageObserver;
impl ProcessStageObserver for ObservabilityProcessStageObserver {
fn observe<'a>(
&'a self,
hint: ProcessPurposeHint,
future: ProcessStageFuture<'a>,
) -> ProcessStageFuture<'a> {
Box::pin(async move {
with_observation_purpose(observation_purpose_for_process(hint), future).await
})
}
}
struct RegistryLLMGetter {
registry: Arc<LLMRegistry>,
}
impl LLMGetter for RegistryLLMGetter {
fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
self.registry.get(alias).ok()
}
}
impl RuntimeAgent {
#[allow(clippy::too_many_arguments)]
pub fn new(
info: AgentInfo,
llm_registry: Arc<LLMRegistry>,
memory: Arc<dyn Memory>,
tools: Arc<ToolRegistry>,
skills: Vec<SkillDefinition>,
system_prompt: String,
max_iterations: u32,
) -> Self {
let (skill_router, skill_executor) = if !skills.is_empty() {
let router_llm = llm_registry.router().ok();
let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
(router, Some(executor))
} else {
(None, None)
};
let context_manager =
ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
Self {
info,
llm_registry,
memory,
tools,
skills,
skill_router,
skill_executor,
base_system_prompt: system_prompt,
max_iterations,
iteration_count: RwLock::new(0),
max_context_tokens: 128000,
memory_token_budget: None,
recovery_manager: RecoveryManager::default(),
tool_security: ToolSecurityEngine::default(),
process_processor: None,
message_filters: RwLock::new(HashMap::new()),
state_machine: None,
transition_evaluator: None,
context_manager: Arc::new(context_manager),
template_renderer: TemplateRenderer::new(),
tool_call_history: RwLock::new(Vec::new()),
parallel_tools: ParallelToolsConfig::default(),
streaming: StreamingConfig::default(),
hooks: Arc::new(NoopHooks),
hitl_engine: None,
approval_handler: Arc::new(RejectAllHandler::new()),
storage_config: StorageConfig::default(),
storage: RwLock::new(None),
storage_init: tokio::sync::Mutex::new(()),
reasoning_config: ReasoningConfig::default(),
reflection_config: ReflectionConfig::default(),
disambiguation_manager: None,
disambiguation_epoch: AtomicU64::new(0),
disambiguation_admission: tokio::sync::RwLock::new(()),
state_transition_reserved: AtomicBool::new(false),
persona_manager: None,
pending_skill_id: RwLock::new(None),
current_plan: RwLock::new(None),
declared_tool_ids: None,
context_initialized: AtomicBool::new(false),
spawner: None,
spawner_registry: None,
redispatch_depth: RwLock::new(0),
active_turn_context: RwLock::new(None),
root_user_message_committed: AtomicBool::new(false),
actor_id: RwLock::new(None),
fact_store: RwLock::new(None),
fact_extractor: RwLock::new(None),
actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
messages_since_extraction: Arc::new(RwLock::new(0)),
actor_memory_config: None,
facts_config: None,
session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
current_session_id: RwLock::new(None),
relationship_manager: None,
observability_manager: None,
runtime_config: RuntimeConfig::default(),
background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
resource_locks: new_tool_resource_locks(),
runtime_control: Arc::new(RuntimeControlState::default()),
root_turn_gate: Arc::new(tokio::sync::Mutex::new(())),
}
}
pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
self.declared_tool_ids = ids;
self
}
pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
self.storage_config = config;
self
}
pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
*self.storage.write() = Some(storage);
self
}
pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
self.resource_locks = locks;
self
}
pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
self.reasoning_config = config;
self
}
pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
self.reflection_config = config;
self
}
pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
self.relationship_manager = Some(manager);
self
}
pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
self.observability_manager = Some(manager);
self
}
pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
let max_tasks = config.optimization.post_turn.max_background_tasks;
self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
self.runtime_config = config;
self
}
pub fn runtime_config(&self) -> &RuntimeConfig {
&self.runtime_config
}
pub async fn flush_background_tasks(&self) -> Result<()> {
self.background_maintenance.flush_all().await
}
pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
self.background_maintenance.flush_scope(actor_id).await
}
pub async fn flush_background_tasks_for_purpose(
&self,
purpose: RuntimeTaskPurpose,
) -> Result<()> {
self.background_maintenance.flush_purpose(purpose).await
}
pub async fn flush_background_tasks_for_actor_purpose(
&self,
actor_id: &str,
purpose: RuntimeTaskPurpose,
) -> Result<()> {
self.background_maintenance
.flush_scope_purpose(actor_id, purpose)
.await
}
pub async fn shutdown_background_tasks(&self) -> Result<()> {
self.flush_background_tasks().await
}
pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
self.observability_manager.clone()
}
async fn export_observability_if_configured(&self) {
let Some(manager) = self.observability_manager.as_ref() else {
return;
};
let export = &manager.config().export;
if !export.write_report && !export.write_raw_events {
return;
}
if let Err(error) = manager.export().await {
warn!(error = %error, "Observability export failed");
}
}
pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
self.relationship_manager.clone()
}
fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
current_turn_actor_context()
}
fn effective_actor_id(&self) -> Option<String> {
self.current_turn_actor_context()
.and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
.or_else(|| self.actor_id.read().clone())
}
fn effective_origin_actor_id(&self) -> Option<String> {
self.current_turn_actor_context()
.and_then(|ctx| ctx.origin_actor_id.clone())
.or_else(|| self.actor_id.read().clone())
}
fn record_session_actor_if_needed(&self) {
if let Some(actor_id) = self.effective_origin_actor_id() {
let mut meta = self.session_metadata.write();
meta.actor_id = Some(actor_id.clone());
if !meta.actors.iter().any(|a| a == &actor_id) {
meta.actors.push(actor_id);
}
}
}
fn outbound_actor_context(&self) -> crate::TurnActorContext {
let mut context = self.current_turn_actor_context().unwrap_or_default();
if context.origin_actor_id.is_none() {
context.origin_actor_id = self.effective_origin_actor_id();
}
context.sender_agent_id = Some(self.info.id.clone());
context
}
fn observation_session_id(&self) -> Option<String> {
let mut current = self.current_session_id.write();
if current.is_none() {
*current = Some(new_observation_session_id());
}
current.clone()
}
fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
let manager = self.observability_manager.as_ref()?;
let context = self.build_context_with_overlays();
let language = resolve_language_from_context(manager.config(), &context);
let context = current_observation_context()
.map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
.unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
Some(
context
.with_actor(actor_id.or_else(|| self.effective_actor_id()))
.with_session(self.observation_session_id())
.with_state(self.current_state())
.with_language(Some(language)),
)
}
fn current_runtime_observation_context(
&self,
purpose: ObservationPurpose,
) -> Option<SpanContext> {
let manager = self.observability_manager.as_ref()?;
let context = self.build_context_with_overlays();
let language = resolve_language_from_context(manager.config(), &context);
let mut observation = current_observation_context()
.unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
observation.agent_id = self.info.id.clone();
observation.actor_id = self.effective_actor_id();
observation.session_id = self.observation_session_id();
observation.state = self.current_state();
observation.language = Some(language);
observation.purpose = purpose;
Some(observation)
}
async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
where
F: Future<Output = T>,
{
if let Some(context) = self.current_runtime_observation_context(purpose) {
with_observation_context(context, future).await
} else {
future.await
}
}
fn chat_with_actor_context_boxed<'a>(
&'a self,
input: &'a str,
actor_context: crate::TurnActorContext,
) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
Box::pin(async move {
let RootTurnAdmission {
guard,
identity_stack,
} = self.acquire_root_turn().await?;
let result = scope_runtime_gate_identity_stack(&identity_stack, async move {
let actor_id = actor_context.effective_actor_id().map(str::to_string);
let run = async move {
scope_actor_context(
actor_context,
Box::pin(async move { self.run_loop(input).await }),
)
.await
};
let result = if let Some(context) = self.build_observation_context(actor_id) {
with_observation_context(context, run).await
} else {
run.await
};
self.export_observability_if_configured().await;
result
})
.await;
drop(guard);
result
})
}
async fn acquire_root_turn(&self) -> Result<RootTurnAdmission> {
let gate_identity = Arc::clone(&self.root_turn_gate);
let current_identity_stack = current_runtime_gate_identity_stack();
if current_identity_stack
.iter()
.any(|owned_gate| Arc::ptr_eq(owned_gate, &gate_identity))
{
return Err(AgentError::Other(format!(
"RuntimeAgent '{}' rejected reentrant root turn ownership",
self.info.id
)));
}
let guard = Arc::clone(&gate_identity).lock_owned().await;
let mut identity_stack = Vec::with_capacity(current_identity_stack.len() + 1);
identity_stack.extend(current_identity_stack.iter().cloned());
identity_stack.push(gate_identity);
Ok(RootTurnAdmission {
guard,
identity_stack: identity_stack.into(),
})
}
pub async fn chat_with_actor_context(
&self,
input: &str,
actor_context: crate::TurnActorContext,
) -> Result<AgentResponse> {
self.chat_with_actor_context_boxed(input, actor_context)
.await
}
pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
self.chat_with_actor_context(input, actor_context).await
}
pub async fn load_actor_relationship(&self) -> Result<()> {
self.maybe_load_actor_relationship().await;
Ok(())
}
pub async fn update_relationship_dimension(
&self,
dimension: &str,
delta: f64,
reason: Option<&str>,
) -> Result<ai_agents_relationships::DimensionChange> {
self.update_relationship_dimension_for_perspective(
ai_agents_relationships::RelationshipPerspective::AgentToActor,
dimension,
delta,
reason,
)
.await
}
pub async fn update_relationship_dimension_for_perspective(
&self,
perspective: ai_agents_relationships::RelationshipPerspective,
dimension: &str,
delta: f64,
reason: Option<&str>,
) -> Result<ai_agents_relationships::DimensionChange> {
let manager = self
.relationship_manager
.as_ref()
.ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
let actor_id = self.effective_actor_id().ok_or_else(|| {
AgentError::Config("No actor ID set. Use set_actor_id() first".into())
})?;
let change = manager.update_dimension_for_perspective(
&actor_id,
perspective,
dimension,
delta,
1.0,
reason.unwrap_or("manual relationship update"),
)?;
self.persist_actor_relationship(&actor_id).await?;
info!(
actor_id = %actor_id,
perspective = %change.perspective,
dimension = %change.dimension,
delta = change.delta,
current = change.current,
"relationship updated manually"
);
self.hooks
.on_relationship_change(&actor_id, std::slice::from_ref(&change))
.await;
Ok(change)
}
pub fn reasoning_config(&self) -> &ReasoningConfig {
&self.reasoning_config
}
pub fn reflection_config(&self) -> &ReflectionConfig {
&self.reflection_config
}
pub fn with_facts_config(
mut self,
actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
facts_config: Option<ai_agents_facts::FactsConfig>,
) -> Self {
self.actor_memory_config = actor_memory_config;
self.facts_config = facts_config;
self
}
pub fn with_facts(
mut self,
store: Arc<ai_agents_facts::FactStore>,
extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
facts_config: Option<ai_agents_facts::FactsConfig>,
) -> Self {
*self.fact_store.write() = Some(store);
*self.fact_extractor.write() = extractor;
self.actor_memory_config = actor_memory_config;
self.facts_config = facts_config;
self
}
pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
self.fact_store.read().clone()
}
pub fn actor_id(&self) -> Option<String> {
self.actor_id.read().clone()
}
pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
*self.actor_id.write() = Some(actor_id.to_string());
{
let mut meta = self.session_metadata.write();
meta.actor_id = Some(actor_id.to_string());
if !meta.actors.iter().any(|a| a == actor_id) {
meta.actors.push(actor_id.to_string());
}
}
Ok(())
}
pub fn clear_actor_id(&self) {
*self.actor_id.write() = None;
self.session_metadata.write().actor_id = None;
}
pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
self.set_actor_id(user_id)
}
pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
let actor_id = match self.effective_actor_id() {
Some(id) => id,
None => return Ok(()),
};
let store_opt = self.fact_store.read().clone();
if let Some(store) = store_opt {
let facts = store.get_facts(&actor_id).await?;
let count = facts.len();
self.actor_facts_cache
.write()
.insert(actor_id.clone(), facts);
self.hooks.on_actor_memory_loaded(&actor_id, count).await;
tracing::debug!("loaded {} facts for actor {}", count, actor_id);
}
Ok(())
}
async fn maybe_load_actor_memory(&self) {
let Some(actor_id) = self.effective_actor_id() else {
return;
};
if self.actor_facts_cache.read().contains_key(&actor_id) {
return;
}
let _ = self.load_actor_memory().await;
}
async fn pre_turn_session_lifecycle(&self) {
if *self.redispatch_depth.read() > 0 {
return;
}
self.resolve_actor_id_from_context();
self.await_background_before_next_turn().await;
self.record_session_actor_if_needed();
self.maybe_load_actor_memory().await;
self.maybe_load_actor_relationship().await;
*self.messages_since_extraction.write() += 1;
}
async fn post_turn_session_lifecycle(&self) -> Result<()> {
if *self.redispatch_depth.read() > 0 {
return Ok(());
}
*self.messages_since_extraction.write() += 1;
self.run_post_turn_maintenance().await
}
fn begin_root_turn(&self) {
if *self.redispatch_depth.read() == 0 {
let mut guard = self.active_turn_context.write();
if guard.is_none() {
self.root_user_message_committed
.store(false, Ordering::SeqCst);
let max_calls = self
.runtime_config
.optimization
.max_speculative_llm_calls_per_turn;
*guard = Some(TurnOptimizationContext::new(
String::new(),
HashMap::new(),
max_calls,
));
}
}
}
fn update_active_turn_context(
&self,
processed_input: &str,
input_context: HashMap<String, Value>,
) {
if *self.redispatch_depth.read() > 0 {
return;
}
let max_calls = self
.runtime_config
.optimization
.max_speculative_llm_calls_per_turn;
let mut guard = self.active_turn_context.write();
match guard.as_mut() {
Some(context) => {
context.processed_input = processed_input.to_string();
context.input_context = input_context;
context.max_speculative_llm_calls = max_calls;
}
None => {
*guard = Some(TurnOptimizationContext::new(
processed_input,
input_context,
max_calls,
));
}
}
}
async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
if *self.redispatch_depth.read() > 0 {
return Ok(());
}
if !self
.root_user_message_committed
.swap(true, Ordering::SeqCst)
{
self.memory
.add_message(ChatMessage::user(processed_input))
.await?;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.mark_user_message_committed();
}
}
Ok(())
}
fn end_root_turn(&self) {
if *self.redispatch_depth.read() == 0 {
self.root_user_message_committed
.store(false, Ordering::SeqCst);
*self.active_turn_context.write() = None;
}
}
fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
self.begin_root_turn();
let mut guard = self.active_turn_context.write();
let Some(context) = guard.as_mut() else {
return false;
};
context.reserve_speculative_llm_call_for(kind)
}
fn branch_context_preview(&self) -> String {
let context = self.build_context_with_overlays();
let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
value = value
.chars()
.take(MAX_CONTEXT_PREVIEW_CHARS)
.collect::<String>();
value.push_str("...");
}
value
}
async fn await_background_before_next_turn(&self) {
let optimization = &self.runtime_config.optimization;
if !optimization.enabled {
return;
}
let actor_id = self.effective_actor_id();
let post = &optimization.post_turn;
self.await_background_task(
post.facts.await_before_next_turn,
RuntimeTaskPurpose::PostTurnFacts,
actor_id.as_deref(),
"facts",
)
.await;
self.await_background_task(
post.relationships.await_before_next_turn,
RuntimeTaskPurpose::PostTurnRelationship,
actor_id.as_deref(),
"relationships",
)
.await;
}
async fn await_background_task(
&self,
policy: AwaitBeforeNextTurn,
purpose: RuntimeTaskPurpose,
actor_id: Option<&str>,
label: &str,
) {
match policy {
AwaitBeforeNextTurn::Never => {}
AwaitBeforeNextTurn::Always => {
if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
warn!(label = label, error = %error, "background maintenance flush failed");
}
}
AwaitBeforeNextTurn::SameActor => {
if let Some(actor_id) = actor_id
&& let Err(error) = self
.flush_background_tasks_for_actor_purpose(actor_id, purpose)
.await
{
warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
}
}
}
}
async fn run_post_turn_maintenance(&self) -> Result<()> {
let optimization = &self.runtime_config.optimization;
if !optimization.enabled {
self.auto_extract_facts().await;
self.auto_update_relationship().await;
return Ok(());
}
let facts_mode = effective_maintenance_mode(
optimization.post_turn.facts.mode,
optimization.parallel_post_turn_memory,
);
let relationships_mode = effective_maintenance_mode(
optimization.post_turn.relationships.mode,
optimization.parallel_post_turn_memory,
);
match (facts_mode, relationships_mode) {
(MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
self.auto_extract_facts().await;
self.auto_update_relationship().await;
}
(MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
let facts = self.auto_extract_facts();
let relationships = self.auto_update_relationship();
tokio::join!(facts, relationships);
}
(MaintenanceMode::Background, MaintenanceMode::Background) => {
self.schedule_facts_background().await?;
self.schedule_relationship_background().await?;
}
(MaintenanceMode::Background, MaintenanceMode::InlineParallel)
| (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
self.schedule_facts_background().await?;
self.auto_update_relationship().await;
}
(MaintenanceMode::InlineParallel, MaintenanceMode::Background)
| (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
self.auto_extract_facts().await;
self.schedule_relationship_background().await?;
}
_ => {
self.auto_extract_facts().await;
self.auto_update_relationship().await;
}
}
Ok(())
}
async fn schedule_facts_background(&self) -> Result<()> {
let policy = self.runtime_config.optimization.post_turn.facts.clone();
let should_extract = self
.facts_config
.as_ref()
.map(|c| c.enabled && c.auto_extract)
.unwrap_or(false);
if !should_extract {
return Ok(());
}
let msgs_since = *self.messages_since_extraction.read();
if msgs_since < 2 {
return Ok(());
}
let Some(actor_id) = self.effective_actor_id() else {
self.record_skipped_maintenance(
"facts",
ObservationPurpose::FactsExtraction,
"missing_actor",
Some(&policy),
);
return Ok(());
};
let Some(extractor) = self.fact_extractor.read().clone() else {
return Ok(());
};
let messages = match self.memory.get_messages(None).await {
Ok(messages) => messages,
Err(error) => {
warn!(error = %error, "failed to snapshot messages for fact extraction");
return Ok(());
}
};
let recent: Vec<_> = messages
.iter()
.rev()
.take(msgs_since)
.rev()
.cloned()
.collect();
if recent.is_empty() {
return Ok(());
}
let existing = self
.actor_facts_cache
.read()
.get(&actor_id)
.cloned()
.unwrap_or_default();
let categories = self
.facts_config
.as_ref()
.map(|c| c.custom_categories.clone())
.unwrap_or_default();
let store = self.fact_store.read().clone();
let cache = Arc::clone(&self.actor_facts_cache);
let counter = Arc::clone(&self.messages_since_extraction);
let hooks = Arc::clone(&self.hooks);
let agent_id = self.info.id.clone();
let observation = current_observation_context();
let key = MaintenanceSequenceKey::actor(
agent_id,
actor_id.clone(),
RuntimeTaskPurpose::PostTurnFacts,
);
let actor_for_task = actor_id.clone();
let task = async move {
let run = async move {
let facts = extractor
.extract(&recent, &existing, Some(&actor_for_task), &categories)
.await?;
if !facts.is_empty() {
if let Some(store) = store {
let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
cache.write().insert(actor_for_task.clone(), authoritative);
} else {
cache
.write()
.entry(actor_for_task.clone())
.or_default()
.extend(facts.clone());
}
{
let mut count = counter.write();
if *count <= msgs_since {
*count = 0;
} else {
*count -= msgs_since;
}
}
hooks.on_facts_extracted(&actor_for_task, &facts).await;
}
Ok(())
};
if let Some(context) = observation {
with_observation_context(
context.with_purpose(ObservationPurpose::FactsExtraction),
run,
)
.await
} else {
run.await
}
};
self.spawn_or_handle_background(Some(key), task, "facts", &policy)
.await
}
async fn schedule_relationship_background(&self) -> Result<()> {
let policy = self
.runtime_config
.optimization
.post_turn
.relationships
.clone();
let Some(manager) = self.relationship_manager.as_ref().cloned() else {
return Ok(());
};
let Some(actor_id) = self.effective_actor_id() else {
self.record_skipped_maintenance(
"relationships",
ObservationPurpose::RelationshipUpdate,
"missing_actor",
Some(&policy),
);
return Ok(());
};
let recent_messages = manager.config().auto_update.recent_messages;
let messages = match self.memory.get_messages(Some(recent_messages)).await {
Ok(messages) => messages,
Err(error) => {
warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
return Ok(());
}
};
let storage = self.storage.read().clone();
let hooks = Arc::clone(&self.hooks);
let agent_id = self.info.id.clone();
let observation = current_observation_context();
let key = MaintenanceSequenceKey::actor(
agent_id.clone(),
actor_id.clone(),
RuntimeTaskPurpose::PostTurnRelationship,
);
let actor_for_task = actor_id.clone();
let task = async move {
let run = async move {
if manager.config().auto_update.enabled {
let update = manager.auto_update(&actor_for_task, &messages).await?;
if !update.changes.is_empty() {
hooks
.on_relationship_change(&actor_for_task, &update.changes)
.await;
}
if let Some(ref event) = update.event {
hooks.on_notable_event(&actor_for_task, event).await;
}
}
if manager.config().persistence.enabled
&& let (Some(storage), Some(value)) =
(storage, manager.relationship_as_value(&actor_for_task)?)
{
storage
.save_relationship(&agent_id, &actor_for_task, &value)
.await?;
}
Ok(())
};
if let Some(context) = observation {
with_observation_context(
context.with_purpose(ObservationPurpose::RelationshipUpdate),
run,
)
.await
} else {
run.await
}
};
self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
.await
}
async fn spawn_or_handle_background<F>(
&self,
key: Option<MaintenanceSequenceKey>,
task: F,
label: &'static str,
policy: &crate::optimization::config::MaintenanceTaskPolicy,
) -> Result<()>
where
F: Future<Output = Result<()>> + Send + 'static,
{
if self.background_maintenance.is_full() {
match self
.runtime_config
.optimization
.post_turn
.on_background_overflow
{
BackgroundOverflowPolicy::RunInline => {
record_background_maintenance_event(
self.observability_manager.as_ref(),
label,
EventStatus::Success,
0,
"inline_overflow",
None,
Some(policy),
);
let start = Instant::now();
match task.await {
Ok(()) => record_background_maintenance_event(
self.observability_manager.as_ref(),
label,
EventStatus::Success,
start.elapsed().as_millis() as u64,
"inline_completed",
None,
Some(policy),
),
Err(error) => {
warn!(label = label, error = %error, "inline maintenance fallback failed");
record_background_maintenance_event(
self.observability_manager.as_ref(),
label,
EventStatus::Error,
start.elapsed().as_millis() as u64,
"inline_failed",
Some(error.to_string()),
Some(policy),
);
return Err(error);
}
}
}
BackgroundOverflowPolicy::Drop => {
self.record_skipped_maintenance(
label,
ObservationPurpose::Other(label.to_string()),
"queue_full",
Some(policy),
);
}
BackgroundOverflowPolicy::Error => {
record_background_maintenance_event(
self.observability_manager.as_ref(),
label,
EventStatus::Error,
0,
"queue_full",
None,
Some(policy),
);
warn!(label = label, "background maintenance queue full");
return Err(AgentError::Other(format!(
"background maintenance queue is full for {}",
label
)));
}
}
return Ok(());
}
record_background_maintenance_event(
self.observability_manager.as_ref(),
label,
EventStatus::Success,
0,
"scheduled",
None,
Some(policy),
);
let manager = self.observability_manager.clone();
let policy_for_task = policy.clone();
let observed_task = async move {
let start = Instant::now();
let result = task.await;
match &result {
Ok(()) => record_background_maintenance_event(
manager.as_ref(),
label,
EventStatus::Success,
start.elapsed().as_millis() as u64,
"completed",
None,
Some(&policy_for_task),
),
Err(error) => record_background_maintenance_event(
manager.as_ref(),
label,
EventStatus::Error,
start.elapsed().as_millis() as u64,
"failed",
Some(error.to_string()),
Some(&policy_for_task),
),
}
result
};
if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
record_background_maintenance_event(
self.observability_manager.as_ref(),
label,
EventStatus::Error,
0,
"spawn_failed",
Some(error.to_string()),
Some(policy),
);
warn!(label = label, error = %error, "background maintenance spawn failed");
return Err(error);
}
Ok(())
}
fn record_skipped_maintenance(
&self,
label: &str,
purpose: ObservationPurpose,
reason: &str,
policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
) {
if let Some(manager) = self.observability_manager.as_ref() {
let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
tags.insert("runtime.skip_reason".to_string(), reason.to_string());
manager.record_lifecycle_event(
EventType::MemoryOperation {
operation: format!("{}_maintenance", label),
},
purpose,
EventStatus::Skipped,
0,
tags,
None,
);
}
}
pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
let Some(actor_id) = self.effective_actor_id() else {
return Vec::new();
};
self.actor_facts_cache
.read()
.get(&actor_id)
.cloned()
.unwrap_or_default()
}
pub fn relationship_memory_text(&self) -> Option<String> {
self.format_relationship_for_context().map(|(_, text)| text)
}
pub async fn extract_facts(
&self,
last_n: usize,
) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
self.extract_facts_with_source(last_n, "manual").await
}
async fn extract_facts_with_source(
&self,
last_n: usize,
source: &'static str,
) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
let extractor = match self.fact_extractor.read().clone() {
Some(e) => e,
None => return Ok(vec![]),
};
let messages = self.memory.get_messages(None).await?;
let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
if recent.is_empty() {
return Ok(vec![]);
}
let actor_id = self.effective_actor_id();
let existing = actor_id
.as_ref()
.and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
.unwrap_or_default();
let categories = self
.facts_config
.as_ref()
.map(|c| c.custom_categories.clone())
.unwrap_or_default();
let facts = self
.observe_purpose(
ObservationPurpose::FactsExtraction,
extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
)
.await?;
if !facts.is_empty() {
let fact_store_opt = self.fact_store.read().clone();
let mut stored_total = 0usize;
let mut cache_updated = false;
if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
let authoritative = store.add_facts(aid, facts.clone()).await?;
stored_total = authoritative.len();
self.actor_facts_cache
.write()
.insert(aid.clone(), authoritative);
cache_updated = true;
} else if let Some(aid) = &actor_id {
let mut cache = self.actor_facts_cache.write();
let entry = cache.entry(aid.clone()).or_default();
entry.extend(facts.clone());
stored_total = entry.len();
cache_updated = true;
}
info!(
actor_id = %actor_id.as_deref().unwrap_or("<none>"),
source = source,
requested_messages = last_n,
message_count = recent.len(),
extracted_count = facts.len(),
cache_updated = cache_updated,
stored_total = stored_total,
"facts extracted"
);
if let Some(ref aid) = actor_id {
self.hooks.on_facts_extracted(aid, &facts).await;
}
}
Ok(facts)
}
fn resolve_actor_id_from_context(&self) {
if self
.current_turn_actor_context()
.and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
.is_some()
{
return;
}
if let Some(ref am_config) = self.actor_memory_config
&& am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
&& let Some(ref path) = am_config.identification.context_path
{
let val = self
.context_manager
.get_path(path)
.or_else(|| self.context_manager.get(path));
if let Some(val) = val
&& let Some(id_str) = val.as_str()
{
let current = self.actor_id.read().clone();
if current.as_deref() != Some(id_str) {
*self.actor_id.write() = Some(id_str.to_string());
let mut meta = self.session_metadata.write();
meta.actor_id = Some(id_str.to_string());
if !meta.actors.iter().any(|a| a == id_str) {
meta.actors.push(id_str.to_string());
}
}
}
}
}
fn format_actor_facts_for_context(&self) -> String {
let should_inject = self
.facts_config
.as_ref()
.map(|c| c.inject_in_context)
.unwrap_or(true);
if !should_inject {
return String::new();
}
let Some(actor_id) = self.effective_actor_id() else {
return String::new();
};
let facts = self
.actor_facts_cache
.read()
.get(&actor_id)
.cloned()
.unwrap_or_default();
if facts.is_empty() {
return String::new();
}
let am_config = self.actor_memory_config.as_ref();
let facts_budget = self
.memory_token_budget
.as_ref()
.map(|b| b.allocation.facts as usize)
.filter(|n| *n > 0);
let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
let max_tokens = facts_budget.unwrap_or(default_max);
let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
return String::new();
}
if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
&& !cfg.injection.categories.is_empty()
{
facts
.iter()
.filter(|f| {
cfg.injection
.categories
.iter()
.any(|c| f.category.to_string() == *c)
})
.cloned()
.collect()
} else {
facts.clone()
}
} else {
facts.clone()
};
if filtered.is_empty() {
return String::new();
}
if let Some(store) = self.fact_store.read().clone() {
store.format_for_context(&filtered, max_tokens)
} else {
String::new()
}
}
fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
let context = self.build_context_with_overlays();
let mut root = Value::Object(context.into_iter().collect());
for (path, value) in staged {
if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
root = updated;
}
}
match root {
Value::Object(obj) => obj.into_iter().collect(),
_ => HashMap::new(),
}
}
fn build_context_with_overlays(&self) -> HashMap<String, Value> {
let mut context = self.context_manager.get_all();
let mut root = Value::Object(context.clone().into_iter().collect());
if let Some(turn_ctx) = self.current_turn_actor_context() {
if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
&& let Ok(updated) = ai_agents_core::set_dot_path(
root.clone(),
"interaction.origin_actor_id",
serde_json::json!(origin_actor_id),
)
{
root = updated;
}
if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
&& let Ok(updated) = ai_agents_core::set_dot_path(
root.clone(),
"interaction.sender_agent_id",
serde_json::json!(sender_agent_id),
)
{
root = updated;
}
}
if let Some(ref actor_id) = self.effective_actor_id()
&& let Ok(updated) = ai_agents_core::set_dot_path(
root.clone(),
"interaction.actor_id",
serde_json::json!(actor_id),
)
{
root = updated;
}
if let Some(manager) = self.relationship_manager.as_ref()
&& let Some(actor_id) = self.effective_actor_id()
&& let Some(value) = manager.to_context_value(&actor_id)
&& let Ok(updated) = ai_agents_core::set_dot_path(
root.clone(),
&manager.config().injection.context_path,
value,
)
{
root = updated;
}
if let Value::Object(obj) = root {
context = obj.into_iter().collect();
}
context
}
fn resolve_actor_name_from_context(&self) -> Option<String> {
for path in ["actor.name", "user.name", "player.name", "customer.name"] {
if let Some(value) = self.context_manager.get_path(path)
&& let Some(name) = value.as_str()
{
return Some(name.to_string());
}
}
None
}
async fn maybe_load_actor_relationship(&self) {
let Some(manager) = self.relationship_manager.as_ref() else {
return;
};
let Some(actor_id) = self.effective_actor_id() else {
return;
};
let mut should_fire_loaded = false;
if manager.get(&actor_id).is_none() {
let mut loaded = false;
if manager.config().persistence.enabled {
let storage = self.storage.read().clone();
if let Some(storage) = storage {
match storage.load_relationship(&self.info.id, &actor_id).await {
Ok(Some(value)) => match manager.insert_from_value(value) {
Ok(_) => loaded = true,
Err(e) => {
warn!(actor = %actor_id, error = %e, "failed to restore relationship")
}
},
Ok(None) => {}
Err(e) => {
warn!(actor = %actor_id, error = %e, "failed to load relationship")
}
}
}
}
if !loaded {
manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
}
should_fire_loaded = true;
}
let actor_name = self.resolve_actor_name_from_context();
let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
if should_fire_loaded {
self.hooks
.on_relationship_loaded(&actor_id, &relationship)
.await;
}
}
fn format_relationship_for_context(&self) -> Option<(String, String)> {
let manager = self.relationship_manager.as_ref()?;
if !manager.config().injection.enabled {
return None;
}
let actor_id = self.effective_actor_id()?;
let relationship = manager.get(&actor_id)?;
let local_cap = manager.config().injection.max_tokens;
let global_cap = self
.memory_token_budget
.as_ref()
.map(|b| b.allocation.relationships as usize)
.filter(|n| *n > 0);
let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
let text = ai_agents_relationships::format_relationship(
&relationship,
&manager.config().injection.format,
max_tokens,
);
if text.is_empty() {
None
} else {
Some((manager.config().injection.prompt_variable.clone(), text))
}
}
async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
let Some(manager) = self.relationship_manager.as_ref() else {
return Ok(());
};
if !manager.config().persistence.enabled {
return Ok(());
}
let storage = self.storage.read().clone();
let Some(storage) = storage else {
return Ok(());
};
if let Some(value) = manager.relationship_as_value(actor_id)? {
storage
.save_relationship(&self.info.id, actor_id, &value)
.await?;
}
Ok(())
}
async fn auto_update_relationship(&self) {
let Some(manager) = self.relationship_manager.as_ref() else {
return;
};
let Some(actor_id) = self.effective_actor_id() else {
return;
};
if !manager.config().auto_update.enabled {
let _ = self.persist_actor_relationship(&actor_id).await;
return;
}
let recent_messages = manager.config().auto_update.recent_messages;
let messages = match self.memory.get_messages(Some(recent_messages)).await {
Ok(messages) => messages,
Err(e) => {
warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
return;
}
};
match self
.observe_purpose(
ObservationPurpose::RelationshipUpdate,
manager.auto_update(&actor_id, &messages),
)
.await
{
Ok(update) => {
if !update.changes.is_empty() {
self.hooks
.on_relationship_change(&actor_id, &update.changes)
.await;
}
if let Some(ref event) = update.event {
self.hooks.on_notable_event(&actor_id, event).await;
}
let persisted = match self.persist_actor_relationship(&actor_id).await {
Ok(()) => true,
Err(e) => {
warn!(actor = %actor_id, error = %e, "failed to persist relationship");
false
}
};
if !update.changes.is_empty() || update.event.is_some() {
let changed_dimensions: Vec<String> = update
.changes
.iter()
.map(|change| format!("{}:{}", change.perspective, change.dimension))
.collect();
info!(
actor_id = %actor_id,
change_count = update.changes.len(),
changed_dimensions = ?changed_dimensions,
event_present = update.event.is_some(),
persisted = persisted,
"relationship updated"
);
} else {
debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
}
}
Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
}
}
async fn auto_extract_facts(&self) {
let should_extract = self
.facts_config
.as_ref()
.map(|c| c.enabled && c.auto_extract)
.unwrap_or(false);
if !should_extract {
debug!("fact extraction skipped because auto extraction is disabled");
return;
}
let msgs_since = *self.messages_since_extraction.read();
if msgs_since < 2 {
debug!(
messages_since_extraction = msgs_since,
"fact extraction skipped until threshold is reached"
);
return;
}
match self.extract_facts_with_source(msgs_since, "auto").await {
Ok(facts) => {
if !facts.is_empty() {
*self.messages_since_extraction.write() = 0;
} else {
debug!("fact extraction ran but found no new facts");
}
}
Err(e) => {
warn!("fact extraction failed: {}", e);
}
}
}
pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
self.persona_manager = Some(manager);
self
}
pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
self.persona_manager.as_ref()
}
pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
if config.is_enabled() {
let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
.with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
self.disambiguation_manager = Some(manager);
}
self
}
pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
self.disambiguation_manager.as_ref()
}
pub fn has_disambiguation(&self) -> bool {
self.disambiguation_manager
.as_ref()
.is_some_and(|m| m.is_enabled())
}
pub async fn init_storage(&self) -> Result<()> {
let _guard = self.storage_init.lock().await;
let mut storage = self.storage.read().clone();
if storage.is_none() && !self.storage_config.is_none() {
let storage_config = self.convert_storage_config();
storage = create_storage(&storage_config).await?;
*self.storage.write() = storage.clone();
}
self.validate_storage_requirements(storage.as_deref())?;
self.complete_facts_init().await;
Ok(())
}
fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
let facts_required = self
.facts_config
.as_ref()
.is_some_and(|config| config.enabled)
|| self
.actor_memory_config
.as_ref()
.is_some_and(|config| config.enabled);
let relationships_required = self
.relationship_manager
.as_ref()
.is_some_and(|manager| manager.config().persistence.enabled);
let Some(storage) = storage else {
let mut requirements = Vec::new();
if facts_required {
requirements.push("actor facts or actor memory");
}
if relationships_required {
requirements.push("persistent relationships");
}
if requirements.is_empty() {
return Ok(());
}
return Err(AgentError::Config(format!(
"Storage is required for enabled {} but none is configured or injected",
requirements.join(" and ")
)));
};
if facts_required && !storage.supports(StorageCapability::ActorFacts) {
return Err(AgentError::UnsupportedStorageCapability(
StorageCapability::ActorFacts,
));
}
if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
return Err(AgentError::UnsupportedStorageCapability(
StorageCapability::ActorRelationships,
));
}
Ok(())
}
async fn complete_facts_init(&self) {
if self.fact_store.read().is_some() {
return;
}
let storage = match self.storage.read().clone() {
Some(s) => s,
None => return,
};
let facts_enabled = self
.facts_config
.as_ref()
.map(|f| f.enabled)
.unwrap_or(false);
let actor_memory_enabled = self
.actor_memory_config
.as_ref()
.map(|a| a.enabled)
.unwrap_or(false);
if !facts_enabled && !actor_memory_enabled {
return;
}
let fc = self.facts_config.clone().unwrap_or_default();
let store = Arc::new(ai_agents_facts::FactStore::new(
storage,
self.info.id.clone(),
fc.clone(),
));
let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
let extractor_llm = fc
.extractor_llm
.as_ref()
.and_then(|alias| self.llm_registry.get(alias).ok())
.or_else(|| self.llm_registry.router().ok())
.or_else(|| self.llm_registry.default().ok());
extractor_llm.map(|llm| {
Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
as Arc<dyn ai_agents_facts::FactExtractor>
})
} else {
None
};
*self.fact_store.write() = Some(store);
*self.fact_extractor.write() = extractor;
debug!(
agent = %self.info.id,
facts_enabled,
actor_memory_enabled,
"facts storage initialized"
);
}
fn convert_storage_config(&self) -> StorageStorageConfig {
crate::spec::storage::to_storage_config(&self.storage_config)
}
pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
self.storage.read().clone()
}
pub fn storage_config(&self) -> &StorageConfig {
&self.storage_config
}
pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
self.spawner.as_ref()
}
pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
self.spawner_registry.as_ref()
}
pub fn has_spawner(&self) -> bool {
self.spawner_registry.is_some()
}
pub fn with_spawner_handles(
mut self,
spawner: Arc<crate::spawner::AgentSpawner>,
registry: Arc<crate::spawner::AgentRegistry>,
) -> Self {
self.spawner = Some(spawner);
self.spawner_registry = Some(registry);
self
}
pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
self.hooks = hooks;
self
}
pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
self.parallel_tools = config;
self
}
pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
self.streaming = config;
self
}
pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
self.hitl_engine = Some(engine);
self.approval_handler = handler;
self
}
pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
self.max_context_tokens = tokens;
self
}
pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
self.memory_token_budget = Some(budget);
self
}
pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
self.recovery_manager = manager;
self
}
pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
self.tool_security = engine;
self
}
pub fn runtime_control(&self) -> RuntimeControlHandle {
RuntimeControlHandle {
state: Arc::clone(&self.runtime_control),
}
}
pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
self.tools.set_question_handler(handler);
}
pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
self.tools.set_diagnostics_provider(provider);
}
pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
self.tools.set_command_runner(runner);
}
pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
self.tools.set_web_search_provider(provider);
}
pub fn todos(&self) -> Vec<TodoItem> {
self.tools.todos()
}
fn active_tool_security(&self) -> ToolSecurityEngine {
self.runtime_control
.tool_security_override
.read()
.clone()
.unwrap_or_else(|| self.tool_security.clone())
}
fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
let _guard = self.runtime_control.snapshot_guard.read();
RuntimeSafetySnapshot {
version: self.runtime_control.version.load(Ordering::SeqCst),
emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
tool_security: self
.runtime_control
.tool_security_override
.read()
.clone()
.unwrap_or_else(|| self.tool_security.clone()),
tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
}
}
fn admit_tool_execution(
&self,
expected_runtime_version: u64,
expected_policy_version: u64,
expected_state_generation: Option<u64>,
canonical_id: &str,
) -> SecurityCheckResult {
let _guard = self.runtime_control.snapshot_guard.read();
if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
return SecurityCheckResult::Block {
reason: "runtime emergency deny is enabled".to_string(),
};
}
let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
let security_engine = self
.runtime_control
.tool_security_override
.read()
.clone()
.unwrap_or_else(|| self.tool_security.clone());
if runtime_version != expected_runtime_version
|| security_engine.policy_version() != expected_policy_version
{
return SecurityCheckResult::Block {
reason: "runtime safety controls changed before admission".to_string(),
};
}
let current_state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
if current_state_generation != expected_state_generation {
return SecurityCheckResult::Block {
reason: "state scope changed before admission".to_string(),
};
}
security_engine.admit_tool_execution(canonical_id)
}
pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
self.process_processor = Some(processor);
self
}
pub fn with_state_machine(
mut self,
state_machine: Arc<StateMachine>,
evaluator: Arc<dyn TransitionEvaluator>,
) -> Self {
self.state_machine = Some(state_machine);
self.transition_evaluator = Some(evaluator);
self
}
pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
self.context_manager = manager;
self
}
pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
self.message_filters.write().insert(name.into(), filter);
}
pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
self.context_manager.update(key, value)
}
pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
self.context_manager.update(path, value)
}
pub fn get_context(&self) -> HashMap<String, Value> {
self.build_context_with_overlays()
}
pub fn remove_context(&self, key: &str) -> Option<Value> {
self.context_manager.remove(key)
}
pub async fn refresh_context(&self, key: &str) -> Result<()> {
self.context_manager.refresh(key).await
}
pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
self.context_manager.register_provider(name, provider);
}
pub fn current_state(&self) -> Option<String> {
self.state_machine.as_ref().map(|sm| sm.current())
}
async fn invalidate_pending_confirmation(&self, reason: &'static str) {
self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
return;
};
if disambiguator.has_pending_confirmation().await {
disambiguator.clear_pending().await;
*self.pending_skill_id.write() = None;
info!(
confirmation_event = "invalidated",
invalidation_reason = reason,
"Runtime invalidated pending confirmation"
);
}
}
async fn admit_disambiguation_redispatch(
&self,
expected_epoch: u64,
expected_state_generation: Option<u64>,
) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
let admission = self.disambiguation_admission.read().await;
let state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
|| state_generation != expected_state_generation
{
return Err(AgentError::Other(
"Disambiguation ownership changed before redispatch admission".to_string(),
));
}
Ok(admission)
}
fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
self.state_transition_reserved
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.ok()
.map(|_| StateTransitionReservation {
reserved: &self.state_transition_reserved,
})
}
async fn admit_optional_disambiguation_ownership(
&self,
ownership: Option<DisambiguationOwnership>,
) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
match ownership {
Some(ownership) => self
.admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
.await
.map(Some),
None => Ok(None),
}
}
pub async fn transition_to(&self, state: &str) -> Result<()> {
let Some(ref sm) = self.state_machine else {
return Ok(());
};
let claim_admission = self.disambiguation_admission.write().await;
let reservation = self.reserve_state_transition().ok_or_else(|| {
AgentError::Other("Another state transition is already in progress".to_string())
})?;
let from_state = sm.current();
let expected_state_generation = sm.generation();
let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
let history_before = sm.history();
drop(claim_admission);
self.execute_state_exit_actions(&from_state).await;
let admission = self.disambiguation_admission.write().await;
if sm.current() != from_state
|| sm.generation() != expected_state_generation
|| self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
{
return Err(AgentError::Other(
"State ownership changed during manual transition preparation".to_string(),
));
}
sm.transition_to(state, "manual transition")?;
self.invalidate_pending_confirmation("state_transition")
.await;
let entered = sm.current();
let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
drop(admission);
self.execute_state_enter_actions(&entered, is_reentry).await;
drop(reservation);
info!(to = %entered, "Manual state transition");
Ok(())
}
pub fn state_history(&self) -> Vec<StateTransitionEvent> {
self.state_machine
.as_ref()
.map(|sm| sm.history())
.unwrap_or_default()
}
pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
self.session_metadata.read().clone()
}
pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
let allowed = self
.actor_memory_config
.as_ref()
.map(|c| c.privacy.allow_deletion)
.unwrap_or(true);
if !allowed {
return Err(AgentError::Config(
"privacy.allow_deletion is false; actor data deletion is not permitted".into(),
));
}
let storage = self.storage.read().clone();
if let Some(storage) = storage {
if !storage.supports(StorageCapability::ActorDataDeletion) {
return Err(AgentError::UnsupportedStorageCapability(
StorageCapability::ActorDataDeletion,
));
}
storage.delete_actor_data(&self.info.id, actor_id).await?;
} else {
let store = { self.fact_store.read().clone() };
if let Some(store) = store {
store.delete_actor_data(actor_id).await?;
}
}
if let Some(manager) = self.relationship_manager.as_ref() {
manager.remove(actor_id);
}
self.actor_facts_cache.write().remove(actor_id);
Ok(())
}
pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
*self.session_metadata.write() = meta;
}
pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
let storage = self.storage.read().clone();
match storage {
Some(s) => {
let count = s.cleanup_expired().await?;
if count > 0 {
self.hooks.on_sessions_expired(count).await;
}
Ok(count)
}
None => Err(AgentError::Config(
"No storage configured. Use with_storage_config() or with_storage() first".into(),
)),
}
}
pub async fn list_sessions_filtered(
&self,
filter: &ai_agents_core::SessionFilter,
) -> Result<Vec<ai_agents_core::SessionSummary>> {
let storage = self.storage.read().clone();
match storage {
Some(s) => s.list_sessions_filtered(filter).await,
None => Err(AgentError::Config(
"No storage configured. Use with_storage_config() or with_storage() first".into(),
)),
}
}
pub async fn save_state(&self) -> Result<AgentSnapshot> {
let memory_snapshot = self.memory.snapshot().await?;
let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
let context_snapshot = self.context_manager.snapshot();
let mut snapshot = AgentSnapshot::new(self.info.id.clone())
.with_memory(memory_snapshot)
.with_context(context_snapshot)
.with_state_machine(
state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
current_state: String::new(),
previous_state: None,
turn_count: 0,
no_transition_count: 0,
history: vec![],
}),
);
if let Some(ref persona) = self.persona_manager {
snapshot.persona = Some(persona.snapshot_as_value()?);
}
if let Some(ref relationships) = self.relationship_manager {
snapshot.relationships = Some(relationships.snapshot_as_value()?);
}
Ok(snapshot)
}
pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
let mut snapshot = self.save_state().await?;
if let Some(ref registry) = self.spawner_registry {
let entries = registry.list_with_specs();
if !entries.is_empty() {
snapshot = snapshot.with_spawned_agents(entries);
}
}
Ok(snapshot)
}
pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
let _admission = self.disambiguation_admission.write().await;
if self.state_transition_reserved.load(Ordering::SeqCst) {
return Err(AgentError::Other(
"Cannot restore state while a state transition is in progress".to_string(),
));
}
self.invalidate_pending_confirmation("state_restore").await;
*self.pending_skill_id.write() = None;
if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
disambiguator.clear_pending().await;
}
self.memory.restore(snapshot.memory).await?;
if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
&& !sm_snapshot.current_state.is_empty()
{
sm.restore(sm_snapshot)?;
}
self.context_manager.restore(snapshot.context);
if let (Some(persona_value), Some(persona_manager)) =
(snapshot.persona, &self.persona_manager)
{
persona_manager.restore_from_value(persona_value)?;
}
if let (Some(relationship_value), Some(relationship_manager)) =
(snapshot.relationships, &self.relationship_manager)
{
relationship_manager.restore_from_value(relationship_value)?;
}
info!(agent_id = %snapshot.agent_id, "State restored");
Ok(())
}
pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
let snapshot = self.save_state().await?;
storage.save(session_id, &snapshot).await
}
async fn load_session_restore(
storage: &dyn AgentStorage,
session_id: &str,
) -> Result<Option<StoredSessionRestore>> {
let Some(snapshot) = storage.load(session_id).await? else {
return Ok(None);
};
let metadata = if storage.supports(StorageCapability::SessionMetadata) {
storage.load_metadata(session_id).await?
} else {
None
};
Ok(Some(StoredSessionRestore { snapshot, metadata }))
}
async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
Ok(RuntimeSessionRestorePoint {
snapshot: self.save_state().await?,
metadata: self.session_metadata(),
actor_id: self.actor_id(),
session_id: self.current_session_id.read().clone(),
})
}
async fn apply_session_restore_unchecked(
&self,
session_id: &str,
stored: StoredSessionRestore,
) -> Result<()> {
self.restore_state(stored.snapshot).await?;
let metadata = stored.metadata.unwrap_or_default();
if let Some(actor_id) = metadata.actor_id.as_deref() {
self.set_actor_id(actor_id)?;
} else {
self.clear_actor_id();
}
self.set_session_metadata(metadata);
*self.current_session_id.write() = Some(session_id.to_string());
Ok(())
}
async fn restore_session_restore_point(
&self,
restore_point: &RuntimeSessionRestorePoint,
) -> Result<()> {
self.restore_state(restore_point.snapshot.clone()).await?;
if let Some(actor_id) = restore_point.actor_id.as_deref() {
self.set_actor_id(actor_id)?;
} else {
self.clear_actor_id();
}
self.set_session_metadata(restore_point.metadata.clone());
*self.current_session_id.write() = restore_point.session_id.clone();
Ok(())
}
async fn apply_session_restore(
&self,
session_id: &str,
stored: StoredSessionRestore,
) -> Result<()> {
let before = self.capture_session_restore_point().await?;
if let Err(error) = self
.apply_session_restore_unchecked(session_id, stored)
.await
{
return match self.restore_session_restore_point(&before).await {
Ok(()) => Err(error),
Err(rollback_error) => Err(AgentError::Other(format!(
"Session restore failed: {error}; rollback failed: {rollback_error}"
))),
};
}
Ok(())
}
async fn rollback_session_restore_set(
parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
) -> Vec<String> {
let mut errors = Vec::new();
if let Some((agent, restore_point)) = parent
&& let Err(error) = agent.restore_session_restore_point(restore_point).await
{
errors.push(format!("parent: {error}"));
}
for (id, agent, restore_point) in children {
if let Err(error) = agent.restore_session_restore_point(restore_point).await {
errors.push(format!("child '{id}': {error}"));
}
}
errors
}
fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
if rollback_errors.is_empty() {
AgentError::Other(format!(
"Session restore failed: {error}; runtime state was rolled back"
))
} else {
AgentError::Other(format!(
"Session restore failed: {error}; rollback also failed for {}",
rollback_errors.join(", ")
))
}
}
pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
return Ok(false);
};
self.apply_session_restore(session_id, stored).await?;
Ok(true)
}
pub async fn save_session(&self, session_id: &str) -> Result<()> {
let storage = self.storage.read().clone();
match storage {
Some(s) => {
let is_new = {
let cur = self.current_session_id.read().clone();
cur.as_deref() != Some(session_id)
};
if is_new {
*self.current_session_id.write() = Some(session_id.to_string());
self.hooks.on_session_created(session_id).await;
}
{
let now = chrono::Utc::now();
let msg_count = self
.memory
.get_messages(None)
.await
.map(|v| v.len())
.unwrap_or(0);
let mut meta = self.session_metadata.write();
meta.last_active = now;
meta.message_count = msg_count;
if meta.actor_id.is_none() {
meta.actor_id = self.actor_id.read().clone();
}
}
let snapshot = self.save_state().await?;
if s.supports(StorageCapability::SessionMetadata) {
let metadata = self.session_metadata.read().clone();
s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
.await
} else {
s.save(session_id, &snapshot).await
}
}
None => Err(AgentError::Config(
"No storage configured. Use with_storage_config() or with_storage() first".into(),
)),
}
}
pub async fn load_session(&self, session_id: &str) -> Result<bool> {
let storage = self.storage.read().clone();
match storage {
Some(storage) => self.load_from(storage.as_ref(), session_id).await,
None => Err(AgentError::Config(
"No storage configured. Use with_storage_config() or with_storage() first".into(),
)),
}
}
pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
self.init_storage().await?;
let storage = self.storage.read().clone().ok_or_else(|| {
AgentError::Config(
"No storage configured. Use with_storage_config() or with_storage() first".into(),
)
})?;
let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
.await?
.ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
let manifest = target_parent
.snapshot
.spawned_agents
.clone()
.unwrap_or_default();
let registry = self.spawner_registry.as_ref().cloned();
let spawner = if manifest.is_empty() {
self.spawner.as_ref().cloned()
} else {
Some(self.spawner.as_ref().cloned().ok_or_else(|| {
AgentError::Config(
"Saved session contains child agents but this runtime has no spawner".into(),
)
})?)
};
let registry = if manifest.is_empty() {
registry
} else {
Some(registry.ok_or_else(|| {
AgentError::Config(
"Saved session contains child agents but this runtime has no registry".into(),
)
})?)
};
let mut target_ids = HashSet::with_capacity(manifest.len());
let mut prepared = Vec::with_capacity(manifest.len());
for entry in manifest {
if !target_ids.insert(entry.id.clone()) {
return Err(AgentError::InvalidSpec(format!(
"Saved child manifest contains duplicate ID: {}",
entry.id
)));
}
let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
spawner
.as_ref()
.expect("non-empty manifests require a spawner")
.validate_explicit_child(&entry.id, &spec)?;
prepared.push((entry.id, spec));
}
let current_ids = registry
.as_ref()
.map(|registry| {
registry
.list()
.into_iter()
.map(|info| info.id)
.collect::<HashSet<_>>()
})
.unwrap_or_default();
let removal_count = current_ids.difference(&target_ids).count();
let additions = prepared
.iter()
.filter(|(id, _)| !current_ids.contains(id))
.cloned()
.collect::<Vec<_>>();
let mut existing = Vec::new();
if let Some(registry) = registry.as_ref() {
for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
let agent = registry.get(id).ok_or_else(|| {
AgentError::Config(format!("Retained child disappeared during restore: {id}"))
})?;
let child_storage = agent.storage().ok_or_else(|| {
AgentError::Config(format!("Child '{id}' has no storage for session restore"))
})?;
let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
.await?
.ok_or_else(|| {
AgentError::Persistence(format!(
"Child '{id}' has no saved session '{session_id}'"
))
})?;
existing.push((id.clone(), agent, stored));
}
}
let mut staged = Vec::with_capacity(additions.len());
if !additions.is_empty() {
let spawner = spawner
.as_ref()
.expect("restored additions require a spawner");
let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
let spawned = spawner
.spawn_with_reserved_capacity(id.clone(), spec, reservation)
.await?;
let child_storage = spawned.agent.storage().ok_or_else(|| {
AgentError::Config(format!("Child '{id}' has no storage for session restore"))
})?;
let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
.await?
.ok_or_else(|| {
AgentError::Persistence(format!(
"Child '{id}' has no saved session '{session_id}'"
))
})?;
staged.push((spawned, stored));
}
} else if let Some(spawner) = spawner.as_ref() {
spawner.reserve_restore_capacity(0, removal_count)?;
}
let parent_before = self.capture_session_restore_point().await?;
let mut existing_before = Vec::with_capacity(existing.len());
for (id, agent, _) in &existing {
existing_before.push((
id.clone(),
Arc::clone(agent),
agent.capture_session_restore_point().await?,
));
}
for (_, agent, stored) in &existing {
if let Err(error) = agent
.apply_session_restore_unchecked(session_id, stored.clone())
.await
{
drop(staged);
let rollback_errors =
Self::rollback_session_restore_set(None, &existing_before).await;
return Err(Self::restore_failure(error, rollback_errors));
}
}
for (spawned, stored) in &staged {
if let Err(error) = spawned
.agent
.apply_session_restore_unchecked(session_id, stored.clone())
.await
{
drop(staged);
let rollback_errors =
Self::rollback_session_restore_set(None, &existing_before).await;
return Err(Self::restore_failure(error, rollback_errors));
}
}
if let Err(error) = self
.apply_session_restore_unchecked(session_id, target_parent)
.await
{
drop(staged);
let rollback_errors =
Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
.await;
return Err(Self::restore_failure(error, rollback_errors));
}
if let Some(registry) = registry.as_ref()
&& let Err(error) = registry
.reconcile(
&target_ids,
staged.into_iter().map(|(spawned, _)| spawned).collect(),
)
.await
{
let rollback_errors =
Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
.await;
return Err(Self::restore_failure(error, rollback_errors));
}
Ok(target_ids.len())
}
pub async fn delete_session(&self, session_id: &str) -> Result<()> {
let storage = self.storage.read().clone();
match storage {
Some(s) => s.delete(session_id).await,
None => Err(AgentError::Config(
"No storage configured. Use with_storage_config() or with_storage() first".into(),
)),
}
}
pub async fn list_sessions(&self) -> Result<Vec<String>> {
let storage = self.storage.read().clone();
match storage {
Some(s) => s.list_sessions().await,
None => Err(AgentError::Config(
"No storage configured. Use with_storage_config() or with_storage() first".into(),
)),
}
}
fn estimate_tokens(&self, text: &str) -> u32 {
(text.len() as f32 / 4.0).ceil() as u32
}
fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
messages
.iter()
.map(|m| self.estimate_tokens(&m.content))
.sum()
}
fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
if messages.len() <= keep_recent + 1 {
return;
}
let system_msg = messages.remove(0);
let to_remove = messages.len().saturating_sub(keep_recent);
messages.drain(..to_remove);
messages.insert(0, system_msg);
}
fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
match config {
FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
FilterConfig::SkipPattern { skip_if_contains } => {
Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
}
FilterConfig::Custom { name } => {
let filters = self.message_filters.read();
filters
.get(name)
.cloned()
.unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
}
}
}
async fn summarize_context(
&self,
messages: &mut Vec<ChatMessage>,
summarizer_llm: Option<&str>,
max_summary_tokens: u32,
custom_prompt: Option<&str>,
keep_recent: usize,
filter: Option<&FilterConfig>,
) -> Result<()> {
let system_msg = messages.remove(0);
let to_summarize_count = messages.len().saturating_sub(keep_recent);
if to_summarize_count == 0 {
messages.insert(0, system_msg);
return Ok(());
}
let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
let mut to_summarize = std::mem::take(messages);
if let Some(filter_config) = filter {
let filter = self.get_filter(filter_config);
to_summarize = filter.filter(to_summarize);
}
if to_summarize.is_empty() {
*messages = recent_msgs;
messages.insert(0, system_msg);
return Ok(());
}
let conversation_text = to_summarize
.iter()
.map(|m| format!("{:?}: {}", m.role, m.content))
.collect::<Vec<_>>()
.join("\n");
let default_prompt = format!(
"Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
max_summary_tokens, conversation_text
);
let summary_prompt = custom_prompt
.map(|p| format!("{}\n\n{}", p, conversation_text))
.unwrap_or(default_prompt);
let summarizer = if let Some(alias) = summarizer_llm {
self.llm_registry
.get(alias)
.map_err(|e| AgentError::Config(e.to_string()))?
} else {
self.llm_registry
.router()
.or_else(|_| self.llm_registry.default())
.map_err(|e| AgentError::Config(e.to_string()))?
};
let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
let response = self
.observe_purpose(
ObservationPurpose::Summarization,
summarizer.complete(&summary_msgs, None),
)
.await?;
let summary_message = ChatMessage::system(format!(
"[Previous conversation summary]\n{}",
response.content
));
*messages = vec![system_msg, summary_message];
messages.extend(recent_msgs);
debug!(
summarized_count = to_summarize_count,
kept_recent = keep_recent,
"Context summarized"
);
Ok(())
}
fn render_system_prompt(&self) -> Result<String> {
let mut context = self.build_context_with_overlays();
let facts_text = self.format_actor_facts_for_context();
if !facts_text.is_empty() {
context.insert(
"actor_facts".to_string(),
serde_json::Value::String(facts_text),
);
}
if let Some((key, text)) = self.format_relationship_for_context() {
context.insert(key, serde_json::Value::String(text));
}
self.template_renderer
.render(&self.base_system_prompt, &context)
}
fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
let mut seen = HashSet::new();
ids.iter()
.filter_map(|id| self.tools.canonical_id(id))
.filter(|canonical_id| seen.insert(canonical_id.clone()))
.collect()
}
fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
let Some(declared) = self.declared_tool_ids.as_deref() else {
return Vec::new();
};
let mut effective = self.canonical_unique_tool_ids(declared);
if let Some(scope) = scope_override {
let scope: HashSet<String> =
self.canonical_unique_tool_ids(scope).into_iter().collect();
effective.retain(|canonical_id| scope.contains(canonical_id));
}
effective
}
async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
}
async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
let scope_override = self.runtime_control.tool_scope_override.read().clone();
self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
.await
}
async fn get_available_tool_ids_snapshot_for_scope(
&self,
scope_override: Option<&[String]>,
) -> Result<AvailableToolIdsSnapshot> {
let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
let (state_generation, state_scopes) = self
.state_machine
.as_ref()
.map(|state_machine| {
let (generation, scopes) = state_machine.current_tool_scope_snapshot();
(Some(generation), scopes)
})
.unwrap_or((None, Vec::new()));
if available.is_empty() || state_scopes.is_empty() {
return Ok(AvailableToolIdsSnapshot {
tool_ids: available,
state_generation,
});
}
let eval_ctx = self.build_evaluation_context().await?;
let llm_getter = RegistryLLMGetter {
registry: self.llm_registry.clone(),
};
let evaluator = ConditionEvaluator::new(llm_getter);
for state_scope in state_scopes {
if state_scope.is_empty() {
available.clear();
break;
}
let mut allowed = HashSet::new();
for tool_ref in &state_scope {
let tool_id = tool_ref.id();
let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
continue;
};
let condition_matches = if let Some(condition) = tool_ref.condition() {
match evaluator.evaluate(condition, &eval_ctx).await {
Ok(matches) => matches,
Err(error) => {
warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
false
}
}
} else {
true
};
if condition_matches {
allowed.insert(canonical_id);
} else {
debug!(tool = tool_id, "Tool condition not met, skipping");
}
}
available.retain(|canonical_id| allowed.contains(canonical_id));
if available.is_empty() {
break;
}
}
Ok(AvailableToolIdsSnapshot {
tool_ids: available,
state_generation,
})
}
async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
let context = self.build_context_with_overlays();
let messages = self.memory.get_messages(Some(10)).await?;
let tool_history = self.tool_call_history.read().clone();
let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
(Some(sm.current()), sm.turn_count(), sm.previous())
} else {
(None, 0, None)
};
Ok(EvaluationContext::default()
.with_context(context)
.with_state(state_name, turn_count, previous_state)
.with_called_tools(tool_history)
.with_messages(messages))
}
fn record_tool_call(&self, tool_id: &str, result: Value) {
self.tool_call_history.write().push(ToolCallRecord {
tool_id: tool_id.to_string(),
result,
timestamp: chrono::Utc::now(),
});
}
async fn get_effective_system_prompt_with_persona_hooks(
&self,
fire_persona_hooks: bool,
include_tool_prompt: bool,
) -> Result<String> {
let rendered_base = self.render_system_prompt()?;
let persona_prefix = if let Some(ref persona) = self.persona_manager {
let context = self.build_context_with_overlays();
if fire_persona_hooks {
let render_result = persona.render_prompt(&context)?;
for content in &render_result.newly_revealed {
self.hooks.on_secret_revealed(content).await;
}
render_result.prompt
} else {
persona.render_prompt_preview(&context)?
}
} else {
String::new()
};
if let Some(ref sm) = self.state_machine
&& let Some(state_def) = sm.current_definition()
{
let state_prompt = if let Some(ref prompt) = state_def.prompt {
let context = self.build_context_with_overlays();
self.template_renderer.render_with_state(
prompt,
&context,
&sm.current(),
sm.previous().as_deref(),
sm.turn_count(),
state_def.max_turns,
)?
} else {
String::new()
};
let combined = match state_def.prompt_mode {
PromptMode::Append => {
if state_prompt.is_empty() {
rendered_base
} else {
format!(
"{}\n\n[Current State: {}]\n{}",
rendered_base,
sm.current(),
state_prompt
)
}
}
PromptMode::Replace => {
if state_prompt.is_empty() {
rendered_base
} else {
state_prompt
}
}
PromptMode::Prepend => {
if state_prompt.is_empty() {
rendered_base
} else {
format!("{}\n\n{}", state_prompt, rendered_base)
}
}
};
let with_persona = if persona_prefix.is_empty() {
combined
} else {
format!("{}\n\n{}", persona_prefix, combined)
};
if include_tool_prompt {
let available_tool_ids = self.get_available_tool_ids().await?;
if !available_tool_ids.is_empty() {
let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
&available_tool_ids,
None,
self.parallel_tools.enabled,
self.runtime_config.tool_schema_prompt_mode,
);
if !tools_prompt.is_empty() {
return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
}
}
}
return Ok(with_persona);
}
let with_persona = if persona_prefix.is_empty() {
rendered_base
} else {
format!("{}\n\n{}", persona_prefix, rendered_base)
};
if include_tool_prompt {
let available_tool_ids = self.get_available_tool_ids().await?;
let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
&available_tool_ids,
None,
self.parallel_tools.enabled,
self.runtime_config.tool_schema_prompt_mode,
);
if !tools_prompt.is_empty() {
return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
}
}
Ok(with_persona)
}
fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
if let Some(ref sm) = self.state_machine
&& let Some(state_def) = sm.current_definition()
&& let Some(ref llm_alias) = state_def.llm
{
return self
.llm_registry
.get(llm_alias)
.map_err(|e| AgentError::Config(e.to_string()));
}
self.llm_registry
.default()
.map_err(|e| AgentError::Config(e.to_string()))
}
fn get_effective_reasoning_config(&self) -> ReasoningConfig {
if let Some(ref sm) = self.state_machine
&& let Some(state_def) = sm.current_definition()
&& let Some(ref state_reasoning) = state_def.reasoning
{
return state_reasoning.clone();
}
self.reasoning_config.clone()
}
fn get_effective_reflection_config(&self) -> ReflectionConfig {
if let Some(ref sm) = self.state_machine
&& let Some(state_def) = sm.current_definition()
&& let Some(ref state_reflection) = state_def.reflection
{
return state_reflection.clone();
}
self.reflection_config.clone()
}
fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
skill
.reasoning
.clone()
.unwrap_or_else(|| self.get_effective_reasoning_config())
}
fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
skill
.reflection
.clone()
.unwrap_or_else(|| self.get_effective_reflection_config())
}
async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
let recent_messages: Vec<String> = self
.memory
.get_messages(Some(5))
.await?
.iter()
.rev()
.map(|m| format!("{:?}: {}", m.role, m.content))
.collect();
let current_state = self.current_state().map(|s| s.to_string());
let state_prompt: Option<String> = self
.state_machine
.as_ref()
.and_then(|sm| sm.current_definition())
.and_then(|def| def.prompt.clone());
let available_tools: Vec<String> = self
.get_available_tool_ids()
.await
.unwrap_or_else(|_| self.tools.list_ids());
let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
let mut user_context = self.build_context_with_overlays();
user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
if let Some(state_generation) = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation())
{
user_context.insert(
DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
serde_json::json!(state_generation),
);
}
let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
sm.current_definition()
.map(|def| {
def.transitions
.iter()
.filter_map(|t| t.intent.clone())
.collect()
})
.unwrap_or_default()
} else {
Vec::new()
};
Ok(DisambiguationContext::from_agent_state(
recent_messages,
current_state,
state_prompt,
available_tools,
available_skills,
available_intents,
user_context,
))
}
fn get_available_skills(&self) -> Vec<&SkillDefinition> {
if let Some(ref sm) = self.state_machine
&& let Some(state_def) = sm.current_definition()
{
let parent_def = sm.get_parent_definition();
let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
if !effective_skills.is_empty() {
return self
.skills
.iter()
.filter(|s| effective_skills.contains(&&s.id))
.collect();
}
}
self.skills.iter().collect()
}
async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
self.build_messages_internal(true, None, true).await
}
async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
self.build_messages_internal(false, Some(user_message), true)
.await
}
async fn build_messages_internal(
&self,
fire_persona_hooks: bool,
ephemeral_user_message: Option<&str>,
include_tool_prompt: bool,
) -> Result<Vec<ChatMessage>> {
let system_prompt = self
.get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
.await?;
let mut messages = vec![ChatMessage::system(&system_prompt)];
let context = self.memory.get_context().await?;
let history = if let Some(ref budget) = self.memory_token_budget {
context.to_llm_messages_with_allocation(&budget.allocation)
} else {
context.to_llm_messages()
};
messages.extend(history);
if let Some(user_message) = ephemeral_user_message {
messages.push(ChatMessage::user(user_message));
}
let total_tokens = self.estimate_total_tokens(&messages);
if total_tokens > self.max_context_tokens {
debug!(
total = total_tokens,
limit = self.max_context_tokens,
"Context overflow"
);
match &self.recovery_manager.config().llm.on_context_overflow {
ContextOverflowAction::Error => {
return Err(AgentError::LLM(format!(
"Context overflow: {} tokens > {} limit",
total_tokens, self.max_context_tokens
)));
}
ContextOverflowAction::Truncate { keep_recent } => {
self.truncate_context(&mut messages, *keep_recent);
}
ContextOverflowAction::Summarize {
summarizer_llm,
max_summary_tokens,
custom_prompt,
keep_recent,
filter,
} => {
self.summarize_context(
&mut messages,
summarizer_llm.as_deref(),
*max_summary_tokens,
custom_prompt.as_deref(),
*keep_recent,
filter.as_ref(),
)
.await?;
}
}
}
Ok(messages)
}
async fn main_tool_protocol(
&self,
llm: &dyn LLMProvider,
ephemeral_new_turn: bool,
) -> Result<MainToolProtocol> {
let mut choice = llm.configured_tool_choice();
if matches!(choice.as_ref(), Some(ToolChoice::None)) {
return Ok(MainToolProtocol {
choice,
tool_ids: Vec::new(),
definitions: Vec::new(),
});
}
let mut tool_ids = self.get_available_tool_ids().await?;
tool_ids.sort();
tool_ids.dedup();
if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
AgentError::Config(format!(
"specific tool choice '{expected}' is not registered"
))
})?;
if canonical != *expected {
return Err(AgentError::Config(format!(
"specific tool choice must use canonical ID '{canonical}', not '{expected}'"
)));
}
if !tool_ids.iter().any(|tool_id| tool_id == expected) {
return Err(AgentError::Config(format!(
"specific tool choice '{expected}' is outside the effective tool grant"
)));
}
}
if matches!(
choice.as_ref(),
Some(ToolChoice::Required | ToolChoice::Specific(_))
) && tool_ids.is_empty()
{
return Err(AgentError::Config(
"required tool choice has no tool inside the effective grant".to_string(),
));
}
if !ephemeral_new_turn
&& let Some(configured_choice) = choice.as_ref()
&& matches!(
configured_choice,
ToolChoice::Required | ToolChoice::Specific(_)
)
&& self
.tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
.await?
{
choice = Some(ToolChoice::Auto);
}
if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
tool_ids.retain(|tool_id| tool_id == expected);
}
let definitions = tool_ids
.iter()
.map(|tool_id| {
let tool = self.tools.get(tool_id).ok_or_else(|| {
AgentError::Config(format!(
"effective tool '{tool_id}' disappeared before provider exposure"
))
})?;
Ok(LLMToolDefinition {
name: tool_id.clone(),
description: tool.description().to_string(),
input_schema: tool.input_schema(),
})
})
.collect::<Result<Vec<_>>>()?;
Ok(MainToolProtocol {
choice,
tool_ids,
definitions,
})
}
async fn tool_choice_satisfied_in_current_turn(
&self,
choice: &ToolChoice,
effective_tool_ids: &[String],
) -> Result<bool> {
let messages = self.memory.get_messages(None).await?;
let mut saw_tool_result = false;
for message in messages.iter().rev() {
match message.role {
ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
saw_tool_result = true;
}
ai_agents_core::Role::Assistant if saw_tool_result => {
let Some(calls) = self.parse_tool_calls(&message.content) else {
continue;
};
let calls_are_effective = !calls.is_empty()
&& calls.iter().all(|call| {
self.tools
.canonical_id(&call.name)
.is_some_and(|canonical| effective_tool_ids.contains(&canonical))
});
return Ok(calls_are_effective
&& match choice {
ToolChoice::Required => true,
ToolChoice::Specific(expected) => calls.iter().all(|call| {
self.tools.canonical_id(&call.name).as_deref()
== Some(expected.as_str())
}),
_ => false,
});
}
ai_agents_core::Role::User => return Ok(false),
_ => {}
}
}
Ok(false)
}
fn provider_can_use_native_tools(
&self,
llm: &dyn LLMProvider,
protocol: &MainToolProtocol,
) -> bool {
let Some(choice) = protocol.choice.as_ref() else {
return false;
};
if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
return false;
}
llm.supports_tool_choice(choice)
&& protocol.definitions.iter().all(|definition| {
!definition.name.is_empty()
&& definition.name.len() <= 64
&& definition
.name
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
})
}
fn prompt_messages_for_tool_protocol(
&self,
messages: &[ChatMessage],
protocol: &MainToolProtocol,
corrective: bool,
) -> Vec<ChatMessage> {
let mut messages = messages.to_vec();
let Some(choice) = protocol.choice.as_ref() else {
return messages;
};
if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
return messages;
}
let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
&protocol.tool_ids,
None,
self.parallel_tools.enabled,
self.runtime_config.tool_schema_prompt_mode,
);
match choice {
ToolChoice::Required => tool_prompt.push_str(
"\n\nYou must call at least one listed tool before giving a final answer.",
),
ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
"\n\nYou must call the '{tool_id}' tool before giving a final answer."
)),
ToolChoice::Auto => {}
ToolChoice::None => return messages,
_ => return messages,
}
if let Some(system) = messages
.iter_mut()
.find(|message| message.role == ai_agents_core::Role::System)
{
system.content.push_str("\n\n");
system.content.push_str(&tool_prompt);
} else {
messages.insert(0, ChatMessage::system(tool_prompt));
}
if corrective {
let instruction = match choice {
ToolChoice::Required => {
"Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
}
ToolChoice::Specific(tool_id) => {
messages.push(ChatMessage::user(format!(
"Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
)));
return messages;
}
_ => return messages,
};
messages.push(ChatMessage::user(instruction));
}
messages
}
async fn invoke_main_provider(
&self,
llm: Arc<dyn LLMProvider>,
messages: &[ChatMessage],
protocol: &MainToolProtocol,
corrective: bool,
) -> std::result::Result<MainProviderResponse, LLMError> {
let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
let response = if use_native {
let request = LLMToolRequest {
tools: protocol.definitions.clone(),
choice: protocol
.choice
.clone()
.expect("native tool requests require an explicit choice"),
};
self.observe_purpose(
ObservationPurpose::MainResponse,
llm.complete_with_tools(messages, None, &request),
)
.await?
} else {
let prompt_messages =
self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
self.observe_purpose(
ObservationPurpose::MainResponse,
llm.complete(&prompt_messages, None),
)
.await?
};
Ok(MainProviderResponse {
response,
used_native_tools: use_native,
})
}
async fn complete_main_attempt_with_recovery(
&self,
llm: Arc<dyn LLMProvider>,
messages: &[ChatMessage],
protocol: &MainToolProtocol,
corrective: bool,
) -> Result<MainProviderResponse> {
let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
self.recovery_manager
.with_retry("llm_call", None, || {
let llm = Arc::clone(&llm);
async move {
self.invoke_main_provider(llm, messages, protocol, corrective)
.await
.map_err(|error| error.classify())
}
})
.await
.map_err(|error| AgentError::LLM(error.to_string()))
} else {
self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
.await
.map_err(|error| AgentError::LLM(error.to_string()))
};
match primary_result {
Ok(response) => Ok(response),
Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
LLMFailureAction::FallbackLlm { fallback_llm } => {
let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
AgentError::Config(format!(
"Fallback LLM '{fallback_llm}' not found: {error}"
))
})?;
self.invoke_main_provider(fallback, messages, protocol, corrective)
.await
.map_err(|error| AgentError::LLM(error.to_string()))
}
LLMFailureAction::FallbackResponse { message } => {
if matches!(
protocol.choice.as_ref(),
Some(ToolChoice::Required | ToolChoice::Specific(_))
) {
Err(AgentError::LLM(format!(
"Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
)))
} else {
Ok(MainProviderResponse {
response: LLMResponse::new(message.clone(), FinishReason::Stop),
used_native_tools: false,
})
}
}
LLMFailureAction::Error => Err(primary_error),
},
}
}
fn normalize_main_provider_response(
&self,
mut response: LLMResponse,
protocol: &MainToolProtocol,
) -> Result<(LLMResponse, bool)> {
let native_calls = response
.tool_calls()
.map_err(|error| AgentError::LLM(error.to_string()))?;
let calls = match native_calls {
Some(calls) => {
let markers = calls
.iter()
.map(|call| {
serde_json::json!({
"_ai_agents_native_tool_call": true,
"id": call.id,
"tool": call.name,
"arguments": call.arguments,
})
})
.collect::<Vec<_>>();
response.content = if markers.len() == 1 {
markers[0].to_string()
} else {
serde_json::Value::Array(markers).to_string()
};
Some(calls)
}
None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
self.parse_tool_calls(response.content.trim())
}
None => None,
};
if protocol.choice.is_some()
&& let Some(calls) = calls.as_ref()
&& calls.iter().any(|call| {
self.tools
.canonical_id(&call.name)
.is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
})
{
return Err(AgentError::LLM(
"Provider returned a tool call outside the effective grant".to_string(),
));
}
let compliant = match protocol.choice.as_ref() {
Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
!calls.is_empty()
&& calls.iter().all(|call| {
self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
})
}),
_ => true,
};
Ok((response, compliant))
}
async fn complete_main_llm_with_recovery(
&self,
llm: Arc<dyn LLMProvider>,
messages: &[ChatMessage],
protocol: &MainToolProtocol,
) -> Result<LLMResponse> {
let first = self
.complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
.await?;
let (response, compliant) =
self.normalize_main_provider_response(first.response, protocol)?;
if compliant {
return Ok(response);
}
if first.used_native_tools {
return Err(AgentError::LLM(
"Provider returned no compliant native call for required tool choice".to_string(),
));
}
let corrected = self
.complete_main_attempt_with_recovery(llm, messages, protocol, true)
.await?;
let (response, compliant) =
self.normalize_main_provider_response(corrected.response, protocol)?;
if compliant {
return Ok(response);
}
Err(AgentError::LLM(
"Provider returned no compliant tool call after one corrective retry".to_string(),
))
}
fn is_native_tool_call_content(content: &str) -> bool {
let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
return false;
};
match value {
serde_json::Value::Array(values) => {
!values.is_empty()
&& values.iter().all(|value| {
value
.get("_ai_agents_native_tool_call")
.and_then(|marker| marker.as_bool())
== Some(true)
})
}
serde_json::Value::Object(map) => {
map.get("_ai_agents_native_tool_call")
.and_then(|marker| marker.as_bool())
== Some(true)
}
_ => false,
}
}
fn tool_result_message(
tool_call: &ToolCall,
output: &str,
native_tool_call: bool,
) -> ChatMessage {
if !native_tool_call {
return ChatMessage::function(&tool_call.name, output);
}
let output = serde_json::from_str::<serde_json::Value>(output)
.unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
ChatMessage::function(
&tool_call.name,
serde_json::json!({
"_ai_agents_native_tool_result": true,
"id": tool_call.id,
"tool": tool_call.name,
"output": output,
})
.to_string(),
)
}
fn parse_main_tool_calls(
&self,
content: &str,
protocol: &MainToolProtocol,
) -> Option<Vec<ToolCall>> {
if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
None
} else {
self.parse_tool_calls(content)
}
}
fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
if let Some(arr) = parsed.as_array() {
let calls: Vec<ToolCall> = arr
.iter()
.filter_map(|v| self.extract_tool_call_from_value(v))
.collect();
if !calls.is_empty() {
return Some(calls);
}
}
if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
return Some(vec![tool_call]);
}
}
if let Some(json_str) = self.extract_json_from_content(content)
&& let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
{
if let Some(arr) = parsed.as_array() {
let calls: Vec<ToolCall> = arr
.iter()
.filter_map(|v| self.extract_tool_call_from_value(v))
.collect();
if !calls.is_empty() {
return Some(calls);
}
}
if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
return Some(vec![tool_call]);
}
}
None
}
fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
let arguments = parsed
.get("arguments")
.cloned()
.unwrap_or(serde_json::json!({}));
return Some(ToolCall {
id: parsed
.get("id")
.and_then(|value| value.as_str())
.filter(|id| !id.is_empty())
.map(str::to_string)
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
name: tool_name.to_string(),
arguments,
});
}
None
}
fn extract_json_from_content(&self, content: &str) -> Option<String> {
if let Some(result) = self.extract_json_array_from_content(content) {
return Some(result);
}
self.extract_json_object_from_content(content)
}
fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
let start = content.find('[')?;
let content_from_start = &content[start..];
let mut depth = 0;
let mut end = 0;
for (i, ch) in content_from_start.char_indices() {
match ch {
'[' => depth += 1,
']' => {
depth -= 1;
if depth == 0 {
end = i + 1;
break;
}
}
_ => {}
}
}
if end > 0 {
let json_str = &content_from_start[..end];
if json_str.contains("\"tool\"") {
return Some(json_str.to_string());
}
}
None
}
fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
let start = content.find('{')?;
let content_from_start = &content[start..];
let mut depth = 0;
let mut end = 0;
for (i, ch) in content_from_start.char_indices() {
match ch {
'{' => depth += 1,
'}' => {
depth -= 1;
if depth == 0 {
end = i + 1;
break;
}
}
_ => {}
}
}
if end > 0 {
let json_str = &content_from_start[..end];
if json_str.contains("\"tool\"") {
return Some(json_str.to_string());
}
}
None
}
#[allow(clippy::too_many_arguments)]
fn record_from_parts(
&self,
request: &ToolExecutionRequest,
canonical_id: String,
executed_arguments: Value,
started_at: chrono::DateTime<chrono::Utc>,
start: Instant,
executed: bool,
success: bool,
output: String,
metadata: HashMap<String, Value>,
policy: ToolPolicyDecisionRecord,
approval: Option<ToolApprovalRecord>,
timed_out: bool,
output_truncated: bool,
) -> ToolExecutionRecord {
let versions = ToolDecisionVersions {
policy: self.active_tool_security().policy_version(),
registry: self.tools.version(),
runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
state: self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation()),
};
self.record_from_parts_at(
request,
canonical_id,
executed_arguments,
started_at,
start,
executed,
success,
output,
metadata,
policy,
approval,
timed_out,
output_truncated,
versions,
)
}
#[allow(clippy::too_many_arguments)]
fn record_from_parts_at(
&self,
request: &ToolExecutionRequest,
canonical_id: String,
executed_arguments: Value,
started_at: chrono::DateTime<chrono::Utc>,
start: Instant,
executed: bool,
success: bool,
output: String,
metadata: HashMap<String, Value>,
policy: ToolPolicyDecisionRecord,
approval: Option<ToolApprovalRecord>,
timed_out: bool,
output_truncated: bool,
versions: ToolDecisionVersions,
) -> ToolExecutionRecord {
ToolExecutionRecord {
call_id: request.call_id.clone(),
requested_name: request.requested_name.clone(),
canonical_id,
source: request.source.clone(),
arguments: request.arguments.clone(),
executed_arguments,
policy_version: versions.policy,
registry_version: versions.registry,
runtime_config_version: versions.runtime_control,
executed,
success,
output,
metadata,
policy,
approval,
started_at,
duration_ms: start.elapsed().as_millis() as u64,
timed_out,
cancelled: false,
cancellation_reason: None,
output_truncated,
}
}
async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
let result = ToolResult {
success: record.success,
output: record.model_output_string(),
metadata: if record.metadata.is_empty() {
None
} else {
Some(record.metadata.clone())
},
};
self.hooks
.on_tool_complete(&record.canonical_id, &result, record.duration_ms)
.await;
self.hooks.on_tool_execution_record(record).await;
self.record_tool_call(&record.canonical_id, record.model_output_value());
if !record.success {
self.hooks
.on_error(&AgentError::Tool(record.output.clone()))
.await;
}
}
async fn finish_tool_record_after_resource_guards(
&self,
resource_guards: ToolResourceGuards,
record: &ToolExecutionRecord,
) {
drop(resource_guards);
self.finish_tool_record(record).await;
}
fn validated_tool_timeout(timeout_ms: u64) -> Result<ValidatedToolTimeout> {
if timeout_ms > MAX_TOOL_TIMEOUT_MS {
return Err(AgentError::Config(format!(
"effective tool timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
)));
}
let timer = Duration::from_millis(timeout_ms);
let deadline_delta = chrono::Duration::from_std(timer).map_err(|_| {
AgentError::Config(format!(
"effective tool timeout_ms cannot be represented as a UTC deadline: {timeout_ms}"
))
})?;
Ok(ValidatedToolTimeout {
timer,
deadline_delta,
})
}
fn effective_tool_limits(
security_engine: &ToolSecurityEngine,
canonical_id: &str,
safety: &ToolSafetyMetadata,
classification: &ToolCallClassification,
recovery_timeout_ms: Option<u64>,
) -> Result<(ToolExecutionLimits, ValidatedToolTimeout)> {
if let Some(timeout_ms) = classification.timeout_ms {
Self::validated_tool_timeout(timeout_ms)?;
}
if let Some(timeout_ms) = recovery_timeout_ms {
Self::validated_tool_timeout(timeout_ms)?;
}
let mut limits = security_engine.effective_limits(canonical_id, safety, classification);
if let Some(recovery_timeout_ms) = recovery_timeout_ms {
limits.timeout_ms = Some(limits.timeout_ms.map_or(recovery_timeout_ms, |timeout_ms| {
timeout_ms.min(recovery_timeout_ms)
}));
}
let timeout_ms = limits
.timeout_ms
.unwrap_or_else(|| security_engine.get_tool_timeout(canonical_id));
let timeout = Self::validated_tool_timeout(timeout_ms)?;
Ok((limits, timeout))
}
async fn execute_resolved_tool_once(
&self,
tool: Arc<dyn ai_agents_core::Tool>,
args: Value,
mut ctx: ToolExecutionContext,
timeout: ValidatedToolTimeout,
) -> Result<(ToolResult, bool, bool, bool)> {
if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
return Ok((
ToolResult::error("Tool execution cancelled by runtime control"),
false,
true,
false,
));
}
ctx.deadline = Some(
chrono::Utc::now()
.checked_add_signed(timeout.deadline_delta)
.ok_or_else(|| {
AgentError::Config(
"effective tool timeout_ms exceeds the current UTC deadline range"
.to_string(),
)
})?,
);
let invoked = Arc::new(AtomicBool::new(false));
let invoked_by_future = Arc::clone(&invoked);
let actor_context = current_turn_actor_context();
let future = async move {
invoked_by_future.store(true, Ordering::SeqCst);
if let Some(actor_context) = actor_context {
scope_actor_context(actor_context, tool.execute(args, ctx)).await
} else {
tool.execute(args, ctx).await
}
};
tokio::pin!(future);
let timer = tokio::time::sleep(timeout.timer);
tokio::pin!(timer);
let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
loop {
tokio::select! {
result = &mut future => return Ok((result, false, false, true)),
_ = &mut timer => {
return Ok((
ToolResult::error("Tool execution timed out"),
true,
false,
invoked.load(Ordering::SeqCst),
));
}
_ = cancel_tick.tick() => {
if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
return Ok((
ToolResult::error("Tool execution cancelled by runtime control"),
false,
true,
invoked.load(Ordering::SeqCst),
));
}
}
}
}
}
fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
let Some(max_chars) = max_chars else {
return (output, false);
};
let mut chars = output.chars();
let truncated: String = chars.by_ref().take(max_chars).collect();
if chars.next().is_some() {
(truncated, true)
} else {
(output, false)
}
}
async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
let locks = {
let mut table = self.resource_locks.write();
table.retain(|_, lock| lock.strong_count() > 0);
keys.iter()
.map(|key| {
if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
lock
} else {
let lock = Arc::new(tokio::sync::Mutex::new(()));
table.insert(key.clone(), Arc::downgrade(&lock));
lock
}
})
.collect::<Vec<_>>()
};
let mut resource_guards = ToolResourceGuards {
guards: Vec::with_capacity(locks.len()),
locks: Arc::clone(&self.resource_locks),
};
let mut locks = locks.into_iter();
while let Some(lock) = locks.next() {
let mut lock = Box::pin(lock.lock_owned());
loop {
tokio::select! {
guard = &mut lock => {
resource_guards.guards.push(guard);
break;
}
_ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
drop(lock);
drop(locks);
drop(resource_guards);
return None;
}
}
}
}
}
Some(resource_guards)
}
async fn run_tool_with_retries(
&self,
canonical_id: &str,
tool: Arc<dyn ai_agents_core::Tool>,
args: Value,
ctx: ToolExecutionContext,
timeout: ValidatedToolTimeout,
max_retries: u32,
) -> Result<(ToolResult, bool, bool, bool)> {
let max_retries = if ctx.classification.safely_retryable {
max_retries
} else {
0
};
let mut attempts = 0;
let mut invoked = false;
loop {
let (result, timed_out, cancelled, attempt_invoked) = self
.execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout)
.await?;
invoked |= attempt_invoked;
if result.success || timed_out || cancelled || attempts >= max_retries {
return Ok((result, timed_out, cancelled, invoked));
}
attempts += 1;
warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
}
}
fn host_tool_unavailability(&self, canonical_id: &str) -> Option<(&'static str, &'static str)> {
match canonical_id {
"command" if !self.tools.command_runner_available() => Some((
"Command runner is unavailable",
"command runner is unavailable",
)),
"diagnostics" if !self.tools.diagnostics_available() => Some((
"Diagnostics provider is unavailable",
"diagnostics provider is unavailable",
)),
"web_search" if !self.tools.web_search_available() => Some((
"Web search provider is unavailable",
"web search provider is unavailable",
)),
_ => None,
}
}
fn execute_tool_record(
&self,
request: ToolExecutionRequest,
) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
Box::pin(self.execute_tool_record_inner(request, ToolFallbackState::default()))
}
async fn execute_tool_record_inner(
&self,
request: ToolExecutionRequest,
fallback_state: ToolFallbackState,
) -> Result<ToolExecutionRecord> {
let started_at = chrono::Utc::now();
let start = Instant::now();
info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
let record = self.record_from_parts(
&request,
request.requested_name.clone(),
request.arguments.clone(),
started_at,
start,
false,
false,
"Tool execution is disabled by runtime control".to_string(),
HashMap::new(),
ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
None,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let Some(resolved) = self.tools.resolve(&request.requested_name) else {
let record = self.record_from_parts(
&request,
request.requested_name.clone(),
request.arguments.clone(),
started_at,
start,
false,
false,
format!("Tool '{}' is unavailable", request.requested_name),
HashMap::new(),
ToolPolicyDecisionRecord::unavailable(format!(
"Tool '{}' is not registered",
request.requested_name
)),
None,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
};
let canonical_id = resolved.identity.canonical_id.clone();
let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
if !initial_scope_snapshot
.tool_ids
.iter()
.any(|id| id == &canonical_id)
{
let record = self.record_from_parts(
&request,
canonical_id.clone(),
request.arguments.clone(),
started_at,
start,
false,
false,
format!(
"Tool '{}' is not available in the current scope",
canonical_id
),
HashMap::new(),
ToolPolicyDecisionRecord::deny(format!(
"Tool '{}' is not granted by the current top-level and state tool scope",
canonical_id
)),
None,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let approval_control_snapshot = self.runtime_safety_snapshot();
let security_engine = approval_control_snapshot.tool_security.clone();
if let Some(reason) = fallback_state.rejection_reason(&canonical_id) {
let mut metadata = HashMap::new();
metadata.insert(
"fallback_chain".to_string(),
serde_json::to_value(&fallback_state.visited_canonical_ids).unwrap_or(Value::Null),
);
let record = self.record_from_parts(
&request,
canonical_id,
request.arguments.clone(),
started_at,
start,
false,
false,
format!("Denied: {reason}"),
metadata,
ToolPolicyDecisionRecord::deny(reason),
None,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let admitted_canonical_id = canonical_id.clone();
let fallback_state = fallback_state.with_current(canonical_id.clone());
let bindings = resolved.tool.policy_bindings();
let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
&canonical_id,
&request.arguments,
&bindings,
);
let mut metadata = HashMap::new();
let safety = resolved.tool.safety_metadata();
let classification = resolved.tool.classify_call(&executed_arguments);
let initial_recovery_timeout_ms = self.recovery_manager.get_tool_timeout(&canonical_id);
let (limits, _) = Self::effective_tool_limits(
&security_engine,
&canonical_id,
&safety,
&classification,
initial_recovery_timeout_ms,
)?;
self.hooks
.on_tool_start(&canonical_id, &executed_arguments)
.await;
metadata.insert(
"classification".to_string(),
serde_json::to_value(&classification).unwrap_or(Value::Null),
);
metadata.insert(
"effective_limits".to_string(),
serde_json::to_value(&limits).unwrap_or(Value::Null),
);
let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
if !policy_snapshot.is_null() {
metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
}
let mut approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::NotRequired,
reason: None,
modified_arguments: None,
});
let mut security_result = security_engine
.validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
.await?;
if (security_result.is_allowed()
|| matches!(
&security_result,
SecurityCheckResult::RequireConfirmation { .. }
))
&& let Some((output, reason)) = self.host_tool_unavailability(&canonical_id)
{
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
output.to_string(),
metadata,
ToolPolicyDecisionRecord::unavailable(reason),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Unavailable,
reason: Some(reason.to_string()),
modified_arguments: None,
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
match &security_result {
SecurityCheckResult::Allow => {}
SecurityCheckResult::Warn { message } => {
warn!(tool = %canonical_id, message = %message, "Tool security warning");
}
SecurityCheckResult::Block { reason } => {
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Denied: {}", reason),
metadata,
ToolPolicyDecisionRecord::deny(reason.clone()),
approval_record,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
SecurityCheckResult::Unavailable { reason } => {
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Unavailable: {}", reason),
metadata,
ToolPolicyDecisionRecord::unavailable(reason.clone()),
approval_record,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
SecurityCheckResult::RequireConfirmation { message } => {
if self.hitl_engine.is_none() {
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Unavailable,
reason: Some("No HITL engine configured".to_string()),
modified_arguments: None,
});
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Approval unavailable: {}", message),
metadata,
ToolPolicyDecisionRecord::approval(message.clone()),
approval_record,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let check_result = HITLCheckResult::required(
ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
HashMap::new(),
message.clone(),
None,
);
match self.request_hitl_approval(check_result).await? {
ApprovalResult::Approved => {
merge_approved_record(&mut approval_record);
}
ApprovalResult::Modified { changes } => {
if let Some(obj) = executed_arguments.as_object_mut() {
for (key, value) in changes {
obj.insert(key, value);
}
}
security_result = security_engine
.validate_tool_execution_with_bindings(
&canonical_id,
&executed_arguments,
&bindings,
)
.await?;
if !matches!(
security_result,
SecurityCheckResult::Allow
| SecurityCheckResult::Warn { .. }
| SecurityCheckResult::RequireConfirmation { .. }
) {
let reason = security_result
.reason()
.unwrap_or("modified arguments failed policy")
.to_string();
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments.clone(),
started_at,
start,
false,
false,
reason.clone(),
metadata,
ToolPolicyDecisionRecord::deny(reason),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(executed_arguments),
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(executed_arguments.clone()),
});
}
ApprovalResult::Rejected { reason } => {
let reason = reason.unwrap_or_else(|| "rejected".to_string());
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Rejected,
reason: Some(reason.clone()),
modified_arguments: None,
});
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Approval rejected: {}", reason),
metadata,
ToolPolicyDecisionRecord::approval(reason),
approval_record,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
ApprovalResult::Timeout => {
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Timeout,
reason: Some("approval timeout".to_string()),
modified_arguments: None,
});
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
"Approval timed out".to_string(),
metadata,
ToolPolicyDecisionRecord::approval("approval timeout"),
approval_record,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
}
}
}
if approval_record
.as_ref()
.is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
&& let Some(message) =
security_engine.classification_approval_message(&canonical_id, &classification)
{
if self.hitl_engine.is_none() {
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Unavailable,
reason: Some("No HITL engine configured".to_string()),
modified_arguments: None,
});
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Approval unavailable: {}", message),
metadata,
ToolPolicyDecisionRecord::approval(message),
approval_record,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let check_result = HITLCheckResult::required(
ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
HashMap::new(),
message.clone(),
None,
);
match self.request_hitl_approval(check_result).await? {
ApprovalResult::Approved => {
merge_approved_record(&mut approval_record);
}
ApprovalResult::Modified { changes } => {
if let Some(obj) = executed_arguments.as_object_mut() {
for (key, value) in changes {
obj.insert(key, value);
}
}
let modified_security = security_engine
.validate_tool_execution_with_bindings(
&canonical_id,
&executed_arguments,
&bindings,
)
.await?;
if !matches!(
modified_security,
SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
) {
let reason = modified_security
.reason()
.unwrap_or("modified arguments failed policy")
.to_string();
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments.clone(),
started_at,
start,
false,
false,
reason.clone(),
metadata,
ToolPolicyDecisionRecord::deny(reason),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(executed_arguments),
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(executed_arguments.clone()),
});
}
ApprovalResult::Rejected { reason } => {
let reason = reason.unwrap_or_else(|| "rejected".to_string());
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Approval rejected: {}", reason),
metadata,
ToolPolicyDecisionRecord::approval(reason.clone()),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Rejected,
reason: Some(reason),
modified_arguments: None,
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
ApprovalResult::Timeout => {
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
"Approval timed out".to_string(),
metadata,
ToolPolicyDecisionRecord::approval("approval timeout"),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Timeout,
reason: Some("approval timeout".to_string()),
modified_arguments: None,
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
}
}
let hitl_lang_ctx = self.build_hitl_language_context();
if let Some(ref hitl_engine) = self.hitl_engine {
let check_result = self
.observe_purpose(
ObservationPurpose::HitlLocalization,
hitl_engine.check_tool_with_localization(
&canonical_id,
&executed_arguments,
&hitl_lang_ctx,
self.approval_handler.as_ref(),
Some(&self.llm_registry),
),
)
.await?;
if check_result.is_required() {
match self.request_hitl_approval(check_result).await? {
ApprovalResult::Approved => {
merge_approved_record(&mut approval_record);
}
ApprovalResult::Modified { changes } => {
if let Some(obj) = executed_arguments.as_object_mut() {
for (key, value) in changes {
obj.insert(key, value);
}
}
let modified_security = security_engine
.validate_tool_execution_with_bindings(
&canonical_id,
&executed_arguments,
&bindings,
)
.await?;
if !matches!(
modified_security,
SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
) {
let reason = modified_security
.reason()
.unwrap_or("modified arguments failed policy")
.to_string();
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments.clone(),
started_at,
start,
false,
false,
reason.clone(),
metadata,
ToolPolicyDecisionRecord::deny(reason),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(executed_arguments),
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(executed_arguments.clone()),
});
}
ApprovalResult::Rejected { reason } => {
let reason = reason.unwrap_or_else(|| "rejected".to_string());
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Approval rejected: {}", reason),
metadata,
ToolPolicyDecisionRecord::approval(reason.clone()),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Rejected,
reason: Some(reason),
modified_arguments: None,
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
ApprovalResult::Timeout => {
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
"Approval timed out".to_string(),
metadata,
ToolPolicyDecisionRecord::approval("approval timeout"),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Timeout,
reason: Some("approval timeout".to_string()),
modified_arguments: None,
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
}
}
let condition_check = self
.observe_purpose(
ObservationPurpose::HitlLocalization,
hitl_engine.check_conditions_with_localization(
&executed_arguments,
&hitl_lang_ctx,
self.approval_handler.as_ref(),
Some(&self.llm_registry),
),
)
.await?;
if condition_check.is_required() {
match self.request_hitl_approval(condition_check).await? {
ApprovalResult::Approved => {
merge_approved_record(&mut approval_record);
}
ApprovalResult::Modified { changes } => {
if let Some(obj) = executed_arguments.as_object_mut() {
for (key, value) in changes {
obj.insert(key, value);
}
}
let modified_security = security_engine
.validate_tool_execution_with_bindings(
&canonical_id,
&executed_arguments,
&bindings,
)
.await?;
if !matches!(
modified_security,
SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
) {
let reason = modified_security
.reason()
.unwrap_or("modified arguments failed policy")
.to_string();
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
reason.clone(),
metadata,
ToolPolicyDecisionRecord::deny(reason),
approval_record,
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
approval_record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(executed_arguments.clone()),
});
}
ApprovalResult::Rejected { reason } => {
let reason = reason.unwrap_or_else(|| "rejected".to_string());
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Approval rejected: {}", reason),
metadata,
ToolPolicyDecisionRecord::approval(reason.clone()),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Rejected,
reason: Some(reason),
modified_arguments: None,
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
ApprovalResult::Timeout => {
let record = self.record_from_parts(
&request,
canonical_id,
executed_arguments,
started_at,
start,
false,
false,
"Approval timed out".to_string(),
metadata,
ToolPolicyDecisionRecord::approval("approval timeout"),
Some(ToolApprovalRecord {
status: ToolApprovalStatus::Timeout,
reason: Some("approval timeout".to_string()),
modified_arguments: None,
}),
false,
false,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
}
}
}
executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
&canonical_id,
&executed_arguments,
&bindings,
);
if let Some(record) = approval_record.as_mut()
&& matches!(record.status, ToolApprovalStatus::Modified)
{
record.modified_arguments = Some(executed_arguments.clone());
}
let binding_security_result = security_engine
.validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
.await?;
let approval_confirmation_required = matches!(
binding_security_result,
SecurityCheckResult::RequireConfirmation { .. }
) || security_engine
.classification_approval_message(
&canonical_id,
&resolved.tool.classify_call(&executed_arguments),
)
.is_some();
let approval_binding = approval_record.as_ref().and_then(|record| {
matches!(
record.status,
ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
)
.then(|| ToolApprovalBinding {
canonical_id: canonical_id.clone(),
arguments: executed_arguments.clone(),
confirmation_required: approval_confirmation_required,
policy_version: security_engine.policy_version(),
runtime_control_version: approval_control_snapshot.version,
state_generation: initial_scope_snapshot.state_generation,
reviewed_tool: Arc::clone(&resolved.tool),
})
});
let control_snapshot = self.runtime_safety_snapshot();
let resolved = self.tools.resolve(&request.requested_name);
let registry_version = self.tools.version();
let mut versions = ToolDecisionVersions {
policy: control_snapshot.tool_security.policy_version(),
registry: registry_version,
runtime_control: control_snapshot.version,
state: None,
};
metadata.insert(
"runtime_scope_snapshot".to_string(),
serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
);
let resolved = match resolved {
Some(resolved) => resolved,
None => {
let reason = format!(
"Tool '{}' became unavailable after approval",
request.requested_name
);
let record = self.record_from_parts_at(
&request,
request.requested_name.clone(),
executed_arguments,
started_at,
start,
false,
false,
reason.clone(),
metadata,
ToolPolicyDecisionRecord::unavailable(reason),
approval_record,
false,
false,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
};
let canonical_id = resolved.identity.canonical_id.clone();
if let Some(reason) =
fallback_state.final_rejection_reason(&admitted_canonical_id, &canonical_id)
{
metadata.insert(
"fallback_chain".to_string(),
serde_json::to_value(&fallback_state.visited_canonical_ids).unwrap_or(Value::Null),
);
metadata.insert(
"final_resolved_canonical_id".to_string(),
Value::String(canonical_id),
);
let record = self.record_from_parts_at(
&request,
admitted_canonical_id,
executed_arguments,
started_at,
start,
false,
false,
format!("Denied: {reason}"),
metadata,
ToolPolicyDecisionRecord::deny(reason),
approval_record,
false,
false,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let bindings = resolved.tool.policy_bindings();
let final_arguments = control_snapshot
.tool_security
.prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
if let Some(record) = approval_record.as_mut()
&& matches!(record.status, ToolApprovalStatus::Modified)
{
record.modified_arguments = Some(final_arguments.clone());
}
let classification = resolved.tool.classify_call(&final_arguments);
let safety = resolved.tool.safety_metadata();
let security_engine = control_snapshot.tool_security;
let tool_config = self.recovery_manager.get_tool_config(&canonical_id).clone();
let recovery_timeout_ms = self.recovery_manager.get_tool_timeout(&canonical_id);
metadata.insert(
"classification".to_string(),
serde_json::to_value(&classification).unwrap_or(Value::Null),
);
let (limits, timeout) = match Self::effective_tool_limits(
&security_engine,
&canonical_id,
&safety,
&classification,
recovery_timeout_ms,
) {
Ok(effective) => effective,
Err(error) => {
let reason = error.to_string();
metadata.insert(
"configuration_error".to_string(),
Value::String(reason.clone()),
);
let record = self.record_from_parts_at(
&request,
canonical_id,
final_arguments,
started_at,
start,
false,
false,
format!("Denied: {reason}"),
metadata,
ToolPolicyDecisionRecord::deny(reason),
approval_record,
false,
false,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
};
let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
let resource_lock_keys =
tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
metadata.insert(
"effective_limits".to_string(),
serde_json::to_value(&limits).unwrap_or(Value::Null),
);
metadata.insert(
"resource_lock_keys".to_string(),
serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
);
if policy_snapshot.is_null() {
metadata.remove("policy_snapshot");
} else {
metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
}
let final_denial = |canonical_id: String,
output: String,
policy: ToolPolicyDecisionRecord,
metadata: HashMap<String, Value>,
decision_versions: ToolDecisionVersions| {
self.record_from_parts_at(
&request,
canonical_id,
final_arguments.clone(),
started_at,
start,
false,
false,
output,
metadata,
policy,
approval_record.clone(),
false,
false,
decision_versions,
)
};
if control_snapshot.emergency_deny {
let reason = "Tool execution is disabled by runtime control".to_string();
let record = final_denial(
canonical_id,
reason.clone(),
ToolPolicyDecisionRecord::deny(reason),
metadata,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let available_snapshot = self
.get_available_tool_ids_snapshot_for_scope(
control_snapshot.tool_scope_override.as_deref(),
)
.await?;
versions.state = available_snapshot.state_generation;
metadata.insert(
"available_tool_ids_snapshot".to_string(),
serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
);
metadata.insert(
"state_generation_snapshot".to_string(),
serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
);
if !available_snapshot
.tool_ids
.iter()
.any(|tool_id| tool_id == &canonical_id)
{
let reason = format!(
"Tool '{}' is not available in the final runtime scope",
canonical_id
);
let record = final_denial(
canonical_id,
reason.clone(),
ToolPolicyDecisionRecord::deny(reason),
metadata,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let final_security_result = security_engine
.validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
.await?;
match &final_security_result {
SecurityCheckResult::Block { reason } => {
let record = final_denial(
canonical_id,
format!("Denied: {}", reason),
ToolPolicyDecisionRecord::deny(reason.clone()),
metadata,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
SecurityCheckResult::Unavailable { reason } => {
let record = final_denial(
canonical_id,
format!("Unavailable: {}", reason),
ToolPolicyDecisionRecord::unavailable(reason.clone()),
metadata,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
SecurityCheckResult::Warn { message } => {
warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
}
SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
}
let final_confirmation_required = matches!(
final_security_result,
SecurityCheckResult::RequireConfirmation { .. }
) || security_engine
.classification_approval_message(&canonical_id, &classification)
.is_some();
let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
binding.is_stale(
&canonical_id,
&final_arguments,
final_confirmation_required,
versions,
&resolved.tool,
)
});
if stale_approval {
let reason = "Approval became stale before final admission".to_string();
let record = final_denial(
canonical_id,
reason.clone(),
ToolPolicyDecisionRecord::deny(reason),
metadata,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
if final_confirmation_required && approval_binding.is_none() {
let reason = "Final policy requires fresh approval".to_string();
let record = final_denial(
canonical_id,
reason.clone(),
ToolPolicyDecisionRecord::approval(reason),
metadata,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
if let Some((_, reason)) = self.host_tool_unavailability(&canonical_id) {
let record = final_denial(
canonical_id,
reason.to_string(),
ToolPolicyDecisionRecord::unavailable(reason),
metadata,
versions,
);
self.finish_tool_record(&record).await;
return Ok(record);
}
let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
else {
let reason = "Tool execution cancelled while waiting for resource locks".to_string();
let mut record = final_denial(
canonical_id,
reason.clone(),
ToolPolicyDecisionRecord::deny(reason),
metadata,
versions,
);
record.cancelled = true;
record.cancellation_reason = Some("runtime control cancellation".to_string());
self.finish_tool_record(&record).await;
return Ok(record);
};
let admission = self.admit_tool_execution(
versions.runtime_control,
versions.policy,
versions.state,
&canonical_id,
);
if !matches!(admission, SecurityCheckResult::Allow) {
let latest_control = self.runtime_safety_snapshot();
let reason = admission
.reason()
.unwrap_or("tool admission was denied")
.to_string();
let policy = if admission.is_unavailable() {
ToolPolicyDecisionRecord::unavailable(reason.clone())
} else {
ToolPolicyDecisionRecord::deny(reason.clone())
};
let record = self.record_from_parts_at(
&request,
canonical_id,
final_arguments,
started_at,
start,
false,
false,
reason,
metadata,
policy,
approval_record,
false,
false,
ToolDecisionVersions {
policy: latest_control.tool_security.policy_version(),
registry: versions.registry,
runtime_control: latest_control.version,
state: self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation()),
},
);
self.finish_tool_record_after_resource_guards(resource_guards, &record)
.await;
return Ok(record);
}
let executed_arguments = final_arguments;
let turn_actor = current_turn_actor_context();
let actor = ToolActorContext {
actor_id: turn_actor
.as_ref()
.and_then(|context| context.effective_actor_id().map(str::to_string))
.or_else(|| self.actor_id()),
origin_actor_id: turn_actor
.as_ref()
.and_then(|context| context.origin_actor_id.clone()),
sender_agent_id: turn_actor
.as_ref()
.and_then(|context| context.sender_agent_id.clone()),
};
let tool_context = ToolExecutionContext {
requested_name: request.requested_name.clone(),
canonical_id: canonical_id.clone(),
display_name: resolved.identity.display_name.clone(),
provider_id: resolved.identity.provider_id.clone(),
registry_version: versions.registry,
policy_version: versions.policy,
runtime_control_version: versions.runtime_control,
call_id: request.call_id.clone(),
source: request.source.clone(),
actor,
cancellation: ToolCancellationToken::new(
Arc::clone(&self.runtime_control.emergency_deny),
Some("runtime control cancellation".to_string()),
),
started_at,
deadline: None,
permission: ToolPolicyDecisionRecord::allow(),
approval: approval_record.clone(),
classification: classification.clone(),
safety,
limits: limits.clone(),
policy_snapshot,
custom_config: security_engine.custom_config(&canonical_id),
};
let (mut result, timed_out, cancelled, invoked) = self
.run_tool_with_retries(
&canonical_id,
resolved.tool.clone(),
executed_arguments.clone(),
tool_context,
timeout,
tool_config.max_retries,
)
.await?;
let fallback_tool = if !result.success && !cancelled {
match &tool_config.on_failure {
ToolFailureAction::Skip => {
result = ToolResult::ok(format!(
"{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
canonical_id
));
None
}
ToolFailureAction::Fallback { fallback_tool } => Some(fallback_tool.clone()),
ToolFailureAction::ReportError => None,
}
} else {
None
};
let output_cap = limits.max_output_chars;
let (output, output_truncated) =
Self::truncate_tool_output(result.output.clone(), output_cap);
if let Some(result_metadata) = result.metadata {
metadata.extend(result_metadata);
}
let mut record = self.record_from_parts_at(
&request,
canonical_id,
executed_arguments,
started_at,
start,
invoked,
result.success,
output,
metadata,
ToolPolicyDecisionRecord::allow(),
approval_record,
timed_out,
output_truncated,
versions,
);
record.cancelled = cancelled;
if cancelled {
record.cancellation_reason = Some("runtime control cancellation".to_string());
}
if let Some(fallback_tool) = fallback_tool {
let fallback_arguments = record.executed_arguments.clone();
let original_tool = record.canonical_id.clone();
self.finish_tool_record_after_resource_guards(resource_guards, &record)
.await;
let fallback_request = ToolExecutionRequest::new(
request.call_id.clone(),
fallback_tool,
fallback_arguments,
ToolCallSource::Fallback { original_tool },
);
return Box::pin(self.execute_tool_record_inner(fallback_request, fallback_state))
.await;
}
self.finish_tool_record_after_resource_guards(resource_guards, &record)
.await;
Ok(record)
}
#[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
let record = self
.execute_tool_record(ToolExecutionRequest::new(
tool_call.id.clone(),
tool_call.name.clone(),
tool_call.arguments.clone(),
ToolCallSource::Model,
))
.await?;
if record.success {
Ok(record.model_output_string())
} else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
Err(AgentError::HITLRejected(record.model_output_string()))
} else {
Err(AgentError::Tool(record.model_output_string()))
}
}
async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
let Some(ref router) = self.skill_router else {
return Ok(None);
};
let available_skills = self.get_available_skills();
if available_skills.is_empty() {
return Ok(None);
}
let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
let Some(skill_id) = self
.observe_purpose(
ObservationPurpose::SkillRouting,
router.select_skill_filtered(input, &skill_ids),
)
.await?
else {
return Ok(None);
};
let skill = router
.get_skill(&skill_id)
.cloned()
.ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
info!(skill_id = %skill_id, "Skill selected");
Ok(Some(SkillCandidate::new(skill_id, skill)))
}
async fn commit_skill_candidate_route_result(
&self,
candidate: SkillCandidate,
input: &str,
) -> Result<SkillRouteResult> {
let skill_id = candidate.skill_id;
let skill = candidate.skill;
let expected_state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
if let Some(ref skill_disambig) = skill.disambiguation
&& skill_disambig.enabled.unwrap_or(false)
&& let Some(ref disambiguator) = self.disambiguation_manager
{
let context = self.build_disambiguation_context().await?;
let state_override = self
.state_machine
.as_ref()
.and_then(|sm| sm.current_definition())
.and_then(|def| def.disambiguation.clone());
let disambiguation_result = self
.observe_purpose(
ObservationPurpose::DisambiguationDetection,
disambiguator.process_input_with_override(
input,
&context,
state_override.as_ref(),
Some(skill_disambig),
),
)
.await?;
let current_state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
if current_state_generation != expected_state_generation
|| self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
{
disambiguator.clear_pending().await;
*self.pending_skill_id.write() = None;
return Err(AgentError::Other(
"State or reset ownership changed during skill disambiguation".to_string(),
));
}
match disambiguation_result {
DisambiguationResult::Clear => {
debug!(skill_id = %skill_id, "Skill disambiguation: clear");
}
DisambiguationResult::NeedsClarification {
question,
detection,
} => {
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
info!(
skill_id = %skill_id,
ambiguity_type = ?detection.ambiguity_type,
confidence = detection.confidence,
"Skill requires clarification before execution"
);
*self.pending_skill_id.write() = Some(skill_id.clone());
let response = AgentResponse::new(&question.question).with_metadata(
"disambiguation",
serde_json::json!({
"status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
"skill_id": skill_id,
"options": question.options,
"clarifying": question.clarifying,
"detection": {
"type": detection.ambiguity_type,
"confidence": detection.confidence,
"what_is_unclear": detection.what_is_unclear,
}
}),
);
drop(admission);
return Ok(SkillRouteResult::NeedsClarification {
response,
ownership: Some(DisambiguationOwnership {
epoch: expected_disambiguation_epoch,
state_generation: expected_state_generation,
}),
});
}
DisambiguationResult::Clarified { enriched_input, .. } => {
info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
drop(admission);
let content = self.execute_skill(&skill, &enriched_input).await?;
return Ok(SkillRouteResult::Response { skill_id, content });
}
DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
info!(skill_id = %skill_id, "Skill disambiguation best guess");
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
drop(admission);
let content = self.execute_skill(&skill, &enriched_input).await?;
return Ok(SkillRouteResult::Response { skill_id, content });
}
DisambiguationResult::GiveUp { reason } => {
warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
let apology = self
.generate_localized_apology(
"Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I'm sorry, I couldn't understand your request: {}", reason)
});
return Ok(SkillRouteResult::NeedsClarification {
response: AgentResponse::new(&apology),
ownership: None,
});
}
DisambiguationResult::Escalate { reason } => {
info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
let apology = self
.generate_localized_apology(
"Explain briefly that you're transferring the user to a human agent for help.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I need human assistance to help with your request: {}", reason)
});
return Ok(SkillRouteResult::NeedsClarification {
response: AgentResponse::new(&apology),
ownership: None,
});
}
DisambiguationResult::Abandoned { .. } => {
debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
return Ok(SkillRouteResult::NoMatch);
}
}
}
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
drop(admission);
let content = self.execute_skill(&skill, input).await?;
Ok(SkillRouteResult::Response { skill_id, content })
}
async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
if let Some(candidate) = self.select_skill_candidate(input).await? {
self.commit_skill_candidate_route_result(candidate, input)
.await
} else {
Ok(SkillRouteResult::NoMatch)
}
}
async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
if let Some(ref executor) = self.skill_executor {
let skill_reasoning = self.get_skill_reasoning_config(skill);
let skill_reflection = self.get_skill_reflection_config(skill);
debug!(
skill_id = %skill.id,
reasoning_mode = ?skill_reasoning.mode,
reflection_enabled = ?skill_reflection.enabled,
"Skill reasoning/reflection config"
);
let response = self
.observe_purpose(
ObservationPurpose::SkillPrompt,
executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
)
.await?;
if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
let should_reflect = self
.should_reflect_with_config(input, &response, &skill_reflection)
.await?;
if should_reflect {
let evaluated = self
.evaluate_and_retry_with_config(input, response, &skill_reflection)
.await?;
return Ok(evaluated);
}
}
return Ok(response);
}
Err(AgentError::Skill(
"No skill executor configured".to_string(),
))
}
async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
let skill = self
.skill_router
.as_ref()
.and_then(|r| r.get_skill(skill_id).cloned())
.ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
self.execute_skill(&skill, input).await
}
async fn should_reflect_with_config(
&self,
input: &str,
response: &str,
config: &ReflectionConfig,
) -> Result<bool> {
if !config.requires_evaluation() {
return Ok(false);
}
if config.is_enabled() {
return Ok(true);
}
let evaluator_llm = config
.evaluator_llm
.as_ref()
.and_then(|alias| self.llm_registry.get(alias).ok())
.or_else(|| self.llm_registry.router().ok())
.or_else(|| self.llm_registry.default().ok());
let Some(llm) = evaluator_llm else {
return Ok(false);
};
let response_preview: String = response.chars().take(500).collect();
let prompt = format!(
r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
User query: "{}"
Response: "{}"
Answer YES or NO only."#,
input, response_preview
);
let messages = vec![ChatMessage::user(&prompt)];
let result = self
.observe_purpose(
ObservationPurpose::ReflectionDecision,
llm.complete(&messages, None),
)
.await;
match result {
Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
Err(_) => Ok(false),
}
}
async fn evaluate_and_retry_with_config(
&self,
input: &str,
mut response: String,
config: &ReflectionConfig,
) -> Result<String> {
let llm = self.get_state_llm()?;
let mut attempts = 0u32;
let max_retries = config.max_retries;
loop {
let evaluation = self
.evaluate_response_with_config(input, &response, config)
.await?;
if evaluation.passed || attempts >= max_retries {
info!(
passed = evaluation.passed,
confidence = evaluation.confidence,
attempts = attempts + 1,
"Skill reflection evaluation complete"
);
return Ok(response);
}
debug!(
attempt = attempts + 1,
failed_criteria = evaluation.failed_criteria().count(),
"Skill response did not meet criteria, retrying"
);
let feedback: Vec<String> = evaluation
.failed_criteria()
.map(|c| format!("- {}", c.criterion))
.collect();
let retry_prompt = format!(
"Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
feedback.join("\n"),
input
);
let messages = vec![ChatMessage::user(&retry_prompt)];
let retry_response = self
.observe_purpose(
ObservationPurpose::ReflectionEvaluation,
llm.complete(&messages, None),
)
.await
.map_err(|e| AgentError::LLM(e.to_string()))?;
response = retry_response.content.trim().to_string();
attempts += 1;
}
}
async fn evaluate_response_with_config(
&self,
input: &str,
response: &str,
config: &ReflectionConfig,
) -> Result<EvaluationResult> {
let evaluator_llm = config
.evaluator_llm
.as_ref()
.and_then(|alias| self.llm_registry.get(alias).ok())
.or_else(|| self.llm_registry.router().ok())
.or_else(|| self.llm_registry.default().ok())
.ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
let criteria = &config.criteria;
let criteria_list = criteria
.iter()
.enumerate()
.map(|(i, c)| format!("{}. {}", i + 1, c))
.collect::<Vec<_>>()
.join("\n");
let prompt = format!(
r#"Evaluate this response against the criteria.
User query: "{}"
Response to evaluate: "{}"
Criteria:
{}
For each criterion, respond with:
- criterion number
- PASS or FAIL
- brief reason
Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
Format:
1. PASS/FAIL - reason
2. PASS/FAIL - reason
...
CONFIDENCE: 0.X
OVERALL: PASS/FAIL"#,
input, response, criteria_list
);
let messages = vec![ChatMessage::user(&prompt)];
let eval_response = self
.observe_purpose(
ObservationPurpose::ReflectionEvaluation,
evaluator_llm.complete(&messages, None),
)
.await
.map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
let content = eval_response.content.to_uppercase();
let llm_pass = content.contains("OVERALL: PASS");
let confidence = content
.lines()
.find(|l| l.contains("CONFIDENCE:"))
.and_then(|l| {
l.split(':')
.nth(1)
.and_then(|v| v.trim().parse::<f32>().ok())
})
.unwrap_or(if llm_pass { 0.8 } else { 0.4 });
let overall_pass = llm_pass && confidence >= config.pass_threshold;
let mut criteria_results = Vec::new();
for (i, criterion) in criteria.iter().enumerate() {
let line_marker = format!("{}.", i + 1);
let passed = eval_response
.content
.lines()
.find(|l| l.contains(&line_marker))
.map(|l| l.to_uppercase().contains("PASS"))
.unwrap_or(overall_pass);
if passed {
criteria_results.push(CriterionResult::pass(criterion));
} else {
criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
}
}
Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
}
async fn process_input(&self, input: &str) -> Result<ProcessData> {
if let Some(processor) = self.get_state_process_processor() {
let purpose = observation_purpose_for_process(processor.input_purpose_hint());
return self
.observe_purpose(purpose, processor.process_input(input))
.await;
}
if let Some(ref processor) = self.process_processor {
let purpose = observation_purpose_for_process(processor.input_purpose_hint());
self.observe_purpose(purpose, processor.process_input(input))
.await
} else {
Ok(ProcessData::new(input))
}
}
async fn process_output(
&self,
output: &str,
input_context: &std::collections::HashMap<String, serde_json::Value>,
) -> Result<ProcessData> {
if let Some(processor) = self.get_state_process_processor() {
let purpose = observation_purpose_for_process(processor.output_purpose_hint());
return self
.observe_purpose(purpose, processor.process_output(output, input_context))
.await;
}
if let Some(ref processor) = self.process_processor {
let purpose = observation_purpose_for_process(processor.output_purpose_hint());
self.observe_purpose(purpose, processor.process_output(output, input_context))
.await
} else {
Ok(ProcessData::new(output))
}
}
fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
let sm = self.state_machine.as_ref()?;
let def = sm.current_definition()?;
let config = def.process.as_ref()?;
let mut processor = ProcessProcessor::new(config.clone());
if let Some(ref registry) = Some(self.llm_registry.clone()) {
processor = processor.with_llm_registry(registry.clone());
}
processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
Some(processor)
}
async fn check_turn_timeout(&self) -> Result<()> {
let Some(ref sm) = self.state_machine else {
return Ok(());
};
let Some(timeout_state) = sm.check_timeout() else {
return Ok(());
};
let claim_admission = self.disambiguation_admission.write().await;
if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
return Ok(());
}
let Some(reservation) = self.reserve_state_transition() else {
return Ok(());
};
let from_state = sm.current();
let expected_state_generation = sm.generation();
let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
let history_before = sm.history();
drop(claim_admission);
self.execute_state_exit_actions(&from_state).await;
let admission = self.disambiguation_admission.write().await;
if sm.current() != from_state
|| sm.generation() != expected_state_generation
|| self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
|| sm.check_timeout().as_deref() != Some(timeout_state.as_str())
{
return Ok(());
}
sm.transition_to(&timeout_state, "max_turns exceeded")?;
self.invalidate_pending_confirmation("state_timeout").await;
let entered = sm.current();
let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
drop(admission);
self.execute_state_enter_actions(&entered, is_reentry).await;
drop(reservation);
info!(to = %entered, "Timeout transition");
Ok(())
}
fn increment_turn(&self) {
if let Some(ref sm) = self.state_machine {
sm.increment_turn();
}
}
fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
let sm = self.state_machine.as_ref()?;
let current = sm.current();
let transitions: Vec<_> = sm
.auto_transitions()
.into_iter()
.filter(|t| match t.cooldown_turns {
Some(cd) if cd > 0 => {
let resolved = sm.config().resolve_full_path(¤t, &t.to);
!sm.is_on_cooldown(&resolved, cd)
}
_ => true,
})
.collect();
Some((transitions, current))
}
fn transition_reason(transition: &Transition) -> String {
if transition.when.is_empty() {
"guard condition met".to_string()
} else {
transition.when.clone()
}
}
fn build_transition_context(
&self,
user_message: &str,
response: &str,
current_state: &str,
staged: Option<&HashMap<String, Value>>,
) -> TransitionContext {
let context_map = staged
.map(|writes| self.build_context_with_staged(writes))
.unwrap_or_else(|| self.build_context_with_overlays());
TransitionContext::new(user_message, response, current_state).with_context(context_map)
}
async fn select_transition_candidate(
&self,
user_message: &str,
response: &str,
) -> Result<Option<TransitionCandidate>> {
let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
return Ok(None);
};
let transitions: Vec<Transition> = transitions
.into_iter()
.filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
.collect();
if transitions.is_empty() {
return Ok(None);
}
let Some(evaluator) = self.transition_evaluator.as_ref() else {
return Ok(None);
};
let context = self.build_transition_context(user_message, response, ¤t_state, None);
let selected = self
.observe_purpose(
ObservationPurpose::StateTransitionEvaluation,
evaluator.select_transition(&transitions, &context),
)
.await?;
Ok(selected.map(|index| {
let transition = transitions[index].clone();
TransitionCandidate::new(
current_state,
transition.clone(),
Self::transition_reason(&transition),
)
}))
}
fn select_deterministic_transition_candidate(
&self,
user_message: &str,
current_state: &str,
transitions: &[Transition],
staged: &HashMap<String, Value>,
) -> Option<TransitionCandidate> {
let context = self.build_transition_context(user_message, "", current_state, Some(staged));
for transition in transitions {
if let Some(guard) = transition.guard.as_ref()
&& evaluate_guard(guard, &context)
{
return Some(TransitionCandidate::new(
current_state,
transition.clone(),
Self::transition_reason(transition),
));
}
}
let resolved_intent = context
.context
.get("resolved_intent")
.and_then(Value::as_str)
.filter(|value| !value.is_empty());
if let Some(resolved_intent) = resolved_intent {
for transition in transitions {
if transition.intent.as_deref() == Some(resolved_intent) {
return Some(TransitionCandidate::new(
current_state,
transition.clone(),
Self::transition_reason(transition),
));
}
}
}
None
}
async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
.await
}
async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
let approved = self.check_state_hitl(Some(from_state), target).await?;
if !approved {
info!(to = %target, "State transition rejected by HITL");
}
Ok(approved)
}
async fn apply_transition_target(
&self,
from_state: &str,
target: &str,
reason: &str,
staged: Option<&HashMap<String, Value>>,
) -> Result<bool> {
let Some(ref sm) = self.state_machine else {
return Ok(false);
};
let claim_admission = self.disambiguation_admission.write().await;
if sm.current() != from_state {
return Ok(false);
}
let Some(reservation) = self.reserve_state_transition() else {
return Ok(false);
};
let expected_state_generation = sm.generation();
let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
let history_before = sm.history();
drop(claim_admission);
self.execute_state_exit_actions(from_state).await;
let admission = self.disambiguation_admission.write().await;
if sm.current() != from_state
|| sm.generation() != expected_state_generation
|| self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
{
return Ok(false);
}
sm.transition_to(target, reason)?;
self.invalidate_pending_confirmation("state_transition")
.await;
sm.reset_no_transition();
if let Some(staged) = staged {
self.commit_staged_context_writes(staged);
}
let entered = sm.current();
let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
drop(admission);
self.execute_state_enter_actions(&entered, is_reentry).await;
drop(reservation);
self.hooks
.on_state_transition(Some(from_state), &entered, reason)
.await;
info!(from = %from_state, to = %entered, "State transition");
Ok(true)
}
async fn commit_transition_target(
&self,
from_state: &str,
target: &str,
reason: &str,
) -> Result<bool> {
if !self.approve_transition_target(from_state, target).await? {
return Ok(false);
}
self.apply_transition_target(from_state, target, reason, None)
.await
}
async fn apply_pre_response_transition_candidate(
&self,
candidate: &TransitionCandidate,
staged: &HashMap<String, Value>,
processed_input: &str,
) -> Result<bool> {
self.commit_root_user_message(processed_input).await?;
self.apply_transition_target(
&candidate.from_state,
candidate.target(),
&candidate.reason,
Some(staged),
)
.await
}
async fn commit_pre_response_transition_candidate(
&self,
candidate: &TransitionCandidate,
staged: &HashMap<String, Value>,
processed_input: &str,
) -> Result<bool> {
if !self
.approve_transition_target(&candidate.from_state, candidate.target())
.await?
{
return Ok(false);
}
self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
.await
}
async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
let Some(ref sm) = self.state_machine else {
return Ok(false);
};
sm.increment_no_transition();
let Some(fallback) = sm.check_fallback() else {
return Ok(false);
};
self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
.await
}
async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
return Ok(false);
};
if transitions.is_empty() {
return Ok(false);
}
if let Some(candidate) = self
.select_transition_candidate(user_message, response)
.await?
{
return self.commit_transition_candidate(&candidate).await;
}
self.handle_transition_miss(¤t_state).await
}
async fn try_pre_response_transition(
&self,
processed_input: &str,
) -> Result<Option<AgentResponse>> {
let optimization = &self.runtime_config.optimization;
if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
return Ok(None);
}
let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
return Ok(None);
};
let eligible: Vec<Transition> = transitions
.into_iter()
.filter(|transition| !transition.requires_response)
.filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
.collect();
if eligible.is_empty() {
return Ok(None);
}
let empty_staged = HashMap::new();
let mut extracted_staged: Option<HashMap<String, Value>> = None;
let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
for transition in &eligible {
let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
let staged_for_eval = if use_extractors {
if extracted_staged.is_none() {
extracted_staged =
Some(self.run_context_extractors_staged(processed_input).await);
}
extracted_staged.as_ref().unwrap_or(&empty_staged)
} else {
&empty_staged
};
if let Some(candidate) = self.select_deterministic_transition_candidate(
processed_input,
¤t_state,
std::slice::from_ref(transition),
staged_for_eval,
) {
let staged_for_commit = if use_extractors {
staged_for_eval.clone()
} else {
HashMap::new()
};
selected = Some((candidate, staged_for_commit));
break;
}
}
let Some((candidate, staged)) = selected else {
return Ok(None);
};
if !self
.commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
.await?
{
return Ok(None);
}
self.redispatch_current_state(processed_input)
.await
.map(Some)
}
async fn try_speculative_branches(
&self,
processed_input: &str,
input_context: &HashMap<String, Value>,
) -> Result<Option<AgentResponse>> {
let optimization = &self.runtime_config.optimization;
if !optimization.enabled {
return Ok(None);
}
let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
if !matches!(
effective_reasoning_mode,
ReasoningMode::None | ReasoningMode::Auto
) {
return Ok(None);
}
let mut transition_enabled =
optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
let mut skill_enabled = optimization.speculative_skill_routing
&& self.skill_router.is_some()
&& self.pending_skill_id.read().is_none();
let mut reasoning_enabled = optimization.speculative_reasoning_auto
&& matches!(effective_reasoning_mode, ReasoningMode::Auto);
if matches!(effective_reasoning_mode, ReasoningMode::Auto)
&& (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
{
return Ok(None);
}
if !transition_enabled && !skill_enabled && !reasoning_enabled {
return Ok(None);
}
let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
let mut speculative_call_slots = optimization
.max_speculative_llm_calls_per_turn
.saturating_sub(1);
if reasoning_enabled {
if optional_slots == 0 || speculative_call_slots == 0 {
return Ok(None);
}
optional_slots -= 1;
speculative_call_slots -= 1;
}
if transition_enabled {
if optional_slots == 0 {
transition_enabled = false;
} else {
optional_slots -= 1;
}
}
if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
skill_enabled = false;
}
if !transition_enabled && !skill_enabled && !reasoning_enabled {
return Ok(None);
}
let main_kind = if transition_enabled {
RuntimeOptimizationKind::ParallelStateTransition
} else if skill_enabled {
RuntimeOptimizationKind::SpeculativeSkillRouting
} else {
RuntimeOptimizationKind::SpeculativeReasoningAuto
};
if !self.reserve_active_speculative_llm_call(main_kind) {
return Ok(None);
}
let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
let main_branch = RuntimeBranch::new(
RuntimeTaskPurpose::MainResponse,
main_kind,
RuntimeTaskPriority::Normal,
RuntimeCommitBehavior::FinalResponse,
);
let transition_branch = RuntimeBranch::new(
RuntimeTaskPurpose::StateTransition,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeTaskPriority::Critical,
RuntimeCommitBehavior::TransitionDecision,
);
let skill_branch = RuntimeBranch::new(
RuntimeTaskPurpose::SkillRouting,
RuntimeOptimizationKind::SpeculativeSkillRouting,
RuntimeTaskPriority::High,
RuntimeCommitBehavior::SkillSelection,
);
let reasoning_branch = RuntimeBranch::new(
RuntimeTaskPurpose::ReasoningJudge,
RuntimeOptimizationKind::SpeculativeReasoningAuto,
RuntimeTaskPriority::Normal,
RuntimeCommitBehavior::ReasoningDecision,
);
let main_id = main_branch.branch_id();
let transition_id = transition_branch.branch_id();
let skill_id = skill_branch.branch_id();
let reasoning_id = reasoning_branch.branch_id();
let main_id_for_future = main_id.clone();
if !branch_set.schedule(
main_branch,
Box::pin(async move {
match crate::optimization::observability::with_branch_observation(
&main_id_for_future,
main_kind,
RuntimeCommitBehavior::FinalResponse,
self.generate_main_response_draft(processed_input, &ReasoningMode::None),
)
.await
{
Ok(draft) => RuntimeBranchResult::MainDraft(draft),
Err(error) => RuntimeBranchResult::Failed(error),
}
}),
) {
return Ok(None);
}
if transition_enabled {
let transition_id_for_future = transition_id.clone();
if !branch_set.schedule(
transition_branch,
Box::pin(async move {
match crate::optimization::observability::with_branch_observation(
&transition_id_for_future,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
self.select_parallel_transition_candidate(processed_input),
)
.await
{
Ok(ParallelTransitionSelection::Candidate(candidate)) => {
RuntimeBranchResult::Transition(Some(candidate))
}
Ok(ParallelTransitionSelection::NoMatch) => {
RuntimeBranchResult::Transition(None)
}
Ok(ParallelTransitionSelection::ReservationExhausted) => {
RuntimeBranchResult::Cancelled
}
Err(error) => RuntimeBranchResult::Failed(error),
}
}),
) {
transition_enabled = false;
}
}
if skill_enabled {
let skill_id_for_future = skill_id.clone();
if !branch_set.schedule(
skill_branch,
Box::pin(async move {
if !self.reserve_active_speculative_llm_call(
RuntimeOptimizationKind::SpeculativeSkillRouting,
) {
return RuntimeBranchResult::Cancelled;
}
match crate::optimization::observability::with_branch_observation(
&skill_id_for_future,
RuntimeOptimizationKind::SpeculativeSkillRouting,
RuntimeCommitBehavior::SkillSelection,
self.select_skill_candidate(processed_input),
)
.await
{
Ok(candidate) => RuntimeBranchResult::Skill(candidate),
Err(error) => RuntimeBranchResult::Failed(error),
}
}),
) {
skill_enabled = false;
}
}
if reasoning_enabled {
let reasoning_id_for_future = reasoning_id.clone();
if !branch_set.schedule(
reasoning_branch,
Box::pin(async move {
if !self.reserve_active_speculative_llm_call(
RuntimeOptimizationKind::SpeculativeReasoningAuto,
) {
return RuntimeBranchResult::Cancelled;
}
match crate::optimization::observability::with_branch_observation(
&reasoning_id_for_future,
RuntimeOptimizationKind::SpeculativeReasoningAuto,
RuntimeCommitBehavior::ReasoningDecision,
self.determine_reasoning_mode_strict(processed_input),
)
.await
{
Ok(mode) => RuntimeBranchResult::Reasoning(mode),
Err(error) => RuntimeBranchResult::Failed(error),
}
}),
) {
reasoning_enabled = false;
}
}
if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
self.finalize_pending_branches(branch_set.cancel_pending());
return Ok(None);
}
if !transition_enabled && !skill_enabled && !reasoning_enabled {
self.finalize_pending_branches(branch_set.cancel_pending());
return Ok(None);
}
let mut main_pending = true;
let mut skill_pending = skill_enabled;
let mut reasoning_pending = reasoning_enabled;
let mut transition_finalized = !transition_enabled;
let mut skill_finalized = !skill_enabled;
let mut reasoning_finalized = !reasoning_enabled;
let mut main_result: Option<Result<MainResponseDraft>> = None;
let mut transition_candidate: Option<TransitionCandidate> = None;
let mut skill_candidate: Option<SkillCandidate> = None;
let mut reasoning_decision: Option<ReasoningMode> = None;
let mut transition_fallback_required = false;
let mut skill_fallback_required = false;
let mut reasoning_fallback_required = false;
loop {
if let Some(candidate) = transition_candidate.take() {
if self
.approve_transition_target(&candidate.from_state, candidate.target())
.await?
{
self.finalize_pending_branches(branch_set.cancel_pending());
if !main_pending {
self.finalize_branch_loss(
&main_id,
main_kind,
RuntimeCommitBehavior::FinalResponse,
false,
main_result.as_ref().map(|result| result.is_err()),
);
}
if skill_enabled && !skill_pending {
self.finalize_branch_loss(
&skill_id,
RuntimeOptimizationKind::SpeculativeSkillRouting,
RuntimeCommitBehavior::SkillSelection,
false,
Some(false),
);
}
if reasoning_enabled && !reasoning_pending {
self.finalize_branch_loss(
&reasoning_id,
RuntimeOptimizationKind::SpeculativeReasoningAuto,
RuntimeCommitBehavior::ReasoningDecision,
false,
Some(false),
);
}
if !self
.apply_pre_response_transition_candidate(
&candidate,
&HashMap::new(),
processed_input,
)
.await?
{
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"discarded",
false,
);
return Ok(None);
}
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"committed",
true,
);
return self
.redispatch_current_state(processed_input)
.await
.map(Some);
}
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"discarded",
false,
);
transition_finalized = true;
}
if transition_finalized && skill_candidate.is_some() {
let candidate = skill_candidate.take().unwrap();
self.finalize_optional_branch(
&skill_id,
RuntimeOptimizationKind::SpeculativeSkillRouting,
RuntimeCommitBehavior::SkillSelection,
"committed",
true,
);
if !main_pending {
self.finalize_branch_loss(
&main_id,
main_kind,
RuntimeCommitBehavior::FinalResponse,
false,
main_result.as_ref().map(|result| result.is_err()),
);
}
if reasoning_enabled && !reasoning_pending {
self.finalize_branch_loss(
&reasoning_id,
RuntimeOptimizationKind::SpeculativeReasoningAuto,
RuntimeCommitBehavior::ReasoningDecision,
false,
Some(false),
);
}
self.finalize_pending_branches(branch_set.cancel_pending());
self.commit_root_user_message(processed_input).await?;
return match self
.commit_skill_candidate_route_result(candidate, processed_input)
.await?
{
SkillRouteResult::Response { skill_id, content } => self
.handle_skill_response(processed_input, &skill_id, content, input_context)
.await
.map(Some),
SkillRouteResult::NeedsClarification {
response,
ownership,
} => {
let admission = self
.admit_optional_disambiguation_ownership(ownership)
.await?;
if response
.metadata
.as_ref()
.and_then(|m| m.get("disambiguation"))
.and_then(|d| d.get("status"))
.and_then(|s| s.as_str())
== Some("awaiting_clarification")
{
self.memory
.add_message(ChatMessage::assistant(&response.content))
.await?;
}
drop(admission);
self.finish_turn_if_root(&response).await?;
Ok(Some(response))
}
SkillRouteResult::NoMatch => Ok(None),
};
}
if transition_finalized
&& skill_finalized
&& let Some(reasoning_mode) = reasoning_decision.take()
{
if !matches!(reasoning_mode, ReasoningMode::None) {
self.finalize_optional_branch(
&reasoning_id,
RuntimeOptimizationKind::SpeculativeReasoningAuto,
RuntimeCommitBehavior::ReasoningDecision,
"committed",
true,
);
if !main_pending {
self.finalize_branch_loss(
&main_id,
main_kind,
RuntimeCommitBehavior::FinalResponse,
false,
main_result.as_ref().map(|result| result.is_err()),
);
}
self.finalize_pending_branches(branch_set.cancel_pending());
self.commit_root_user_message(processed_input).await?;
return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
self.handle_plan_and_execute(processed_input, input_context, true)
.await
.map(Some)
} else {
self.run_committed_response_loop_with_reasoning(
processed_input,
input_context,
reasoning_mode,
true,
)
.await
.map(Some)
};
}
self.finalize_optional_branch(
&reasoning_id,
RuntimeOptimizationKind::SpeculativeReasoningAuto,
RuntimeCommitBehavior::ReasoningDecision,
"committed",
true,
);
reasoning_finalized = true;
}
if transition_finalized && skill_finalized && reasoning_finalized {
if transition_fallback_required
|| skill_fallback_required
|| reasoning_fallback_required
{
if !main_pending {
self.finalize_branch_loss(
&main_id,
main_kind,
RuntimeCommitBehavior::FinalResponse,
false,
main_result.as_ref().map(|result| result.is_err()),
);
}
self.finalize_pending_branches(branch_set.cancel_pending());
return Ok(None);
}
if let Some(result) = main_result.take() {
let draft = match result {
Ok(draft) => draft,
Err(error) => {
self.finalize_optional_branch(
&main_id,
main_kind,
RuntimeCommitBehavior::FinalResponse,
"failed",
false,
);
self.finalize_pending_branches(branch_set.cancel_pending());
return Err(error);
}
};
self.finalize_optional_branch(
&main_id,
main_kind,
RuntimeCommitBehavior::FinalResponse,
"committed",
true,
);
self.finalize_pending_branches(branch_set.cancel_pending());
return self
.commit_main_response_draft(
processed_input,
input_context,
draft,
ReasoningMode::None,
reasoning_enabled,
)
.await
.map(Some);
}
}
if branch_set.is_empty() {
return Ok(None);
}
let Some(outcome) = branch_set.next_completed().await else {
return Ok(None);
};
let branch_id = outcome.branch.branch_id();
match outcome.result {
RuntimeBranchResult::MainDraft(draft) => {
main_pending = false;
main_result = Some(Ok(draft));
}
RuntimeBranchResult::Transition(candidate) => {
if let Some(candidate) = candidate {
transition_candidate = Some(candidate);
} else {
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"discarded",
false,
);
transition_finalized = true;
}
}
RuntimeBranchResult::Skill(candidate) => {
skill_pending = false;
if let Some(candidate) = candidate {
skill_candidate = Some(candidate);
} else {
self.finalize_optional_branch(
&skill_id,
RuntimeOptimizationKind::SpeculativeSkillRouting,
RuntimeCommitBehavior::SkillSelection,
"discarded",
false,
);
skill_finalized = true;
}
}
RuntimeBranchResult::Reasoning(mode) => {
reasoning_pending = false;
reasoning_decision = Some(mode);
}
RuntimeBranchResult::Failed(error) => {
if branch_id == main_id {
main_pending = false;
main_result = Some(Err(error));
} else if branch_id == transition_id {
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"failed",
false,
);
transition_finalized = true;
} else if branch_id == skill_id {
skill_pending = false;
self.finalize_optional_branch(
&skill_id,
RuntimeOptimizationKind::SpeculativeSkillRouting,
RuntimeCommitBehavior::SkillSelection,
"failed",
false,
);
skill_finalized = true;
} else if branch_id == reasoning_id {
reasoning_pending = false;
self.finalize_optional_branch(
&reasoning_id,
RuntimeOptimizationKind::SpeculativeReasoningAuto,
RuntimeCommitBehavior::ReasoningDecision,
"failed",
false,
);
reasoning_finalized = true;
}
}
RuntimeBranchResult::Cancelled => {
self.finalize_optional_branch(
&branch_id,
outcome.branch.optimization,
outcome.branch.commit_behavior,
"cancelled",
false,
);
if branch_id == main_id {
main_pending = false;
main_result =
Some(Err(AgentError::Other("main branch cancelled".to_string())));
} else if branch_id == transition_id {
transition_finalized = true;
transition_fallback_required = true;
} else if branch_id == skill_id {
skill_pending = false;
skill_finalized = true;
skill_fallback_required = true;
} else if branch_id == reasoning_id {
reasoning_pending = false;
reasoning_finalized = true;
reasoning_fallback_required = true;
}
}
}
}
}
fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
for branch in branches {
self.finalize_optional_branch(
&branch.branch_id(),
branch.optimization,
branch.commit_behavior,
"cancelled",
false,
);
}
}
fn finalize_branch_loss(
&self,
branch_id: &str,
optimization: RuntimeOptimizationKind,
commit_behavior: RuntimeCommitBehavior,
pending: bool,
completed_failed: Option<bool>,
) {
let status = if pending {
"cancelled"
} else if completed_failed.unwrap_or(false) {
"failed"
} else {
"discarded"
};
self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
}
fn finalize_optional_branch(
&self,
branch_id: &str,
optimization: RuntimeOptimizationKind,
commit_behavior: RuntimeCommitBehavior,
status: &str,
winner: bool,
) {
crate::optimization::observability::finalize_branch(
self.observability_manager.as_ref(),
branch_id,
status,
winner,
optimization,
commit_behavior,
);
}
fn has_parallel_transition_candidates(&self) -> bool {
self.transitions_available_for_commit()
.map(|(transitions, _)| {
transitions
.iter()
.any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
})
.unwrap_or(false)
}
async fn select_parallel_transition_candidate(
&self,
processed_input: &str,
) -> Result<ParallelTransitionSelection> {
let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
return Ok(ParallelTransitionSelection::NoMatch);
};
let parallel: Vec<Transition> = transitions
.into_iter()
.filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
.filter(|transition| !transition.requires_response)
.collect();
if parallel.is_empty() {
return Ok(ParallelTransitionSelection::NoMatch);
}
let empty_staged = HashMap::new();
if let Some(candidate) = self.select_deterministic_transition_candidate(
processed_input,
¤t_state,
¶llel,
&empty_staged,
) {
return Ok(ParallelTransitionSelection::Candidate(candidate));
}
let when_transitions: Vec<(usize, &Transition)> = parallel
.iter()
.enumerate()
.filter(|(_, transition)| !transition.when.trim().is_empty())
.collect();
if when_transitions.is_empty() {
return Ok(ParallelTransitionSelection::NoMatch);
}
let llm = self
.llm_registry
.router()
.or_else(|_| self.llm_registry.default())
.map_err(|e| AgentError::Config(e.to_string()))?;
let conditions = when_transitions
.iter()
.enumerate()
.map(|(display_idx, (_, transition))| {
format!("{}. {}", display_idx + 1, transition.when)
})
.collect::<Vec<_>>()
.join("\n");
if !self
.reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
{
return Ok(ParallelTransitionSelection::ReservationExhausted);
}
let context_preview = self.branch_context_preview();
let prompt = format!(
"Based only on the current user message and context, which transition condition is met?\n\nCurrent state: {}\nUser message: {}\nContext:\n{}\n\nConditions:\n{}\n0. None of the above\n\nReply with ONLY the number (0-{}).",
current_state,
processed_input,
context_preview,
conditions,
when_transitions.len()
);
let response = self
.observe_purpose(
ObservationPurpose::StateTransitionEvaluation,
llm.complete(&[ChatMessage::user(prompt)], None),
)
.await
.map_err(|e| AgentError::LLM(e.to_string()))?;
let choice = response.content.trim().parse::<usize>().unwrap_or(0);
if choice == 0 || choice > when_transitions.len() {
return Ok(ParallelTransitionSelection::NoMatch);
}
let transition = when_transitions[choice - 1].1.clone();
Ok(ParallelTransitionSelection::Candidate(
TransitionCandidate::new(
current_state,
transition.clone(),
Self::transition_reason(&transition),
),
))
}
async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
const MAX_REDISPATCH_DEPTH: u32 = 3;
let current_depth = *self.redispatch_depth.read();
if current_depth >= MAX_REDISPATCH_DEPTH {
warn!(depth = current_depth, "Re-dispatch depth limit reached");
let response = AgentResponse::new("");
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
*self.redispatch_depth.write() += 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.enter_redispatch();
}
let result = Box::pin(self.run_loop_internal(processed_input)).await;
*self.redispatch_depth.write() -= 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.exit_redispatch();
}
let response = result?;
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
if *self.redispatch_depth.read() == 0 {
self.post_turn_session_lifecycle().await?;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.mark_post_turn_lifecycle_completed();
}
self.hooks.on_response(response).await;
self.end_root_turn();
}
Ok(())
}
async fn execute_state_exit_actions(&self, state_path: &str) {
if let Some(ref sm) = self.state_machine
&& let Some(def) = sm.get_definition(state_path)
&& !def.on_exit.is_empty()
{
debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
self.execute_state_actions(&def.on_exit).await;
}
}
fn state_was_previously_entered(
state_path: &str,
from_state: &str,
history_before: &[StateTransitionEvent],
) -> bool {
state_path == from_state
|| history_before
.iter()
.any(|event| event.from == state_path || event.to == state_path)
}
async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
if let Some(ref sm) = self.state_machine
&& let Some(def) = sm.get_definition(state_path)
{
if is_reentry && !def.on_reenter.is_empty() {
debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
self.execute_state_actions(&def.on_reenter).await;
} else if !def.on_enter.is_empty() {
debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
self.execute_state_actions(&def.on_enter).await;
}
}
}
async fn execute_state_actions(&self, actions: &[StateAction]) {
for (action_index, action) in actions.iter().enumerate() {
match action {
StateAction::Tool { tool, args } => {
let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
let args_value = self.render_action_args(&raw_args);
let state = self.state_machine.as_ref().map(|sm| sm.current());
let request = ToolExecutionRequest::new(
uuid::Uuid::new_v4().to_string(),
tool.clone(),
args_value,
ToolCallSource::StateAction {
state,
action_index,
},
);
match self.execute_tool_record(request).await {
Ok(record) if record.success => {
debug!(tool = %record.canonical_id, "State action: tool executed");
let _ = self.context_manager.set(
"last_tool_result",
serde_json::Value::String(record.model_output_string()),
);
let _ = self.context_manager.set(
"last_tool_record",
serde_json::to_value(record).unwrap_or(Value::Null),
);
}
Ok(record) => {
warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
}
Err(e) => {
warn!(tool = %tool, error = %e, "State action: tool failed")
}
}
}
StateAction::Skill { skill } => {
if let Some(ref executor) = self.skill_executor {
if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
match executor
.execute_with_invoker(def, "", serde_json::json!({}), self)
.await
{
Ok(_) => debug!(skill = %skill, "State action: skill executed"),
Err(e) => {
warn!(skill = %skill, error = %e, "State action: skill failed")
}
}
} else {
warn!(skill = %skill, "State action: skill not found");
}
}
}
StateAction::SetContext { set_context } => {
for (key, value) in set_context {
if let Err(e) = self.context_manager.set(key, value.clone()) {
warn!(key = %key, error = %e, "State action: set_context failed");
} else {
debug!(key = %key, "State action: context set");
}
}
}
StateAction::Prompt {
prompt,
llm,
store_as,
} => {
let llm_result = if let Some(alias) = llm {
self.llm_registry.get(alias)
} else {
self.llm_registry.default()
};
match llm_result {
Ok(llm_provider) => {
let context = self.build_context_with_overlays();
let rendered_prompt = self
.template_renderer
.render(prompt, &context)
.unwrap_or_else(|_| prompt.clone());
let recent =
self.memory.get_messages(Some(5)).await.unwrap_or_default();
let mut messages: Vec<ChatMessage> = recent;
messages.push(ChatMessage::user(&rendered_prompt));
match self
.observe_purpose(
ObservationPurpose::StateAction,
llm_provider.complete(&messages, None),
)
.await
{
Ok(response) => {
if let Some(key) = store_as {
let _ = self
.context_manager
.set(key, Value::String(response.content));
debug!(key = %key, "State action: prompt result stored");
}
}
Err(e) => {
warn!(error = %e, "State action: prompt LLM call failed");
}
}
}
Err(e) => {
warn!(error = %e, "State action: LLM not found for prompt");
}
}
}
}
}
}
async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
let extractors = match &self.state_machine {
Some(sm) => match sm.current_definition() {
Some(def) if !def.extract.is_empty() => def.extract.clone(),
_ => return HashMap::new(),
},
None => return HashMap::new(),
};
let mut staged = HashMap::new();
for extractor in &extractors {
let prompt = if let Some(ref custom) = extractor.llm_extract {
format!(
"User message:\n\"{}\"\n\nInstruction:\n{}",
user_message, custom
)
} else if let Some(ref desc) = extractor.description {
format!(
"From the following message, extract: {}\n\n\
Message: \"{}\"\n\n\
If the information is present, return ONLY the extracted value.\n\
If NOT present, return exactly: __NONE__",
desc, user_message
)
} else {
continue;
};
let llm = match self
.llm_registry
.get(&extractor.llm)
.or_else(|_| self.llm_registry.get("router"))
.or_else(|_| self.llm_registry.get("default"))
{
Ok(llm) => llm,
Err(e) => {
warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
continue;
}
};
let messages = vec![ChatMessage::user(&prompt)];
match self
.observe_purpose(
ObservationPurpose::ContextExtraction,
llm.complete(&messages, None),
)
.await
{
Ok(response) => {
let value = response.content.trim().to_string();
if value != "__NONE__" && !value.is_empty() {
staged.insert(
extractor.key.clone(),
serde_json::Value::String(value.clone()),
);
debug!(key = %extractor.key, value = %value, "Context extracted");
} else if extractor.required {
warn!(key = %extractor.key, "Required extraction returned no value");
}
}
Err(e) => {
warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
}
}
}
staged
}
fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
for (key, value) in staged {
if let Err(error) = self.context_manager.update(key, value.clone()) {
warn!(key = %key, error = %error, "staged context write failed");
}
}
}
async fn run_context_extractors(&self, user_message: &str) {
let staged = self.run_context_extractors_staged(user_message).await;
self.commit_staged_context_writes(&staged);
}
async fn check_memory_compression(&self) -> Result<()> {
if self.memory.needs_compression() {
let result = self.memory.compress(None).await?;
if let CompressResult::Compressed {
messages_summarized,
new_summary_length,
tokens_saved,
} = result
{
let event = MemoryCompressEvent::new(
messages_summarized,
tokens_saved,
new_summary_length as u32,
);
self.hooks.on_memory_compress(&event).await;
debug!(
messages = messages_summarized,
tokens_saved = tokens_saved,
"Memory compressed"
);
}
}
self.handle_memory_overflow().await?;
self.check_memory_budget().await;
Ok(())
}
async fn check_memory_budget(&self) {
let Some(ref budget) = self.memory_token_budget else {
return;
};
let context = match self.memory.get_context().await {
Ok(ctx) => ctx,
Err(_) => return,
};
let used_tokens = context.estimated_tokens();
if budget.is_over_warn_threshold(used_tokens) {
let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
self.hooks.on_memory_budget_warning(&event).await;
debug!(
used = used_tokens,
total = budget.total,
percent = event.usage_percent,
"Memory budget warning"
);
}
if let Some(ref summary) = context.summary {
let summary_tokens = ai_agents_memory::estimate_tokens(summary);
let summary_budget = budget.allocation.summary;
if summary_budget > 0 {
let warn_threshold =
(summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
if summary_tokens >= warn_threshold {
let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
self.hooks.on_memory_budget_warning(&event).await;
}
}
}
let recent_tokens: u32 = context
.messages
.iter()
.map(ai_agents_memory::estimate_message_tokens)
.sum();
let recent_budget = budget.allocation.recent_messages;
if recent_budget > 0 {
let warn_threshold =
(recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
if recent_tokens >= warn_threshold {
let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
self.hooks.on_memory_budget_warning(&event).await;
}
}
let relationship_budget = budget.allocation.relationships;
if relationship_budget > 0 {
let relationship_tokens = self
.relationship_memory_text()
.map(|text| ai_agents_memory::estimate_tokens(&text))
.unwrap_or(0);
let warn_threshold =
(relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
if relationship_tokens >= warn_threshold {
let event = MemoryBudgetEvent::new(
"relationships",
relationship_tokens,
relationship_budget,
);
self.hooks.on_memory_budget_warning(&event).await;
}
}
}
async fn handle_memory_overflow(&self) -> Result<()> {
let Some(ref budget) = self.memory_token_budget else {
return Ok(());
};
let context = self.memory.get_context().await?;
let used_tokens = context.estimated_tokens();
if used_tokens <= budget.total {
return Ok(());
}
match budget.overflow_strategy {
OverflowStrategy::TruncateOldest => {
let tokens_to_free = used_tokens - budget.total;
let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
if messages_to_evict > 0 {
self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
.await?;
}
}
OverflowStrategy::SummarizeMore => {
let max_attempts = context.total_messages.max(1);
for _ in 0..max_attempts {
match self.memory.compress(None).await? {
CompressResult::Compressed {
messages_summarized,
..
} if messages_summarized > 0 => {
let context = self.memory.get_context().await?;
if context.estimated_tokens() <= budget.total {
return Ok(());
}
}
_ => break,
}
}
let context = self.memory.get_context().await?;
let used_tokens = context.estimated_tokens();
if used_tokens > budget.total {
return Err(AgentError::MemoryBudgetExceeded {
used: used_tokens,
budget: budget.total,
});
}
}
OverflowStrategy::Error => {
return Err(AgentError::MemoryBudgetExceeded {
used: used_tokens,
budget: budget.total,
});
}
}
Ok(())
}
fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
}
async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
let evicted = self.memory.evict_oldest(count).await?;
if !evicted.is_empty() {
let event = MemoryEvictEvent {
reason,
messages_evicted: evicted.len(),
importance_scores: vec![],
};
self.hooks.on_memory_evict(&event).await;
debug!(count = evicted.len(), "Messages evicted from memory");
}
Ok(())
}
#[instrument(skip(self, input), fields(agent = %self.info.name))]
async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
match self.determine_reasoning_mode_strict(input).await {
Ok(mode) => Ok(mode),
Err(_) => Ok(ReasoningMode::None),
}
}
async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
let effective_config = self.get_effective_reasoning_config();
if !matches!(effective_config.mode, ReasoningMode::Auto) {
return Ok(effective_config.mode.clone());
}
let judge_llm = effective_config
.judge_llm
.as_ref()
.and_then(|alias| self.llm_registry.get(alias).ok())
.or_else(|| self.llm_registry.router().ok())
.or_else(|| self.llm_registry.default().ok());
let Some(llm) = judge_llm else {
return Ok(ReasoningMode::None);
};
let prompt = format!(
r#"Analyze this user request and determine the appropriate reasoning mode.
User request: "{}"
Choose ONE of these modes:
- none: Simple queries, greetings, direct answers (fastest)
- cot: Complex analysis, multi-step reasoning, math problems
- react: Tasks requiring multiple tool calls with observation
- plan_and_execute: Complex multi-step tasks requiring coordination
Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
input
);
let messages = vec![ChatMessage::user(&prompt)];
let response = self
.observe_purpose(
ObservationPurpose::ReflectionDecision,
llm.complete(&messages, None),
)
.await
.map_err(|e| AgentError::LLM(e.to_string()))?;
let mode_str = response.content.trim().to_lowercase();
Ok(match mode_str.as_str() {
"cot" => ReasoningMode::CoT,
"react" => ReasoningMode::React,
"plan_and_execute" => ReasoningMode::PlanAndExecute,
_ => ReasoningMode::None,
})
}
async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
let effective_config = self.get_effective_reflection_config();
if !effective_config.requires_evaluation() {
return Ok(false);
}
if effective_config.is_enabled() {
return Ok(true);
}
let evaluator_llm = effective_config
.evaluator_llm
.as_ref()
.and_then(|alias| self.llm_registry.get(alias).ok())
.or_else(|| self.llm_registry.router().ok())
.or_else(|| self.llm_registry.default().ok());
let Some(llm) = evaluator_llm else {
return Ok(false);
};
let response_preview: String = response.chars().take(500).collect();
let prompt = format!(
r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
User query: "{}"
Response: "{}"
Answer YES or NO only."#,
input, response_preview
);
let messages = vec![ChatMessage::user(&prompt)];
let result = self
.observe_purpose(
ObservationPurpose::ReflectionDecision,
llm.complete(&messages, None),
)
.await;
match result {
Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
Err(_) => Ok(false),
}
}
fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
format!(
"{}\n\n<instruction>\nThink through this step by step before answering:\n1. Understand what is being asked\n2. Break down the problem\n3. Work through each part\n4. Provide your final answer\n\nShow your thinking process, then give your final answer.\n</instruction>",
base_prompt
)
}
fn build_react_system_prompt(&self, base_prompt: &str) -> String {
format!(
"{}\n\n<instruction>\nUse the Reason-Act-Observe pattern:\n1. Thought: Think about what to do\n2. Action: Use a tool if needed\n3. Observation: Analyze the result\n4. Repeat until you have the answer\n\nFormat your response showing Thought/Action/Observation steps.\n</instruction>",
base_prompt
)
}
async fn generate_plan(&self, input: &str) -> Result<Plan> {
let effective = self.get_effective_reasoning_config();
let planning_config = effective.get_planning();
let planner_llm = planning_config
.and_then(|c| c.planner_llm.as_ref())
.and_then(|alias| self.llm_registry.get(alias).ok())
.or_else(|| self.llm_registry.router().ok())
.or_else(|| self.llm_registry.default().ok())
.ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
let mut available_tool_ids: Vec<String> = self
.get_available_tool_ids()
.await
.unwrap_or_else(|_| self.tools.list_ids());
let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
if let Some(config) = planning_config {
if !config.available.tools.is_all() {
available_tool_ids.retain(|t| config.available.tools.allows(t));
}
if !config.available.skills.is_all() {
available_skills.retain(|s| config.available.skills.allows(s));
}
}
let tool_descriptions: Vec<String> = available_tool_ids
.iter()
.filter_map(|id| {
self.tools.get(id).map(|tool| {
let schema = tool.input_schema();
let args_desc = schema
.get("properties")
.and_then(|p| serde_json::to_string(p).ok())
.unwrap_or_else(|| "{}".to_string());
format!(
"- {} ({}): {}\n Arguments: {}",
id,
tool.name(),
tool.description(),
args_desc
)
})
})
.collect();
let tools_section = if tool_descriptions.is_empty() {
"Available tools: none".to_string()
} else {
format!("Available tools:\n{}", tool_descriptions.join("\n"))
};
let skills_section = if available_skills.is_empty() {
"Available skills: none".to_string()
} else {
format!("Available skills: {}", available_skills.join(", "))
};
let prompt = format!(
r#"Create a step-by-step plan to accomplish this goal.
Goal: "{}"
{}
{}
Create a plan with clear steps. For each step, specify:
- description: What this step accomplishes
- action_type: "tool", "skill", "think", or "respond"
- action_target: The tool/skill id (if applicable)
- args: The arguments object matching the tool's schema (if action_type is "tool")
- dependencies: List of step IDs this depends on (empty if none)
Respond in JSON format:
{{
"steps": [
{{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
{{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
]
}}"#,
input, tools_section, skills_section,
);
let messages = vec![ChatMessage::user(&prompt)];
let response = self
.observe_purpose(
ObservationPurpose::PlanGeneration,
planner_llm.complete(&messages, None),
)
.await
.map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
let mut plan = Plan::new(input);
if let Some(json_start) = response.content.find('{')
&& let Some(json_end) = response.content.rfind('}')
{
let json_str = &response.content[json_start..=json_end];
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
&& let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
{
for step_value in steps {
let id = step_value
.get("id")
.and_then(|v| v.as_str())
.unwrap_or("step");
let desc = step_value
.get("description")
.and_then(|v| v.as_str())
.unwrap_or("");
let action_type = step_value
.get("action_type")
.and_then(|v| v.as_str())
.unwrap_or("think");
let action_target = step_value
.get("action_target")
.and_then(|v| v.as_str())
.unwrap_or("");
let args = step_value
.get("args")
.cloned()
.unwrap_or(serde_json::json!({}));
let deps: Vec<String> = step_value
.get("dependencies")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(String::from))
.collect()
})
.unwrap_or_default();
let action = match action_type {
"tool" => PlanAction::tool(action_target, args),
"skill" => PlanAction::skill(action_target),
"respond" => PlanAction::respond(action_target),
_ => PlanAction::think(desc),
};
let step = PlanStep::new(desc, action)
.with_id(id)
.with_dependencies(deps);
plan.add_step(step);
}
}
}
if plan.steps.is_empty() {
plan.add_step(PlanStep::new(
"Process the request",
PlanAction::think(input),
));
plan.add_step(PlanStep::new(
"Provide response",
PlanAction::respond("Answer based on analysis"),
));
}
Ok(plan)
}
async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
let llm = self.get_state_llm()?;
let mut results: HashMap<String, serde_json::Value> = HashMap::new();
let effective = self.get_effective_reasoning_config();
let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
plan.status = PlanStatus::InProgress;
for step_idx in 0..plan.steps.len().min(max_steps as usize) {
let step = &plan.steps[step_idx];
let deps_satisfied = step.dependencies.iter().all(|dep| {
plan.steps
.iter()
.find(|s| &s.id == dep)
.map(|s| s.status.is_completed())
.unwrap_or(false)
});
if !deps_satisfied {
continue;
}
plan.steps[step_idx].mark_running();
let result = match &plan.steps[step_idx].action {
PlanAction::Tool { tool, args } => {
let has_dep_results = plan.steps[step_idx]
.dependencies
.iter()
.any(|dep| results.contains_key(dep));
let final_args = if has_dep_results {
let dep_context: String = plan.steps[step_idx]
.dependencies
.iter()
.filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
.collect::<Vec<_>>()
.join("\n");
let tool_schema = self
.tools
.get(tool)
.map(|t| {
let schema = t.input_schema();
let props = schema
.get("properties")
.and_then(|p| serde_json::to_string(p).ok())
.unwrap_or_else(|| "{}".to_string());
format!(
"{}: {}\nArguments schema: {}",
t.id(),
t.description(),
props
)
})
.unwrap_or_default();
let step_desc = &plan.steps[step_idx].description;
let arg_prompt = format!(
"Generate the JSON arguments for a tool call.\n\n\
Tool: {}\n\n\
Task: {}\n\n\
Previous step results:\n{}\n\n\
Planner's draft arguments: {}\n\n\
Produce ONLY a valid JSON object with the correct argument values.\n\
Use actual values from the previous step results, not template references.",
tool_schema,
step_desc,
dep_context,
serde_json::to_string(args).unwrap_or_default()
);
let messages = vec![ChatMessage::user(&arg_prompt)];
match self
.observe_purpose(
ObservationPurpose::PlanStep,
llm.complete(&messages, None),
)
.await
{
Ok(resp) => {
let content = resp.content.trim();
let json_start = content.find('{');
let json_end = content.rfind('}');
if let (Some(start), Some(end)) = (json_start, json_end) {
serde_json::from_str(&content[start..=end])
.unwrap_or_else(|_| args.clone())
} else {
args.clone()
}
}
Err(_) => args.clone(),
}
} else {
args.clone()
};
let request = ToolExecutionRequest::new(
uuid::Uuid::new_v4().to_string(),
tool.clone(),
final_args,
ToolCallSource::Plan {
step_index: step_idx,
},
);
match self.execute_tool_record(request).await {
Ok(record) if record.success => {
serde_json::json!({ "output": record.model_output_string() })
}
Ok(record) => {
plan.steps[step_idx].mark_failed(record.model_output_string());
continue;
}
Err(e) => {
plan.steps[step_idx].mark_failed(e.to_string());
continue;
}
}
}
PlanAction::Skill { skill } => {
if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
if let Some(ref executor) = self.skill_executor {
match executor
.execute_with_invoker(skill_def, "", serde_json::json!({}), self)
.await
{
Ok(output) => serde_json::json!({ "output": output }),
Err(e) => {
plan.steps[step_idx].mark_failed(e.to_string());
continue;
}
}
} else {
serde_json::json!({ "output": "Skill executor not available" })
}
} else {
plan.steps[step_idx].mark_failed("Skill not found");
continue;
}
}
PlanAction::Think { prompt } => {
let context: String = results
.iter()
.map(|(k, v)| format!("{}: {}", k, v))
.collect::<Vec<_>>()
.join("\n");
let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
let messages = vec![ChatMessage::user(&think_prompt)];
match self
.observe_purpose(
ObservationPurpose::PlanStep,
llm.complete(&messages, None),
)
.await
{
Ok(resp) => serde_json::json!({ "output": resp.content }),
Err(e) => {
plan.steps[step_idx].mark_failed(e.to_string());
continue;
}
}
}
PlanAction::Respond { template } => {
let context: String = results
.iter()
.map(|(k, v)| format!("{}: {}", k, v))
.collect::<Vec<_>>()
.join("\n");
let respond_prompt = format!(
"Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
context, template
);
let messages = vec![ChatMessage::user(&respond_prompt)];
match self
.observe_purpose(
ObservationPurpose::PlanStep,
llm.complete(&messages, None),
)
.await
{
Ok(resp) => serde_json::json!({ "output": resp.content }),
Err(e) => {
plan.steps[step_idx].mark_failed(e.to_string());
continue;
}
}
}
};
results.insert(plan.steps[step_idx].id.clone(), result.clone());
plan.steps[step_idx].mark_completed(Some(result));
}
let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
if has_failures {
let failed_ids: Vec<String> = plan
.steps
.iter()
.filter(|s| s.status.is_failed())
.map(|s| s.id.clone())
.collect();
plan.status = PlanStatus::Failed {
error: format!("Steps failed: {}", failed_ids.join(", ")),
};
} else {
plan.status = PlanStatus::Completed;
}
let all_outputs: Vec<String> = plan
.steps
.iter()
.filter(|s| s.status.is_completed())
.filter_map(|s| {
s.result
.as_ref()
.and_then(|r| r.get("output"))
.and_then(|o| o.as_str())
.map(|o| format!("{}: {}", s.description, o))
})
.collect();
if all_outputs.is_empty() {
return Ok("Plan execution completed but produced no results.".to_string());
}
if all_outputs.len() == 1 {
return Ok(all_outputs.into_iter().next().unwrap());
}
let context = all_outputs.join("\n\n");
let prompt = format!(
"You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
plan.goal, context
);
let messages = vec![ChatMessage::user(&prompt)];
match self
.observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
.await
{
Ok(resp) => Ok(resp.content.trim().to_string()),
Err(_) => Ok(context),
}
}
async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
let effective_config = self.get_effective_reflection_config();
self.evaluate_response_with_config(input, response, &effective_config)
.await
}
fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
if let Some(start) = content.find("<thinking>")
&& let Some(end) = content.find("</thinking>")
{
let thinking = content[start + 10..end].trim().to_string();
let answer = content[end + 11..].trim().to_string();
return (Some(thinking), answer);
}
(None, content.to_string())
}
fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
match self.get_effective_reasoning_config().output {
ReasoningOutput::Hidden => answer.to_string(),
ReasoningOutput::Visible => {
if let Some(t) = thinking {
format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
} else {
answer.to_string()
}
}
ReasoningOutput::Tagged => {
if let Some(t) = thinking {
format!("<thinking>{}</thinking>\n{}", t, answer)
} else {
answer.to_string()
}
}
}
}
async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
self.init_storage().await?;
self.begin_root_turn();
let _root_cleanup = RootTurnCleanup::new(self);
info!(input_len = input.len(), "Starting chat");
self.hooks.on_message_received(input).await;
if !self.context_initialized.swap(true, Ordering::SeqCst) {
self.context_manager.initialize().await?;
debug!("Context manager initialized (defaults, env, builtins)");
}
self.check_turn_timeout().await?;
self.context_manager.refresh_per_turn().await?;
self.clear_disambiguation_context();
if let Some(ref disambiguator) = self.disambiguation_manager {
let disambiguation_context = self.build_disambiguation_context().await?;
let state_override = self
.state_machine
.as_ref()
.and_then(|sm| sm.current_definition())
.and_then(|def| def.disambiguation.clone());
let state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
let mut disambiguation_result = self
.observe_purpose(
ObservationPurpose::DisambiguationDetection,
disambiguator.process_input_with_override(
input,
&disambiguation_context,
state_override.as_ref(),
None,
),
)
.await?;
let current_state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
if current_state_generation != state_generation
|| self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
{
disambiguator.clear_pending().await;
*self.pending_skill_id.write() = None;
disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
info!(
confirmation_event = "invalidated",
invalidation_reason = "state_generation_changed",
"Disambiguation result invalidated before redispatch"
);
}
match disambiguation_result {
DisambiguationResult::Clear => {
debug!("Input is clear, proceeding normally");
}
DisambiguationResult::NeedsClarification {
question,
detection,
} => {
let admission = self
.admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
.await?;
let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
info!(
ambiguity_type = ?detection.ambiguity_type,
confidence = detection.confidence,
"Input requires clarification"
);
self.commit_root_user_message(input).await?;
self.memory
.add_message(ChatMessage::assistant(&question.question))
.await?;
let status = if awaiting_confirmation {
"awaiting_confirmation"
} else {
"awaiting_clarification"
};
let response = AgentResponse::new(&question.question).with_metadata(
"disambiguation",
serde_json::json!({
"status": status,
"options": question.options,
"clarifying": question.clarifying,
"detection": {
"type": detection.ambiguity_type,
"confidence": detection.confidence,
"what_is_unclear": detection.what_is_unclear,
}
}),
);
drop(admission);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
DisambiguationResult::Clarified {
enriched_input,
resolved,
..
} => {
let admission = match self
.admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
.await
{
Ok(admission) => admission,
Err(error) => {
*self.pending_skill_id.write() = None;
return Err(error);
}
};
info!(
resolved_count = resolved.len(),
enriched = %enriched_input,
"Input clarified, injecting resolved intent into context"
);
for (key, value) in &resolved {
let context_key = format!("disambiguation.{}", key);
let _ = self.context_manager.set(&context_key, value.clone());
}
if let Some(intent) = resolved.get("intent") {
let _ = self.context_manager.set("resolved_intent", intent.clone());
}
let _ = self
.context_manager
.set("disambiguation.resolved", serde_json::Value::Bool(true));
let skill_id = self.pending_skill_id.read().clone();
if let Some(skill_id) = skill_id {
info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
drop(admission);
return self
.recheck_skill_disambiguation(
&skill_id,
&enriched_input,
disambiguation_epoch,
state_generation,
)
.await;
}
drop(admission);
return self.run_loop_internal(&enriched_input).await;
}
DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
info!("Proceeding with best guess interpretation");
let skill_id = self.pending_skill_id.read().clone();
if let Some(skill_id) = skill_id {
info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
return self
.recheck_skill_disambiguation(
&skill_id,
&enriched_input,
disambiguation_epoch,
state_generation,
)
.await;
}
return self.run_loop_internal(&enriched_input).await;
}
DisambiguationResult::GiveUp { reason } => {
*self.pending_skill_id.write() = None;
warn!(reason = %reason, "Disambiguation gave up");
let apology = self
.generate_localized_apology(
"Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I'm sorry, I couldn't understand your request: {}", reason)
});
let response = AgentResponse::new(&apology);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
DisambiguationResult::Escalate { reason } => {
*self.pending_skill_id.write() = None;
info!(reason = %reason, "Escalating to human");
if let Some(ref hitl) = self.hitl_engine {
let trigger =
ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
let mut context_map = HashMap::new();
context_map.insert("original_input".to_string(), serde_json::json!(input));
context_map.insert("reason".to_string(), serde_json::json!(&reason));
let check_result = HITLCheckResult::required(
trigger,
context_map,
format!("User request needs human assistance: {}", reason),
Some(hitl.config().default_timeout_seconds),
);
let result = self.request_hitl_approval(check_result).await?;
if matches!(
result,
ApprovalResult::Approved | ApprovalResult::Modified { .. }
) {
return self.run_loop_internal(input).await;
}
}
let apology = self
.generate_localized_apology(
"Explain briefly that you're transferring the user to a human agent for help.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I need human assistance to help with your request: {}", reason)
});
let response = AgentResponse::new(&apology);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
DisambiguationResult::Abandoned { new_input } => {
*self.pending_skill_id.write() = None;
info!(
has_new_input = new_input.is_some(),
"Clarification abandoned by user"
);
self.commit_root_user_message(input).await?;
match new_input {
Some(fresh_input) => {
return self.run_loop_internal(&fresh_input).await;
}
None => {
let ack = self
.generate_localized_apology(
"The user changed their mind about their previous request. \
Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
Do NOT apologize excessively. Be concise.",
"User abandoned clarification",
)
.await
.unwrap_or_else(|_| {
"OK, no problem. What else can I help with?".to_string()
});
self.memory
.add_message(ChatMessage::assistant(&ack))
.await?;
let response = AgentResponse::new(&ack);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
}
}
}
}
self.run_loop_internal(input).await
}
async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
let llm = self.llm_registry.router().map_err(|e| {
AgentError::LLM(format!(
"Router LLM not available for localized response: {}",
e
))
})?;
let recent: Vec<String> = self
.memory
.get_messages(Some(3))
.await?
.iter()
.map(|m| m.content.clone())
.collect();
let context_hint = if recent.is_empty() {
String::new()
} else {
format!(
"\nRecent conversation (detect the user's language from this):\n{}\n",
recent.join("\n")
)
};
let prompt = format!(
"{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
instruction, reason, context_hint
);
let messages = vec![ChatMessage::user(&prompt)];
let response = self
.observe_purpose(
ObservationPurpose::DisambiguationClarification,
llm.complete(&messages, None),
)
.await
.map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
Ok(response.content.trim().to_string())
}
fn render_action_args(&self, args: &Value) -> Value {
let context = self.build_context_with_overlays();
match args {
Value::Object(map) => {
let mut rendered = serde_json::Map::new();
for (k, v) in map {
match v {
Value::String(s) if s.contains("{{") => {
match self.template_renderer.render(s, &context) {
Ok(rendered_str) => {
rendered.insert(k.clone(), Value::String(rendered_str));
}
Err(_) => {
rendered.insert(k.clone(), v.clone());
}
}
}
_ => {
rendered.insert(k.clone(), v.clone());
}
}
}
Value::Object(rendered)
}
_ => args.clone(),
}
}
fn clear_disambiguation_context(&self) {
let _ = self
.context_manager
.set("resolved_intent", serde_json::Value::Null);
let all = self.context_manager.get_all();
for key in all.keys() {
if key.starts_with("disambiguation.") {
let _ = self.context_manager.set(key, serde_json::Value::Null);
}
}
}
async fn recheck_skill_disambiguation(
&self,
skill_id: &str,
enriched_input: &str,
expected_disambiguation_epoch: u64,
expected_state_generation: Option<u64>,
) -> Result<AgentResponse> {
let skill = self
.skill_router
.as_ref()
.and_then(|r| r.get_skill(skill_id).cloned());
if let Some(ref skill) = skill
&& let Some(ref skill_disambig) = skill.disambiguation
&& skill_disambig.enabled.unwrap_or(false)
&& let Some(ref disambiguator) = self.disambiguation_manager
{
let context = self.build_disambiguation_context().await?;
let state_override = self
.state_machine
.as_ref()
.and_then(|sm| sm.current_definition())
.and_then(|def| def.disambiguation.clone());
let disambiguation_result = self
.observe_purpose(
ObservationPurpose::DisambiguationDetection,
disambiguator.process_input_with_override(
enriched_input,
&context,
state_override.as_ref(),
Some(skill_disambig),
),
)
.await?;
let current_state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
if current_state_generation != expected_state_generation
|| self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
{
disambiguator.clear_pending().await;
*self.pending_skill_id.write() = None;
return Err(AgentError::Other(
"State or reset ownership changed during skill disambiguation recheck"
.to_string(),
));
}
match disambiguation_result {
DisambiguationResult::Clear => {
debug!(skill_id = %skill_id, "Skill re-check: all fields present");
}
DisambiguationResult::NeedsClarification {
question,
detection,
} => {
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
info!(
skill_id = %skill_id,
ambiguity_type = ?detection.ambiguity_type,
what_is_unclear = ?detection.what_is_unclear,
"Skill re-check: still missing fields, asking again"
);
self.memory
.add_message(ChatMessage::user(enriched_input))
.await?;
self.memory
.add_message(ChatMessage::assistant(&question.question))
.await?;
let response = AgentResponse::new(&question.question).with_metadata(
"disambiguation",
serde_json::json!({
"status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
"skill_id": skill_id,
"options": question.options,
"clarifying": question.clarifying,
"detection": {
"type": detection.ambiguity_type,
"confidence": detection.confidence,
"what_is_unclear": detection.what_is_unclear,
}
}),
);
drop(admission);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
DisambiguationResult::Clarified {
enriched_input: re_enriched,
..
} => {
debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
*self.pending_skill_id.write() = None;
drop(admission);
let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
self.memory
.add_message(ChatMessage::user(&re_enriched))
.await?;
return self
.handle_skill_response(
&re_enriched,
skill_id,
skill_response,
&HashMap::new(),
)
.await;
}
DisambiguationResult::ProceedWithBestGuess {
enriched_input: re_enriched,
} => {
debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
*self.pending_skill_id.write() = None;
drop(admission);
let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
self.memory
.add_message(ChatMessage::user(&re_enriched))
.await?;
return self
.handle_skill_response(
&re_enriched,
skill_id,
skill_response,
&HashMap::new(),
)
.await;
}
DisambiguationResult::GiveUp { reason } => {
*self.pending_skill_id.write() = None;
let apology = self
.generate_localized_apology(
"Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I'm sorry, I couldn't understand your request: {}", reason)
});
let response = AgentResponse::new(&apology);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
DisambiguationResult::Escalate { reason } => {
*self.pending_skill_id.write() = None;
let apology = self
.generate_localized_apology(
"Explain briefly that you're transferring the user to a human agent for help.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I need human assistance to help with your request: {}", reason)
});
let response = AgentResponse::new(&apology);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
DisambiguationResult::Abandoned { new_input } => {
*self.pending_skill_id.write() = None;
debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
if let Some(fresh) = new_input {
return self.run_loop_internal(&fresh).await;
}
let ack = self
.generate_localized_apology(
"The user changed their mind about their previous request. \
Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
Do NOT apologize excessively. Be concise.",
"User abandoned clarification",
)
.await
.unwrap_or_else(|_| {
"OK, no problem. What else can I help with?".to_string()
});
self.memory
.add_message(ChatMessage::assistant(&ack))
.await?;
let response = AgentResponse::new(&ack);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
}
}
let admission = self
.admit_disambiguation_redispatch(
expected_disambiguation_epoch,
expected_state_generation,
)
.await?;
*self.pending_skill_id.write() = None;
drop(admission);
let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
self.memory
.add_message(ChatMessage::user(enriched_input))
.await?;
self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
.await
}
async fn handle_skill_response(
&self,
processed_input: &str,
skill_id: &str,
skill_response: String,
input_context: &HashMap<String, Value>,
) -> Result<AgentResponse> {
let output_data = self.process_output(&skill_response, input_context).await?;
let final_response = output_data.content;
self.memory
.add_message(ChatMessage::assistant(&final_response))
.await?;
self.check_memory_compression().await?;
self.increment_turn();
self.evaluate_transitions(processed_input, &final_response)
.await?;
let response = AgentResponse::new(final_response)
.with_metadata("skill_id", serde_json::json!(skill_id));
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn handle_plan_and_execute(
&self,
processed_input: &str,
input_context: &HashMap<String, Value>,
auto_detected: bool,
) -> Result<AgentResponse> {
let effective = self.get_effective_reasoning_config();
let plan_reflection = effective
.get_planning()
.map(|c| c.reflection.clone())
.unwrap_or_default();
let max_attempts = if plan_reflection.enabled {
1 + plan_reflection.max_replans
} else {
1
};
let mut plan = self.generate_plan(processed_input).await?;
info!(
plan_id = %plan.id,
steps = plan.steps.len(),
"Plan generated"
);
let mut plan_result = String::new();
for attempt in 0..max_attempts {
*self.current_plan.write() = Some(plan.clone());
plan_result = self.execute_plan(&mut plan).await?;
info!(
plan_status = ?plan.status,
completed_steps = plan.completed_steps().count(),
attempt = attempt + 1,
"Plan execution completed"
);
if !plan_reflection.enabled {
break;
}
let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
if !has_failures {
break;
}
if attempt + 1 >= max_attempts {
break;
}
match plan_reflection.on_step_failure {
StepFailureAction::Replan => {
info!(attempt = attempt + 1, "Plan had failures, replanning");
plan = self.generate_plan(processed_input).await?;
}
StepFailureAction::Abort => {
warn!("Plan step failed, aborting");
break;
}
StepFailureAction::Skip | StepFailureAction::Continue => {
break;
}
}
}
*self.current_plan.write() = Some(plan);
let output_data = self.process_output(&plan_result, input_context).await?;
let final_content = output_data.content;
self.memory
.add_message(ChatMessage::assistant(&final_content))
.await?;
self.check_memory_compression().await?;
self.increment_turn();
self.evaluate_transitions(processed_input, &final_content)
.await?;
let reasoning_metadata =
ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
let response = AgentResponse::new(&final_content).with_metadata(
"reasoning",
serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
);
self.finish_turn_if_root(&response).await?;
Ok(response)
}
fn inject_reasoning_prompt(
&self,
messages: &mut [ChatMessage],
reasoning_mode: &ReasoningMode,
is_first_iteration: bool,
) {
if !is_first_iteration {
return;
}
match reasoning_mode {
ReasoningMode::CoT => {
if let Some(msg) = messages.first_mut()
&& matches!(msg.role, ai_agents_core::Role::System)
{
msg.content = self.build_cot_system_prompt(&msg.content);
debug!("Applied Chain-of-Thought system prompt");
}
}
ReasoningMode::React => {
if let Some(msg) = messages.first_mut()
&& matches!(msg.role, ai_agents_core::Role::System)
{
msg.content = self.build_react_system_prompt(&msg.content);
debug!("Applied ReAct system prompt");
}
}
_ => {}
}
}
async fn generate_main_response_draft(
&self,
processed_input: &str,
reasoning_mode: &ReasoningMode,
) -> Result<MainResponseDraft> {
let llm = self.get_state_llm()?;
let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
let mut messages = self
.build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
.await?;
self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
let response = self
.complete_main_llm_with_recovery(llm, &messages, &protocol)
.await?;
let content = response.content.trim().to_string();
let (thinking, answer) = self.extract_thinking(&content);
if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
return Ok(MainResponseDraft::ToolCalls {
raw_content: content,
calls,
thinking,
});
}
Ok(MainResponseDraft::Text {
raw_content: answer,
thinking,
})
}
async fn commit_main_response_draft(
&self,
processed_input: &str,
input_context: &HashMap<String, Value>,
draft: MainResponseDraft,
reasoning_mode: ReasoningMode,
auto_detected: bool,
) -> Result<AgentResponse> {
self.commit_root_user_message(processed_input).await?;
match draft {
MainResponseDraft::Text {
raw_content,
thinking,
} => {
self.finish_text_response_from_model(CommittedTextResponse {
processed_input,
input_context,
answer: raw_content,
reasoning_mode,
auto_detected,
iterations: 1,
thinking_content: thinking,
all_tool_calls: Vec::new(),
})
.await
}
MainResponseDraft::ToolCalls {
raw_content,
calls,
thinking: _,
} => {
let mut all_tool_calls = Vec::new();
match self
.handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
.await?
{
ToolCallOutcome::Rejected(response) => {
self.finish_turn_if_root(&response).await?;
Ok(response)
}
ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
self.continue_after_committed_tool_draft(processed_input)
.await
}
}
}
}
}
async fn continue_after_committed_tool_draft(
&self,
processed_input: &str,
) -> Result<AgentResponse> {
*self.redispatch_depth.write() += 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.enter_redispatch();
}
let result = Box::pin(self.run_loop_internal(processed_input)).await;
*self.redispatch_depth.write() -= 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.exit_redispatch();
}
let response = result?;
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn finish_text_response_from_model(
&self,
response: CommittedTextResponse<'_>,
) -> Result<AgentResponse> {
let CommittedTextResponse {
processed_input,
input_context,
answer,
reasoning_mode,
auto_detected,
iterations,
thinking_content,
all_tool_calls,
} = response;
let output_data = self.process_output(&answer, input_context).await?;
let mut final_content = if output_data.metadata.rejected {
output_data
.metadata
.rejection_reason
.unwrap_or_else(|| answer.to_string())
} else {
output_data.content
};
let llm = self.get_state_llm()?;
let reflection_metadata;
(final_content, reflection_metadata) = self
.run_reflection(&*llm, processed_input, final_content)
.await?;
final_content =
self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
let final_content = {
let result = self
.post_loop_processing(processed_input, final_content)
.await?;
self.apply_post_loop_result(processed_input, result).await?
};
let response = self.build_agent_response(AgentResponseParts {
content: final_content,
all_tool_calls,
reasoning_mode,
auto_detected,
iterations,
thinking: thinking_content,
reflection_metadata,
});
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn run_committed_response_loop_with_reasoning(
&self,
processed_input: &str,
input_context: &HashMap<String, Value>,
reasoning_mode: ReasoningMode,
auto_detected: bool,
) -> Result<AgentResponse> {
self.commit_root_user_message(processed_input).await?;
let llm = self.get_state_llm()?;
let mut iterations = 0u32;
let mut all_tool_calls = Vec::new();
let mut thinking_content = None;
loop {
let effective_max = if reasoning_mode != ReasoningMode::None {
let rc = self.get_effective_reasoning_config();
self.max_iterations.min(rc.max_iterations)
} else {
self.max_iterations
};
if iterations >= effective_max {
return Err(AgentError::Other(format!(
"Max iterations ({}) exceeded",
effective_max
)));
}
iterations += 1;
*self.iteration_count.write() = iterations;
let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
let mut messages = self
.build_messages_internal(true, None, protocol.choice.is_none())
.await?;
self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
self.hooks.on_llm_start(&messages).await;
let llm_start = Instant::now();
let response = self
.complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
.await?;
let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
self.hooks.on_llm_complete(&response, llm_duration_ms).await;
let content = response.content.trim();
if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
match self
.handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
.await?
{
ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
ToolCallOutcome::Rejected(resp) => {
self.finish_turn_if_root(&resp).await?;
return Ok(resp);
}
}
}
let (extracted_thinking, answer) = self.extract_thinking(content);
if extracted_thinking.is_some() {
thinking_content = extracted_thinking;
}
return self
.finish_text_response_from_model(CommittedTextResponse {
processed_input,
input_context,
answer,
reasoning_mode,
auto_detected,
iterations,
thinking_content,
all_tool_calls,
})
.await;
}
}
async fn handle_tool_calls(
&self,
processed_input: &str,
content: &str,
tool_calls: Vec<ToolCall>,
all_tool_calls: &mut Vec<ToolCall>,
) -> Result<ToolCallOutcome> {
let transition_fired = self.evaluate_transitions(processed_input, content).await?;
if transition_fired {
self.memory
.add_message(ChatMessage::assistant(
"(Transitioned to new state — tool call handled by workflow)",
))
.await?;
return Ok(ToolCallOutcome::TransitionFired);
}
self.memory
.add_message(ChatMessage::assistant(content))
.await?;
let native_tool_call = Self::is_native_tool_call_content(content);
let results = self.execute_tools_parallel(&tool_calls).await;
for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
match result {
Ok(output) => {
self.memory
.add_message(Self::tool_result_message(
tool_call,
&output,
native_tool_call,
))
.await?;
}
Err(e) => {
if matches!(e, AgentError::HITLRejected(_)) {
self.memory
.add_message(ChatMessage::assistant(format!(
"The operation was rejected by the approver: {}",
e
)))
.await?;
return Ok(ToolCallOutcome::Rejected(AgentResponse {
content: format!("Operation cancelled: {}", e),
metadata: None,
tool_calls: Some(all_tool_calls.clone()),
}));
}
self.memory
.add_message(Self::tool_result_message(
tool_call,
&format!("Error: {}", e),
native_tool_call,
))
.await?;
}
}
all_tool_calls.push(tool_call.clone());
}
Ok(ToolCallOutcome::Continue)
}
async fn run_reflection(
&self,
llm: &dyn LLMProvider,
processed_input: &str,
mut content: String,
) -> Result<(String, Option<ReflectionMetadata>)> {
let should_reflect = self.should_reflect(processed_input, &content).await?;
if !should_reflect {
return Ok((content, None));
}
info!("Starting response reflection evaluation");
let mut attempts = 0u32;
let max_retries = self.reflection_config.max_retries;
let mut history: Vec<ReflectionAttempt> = Vec::new();
loop {
let evaluation = self.evaluate_response(processed_input, &content).await?;
if evaluation.passed || attempts >= max_retries {
info!(
passed = evaluation.passed,
confidence = evaluation.confidence,
attempts = attempts + 1,
"Reflection evaluation complete"
);
let reflection_metadata = Some(
ReflectionMetadata::new(evaluation)
.with_attempts(attempts + 1)
.with_history(history),
);
return Ok((content, reflection_metadata));
}
debug!(
attempt = attempts + 1,
failed_criteria = evaluation.failed_criteria().count(),
"Response did not meet criteria, retrying"
);
history.push(
ReflectionAttempt::new(&content, evaluation.clone())
.with_feedback("Response did not meet quality criteria"),
);
let feedback: Vec<String> = evaluation
.failed_criteria()
.map(|c| format!("- {}", c.criterion))
.collect();
let retry_prompt = format!(
"Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
feedback.join("\n")
);
self.memory
.add_message(ChatMessage::user(&retry_prompt))
.await?;
let retry_messages = self.build_messages().await?;
let retry_response = self
.observe_purpose(
ObservationPurpose::ReflectionEvaluation,
llm.complete(&retry_messages, None),
)
.await
.map_err(|e| AgentError::LLM(e.to_string()))?;
content = retry_response.content.trim().to_string();
attempts += 1;
}
}
async fn post_loop_processing(
&self,
processed_input: &str,
content: String,
) -> Result<PostLoopResult> {
self.increment_turn();
self.run_context_extractors(processed_input).await;
let transitioned = self.evaluate_transitions(processed_input, &content).await?;
if !transitioned {
self.memory
.add_message(ChatMessage::assistant(&content))
.await?;
self.check_memory_compression().await?;
return Ok(PostLoopResult::NoTransition(content));
}
if !self.should_regenerate_after_transition() {
self.memory
.add_message(ChatMessage::assistant(&content))
.await?;
self.check_memory_compression().await?;
return Ok(PostLoopResult::Transitioned(content));
}
if self.needs_redispatch_for_new_state() {
info!("Post-transition NeedsRedispatch: new state requires full dispatch");
return Ok(PostLoopResult::NeedsRedispatch);
}
self.memory
.add_message(ChatMessage::assistant(&content))
.await?;
self.check_memory_compression().await?;
let new_llm = self.get_state_llm()?;
let mut final_content;
for post_iter in 0..self.max_iterations {
let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
let new_messages = self
.build_messages_internal(true, None, protocol.choice.is_none())
.await?;
if post_iter == 0
&& let Some(system_msg) = new_messages.first()
&& system_msg.role == ai_agents_core::Role::System
{
debug!(
prompt_preview =
&system_msg.content[system_msg.content.len().saturating_sub(200)..],
"Post-transition system prompt (last 200 chars)"
);
}
let new_response = self
.complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
.await?;
final_content = new_response.content.trim().to_string();
if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
let native_tool_call = Self::is_native_tool_call_content(&final_content);
debug!(
post_iter = post_iter,
tools = tool_calls.len(),
"Post-transition tool call detected, executing"
);
self.memory
.add_message(ChatMessage::assistant(&final_content))
.await?;
let results = self.execute_tools_parallel(&tool_calls).await;
for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
match result {
Ok(output) => {
self.memory
.add_message(Self::tool_result_message(
tool_call,
&output,
native_tool_call,
))
.await?;
}
Err(e) => {
self.memory
.add_message(Self::tool_result_message(
tool_call,
&format!("Error: {}", e),
native_tool_call,
))
.await?;
}
}
}
continue;
}
self.memory
.add_message(ChatMessage::assistant(&final_content))
.await?;
return Ok(PostLoopResult::Transitioned(final_content));
}
final_content = "Post-transition processing completed.".to_string();
self.memory
.add_message(ChatMessage::assistant(&final_content))
.await?;
Ok(PostLoopResult::Transitioned(final_content))
}
fn should_regenerate_after_transition(&self) -> bool {
if let Some(ref sm) = self.state_machine {
if !sm.config().regenerate_on_transition {
return false;
}
if let Some(def) = sm.current_definition()
&& let Some(regen) = def.regenerate_on_enter
{
return regen;
}
}
true
}
fn needs_redispatch_for_new_state(&self) -> bool {
if let Some(ref sm) = self.state_machine
&& let Some(def) = sm.current_definition()
{
if def.concurrent.is_some()
|| def.group_chat.is_some()
|| def.pipeline.is_some()
|| def.handoff.is_some()
|| def.delegate.is_some()
{
return true;
}
let effective = self.get_effective_reasoning_config();
if !matches!(effective.mode, ReasoningMode::None) {
return true;
}
}
false
}
async fn apply_post_loop_result(
&self,
processed_input: &str,
result: PostLoopResult,
) -> Result<String> {
match result {
PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
Ok(content)
}
PostLoopResult::NeedsRedispatch => {
const MAX_REDISPATCH_DEPTH: u32 = 3;
let current_depth = *self.redispatch_depth.read();
if current_depth >= MAX_REDISPATCH_DEPTH {
warn!(
depth = current_depth,
"Post-transition re-dispatch depth limit reached, returning empty response"
);
let content = String::new();
self.memory
.add_message(ChatMessage::assistant(&content))
.await?;
return Ok(content);
}
*self.redispatch_depth.write() += 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.enter_redispatch();
}
info!(
depth = current_depth + 1,
"Re-dispatching for new state after transition"
);
let resp = Box::pin(self.run_loop_internal(processed_input)).await;
*self.redispatch_depth.write() -= 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.exit_redispatch();
}
resp.map(|r| r.content)
}
}
}
fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
let AgentResponseParts {
content,
all_tool_calls,
reasoning_mode,
auto_detected,
iterations,
thinking,
reflection_metadata,
} = parts;
let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
.with_thinking(thinking.clone().unwrap_or_default())
.with_iterations(iterations)
.with_auto_detected(auto_detected);
let mut response = AgentResponse::new(&content);
if !all_tool_calls.is_empty() {
response = response.with_tool_calls(all_tool_calls);
}
if let Some(state) = self.current_state() {
response = response.with_metadata("current_state", serde_json::json!(state));
}
response = response.with_metadata(
"reasoning",
serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
);
if let Some(ref refl_meta) = reflection_metadata {
response = response.with_metadata(
"reflection",
serde_json::to_value(refl_meta).unwrap_or_default(),
);
}
response
}
async fn handle_delegated_state(
&self,
input: &str,
delegate_id: &str,
state_def: &ai_agents_state::StateDefinition,
) -> Result<AgentResponse> {
use std::time::Instant;
let registry = self.spawner_registry.as_ref().ok_or_else(|| {
AgentError::Config(format!(
"State delegates to '{}' but no agent registry is configured. \
Add a spawner section with auto_spawn to your YAML.",
delegate_id
))
})?;
let state_name = self
.state_machine
.as_ref()
.map(|sm| sm.current())
.unwrap_or_else(|| "unknown".to_string());
self.hooks.on_delegate_start(delegate_id, &state_name).await;
let start = Instant::now();
let delegate = registry.get(delegate_id).ok_or_else(|| {
AgentError::Other(format!(
"State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
state_name, delegate_id
))
})?;
let context_mode = state_def.delegate_context.clone().unwrap_or_default();
let effective_input = self
.observe_purpose(
ObservationPurpose::OrchestrationRouting,
crate::orchestration::context::prepare_delegate_input(
input,
&context_mode,
&*self.memory,
self.llm_registry.get("router").ok().as_deref(),
),
)
.await?;
let response = delegate
.chat_with_actor_context(&effective_input, self.outbound_actor_context())
.await?;
let duration_ms = start.elapsed().as_millis() as u64;
self.hooks
.on_delegate_complete(delegate_id, &state_name, duration_ms)
.await;
let ctx_key = format!("delegation.{}.last_response", delegate_id);
let _ = self.context_manager.set(
&ctx_key,
serde_json::Value::String(response.content.clone()),
);
let _ = self.context_manager.set(
"orchestration",
serde_json::json!({
"type": "delegate",
"agent": delegate_id,
"state": state_name,
"response": response.content,
"duration_ms": duration_ms,
}),
);
self.commit_root_user_message(input).await?;
let post_result = self
.post_loop_processing(
input,
format!("[Delegated to {}]: {}", delegate_id, response.content),
)
.await?;
let final_content = self.apply_post_loop_result(input, post_result).await?;
let mut result = AgentResponse::new(final_content);
let metadata = serde_json::json!({
"orchestration": {
"type": "delegate",
"agent": delegate_id,
"state": state_name,
"response": response.content,
"duration_ms": duration_ms,
}
});
result.metadata = Some(
serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
metadata,
)
.unwrap_or_default(),
);
self.finish_turn_if_root(&result).await?;
Ok(result)
}
async fn handle_concurrent_state(
&self,
input: &str,
config: &ai_agents_state::ConcurrentStateConfig,
) -> Result<AgentResponse> {
use std::time::Instant;
let registry = self.spawner_registry.as_ref().ok_or_else(|| {
AgentError::Config(
"Concurrent state requires an agent registry. Add a spawner section.".into(),
)
})?;
let context_mode = config.context_mode.clone().unwrap_or_default();
let context_input = self
.observe_purpose(
ObservationPurpose::OrchestrationRouting,
crate::orchestration::context::prepare_delegate_input(
input,
&context_mode,
&*self.memory,
self.llm_registry.get("router").ok().as_deref(),
),
)
.await?;
let effective_input = if let Some(ref tmpl) = config.input {
render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
.unwrap_or_else(|_| context_input.clone())
} else {
context_input
};
let start = Instant::now();
let llm_name = config
.aggregation
.synthesizer_llm
.as_deref()
.unwrap_or("router");
let llm_provider = self.llm_registry.get(llm_name).ok();
let vote_parallelism = if self.runtime_config.optimization.enabled
&& self
.runtime_config
.optimization
.parallel_orchestration_vote_extraction
{
Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
} else {
None
};
let result = self
.observe_purpose(
ObservationPurpose::OrchestrationAggregation,
scope_actor_context(
self.outbound_actor_context(),
crate::orchestration::concurrent(
registry,
&effective_input,
&config.agents,
&config.aggregation,
llm_provider.as_deref(),
config.min_required,
config.timeout_ms,
config.on_partial_failure.clone(),
vote_parallelism,
),
),
)
.await?;
let duration_ms = start.elapsed().as_millis() as u64;
let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
let strategy = format!("{:?}", config.aggregation.strategy);
self.hooks
.on_concurrent_complete(&agent_ids, &strategy, duration_ms)
.await;
let _ = self.context_manager.set(
"concurrent.result",
serde_json::Value::String(result.response.content.clone()),
);
let agents_json: Vec<serde_json::Value> = result
.agent_results
.iter()
.map(|ar| {
serde_json::json!({
"id": ar.agent_id,
"response": ar.response.as_ref().map(|r| r.content.as_str()),
"success": ar.success,
"error": ar.error,
"duration_ms": ar.duration_ms,
})
})
.collect();
let _ = self.context_manager.set(
"orchestration",
serde_json::json!({
"type": "concurrent",
"result": result.response.content,
"strategy": strategy,
"agents": agents_json,
"duration_ms": duration_ms,
}),
);
self.commit_root_user_message(input).await?;
let post_result = self
.post_loop_processing(input, result.response.content.clone())
.await?;
let final_content = self.apply_post_loop_result(input, post_result).await?;
let mut response = AgentResponse::new(final_content);
let metadata = serde_json::json!({
"orchestration": {
"type": "concurrent",
"result": result.response.content,
"strategy": strategy,
"agents": agents_json,
"duration_ms": duration_ms,
}
});
response.metadata = Some(
serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
metadata,
)
.unwrap_or_default(),
);
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn handle_group_chat_state(
&self,
input: &str,
config: &ai_agents_state::GroupChatStateConfig,
) -> Result<AgentResponse> {
use std::time::Instant;
let registry = self.spawner_registry.as_ref().ok_or_else(|| {
AgentError::Config(
"Group chat state requires an agent registry. Add a spawner section.".into(),
)
})?;
let start = Instant::now();
let llm_provider = self.llm_registry.get("router").ok();
let context_mode = config.context_mode.clone().unwrap_or_default();
let context_input = self
.observe_purpose(
ObservationPurpose::OrchestrationRouting,
crate::orchestration::context::prepare_delegate_input(
input,
&context_mode,
&*self.memory,
self.llm_registry.get("router").ok().as_deref(),
),
)
.await?;
let effective_topic = if let Some(ref tmpl) = config.input {
render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
.unwrap_or_else(|_| context_input.clone())
} else {
context_input
};
let result = self
.observe_purpose(
ObservationPurpose::OrchestrationConversation,
scope_actor_context(
self.outbound_actor_context(),
crate::orchestration::group_chat(
registry,
&effective_topic,
config,
llm_provider.as_deref(),
Some(&*self.hooks),
),
),
)
.await?;
let duration_ms = start.elapsed().as_millis() as u64;
let _ = self.context_manager.set(
"group_chat.conclusion",
serde_json::Value::String(result.response.content.clone()),
);
let transcript_json: Vec<serde_json::Value> = result
.transcript
.iter()
.map(|t| {
serde_json::json!({
"speaker": t.speaker,
"round": t.round,
"content": t.content,
})
})
.collect();
let _ = self.context_manager.set(
"orchestration",
serde_json::json!({
"type": "group_chat",
"conclusion": result.response.content,
"transcript": transcript_json,
"rounds": result.rounds_completed,
"termination": result.termination_reason,
"duration_ms": duration_ms,
}),
);
self.commit_root_user_message(input).await?;
let post_result = self
.post_loop_processing(input, result.response.content.clone())
.await?;
let final_content = self.apply_post_loop_result(input, post_result).await?;
let mut response = AgentResponse::new(final_content);
let metadata = serde_json::json!({
"orchestration": {
"type": "group_chat",
"conclusion": result.response.content,
"transcript": transcript_json,
"rounds": result.rounds_completed,
"termination": result.termination_reason,
"duration_ms": duration_ms,
}
});
response.metadata = Some(
serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
metadata,
)
.unwrap_or_default(),
);
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn handle_pipeline_state(
&self,
input: &str,
config: &ai_agents_state::PipelineStateConfig,
) -> Result<AgentResponse> {
use std::time::Instant;
let registry = self.spawner_registry.as_ref().ok_or_else(|| {
AgentError::Config(
"Pipeline state requires an agent registry. Add a spawner section.".into(),
)
})?;
let start = Instant::now();
let stages: Vec<crate::orchestration::PipelineStage> = config
.stages
.iter()
.map(|entry| {
let mut stage = crate::orchestration::PipelineStage::id(entry.id());
if let Some(tmpl) = entry.input() {
stage = stage.with_input(tmpl);
}
stage
})
.collect();
let context_mode = config.context_mode.clone().unwrap_or_default();
let context_input = self
.observe_purpose(
ObservationPurpose::OrchestrationRouting,
crate::orchestration::context::prepare_delegate_input(
input,
&context_mode,
&*self.memory,
self.llm_registry.get("router").ok().as_deref(),
),
)
.await?;
let context_values = self.build_context_with_overlays();
let result = self
.observe_purpose(
ObservationPurpose::OrchestrationRouting,
scope_actor_context(
self.outbound_actor_context(),
crate::orchestration::pipeline(
registry,
&context_input,
&stages,
config.timeout_ms,
Some(&*self.hooks),
Some(&context_values),
),
),
)
.await?;
let duration_ms = start.elapsed().as_millis() as u64;
let _ = self.context_manager.set(
"pipeline.result",
serde_json::Value::String(result.response.content.clone()),
);
let stages_json: Vec<serde_json::Value> = result
.stage_outputs
.iter()
.map(|s| {
serde_json::json!({
"agent_id": s.agent_id,
"output": s.output,
"duration_ms": s.duration_ms,
"skipped": s.skipped,
})
})
.collect();
let _ = self.context_manager.set(
"orchestration",
serde_json::json!({
"type": "pipeline",
"result": result.response.content,
"stages": stages_json,
"duration_ms": duration_ms,
}),
);
self.commit_root_user_message(input).await?;
let post_result = self
.post_loop_processing(input, result.response.content.clone())
.await?;
let final_content = self.apply_post_loop_result(input, post_result).await?;
let mut response = AgentResponse::new(final_content);
let metadata = serde_json::json!({
"orchestration": {
"type": "pipeline",
"result": result.response.content,
"stages": stages_json,
"duration_ms": duration_ms,
}
});
response.metadata = Some(
serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
metadata,
)
.unwrap_or_default(),
);
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn handle_handoff_state(
&self,
input: &str,
config: &ai_agents_state::HandoffStateConfig,
) -> Result<AgentResponse> {
use std::time::Instant;
let registry = self.spawner_registry.as_ref().ok_or_else(|| {
AgentError::Config(
"Handoff state requires an agent registry. Add a spawner section.".into(),
)
})?;
let llm = self
.llm_registry
.get("router")
.map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
let start = Instant::now();
let context_mode = config.context_mode.clone().unwrap_or_default();
let context_input = self
.observe_purpose(
ObservationPurpose::OrchestrationRouting,
crate::orchestration::context::prepare_delegate_input(
input,
&context_mode,
&*self.memory,
self.llm_registry.get("router").ok().as_deref(),
),
)
.await?;
let effective_input = if let Some(ref tmpl) = config.input {
render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
.unwrap_or_else(|_| context_input.clone())
} else {
context_input
};
let result = self
.observe_purpose(
ObservationPurpose::OrchestrationRouting,
scope_actor_context(
self.outbound_actor_context(),
crate::orchestration::handoff(
registry,
&effective_input,
&config.initial_agent,
&config.available_agents,
config.max_handoffs,
llm.as_ref(),
Some(&*self.hooks),
),
),
)
.await?;
let duration_ms = start.elapsed().as_millis() as u64;
let _ = self.context_manager.set(
"handoff.result",
serde_json::Value::String(result.response.content.clone()),
);
let chain_json: Vec<serde_json::Value> = result
.handoff_chain
.iter()
.map(|h| {
serde_json::json!({
"from": h.from_agent,
"to": h.to_agent,
"reason": h.reason,
})
})
.collect();
let _ = self.context_manager.set(
"orchestration",
serde_json::json!({
"type": "handoff",
"result": result.response.content,
"final_agent": result.final_agent,
"handoff_chain": chain_json,
"duration_ms": duration_ms,
}),
);
self.commit_root_user_message(input).await?;
let post_result = self
.post_loop_processing(input, result.response.content.clone())
.await?;
let final_content = self.apply_post_loop_result(input, post_result).await?;
let mut response = AgentResponse::new(final_content);
let metadata = serde_json::json!({
"orchestration": {
"type": "handoff",
"result": result.response.content,
"final_agent": result.final_agent,
"handoff_chain": chain_json,
"duration_ms": duration_ms,
}
});
response.metadata = Some(
serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
metadata,
)
.unwrap_or_default(),
);
self.finish_turn_if_root(&response).await?;
Ok(response)
}
async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
self.begin_root_turn();
self.pre_turn_session_lifecycle().await;
let input_data = self.process_input(input).await?;
self.update_active_turn_context(&input_data.content, input_data.context.clone());
for (key, value) in &input_data.context {
let _ = self.context_manager.set(key, value.clone());
}
if input_data.metadata.rejected {
let reason = input_data
.metadata
.rejection_reason
.unwrap_or_else(|| "Input rejected".to_string());
warn!(reason = %reason, "Input rejected");
let response = AgentResponse::new(reason);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
let processed_input = &input_data.content;
if let Some(response) = self.try_pre_response_transition(processed_input).await? {
return Ok(response);
}
if let Some(ref sm) = self.state_machine
&& let Some(def) = sm.current_definition()
{
if let Some(ref delegate_id) = def.delegate {
return self
.handle_delegated_state(processed_input, delegate_id, &def)
.await;
}
if let Some(ref concurrent_config) = def.concurrent {
return self
.handle_concurrent_state(processed_input, concurrent_config)
.await;
}
if let Some(ref group_chat_config) = def.group_chat {
return self
.handle_group_chat_state(processed_input, group_chat_config)
.await;
}
if let Some(ref pipeline_config) = def.pipeline {
return self
.handle_pipeline_state(processed_input, pipeline_config)
.await;
}
if let Some(ref handoff_config) = def.handoff {
return self
.handle_handoff_state(processed_input, handoff_config)
.await;
}
}
if let Some(response) =
Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
{
return Ok(response);
}
match self.try_skill_route(processed_input).await? {
SkillRouteResult::Response { skill_id, content } => {
self.commit_root_user_message(processed_input).await?;
return self
.handle_skill_response(processed_input, &skill_id, content, &input_data.context)
.await;
}
SkillRouteResult::NeedsClarification {
response,
ownership,
} => {
let admission = self
.admit_optional_disambiguation_ownership(ownership)
.await?;
self.commit_root_user_message(processed_input).await?;
if let Some(q) = response
.metadata
.as_ref()
.and_then(|m| m.get("disambiguation"))
.and_then(|d| d.get("status"))
.and_then(|s| s.as_str())
&& q == "awaiting_clarification"
{
self.memory
.add_message(ChatMessage::assistant(&response.content))
.await?;
}
drop(admission);
self.finish_turn_if_root(&response).await?;
return Ok(response);
}
SkillRouteResult::NoMatch => {} }
let effective_reasoning = self.get_effective_reasoning_config();
let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
info!(
reasoning_mode = ?reasoning_mode,
auto_detected = auto_detected,
reflection_enabled = ?self.reflection_config.enabled,
"Reasoning mode determined"
);
if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
self.commit_root_user_message(processed_input).await?;
return self
.handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
.await;
}
self.commit_root_user_message(processed_input).await?;
let mut iterations = 0u32;
let mut all_tool_calls: Vec<ToolCall> = Vec::new();
let mut thinking_content: Option<String> = None;
let llm = self.get_state_llm()?;
loop {
let effective_max = if reasoning_mode != ReasoningMode::None {
let rc = self.get_effective_reasoning_config();
self.max_iterations.min(rc.max_iterations)
} else {
self.max_iterations
};
if iterations >= effective_max {
let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
self.hooks.on_error(&err).await;
error!(iterations = iterations, "Max iterations exceeded");
return Err(err);
}
iterations += 1;
*self.iteration_count.write() = iterations;
debug!(iteration = iterations, max = effective_max, "LLM call");
let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
let mut messages = self
.build_messages_internal(true, None, protocol.choice.is_none())
.await?;
self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
self.hooks.on_llm_start(&messages).await;
let llm_start = Instant::now();
let response = self
.complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
.await?;
let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
self.hooks.on_llm_complete(&response, llm_duration_ms).await;
let content = response.content.trim();
if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
match self
.handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
.await?
{
ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
ToolCallOutcome::Rejected(resp) => {
self.finish_turn_if_root(&resp).await?;
return Ok(resp);
}
}
}
let (extracted_thinking, answer) = self.extract_thinking(content);
if extracted_thinking.is_some() {
thinking_content = extracted_thinking;
}
let output_data = self.process_output(&answer, &input_data.context).await?;
let mut final_content = if output_data.metadata.rejected {
output_data
.metadata
.rejection_reason
.unwrap_or_else(|| answer.to_string())
} else {
output_data.content
};
let reflection_metadata;
(final_content, reflection_metadata) = self
.run_reflection(&*llm, processed_input, final_content)
.await?;
final_content =
self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
let final_content = {
let result = self
.post_loop_processing(processed_input, final_content)
.await?;
self.apply_post_loop_result(processed_input, result).await?
};
let reflected = reflection_metadata.is_some();
let reasoning_mode_debug = format!("{:?}", reasoning_mode);
let response = self.build_agent_response(AgentResponseParts {
content: final_content,
all_tool_calls,
reasoning_mode,
auto_detected,
iterations,
thinking: thinking_content,
reflection_metadata,
});
self.finish_turn_if_root(&response).await?;
let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
info!(
tool_calls = tool_call_count,
response_len = response.content.len(),
reasoning_mode = %reasoning_mode_debug,
reflected = reflected,
"Chat completed"
);
return Ok(response);
}
}
async fn generate_buffered_streaming_draft(
&self,
processed_input: &str,
routing_resolved: Arc<AtomicBool>,
) -> Result<StreamingDraftResult> {
let llm = self.get_state_llm()?;
if llm.configured_tool_choice().is_some() {
let draft = self
.generate_main_response_draft(processed_input, &ReasoningMode::None)
.await?;
return Ok(StreamingDraftResult::new(draft, Vec::new()));
}
let messages = self.build_messages_for_draft(processed_input).await?;
let mut stream = self
.observe_purpose(
ObservationPurpose::MainResponse,
llm.complete_stream(&messages, None),
)
.await
.map_err(|e| AgentError::LLM(e.to_string()))?;
let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
let mut chunks = Vec::new();
let mut accumulated = String::new();
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
accumulated.push_str(&chunk.delta);
let stream_chunk = StreamChunk::content(chunk.delta);
if routing_resolved.load(Ordering::SeqCst) {
chunks.push(stream_chunk);
} else {
buffer.push(stream_chunk)?;
}
}
chunks.splice(0..0, buffer.drain());
let content = accumulated.trim().to_string();
let draft = if let Some(calls) = self.parse_tool_calls(&content) {
MainResponseDraft::ToolCalls {
raw_content: content,
calls,
thinking: None,
}
} else {
MainResponseDraft::Text {
raw_content: content,
thinking: None,
}
};
Ok(StreamingDraftResult::new(draft, chunks))
}
async fn try_buffered_streaming_branches(
&self,
processed_input: &str,
input_context: &HashMap<String, Value>,
) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
let optimization = &self.runtime_config.optimization;
if !optimization.enabled {
return Ok(None);
}
let transition_enabled =
optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
if !transition_enabled {
return Ok(None);
}
let mut branch_scheduler =
TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
if !branch_scheduler.reserve_task() {
return Ok(None);
}
if !self
.reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
{
branch_scheduler.release_task();
return Ok(None);
}
if !branch_scheduler.reserve_task() {
branch_scheduler.release_task();
return Ok(None);
}
let mut main_branch = RuntimeBranch::new(
RuntimeTaskPurpose::MainResponse,
RuntimeOptimizationKind::BufferedStreamingRouting,
RuntimeTaskPriority::Normal,
RuntimeCommitBehavior::FinalResponse,
);
let mut transition_branch = RuntimeBranch::new(
RuntimeTaskPurpose::StateTransition,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeTaskPriority::Critical,
RuntimeCommitBehavior::TransitionDecision,
);
let main_id = main_branch.branch_id();
let transition_id = transition_branch.branch_id();
let routing_resolved = Arc::new(AtomicBool::new(false));
let mut main_future =
Box::pin(crate::optimization::observability::with_branch_observation(
&main_id,
RuntimeOptimizationKind::BufferedStreamingRouting,
RuntimeCommitBehavior::FinalResponse,
self.generate_buffered_streaming_draft(
processed_input,
Arc::clone(&routing_resolved),
),
));
let mut transition_future =
Box::pin(crate::optimization::observability::with_branch_observation(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
self.select_parallel_transition_candidate(processed_input),
));
let mut main_pending = true;
let mut transition_pending = true;
let mut main_result: Option<Result<StreamingDraftResult>> = None;
let mut transition_finalized = false;
let mut transition_candidate: Option<TransitionCandidate> = None;
loop {
if let Some(candidate) = transition_candidate.take() {
if self
.approve_transition_target(&candidate.from_state, candidate.target())
.await?
{
drop(main_future);
drop(transition_future);
self.finalize_branch_loss(
&main_id,
RuntimeOptimizationKind::BufferedStreamingRouting,
RuntimeCommitBehavior::FinalResponse,
main_pending,
main_result.as_ref().map(|result| result.is_err()),
);
if !self
.apply_pre_response_transition_candidate(
&candidate,
&HashMap::new(),
processed_input,
)
.await?
{
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"discarded",
false,
);
return Ok(None);
}
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"committed",
true,
);
let response = self.redispatch_current_state(processed_input).await?;
return Ok(Some((
response.clone(),
vec![StreamChunk::content(response.content)],
)));
}
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"discarded",
false,
);
routing_resolved.store(true, Ordering::SeqCst);
transition_finalized = true;
}
if transition_finalized && let Some(result) = main_result.take() {
let stream_draft = match result {
Ok(stream_draft) => stream_draft,
Err(error) => {
self.finalize_optional_branch(
&main_id,
RuntimeOptimizationKind::BufferedStreamingRouting,
RuntimeCommitBehavior::FinalResponse,
"failed",
false,
);
return Err(error);
}
};
let raw_draft_content = stream_draft.draft.raw_content().to_string();
let buffered_chunks = stream_draft.chunks;
self.finalize_optional_branch(
&main_id,
RuntimeOptimizationKind::BufferedStreamingRouting,
RuntimeCommitBehavior::FinalResponse,
"committed",
true,
);
let response = self
.commit_main_response_draft(
processed_input,
input_context,
stream_draft.draft,
ReasoningMode::None,
false,
)
.await?;
let chunks = if response.content == raw_draft_content {
buffered_chunks
} else {
vec![StreamChunk::content(response.content.clone())]
};
return Ok(Some((response, chunks)));
}
tokio::select! {
result = &mut main_future, if main_pending => {
main_pending = false;
main_branch.transition_to(RuntimeBranchStatus::Completed)?;
main_result = Some(result);
}
result = &mut transition_future, if transition_pending => {
transition_pending = false;
transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
match result {
Ok(ParallelTransitionSelection::Candidate(candidate)) => {
transition_candidate = Some(candidate)
}
Ok(ParallelTransitionSelection::NoMatch) => {
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"discarded",
false,
);
routing_resolved.store(true, Ordering::SeqCst);
transition_finalized = true;
}
Ok(ParallelTransitionSelection::ReservationExhausted) => {
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"cancelled",
false,
);
routing_resolved.store(true, Ordering::SeqCst);
self.finalize_branch_loss(
&main_id,
RuntimeOptimizationKind::BufferedStreamingRouting,
RuntimeCommitBehavior::FinalResponse,
main_pending,
main_result.as_ref().map(|result| result.is_err()),
);
return Ok(None);
}
Err(_) => {
self.finalize_optional_branch(
&transition_id,
RuntimeOptimizationKind::ParallelStateTransition,
RuntimeCommitBehavior::TransitionDecision,
"failed",
false,
);
routing_resolved.store(true, Ordering::SeqCst);
transition_finalized = true;
}
}
}
}
}
}
fn run_loop_internal_stream<'a>(
&'a self,
input: &'a str,
terminal: RuntimeStreamTerminalSlot,
) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
let include_tool_events = self.streaming.include_tool_events;
let include_state_events = self.streaming.include_state_events;
Box::pin(async_stream::stream! {
self.begin_root_turn();
self.pre_turn_session_lifecycle().await;
let input_data = match self.process_input(input).await {
Ok(data) => data,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
self.update_active_turn_context(&input_data.content, input_data.context.clone());
for (key, value) in &input_data.context {
let _ = self.context_manager.set(key, value.clone());
}
if input_data.metadata.rejected {
let reason = input_data
.metadata
.rejection_reason
.unwrap_or_else(|| "Input rejected".to_string());
warn!(reason = %reason, "Input rejected (stream)");
yield StreamChunk::error(reason);
return;
}
let processed_input = &input_data.content;
if self.runtime_config.optimization.enabled
&& matches!(
self.runtime_config.optimization.streaming_policy,
crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
)
{
match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
Ok(Some((response, chunks))) => {
for chunk in chunks {
yield chunk;
}
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
return;
}
Ok(None) => {}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
if self.runtime_config.optimization.enabled
&& matches!(
self.runtime_config.optimization.streaming_policy,
crate::optimization::StreamingOptimizationPolicy::PreflightOnly
)
{
match self.try_pre_response_transition(processed_input).await {
Ok(Some(response)) => {
yield StreamChunk::content(&response.content);
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
return;
}
Ok(None) => {}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
if let Some(ref sm) = self.state_machine
&& let Some(def) = sm.current_definition()
{
let orchestration_result = if let Some(ref delegate_id) = def.delegate {
Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
} else if let Some(ref concurrent_config) = def.concurrent {
Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
} else if let Some(ref group_chat_config) = def.group_chat {
Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
} else if let Some(ref pipeline_config) = def.pipeline {
Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
} else if let Some(ref handoff_config) = def.handoff {
Some(self.handle_handoff_state(processed_input, handoff_config).await)
} else {
None
};
if let Some(result) = orchestration_result {
match result {
Ok(response) => {
yield StreamChunk::content(&response.content);
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
}
Err(e) => {
yield StreamChunk::error(e.to_string());
}
}
return;
}
}
match self.try_skill_route(processed_input).await {
Ok(SkillRouteResult::Response { skill_id, content }) => {
if let Err(e) = self.commit_root_user_message(processed_input).await {
yield StreamChunk::error(e.to_string());
return;
}
match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
Ok(resp) => {
yield StreamChunk::content(&resp.content);
record_runtime_stream_final(&terminal, resp);
yield StreamChunk::Done {};
return;
}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
Ok(SkillRouteResult::NeedsClarification {
response,
ownership,
}) => {
let admission = match self
.admit_optional_disambiguation_ownership(ownership)
.await
{
Ok(admission) => admission,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
if let Err(e) = self.commit_root_user_message(processed_input).await {
yield StreamChunk::error(e.to_string());
return;
}
let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
drop(admission);
if let Err(e) = self.finish_turn_if_root(&response).await {
yield StreamChunk::error(e.to_string());
return;
}
yield StreamChunk::content(&response.content);
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
return;
}
Ok(SkillRouteResult::NoMatch) => {} Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
let effective_reasoning = self.get_effective_reasoning_config();
let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
Ok(mode) => mode,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
info!(
reasoning_mode = ?reasoning_mode,
auto_detected = auto_detected,
"Reasoning mode determined (stream)"
);
if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
if let Err(e) = self.commit_root_user_message(processed_input).await {
yield StreamChunk::error(e.to_string());
return;
}
match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
Ok(resp) => {
yield StreamChunk::content(&resp.content);
record_runtime_stream_final(&terminal, resp);
yield StreamChunk::Done {};
return;
}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
if let Err(e) = self.commit_root_user_message(processed_input).await {
yield StreamChunk::error(e.to_string());
return;
}
let llm = match self.get_state_llm() {
Ok(llm) => llm,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let mut iterations = 0u32;
let mut all_tool_calls: Vec<ToolCall> = Vec::new();
let mut thinking_content: Option<String> = None;
loop {
let effective_max = if reasoning_mode != ReasoningMode::None {
let rc = self.get_effective_reasoning_config();
self.max_iterations.min(rc.max_iterations)
} else {
self.max_iterations
};
if iterations >= effective_max {
let err_msg = format!("Max iterations ({}) exceeded", effective_max);
let err = AgentError::Other(err_msg.clone());
self.hooks.on_error(&err).await;
error!(iterations = iterations, "Max iterations exceeded (stream)");
yield StreamChunk::error(err_msg);
return;
}
iterations += 1;
*self.iteration_count.write() = iterations;
debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
Ok(protocol) => protocol,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let mut messages = match self
.build_messages_internal(true, None, protocol.choice.is_none())
.await
{
Ok(m) => m,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
self.hooks.on_llm_start(&messages).await;
let llm_start = Instant::now();
let reflection_active = self
.should_reflect(processed_input, "")
.await
.unwrap_or_default();
let buffered_decision = reflection_active || protocol.choice.is_some();
let content = if buffered_decision {
let response = match self
.complete_main_llm_with_recovery(
Arc::clone(&llm),
&messages,
&protocol,
)
.await
{
Ok(r) => r,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
self.hooks.on_llm_complete(&response, llm_duration_ms).await;
response.content.trim().to_string()
} else {
let llm_stream = match self
.observe_purpose(
ObservationPurpose::MainResponse,
llm.complete_stream(&messages, None),
)
.await
{
Ok(s) => s,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let mut accumulated = String::new();
let mut stream_inner = llm_stream;
while let Some(chunk_result) = stream_inner.next().await {
match chunk_result {
Ok(chunk) => {
accumulated.push_str(&chunk.delta);
yield StreamChunk::content(chunk.delta);
}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
let llm_response = ai_agents_core::LLMResponse::new(
accumulated.trim(),
ai_agents_core::FinishReason::Stop,
);
self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
accumulated.trim().to_string()
};
if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
let native_tool_call = Self::is_native_tool_call_content(&content);
let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
Ok(v) => v,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
if transition_fired {
let _ = self.memory.add_message(ChatMessage::assistant(
"(Transitioned to new state — tool call handled by workflow)",
)).await;
if include_state_events
&& let Some(state) = self.current_state()
{
yield StreamChunk::state_transition(None, state);
}
continue;
}
let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
let results = self.execute_tools_parallel(&tool_calls).await;
for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
if include_tool_events {
yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
}
match result {
Ok(output) => {
if include_tool_events {
yield StreamChunk::tool_result(
&tool_call.id,
&tool_call.name,
&output,
true,
);
}
let _ = self.memory
.add_message(Self::tool_result_message(
tool_call,
&output,
native_tool_call,
))
.await;
}
Err(e) => {
if matches!(e, AgentError::HITLRejected(_)) {
let _ = self.memory.add_message(ChatMessage::assistant(
format!("The operation was rejected by the approver: {}", e),
)).await;
let response = AgentResponse {
content: format!("Operation cancelled: {}", e),
metadata: None,
tool_calls: Some(all_tool_calls.clone()),
};
if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
yield StreamChunk::error(finalize_error.to_string());
return;
}
let legacy_error = response.content.clone();
record_runtime_stream_final(&terminal, response);
yield StreamChunk::error(legacy_error);
yield StreamChunk::Done {};
return;
}
if include_tool_events {
yield StreamChunk::tool_result(
&tool_call.id,
&tool_call.name,
e.to_string(),
false,
);
}
let _ = self.memory
.add_message(Self::tool_result_message(
tool_call,
&format!("Error: {}", e),
native_tool_call,
))
.await;
}
}
all_tool_calls.push(tool_call.clone());
if include_tool_events {
yield StreamChunk::tool_end(&tool_call.id);
}
}
continue;
}
let (extracted_thinking, answer) = self.extract_thinking(&content);
if extracted_thinking.is_some() {
thinking_content = extracted_thinking;
}
let output_data = match self.process_output(&answer, &input_data.context).await {
Ok(d) => d,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let final_content = if output_data.metadata.rejected {
output_data
.metadata
.rejection_reason
.unwrap_or_else(|| answer.to_string())
} else {
output_data.content
};
let (final_content, reflection_metadata) = match self
.run_reflection(&*llm, processed_input, final_content)
.await
{
Ok(r) => r,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let final_content = self.format_response_with_thinking(
thinking_content.as_deref(),
&final_content,
);
if buffered_decision {
yield StreamChunk::content(&final_content);
}
let post_result = match self
.post_loop_processing(processed_input, final_content)
.await
{
Ok(r) => r,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let (final_content, transitioned) = match post_result {
PostLoopResult::NoTransition(content) => (content, false),
PostLoopResult::Transitioned(content) => (content, true),
PostLoopResult::NeedsRedispatch => {
const MAX_REDISPATCH_DEPTH: u32 = 3;
let current_depth = *self.redispatch_depth.read();
let content = if current_depth >= MAX_REDISPATCH_DEPTH {
warn!(
depth = current_depth,
"Post-transition re-dispatch depth limit reached (stream)"
);
let c = String::new();
let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
c
} else {
*self.redispatch_depth.write() += 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.enter_redispatch();
}
info!(
depth = current_depth + 1,
"Re-dispatching for new state after transition (stream)"
);
let result = self.run_loop_internal(processed_input).await;
*self.redispatch_depth.write() -= 1;
if let Some(context) = self.active_turn_context.write().as_mut() {
context.exit_redispatch();
}
match result {
Ok(resp) => resp.content,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
};
(content, true)
}
};
if transitioned {
if include_state_events
&& let Some(state) = self.current_state()
{
yield StreamChunk::state_transition(None, state);
}
yield StreamChunk::content(&final_content);
}
let final_response = self.build_agent_response(AgentResponseParts {
content: final_content,
all_tool_calls,
reasoning_mode,
auto_detected,
iterations,
thinking: thinking_content,
reflection_metadata,
});
if let Err(e) = self.finish_turn_if_root(&final_response).await {
yield StreamChunk::error(e.to_string());
return;
}
record_runtime_stream_final(&terminal, final_response);
yield StreamChunk::Done {};
return;
}
})
}
fn run_loop_stream<'a>(
&'a self,
input: &'a str,
terminal: RuntimeStreamTerminalSlot,
) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
Box::pin(async_stream::stream! {
self.begin_root_turn();
let _root_cleanup = RootTurnCleanup::new(self);
self.hooks.on_message_received(input).await;
if !self.context_initialized.swap(true, Ordering::SeqCst) {
if let Err(e) = self.context_manager.initialize().await {
yield StreamChunk::error(e.to_string());
return;
}
debug!("Context manager initialized (defaults, env, builtins)");
}
if let Err(e) = self.check_turn_timeout().await {
yield StreamChunk::error(e.to_string());
return;
}
if let Err(e) = self.context_manager.refresh_per_turn().await {
yield StreamChunk::error(e.to_string());
return;
}
self.clear_disambiguation_context();
if let Some(ref disambiguator) = self.disambiguation_manager {
let disambiguation_context = match self.build_disambiguation_context().await {
Ok(ctx) => ctx,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let state_override = self
.state_machine
.as_ref()
.and_then(|sm| sm.current_definition())
.and_then(|def| def.disambiguation.clone());
let state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
let mut result = match self
.observe_purpose(
ObservationPurpose::DisambiguationDetection,
disambiguator.process_input_with_override(
input,
&disambiguation_context,
state_override.as_ref(),
None,
),
)
.await
{
Ok(r) => r,
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
};
let current_state_generation = self
.state_machine
.as_ref()
.map(|state_machine| state_machine.generation());
if current_state_generation != state_generation
|| self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
{
disambiguator.clear_pending().await;
*self.pending_skill_id.write() = None;
result = DisambiguationResult::Abandoned { new_input: None };
info!(
confirmation_event = "invalidated",
invalidation_reason = "state_generation_changed",
"Streaming disambiguation result invalidated before redispatch"
);
}
match result {
DisambiguationResult::Clear => {
debug!("Input is clear, proceeding normally (stream)");
}
DisambiguationResult::NeedsClarification {
question,
detection,
} => {
let admission = match self
.admit_disambiguation_redispatch(
disambiguation_epoch,
state_generation,
)
.await
{
Ok(admission) => admission,
Err(error) => {
*self.pending_skill_id.write() = None;
yield StreamChunk::error(error.to_string());
return;
}
};
let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
info!(
ambiguity_type = ?detection.ambiguity_type,
confidence = detection.confidence,
"Input requires clarification (stream)"
);
if let Err(e) = self.commit_root_user_message(input).await {
yield StreamChunk::error(e.to_string());
return;
}
let _ = self
.memory
.add_message(ChatMessage::assistant(&question.question))
.await;
let status = if awaiting_confirmation {
"awaiting_confirmation"
} else {
"awaiting_clarification"
};
let response = AgentResponse::new(&question.question).with_metadata(
"disambiguation",
serde_json::json!({ "status": status }),
);
drop(admission);
if let Err(e) = self.finish_turn_if_root(&response).await {
yield StreamChunk::error(e.to_string());
return;
}
yield StreamChunk::content(&question.question);
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
return;
}
DisambiguationResult::Clarified {
enriched_input,
resolved,
..
} => {
let admission = match self
.admit_disambiguation_redispatch(
disambiguation_epoch,
state_generation,
)
.await
{
Ok(admission) => admission,
Err(error) => {
*self.pending_skill_id.write() = None;
yield StreamChunk::error(error.to_string());
return;
}
};
info!(
resolved_count = resolved.len(),
enriched = %enriched_input,
"Input clarified (stream)"
);
for (key, value) in &resolved {
let context_key = format!("disambiguation.{}", key);
let _ = self.context_manager.set(&context_key, value.clone());
}
if let Some(intent) = resolved.get("intent") {
let _ = self.context_manager.set("resolved_intent", intent.clone());
}
let _ = self
.context_manager
.set("disambiguation.resolved", serde_json::Value::Bool(true));
let skill_id = self.pending_skill_id.read().clone();
if let Some(skill_id) = skill_id {
info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
drop(admission);
match self
.recheck_skill_disambiguation(
&skill_id,
&enriched_input,
disambiguation_epoch,
state_generation,
)
.await
{
Ok(resp) => {
yield StreamChunk::content(&resp.content);
record_runtime_stream_final(&terminal, resp);
yield StreamChunk::Done {};
return;
}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
drop(admission);
let mut inner = self.run_loop_internal_stream(
&enriched_input,
Arc::clone(&terminal),
);
while let Some(chunk) = inner.next().await {
yield chunk;
}
return;
}
DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
info!("Proceeding with best guess (stream)");
let skill_id = self.pending_skill_id.read().clone();
if let Some(skill_id) = skill_id {
info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
match self
.recheck_skill_disambiguation(
&skill_id,
&enriched_input,
disambiguation_epoch,
state_generation,
)
.await
{
Ok(resp) => {
yield StreamChunk::content(&resp.content);
record_runtime_stream_final(&terminal, resp);
yield StreamChunk::Done {};
return;
}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
let mut inner = self.run_loop_internal_stream(
&enriched_input,
Arc::clone(&terminal),
);
while let Some(chunk) = inner.next().await {
yield chunk;
}
return;
}
DisambiguationResult::GiveUp { reason } => {
*self.pending_skill_id.write() = None;
warn!(reason = %reason, "Disambiguation gave up (stream)");
let apology = self
.generate_localized_apology(
"Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I'm sorry, I couldn't understand your request: {}", reason)
});
let response = AgentResponse::new(&apology);
if let Err(e) = self.finish_turn_if_root(&response).await {
yield StreamChunk::error(e.to_string());
return;
}
yield StreamChunk::content(&apology);
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
return;
}
DisambiguationResult::Escalate { reason } => {
*self.pending_skill_id.write() = None;
info!(reason = %reason, "Escalating to human (stream)");
if let Some(ref hitl) = self.hitl_engine {
let trigger =
ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
let mut context_map = HashMap::new();
context_map.insert("original_input".to_string(), serde_json::json!(input));
context_map.insert("reason".to_string(), serde_json::json!(&reason));
let check_result = HITLCheckResult::required(
trigger,
context_map,
format!("User request needs human assistance: {}", reason),
Some(hitl.config().default_timeout_seconds),
);
match self.request_hitl_approval(check_result).await {
Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
let mut inner = self.run_loop_internal_stream(
input,
Arc::clone(&terminal),
);
while let Some(chunk) = inner.next().await {
yield chunk;
}
return;
}
Ok(_) => {}
Err(e) => {
yield StreamChunk::error(e.to_string());
return;
}
}
}
let apology = self
.generate_localized_apology(
"Explain briefly that you're transferring the user to a human agent for help.",
&reason,
)
.await
.unwrap_or_else(|_| {
format!("I need human assistance to help with your request: {}", reason)
});
let response = AgentResponse::new(&apology);
if let Err(e) = self.finish_turn_if_root(&response).await {
yield StreamChunk::error(e.to_string());
return;
}
yield StreamChunk::content(&apology);
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
return;
}
DisambiguationResult::Abandoned { new_input } => {
*self.pending_skill_id.write() = None;
info!(
has_new_input = new_input.is_some(),
"Clarification abandoned by user (stream)"
);
if let Err(e) = self.commit_root_user_message(input).await {
yield StreamChunk::error(e.to_string());
return;
}
match new_input {
Some(fresh_input) => {
let mut inner = self.run_loop_internal_stream(
&fresh_input,
Arc::clone(&terminal),
);
while let Some(chunk) = inner.next().await {
yield chunk;
}
return;
}
None => {
let ack = self
.generate_localized_apology(
"The user changed their mind about their previous request. \
Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
Do NOT apologize excessively. Be concise.",
"User abandoned clarification",
)
.await
.unwrap_or_else(|_| {
"OK, no problem. What else can I help with?".to_string()
});
let _ = self
.memory
.add_message(ChatMessage::assistant(&ack))
.await;
let response = AgentResponse::new(&ack);
if let Err(e) = self.finish_turn_if_root(&response).await {
yield StreamChunk::error(e.to_string());
return;
}
yield StreamChunk::content(&ack);
record_runtime_stream_final(&terminal, response);
yield StreamChunk::Done {};
return;
}
}
}
}
}
let mut inner = self.run_loop_internal_stream(input, Arc::clone(&terminal));
while let Some(chunk) = inner.next().await {
yield chunk;
}
})
}
pub fn info(&self) -> AgentInfo {
self.info.clone()
}
pub fn skills(&self) -> &[SkillDefinition] {
&self.skills
}
async fn reset_runtime_state(&self) -> Result<()> {
let _admission = self.disambiguation_admission.write().await;
if self.state_transition_reserved.load(Ordering::SeqCst) {
return Err(AgentError::Other(
"Cannot reset while a state transition is in progress".to_string(),
));
}
self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
*self.pending_skill_id.write() = None;
if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
disambiguator.clear_pending().await;
}
self.memory.clear().await?;
*self.iteration_count.write() = 0;
self.tool_call_history.write().clear();
if let Some(ref sm) = self.state_machine {
sm.reset();
}
Ok(())
}
pub async fn reset(&self) -> Result<()> {
self.reset_runtime_state().await
}
pub fn max_context_tokens(&self) -> u32 {
self.max_context_tokens
}
pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
&self.llm_registry
}
pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
self.state_machine.as_ref()
}
pub fn context_manager(&self) -> &Arc<ContextManager> {
&self.context_manager
}
pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
self.tool_call_history.read().clone()
}
pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
self.memory_token_budget.as_ref()
}
pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
&self.parallel_tools
}
pub fn streaming_config(&self) -> &StreamingConfig {
&self.streaming
}
pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
&self.hooks
}
pub fn hitl_engine(&self) -> Option<&HITLEngine> {
self.hitl_engine.as_ref()
}
pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
&self.approval_handler
}
fn build_hitl_language_context(&self) -> HashMap<String, Value> {
let mut ctx = HashMap::new();
for key in &["user.language", "input.detected.language", "language"] {
if let Some(val) = self.context_manager.get(key) {
ctx.insert(key.to_string(), val);
}
}
ctx
}
async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
let Some(request) = check_result.into_request() else {
return Ok(ApprovalResult::Approved);
};
self.hooks.on_approval_requested(&request).await;
let timeout = request.timeout;
let raw_result = if let Some(duration) = timeout {
match tokio::time::timeout(
duration,
self.approval_handler.request_approval(request.clone()),
)
.await
{
Ok(result) => result,
Err(_) => ApprovalResult::timeout(),
}
} else {
self.approval_handler
.request_approval(request.clone())
.await
};
self.hooks
.on_approval_result(&request.id, &raw_result)
.await;
let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
match &raw_result {
ApprovalResult::Approved => (
ApprovalResolvedOutcome::Approved,
Ok(ApprovalResult::Approved),
),
ApprovalResult::Rejected { reason } => (
ApprovalResolvedOutcome::Rejected {
reason: reason.clone(),
},
Ok(ApprovalResult::Rejected {
reason: reason.clone(),
}),
),
ApprovalResult::Modified { changes } => (
ApprovalResolvedOutcome::Modified {
changes: changes.clone(),
},
Ok(ApprovalResult::Modified {
changes: changes.clone(),
}),
),
ApprovalResult::Timeout => {
if let Some(ref engine) = self.hitl_engine {
match engine.config().on_timeout {
TimeoutAction::Approve => (
ApprovalResolvedOutcome::Approved,
Ok(ApprovalResult::Approved),
),
TimeoutAction::Reject => {
let reason = Some("Timeout".to_string());
(
ApprovalResolvedOutcome::Rejected {
reason: reason.clone(),
},
Ok(ApprovalResult::Rejected { reason }),
)
}
TimeoutAction::Error => {
let message = "HITL approval timeout".to_string();
(
ApprovalResolvedOutcome::Error {
message: message.clone(),
},
Err(AgentError::Other(message)),
)
}
}
} else {
let reason = Some("Timeout (no engine)".to_string());
(
ApprovalResolvedOutcome::Rejected {
reason: reason.clone(),
},
Ok(ApprovalResult::Rejected { reason }),
)
}
}
};
self.hooks
.on_approval_resolved(&request, &raw_result, &outcome)
.await;
effective_result
}
pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
if let Some(ref hitl_engine) = self.hitl_engine {
let hitl_lang_ctx = self.build_hitl_language_context();
let check_result = self
.observe_purpose(
ObservationPurpose::HitlLocalization,
hitl_engine.check_state_transition_with_localization(
from,
to,
&hitl_lang_ctx,
self.approval_handler.as_ref(),
Some(&self.llm_registry),
),
)
.await?;
if check_result.is_required() {
let result = self.request_hitl_approval(check_result).await?;
return Ok(matches!(
result,
ApprovalResult::Approved | ApprovalResult::Modified { .. }
));
}
}
Ok(true)
}
async fn execute_tools_parallel(
&self,
tool_calls: &[ToolCall],
) -> Vec<(String, Result<String>)> {
let can_run_parallel = tool_calls.iter().all(|tc| {
self.tools
.resolve(&tc.name)
.map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
.unwrap_or(false)
});
if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
let mut results = Vec::new();
for tc in tool_calls {
let result = self
.observe_purpose(
current_observation_context()
.map(|context| context.purpose)
.unwrap_or_default(),
self.execute_tool_smart(tc),
)
.await;
results.push((tc.id.clone(), result));
}
return results;
}
let chunks: Vec<_> = tool_calls
.chunks(self.parallel_tools.max_parallel)
.collect();
let mut all_results = Vec::new();
for chunk in chunks {
let futures: Vec<_> = chunk
.iter()
.map(|tc| {
let tc = tc.clone();
async move {
let result = self.execute_tool_smart(&tc).await;
(tc.id.clone(), result)
}
})
.collect();
let results = futures::future::join_all(futures).await;
all_results.extend(results);
}
all_results
}
pub async fn chat_stream<'a>(
&'a self,
input: &'a str,
) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
let RootTurnAdmission {
guard: root_turn_guard,
identity_stack,
} = self.acquire_root_turn().await?;
scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
info!(input_len = input.len(), "Starting streaming chat");
let terminal = new_runtime_stream_terminal_slot();
let inner = self.run_loop_stream(input, terminal);
let observation_context = self.build_observation_context(None);
let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
Box::pin(async_stream::stream! {
let mut root_turn_guard = Some(root_turn_guard);
let mut inner = inner;
loop {
let next = scope_runtime_gate_identity_stack(&identity_stack, async {
if let Some(context) = observation_context.as_ref() {
with_observation_context(context.clone(), inner.next()).await
} else {
inner.next().await
}
})
.await;
match next {
Some(StreamChunk::Done {}) => {
while scope_runtime_gate_identity_stack(&identity_stack, async {
if let Some(context) = observation_context.as_ref() {
with_observation_context(context.clone(), inner.next())
.await
.is_some()
} else {
inner.next().await.is_some()
}
})
.await
{}
if observation_context.is_some() {
scope_runtime_gate_identity_stack(
&identity_stack,
self.export_observability_if_configured(),
)
.await;
}
drop(root_turn_guard.take());
yield StreamChunk::Done {};
return;
}
Some(chunk) => yield chunk,
None => {
if observation_context.is_some() {
scope_runtime_gate_identity_stack(
&identity_stack,
self.export_observability_if_configured(),
)
.await;
}
drop(root_turn_guard.take());
return;
}
}
}
});
Ok(stream)
}
pub async fn chat_stream_events<'a>(
&'a self,
input: &'a str,
) -> Result<Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>>> {
let RootTurnAdmission {
guard: root_turn_guard,
identity_stack,
} = self.acquire_root_turn().await?;
scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
info!(input_len = input.len(), "Starting streaming chat events");
let terminal = new_runtime_stream_terminal_slot();
let mut inner = self.run_loop_stream(input, Arc::clone(&terminal));
let observation_context = self.build_observation_context(None);
let stream: Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>> =
Box::pin(async_stream::stream! {
let mut root_turn_guard = Some(root_turn_guard);
loop {
let next = scope_runtime_gate_identity_stack(&identity_stack, async {
if let Some(context) = observation_context.as_ref() {
with_observation_context(context.clone(), inner.next()).await
} else {
inner.next().await
}
})
.await;
match next {
Some(StreamChunk::Done {}) => {
let terminal_event = { terminal.write().take() };
if let Some(response) = terminal_event {
while scope_runtime_gate_identity_stack(&identity_stack, async {
if let Some(context) = observation_context.as_ref() {
with_observation_context(context.clone(), inner.next())
.await
.is_some()
} else {
inner.next().await.is_some()
}
})
.await
{}
if observation_context.is_some() {
scope_runtime_gate_identity_stack(
&identity_stack,
self.export_observability_if_configured(),
)
.await;
}
drop(root_turn_guard.take());
yield AgentStreamEvent::Final(response);
return;
}
}
Some(StreamChunk::Error { message }) => {
let finalized = { terminal.read().is_some() };
if finalized {
continue;
}
while scope_runtime_gate_identity_stack(&identity_stack, async {
if let Some(context) = observation_context.as_ref() {
with_observation_context(context.clone(), inner.next())
.await
.is_some()
} else {
inner.next().await.is_some()
}
})
.await
{}
if observation_context.is_some() {
scope_runtime_gate_identity_stack(
&identity_stack,
self.export_observability_if_configured(),
)
.await;
}
drop(root_turn_guard.take());
yield AgentStreamEvent::Chunk(StreamChunk::Error { message });
return;
}
Some(chunk) => yield AgentStreamEvent::Chunk(chunk),
None => {
if observation_context.is_some() {
scope_runtime_gate_identity_stack(
&identity_stack,
self.export_observability_if_configured(),
)
.await;
}
drop(root_turn_guard.take());
return;
}
}
}
});
Ok(stream)
}
}
#[async_trait]
impl ToolInvoker for RuntimeAgent {
async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
self.execute_tool_record(request).await
}
}
#[async_trait]
impl Agent for RuntimeAgent {
async fn chat(&self, input: &str) -> Result<AgentResponse> {
let RootTurnAdmission {
guard,
identity_stack,
} = self.acquire_root_turn().await?;
let result = scope_runtime_gate_identity_stack(&identity_stack, async {
let result = if let Some(context) = self.build_observation_context(None) {
with_observation_context(context, self.run_loop(input)).await
} else {
self.run_loop(input).await
};
self.export_observability_if_configured().await;
result
})
.await;
drop(guard);
result
}
fn info(&self) -> AgentInfo {
self.info.clone()
}
async fn reset(&self) -> Result<()> {
self.reset_runtime_state().await
}
}
fn background_maintenance_tags(
label: &str,
stage: &str,
reason: Option<&str>,
policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
) -> HashMap<String, String> {
let mut tags = HashMap::new();
tags.insert("runtime.background".to_string(), "true".to_string());
tags.insert("runtime.maintenance".to_string(), label.to_string());
tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
if let Some(policy) = policy {
tags.insert(
"runtime.await_before_next_turn".to_string(),
await_before_next_turn_label(policy.await_before_next_turn).to_string(),
);
tags.insert(
"runtime.maintenance_mode".to_string(),
maintenance_mode_label(policy.mode).to_string(),
);
}
if let Some(reason) = reason {
tags.insert("runtime.reason".to_string(), reason.to_string());
}
tags
}
fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
match policy {
AwaitBeforeNextTurn::Never => "never",
AwaitBeforeNextTurn::SameActor => "same_actor",
AwaitBeforeNextTurn::Always => "always",
}
}
fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
match mode {
MaintenanceMode::InlineSerial => "inline_serial",
MaintenanceMode::InlineParallel => "inline_parallel",
MaintenanceMode::Background => "background",
}
}
fn record_background_maintenance_event(
manager: Option<&Arc<ObservabilityManager>>,
label: &str,
status: EventStatus,
duration_ms: u64,
stage: &str,
reason: Option<String>,
policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
) {
if let Some(manager) = manager {
manager.record_lifecycle_event(
EventType::MemoryOperation {
operation: format!("{}_background_{}", label, stage),
},
ObservationPurpose::Other(format!("{}_maintenance", label)),
status,
duration_ms,
background_maintenance_tags(label, stage, reason.as_deref(), policy),
None,
);
}
}
fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
MaintenanceMode::InlineParallel
} else {
mode
}
}
fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
match hint {
ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
ObservationPurpose::ProcessTransform
}
}
}
fn new_tool_resource_locks() -> ToolResourceLocks {
Arc::new(RwLock::new(HashMap::new()))
}
fn tool_resource_lock_keys(
_canonical_id: &str,
args: &Value,
bindings: &ai_agents_core::ToolPolicyBindings,
classification: &ai_agents_core::ToolCallClassification,
) -> Vec<String> {
if classification.concurrency_safe {
return Vec::new();
}
let mut keys = Vec::new();
let mut has_path_resource = false;
for binding in &bindings.path_fields {
let value = value_at_argument_path(args, &binding.field)
.cloned()
.or_else(|| {
binding
.default_path
.as_ref()
.map(|path| Value::String(path.clone()))
});
if let Some(value) = value {
collect_resource_strings(&value, |_| {
has_path_resource = true;
});
}
}
for binding in &bindings.domain_fields {
if let Some(value) = value_at_argument_path(args, &binding.field) {
collect_resource_strings(value, |domain| {
let normalized = if binding.is_url {
normalized_url_resource_key(domain)
} else {
domain.trim().trim_end_matches('.').to_ascii_lowercase()
};
keys.push(format!("domain:{}", normalized));
});
}
}
for binding in &bindings.command_fields {
if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
continue;
}
if let Some(value) = value_at_argument_path(args, &binding.field) {
collect_resource_strings(value, |_| {
has_path_resource = true;
});
}
}
if has_path_resource {
keys.push("path-mutation:global".to_string());
}
if keys.is_empty() {
keys.push("side-effect:unbound".to_string());
}
keys.sort();
keys.dedup();
keys
}
fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
let mut current = value;
for segment in field.split('.') {
if segment.is_empty() {
return None;
}
current = current.get(segment)?;
}
Some(current)
}
fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
match value {
Value::String(value) => collect(value),
Value::Array(values) => {
for value in values {
if let Some(value) = value.as_str() {
collect(value);
}
}
}
_ => {}
}
}
fn normalized_url_resource_key(value: &str) -> String {
let value = value.trim();
let Some((scheme, remainder)) = value.split_once("://") else {
return value.to_ascii_lowercase();
};
let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
let (authority, suffix) = remainder.split_at(authority_end);
format!(
"{}://{}{}",
scheme.to_ascii_lowercase(),
authority.to_ascii_lowercase(),
suffix
)
}
fn render_concurrent_template(
template: &str,
user_input: &str,
context_values: &std::collections::HashMap<String, serde_json::Value>,
) -> Result<String> {
let mut env = minijinja::Environment::new();
env.add_template("concurrent", template)
.map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
let mut ctx = std::collections::BTreeMap::new();
ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
let context_obj = minijinja::Value::from_serialize(context_values);
ctx.insert("context".to_string(), context_obj);
let tmpl = env
.get_template("concurrent")
.map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
tmpl.render(minijinja::Value::from_serialize(&ctx))
.map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::AgentBuilder;
use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
use ai_agents_llm::mock::MockLLMProvider;
use ai_agents_skills::{SkillDefinition, SkillStep};
use ai_agents_tools::{
CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool, ToolAliases,
ToolDescriptor, ToolProvider, ToolProviderError, ToolProviderType, WebFetchResolver,
WebFetchTool, WebFetchTransport, WebFetchTransportRequest, WebFetchTransportResponse,
};
fn mock_with_response(response: &str) -> MockLLMProvider {
let mut mock = MockLLMProvider::new("test");
mock.set_response(response);
mock
}
fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
let mut mock = MockLLMProvider::new("test");
mock.set_responses(responses.into_iter().map(String::from).collect(), true);
mock
}
fn disambiguation_state_machine(
state_enabled: Option<bool>,
require_confirmation: bool,
) -> Arc<StateMachine> {
let definition = ai_agents_state::StateDefinition {
prompt: Some("Handle the resolved request.".to_string()),
disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
enabled: state_enabled,
require_confirmation,
..Default::default()
}),
..Default::default()
};
let review = ai_agents_state::StateDefinition {
prompt: Some("Review a fresh request.".to_string()),
..Default::default()
};
Arc::new(
StateMachine::new(ai_agents_state::StateConfig {
initial: "active".to_string(),
states: std::collections::HashMap::from([
("active".to_string(), definition),
("review".to_string(), review),
]),
global_transitions: Vec::new(),
fallback: None,
max_no_transition: None,
regenerate_on_transition: true,
})
.unwrap(),
)
}
fn state_disambiguation_agent(
responses: Vec<&str>,
manager_enabled: bool,
state_enabled: Option<bool>,
require_confirmation: bool,
) -> (RuntimeAgent, MockLLMProvider) {
state_disambiguation_agent_with_skills(
responses,
manager_enabled,
state_enabled,
require_confirmation,
Vec::new(),
)
}
fn state_disambiguation_agent_with_skills(
responses: Vec<&str>,
manager_enabled: bool,
state_enabled: Option<bool>,
require_confirmation: bool,
skills: Vec<SkillDefinition>,
) -> (RuntimeAgent, MockLLMProvider) {
let mut mock = MockLLMProvider::new("state-confirmation");
mock.set_responses(responses.into_iter().map(String::from).collect(), false);
let observed = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("Handle requests.")
.llm(Arc::new(mock.clone()))
.llm_alias("router", Arc::new(mock))
.state_machine(disambiguation_state_machine(
state_enabled,
require_confirmation,
))
.skills(skills)
.build()
.unwrap()
.with_disambiguation(DisambiguationConfig {
enabled: manager_enabled,
..Default::default()
});
(agent, observed)
}
fn confirmation_skill() -> SkillDefinition {
SkillDefinition {
id: "send_report".to_string(),
description: "Send a report after clarification".to_string(),
trigger: "When the user asks to send a report".to_string(),
steps: vec![SkillStep::Prompt {
prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
llm: None,
}],
reasoning: None,
reflection: None,
disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
enabled: Some(true),
..Default::default()
}),
}
}
fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
observed
.call_history()
.iter()
.filter(|call| {
call.messages
.iter()
.any(|message| message.content.contains("Execute confirmed report skill"))
})
.count()
}
struct BlockingRuntimeConfirmationObserver {
entered: tokio::sync::Barrier,
release: tokio::sync::Notify,
}
impl BlockingRuntimeConfirmationObserver {
fn new() -> Self {
Self {
entered: tokio::sync::Barrier::new(2),
release: tokio::sync::Notify::new(),
}
}
}
struct ResetOnTransitionHooks {
agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
invoked: AtomicBool,
}
#[async_trait]
impl AgentHooks for ResetOnTransitionHooks {
async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
if self.invoked.swap(true, Ordering::SeqCst) {
return;
}
let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
if let Some(agent) = agent {
agent.reset().await.unwrap();
}
}
}
impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
fn observe_question<'a>(
&'a self,
future: ClarificationQuestionFuture<'a>,
) -> ClarificationQuestionFuture<'a> {
future
}
fn observe_parse<'a>(
&'a self,
future: ClarificationParseFuture<'a>,
) -> ClarificationParseFuture<'a> {
future
}
fn observe_confirmation_parse<'a>(
&'a self,
future: ConfirmationParseFuture<'a>,
) -> ConfirmationParseFuture<'a> {
Box::pin(async move {
self.entered.wait().await;
self.release.notified().await;
future.await
})
}
}
#[tokio::test]
async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
let (agent, observed) = state_disambiguation_agent(
vec![
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"status":"confirmed"}"#,
"Request executed.",
],
true,
None,
true,
);
let clarification = agent.chat("Send it").await.unwrap();
assert_eq!(clarification.content, "What should I send?");
assert_eq!(observed.call_count(), 2);
let confirmation = agent.chat("The report to Ada").await.unwrap();
assert_eq!(confirmation.content, "Should I send the report to Ada?");
assert_eq!(
confirmation
.metadata
.as_ref()
.and_then(|metadata| metadata.get("disambiguation"))
.and_then(|metadata| metadata.get("status"))
.and_then(Value::as_str),
Some("awaiting_confirmation")
);
assert_eq!(observed.call_count(), 4);
let completed = agent.chat("Yes").await.unwrap();
assert_eq!(completed.content, "Request executed.");
assert_eq!(observed.call_count(), 6);
}
#[tokio::test]
async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
let (agent, observed) = state_disambiguation_agent(
vec![
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"status":"confirmed"}"#,
"Request executed.",
],
true,
None,
true,
);
let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
let mut clarification = String::new();
while let Some(chunk) = clarification_stream.next().await {
match chunk {
StreamChunk::Content { text } => clarification.push_str(&text),
StreamChunk::Done {} => break,
StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
_ => {}
}
}
assert_eq!(clarification, "What should I send?");
assert_eq!(observed.call_count(), 2);
let mut confirmation_stream = agent.chat_stream_events("The report to Ada").await.unwrap();
let mut confirmation = None;
while let Some(event) = confirmation_stream.next().await {
match event {
AgentStreamEvent::Final(response) => confirmation = Some(response),
AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
panic!("unexpected stream error: {message}")
}
AgentStreamEvent::Chunk(_) => {}
}
}
let confirmation = confirmation.expect("confirmation must finalize");
assert_eq!(confirmation.content, "Should I send the report to Ada?");
assert_eq!(
confirmation
.metadata
.as_ref()
.and_then(|metadata| metadata.get("disambiguation"))
.and_then(|metadata| metadata.get("status"))
.and_then(Value::as_str),
Some("awaiting_confirmation")
);
assert_eq!(observed.call_count(), 4);
let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
let mut completed = String::new();
while let Some(chunk) = completed_stream.next().await {
match chunk {
StreamChunk::Content { text } => completed.push_str(&text),
StreamChunk::Done {} => break,
StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
_ => {}
}
}
assert_eq!(completed, "Request executed.");
assert_eq!(observed.call_count(), 6);
}
#[tokio::test]
async fn root_turn_gate_serializes_blocking_and_streaming_entry_points() {
let (complete_entered, mut complete_events) = tokio::sync::mpsc::unbounded_channel();
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Serialize root turns.")
.llm(Arc::new(RootTurnProbeProvider { complete_entered }))
.build()
.unwrap(),
);
let blocking_agent = Arc::clone(&agent);
let legacy_stream = agent.chat_stream("stream owner").await.unwrap();
assert!(agent.root_turn_gate.try_lock().is_err());
let blocking = tokio::spawn(async move { blocking_agent.chat("blocked").await.unwrap() });
assert!(
tokio::time::timeout(std::time::Duration::from_millis(50), complete_events.recv())
.await
.is_err(),
"blocking turn reached the provider while the legacy stream owned the root gate"
);
drop(legacy_stream);
assert_eq!(
tokio::time::timeout(std::time::Duration::from_secs(2), complete_events.recv())
.await
.expect("blocking turn did not enter after stream drop"),
Some(())
);
let response = tokio::time::timeout(std::time::Duration::from_secs(2), blocking)
.await
.expect("blocking turn did not finish after stream drop")
.unwrap();
assert_eq!(response.content, "blocking complete");
let mut event_stream = agent.chat_stream_events("event terminal").await.unwrap();
assert!(agent.root_turn_gate.try_lock().is_err());
let mut saw_final = false;
while let Some(event) = event_stream.next().await {
if matches!(event, AgentStreamEvent::Final(_)) {
saw_final = true;
break;
}
}
assert!(saw_final);
assert!(
agent.root_turn_gate.try_lock().is_ok(),
"authoritative terminal event retained the root gate"
);
}
#[tokio::test]
async fn response_hook_rejects_same_runtime_chat_reentry() {
let hooks = Arc::new(ResponseChatHooks {
target: parking_lot::Mutex::new(None),
invoked: AtomicBool::new(false),
nested_result: parking_lot::Mutex::new(None),
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Reject response hook reentry.")
.llm(Arc::new(mock_with_response("outer response")))
.hooks(hooks.clone())
.build()
.unwrap(),
);
*hooks.target.lock() = Some(Arc::downgrade(&agent));
let response = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.chat("outer request"),
)
.await
.expect("same-runtime response hook reentry must fail without deadlocking")
.unwrap();
assert_eq!(response.content, "outer response");
let nested_result = hooks
.nested_result
.lock()
.clone()
.expect("response hook must record its nested call");
let error = nested_result.expect_err("same-runtime nested chat must be rejected");
assert!(error.contains("reentrant root turn ownership"));
}
#[tokio::test]
async fn root_turn_gate_allows_nested_runtime_and_rejects_cycles() {
let agent_a = AgentBuilder::new()
.system_prompt("Runtime A.")
.llm(Arc::new(mock_with_response("response A")))
.build()
.unwrap();
let agent_b = AgentBuilder::new()
.system_prompt("Runtime B.")
.llm(Arc::new(mock_with_response("response B")))
.build()
.unwrap();
let RootTurnAdmission {
guard: guard_a,
identity_stack: stack_a,
} = agent_a.acquire_root_turn().await.unwrap();
let cycle_error = scope_runtime_gate_identity_stack(&stack_a, async {
let RootTurnAdmission {
guard: guard_b,
identity_stack: stack_b,
} = agent_b
.acquire_root_turn()
.await
.expect("runtime B must acquire a different gate");
let result =
scope_runtime_gate_identity_stack(&stack_b, agent_a.acquire_root_turn()).await;
drop(guard_b);
match result {
Err(error) => error,
Ok(_) => panic!("runtime A accepted a repeated gate identity"),
}
})
.await;
drop(guard_a);
assert!(
cycle_error
.to_string()
.contains("reentrant root turn ownership")
);
}
#[tokio::test]
async fn concurrent_orchestration_propagates_root_gate_ancestry() {
let registry = Arc::new(crate::spawner::AgentRegistry::new());
let hooks_a = Arc::new(ConcurrentResponseHooks {
registry: Arc::downgrade(®istry),
child_id: "runtime-b".to_string(),
invoked: AtomicBool::new(false),
nested_result: parking_lot::Mutex::new(None),
});
let hooks_b = Arc::new(ResponseChatHooks {
target: parking_lot::Mutex::new(None),
invoked: AtomicBool::new(false),
nested_result: parking_lot::Mutex::new(None),
});
let agent_a = AgentBuilder::new()
.system_prompt("Runtime A dispatches runtime B concurrently.")
.llm(Arc::new(mock_with_response("response A")))
.hooks(hooks_a.clone())
.build()
.unwrap();
let agent_b = AgentBuilder::new()
.system_prompt("Runtime B attempts to re-enter runtime A.")
.llm(Arc::new(mock_with_response("response B")))
.hooks(hooks_b.clone())
.build()
.unwrap();
let spec_a = crate::spec::AgentSpec {
name: "runtime-a".to_string(),
system_prompt: "Runtime A dispatches runtime B concurrently.".to_string(),
..crate::spec::AgentSpec::default()
};
let spec_b = crate::spec::AgentSpec {
name: "runtime-b".to_string(),
system_prompt: "Runtime B attempts to re-enter runtime A.".to_string(),
..crate::spec::AgentSpec::default()
};
registry
.register(crate::spawner::SpawnedAgent::from_runtime(
"runtime-a".to_string(),
agent_a,
spec_a,
))
.await
.unwrap();
registry
.register(crate::spawner::SpawnedAgent::from_runtime(
"runtime-b".to_string(),
agent_b,
spec_b,
))
.await
.unwrap();
let runtime_a = registry.get("runtime-a").unwrap();
*hooks_b.target.lock() = Some(Arc::downgrade(&runtime_a));
let response = tokio::time::timeout(
std::time::Duration::from_secs(2),
runtime_a.chat("outer concurrent request"),
)
.await
.expect("concurrent orchestration cycle must fail without deadlocking")
.unwrap();
assert_eq!(response.content, "response A");
let child_result = hooks_a
.nested_result
.lock()
.clone()
.expect("runtime A hook must record runtime B completion");
assert_eq!(child_result.unwrap(), "response B");
let cycle_result = hooks_b
.nested_result
.lock()
.clone()
.expect("runtime B hook must record runtime A reentry");
assert!(
cycle_result
.expect_err("runtime A accepted a repeated gate identity")
.contains("reentrant root turn ownership")
);
}
#[tokio::test]
async fn confirmed_skill_route_executes_exactly_once() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"status":"confirmed"}"#,
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
"Report skill executed.",
],
true,
None,
true,
vec![confirmation_skill()],
);
let clarification = agent.chat("Send it").await.unwrap();
assert_eq!(clarification.content, "What should I send?");
assert_eq!(confirmation_skill_call_count(&observed), 0);
let confirmation = agent.chat("The report to Ada").await.unwrap();
assert_eq!(confirmation.content, "Should I send the report to Ada?");
assert_eq!(
confirmation
.metadata
.as_ref()
.and_then(|metadata| metadata.get("disambiguation"))
.and_then(|metadata| metadata.get("status"))
.and_then(Value::as_str),
Some("awaiting_confirmation")
);
assert_eq!(confirmation_skill_call_count(&observed), 0);
let completed = agent.chat("Yes").await.unwrap();
assert_eq!(completed.content, "Report skill executed.");
assert_eq!(confirmation_skill_call_count(&observed), 1);
assert!(agent.pending_skill_id.read().is_none());
let messages = agent.memory.get_messages(None).await.unwrap();
assert!(!messages.iter().any(|message| message.content == "Yes"));
}
#[tokio::test]
async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"status":"confirmed"}"#,
r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
r#"{"question":"When should I send it?","options":null}"#,
],
true,
None,
true,
vec![confirmation_skill()],
);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
let follow_up = agent.chat("Yes").await.unwrap();
assert_eq!(follow_up.content, "When should I send it?");
let metadata = follow_up
.metadata
.as_ref()
.and_then(|metadata| metadata.get("disambiguation"))
.unwrap();
assert_eq!(
metadata.get("status").and_then(Value::as_str),
Some("awaiting_clarification")
);
assert_eq!(
metadata.get("skill_id").and_then(Value::as_str),
Some("send_report")
);
assert!(metadata.get("detection").is_some());
assert_eq!(confirmation_skill_call_count(&observed), 0);
}
#[tokio::test]
async fn rejected_skill_confirmation_never_executes() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"status":"rejected"}"#,
"Confirmation rejected.",
],
true,
None,
true,
vec![confirmation_skill()],
);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
let rejected = agent.chat("No").await.unwrap();
assert_eq!(rejected.content, "Confirmation rejected.");
assert_eq!(confirmation_skill_call_count(&observed), 0);
assert!(agent.pending_skill_id.read().is_none());
}
#[tokio::test]
async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
"none",
"Fresh response.",
],
true,
None,
true,
vec![confirmation_skill()],
);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
agent.reset().await.unwrap();
assert!(agent.pending_skill_id.read().is_none());
assert!(
!agent
.disambiguation_manager()
.unwrap()
.has_pending_clarification()
.await
);
let mut stream = agent.chat_stream("Yes").await.unwrap();
let mut content = String::new();
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Content { text } => content.push_str(&text),
StreamChunk::Done {} => break,
StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
_ => {}
}
}
assert_eq!(content, "Fresh response.");
assert_eq!(confirmation_skill_call_count(&observed), 0);
}
#[tokio::test]
async fn trait_reset_clears_pending_skill_confirmation() {
let (agent, _) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
],
true,
None,
true,
vec![confirmation_skill()],
);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
<RuntimeAgent as Agent>::reset(&agent).await.unwrap();
assert!(agent.pending_skill_id.read().is_none());
assert!(
!agent
.disambiguation_manager()
.unwrap()
.has_pending_clarification()
.await
);
}
#[tokio::test]
async fn state_change_invalidates_pending_skill_confirmation() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
"none",
"Fresh response.",
],
true,
None,
true,
vec![confirmation_skill()],
);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
agent.transition_to("review").await.unwrap();
let cancelled = agent.chat("Yes").await.unwrap();
assert_eq!(cancelled.content, "Fresh response.");
assert_eq!(confirmation_skill_call_count(&observed), 0);
assert!(agent.pending_skill_id.read().is_none());
}
#[tokio::test]
async fn in_flight_confirmation_cannot_redispatch_after_reset() {
let (mut agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
r#"{"status":"confirmed"}"#,
"Confirmation cancelled.",
],
true,
None,
true,
vec![confirmation_skill()],
);
let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
let manager = agent
.disambiguation_manager
.take()
.unwrap()
.with_clarification_observer(observer.clone());
agent.disambiguation_manager = Some(manager);
let agent = Arc::new(agent);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
let confirming_agent = Arc::clone(&agent);
let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
observer.entered.wait().await;
agent.reset().await.unwrap();
observer.release.notify_one();
let response = confirmation.await.unwrap().unwrap();
assert_eq!(response.content, "Confirmation cancelled.");
assert_eq!(confirmation_skill_call_count(&observed), 0);
assert!(agent.pending_skill_id.read().is_none());
}
#[tokio::test]
async fn queued_reset_prevents_stale_confirmation_question_publication() {
let (agent, observed) = state_disambiguation_agent(
vec![
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
],
true,
None,
true,
);
let agent = Arc::new(agent);
agent.chat("Send it").await.unwrap();
let admission = agent.disambiguation_admission.write().await;
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let resetting_agent = Arc::clone(&agent);
let reset = tokio::spawn(async move {
let _ = started_tx.send(());
resetting_agent.reset().await
});
started_rx.await.unwrap();
tokio::task::yield_now().await;
let responding_agent = Arc::clone(&agent);
let response =
tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
tokio::time::timeout(std::time::Duration::from_secs(2), async {
while observed.call_count() < 4 {
tokio::task::yield_now().await;
}
})
.await
.expect("clarification processing must reach terminal publication");
drop(admission);
reset.await.unwrap().unwrap();
let error = response.await.unwrap().unwrap_err();
assert!(error.to_string().contains("ownership changed"));
assert!(
!agent
.disambiguation_manager()
.unwrap()
.has_pending_clarification()
.await
);
assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
}
#[tokio::test]
async fn queued_reset_prevents_stale_skill_clarification_publication() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
],
true,
None,
true,
vec![confirmation_skill()],
);
let agent = Arc::new(agent);
let admission = agent.disambiguation_admission.write().await;
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let resetting_agent = Arc::clone(&agent);
let reset = tokio::spawn(async move {
let _ = started_tx.send(());
resetting_agent.reset().await
});
started_rx.await.unwrap();
tokio::task::yield_now().await;
let responding_agent = Arc::clone(&agent);
let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
tokio::time::timeout(std::time::Duration::from_secs(2), async {
while observed.call_count() < 4 {
tokio::task::yield_now().await;
}
})
.await
.expect("skill clarification must reach terminal publication");
drop(admission);
reset.await.unwrap().unwrap();
let error = response.await.unwrap().unwrap_err();
assert!(error.to_string().contains("ownership changed"));
assert_eq!(confirmation_skill_call_count(&observed), 0);
assert!(agent.pending_skill_id.read().is_none());
assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
}
#[tokio::test]
async fn transition_hook_can_reset_without_admission_deadlock() {
let hooks = Arc::new(ResetOnTransitionHooks {
agent: parking_lot::Mutex::new(None),
invoked: AtomicBool::new(false),
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test transition hook reentrancy.")
.llm(Arc::new(mock_with_response("done")))
.state_machine(disambiguation_state_machine(None, false))
.build()
.unwrap()
.with_hooks(hooks.clone()),
);
*hooks.agent.lock() = Some(Arc::downgrade(&agent));
let transitioned = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.apply_transition_target("active", "review", "test transition", None),
)
.await
.expect("transition hook reset must not deadlock")
.unwrap();
assert!(transitioned);
assert!(hooks.invoked.load(Ordering::SeqCst));
assert_eq!(agent.current_state().as_deref(), Some("active"));
}
#[tokio::test]
async fn concurrent_transition_cannot_duplicate_exit_actions() {
let gate = PathMutationGate::new();
let active = ai_agents_state::StateDefinition {
on_exit: vec![StateAction::Tool {
tool: "transition_exit".to_string(),
args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
}],
..Default::default()
};
let state_machine = Arc::new(
StateMachine::new(ai_agents_state::StateConfig {
initial: "active".to_string(),
states: HashMap::from([
("active".to_string(), active),
(
"review".to_string(),
ai_agents_state::StateDefinition::default(),
),
]),
global_transitions: Vec::new(),
fallback: None,
max_no_transition: None,
regenerate_on_transition: true,
})
.unwrap(),
);
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test transition reservation.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(BlockingPathMutationTool {
id: "transition_exit",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: gate.clone(),
}))
.state_machine(state_machine)
.build()
.unwrap(),
);
let first_agent = Arc::clone(&agent);
let first = tokio::spawn(async move { first_agent.transition_to("review").await });
tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
.await
.expect("reserved transition must enter its exit action");
let second = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.transition_to("review"),
)
.await
.expect("competing transition must fail without waiting for the exit action")
.unwrap_err();
assert!(second.to_string().contains("already in progress"));
gate.release();
first.await.unwrap().unwrap();
assert_eq!(agent.current_state().as_deref(), Some("review"));
}
#[tokio::test]
async fn concurrent_transition_cannot_overtake_enter_actions() {
let gate = PathMutationGate::new();
let review = ai_agents_state::StateDefinition {
on_enter: vec![StateAction::Tool {
tool: "transition_enter".to_string(),
args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
}],
..Default::default()
};
let state_machine = Arc::new(
StateMachine::new(ai_agents_state::StateConfig {
initial: "active".to_string(),
states: HashMap::from([
(
"active".to_string(),
ai_agents_state::StateDefinition::default(),
),
("review".to_string(), review),
]),
global_transitions: Vec::new(),
fallback: None,
max_no_transition: None,
regenerate_on_transition: true,
})
.unwrap(),
);
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test transition lifecycle reservation.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(BlockingPathMutationTool {
id: "transition_enter",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: gate.clone(),
}))
.state_machine(state_machine)
.build()
.unwrap(),
);
let first_agent = Arc::clone(&agent);
let first = tokio::spawn(async move { first_agent.transition_to("review").await });
tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
.await
.expect("committed transition must enter its destination action");
let second = agent.transition_to("active").await.unwrap_err();
assert!(second.to_string().contains("already in progress"));
assert!(agent.reset().await.is_err());
gate.release();
first.await.unwrap().unwrap();
assert_eq!(agent.current_state().as_deref(), Some("review"));
}
#[tokio::test]
async fn same_state_restore_invalidates_pending_skill_confirmation() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
],
true,
None,
true,
vec![confirmation_skill()],
);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
let snapshot = agent.save_state().await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("active"));
agent.restore_state(snapshot).await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("active"));
assert!(agent.pending_skill_id.read().is_none());
assert!(
!agent
.disambiguation_manager()
.unwrap()
.has_pending_clarification()
.await
);
assert_eq!(confirmation_skill_call_count(&observed), 0);
}
#[tokio::test]
async fn direct_state_generation_change_invalidates_confirmation() {
let (agent, observed) = state_disambiguation_agent_with_skills(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
"send_report",
r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
r#"{"question":"What should I send?","options":null}"#,
r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
r#"{"question":"Should I send the report to Ada?"}"#,
"Confirmation cancelled.",
],
true,
None,
true,
vec![confirmation_skill()],
);
agent.chat("Send it").await.unwrap();
agent.chat("The report to Ada").await.unwrap();
let state_machine = agent.state_machine().unwrap();
state_machine
.transition_to("review", "external test")
.unwrap();
state_machine
.transition_to("active", "external test")
.unwrap();
let response = agent.chat("Yes").await.unwrap();
assert_eq!(response.content, "Confirmation cancelled.");
assert_eq!(confirmation_skill_call_count(&observed), 0);
assert!(agent.pending_skill_id.read().is_none());
}
#[tokio::test]
async fn state_confirmation_does_not_add_a_question_for_clear_input() {
let (agent, observed) = state_disambiguation_agent(
vec![
r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
"Request executed.",
],
true,
None,
true,
);
let response = agent.chat("Send the report to Ada").await.unwrap();
assert_eq!(response.content, "Request executed.");
assert_eq!(observed.call_count(), 2);
}
#[tokio::test]
async fn state_override_cannot_activate_a_disabled_top_level_manager() {
let (agent, observed) =
state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
assert!(!agent.has_disambiguation());
let response = agent.chat("Send it").await.unwrap();
assert_eq!(response.content, "Request executed.");
assert_eq!(observed.call_count(), 1);
}
#[tokio::test]
async fn native_required_choice_executes_through_the_shared_tool_path() {
let mut mock = MockLLMProvider::new("native-required");
mock.set_tool_choice(Some(ToolChoice::Required));
mock.add_response(
LLMResponse::new("", FinishReason::ToolCall)
.with_tool_calls(vec![ToolCall {
id: "provider-call-1".to_string(),
name: "calculator".to_string(),
arguments: serde_json::json!({"expression": "2 + 2"}),
}])
.unwrap(),
);
mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
let observed = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("Use the calculator when needed.")
.llm(Arc::new(mock))
.tool(Arc::new(CalculatorTool::new()))
.build()
.unwrap();
let response = agent.chat("What is 2 + 2?").await.unwrap();
assert_eq!(response.content, "The answer is 4.");
assert_eq!(
response.tool_calls.as_ref().unwrap()[0].id,
"provider-call-1"
);
let calls = observed.call_history();
assert_eq!(calls.len(), 2);
assert!(matches!(
calls[0].request.as_ref().map(|request| &request.choice),
Some(ToolChoice::Required)
));
assert!(matches!(
calls[1].request.as_ref().map(|request| &request.choice),
Some(ToolChoice::Auto)
));
}
#[tokio::test]
async fn prompt_fallback_uses_one_corrective_retry() {
let mut mock = MockLLMProvider::new("prompt-required");
mock.set_tool_choice(Some(ToolChoice::Required));
mock.set_native_tool_support(false);
mock.set_responses(
vec![
"I can calculate that.".to_string(),
r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
"The answer is 4.".to_string(),
],
false,
);
let observed = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("Use tools.")
.llm(Arc::new(mock))
.tool(Arc::new(CalculatorTool::new()))
.build()
.unwrap();
let response = agent.chat("What is 2 + 2?").await.unwrap();
assert_eq!(response.content, "The answer is 4.");
assert_eq!(observed.call_count(), 3);
let corrective = &observed.call_history()[1].messages;
assert!(
corrective
.last()
.unwrap()
.content
.contains("previous response")
);
}
#[tokio::test]
async fn prompt_fallback_fails_after_one_noncompliant_retry() {
let mut mock = MockLLMProvider::new("prompt-required-failure");
mock.set_tool_choice(Some(ToolChoice::Required));
mock.set_native_tool_support(false);
mock.set_responses(
vec!["No tool.".to_string(), "Still no tool.".to_string()],
false,
);
let observed = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("Use tools.")
.llm(Arc::new(mock))
.tool(Arc::new(CalculatorTool::new()))
.build()
.unwrap();
let error = agent.chat("What is 2 + 2?").await.unwrap_err();
assert!(error.to_string().contains("one corrective retry"));
assert_eq!(observed.call_count(), 2);
}
#[tokio::test]
async fn specific_choice_cannot_widen_the_effective_grant() {
let mut mock = MockLLMProvider::new("specific-outside-grant");
mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
let observed = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("Use tools.")
.llm(Arc::new(mock))
.tool(Arc::new(CalculatorTool::new()))
.build()
.unwrap();
let error = agent.chat("Generate a value.").await.unwrap_err();
assert!(error.to_string().contains("is not registered"));
assert_eq!(observed.call_count(), 0);
}
#[tokio::test]
async fn none_choice_exposes_no_tool_protocol() {
let mut mock = MockLLMProvider::new("no-tools");
mock.set_tool_choice(Some(ToolChoice::None));
mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
let observed = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("Answer directly.")
.llm(Arc::new(mock))
.tool(Arc::new(CalculatorTool::new()))
.build()
.unwrap();
let response = agent.chat("Hello").await.unwrap();
assert!(response.tool_calls.is_none());
assert_eq!(observed.call_count(), 1);
let call = observed.last_call().unwrap();
assert!(call.request.is_none());
assert!(
call.messages
.iter()
.all(|message| !message.content.contains("Available tools:"))
);
}
struct RuntimeStorage {
capabilities: Box<[StorageCapability]>,
snapshots: RwLock<HashMap<String, AgentSnapshot>>,
metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
metadata_save_calls: AtomicU64,
metadata_load_calls: AtomicU64,
fail_metadata_save: AtomicBool,
fail_metadata_load: AtomicBool,
}
impl RuntimeStorage {
fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
Self {
capabilities: capabilities.into_iter().collect(),
snapshots: RwLock::new(HashMap::new()),
metadata: RwLock::new(HashMap::new()),
metadata_save_calls: AtomicU64::new(0),
metadata_load_calls: AtomicU64::new(0),
fail_metadata_save: AtomicBool::new(false),
fail_metadata_load: AtomicBool::new(false),
}
}
}
#[async_trait]
impl AgentStorage for RuntimeStorage {
fn supports(&self, capability: StorageCapability) -> bool {
self.capabilities.contains(&capability)
}
async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
self.snapshots
.write()
.insert(session_id.to_string(), snapshot.clone());
Ok(())
}
async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
Ok(self.snapshots.read().get(session_id).cloned())
}
async fn delete(&self, session_id: &str) -> Result<()> {
self.snapshots.write().remove(session_id);
Ok(())
}
async fn list_sessions(&self) -> Result<Vec<String>> {
Ok(self.snapshots.read().keys().cloned().collect())
}
async fn save_snapshot_with_metadata(
&self,
session_id: &str,
snapshot: &AgentSnapshot,
metadata: &ai_agents_core::SessionMetadata,
) -> Result<()> {
self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
if self.fail_metadata_save.load(Ordering::SeqCst) {
return Err(AgentError::Persistence("metadata save failed".into()));
}
self.snapshots
.write()
.insert(session_id.to_string(), snapshot.clone());
self.metadata
.write()
.insert(session_id.to_string(), metadata.clone());
Ok(())
}
async fn save_metadata(
&self,
session_id: &str,
metadata: &ai_agents_core::SessionMetadata,
) -> Result<()> {
self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
if self.fail_metadata_save.load(Ordering::SeqCst) {
return Err(AgentError::Persistence("metadata save failed".into()));
}
self.metadata
.write()
.insert(session_id.to_string(), metadata.clone());
Ok(())
}
async fn load_metadata(
&self,
session_id: &str,
) -> Result<Option<ai_agents_core::SessionMetadata>> {
self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
if self.fail_metadata_load.load(Ordering::SeqCst) {
return Err(AgentError::Persistence("metadata load failed".into()));
}
Ok(self.metadata.read().get(session_id).cloned())
}
}
fn runtime_storage_agent() -> RuntimeAgent {
AgentBuilder::new()
.system_prompt("Test runtime storage integration.")
.llm(Arc::new(mock_with_response("done")))
.build()
.unwrap()
}
fn restore_spec(id: &str) -> crate::spec::AgentSpec {
crate::spec::AgentSpec {
name: id.to_string(),
system_prompt: format!("Restore child {id}."),
..crate::spec::AgentSpec::default()
}
}
fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
ai_agents_core::SpawnedAgentEntry {
id: id.to_string(),
name: id.to_string(),
spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
}
}
fn restore_spawner(
storage: Arc<RuntimeStorage>,
max_agents: usize,
) -> (
Arc<crate::spawner::AgentSpawner>,
Arc<crate::spawner::AgentRegistry>,
) {
let mut llms = LLMRegistry::new();
llms.register("default", Arc::new(mock_with_response("done")));
(
Arc::new(
crate::spawner::AgentSpawner::new()
.with_shared_llms(llms)
.with_shared_storage(storage)
.with_max_agents(max_agents),
),
Arc::new(crate::spawner::AgentRegistry::new()),
)
}
async fn save_restore_target(
parent: &RuntimeAgent,
storage: &RuntimeStorage,
session_id: &str,
entries: Vec<ai_agents_core::SpawnedAgentEntry>,
) {
let mut snapshot = parent.save_state().await.unwrap();
snapshot.spawned_agents = Some(entries);
storage.save(session_id, &snapshot).await.unwrap();
storage
.save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
.await
.unwrap();
}
#[tokio::test]
async fn storage_init_requires_storage_for_actor_facts() {
let facts = ai_agents_facts::FactsConfig {
enabled: true,
..Default::default()
};
let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
let error = agent.init_storage().await.unwrap_err();
assert!(matches!(
error,
AgentError::Config(message)
if message.contains("actor facts or actor memory")
&& message.contains("none is configured or injected")
));
}
#[tokio::test]
async fn storage_init_validates_actor_facts_capability() {
let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
let actor_memory = ai_agents_facts::ActorMemoryConfig {
enabled: true,
..Default::default()
};
let agent = runtime_storage_agent()
.with_storage(storage)
.with_facts_config(Some(actor_memory), None);
assert!(matches!(
agent.init_storage().await,
Err(AgentError::UnsupportedStorageCapability(
StorageCapability::ActorFacts
))
));
}
#[tokio::test]
async fn blocking_chat_rejects_unsupported_required_storage() {
let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
let facts = ai_agents_facts::FactsConfig {
enabled: true,
..Default::default()
};
let agent = runtime_storage_agent()
.with_storage(storage)
.with_facts_config(None, Some(facts));
assert!(matches!(
agent.chat("hello").await,
Err(AgentError::UnsupportedStorageCapability(
StorageCapability::ActorFacts
))
));
}
#[tokio::test]
async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
let config = ai_agents_relationships::RelationshipConfig {
enabled: true,
..Default::default()
};
let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
let agent = runtime_storage_agent()
.with_storage(storage)
.with_relationships(manager);
assert!(matches!(
agent.chat_stream("hello").await,
Err(AgentError::UnsupportedStorageCapability(
StorageCapability::ActorRelationships
))
));
}
#[tokio::test]
async fn storage_init_completes_facts_for_injected_storage() {
let storage = Arc::new(RuntimeStorage::new([
StorageCapability::Snapshot,
StorageCapability::ActorFacts,
]));
let facts = ai_agents_facts::FactsConfig {
enabled: true,
..Default::default()
};
let agent = runtime_storage_agent()
.with_storage(storage)
.with_facts_config(None, Some(facts));
agent.init_storage().await.unwrap();
assert!(agent.fact_store().is_some());
}
#[tokio::test]
async fn storage_init_requires_storage_for_persistent_relationships() {
let config = ai_agents_relationships::RelationshipConfig {
enabled: true,
..Default::default()
};
let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
let agent = runtime_storage_agent().with_relationships(manager);
let error = agent.init_storage().await.unwrap_err();
assert!(matches!(
error,
AgentError::Config(message)
if message.contains("persistent relationships")
&& message.contains("none is configured or injected")
));
}
#[tokio::test]
async fn storage_init_validates_persistent_relationships_capability() {
let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
let config = ai_agents_relationships::RelationshipConfig {
enabled: true,
..Default::default()
};
let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
let agent = runtime_storage_agent()
.with_storage(storage)
.with_relationships(manager);
assert!(matches!(
agent.init_storage().await,
Err(AgentError::UnsupportedStorageCapability(
StorageCapability::ActorRelationships
))
));
}
#[tokio::test]
async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
let storage = Arc::new(RuntimeStorage::new([
StorageCapability::Snapshot,
StorageCapability::SessionMetadata,
]));
let agent = runtime_storage_agent().with_storage(storage.clone());
agent.set_actor_id("old-actor").unwrap();
agent.save_session("old").await.unwrap();
storage
.save("target", &agent.save_state().await.unwrap())
.await
.unwrap();
storage
.save_metadata("target", &ai_agents_core::SessionMetadata::default())
.await
.unwrap();
assert!(agent.load_session("target").await.unwrap());
assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
assert_eq!(agent.actor_id(), None);
}
#[tokio::test]
async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
let storage = Arc::new(RuntimeStorage::new([
StorageCapability::Snapshot,
StorageCapability::SessionMetadata,
]));
let (spawner, registry) = restore_spawner(storage.clone(), 3);
let parent = runtime_storage_agent()
.with_storage(storage.clone())
.with_spawner_handles(Arc::clone(&spawner), Arc::clone(®istry));
for id in ["a", "b"] {
let spawned = spawner
.spawn_with_id(id.to_string(), restore_spec(id))
.await
.unwrap();
spawned.agent.save_session("grow").await.unwrap();
registry.register(spawned).await.unwrap();
}
let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
staged_c
.save("grow", &AgentSnapshot::new("c".into()))
.await
.unwrap();
staged_c
.save_metadata("grow", &ai_agents_core::SessionMetadata::default())
.await
.unwrap();
save_restore_target(
&parent,
storage.as_ref(),
"grow",
vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
)
.await;
assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
assert_eq!(registry.count(), 3);
assert!(registry.contains("c"));
assert_eq!(spawner.spawned_count(), 3);
for id in ["a", "b"] {
registry
.get(id)
.unwrap()
.save_session("shrink")
.await
.unwrap();
}
save_restore_target(
&parent,
storage.as_ref(),
"shrink",
vec![restore_entry("a"), restore_entry("b")],
)
.await;
assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
assert_eq!(registry.count(), 2);
assert!(!registry.contains("c"));
assert_eq!(spawner.spawned_count(), 2);
save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
assert_eq!(registry.count(), 0);
assert_eq!(spawner.spawned_count(), 0);
assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
}
#[tokio::test]
async fn storage_session_metadata_is_called_only_when_advertised() {
let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
storage.fail_metadata_save.store(true, Ordering::SeqCst);
storage.fail_metadata_load.store(true, Ordering::SeqCst);
let agent = runtime_storage_agent().with_storage(storage.clone());
agent.save_session("session").await.unwrap();
assert!(agent.load_session("session").await.unwrap());
assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
}
#[cfg(feature = "sqlite")]
#[tokio::test]
async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
let directory =
std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
let path = directory.join("sessions.sqlite");
let path_string = path.to_string_lossy().into_owned();
let storage = Arc::new(
ai_agents_storage::SqliteStorage::new(&path_string)
.await
.unwrap(),
);
let agent = runtime_storage_agent().with_storage(storage.clone());
agent.set_session_metadata(ai_agents_core::SessionMetadata {
tags: vec!["initial".into()],
..Default::default()
});
agent.chat("persist this turn").await.unwrap();
agent.save_session("session").await.unwrap();
agent.set_session_metadata(ai_agents_core::SessionMetadata {
tags: vec!["updated".into()],
..Default::default()
});
agent.save_session("session").await.unwrap();
assert!(
agent
.list_sessions_filtered(&ai_agents_core::SessionFilter {
tags: Some(vec!["initial".into()]),
..Default::default()
})
.await
.unwrap()
.is_empty()
);
assert_eq!(
agent
.list_sessions_filtered(&ai_agents_core::SessionFilter {
tags: Some(vec!["updated".into()]),
..Default::default()
})
.await
.unwrap()
.len(),
1
);
drop(agent);
storage.close().await;
drop(storage);
let reopened_storage = Arc::new(
ai_agents_storage::SqliteStorage::new(&path_string)
.await
.unwrap(),
);
let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
assert!(restored.load_session("session").await.unwrap());
assert_eq!(restored.session_metadata().tags, vec!["updated"]);
assert_eq!(
restored.current_session_id.read().as_deref(),
Some("session")
);
assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
assert_eq!(
restored
.list_sessions_filtered(&ai_agents_core::SessionFilter {
tags: Some(vec!["updated".into()]),
..Default::default()
})
.await
.unwrap()
.len(),
1
);
drop(restored);
reopened_storage.close().await;
drop(reopened_storage);
crate::remove_sqlite_test_directory(&directory)
.await
.unwrap();
}
#[tokio::test]
async fn storage_session_metadata_backend_failures_propagate() {
let storage = Arc::new(RuntimeStorage::new([
StorageCapability::Snapshot,
StorageCapability::SessionMetadata,
]));
let agent = runtime_storage_agent().with_storage(storage.clone());
agent.save_session("session").await.unwrap();
storage
.save("target", &agent.save_state().await.unwrap())
.await
.unwrap();
storage.fail_metadata_load.store(true, Ordering::SeqCst);
assert!(matches!(
agent.load_session("target").await,
Err(AgentError::Persistence(message)) if message == "metadata load failed"
));
assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
storage.fail_metadata_save.store(true, Ordering::SeqCst);
assert!(matches!(
agent.save_session("session").await,
Err(AgentError::Persistence(message)) if message == "metadata save failed"
));
}
struct ProviderFutureDropSignal {
dropped: Arc<AtomicBool>,
}
impl Drop for ProviderFutureDropSignal {
fn drop(&mut self) {
self.dropped.store(true, Ordering::SeqCst);
}
}
struct BufferedLockingProvider {
lock: Arc<tokio::sync::Mutex<()>>,
stream_started: Arc<tokio::sync::Notify>,
stream_dropped: Arc<AtomicBool>,
committed_after_drop: Arc<AtomicBool>,
}
#[async_trait]
impl LLMProvider for BufferedLockingProvider {
async fn complete(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<LLMResponse, LLMError> {
let _guard = self.lock.lock().await;
self.committed_after_drop
.store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
Ok(LLMResponse::new(
"Committed technical response.",
FinishReason::Stop,
))
}
async fn complete_stream(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<
Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
LLMError,
> {
let _guard = self.lock.lock().await;
let _drop_signal = ProviderFutureDropSignal {
dropped: Arc::clone(&self.stream_dropped),
};
self.stream_started.notify_one();
std::future::pending().await
}
fn provider_name(&self) -> &str {
"buffered-locking"
}
fn supports(&self, _feature: LLMFeature) -> bool {
false
}
}
struct PendingDropStream {
dropped: Arc<AtomicBool>,
dropped_notify: Arc<tokio::sync::Notify>,
}
impl Stream for PendingDropStream {
type Item = std::result::Result<LLMChunk, LLMError>;
fn poll_next(
self: Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
std::task::Poll::Pending
}
}
impl Drop for PendingDropStream {
fn drop(&mut self) {
self.dropped.store(true, Ordering::SeqCst);
self.dropped_notify.notify_one();
}
}
struct EstablishedStreamProvider {
stream_started: Arc<tokio::sync::Notify>,
stream_dropped: Arc<AtomicBool>,
stream_dropped_notify: Arc<tokio::sync::Notify>,
committed_after_drop: Arc<AtomicBool>,
}
#[async_trait]
impl LLMProvider for EstablishedStreamProvider {
async fn complete(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<LLMResponse, LLMError> {
if !self.stream_dropped.load(Ordering::SeqCst) {
self.stream_dropped_notify.notified().await;
}
self.committed_after_drop
.store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
Ok(LLMResponse::new(
"Committed technical response.",
FinishReason::Stop,
))
}
async fn complete_stream(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<
Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
LLMError,
> {
self.stream_started.notify_one();
Ok(Box::new(PendingDropStream {
dropped: Arc::clone(&self.stream_dropped),
dropped_notify: Arc::clone(&self.stream_dropped_notify),
}))
}
fn provider_name(&self) -> &str {
"established-stream"
}
fn supports(&self, _feature: LLMFeature) -> bool {
false
}
}
struct FirstCallLockingProvider {
lock: Arc<tokio::sync::Mutex<()>>,
first_started: Arc<tokio::sync::Notify>,
first_dropped: Arc<AtomicBool>,
committed_after_drop: Arc<AtomicBool>,
calls: AtomicU64,
}
#[async_trait]
impl LLMProvider for FirstCallLockingProvider {
async fn complete(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<LLMResponse, LLMError> {
let _guard = self.lock.lock().await;
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if call == 0 {
let _drop_signal = ProviderFutureDropSignal {
dropped: Arc::clone(&self.first_dropped),
};
self.first_started.notify_one();
return std::future::pending().await;
}
self.committed_after_drop
.store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
Ok(LLMResponse::new(
"Committed technical response.",
FinishReason::Stop,
))
}
async fn complete_stream(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<
Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
LLMError,
> {
Err(LLMError::Other(
"streaming is not used in this test".to_string(),
))
}
fn provider_name(&self) -> &str {
"first-call-locking"
}
fn supports(&self, _feature: LLMFeature) -> bool {
false
}
}
struct RoutingAfterProviderStart {
provider_started: Arc<tokio::sync::Notify>,
}
#[async_trait]
impl LLMProvider for RoutingAfterProviderStart {
async fn complete(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<LLMResponse, LLMError> {
self.provider_started.notified().await;
Ok(LLMResponse::new("1", FinishReason::Stop))
}
async fn complete_stream(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<
Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
LLMError,
> {
Err(LLMError::Other(
"streaming is not used in this test".to_string(),
))
}
fn provider_name(&self) -> &str {
"routing-after-start"
}
fn supports(&self, _feature: LLMFeature) -> bool {
false
}
}
struct ResponseCountingHooks {
responses: Arc<std::sync::atomic::AtomicUsize>,
}
struct RootTurnProbeProvider {
complete_entered: tokio::sync::mpsc::UnboundedSender<()>,
}
struct ResponseChatHooks {
target: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
invoked: AtomicBool,
nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
}
struct ConcurrentResponseHooks {
registry: Weak<crate::spawner::AgentRegistry>,
child_id: String,
invoked: AtomicBool,
nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
}
struct RetryDeadlineTool {
calls: Arc<std::sync::atomic::AtomicUsize>,
deadlines: Arc<parking_lot::Mutex<Vec<chrono::DateTime<chrono::Utc>>>>,
remaining_ms: Arc<parking_lot::Mutex<Vec<i64>>>,
}
struct ToolLifecycleRecordingHooks {
events: parking_lot::Mutex<Vec<String>>,
records: parking_lot::Mutex<Vec<ToolExecutionRecord>>,
}
impl ToolLifecycleRecordingHooks {
fn new() -> Self {
Self {
events: parking_lot::Mutex::new(Vec::new()),
records: parking_lot::Mutex::new(Vec::new()),
}
}
fn events(&self) -> Vec<String> {
self.events.lock().clone()
}
fn records(&self) -> Vec<ToolExecutionRecord> {
self.records.lock().clone()
}
}
struct ContextEchoTool;
#[async_trait]
impl LLMProvider for RootTurnProbeProvider {
async fn complete(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<LLMResponse, LLMError> {
let _ = self.complete_entered.send(());
Ok(LLMResponse::new("blocking complete", FinishReason::Stop))
}
async fn complete_stream(
&self,
_messages: &[ChatMessage],
_config: Option<&LLMConfig>,
) -> std::result::Result<
Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
LLMError,
> {
Ok(Box::new(futures::stream::iter(vec![Ok(
LLMChunk::final_chunk("stream complete", FinishReason::Stop, None),
)])))
}
fn provider_name(&self) -> &str {
"root-turn-probe"
}
fn supports(&self, feature: LLMFeature) -> bool {
matches!(feature, LLMFeature::Streaming)
}
}
#[async_trait]
impl ai_agents_core::Tool for ContextEchoTool {
fn id(&self) -> &str {
"context_echo"
}
fn name(&self) -> &str {
"Context Echo"
}
fn description(&self) -> &str {
"Returns selected execution context fields."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
ai_agents_core::ToolPolicyBindings {
path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
"max_results",
ai_agents_core::ResultLimitKind::MaxResults,
)],
..Default::default()
}
}
async fn execute(
&self,
_args: Value,
ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
ToolResult::ok(
serde_json::json!({
"requested_name": ctx.requested_name,
"canonical_id": ctx.canonical_id,
"display_name": ctx.display_name,
"max_results": ctx.limits.max_results,
"custom_config": ctx.custom_config,
})
.to_string(),
)
}
}
#[async_trait]
impl ai_agents_core::Tool for RetryDeadlineTool {
fn id(&self) -> &str {
"retry_deadline"
}
fn name(&self) -> &str {
"Retry Deadline"
}
fn description(&self) -> &str {
"Records one deadline per retry invocation."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
ai_agents_core::ToolSafetyMetadata::compute()
}
fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
let mut classification =
ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
classification.timeout_ms = Some(1_000);
classification.safely_retryable = true;
classification
}
async fn execute(
&self,
_args: Value,
ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
let deadline = ctx
.deadline
.expect("each invocation must receive a deadline");
self.remaining_ms.lock().push(
deadline
.signed_duration_since(chrono::Utc::now())
.num_milliseconds(),
);
self.deadlines.lock().push(deadline);
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if call == 0 {
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
ToolResult::error("retry")
} else {
ToolResult::ok("done")
}
}
}
struct ClassifiedTimeoutTool {
id: &'static str,
calls: Arc<std::sync::atomic::AtomicUsize>,
timeout_ms: u64,
sleep_ms: u64,
requires_approval: bool,
remaining_ms: Arc<parking_lot::Mutex<Vec<i64>>>,
}
struct ApprovalModifiedTimeoutTool {
calls: Arc<std::sync::atomic::AtomicUsize>,
}
struct SlowTool;
struct FlakyWriteTool {
calls: Arc<std::sync::atomic::AtomicUsize>,
}
struct LockedWriteTool {
active: Arc<std::sync::atomic::AtomicUsize>,
max_active: Arc<std::sync::atomic::AtomicUsize>,
}
struct MultiResourceWriteTool {
active: Arc<std::sync::atomic::AtomicUsize>,
max_active: Arc<std::sync::atomic::AtomicUsize>,
}
#[derive(Clone)]
struct PathMutationGate {
entered: Arc<AtomicBool>,
entered_notify: Arc<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
}
impl PathMutationGate {
fn new() -> Self {
Self {
entered: Arc::new(AtomicBool::new(false)),
entered_notify: Arc::new(tokio::sync::Notify::new()),
release: Arc::new(tokio::sync::Notify::new()),
}
}
async fn wait_until_entered(&self) {
if !self.entered.load(Ordering::SeqCst) {
self.entered_notify.notified().await;
}
}
fn release(&self) {
self.release.notify_one();
}
}
struct BlockingPathMutationTool {
id: &'static str,
path_fields: Vec<ai_agents_core::PathPolicyBinding>,
gate: PathMutationGate,
}
struct NoBindingWriteTool {
active: Arc<std::sync::atomic::AtomicUsize>,
max_active: Arc<std::sync::atomic::AtomicUsize>,
}
struct RecoveryTestTool {
id: String,
succeeds: bool,
calls: Arc<std::sync::atomic::AtomicUsize>,
max_output_chars: Option<usize>,
}
struct BlockingApprovalHandler {
entered: Arc<tokio::sync::Barrier>,
release: Arc<tokio::sync::Notify>,
result: ApprovalResult,
}
struct CountingApprovalHandler {
calls: Arc<std::sync::atomic::AtomicUsize>,
}
struct DriftingFallbackProvider {
refreshed: AtomicBool,
primary_calls: Arc<std::sync::atomic::AtomicUsize>,
secondary_calls: Arc<std::sync::atomic::AtomicUsize>,
}
struct RefreshFallbackProviderHooks {
agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
lifecycle: Arc<ToolLifecycleRecordingHooks>,
}
struct RuntimeWebFetchTransport {
calls: Arc<std::sync::atomic::AtomicUsize>,
}
struct RuntimeWebFetchResolver;
struct ReentrantToolHooks {
agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
invoked: AtomicBool,
nested_success: AtomicBool,
}
#[async_trait]
impl ai_agents_core::Tool for ClassifiedTimeoutTool {
fn id(&self) -> &str {
self.id
}
fn name(&self) -> &str {
"Classified Timeout"
}
fn description(&self) -> &str {
"Records and waits under one call-level timeout."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
let mut classification =
ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
classification.timeout_ms = Some(self.timeout_ms);
classification.requires_approval = self.requires_approval;
classification
}
async fn execute(
&self,
_args: Value,
ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
self.calls.fetch_add(1, Ordering::SeqCst);
let deadline = ctx
.deadline
.expect("each invocation must receive a deadline");
self.remaining_ms.lock().push(
deadline
.signed_duration_since(chrono::Utc::now())
.num_milliseconds(),
);
tokio::time::sleep(Duration::from_millis(self.sleep_ms)).await;
ToolResult::ok("done")
}
}
#[async_trait]
impl ai_agents_core::Tool for ApprovalModifiedTimeoutTool {
fn id(&self) -> &str {
"approval_modified_timeout"
}
fn name(&self) -> &str {
"Approval Modified Timeout"
}
fn description(&self) -> &str {
"Becomes invalid only after approval modifies its arguments."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
ai_agents_core::ToolPolicyBindings {
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
..Default::default()
}
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
ai_agents_core::ToolSafetyMetadata {
read_only: false,
concurrency_safe: false,
operation: ai_agents_core::ToolOperationKind::Write,
side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
requires_network: false,
destructive: false,
open_world: false,
host_dependent: false,
requires_user_interaction: false,
supports_cancellation: true,
default_requires_approval: true,
should_defer_schema: false,
max_output_chars: Some(1024),
max_result_size_chars: Some(1024),
}
}
fn classify_call(&self, args: &Value) -> ai_agents_core::ToolCallClassification {
let mut classification =
ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
classification.timeout_ms = Some(if args["invalid_timeout"].as_bool() == Some(true) {
u64::MAX
} else {
1_000
});
classification
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
self.calls.fetch_add(1, Ordering::SeqCst);
ToolResult::ok("unexpected")
}
}
#[async_trait]
impl ai_agents_core::Tool for SlowTool {
fn id(&self) -> &str {
"slow"
}
fn name(&self) -> &str {
"Slow"
}
fn description(&self) -> &str {
"Waits until cancelled or timed out."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
ToolResult::ok("done")
}
}
#[async_trait]
impl ai_agents_core::Tool for FlakyWriteTool {
fn id(&self) -> &str {
"flaky_write"
}
fn name(&self) -> &str {
"Flaky Write"
}
fn description(&self) -> &str {
"Fails on the first write attempt."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
}
fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
ai_agents_core::ToolPolicyBindings {
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
..Default::default()
}
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
ai_agents_core::ToolSafetyMetadata {
read_only: false,
concurrency_safe: false,
operation: ai_agents_core::ToolOperationKind::Write,
side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
requires_network: false,
destructive: false,
open_world: false,
host_dependent: false,
requires_user_interaction: false,
supports_cancellation: true,
default_requires_approval: false,
should_defer_schema: false,
max_output_chars: Some(1024),
max_result_size_chars: Some(1024),
}
}
fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
let mut classification =
ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
classification.safely_retryable = false;
classification
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if call == 0 {
ToolResult::error("first failure")
} else {
ToolResult::ok("second success")
}
}
}
#[async_trait]
impl ai_agents_core::Tool for LockedWriteTool {
fn id(&self) -> &str {
"locked_write"
}
fn name(&self) -> &str {
"Locked Write"
}
fn description(&self) -> &str {
"Tracks concurrent execution on one resource."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
}
fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
ai_agents_core::ToolPolicyBindings {
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
..Default::default()
}
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
ai_agents_core::ToolSafetyMetadata {
read_only: false,
concurrency_safe: false,
operation: ai_agents_core::ToolOperationKind::Write,
side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
requires_network: false,
destructive: false,
open_world: false,
host_dependent: false,
requires_user_interaction: false,
supports_cancellation: true,
default_requires_approval: false,
should_defer_schema: false,
max_output_chars: Some(1024),
max_result_size_chars: Some(1024),
}
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
loop {
let current_max = self.max_active.load(Ordering::SeqCst);
if active <= current_max {
break;
}
if self
.max_active
.compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
.is_ok()
{
break;
}
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
self.active.fetch_sub(1, Ordering::SeqCst);
ToolResult::ok("done")
}
}
#[async_trait]
impl ai_agents_core::Tool for MultiResourceWriteTool {
fn id(&self) -> &str {
"multi_resource_write"
}
fn name(&self) -> &str {
"Multi Resource Write"
}
fn description(&self) -> &str {
"Tracks concurrent execution across source and destination resources."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
ai_agents_core::ToolPolicyBindings {
path_fields: vec![
ai_agents_core::PathPolicyBinding::read_write("source_path"),
ai_agents_core::PathPolicyBinding::write("destination_path"),
],
..Default::default()
}
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
LockedWriteTool {
active: Arc::clone(&self.active),
max_active: Arc::clone(&self.max_active),
}
.safety_metadata()
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
self.max_active.fetch_max(active, Ordering::SeqCst);
tokio::time::sleep(std::time::Duration::from_millis(75)).await;
self.active.fetch_sub(1, Ordering::SeqCst);
ToolResult::ok("done")
}
}
#[async_trait]
impl ai_agents_core::Tool for BlockingPathMutationTool {
fn id(&self) -> &str {
self.id
}
fn name(&self) -> &str {
self.id
}
fn description(&self) -> &str {
"Blocks a path mutation until the test releases it."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
ai_agents_core::ToolPolicyBindings {
path_fields: self.path_fields.clone(),
..Default::default()
}
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
ai_agents_core::ToolSafetyMetadata {
read_only: false,
concurrency_safe: false,
operation: ai_agents_core::ToolOperationKind::Write,
side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
requires_network: false,
destructive: false,
open_world: false,
host_dependent: false,
requires_user_interaction: false,
supports_cancellation: true,
default_requires_approval: false,
should_defer_schema: false,
max_output_chars: Some(1024),
max_result_size_chars: Some(1024),
}
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
self.gate.entered.store(true, Ordering::SeqCst);
self.gate.entered_notify.notify_one();
self.gate.release.notified().await;
ToolResult::ok("done")
}
}
#[async_trait]
impl ai_agents_core::Tool for NoBindingWriteTool {
fn id(&self) -> &str {
"no_binding_write"
}
fn name(&self) -> &str {
"No Binding Write"
}
fn description(&self) -> &str {
"Tracks concurrent execution without resource bindings."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
LockedWriteTool {
active: Arc::clone(&self.active),
max_active: Arc::clone(&self.max_active),
}
.safety_metadata()
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
self.max_active.fetch_max(active, Ordering::SeqCst);
tokio::time::sleep(std::time::Duration::from_millis(75)).await;
self.active.fetch_sub(1, Ordering::SeqCst);
ToolResult::ok("done")
}
}
#[async_trait]
impl ai_agents_core::Tool for RecoveryTestTool {
fn id(&self) -> &str {
&self.id
}
fn name(&self) -> &str {
&self.id
}
fn description(&self) -> &str {
"Records recovery execution and returns a configured result."
}
fn input_schema(&self) -> Value {
serde_json::json!({"type": "object"})
}
fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
ai_agents_core::ToolPolicyBindings {
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
..Default::default()
}
}
fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
ai_agents_core::ToolSafetyMetadata {
read_only: false,
concurrency_safe: false,
operation: ai_agents_core::ToolOperationKind::Write,
side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
requires_network: false,
destructive: false,
open_world: false,
host_dependent: false,
requires_user_interaction: false,
supports_cancellation: true,
default_requires_approval: false,
should_defer_schema: false,
max_output_chars: Some(self.max_output_chars.unwrap_or(1024)),
max_result_size_chars: Some(1024),
}
}
async fn execute(
&self,
_args: Value,
_ctx: ai_agents_core::ToolExecutionContext,
) -> ToolResult {
self.calls.fetch_add(1, Ordering::SeqCst);
let mut result = if self.succeeds {
ToolResult::ok(format!("{} succeeded", self.id))
} else {
ToolResult::error(format!("{} failed", self.id))
};
result.metadata = Some(HashMap::from([(
"recovery_test_tool".to_string(),
Value::String(self.id.clone()),
)]));
result
}
}
#[async_trait]
impl WebFetchTransport for RuntimeWebFetchTransport {
async fn send(
&self,
_request: WebFetchTransportRequest,
) -> std::result::Result<WebFetchTransportResponse, String> {
Err("validated addresses are required".to_string())
}
async fn send_validated(
&self,
_request: WebFetchTransportRequest,
_addresses: &[std::net::SocketAddr],
) -> std::result::Result<WebFetchTransportResponse, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(WebFetchTransportResponse {
status: 200,
content_type: Some("text/plain".to_string()),
location: None,
body: b"approved".to_vec(),
})
}
}
#[async_trait]
impl WebFetchResolver for RuntimeWebFetchResolver {
async fn resolve(
&self,
_host: &str,
_port: u16,
) -> std::result::Result<Vec<std::net::IpAddr>, String> {
Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
93, 184, 216, 34,
))])
}
}
#[async_trait]
impl ToolProvider for DriftingFallbackProvider {
fn id(&self) -> &str {
"drifting_fallback"
}
fn name(&self) -> &str {
"Drifting Fallback"
}
fn provider_type(&self) -> ToolProviderType {
ToolProviderType::Custom
}
async fn list_tools(&self) -> Vec<ToolDescriptor> {
let alias = ToolAliases::new().with_name("en", "fallback alias");
let mut primary = ToolDescriptor::new(
"primary",
"Primary",
"Fails before fallback.",
serde_json::json!({"type": "object"}),
);
let mut secondary = ToolDescriptor::new(
"secondary",
"Secondary",
"Must not execute after final canonical drift.",
serde_json::json!({"type": "object"}),
);
if self.refreshed.load(Ordering::SeqCst) {
primary = primary.with_aliases(alias);
} else {
secondary = secondary.with_aliases(alias);
}
vec![primary, secondary]
}
async fn get_tool(&self, tool_id: &str) -> Option<Arc<dyn Tool>> {
let calls = match tool_id {
"primary" => Arc::clone(&self.primary_calls),
"secondary" => Arc::clone(&self.secondary_calls),
_ => return None,
};
Some(Arc::new(RecoveryTestTool {
id: tool_id.to_string(),
succeeds: false,
calls,
max_output_chars: None,
}))
}
fn supports_refresh(&self) -> bool {
true
}
async fn refresh(&self) -> std::result::Result<(), ToolProviderError> {
self.refreshed.store(true, Ordering::SeqCst);
Ok(())
}
}
#[async_trait]
impl AgentHooks for RefreshFallbackProviderHooks {
async fn on_tool_start(&self, tool: &str, args: &Value) {
self.lifecycle.on_tool_start(tool, args).await;
if tool != "secondary" {
return;
}
let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
if let Some(agent) = agent {
agent
.tools
.refresh_provider("drifting_fallback")
.await
.unwrap();
}
}
async fn on_tool_complete(&self, tool: &str, result: &ToolResult, duration_ms: u64) {
self.lifecycle
.on_tool_complete(tool, result, duration_ms)
.await;
}
async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
self.lifecycle.on_tool_execution_record(record).await;
}
async fn on_error(&self, error: &AgentError) {
self.lifecycle.on_error(error).await;
}
}
#[async_trait]
impl ApprovalHandler for BlockingApprovalHandler {
async fn request_approval(
&self,
_request: ai_agents_hitl::ApprovalRequest,
) -> ApprovalResult {
self.entered.wait().await;
self.release.notified().await;
self.result.clone()
}
}
#[async_trait]
impl ApprovalHandler for CountingApprovalHandler {
async fn request_approval(
&self,
_request: ai_agents_hitl::ApprovalRequest,
) -> ApprovalResult {
self.calls.fetch_add(1, Ordering::SeqCst);
ApprovalResult::Approved
}
}
#[async_trait]
impl AgentHooks for ReentrantToolHooks {
async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
return;
}
let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
if let Some(agent) = agent {
let result = agent
.invoke_tool(ToolExecutionRequest::new(
"nested-hook-call",
"reentrant_write",
serde_json::json!({"path": "./hook.txt"}),
ToolCallSource::Manual,
))
.await;
self.nested_success
.store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
}
}
}
#[async_trait]
impl AgentHooks for ResponseCountingHooks {
async fn on_response(&self, _response: &AgentResponse) {
self.responses.fetch_add(1, Ordering::SeqCst);
}
}
#[async_trait]
impl AgentHooks for ResponseChatHooks {
async fn on_response(&self, _response: &AgentResponse) {
if self.invoked.swap(true, Ordering::SeqCst) {
return;
}
let target = self.target.lock().as_ref().and_then(Weak::upgrade);
let result = if let Some(target) = target {
target
.chat("nested response hook call")
.await
.map(|response| response.content)
.map_err(|error| error.to_string())
} else {
Err("response hook target is unavailable".to_string())
};
*self.nested_result.lock() = Some(result);
}
}
#[async_trait]
impl AgentHooks for ConcurrentResponseHooks {
async fn on_response(&self, _response: &AgentResponse) {
if self.invoked.swap(true, Ordering::SeqCst) {
return;
}
let Some(registry) = self.registry.upgrade() else {
*self.nested_result.lock() =
Some(Err("concurrent registry is unavailable".to_string()));
return;
};
let agents = [ai_agents_state::ConcurrentAgentRef::Id(
self.child_id.clone(),
)];
let aggregation = ai_agents_state::AggregationConfig {
strategy: ai_agents_state::AggregationStrategy::FirstWins,
synthesizer_llm: None,
synthesizer_prompt: None,
vote: None,
};
let result = crate::orchestration::concurrent(
®istry,
"nested concurrent response hook call",
&agents,
&aggregation,
None,
Some(1),
None,
ai_agents_state::PartialFailureAction::Abort,
None,
)
.await
.map(|result| result.response.content)
.map_err(|error| error.to_string());
*self.nested_result.lock() = Some(result);
}
}
#[async_trait]
impl AgentHooks for ToolLifecycleRecordingHooks {
async fn on_tool_start(&self, tool: &str, _args: &Value) {
self.events.lock().push(format!("start:{tool}"));
}
async fn on_tool_complete(&self, tool: &str, result: &ToolResult, _duration_ms: u64) {
self.events
.lock()
.push(format!("complete:{tool}:{}", result.success));
}
async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
self.events.lock().push(format!(
"record:{}:{}",
record.canonical_id, record.executed
));
self.records.lock().push(record.clone());
}
async fn on_error(&self, _error: &AgentError) {
self.events.lock().push("error".to_string());
}
}
struct ApprovalRecordingHooks {
events: parking_lot::Mutex<Vec<String>>,
}
impl ApprovalRecordingHooks {
fn new() -> Self {
Self {
events: parking_lot::Mutex::new(Vec::new()),
}
}
fn events(&self) -> Vec<String> {
self.events.lock().clone()
}
}
#[async_trait]
impl AgentHooks for ApprovalRecordingHooks {
async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
self.events.lock().push(format!(
"raw:{}:{}",
request_id,
approval_result_name(result)
));
}
async fn on_approval_resolved(
&self,
request: &ai_agents_hitl::ApprovalRequest,
raw_result: &ApprovalResult,
outcome: &ApprovalResolvedOutcome,
) {
self.events.lock().push(format!(
"resolved:{}:{}:{}",
request.id,
approval_result_name(raw_result),
approval_outcome_name(outcome)
));
}
}
fn approval_result_name(result: &ApprovalResult) -> &'static str {
match result {
ApprovalResult::Approved => "approved",
ApprovalResult::Rejected { .. } => "rejected",
ApprovalResult::Modified { .. } => "modified",
ApprovalResult::Timeout => "timeout",
}
}
fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
match outcome {
ApprovalResolvedOutcome::Approved => "approved",
ApprovalResolvedOutcome::Rejected { .. } => "rejected",
ApprovalResolvedOutcome::Modified { .. } => "modified",
ApprovalResolvedOutcome::Error { .. } => "error",
}
}
fn assert_correlated_approval_events(
events: &[String],
raw_status: &str,
outcome_status: &str,
) {
assert_eq!(events.len(), 2);
let raw: Vec<_> = events[0].split(':').collect();
let resolved: Vec<_> = events[1].split(':').collect();
assert_eq!(raw[0], "raw");
assert_eq!(resolved[0], "resolved");
assert_eq!(raw[1], resolved[1]);
assert_eq!(raw[2], raw_status);
assert_eq!(resolved[2], raw_status);
assert_eq!(resolved[3], outcome_status);
}
fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
let policy = ai_agents_tools::ToolPolicyConfig {
enabled: policy_enabled,
write_paths: vec![".".to_string()],
require_confirmation: true,
..Default::default()
};
security.tools.insert("locked_write".to_string(), policy);
security
}
struct MutationTestWorkspace {
root: std::path::PathBuf,
}
impl MutationTestWorkspace {
fn new() -> Self {
let root = std::env::temp_dir().join(format!(
"ai-agents-runtime-mutation-{}",
uuid::Uuid::new_v4()
));
std::fs::create_dir_all(&root).unwrap();
Self { root }
}
}
impl Drop for MutationTestWorkspace {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.root);
}
}
async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
let strong_count = locks
.read()
.get("path-mutation:global")
.map_or(0, |lock| lock.strong_count());
if strong_count >= minimum {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("path mutation call did not reach the shared lock");
}
async fn assert_path_mutation_pair_serialized(
first_id: &'static str,
first_fields: Vec<ai_agents_core::PathPolicyBinding>,
first_args: Value,
second_id: &'static str,
second_fields: Vec<ai_agents_core::PathPolicyBinding>,
second_args: Value,
) {
let locks = new_tool_resource_locks();
let first_gate = PathMutationGate::new();
let second_gate = PathMutationGate::new();
second_gate.release();
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test global path mutation locking.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(BlockingPathMutationTool {
id: first_id,
path_fields: first_fields,
gate: first_gate.clone(),
}))
.tool(Arc::new(BlockingPathMutationTool {
id: second_id,
path_fields: second_fields,
gate: second_gate.clone(),
}))
.build()
.unwrap()
.with_shared_resource_locks(Arc::clone(&locks)),
);
let first = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
format!("{}-first", first_id),
first_id,
first_args,
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
first_gate.wait_until_entered().await;
let second = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
format!("{}-second", second_id),
second_id,
second_args,
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
wait_for_resource_lock_strong_count(&locks, 2).await;
assert!(!second_gate.entered.load(Ordering::SeqCst));
assert!(!second.is_finished());
first_gate.release();
let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
tokio::join!(first, second)
})
.await
.expect("serialized path mutation calls did not finish");
assert!(first.unwrap().success);
assert!(second.unwrap().success);
assert!(second_gate.entered.load(Ordering::SeqCst));
assert!(locks.read().is_empty());
}
#[derive(Clone, Copy)]
enum MutationDenial {
Policy,
Approval,
}
fn mutation_denial_security_config(
tool_id: &str,
workspace: &std::path::Path,
denial: MutationDenial,
) -> ToolSecurityConfig {
let workspace = workspace.to_string_lossy().into_owned();
let mut policy = ai_agents_tools::ToolPolicyConfig {
read_paths: vec![workspace.clone()],
write_paths: vec![workspace.clone()],
..Default::default()
};
match denial {
MutationDenial::Policy => policy.blocked_paths = vec![workspace],
MutationDenial::Approval => policy.require_confirmation = true,
}
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
security.tools.insert(tool_id.to_string(), policy);
security
}
async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
let workspace = MutationTestWorkspace::new();
let tool_id = tool.id().to_string();
let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
std::fs::write(&preserved, "preserved").unwrap();
let arguments = match tool_id.as_str() {
"copy_path" | "move_path" => serde_json::json!({
"source_path": preserved.to_string_lossy(),
"destination_path": destination.to_string_lossy(),
"dry_run": false
}),
"delete_path" => serde_json::json!({
"path": preserved.to_string_lossy(),
"recursive": false,
"dry_run": false
}),
_ => panic!("unsupported mutation tool: {}", tool_id),
};
let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
let builder = AgentBuilder::new()
.system_prompt("Test mutation denial.")
.llm(Arc::new(mock_with_response("done")))
.tool(tool)
.tool_security(ToolSecurityEngine::new(security));
let builder = match denial {
MutationDenial::Policy => builder,
MutationDenial::Approval => builder
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(Arc::new(RejectAllHandler::new())),
};
let agent = builder.build().unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
format!("{}-denied", tool_id),
tool_id.clone(),
arguments,
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(!record.executed, "{} must not be invoked", tool_id);
assert!(!record.success);
match denial {
MutationDenial::Policy => {
assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
assert!(record.approval.as_ref().is_some_and(|approval| matches!(
&approval.status,
ToolApprovalStatus::NotRequired
)));
}
MutationDenial::Approval => {
assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
assert!(record.approval.as_ref().is_some_and(|approval| matches!(
&approval.status,
ToolApprovalStatus::Rejected
)));
}
}
assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
assert!(!destination.exists());
}
fn recovery_manager_with_fallbacks(
fallbacks: impl IntoIterator<Item = (String, String)>,
) -> RecoveryManager {
use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
let per_tool = fallbacks
.into_iter()
.map(|(tool, fallback_tool)| {
(
tool,
ToolRetryConfig {
max_retries: 0,
timeout_ms: Some(1_000),
on_failure: ToolFailureAction::Fallback { fallback_tool },
},
)
})
.collect();
RecoveryManager::new(ErrorRecoveryConfig {
tools: ToolRecoveryConfig {
per_tool,
..Default::default()
},
..Default::default()
})
}
fn approval_check() -> HITLCheckResult {
HITLCheckResult::required(
ApprovalTrigger::tool("test", serde_json::json!({})),
HashMap::new(),
"Approve?",
None,
)
}
fn agent_with_approval_result(
raw_result: ApprovalResult,
timeout_action: TimeoutAction,
hooks: Arc<ApprovalRecordingHooks>,
) -> RuntimeAgent {
use ai_agents_hitl::{CallbackHandler, HITLConfig};
let config = HITLConfig {
on_timeout: timeout_action,
..Default::default()
};
let handler = CallbackHandler::new(move |_| raw_result.clone());
AgentBuilder::new()
.system_prompt("Test HITL hooks.")
.llm(Arc::new(mock_with_response("done")))
.build()
.unwrap()
.with_hooks(hooks)
.with_hitl(HITLEngine::new(config), Arc::new(handler))
}
#[tokio::test]
async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
let cases = vec![
(ApprovalResult::Approved, "approved"),
(
ApprovalResult::Rejected {
reason: Some("denied".to_string()),
},
"rejected",
),
(
ApprovalResult::Modified {
changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
},
"modified",
),
];
for (raw_result, expected) in cases {
let hooks = Arc::new(ApprovalRecordingHooks::new());
let agent =
agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
let result = agent.request_hitl_approval(approval_check()).await.unwrap();
assert_eq!(approval_result_name(&result), expected);
assert_correlated_approval_events(&hooks.events(), expected, expected);
}
}
#[tokio::test]
async fn approval_hooks_expose_timeout_policy_decisions() {
for (timeout_action, expected) in [
(TimeoutAction::Approve, "approved"),
(TimeoutAction::Reject, "rejected"),
] {
let hooks = Arc::new(ApprovalRecordingHooks::new());
let agent =
agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
let result = agent.request_hitl_approval(approval_check()).await.unwrap();
assert_eq!(approval_result_name(&result), expected);
assert_correlated_approval_events(&hooks.events(), "timeout", expected);
}
}
#[tokio::test]
async fn timeout_error_fires_correlated_resolved_error_before_returning() {
let hooks = Arc::new(ApprovalRecordingHooks::new());
let agent = agent_with_approval_result(
ApprovalResult::Timeout,
TimeoutAction::Error,
hooks.clone(),
);
let error = agent
.request_hitl_approval(approval_check())
.await
.unwrap_err();
assert!(error.to_string().contains("HITL approval timeout"));
assert_correlated_approval_events(&hooks.events(), "timeout", "error");
}
#[tokio::test]
async fn test_integration_yaml_to_chat_basic() {
let mock = mock_with_response("Hello! How can I help you?");
let agent = AgentBuilder::new()
.system_prompt("You are a test assistant.")
.llm(Arc::new(mock))
.build()
.unwrap();
let response = agent.chat("Hi").await.unwrap();
assert!(!response.content.is_empty());
assert_eq!(response.content, "Hello! How can I help you?");
}
#[tokio::test]
async fn stream_events_emit_one_authoritative_final_without_legacy_done() {
let agent = AgentBuilder::new()
.system_prompt("You are a test assistant.")
.llm(Arc::new(mock_with_response(
"Hello from the final response.",
)))
.build()
.unwrap();
let mut stream = agent.chat_stream_events("Hi").await.unwrap();
let mut final_responses = Vec::new();
let mut legacy_done = 0;
while let Some(event) = stream.next().await {
match event {
AgentStreamEvent::Chunk(StreamChunk::Done {}) => legacy_done += 1,
AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
panic!("unexpected stream error: {message}")
}
AgentStreamEvent::Final(response) => final_responses.push(response),
AgentStreamEvent::Chunk(_) => {}
}
}
assert_eq!(legacy_done, 0);
assert_eq!(final_responses.len(), 1);
let response = final_responses.pop().unwrap();
assert_eq!(response.content, "Hello from the final response.");
assert!(
response
.metadata
.as_ref()
.is_some_and(|metadata| { metadata.contains_key("reasoning") })
);
}
#[tokio::test]
async fn stream_final_content_includes_output_processing_after_provisional_chunks() {
let yaml = r#"
name: ProcessedStreamAgent
system_prompt: "Answer directly."
process:
output:
- type: format
config:
template: "{{ response }} [finalized]"
streaming:
enabled: true
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("provisional answer")))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let mut stream = agent.chat_stream_events("Hi").await.unwrap();
let mut provisional = String::new();
let mut final_content = None;
while let Some(event) = stream.next().await {
match event {
AgentStreamEvent::Chunk(StreamChunk::Content { text }) => {
provisional.push_str(&text)
}
AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
panic!("unexpected stream error: {message}")
}
AgentStreamEvent::Final(response) => final_content = Some(response.content),
AgentStreamEvent::Chunk(_) => {}
}
}
assert_eq!(provisional, "provisional answer");
assert_eq!(
final_content.as_deref(),
Some("provisional answer [finalized]")
);
}
#[tokio::test]
async fn stream_events_preserve_tool_progress_and_final_tool_calls() {
let agent = AgentBuilder::new()
.system_prompt("Use the echo tool once, then answer.")
.llm(Arc::new(mock_with_responses(vec![
r#"{"tool":"echo","arguments":{"message":"hello"}}"#,
"Echo completed.",
])))
.tool(Arc::new(ai_agents_tools::EchoTool::new()))
.build()
.unwrap();
let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
let mut starts = 0;
let mut results = 0;
let mut ends = 0;
let mut final_response = None;
while let Some(event) = stream.next().await {
match event {
AgentStreamEvent::Chunk(StreamChunk::ToolCallStart { name, .. }) => {
assert_eq!(name, "echo");
starts += 1;
}
AgentStreamEvent::Chunk(StreamChunk::ToolResult { name, success, .. }) => {
assert_eq!(name, "echo");
assert!(success);
results += 1;
}
AgentStreamEvent::Chunk(StreamChunk::ToolCallEnd { .. }) => ends += 1,
AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
panic!("unexpected stream error: {message}")
}
AgentStreamEvent::Final(response) => final_response = Some(response),
AgentStreamEvent::Chunk(_) => {}
}
}
assert_eq!((starts, results, ends), (1, 1, 1));
let response = final_response.expect("tool stream must finalize");
assert_eq!(response.content, "Echo completed.");
assert_eq!(
response.tool_calls.as_ref().map(|calls| calls
.iter()
.map(|call| call.name.as_str())
.collect::<Vec<_>>()),
Some(vec!["echo"])
);
}
#[tokio::test]
async fn legacy_stream_still_emits_one_done_chunk() {
let agent = AgentBuilder::new()
.system_prompt("You are a test assistant.")
.llm(Arc::new(mock_with_response(
"Hello from the legacy stream.",
)))
.build()
.unwrap();
let mut stream = agent.chat_stream("Hi").await.unwrap();
let mut done = 0;
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Done {} => done += 1,
StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
_ => {}
}
}
assert_eq!(done, 1);
}
#[tokio::test]
async fn test_integration_multi_turn_conversation() {
let mock = mock_with_responses(vec![
"Hello! I'm your assistant.",
"The weather is sunny today.",
"Goodbye!",
]);
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.build()
.unwrap();
let r1 = agent.chat("Hi").await.unwrap();
assert_eq!(r1.content, "Hello! I'm your assistant.");
let r2 = agent.chat("What's the weather?").await.unwrap();
assert_eq!(r2.content, "The weather is sunny today.");
let r3 = agent.chat("Bye").await.unwrap();
assert_eq!(r3.content, "Goodbye!");
let messages = agent.memory.get_messages(None).await.unwrap();
assert_eq!(messages.len(), 6);
}
#[test]
fn later_approval_preserves_modified_evidence() {
let arguments = serde_json::json!({"dry_run": true});
let mut record = Some(ToolApprovalRecord {
status: ToolApprovalStatus::Modified,
reason: None,
modified_arguments: Some(arguments.clone()),
});
merge_approved_record(&mut record);
let record = record.unwrap();
assert!(matches!(record.status, ToolApprovalStatus::Modified));
assert_eq!(record.modified_arguments, Some(arguments));
}
#[test]
fn approval_binding_rejects_replaced_tool_implementation() {
let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
let same_tool = Arc::clone(&reviewed_tool);
let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
let arguments = serde_json::json!({"path": "."});
let versions = ToolDecisionVersions {
policy: 2,
registry: 3,
runtime_control: 4,
state: Some(5),
};
let binding = ToolApprovalBinding {
canonical_id: "context_echo".to_string(),
arguments: arguments.clone(),
confirmation_required: true,
policy_version: versions.policy,
runtime_control_version: versions.runtime_control,
state_generation: versions.state,
reviewed_tool,
};
assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
assert!(binding.is_stale(
"context_echo",
&arguments,
true,
versions,
&replacement_tool,
));
}
#[tokio::test]
async fn approved_mutation_to_dry_run_remains_executable() {
use ai_agents_hitl::CallbackHandler;
let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
});
let agent = AgentBuilder::new()
.system_prompt("Test safer approval modifications.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(Arc::new(handler))
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"approved-dry-run",
"file_write",
serde_json::json!({
"path": "./approval-dry-run.txt",
"content": "not written"
}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.executed);
assert!(record.success);
assert_eq!(record.executed_arguments["dry_run"], true);
assert!(matches!(
record.approval.as_ref().map(|approval| &approval.status),
Some(ToolApprovalStatus::Modified)
));
let output: Value = serde_json::from_str(&record.output).unwrap();
assert_eq!(output["mutation_performed"], false);
}
#[tokio::test]
async fn shared_executor_approval_reaches_web_fetch_transport() {
use ai_agents_hitl::{CallbackHandler, HITLConfig};
use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let tool = WebFetchTool::with_transport_and_resolver(
Arc::new(RuntimeWebFetchTransport {
calls: Arc::clone(&calls),
}),
Arc::new(RuntimeWebFetchResolver),
);
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
security.tools.insert(
"web_fetch".to_string(),
ToolPolicyConfig {
domains: DomainPolicyConfig {
requires_approval: vec!["approval.test".to_string()],
..Default::default()
},
allowed_schemes: vec!["https".to_string()],
allowed_ports: vec![443],
..Default::default()
},
);
let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
let agent = AgentBuilder::new()
.system_prompt("Test approved web fetch execution.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(tool))
.tool_security(ToolSecurityEngine::new(security))
.build()
.unwrap()
.with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"approved-web-fetch",
"web_fetch",
serde_json::json!({
"url": "https://approval.test/page",
"cache_ttl_seconds": 0
}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.success);
assert!(
record
.approval
.as_ref()
.is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
);
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn context_preserves_requested_and_canonical_identity() {
let mock = mock_with_response("hello");
let mut tools = ai_agents_tools::ToolRegistry::new();
tools.register(Arc::new(ContextEchoTool)).unwrap();
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
let mut policy = ai_agents_tools::ToolPolicyConfig {
read_paths: vec![".".to_string()],
max_results: Some(7),
..Default::default()
};
policy
.config
.insert("backend".to_string(), serde_json::json!("memory"));
security.tools.insert("context_echo".to_string(), policy);
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.tools(tools)
.tool_security(ToolSecurityEngine::new(security))
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"ctx-call",
"Context Echo",
serde_json::json!({"path": ".", "max_results": 99}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.success);
assert!(matches!(&record.source, ToolCallSource::Manual));
assert_eq!(record.requested_name, "Context Echo");
assert_eq!(record.canonical_id, "context_echo");
assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
assert_eq!(record.executed_arguments["max_results"], 7);
let output: Value = serde_json::from_str(&record.output).unwrap();
assert_eq!(output["requested_name"], "Context Echo");
assert_eq!(output["canonical_id"], "context_echo");
assert_eq!(output["max_results"], 7);
assert_eq!(output["custom_config"]["backend"], "memory");
assert!(record.metadata.contains_key("effective_limits"));
assert!(record.metadata.contains_key("policy_snapshot"));
}
#[tokio::test]
async fn test_runtime_control_cancels_active_tool_call() {
let mock = mock_with_response("hello");
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.tool(Arc::new(SlowTool))
.build()
.unwrap(),
);
let control = agent.runtime_control();
let running_agent = Arc::clone(&agent);
let handle = tokio::spawn(async move {
running_agent
.invoke_tool(ToolExecutionRequest::new(
"slow-call",
"slow",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap()
});
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
control.cancel_all();
let record = handle.await.unwrap();
assert!(record.executed);
assert!(record.cancelled);
assert!(!record.success);
assert!(record.cancellation_reason.is_some());
}
#[tokio::test]
async fn cancelled_tool_does_not_enter_fallback() {
let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test cancellation before fallback.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(SlowTool))
.tool(Arc::new(RecoveryTestTool {
id: "fallback".to_string(),
succeeds: true,
calls: Arc::clone(&fallback_calls),
max_output_chars: None,
}))
.recovery_manager(recovery_manager_with_fallbacks([(
"slow".to_string(),
"fallback".to_string(),
)]))
.build()
.unwrap(),
);
let control = agent.runtime_control();
let running_agent = Arc::clone(&agent);
let handle = tokio::spawn(async move {
running_agent
.invoke_tool(ToolExecutionRequest::new(
"cancelled-fallback-call",
"slow",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap()
});
tokio::time::sleep(Duration::from_millis(100)).await;
control.cancel_all();
let record = handle.await.unwrap();
assert!(record.executed);
assert!(record.cancelled);
assert!(!record.success);
assert_eq!(record.canonical_id, "slow");
assert_eq!(fallback_calls.load(Ordering::SeqCst), 0);
assert_eq!(agent.tool_call_history().len(), 1);
}
#[tokio::test]
async fn non_idempotent_tool_calls_are_not_retried() {
use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
let mock = mock_with_response("hello");
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.tool(Arc::new(FlakyWriteTool {
calls: Arc::clone(&calls),
}))
.recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
tools: ToolRecoveryConfig {
default: ToolRetryConfig {
max_retries: 2,
..Default::default()
},
..Default::default()
},
..Default::default()
}))
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"flaky-call",
"flaky_write",
serde_json::json!({"path": "./tmp.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(!record.success);
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn safely_retryable_tool_receives_a_fresh_deadline_per_attempt() {
use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let deadlines = Arc::new(parking_lot::Mutex::new(Vec::new()));
let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
let agent = AgentBuilder::new()
.system_prompt("Test retry deadlines.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(RetryDeadlineTool {
calls: Arc::clone(&calls),
deadlines: Arc::clone(&deadlines),
remaining_ms: Arc::clone(&remaining_ms),
}))
.recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
tools: ToolRecoveryConfig {
per_tool: HashMap::from([(
"retry_deadline".to_string(),
ToolRetryConfig {
max_retries: 1,
..Default::default()
},
)]),
..Default::default()
},
..Default::default()
}))
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"retry-deadline-call",
"retry_deadline",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.executed);
assert!(record.success);
assert_eq!(calls.load(Ordering::SeqCst), 2);
let deadlines = deadlines.lock();
assert_eq!(deadlines.len(), 2);
assert!(
deadlines[1] > deadlines[0],
"retry inherited the first invocation deadline"
);
let remaining_ms = remaining_ms.lock();
assert_eq!(remaining_ms.len(), 2);
assert!(
remaining_ms
.iter()
.all(|remaining| (800..=1_000).contains(remaining))
);
}
#[tokio::test]
async fn call_classification_timeout_controls_deadline_and_timer() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
let agent = AgentBuilder::new()
.system_prompt("Test call-level timeout.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(ClassifiedTimeoutTool {
id: "classified_timeout",
calls: Arc::clone(&calls),
timeout_ms: 100,
sleep_ms: 150,
requires_approval: false,
remaining_ms: Arc::clone(&remaining_ms),
}))
.build()
.unwrap();
let started = Instant::now();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"classified-timeout-call",
"classified_timeout",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.executed);
assert!(record.timed_out);
assert!(!record.success);
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert!(started.elapsed() < Duration::from_secs(1));
let remaining_ms = remaining_ms.lock();
assert_eq!(remaining_ms.len(), 1);
assert!((1..=100).contains(&remaining_ms[0]));
}
#[tokio::test]
async fn recovery_timeout_only_lowers_call_and_policy_timeouts() {
use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
let agent = AgentBuilder::new()
.system_prompt("Test recovery timeout.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(ClassifiedTimeoutTool {
id: "recovery_timeout",
calls: Arc::clone(&calls),
timeout_ms: 1_000,
sleep_ms: 150,
requires_approval: false,
remaining_ms: Arc::clone(&remaining_ms),
}))
.recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
tools: ToolRecoveryConfig {
per_tool: HashMap::from([(
"recovery_timeout".to_string(),
ToolRetryConfig {
timeout_ms: Some(100),
..Default::default()
},
)]),
..Default::default()
},
..Default::default()
}))
.build()
.unwrap();
let started = Instant::now();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"recovery-timeout-call",
"recovery_timeout",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.executed);
assert!(record.timed_out);
assert!(!record.success);
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert!(started.elapsed() < Duration::from_secs(1));
assert_eq!(record.metadata["effective_limits"]["timeout_ms"], 100);
let remaining_ms = remaining_ms.lock();
assert_eq!(remaining_ms.len(), 1);
assert!((1..=100).contains(&remaining_ms[0]));
}
#[tokio::test]
async fn recovery_default_timeout_controls_deadline_and_timer() {
use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
let agent = AgentBuilder::new()
.system_prompt("Test default recovery timeout.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(ClassifiedTimeoutTool {
id: "default_recovery_timeout",
calls: Arc::clone(&calls),
timeout_ms: 1_000,
sleep_ms: 150,
requires_approval: false,
remaining_ms: Arc::clone(&remaining_ms),
}))
.recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
tools: ToolRecoveryConfig {
default: ToolRetryConfig {
timeout_ms: Some(100),
..Default::default()
},
..Default::default()
},
..Default::default()
}))
.build()
.unwrap();
let started = Instant::now();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"default-recovery-timeout-call",
"default_recovery_timeout",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.executed);
assert!(record.timed_out);
assert!(!record.success);
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert!(started.elapsed() < Duration::from_secs(1));
assert_eq!(record.metadata["effective_limits"]["timeout_ms"], 100);
let remaining_ms = remaining_ms.lock();
assert_eq!(remaining_ms.len(), 1);
assert!((1..=100).contains(&remaining_ms[0]));
}
#[test]
fn recovery_timeout_cannot_widen_security_baseline() {
let security_engine = ToolSecurityEngine::new(ToolSecurityConfig {
default_timeout_ms: 100,
..Default::default()
});
let safety = ToolSafetyMetadata::compute();
let mut classification = ToolCallClassification::from_metadata(&safety);
classification.timeout_ms = Some(500);
let (limits, timeout) = RuntimeAgent::effective_tool_limits(
&security_engine,
"recovery_cannot_widen",
&safety,
&classification,
Some(1_000),
)
.unwrap();
assert_eq!(limits.timeout_ms, Some(100));
assert_eq!(timeout.timer, Duration::from_millis(100));
}
#[tokio::test]
async fn invalid_call_timeout_stops_before_approval_or_tool_invocation() {
let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
let mut security = ToolSecurityConfig {
enabled: true,
..Default::default()
};
security.tools.insert(
"invalid_call_timeout".to_string(),
ai_agents_tools::ToolPolicyConfig {
require_confirmation: true,
..Default::default()
},
);
let agent = AgentBuilder::new()
.system_prompt("Test invalid call timeout.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(ClassifiedTimeoutTool {
id: "invalid_call_timeout",
calls: Arc::clone(&tool_calls),
timeout_ms: u64::MAX,
sleep_ms: 0,
requires_approval: false,
remaining_ms,
}))
.tool_security(ToolSecurityEngine::new(security))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(Arc::new(CountingApprovalHandler {
calls: Arc::clone(&approval_calls),
}))
.build()
.unwrap();
let error = agent
.invoke_tool(ToolExecutionRequest::new(
"invalid-call-timeout",
"invalid_call_timeout",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap_err();
assert!(error.to_string().contains(
"effective tool timeout_ms must be no greater than 3153600000000000 milliseconds"
));
assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn invalid_modified_call_timeout_stops_before_lock_or_invocation() {
use ai_agents_hitl::CallbackHandler;
let blocker_gate = PathMutationGate::new();
let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
changes: HashMap::from([("invalid_timeout".to_string(), Value::Bool(true))]),
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test final call timeout validation.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(BlockingPathMutationTool {
id: "timeout_lock_blocker",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: blocker_gate.clone(),
}))
.tool(Arc::new(ApprovalModifiedTimeoutTool {
calls: Arc::clone(&tool_calls),
}))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(Arc::new(handler))
.hooks(hooks.clone())
.build()
.unwrap(),
);
let blocking_agent = Arc::clone(&agent);
let blocker = tokio::spawn(async move {
blocking_agent
.invoke_tool(ToolExecutionRequest::new(
"timeout-lock-blocker",
"timeout_lock_blocker",
serde_json::json!({"path": "./shared-timeout.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
});
blocker_gate.wait_until_entered().await;
let record = tokio::time::timeout(
Duration::from_millis(500),
agent.invoke_tool(ToolExecutionRequest::new(
"invalid-modified-timeout",
"approval_modified_timeout",
serde_json::json!({
"path": "./shared-timeout.txt",
"invalid_timeout": false
}),
ToolCallSource::Manual,
)),
)
.await
.expect("final timeout validation must not wait for the held path lock")
.unwrap();
blocker_gate.release();
assert!(blocker.await.unwrap().success);
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
assert!(record.output.contains(
"effective tool timeout_ms must be no greater than 3153600000000000 milliseconds"
));
assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
let invalid_request_events = hooks
.events()
.into_iter()
.filter(|event| event.contains("approval_modified_timeout") || event == "error")
.collect::<Vec<_>>();
assert_eq!(
invalid_request_events,
vec![
"start:approval_modified_timeout",
"complete:approval_modified_timeout:false",
"record:approval_modified_timeout:false",
"error"
]
);
}
#[tokio::test]
async fn side_effecting_tools_are_serialized_per_resource() {
let mock = mock_with_response("hello");
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.tool(Arc::new(LockedWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.build()
.unwrap(),
);
let left = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"lock-1",
"locked_write",
serde_json::json!({"path": "./same.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
let right = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"lock-2",
"locked_write",
serde_json::json!({"path": "./same.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
let left = left.await.unwrap();
let right = right.await.unwrap();
assert!(left.success);
assert!(right.success);
assert_eq!(max_active.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn path_resources_use_shared_global_lock_and_cleanup() {
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let bindings = ai_agents_core::ToolPolicyBindings {
path_fields: vec![
ai_agents_core::PathPolicyBinding::read_write("source_path"),
ai_agents_core::PathPolicyBinding::write("destination_path"),
],
..Default::default()
};
let classification = ai_agents_core::ToolCallClassification::from_metadata(
&MultiResourceWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}
.safety_metadata(),
);
let left_args = serde_json::json!({
"source_path": "./a/../first.txt",
"destination_path": "./second.txt"
});
let right_args = serde_json::json!({
"source_path": "./second.txt",
"destination_path": "./first.txt"
});
let left_keys = tool_resource_lock_keys(
"multi_resource_write",
&left_args,
&bindings,
&classification,
);
let right_keys = tool_resource_lock_keys(
"multi_resource_write",
&right_args,
&bindings,
&classification,
);
assert_eq!(left_keys, right_keys);
assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
let locks = new_tool_resource_locks();
let build_agent = || {
AgentBuilder::new()
.system_prompt("Test shared resource locks.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(MultiResourceWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.build()
.unwrap()
.with_shared_resource_locks(Arc::clone(&locks))
};
let left_agent = Arc::new(build_agent());
let right_agent = Arc::new(build_agent());
let left = tokio::spawn(async move {
left_agent
.invoke_tool(ToolExecutionRequest::new(
"multi-left",
"multi_resource_write",
left_args,
ToolCallSource::Manual,
))
.await
.unwrap()
});
let right = tokio::spawn(async move {
right_agent
.invoke_tool(ToolExecutionRequest::new(
"multi-right",
"multi_resource_write",
right_args,
ToolCallSource::Manual,
))
.await
.unwrap()
});
let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
tokio::join!(left, right)
})
.await
.expect("reversed resource acquisition must not deadlock");
assert!(left.unwrap().success);
assert!(right.unwrap().success);
assert_eq!(max_active.load(Ordering::SeqCst), 1);
assert!(locks.read().is_empty());
}
#[tokio::test]
async fn global_path_lock_serializes_copy_destination_with_file_write() {
assert_path_mutation_pair_serialized(
"copy_path",
CopyPathTool::new().policy_bindings().path_fields,
serde_json::json!({
"source_path": "./source.txt",
"destination_path": "./shared.txt"
}),
"file_write",
FileWriteTool::new().policy_bindings().path_fields,
serde_json::json!({"path": "./shared.txt"}),
)
.await;
}
#[tokio::test]
async fn parent_and_spawned_runtime_share_global_path_lock() {
let workspace = MutationTestWorkspace::new();
let destination = workspace.root.join("spawned.txt");
let parent_gate = PathMutationGate::new();
let parent = Arc::new(
AgentBuilder::from_yaml(
r#"
name: LockParent
system_prompt: parent
llm:
default: default
tools:
- parent_path_write
spawner:
shared_llms: true
"#,
)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.auto_configure_spawner()
.await
.unwrap()
.tool(Arc::new(BlockingPathMutationTool {
id: "parent_path_write",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: parent_gate.clone(),
}))
.build()
.unwrap(),
);
let mut child_spec = crate::spec::AgentSpec {
name: "LockChild".to_string(),
system_prompt: "child".to_string(),
tools: Some(vec![crate::spec::ToolEntry::Simple(
"file_write".to_string(),
)]),
..Default::default()
};
child_spec.tool_security.enabled = true;
child_spec.tool_security.fail_closed = true;
let file_write_policy = ai_agents_tools::ToolPolicyConfig {
write_paths: vec![workspace.root.to_string_lossy().into_owned()],
allow_without_confirmation: true,
..Default::default()
};
child_spec
.tool_security
.tools
.insert("file_write".to_string(), file_write_policy);
let spawned = parent
.spawner()
.unwrap()
.spawn_from_spec(child_spec)
.await
.unwrap();
assert!(Arc::ptr_eq(
&parent.resource_locks,
&spawned.agent.resource_locks
));
assert!(!Arc::ptr_eq(
&parent.runtime_control,
&spawned.agent.runtime_control
));
let parent_call = {
let parent = Arc::clone(&parent);
let destination = destination.clone();
tokio::spawn(async move {
parent
.invoke_tool(ToolExecutionRequest::new(
"parent-lock-holder",
"parent_path_write",
serde_json::json!({"path": destination}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
parent_gate.wait_until_entered().await;
let child_call = {
let child = Arc::clone(&spawned.agent);
let destination = destination.clone();
tokio::spawn(async move {
child
.invoke_tool(ToolExecutionRequest::new(
"spawned-file-write",
"file_write",
serde_json::json!({
"path": destination,
"content": "spawned",
"dry_run": false
}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
assert!(!child_call.is_finished());
parent_gate.release();
let (parent_record, child_record) =
tokio::time::timeout(std::time::Duration::from_secs(2), async {
tokio::join!(parent_call, child_call)
})
.await
.expect("parent and spawned path mutations did not finish");
assert!(parent_record.unwrap().success);
assert!(child_record.unwrap().success);
assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
assert!(parent.resource_locks.read().is_empty());
}
#[tokio::test]
async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
let locks = new_tool_resource_locks();
let holder_gate = PathMutationGate::new();
let waiter_gate = PathMutationGate::new();
waiter_gate.release();
let holder = Arc::new(
AgentBuilder::new()
.system_prompt("Hold the global path lock.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(BlockingPathMutationTool {
id: "holder_write",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: holder_gate.clone(),
}))
.build()
.unwrap()
.with_shared_resource_locks(Arc::clone(&locks)),
);
let waiter = Arc::new(
AgentBuilder::new()
.system_prompt("Wait for the global path lock.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(BlockingPathMutationTool {
id: "waiter_write",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: waiter_gate.clone(),
}))
.build()
.unwrap()
.with_shared_resource_locks(Arc::clone(&locks)),
);
let holder_call = {
let holder = Arc::clone(&holder);
tokio::spawn(async move {
holder
.invoke_tool(ToolExecutionRequest::new(
"holder-call",
"holder_write",
serde_json::json!({"path": "./shared.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
holder_gate.wait_until_entered().await;
let waiter_call = {
let waiter = Arc::clone(&waiter);
tokio::spawn(async move {
waiter
.invoke_tool(ToolExecutionRequest::new(
"waiter-call",
"waiter_write",
serde_json::json!({"path": "./shared.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
wait_for_resource_lock_strong_count(&locks, 2).await;
waiter.runtime_control().cancel_all();
let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
.await
.expect("cancelled lock waiter did not finish")
.unwrap();
assert!(!waiter_record.success);
assert!(!waiter_record.executed);
assert!(waiter_record.cancelled);
assert_eq!(
waiter_record.cancellation_reason.as_deref(),
Some("runtime control cancellation")
);
assert!(!waiter_gate.entered.load(Ordering::SeqCst));
assert_eq!(
locks
.read()
.get("path-mutation:global")
.map_or(0, |lock| lock.strong_count()),
1
);
holder_gate.release();
let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
.await
.expect("lock holder did not finish")
.unwrap();
assert!(holder_record.success);
assert!(locks.read().is_empty());
}
#[tokio::test]
async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
for denial in [MutationDenial::Policy, MutationDenial::Approval] {
let tools: [Arc<dyn Tool>; 3] = [
Arc::new(CopyPathTool::new()),
Arc::new(MovePathTool::new()),
Arc::new(DeletePathTool::new()),
];
for tool in tools {
assert_path_mutation_denied(tool, denial).await;
}
}
}
#[tokio::test]
async fn policy_denial_keeps_executor_hook_lifecycle_and_record_authority() {
let workspace = MutationTestWorkspace::new();
let target = workspace.root.join("denied.txt");
let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
let agent = AgentBuilder::new()
.system_prompt("Test denied tool hooks.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(FileWriteTool::new()))
.tool_security(ToolSecurityEngine::new(mutation_denial_security_config(
"file_write",
&workspace.root,
MutationDenial::Policy,
)))
.hooks(hooks.clone())
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"denied-hook-call",
"file_write",
serde_json::json!({
"path": target.to_string_lossy(),
"content": "blocked"
}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
assert_eq!(
hooks.events(),
vec![
"start:file_write",
"complete:file_write:false",
"record:file_write:false",
"error"
]
);
assert!(!target.exists());
}
#[tokio::test]
async fn approval_argument_changes_are_rechecked_against_final_scope() {
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let entered = Arc::new(tokio::sync::Barrier::new(2));
let release = Arc::new(tokio::sync::Notify::new());
let handler = Arc::new(BlockingApprovalHandler {
entered: Arc::clone(&entered),
release: Arc::clone(&release),
result: ApprovalResult::Modified {
changes: HashMap::from([(
"path".to_string(),
Value::String("./after-approval.txt".to_string()),
)]),
},
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test final scope validation.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(LockedWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.tool_security(ToolSecurityEngine::new(approval_security_config(true)))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(handler)
.build()
.unwrap(),
);
let control = agent.runtime_control();
let running = Arc::clone(&agent);
let call = tokio::spawn(async move {
running
.invoke_tool(ToolExecutionRequest::new(
"approval-scope",
"locked_write",
serde_json::json!({"path": "./before-approval.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
});
entered.wait().await;
let expected_version = control.set_tool_scope(Vec::new());
release.notify_one();
let record = call.await.unwrap();
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.runtime_config_version, expected_version);
assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
assert_eq!(max_active.load(Ordering::SeqCst), 0);
assert_eq!(
record.metadata["runtime_scope_snapshot"],
serde_json::json!([])
);
}
#[tokio::test]
async fn approval_is_rechecked_against_final_policy_snapshot() {
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let entered = Arc::new(tokio::sync::Barrier::new(2));
let release = Arc::new(tokio::sync::Notify::new());
let handler = Arc::new(BlockingApprovalHandler {
entered: Arc::clone(&entered),
release: Arc::clone(&release),
result: ApprovalResult::Approved,
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test final policy validation.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(LockedWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.tool_security(ToolSecurityEngine::new(approval_security_config(true)))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(handler)
.build()
.unwrap(),
);
let control = agent.runtime_control();
let running = Arc::clone(&agent);
let call = tokio::spawn(async move {
running
.invoke_tool(ToolExecutionRequest::new(
"approval-policy",
"locked_write",
serde_json::json!({"path": "./policy.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
});
entered.wait().await;
let expected_version = control.set_tool_security(approval_security_config(false));
release.notify_one();
let record = call.await.unwrap();
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.runtime_config_version, expected_version);
assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
assert_eq!(max_active.load(Ordering::SeqCst), 0);
assert!(record.metadata.contains_key("policy_snapshot"));
}
#[test]
fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
let agent = AgentBuilder::new()
.system_prompt("Test runtime policy validation.")
.llm(Arc::new(mock_with_response("done")))
.build()
.unwrap();
let control = agent.runtime_control();
let mut valid = ToolSecurityConfig::default();
valid.tools.insert(
"web_search".to_string(),
ai_agents_tools::ToolPolicyConfig {
max_results: Some(5),
..Default::default()
},
);
let generation = control.try_set_tool_security(valid).unwrap();
let mut invalid = ToolSecurityConfig::default();
invalid.tools.insert(
"web_search".to_string(),
ai_agents_tools::ToolPolicyConfig {
max_results: Some(0),
..Default::default()
},
);
let error = control.try_set_tool_security(invalid).unwrap_err();
assert!(
error
.to_string()
.contains("max_results must be greater than 0")
);
assert_eq!(control.version(), generation);
assert_eq!(
control
.state
.tool_security_override
.read()
.as_ref()
.unwrap()
.config()
.tools["web_search"]
.max_results,
Some(5)
);
}
#[test]
fn invalid_timeout_config_stops_before_approval_or_tool_invocation() {
let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let spec = crate::spec::AgentSpec {
tool_security: ToolSecurityConfig {
enabled: true,
default_timeout_ms: u64::MAX,
..Default::default()
},
..Default::default()
};
let result = AgentBuilder::from_spec(spec)
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(FlakyWriteTool {
calls: Arc::clone(&tool_calls),
}))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(Arc::new(CountingApprovalHandler {
calls: Arc::clone(&approval_calls),
}))
.build();
assert!(result.is_err());
assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
}
#[test]
fn invalid_recovery_timeout_config_stops_before_approval_or_tool_invocation() {
use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let spec = crate::spec::AgentSpec {
error_recovery: ErrorRecoveryConfig {
tools: ToolRecoveryConfig {
default: ToolRetryConfig {
timeout_ms: Some(u64::MAX),
..Default::default()
},
..Default::default()
},
..Default::default()
},
..Default::default()
};
let result = AgentBuilder::from_spec(spec)
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(FlakyWriteTool {
calls: Arc::clone(&tool_calls),
}))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(Arc::new(CountingApprovalHandler {
calls: Arc::clone(&approval_calls),
}))
.build();
assert!(result.is_err());
assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
}
#[test]
fn invalid_timeout_policy_does_not_replace_snapshot_or_generation() {
let agent = AgentBuilder::new()
.system_prompt("Test runtime timeout policy validation.")
.llm(Arc::new(mock_with_response("done")))
.build()
.unwrap();
let control = agent.runtime_control();
let valid = ToolSecurityConfig {
default_timeout_ms: 5_000,
..Default::default()
};
let generation = control.try_set_tool_security(valid).unwrap();
let invalid = ToolSecurityConfig {
default_timeout_ms: MAX_TOOL_TIMEOUT_MS + 1,
..Default::default()
};
let error = control.try_set_tool_security(invalid).unwrap_err();
assert!(error.to_string().contains(&format!(
"tool_security.default_timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
)));
assert_eq!(control.version(), generation);
assert_eq!(
control
.state
.tool_security_override
.read()
.as_ref()
.unwrap()
.config()
.default_timeout_ms,
5_000
);
}
#[test]
fn runtime_tool_timeout_conversion_enforces_the_stable_boundary() {
let timeout = RuntimeAgent::validated_tool_timeout(MAX_TOOL_TIMEOUT_MS).unwrap();
assert_eq!(timeout.timer, Duration::from_millis(MAX_TOOL_TIMEOUT_MS));
assert_eq!(
timeout.deadline_delta,
chrono::Duration::milliseconds(MAX_TOOL_TIMEOUT_MS as i64)
);
for timeout_ms in [MAX_TOOL_TIMEOUT_MS + 1, u64::MAX] {
let error = RuntimeAgent::validated_tool_timeout(timeout_ms).unwrap_err();
assert!(error.to_string().contains(&format!(
"effective tool timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
)));
}
}
#[tokio::test]
async fn persistent_override_preserves_rate_history_within_generation() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let agent = AgentBuilder::new()
.system_prompt("Test persistent policy overrides.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(RecoveryTestTool {
id: "limited_override".to_string(),
succeeds: true,
calls: Arc::clone(&calls),
max_output_chars: None,
}))
.build()
.unwrap();
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
let policy = ai_agents_tools::ToolPolicyConfig {
write_paths: vec![".".to_string()],
rate_limit: Some(1),
..Default::default()
};
security
.tools
.insert("limited_override".to_string(), policy);
let generation = agent.runtime_control().set_tool_security(security);
let first = agent
.invoke_tool(ToolExecutionRequest::new(
"limited-first",
"limited_override",
serde_json::json!({"path": "./limited.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap();
let second = agent
.invoke_tool(ToolExecutionRequest::new(
"limited-second",
"limited_override",
serde_json::json!({"path": "./limited.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(first.success);
assert_eq!(first.policy_version, generation);
assert!(!second.executed);
assert!(second.output.contains("Rate limit exceeded"));
assert_eq!(second.policy_version, generation);
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn concurrent_rate_admission_consumes_capacity_atomically() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let tool = Arc::new(RecoveryTestTool {
id: "atomic_rate".to_string(),
succeeds: true,
calls: Arc::clone(&calls),
max_output_chars: None,
});
let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
let bindings = tool.policy_bindings();
let classification = tool.classify_call(&arguments);
let resource_keys =
tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
let policy = ai_agents_tools::ToolPolicyConfig {
write_paths: vec![".".to_string()],
rate_limit: Some(1),
..Default::default()
};
security.tools.insert(tool.id().to_string(), policy);
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test atomic rate admission.")
.llm(Arc::new(mock_with_response("done")))
.tool(tool)
.tool_security(ToolSecurityEngine::new(security))
.build()
.unwrap(),
);
let held = agent
.acquire_tool_resource_locks(&resource_keys)
.await
.unwrap();
let left = {
let agent = Arc::clone(&agent);
let arguments = arguments.clone();
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"atomic-rate-left",
"atomic_rate",
arguments,
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
let right = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"atomic-rate-right",
"atomic_rate",
arguments,
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
drop(held);
let (left, right) = tokio::join!(left, right);
let records = [left.unwrap(), right.unwrap()];
assert_eq!(records.iter().filter(|record| record.success).count(), 1);
assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
assert!(
records.iter().any(|record| {
!record.executed && record.output.contains("Rate limit exceeded")
})
);
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn changed_policy_generation_invalidates_pending_approval() {
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let entered = Arc::new(tokio::sync::Barrier::new(2));
let release = Arc::new(tokio::sync::Notify::new());
let handler = Arc::new(BlockingApprovalHandler {
entered: Arc::clone(&entered),
release: Arc::clone(&release),
result: ApprovalResult::Approved,
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test stale approval denial.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(LockedWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.tool_security(ToolSecurityEngine::new(approval_security_config(true)))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(handler)
.build()
.unwrap(),
);
let running = Arc::clone(&agent);
let call = tokio::spawn(async move {
running
.invoke_tool(ToolExecutionRequest::new(
"stale-approval",
"locked_write",
serde_json::json!({"path": "./stale.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
});
entered.wait().await;
let generation = agent
.runtime_control()
.set_tool_security(approval_security_config(true));
release.notify_one();
let record = call.await.unwrap();
assert!(!record.executed);
assert!(record.output.contains("Approval became stale"));
assert_eq!(record.policy_version, generation);
assert_eq!(max_active.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn final_policy_reapplies_argument_caps_after_approval_changes() {
use ai_agents_hitl::CallbackHandler;
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
let policy = ai_agents_tools::ToolPolicyConfig {
read_paths: vec![".".to_string()],
max_results: Some(5),
require_confirmation: true,
..Default::default()
};
security.tools.insert("context_echo".to_string(), policy);
let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
});
let agent = AgentBuilder::new()
.system_prompt("Test final argument caps.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(ContextEchoTool))
.tool_security(ToolSecurityEngine::new(security))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(Arc::new(handler))
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"final-cap",
"context_echo",
serde_json::json!({"path": ".", "max_results": 1}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(record.success);
assert_eq!(record.executed_arguments["max_results"], 5);
assert_eq!(
record.approval.unwrap().modified_arguments.unwrap()["max_results"],
5
);
}
#[tokio::test]
async fn no_binding_writes_use_canonical_fallback_lock() {
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test fallback resource locks.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(NoBindingWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.build()
.unwrap(),
);
let left = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"no-binding-left",
"no_binding_write",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
let right = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"no-binding-right",
"no_binding_write",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
let (left, right) = tokio::join!(left, right);
assert!(left.unwrap().success);
assert!(right.unwrap().success);
assert_eq!(max_active.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn parent_and_child_paths_share_a_resource_lock() {
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test parent child resource locks.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(LockedWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.build()
.unwrap(),
);
let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
let child = format!("{}/child.txt", parent);
let left = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"parent-lock",
"locked_write",
serde_json::json!({"path": parent}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
let right = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"child-lock",
"locked_write",
serde_json::json!({"path": child}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
let (left, right) = tokio::join!(left, right);
assert!(left.unwrap().success);
assert!(right.unwrap().success);
assert_eq!(max_active.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hooks = Arc::new(ReentrantToolHooks {
agent: parking_lot::Mutex::new(None),
invoked: AtomicBool::new(false),
nested_success: AtomicBool::new(false),
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test hook reentrancy.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(RecoveryTestTool {
id: "reentrant_write".to_string(),
succeeds: true,
calls: Arc::clone(&calls),
max_output_chars: None,
}))
.hooks(hooks.clone())
.build()
.unwrap(),
);
*hooks.agent.lock() = Some(Arc::downgrade(&agent));
let record = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.invoke_tool(ToolExecutionRequest::new(
"outer-hook-call",
"reentrant_write",
serde_json::json!({"path": "./hook.txt"}),
ToolCallSource::Manual,
)),
)
.await
.expect("tool completion hook must not retain resource guards")
.unwrap();
assert!(record.success);
assert!(hooks.nested_success.load(Ordering::SeqCst));
assert_eq!(calls.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn fallback_finalizes_original_record_before_shared_execution() {
let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
let agent = AgentBuilder::new()
.system_prompt("Test fallback execution.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(RecoveryTestTool {
id: "primary".to_string(),
succeeds: false,
calls: Arc::clone(&primary_calls),
max_output_chars: None,
}))
.tool(Arc::new(RecoveryTestTool {
id: "fallback".to_string(),
succeeds: true,
calls: Arc::clone(&fallback_calls),
max_output_chars: None,
}))
.recovery_manager(recovery_manager_with_fallbacks([(
"primary".to_string(),
"fallback".to_string(),
)]))
.hooks(hooks.clone())
.build()
.unwrap();
let record = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.invoke_tool(ToolExecutionRequest::new(
"fallback-call",
"primary",
serde_json::json!({"path": "./shared.txt"}),
ToolCallSource::Manual,
)),
)
.await
.expect("fallback must not retain the primary resource guard")
.unwrap();
assert_eq!(
hooks.events(),
vec![
"start:primary",
"complete:primary:false",
"record:primary:true",
"error",
"start:fallback",
"complete:fallback:true",
"record:fallback:true",
]
);
let records = hooks.records();
assert_eq!(records.len(), 2);
let original = &records[0];
assert_eq!(original.canonical_id, "primary");
assert!(matches!(original.source, ToolCallSource::Manual));
assert!(original.executed);
assert!(!original.success);
let fallback = &records[1];
assert_eq!(fallback.canonical_id, "fallback");
assert_eq!(fallback.call_id, "fallback-call");
assert!(matches!(
&fallback.source,
ToolCallSource::Fallback { original_tool } if original_tool == "primary"
));
assert!(fallback.executed);
assert!(fallback.success);
assert_eq!(record.canonical_id, fallback.canonical_id);
assert_eq!(record.output, fallback.output);
let history = agent.tool_call_history();
assert_eq!(
history
.iter()
.map(|entry| entry.tool_id.as_str())
.collect::<Vec<_>>(),
vec!["primary", "fallback"]
);
assert_eq!(history[0].result.get("success"), Some(&Value::Bool(false)));
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn self_fallback_cycle_is_denied_before_reinvocation() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
let agent = AgentBuilder::new()
.system_prompt("Test self-fallback cycle admission.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(RecoveryTestTool {
id: "primary".to_string(),
succeeds: false,
calls: Arc::clone(&calls),
max_output_chars: None,
}))
.recovery_manager(recovery_manager_with_fallbacks([(
"primary".to_string(),
"primary".to_string(),
)]))
.hooks(hooks.clone())
.build()
.unwrap();
let record = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.invoke_tool(ToolExecutionRequest::new(
"self-fallback-call",
"primary",
serde_json::json!({"path": "./shared.txt"}),
ToolCallSource::Manual,
)),
)
.await
.expect("self fallback must terminate without recursive execution")
.unwrap();
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert_eq!(record.canonical_id, "primary");
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
assert!(record.output.contains("fallback cycle"));
assert!(matches!(
record.source,
ToolCallSource::Fallback { ref original_tool } if original_tool == "primary"
));
assert_eq!(
record.metadata.get("fallback_chain"),
Some(&serde_json::json!(["primary"]))
);
assert_eq!(
hooks.events(),
vec![
"start:primary",
"complete:primary:false",
"record:primary:true",
"error",
"complete:primary:false",
"record:primary:false",
"error",
]
);
assert_eq!(agent.tool_call_history().len(), 2);
}
#[tokio::test]
async fn alias_mediated_fallback_cycle_is_denied_canonically() {
let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let secondary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
let agent = AgentBuilder::new()
.system_prompt("Test canonical fallback cycle admission.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(RecoveryTestTool {
id: "primary".to_string(),
succeeds: false,
calls: Arc::clone(&primary_calls),
max_output_chars: None,
}))
.tool(Arc::new(RecoveryTestTool {
id: "secondary".to_string(),
succeeds: false,
calls: Arc::clone(&secondary_calls),
max_output_chars: None,
}))
.recovery_manager(recovery_manager_with_fallbacks([
("primary".to_string(), "secondary".to_string()),
("secondary".to_string(), "primary alias".to_string()),
]))
.hooks(hooks.clone())
.build()
.unwrap();
agent.tools.set_tool_aliases(
"primary",
ToolAliases::new().with_name("en", "primary alias"),
);
let record = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.invoke_tool(ToolExecutionRequest::new(
"alias-fallback-call",
"primary",
serde_json::json!({"path": "./shared.txt"}),
ToolCallSource::Manual,
)),
)
.await
.expect("alias-mediated fallback cycle must terminate")
.unwrap();
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(secondary_calls.load(Ordering::SeqCst), 1);
assert_eq!(record.requested_name, "primary alias");
assert_eq!(record.canonical_id, "primary");
assert!(!record.executed);
assert!(record.output.contains("fallback cycle"));
assert_eq!(
record.metadata.get("fallback_chain"),
Some(&serde_json::json!(["primary", "secondary"]))
);
assert_eq!(hooks.records().len(), 3);
assert_eq!(agent.tool_call_history().len(), 3);
}
#[tokio::test]
async fn final_canonical_drift_cannot_bypass_fallback_ancestry() {
let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let secondary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let provider = Arc::new(DriftingFallbackProvider {
refreshed: AtomicBool::new(false),
primary_calls: Arc::clone(&primary_calls),
secondary_calls: Arc::clone(&secondary_calls),
});
let registry = ToolRegistry::new();
registry.register_provider(provider).await.unwrap();
let lifecycle = Arc::new(ToolLifecycleRecordingHooks::new());
let hooks = Arc::new(RefreshFallbackProviderHooks {
agent: parking_lot::Mutex::new(None),
lifecycle: Arc::clone(&lifecycle),
});
let agent = Arc::new(
AgentBuilder::new()
.system_prompt("Test final canonical fallback admission.")
.llm(Arc::new(mock_with_response("done")))
.tools(registry)
.recovery_manager(recovery_manager_with_fallbacks([(
"primary".to_string(),
"fallback alias".to_string(),
)]))
.hooks(hooks.clone())
.build()
.unwrap(),
);
*hooks.agent.lock() = Some(Arc::downgrade(&agent));
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"drifting-fallback-call",
"primary",
serde_json::json!({"path": "./shared.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
assert_eq!(secondary_calls.load(Ordering::SeqCst), 0);
assert_eq!(record.canonical_id, "secondary");
assert!(!record.executed);
assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
assert!(record.output.contains("fallback cycle"));
assert_eq!(
record.metadata.get("fallback_chain"),
Some(&serde_json::json!(["primary", "secondary"]))
);
assert_eq!(
record.metadata.get("final_resolved_canonical_id"),
Some(&serde_json::json!("primary"))
);
assert_eq!(
lifecycle.events(),
vec![
"start:primary",
"complete:primary:false",
"record:primary:true",
"error",
"start:secondary",
"complete:secondary:false",
"record:secondary:false",
"error",
]
);
let records = lifecycle.records();
assert_eq!(records.len(), 2);
assert_eq!(records[1].canonical_id, "secondary");
assert_eq!(
records[1].metadata.get("final_resolved_canonical_id"),
Some(&serde_json::json!("primary"))
);
let history = agent.tool_call_history();
assert_eq!(
history
.iter()
.map(|entry| entry.tool_id.as_str())
.collect::<Vec<_>>(),
vec!["primary", "secondary"]
);
}
#[tokio::test]
async fn acyclic_fallback_chain_is_denied_after_the_hop_limit() {
let tool_count = MAX_TOOL_FALLBACK_HOPS + 2;
let calls = (0..tool_count)
.map(|_| Arc::new(std::sync::atomic::AtomicUsize::new(0)))
.collect::<Vec<_>>();
let mut builder = AgentBuilder::new()
.system_prompt("Test bounded acyclic fallback admission.")
.llm(Arc::new(mock_with_response("done")));
for (index, counter) in calls.iter().enumerate() {
builder = builder.tool(Arc::new(RecoveryTestTool {
id: format!("fallback_{index}"),
succeeds: false,
calls: Arc::clone(counter),
max_output_chars: None,
}));
}
let fallbacks = (0..tool_count - 1).map(|index| {
(
format!("fallback_{index}"),
format!("fallback_{}", index + 1),
)
});
let agent = builder
.recovery_manager(recovery_manager_with_fallbacks(fallbacks))
.build()
.unwrap();
let record = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.invoke_tool(ToolExecutionRequest::new(
"bounded-fallback-call",
"fallback_0",
serde_json::json!({"path": "./shared.txt"}),
ToolCallSource::Manual,
)),
)
.await
.expect("bounded fallback chain must terminate")
.unwrap();
for counter in calls.iter().take(MAX_TOOL_FALLBACK_HOPS + 1) {
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
assert_eq!(calls[MAX_TOOL_FALLBACK_HOPS + 1].load(Ordering::SeqCst), 0);
assert_eq!(
record.canonical_id,
format!("fallback_{}", MAX_TOOL_FALLBACK_HOPS + 1)
);
assert!(!record.executed);
assert!(record.output.contains("maximum of 16 hops"));
assert_eq!(agent.tool_call_history().len(), tool_count);
}
#[tokio::test]
async fn diagnostics_without_provider_records_unavailable_without_execution() {
let mock = mock_with_response("hello");
let yaml = r#"
name: DiagnosticsNoProviderAgent
system_prompt: "Review diagnostics."
tools: [diagnostics]
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"diagnostics-call",
"diagnostics",
serde_json::json!({}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
}
#[tokio::test]
async fn web_search_without_provider_records_unavailable_without_execution() {
let mock = mock_with_response("hello");
let yaml = r#"
name: WebSearchNoProviderAgent
system_prompt: "You search the web."
tools: [web_search]
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"web-search-call",
"web_search",
serde_json::json!({"query": "rust async"}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
}
#[tokio::test]
async fn unavailable_host_tool_does_not_request_approval() {
let approvals = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let handler = Arc::new(CountingApprovalHandler {
calls: Arc::clone(&approvals),
});
let mut security = ToolSecurityConfig {
enabled: true,
fail_closed: true,
..Default::default()
};
security.tools.insert(
"web_search".to_string(),
ai_agents_tools::ToolPolicyConfig {
enabled: true,
require_confirmation: true,
..Default::default()
},
);
let yaml = r#"
name: UnavailableApprovalAgent
system_prompt: "Search only with approval."
tools: [web_search]
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.auto_configure_features()
.unwrap()
.tool_security(ToolSecurityEngine::new(security))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(handler)
.build()
.unwrap();
let record = agent
.invoke_tool(ToolExecutionRequest::new(
"unavailable-before-approval",
"web_search",
serde_json::json!({"query": "rust async"}),
ToolCallSource::Manual,
))
.await
.unwrap();
assert_eq!(approvals.load(Ordering::SeqCst), 0);
assert!(!record.executed);
assert!(!record.success);
assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
assert!(
record
.approval
.as_ref()
.is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Unavailable))
);
}
#[tokio::test]
async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
let mock = mock_with_response("hello");
let yaml = r#"
name: SpawnerNoGrantAgent
system_prompt: "You manage agents."
spawner:
max_agents: 2
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.auto_configure_spawner()
.await
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert!(available.is_empty());
}
#[tokio::test]
async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
let mock = mock_with_response("hello");
let yaml = r#"
name: EmptySpawnerNoGrantAgent
system_prompt: "You manage agents."
tools: []
spawner:
max_agents: 2
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.auto_configure_spawner()
.await
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert!(available.is_empty());
}
#[tokio::test]
async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
let mock = mock_with_response("hello");
let yaml = r#"
name: ManagementGrantAgent
system_prompt: "You manage agents."
tools: []
spawner:
management_tools: true
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.auto_configure_spawner()
.await
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert_eq!(available.len(), 4);
assert!(available.contains(&"spawn_agent".to_string()));
assert!(available.contains(&"send_agent_message".to_string()));
assert!(available.contains(&"list_agents".to_string()));
assert!(available.contains(&"remove_agent".to_string()));
}
#[tokio::test]
async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
let mock = mock_with_response("hello");
let yaml = r#"
name: ManagementOmittedToolsGrantAgent
system_prompt: "You manage agents."
spawner:
management_tools: true
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.auto_configure_spawner()
.await
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert_eq!(available.len(), 4);
assert!(available.contains(&"spawn_agent".to_string()));
assert!(available.contains(&"send_agent_message".to_string()));
assert!(available.contains(&"list_agents".to_string()));
assert!(available.contains(&"remove_agent".to_string()));
}
#[tokio::test]
async fn test_management_tools_selected_grants_only_selected_tools() {
let mock = mock_with_response("hello");
let yaml = r#"
name: ManagementSelectedGrantAgent
system_prompt: "You manage agents."
tools: []
spawner:
management_tools:
- spawn_agent
- send_agent_message
- list_agents
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.auto_configure_spawner()
.await
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert_eq!(available.len(), 3);
assert!(available.contains(&"spawn_agent".to_string()));
assert!(available.contains(&"send_agent_message".to_string()));
assert!(available.contains(&"list_agents".to_string()));
assert!(!available.contains(&"remove_agent".to_string()));
}
#[tokio::test]
async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
let mock = mock_with_response("hello");
let yaml = r#"
name: OrchestrationGrantAgent
system_prompt: "You coordinate agents."
llms:
default:
provider: openai
model: gpt-4
router:
provider: openai
model: gpt-4
llm:
default: default
router: router
tools: []
spawner:
orchestration_tools: true
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.auto_configure_spawner()
.await
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert_eq!(available.len(), 5);
assert!(available.contains(&"route_to_agent".to_string()));
assert!(available.contains(&"pipeline_process".to_string()));
assert!(available.contains(&"concurrent_ask".to_string()));
assert!(available.contains(&"group_discussion".to_string()));
assert!(available.contains(&"handoff_conversation".to_string()));
}
#[tokio::test]
async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
let mock = mock_with_response("hello");
let yaml = r#"
name: PersonaGrantAgent
system_prompt: "You can evolve persona."
llm:
provider: openai
model: gpt-4
tools: []
persona:
identity:
name: "Guide"
role: "Helper"
evolution:
enabled: true
allow_llm_evolve: true
mutable_fields:
- traits.personality
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert_eq!(available, vec!["persona_evolve".to_string()]);
}
#[tokio::test]
async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
let mock = mock_with_response("hello");
let yaml = r#"
name: PersonaOmittedToolsGrantAgent
system_prompt: "You can evolve persona."
llm:
provider: openai
model: gpt-4
persona:
identity:
name: "Guide"
role: "Helper"
evolution:
enabled: true
allow_llm_evolve: true
mutable_fields:
- traits.personality
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert_eq!(available, vec!["persona_evolve".to_string()]);
}
#[tokio::test]
async fn test_omitted_yaml_tools_exposes_no_tools() {
let mock = mock_with_response("hello");
let yaml = r#"
name: NoToolsAgent
system_prompt: "You are helpful."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert!(available.is_empty());
}
#[tokio::test]
async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
for tools in ["", "tools: []"] {
let yaml = format!(
r#"
name: RuntimeScopeNoGrantAgent
system_prompt: "No ordinary tools are granted."
{tools}
"#
);
let agent = AgentBuilder::from_yaml(&yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
agent
.runtime_control()
.set_tool_scope(vec!["calculator".to_string()]);
assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
}
}
#[tokio::test]
async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
let yaml = r#"
name: RuntimeScopeWideningAgent
system_prompt: "Runtime scope cannot add authority."
tools: [calculator]
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
agent
.runtime_control()
.set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
assert_eq!(
agent.get_available_tool_ids().await.unwrap(),
vec!["calculator".to_string()]
);
}
#[tokio::test]
async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
let yaml = r#"
name: RuntimeScopeIntersectionAgent
system_prompt: "Use only declared tools."
tools: [calculator, datetime]
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let mut aliases = ai_agents_tools::ToolAliases::default();
aliases
.names
.insert("en".to_string(), "calculate_alias".to_string());
agent.tools.set_tool_aliases("calculator", aliases);
let control = agent.runtime_control();
control.set_tool_scope(vec![
"datetime".to_string(),
"calculate_alias".to_string(),
"calculator".to_string(),
"unknown".to_string(),
"datetime".to_string(),
]);
assert_eq!(
agent.get_available_tool_ids().await.unwrap(),
vec!["calculator".to_string(), "datetime".to_string()]
);
control.set_tool_scope(vec!["datetime".to_string()]);
assert_eq!(
agent.get_available_tool_ids().await.unwrap(),
vec!["datetime".to_string()]
);
control.clear_tool_scope_override();
assert_eq!(
agent.get_available_tool_ids().await.unwrap(),
vec!["calculator".to_string(), "datetime".to_string()]
);
}
#[tokio::test]
async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
let agent = AgentBuilder::new()
.system_prompt("Use registered tools.")
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(ContextEchoTool))
.tool(Arc::new(SlowTool))
.build()
.unwrap();
agent.runtime_control().set_tool_scope(vec![
"Context Echo".to_string(),
"context_echo".to_string(),
"unknown".to_string(),
]);
assert_eq!(
agent.get_available_tool_ids().await.unwrap(),
vec!["context_echo".to_string()]
);
}
#[tokio::test]
async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
let yaml = r#"
name: NestedStateScopeAgent
system_prompt: "Honor every state scope."
tools: [calculator, datetime, echo]
states:
initial: root
states:
root:
tools: [calculate_alias, datetime]
initial: middle
states:
middle:
initial: leaf
states:
leaf:
tools: [datetime_alias, echo]
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
calculator_aliases
.names
.insert("en".to_string(), "calculate_alias".to_string());
agent
.tools
.set_tool_aliases("calculator", calculator_aliases);
let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
datetime_aliases
.names
.insert("en".to_string(), "datetime_alias".to_string());
agent.tools.set_tool_aliases("datetime", datetime_aliases);
agent.runtime_control().set_tool_scope(vec![
"unknown".to_string(),
"datetime_alias".to_string(),
"calculate_alias".to_string(),
"datetime".to_string(),
]);
assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
assert_eq!(
agent.get_available_tool_ids().await.unwrap(),
vec!["datetime".to_string()]
);
}
#[tokio::test]
async fn ancestor_empty_state_scope_denies_omitted_descendants() {
let yaml = r#"
name: NestedEmptyStateScopeAgent
system_prompt: "An empty ancestor scope denies all tools."
tools: [calculator]
states:
initial: root
states:
root:
tools: []
initial: middle
states:
middle:
initial: leaf
states:
leaf: {}
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
}
#[tokio::test]
async fn state_change_during_approval_invalidates_the_reviewed_authority() {
let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let entered = Arc::new(tokio::sync::Barrier::new(2));
let release = Arc::new(tokio::sync::Notify::new());
let handler = Arc::new(BlockingApprovalHandler {
entered: Arc::clone(&entered),
release: Arc::clone(&release),
result: ApprovalResult::Approved,
});
let yaml = r#"
name: ApprovalStateGenerationAgent
system_prompt: "State authority may change during approval."
tools: [locked_write]
states:
initial: first
states:
first:
tools: [locked_write]
second:
tools: [locked_write]
"#;
let agent = Arc::new(
AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(LockedWriteTool {
active: Arc::clone(&active),
max_active: Arc::clone(&max_active),
}))
.tool_security(ToolSecurityEngine::new(approval_security_config(true)))
.hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
.approval_handler(handler)
.build()
.unwrap(),
);
let running = Arc::clone(&agent);
let call = tokio::spawn(async move {
running
.invoke_tool(ToolExecutionRequest::new(
"approval-state-generation",
"locked_write",
serde_json::json!({"path": "./state-generation.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
});
entered.wait().await;
agent.transition_to("second").await.unwrap();
release.notify_one();
let record = call.await.unwrap();
assert!(!record.executed);
assert!(record.output.contains("Approval became stale"));
assert_eq!(max_active.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
let holder_gate = PathMutationGate::new();
let waiter_gate = PathMutationGate::new();
let yaml = r#"
name: LockedStateGenerationAgent
system_prompt: "State authority must remain stable through admission."
tools: [state_lock_holder, state_lock_waiter]
states:
initial: first
states:
first:
tools: [state_lock_holder, state_lock_waiter]
second:
tools: [state_lock_holder, state_lock_waiter]
"#;
let agent = Arc::new(
AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("done")))
.tool(Arc::new(BlockingPathMutationTool {
id: "state_lock_holder",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: holder_gate.clone(),
}))
.tool(Arc::new(BlockingPathMutationTool {
id: "state_lock_waiter",
path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
gate: waiter_gate.clone(),
}))
.build()
.unwrap(),
);
let holder_call = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"state-lock-holder",
"state_lock_holder",
serde_json::json!({"path": "./shared-state-path.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
holder_gate.wait_until_entered().await;
let waiter_call = {
let agent = Arc::clone(&agent);
tokio::spawn(async move {
agent
.invoke_tool(ToolExecutionRequest::new(
"state-lock-waiter",
"state_lock_waiter",
serde_json::json!({"path": "./shared-state-path.txt"}),
ToolCallSource::Manual,
))
.await
.unwrap()
})
};
wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
agent.transition_to("second").await.unwrap();
holder_gate.release();
let holder_record = holder_call.await.unwrap();
let waiter_record = waiter_call.await.unwrap();
assert!(holder_record.success);
assert!(!waiter_record.executed);
assert!(
waiter_record
.output
.contains("state scope changed before admission")
);
assert!(!waiter_gate.entered.load(Ordering::SeqCst));
}
#[tokio::test]
async fn test_state_tools_cannot_widen_top_level_grant() {
let mock = mock_with_response("hello");
let yaml = r#"
name: NarrowToolsAgent
system_prompt: "You are helpful."
tools:
- calculator
states:
initial: current
states:
current:
tools: [datetime]
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let available = agent.get_available_tool_ids().await.unwrap();
assert!(available.is_empty());
}
#[tokio::test]
async fn test_integration_tool_execution() {
let mock = mock_with_responses(vec![
r#"I'll calculate that for you.
[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
"The answer is 4.",
]);
let mut tools = ai_agents_tools::ToolRegistry::new();
tools
.register(Arc::new(ai_agents_tools::CalculatorTool))
.unwrap();
let agent = AgentBuilder::new()
.system_prompt("You are a calculator assistant.")
.llm(Arc::new(mock))
.tools(tools)
.build()
.unwrap();
let response = agent.chat("What is 2+2?").await.unwrap();
assert!(!response.content.is_empty());
}
#[tokio::test]
async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hooks = Arc::new(ResponseCountingHooks {
responses: Arc::clone(&responses),
});
let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
let yaml = r#"
name: ToolRejectAgent
system_prompt: "You use tools when requested."
tools:
- echo
hitl:
tools:
echo:
require_approval: true
approval_message: "Approve echo?"
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.hooks(hooks)
.build()
.unwrap();
let response = agent.chat("echo hello").await.unwrap();
assert!(
response.content.contains("Operation cancelled"),
"unexpected response: {}",
response.content
);
assert_eq!(responses.load(Ordering::SeqCst), 1);
let messages = agent.memory.get_messages(None).await.unwrap();
assert_eq!(messages.len(), 3);
assert_eq!(messages[0].content, "echo hello");
assert!(messages[1].content.contains("\"tool\":\"echo\""));
assert!(messages[2].content.contains("rejected by the approver"));
}
#[tokio::test]
async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
use futures::StreamExt;
let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let hooks = Arc::new(ResponseCountingHooks {
responses: Arc::clone(&responses),
});
let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
let yaml = r#"
name: ToolRejectStreamingAgent
system_prompt: "You use tools when requested."
tools:
- echo
streaming:
enabled: true
hitl:
tools:
echo:
require_approval: true
approval_message: "Approve echo?"
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.hooks(hooks)
.build()
.unwrap();
let mut stream = agent.chat_stream("echo hello").await.unwrap();
let mut terminal_error = String::new();
let mut done = false;
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Error { message } => terminal_error = message,
StreamChunk::Done {} => {
done = true;
break;
}
_ => {}
}
}
assert!(done);
assert!(
terminal_error.contains("Operation cancelled"),
"unexpected terminal error: {}",
terminal_error
);
assert_eq!(responses.load(Ordering::SeqCst), 1);
let messages = agent.memory.get_messages(None).await.unwrap();
assert_eq!(messages.len(), 3);
assert_eq!(messages[0].content, "echo hello");
assert!(messages[1].content.contains("\"tool\":\"echo\""));
assert!(messages[2].content.contains("rejected by the approver"));
}
#[tokio::test]
async fn tool_hitl_rejection_preserves_legacy_error_but_finalizes_event_stream() {
let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
let yaml = r#"
name: ToolRejectEventAgent
system_prompt: "You use tools when requested."
tools:
- echo
streaming:
enabled: true
hitl:
tools:
echo:
require_approval: true
approval_message: "Approve echo?"
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.auto_configure_features()
.unwrap()
.build()
.unwrap();
let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
let mut error_seen = false;
let mut final_response = None;
while let Some(event) = stream.next().await {
match event {
AgentStreamEvent::Chunk(StreamChunk::Error { .. }) => error_seen = true,
AgentStreamEvent::Final(response) => final_response = Some(response),
AgentStreamEvent::Chunk(_) => {}
}
}
assert!(!error_seen);
assert!(
final_response
.is_some_and(|response| { response.content.contains("Operation cancelled") })
);
}
#[tokio::test]
async fn test_pre_response_guard_transition_skips_old_state_llm() {
let mock = mock_with_response("Billing state response");
let call_counter = mock.clone();
let yaml = r#"
name: OptimizedStateAgent
system_prompt: "You route before answering."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt that should be skipped."
transitions:
- to: billing
guard:
context:
topic:
eq: billing
timing: pre_response
billing:
prompt: "Answer from the billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
agent
.set_context("topic", serde_json::json!("billing"))
.unwrap();
let response = agent.chat("I need billing help").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert_eq!(response.content, "Billing state response");
assert_eq!(call_counter.call_count(), 1);
assert_eq!(agent.actor_facts().len(), 0);
}
#[tokio::test]
async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
let mock = mock_with_response("Billing state response");
let call_counter = mock.clone();
let yaml = r#"
name: OptimizedStateAgent
system_prompt: "You route before answering."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
context:
request:
type: runtime
default:
topic: general
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt that should be skipped."
transitions:
- to: billing
guard:
context:
request.topic:
eq: billing
timing: pre_response
billing:
prompt: "Answer from the billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
agent
.set_context("request.topic", serde_json::json!("billing"))
.unwrap();
let response = agent.chat("I need billing help").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert_eq!(response.content, "Billing state response");
assert_eq!(call_counter.call_count(), 1);
assert_eq!(
agent.get_context().get("request"),
Some(&serde_json::json!({"topic": "billing"}))
);
}
#[tokio::test]
async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
let mock = mock_with_response("billing");
let yaml = r#"
name: OptimizedStateAgent
system_prompt: "You route before answering."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
hitl:
states:
billing:
on_enter: require_approval
approval_message: "Approve billing route?"
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt."
extract:
- key: topic
description: "Support topic"
transitions:
- to: billing
guard:
context:
topic:
eq: billing
timing: pre_response
run_extractors: true
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
let response = agent
.try_pre_response_transition("billing please")
.await
.unwrap();
assert!(response.is_none());
assert_eq!(agent.current_state().as_deref(), Some("greeting"));
assert!(!agent.get_context().contains_key("topic"));
assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
}
#[tokio::test]
async fn test_pre_response_extractor_commits_context_on_winning_path() {
let mock = mock_with_responses(vec!["billing", "Billing response"]);
let yaml = r#"
name: OptimizedStateAgent
system_prompt: "You route before answering."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt."
extract:
- key: topic
description: "Support topic"
transitions:
- to: billing
guard:
context:
topic:
eq: billing
timing: pre_response
run_extractors: true
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
let response = agent.chat("billing please").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert_eq!(response.content, "Billing response");
assert_eq!(
agent.get_context().get("topic"),
Some(&serde_json::json!("billing"))
);
}
#[tokio::test]
async fn test_pre_response_extractor_miss_does_not_mutate_context() {
let mock = mock_with_response("__NONE__");
let yaml = r#"
name: OptimizedStateAgent
system_prompt: "You route before answering."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt."
extract:
- key: topic
description: "Support topic"
transitions:
- to: billing
guard:
context:
topic:
eq: billing
timing: pre_response
run_extractors: true
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
let response = agent.try_pre_response_transition("hello").await.unwrap();
assert!(response.is_none());
assert_eq!(agent.current_state().as_deref(), Some("greeting"));
assert!(!agent.get_context().contains_key("topic"));
}
#[tokio::test]
async fn test_default_guard_transition_stays_post_response() {
let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
let call_counter = mock.clone();
let yaml = r#"
name: TimingAgent
system_prompt: "You route carefully."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt."
transitions:
- to: billing
guard:
context:
topic:
eq: billing
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
agent
.set_context("topic", serde_json::json!("billing"))
.unwrap();
let response = agent.chat("billing please").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert_eq!(response.content, "Billing response");
assert_eq!(call_counter.call_count(), 2);
}
#[tokio::test]
async fn test_explicit_post_response_guard_transition_stays_post_response() {
let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
let call_counter = mock.clone();
let yaml = r#"
name: TimingAgent
system_prompt: "You route carefully."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt."
transitions:
- to: billing
guard:
context:
topic:
eq: billing
timing: post_response
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
agent
.set_context("topic", serde_json::json!("billing"))
.unwrap();
let response = agent.chat("billing please").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert_eq!(response.content, "Billing response");
assert_eq!(call_counter.call_count(), 2);
}
#[tokio::test]
async fn test_pre_response_extractors_are_transition_scoped() {
let mock = mock_with_responses(vec!["billing", "Billing response"]);
let yaml = r#"
name: ScopedExtractorAgent
system_prompt: "You route carefully."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt."
extract:
- key: topic
description: "Support topic"
transitions:
- to: wrong
guard:
context:
topic:
eq: billing
timing: pre_response
- to: billing
guard:
context:
topic:
eq: billing
timing: pre_response
run_extractors: true
wrong:
prompt: "Wrong state."
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
let response = agent.chat("billing please").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert_eq!(response.content, "Billing response");
}
#[tokio::test]
async fn test_pre_response_resolved_intent_routes_early() {
let mock = mock_with_response("Billing response");
let yaml = r#"
name: IntentAgent
system_prompt: "You route carefully."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
states:
initial: greeting
states:
greeting:
prompt: "Old state prompt."
transitions:
- to: billing
intent: billing
timing: pre_response
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
agent
.set_context("resolved_intent", serde_json::json!("billing"))
.unwrap();
let response = agent
.try_pre_response_transition("I need billing help")
.await
.unwrap()
.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert_eq!(response.content, "Billing response");
}
#[tokio::test]
async fn test_background_overflow_error_surfaces() {
let mut config = RuntimeConfig::default();
config.optimization.enabled = true;
config.optimization.post_turn.max_background_tasks = 1;
config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
let policy = crate::optimization::MaintenanceTaskPolicy {
mode: MaintenanceMode::Background,
await_before_next_turn: AwaitBeforeNextTurn::Always,
};
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock_with_response("ok")))
.build()
.unwrap()
.with_runtime_config(config);
agent
.background_maintenance
.spawn(None, async { std::future::pending::<Result<()>>().await })
.unwrap();
let result = agent
.spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
let default_mock = mock_with_response("Plain draft response");
let router_mock = mock_with_response("cot");
let router_counter = router_mock.clone();
let yaml = r#"
name: ReasoningReservationAgent
system_prompt: "You answer plainly unless reasoning wins."
llm:
default: default
router: router
observability:
enabled: true
export:
write_raw_events: true
reasoning:
mode: auto
judge_llm: router
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 1
speculative_reasoning_auto: true
max_parallel_runtime_tasks: 2
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", Arc::new(default_mock))
.llm_alias("router", Arc::new(router_mock))
.build()
.unwrap();
let response = agent.chat("hello").await.unwrap();
assert_eq!(response.content, "Plain draft response");
assert_eq!(router_counter.call_count(), 1);
let events = agent.observability().unwrap().raw_events();
assert!(!events.iter().any(|event| {
event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
}));
}
#[tokio::test]
async fn test_forced_reasoning_skips_plain_speculative_draft() {
let mock = mock_with_response("Reasoned response");
let yaml = r#"
name: ForcedReasoningAgent
system_prompt: "You reason before answering."
observability:
enabled: true
export:
write_raw_events: true
reasoning:
mode: cot
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 2
speculative_state_transitions: true
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Answer from triage."
transitions:
- to: billing
guard:
context:
route:
eq: billing
timing: parallel
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
let response = agent.chat("hello").await.unwrap();
assert_eq!(response.content, "Reasoned response");
let events = agent.observability().unwrap().raw_events();
assert!(
!events
.iter()
.any(|event| event.dimensions.contains_key("branch_status"))
);
}
#[tokio::test]
async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
let default_mock = mock_with_response("Skill committed response");
let router_mock = mock_with_response("helper");
let router_counter = router_mock.clone();
let yaml = r#"
name: SkillReservationAgent
system_prompt: "Use skills when they match."
llm:
default: default
router: router
observability:
enabled: true
export:
write_raw_events: true
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 1
speculative_skill_routing: true
max_parallel_runtime_tasks: 2
skills:
- id: helper
description: "Answer helper requests"
trigger: "User asks for helper"
steps:
- prompt: "Answer the helper request: {{ user_input }}"
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", Arc::new(default_mock))
.llm_alias("router", Arc::new(router_mock))
.build()
.unwrap();
let response = agent.chat("please use helper").await.unwrap();
assert_eq!(response.content, "Skill committed response");
assert_eq!(router_counter.call_count(), 1);
let events = agent.observability().unwrap().raw_events();
assert!(
!events
.iter()
.any(|event| event.dimensions.contains_key("branch_status"))
);
}
#[tokio::test]
async fn test_parallel_transition_low_cap_allows_deterministic_route() {
let mock = mock_with_response("unused");
let call_counter = mock.clone();
let yaml = r#"
name: ParallelTransitionLowCapAgent
system_prompt: "Route before stale responses when safe."
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 1
speculative_state_transitions: true
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Triage state."
transitions:
- to: billing
guard:
context:
route:
eq: billing
timing: parallel
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
agent
.set_context("route", serde_json::json!("billing"))
.unwrap();
agent.update_active_turn_context("billing help", HashMap::new());
assert!(
agent.reserve_active_speculative_llm_call(
RuntimeOptimizationKind::ParallelStateTransition
)
);
let selection = agent
.select_parallel_transition_candidate("billing help")
.await
.unwrap();
agent.end_root_turn();
match selection {
ParallelTransitionSelection::Candidate(candidate) => {
assert_eq!(candidate.target(), "billing");
}
ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
ParallelTransitionSelection::ReservationExhausted => {
panic!("deterministic route consumed LLM budget")
}
}
assert_eq!(call_counter.call_count(), 0);
}
#[tokio::test]
async fn speculative_transition_drops_loser_before_state_actions() {
let lock = Arc::new(tokio::sync::Mutex::new(()));
let first_started = Arc::new(tokio::sync::Notify::new());
let first_dropped = Arc::new(AtomicBool::new(false));
let committed_after_drop = Arc::new(AtomicBool::new(false));
let default = Arc::new(FirstCallLockingProvider {
lock,
first_started: Arc::clone(&first_started),
first_dropped: Arc::clone(&first_dropped),
committed_after_drop: Arc::clone(&committed_after_drop),
calls: AtomicU64::new(0),
});
let router = Arc::new(RoutingAfterProviderStart {
provider_started: first_started,
});
let yaml = r#"
name: SpeculativeCancellationAgent
system_prompt: "Route before committed work."
llm:
default: default
router: router
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 2
speculative_state_transitions: true
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Triage state."
transitions:
- to: technical
when: "The request needs technical support"
timing: parallel
technical:
prompt: "Technical state."
on_enter:
- prompt: "Prepare technical context."
llm: default
store_as: preparation
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", default)
.llm_alias("router", router)
.build()
.unwrap();
let response = tokio::time::timeout(
std::time::Duration::from_secs(2),
agent.chat("I cannot log in because of AUTH-17."),
)
.await
.expect("committed work must not wait on the losing provider future")
.unwrap();
assert_eq!(response.content, "Committed technical response.");
assert_eq!(agent.current_state().as_deref(), Some("technical"));
assert!(first_dropped.load(Ordering::SeqCst));
assert!(committed_after_drop.load(Ordering::SeqCst));
}
#[tokio::test]
async fn buffered_transition_drops_stale_stream_before_redispatch() {
use futures::StreamExt;
let lock = Arc::new(tokio::sync::Mutex::new(()));
let stream_started = Arc::new(tokio::sync::Notify::new());
let stream_dropped = Arc::new(AtomicBool::new(false));
let committed_after_drop = Arc::new(AtomicBool::new(false));
let default = Arc::new(BufferedLockingProvider {
lock,
stream_started: Arc::clone(&stream_started),
stream_dropped: Arc::clone(&stream_dropped),
committed_after_drop: Arc::clone(&committed_after_drop),
});
let router = Arc::new(RoutingAfterProviderStart {
provider_started: stream_started,
});
let yaml = r#"
name: BufferedCancellationAgent
system_prompt: "Hide stale streamed output."
llm:
default: default
router: router
streaming:
enabled: true
buffer_size: 8
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 2
speculative_state_transitions: true
streaming_policy: buffer_until_routing_done
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Triage state."
transitions:
- to: technical
when: "The request needs technical support"
timing: parallel
technical:
prompt: "Technical state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", default)
.llm_alias("router", router)
.build()
.unwrap();
let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
let mut stream = agent
.chat_stream("AUTH-17 needs technical help.")
.await
.unwrap();
let mut content = String::new();
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Content { text } => content.push_str(&text),
StreamChunk::Done {} => break,
StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
_ => {}
}
}
content
})
.await
.expect("redispatch must not wait on the stale streaming future");
assert_eq!(content, "Committed technical response.");
assert_eq!(agent.current_state().as_deref(), Some("technical"));
assert!(stream_dropped.load(Ordering::SeqCst));
assert!(committed_after_drop.load(Ordering::SeqCst));
}
#[tokio::test]
async fn buffered_transition_drops_established_stream_before_redispatch() {
use futures::StreamExt;
let stream_started = Arc::new(tokio::sync::Notify::new());
let stream_dropped = Arc::new(AtomicBool::new(false));
let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
let committed_after_drop = Arc::new(AtomicBool::new(false));
let default = Arc::new(EstablishedStreamProvider {
stream_started: Arc::clone(&stream_started),
stream_dropped: Arc::clone(&stream_dropped),
stream_dropped_notify,
committed_after_drop: Arc::clone(&committed_after_drop),
});
let router = Arc::new(RoutingAfterProviderStart {
provider_started: stream_started,
});
let yaml = r#"
name: EstablishedStreamCancellationAgent
system_prompt: "Hide stale streamed output."
llm:
default: default
router: router
streaming:
enabled: true
buffer_size: 8
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 2
speculative_state_transitions: true
streaming_policy: buffer_until_routing_done
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Triage state."
transitions:
- to: technical
when: "The request needs technical support"
timing: parallel
technical:
prompt: "Technical state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", default)
.llm_alias("router", router)
.build()
.unwrap();
let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
let mut stream = agent
.chat_stream("AUTH-17 needs technical help.")
.await
.unwrap();
let mut content = String::new();
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Content { text } => content.push_str(&text),
StreamChunk::Done {} => break,
StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
_ => {}
}
}
content
})
.await
.expect("redispatch must wait for the established stale stream to be dropped");
assert_eq!(content, "Committed technical response.");
assert_eq!(agent.current_state().as_deref(), Some("technical"));
assert!(stream_dropped.load(Ordering::SeqCst));
assert!(committed_after_drop.load(Ordering::SeqCst));
}
#[tokio::test]
async fn test_buffered_streaming_transition_reservation_falls_back() {
use futures::StreamExt;
let mock = mock_with_responses(vec![
"Serial streaming response",
"Serial streaming response",
]);
let router_mock = mock_with_response("1");
let router_counter = router_mock.clone();
let yaml = r#"
name: BufferedReservationFallbackAgent
system_prompt: "Stream normally if speculative routing cannot be evaluated."
llm:
default: default
router: router
observability:
enabled: true
export:
write_raw_events: true
streaming:
enabled: true
buffer_size: 8
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 1
speculative_state_transitions: true
streaming_policy: buffer_until_routing_done
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Triage state."
transitions:
- to: billing
guard:
context:
route:
eq: billing
when: "User asks about billing"
timing: parallel
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", Arc::new(mock))
.llm_alias("router", Arc::new(router_mock))
.build()
.unwrap();
let mut stream = agent.chat_stream("hello").await.unwrap();
let mut content = String::new();
let mut error = None;
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Content { text } => content.push_str(&text),
StreamChunk::Error { message } => error = Some(message),
StreamChunk::Done {} => break,
_ => {}
}
}
assert_eq!(error, None);
assert_eq!(content, "Serial streaming response");
assert_eq!(router_counter.call_count(), 0);
let events = agent.observability().unwrap().raw_events();
assert!(events.iter().any(|event| {
event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
&& event.dimensions.get("commit_behavior")
== Some(&"transition_decision".to_string())
}));
}
#[tokio::test]
async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
let mut mock = mock_with_response("Recovered response");
mock.set_error("boom");
let mut handle = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.build()
.unwrap();
assert!(agent.chat("first").await.is_err());
handle.clear_error();
let response = agent.chat("second").await.unwrap();
assert_eq!(response.content, "Recovered response");
let messages = agent.memory.get_messages(None).await.unwrap();
let user_count = messages
.iter()
.filter(|message| message.role == ai_agents_core::Role::User)
.count();
assert_eq!(user_count, 2);
}
#[tokio::test]
async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
use futures::StreamExt;
let mut mock = mock_with_response("Recovered response");
mock.set_error("stream boom");
let mut handle = mock.clone();
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.build()
.unwrap();
let mut stream = agent.chat_stream("first").await.unwrap();
let mut saw_error = false;
while let Some(chunk) = stream.next().await {
if matches!(chunk, StreamChunk::Error { .. }) {
saw_error = true;
}
}
assert!(saw_error);
handle.clear_error();
let response = agent.chat("second").await.unwrap();
assert_eq!(response.content, "Recovered response");
let messages = agent.memory.get_messages(None).await.unwrap();
let user_count = messages
.iter()
.filter(|message| message.role == ai_agents_core::Role::User)
.count();
assert_eq!(user_count, 2);
}
#[tokio::test]
async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
use futures::StreamExt;
let mut mock = mock_with_response("one two three");
mock.set_latency(10);
let yaml = r#"
name: BufferedMissAgent
system_prompt: "You stream safely."
llm:
default: default
streaming:
enabled: true
buffer_size: 1
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 2
speculative_state_transitions: true
streaming_policy: buffer_until_routing_done
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Answer from triage."
transitions:
- to: billing
guard:
context:
route:
eq: billing
timing: parallel
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", Arc::new(mock))
.build()
.unwrap();
let mut stream = agent.chat_stream("hello").await.unwrap();
let mut content = String::new();
let mut error = None;
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Content { text } => content.push_str(&text),
StreamChunk::Error { message } => error = Some(message),
StreamChunk::Done {} => break,
_ => {}
}
}
assert_eq!(error, None);
assert_eq!(content, "one two three");
}
#[tokio::test]
async fn test_buffered_streaming_main_failure_finalizes_branch() {
use futures::StreamExt;
let mock = mock_with_response("one two");
let mut router_mock = mock_with_response("0");
router_mock.set_latency(50);
let yaml = r#"
name: BufferedFailureAgent
system_prompt: "You stream safely."
llm:
default: default
router: router
observability:
enabled: true
export:
write_raw_events: true
streaming:
enabled: true
buffer_size: 1
runtime:
optimization:
enabled: true
max_speculative_llm_calls_per_turn: 2
speculative_state_transitions: true
streaming_policy: buffer_until_routing_done
max_parallel_runtime_tasks: 2
states:
initial: triage
states:
triage:
prompt: "Ask for the category."
transitions:
- to: billing
when: "User asks about billing"
timing: parallel
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm_alias("default", Arc::new(mock))
.llm_alias("router", Arc::new(router_mock))
.build()
.unwrap();
let mut stream = agent.chat_stream("hello").await.unwrap();
let mut error = String::new();
while let Some(chunk) = stream.next().await {
if let StreamChunk::Error { message } = chunk {
error = message;
}
}
assert!(
error.contains("stream buffer filled"),
"unexpected stream error: {}",
error
);
let events = agent.observability().unwrap().raw_events();
assert!(events.iter().any(|event| {
event.dimensions.get("branch_status") == Some(&"failed".to_string())
&& event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
&& event.dimensions.get("optimization")
== Some(&"buffered_streaming_routing".to_string())
}));
}
#[tokio::test]
async fn test_streaming_preflight_does_not_emit_old_state_content() {
use futures::StreamExt;
let mock = mock_with_response("Billing streamed response");
let yaml = r#"
name: StreamingOptimizedAgent
system_prompt: "You route before streaming."
runtime:
optimization:
enabled: true
pre_response_deterministic_transitions: true
streaming:
enabled: true
states:
initial: greeting
states:
greeting:
prompt: "OLD_STATE_SENTINEL"
transitions:
- to: billing
guard:
context:
topic:
eq: billing
timing: pre_response
billing:
prompt: "Billing state."
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock))
.build()
.unwrap();
agent
.set_context("topic", serde_json::json!("billing"))
.unwrap();
let mut stream = agent.chat_stream("billing please").await.unwrap();
let mut content = String::new();
while let Some(chunk) = stream.next().await {
match chunk {
StreamChunk::Content { text } => content.push_str(&text),
StreamChunk::Error { message } => panic!("stream error: {}", message),
StreamChunk::Done {} => break,
_ => {}
}
}
assert_eq!(agent.current_state().as_deref(), Some("billing"));
assert!(content.contains("Billing streamed response"));
assert!(!content.contains("OLD_STATE_SENTINEL"));
}
#[tokio::test]
async fn test_integration_state_machine_basic() {
let yaml = r#"
name: StateAgent
system_prompt: "You are a support agent."
states:
initial: greeting
states:
greeting:
prompt: "Welcome the user warmly."
transitions:
- to: support
when: "User needs help"
auto: true
support:
prompt: "Help solve the user's problem."
"#;
let mock = mock_with_responses(vec![
"Welcome! How can I help?", "1", "I'll help you with that.", ]);
let builder = AgentBuilder::from_yaml(yaml).unwrap();
let agent = builder.llm(Arc::new(mock)).build().unwrap();
assert_eq!(agent.current_state(), Some("greeting".to_string()));
let _ = agent.chat("I need help").await.unwrap();
}
#[tokio::test]
async fn test_integration_state_on_enter_set_context() {
let yaml = r#"
name: ActionAgent
system_prompt: "You are helpful."
states:
initial: step1
states:
step1:
prompt: "Step 1"
on_exit:
- set_context:
step1_exited: true
transitions:
- to: step2
when: "always"
auto: true
step2:
prompt: "Step 2"
on_enter:
- set_context:
step2_entered: true
"#;
let mock = mock_with_responses(vec![
"Processing step 1.",
"0", ]);
let builder = AgentBuilder::from_yaml(yaml).unwrap();
let agent = builder.llm(Arc::new(mock)).build().unwrap();
assert_eq!(agent.current_state(), Some("step1".to_string()));
agent.transition_to("step2").await.unwrap();
assert_eq!(agent.current_state(), Some("step2".to_string()));
let ctx = agent.get_context();
assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
}
#[tokio::test]
async fn state_action_tool_preserves_source_in_stored_record() {
let yaml = r#"
name: StateActionToolAgent
system_prompt: "You are helpful."
tools:
- context_echo
states:
initial: idle
states:
idle:
prompt: "Idle"
active:
prompt: "Active"
on_enter:
- set_context:
action_started: true
- tool: context_echo
args: {}
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("unused")))
.tool(Arc::new(ContextEchoTool))
.build()
.unwrap();
agent.transition_to("active").await.unwrap();
let record: ToolExecutionRecord = serde_json::from_value(
agent
.get_context()
.get("last_tool_record")
.cloned()
.expect("successful state action must store its execution record"),
)
.unwrap();
assert!(record.executed);
assert!(record.success);
assert_eq!(record.canonical_id, "context_echo");
assert!(matches!(
&record.source,
ToolCallSource::StateAction {
state: Some(state),
action_index: 1,
} if state == "active"
));
}
#[tokio::test]
async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
let yaml = r#"
name: OrdinaryLifecycleAgent
system_prompt: "You are helpful."
states:
initial: intake
regenerate_on_transition: false
states:
intake:
prompt: "Intake"
transitions:
- to: drafting
guard:
context:
route:
eq: drafting
drafting:
prompt: "Drafting"
on_enter:
- set_context:
draft_version: 1
on_reenter:
- set_context:
draft_version: 2
transitions:
- to: review
guard:
context:
route:
eq: review
review:
prompt: "Review"
on_enter:
- set_context:
review_entry: first
transitions:
- to: drafting
guard:
context:
route:
eq: drafting
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_responses(vec![
"Intake response",
"Draft response",
"Review response",
])))
.build()
.unwrap();
agent
.set_context("route", serde_json::json!("drafting"))
.unwrap();
agent.chat("Start a draft").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("drafting"));
assert_eq!(
agent.get_context().get("draft_version"),
Some(&serde_json::json!(1))
);
agent
.set_context("route", serde_json::json!("review"))
.unwrap();
agent.chat("Review this").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("review"));
assert_eq!(
agent.get_context().get("review_entry"),
Some(&serde_json::json!("first"))
);
agent
.set_context("route", serde_json::json!("drafting"))
.unwrap();
agent.chat("Revise this").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("drafting"));
assert_eq!(
agent.get_context().get("draft_version"),
Some(&serde_json::json!(2))
);
}
#[tokio::test]
async fn test_manual_transition_uses_on_enter_then_on_reenter() {
let yaml = r#"
name: ManualLifecycleAgent
system_prompt: "You are helpful."
states:
initial: intake
states:
intake:
prompt: "Intake"
drafting:
prompt: "Drafting"
on_enter:
- set_context:
draft_version: 1
on_reenter:
- set_context:
draft_version: 2
review:
prompt: "Review"
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_response("unused")))
.build()
.unwrap();
assert!(!agent.get_context().contains_key("draft_version"));
agent.transition_to("drafting").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("drafting"));
assert_eq!(
agent.get_context().get("draft_version"),
Some(&serde_json::json!(1))
);
agent.transition_to("review").await.unwrap();
agent.transition_to("drafting").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("drafting"));
assert_eq!(
agent.get_context().get("draft_version"),
Some(&serde_json::json!(2))
);
}
#[tokio::test]
async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
let yaml = r#"
name: TimeoutLifecycleAgent
system_prompt: "You are helpful."
states:
initial: intake
regenerate_on_transition: false
states:
intake:
prompt: "Intake"
max_turns: 1
timeout_to: drafting
drafting:
prompt: "Drafting"
max_turns: 1
timeout_to: review
on_enter:
- set_context:
draft_version: 1
on_reenter:
- set_context:
draft_version: 2
review:
prompt: "Review"
max_turns: 1
timeout_to: drafting
on_enter:
- set_context:
review_entry: first
"#;
let agent = AgentBuilder::from_yaml(yaml)
.unwrap()
.llm(Arc::new(mock_with_responses(vec![
"Intake",
"First draft",
"Review",
"Revised draft",
])))
.build()
.unwrap();
agent.chat("First turn").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("intake"));
assert!(!agent.get_context().contains_key("draft_version"));
agent.chat("Second turn").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("drafting"));
assert_eq!(
agent.get_context().get("draft_version"),
Some(&serde_json::json!(1))
);
agent.chat("Third turn").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("review"));
assert_eq!(
agent.get_context().get("review_entry"),
Some(&serde_json::json!("first"))
);
agent.chat("Fourth turn").await.unwrap();
assert_eq!(agent.current_state().as_deref(), Some("drafting"));
assert_eq!(
agent.get_context().get("draft_version"),
Some(&serde_json::json!(2))
);
}
#[tokio::test]
async fn test_integration_process_normalize() {
let yaml = r#"
name: ProcessAgent
system_prompt: "You are helpful."
process:
input:
- type: normalize
config:
trim: true
collapse_whitespace: true
"#;
let mock = mock_with_response("Got your message.");
let builder = AgentBuilder::from_yaml(yaml).unwrap();
let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
let _ = agent.chat(" hello world ").await.unwrap();
let history = mock.call_history();
assert!(!history.is_empty());
let last_call = history.last().unwrap();
let user_msg = last_call
.messages
.iter()
.find(|m| m.role == ai_agents_core::Role::User)
.unwrap();
assert_eq!(user_msg.content, "hello world");
}
#[tokio::test]
async fn test_integration_memory_compression() {
let yaml = r#"
name: MemoryAgent
system_prompt: "You are helpful."
memory:
type: compacting
max_messages: 100
compress_threshold: 5
max_recent_messages: 3
summarize_batch_size: 2
"#;
let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
let mock = mock_with_responses(responses);
let builder = AgentBuilder::from_yaml(yaml).unwrap();
let agent = builder.llm(Arc::new(mock)).build().unwrap();
for i in 0..6 {
let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
}
let messages = agent.memory.get_messages(None).await.unwrap();
assert!(messages.len() <= 12); }
#[tokio::test]
async fn test_integration_multi_llm_registry() {
let mut mock_default = MockLLMProvider::new("default");
mock_default.set_response("Default LLM response.");
let mut mock_router = MockLLMProvider::new("router");
mock_router.set_response("Router response.");
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm_alias("default", Arc::new(mock_default))
.llm_alias("router", Arc::new(mock_router))
.build()
.unwrap();
let response = agent.chat("Hello").await.unwrap();
assert_eq!(response.content, "Default LLM response.");
}
#[tokio::test]
async fn test_integration_agent_reset() {
let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.build()
.unwrap();
let _ = agent.chat("Hi").await.unwrap();
let messages = agent.memory.get_messages(None).await.unwrap();
assert_eq!(messages.len(), 2);
agent.reset().await.unwrap();
let messages = agent.memory.get_messages(None).await.unwrap();
assert_eq!(messages.len(), 0);
}
#[tokio::test]
async fn test_integration_process_validate_reject() {
use ai_agents_process::{ProcessConfig, ProcessProcessor};
let validate_config = ai_agents_process::ValidateStage {
id: Some("length_check".to_string()),
condition: None,
config: ai_agents_process::ValidateConfig {
rules: vec![ai_agents_process::ValidationRule::MinLength {
min_length: 10,
on_fail: ai_agents_process::ValidationAction {
action: ai_agents_process::ValidationActionType::Reject,
message: None,
},
}],
..Default::default()
},
};
let process_config = ProcessConfig {
input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
..Default::default()
};
let processor = ProcessProcessor::new(process_config);
let mock = mock_with_response("Should not reach here.");
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.process_processor(processor)
.build()
.unwrap();
let response = agent.chat("Hi").await.unwrap();
assert!(
response.content.contains("rejected")
|| response.content.contains("Input rejected")
|| response.content.contains("too short")
|| response.content.contains("Too short")
|| response.content.len() < 50, "Expected rejection response, got: {}",
response.content
);
}
#[tokio::test]
async fn test_llm_fallback_on_failure() {
use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
let mut primary = MockLLMProvider::new("primary");
primary.set_error("Primary LLM is unavailable");
let mut fallback = MockLLMProvider::new("fallback");
fallback.set_response("Fallback response works!");
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm_alias("default", Arc::new(primary))
.llm_alias("backup", Arc::new(fallback))
.recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
llm: LLMRecoveryConfig {
on_failure: LLMFailureAction::FallbackLlm {
fallback_llm: "backup".to_string(),
},
..Default::default()
},
..Default::default()
}))
.build()
.unwrap();
let response = agent.chat("Hello").await.unwrap();
assert!(
response.content.contains("Fallback response"),
"Expected fallback response, got: {}",
response.content
);
}
#[tokio::test]
async fn test_llm_fallback_response_static_message() {
use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
let mut primary = MockLLMProvider::new("primary");
primary.set_error("Primary LLM is unavailable");
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(primary))
.recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
llm: LLMRecoveryConfig {
on_failure: LLMFailureAction::FallbackResponse {
message: "I am temporarily unavailable. Please try again later."
.to_string(),
},
..Default::default()
},
..Default::default()
}))
.build()
.unwrap();
let response = agent.chat("Hello").await.unwrap();
assert!(
response.content.contains("temporarily unavailable"),
"Expected static fallback message, got: {}",
response.content
);
}
#[tokio::test]
async fn test_tool_failure_skip() {
use ai_agents_recovery::{
ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
};
let mock = mock_with_responses(vec![
r#"I'll use the nonexistent tool.
[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
"The tool was unavailable, but I can still help you.",
]);
let agent = AgentBuilder::new()
.system_prompt("You are helpful.")
.llm(Arc::new(mock))
.recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
tools: ToolRecoveryConfig {
default: ToolRetryConfig {
max_retries: 0,
timeout_ms: None,
on_failure: ToolFailureAction::Skip,
},
..Default::default()
},
..Default::default()
}))
.build()
.unwrap();
let response = agent.chat("Use the nonexistent tool").await;
assert!(
response.is_ok(),
"Expected Ok with skip policy, got: {:?}",
response
);
}
}