1use async_trait::async_trait;
2use futures::stream::{Stream, StreamExt};
3use parking_lot::RwLock;
4use serde_json::Value;
5use std::collections::{HashMap, HashSet};
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::{Arc, Weak};
10use std::time::Instant;
11use tracing::{debug, error, info, instrument, warn};
12
13const DISAMBIGUATION_STATE_GENERATION_KEY: &str = "_runtime.disambiguation_state_generation";
14
15pub(crate) type RootTurnGate = Arc<tokio::sync::Mutex<()>>;
17
18pub(crate) type RootTurnGateIdentityStack = Arc<[RootTurnGate]>;
20
21tokio::task_local! {
22 static RUNTIME_GATE_IDENTITY_STACK: RootTurnGateIdentityStack;
23}
24
25pub(crate) fn current_runtime_gate_identity_stack() -> RootTurnGateIdentityStack {
27 RUNTIME_GATE_IDENTITY_STACK
28 .try_with(Arc::clone)
29 .unwrap_or_default()
30}
31
32pub(crate) async fn scope_runtime_gate_identity_stack<F, T>(
34 identity_stack: &RootTurnGateIdentityStack,
35 future: F,
36) -> T
37where
38 F: Future<Output = T>,
39{
40 RUNTIME_GATE_IDENTITY_STACK
41 .scope(Arc::clone(identity_stack), future)
42 .await
43}
44
45pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Weak<tokio::sync::Mutex<()>>>>>;
47
48struct ToolResourceGuards {
52 guards: Vec<tokio::sync::OwnedMutexGuard<()>>,
53 locks: ToolResourceLocks,
54}
55
56struct RootTurnAdmission {
60 guard: tokio::sync::OwnedMutexGuard<()>,
61 identity_stack: RootTurnGateIdentityStack,
62}
63
64#[derive(Clone)]
65struct StoredSessionRestore {
66 snapshot: AgentSnapshot,
67 metadata: Option<ai_agents_core::SessionMetadata>,
68}
69
70struct RuntimeSessionRestorePoint {
71 snapshot: AgentSnapshot,
72 metadata: ai_agents_core::SessionMetadata,
73 actor_id: Option<String>,
74 session_id: Option<String>,
75}
76
77impl Drop for ToolResourceGuards {
78 fn drop(&mut self) {
79 self.guards.clear();
80 self.locks.write().retain(|_, lock| lock.strong_count() > 0);
81 }
82}
83
84#[derive(Clone)]
88struct RuntimeSafetySnapshot {
89 version: u64,
90 emergency_deny: bool,
91 tool_security: ToolSecurityEngine,
92 tool_scope_override: Option<Vec<String>>,
93}
94
95#[derive(Clone, Copy)]
99struct ToolDecisionVersions {
100 policy: u64,
101 registry: u64,
102 runtime_control: u64,
103 state: Option<u64>,
104}
105
106struct AvailableToolIdsSnapshot {
110 tool_ids: Vec<String>,
111 state_generation: Option<u64>,
112}
113
114#[derive(Clone)]
118struct ToolApprovalBinding {
119 canonical_id: String,
120 arguments: Value,
121 confirmation_required: bool,
122 policy_version: u64,
123 runtime_control_version: u64,
124 state_generation: Option<u64>,
125 reviewed_tool: Arc<dyn ai_agents_core::Tool>,
126}
127
128fn merge_approved_record(record: &mut Option<ToolApprovalRecord>) {
132 if record
133 .as_ref()
134 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::Modified))
135 {
136 return;
137 }
138 *record = Some(ToolApprovalRecord {
139 status: ToolApprovalStatus::Approved,
140 reason: None,
141 modified_arguments: None,
142 });
143}
144
145impl ToolApprovalBinding {
146 fn is_stale(
148 &self,
149 canonical_id: &str,
150 arguments: &Value,
151 confirmation_required: bool,
152 versions: ToolDecisionVersions,
153 resolved_tool: &Arc<dyn ai_agents_core::Tool>,
154 ) -> bool {
155 self.canonical_id != canonical_id
156 || self.arguments != *arguments
157 || self.confirmation_required != confirmation_required
158 || self.policy_version != versions.policy
159 || self.runtime_control_version != versions.runtime_control
160 || self.state_generation != versions.state
161 || !Arc::ptr_eq(&self.reviewed_tool, resolved_tool)
162 }
163}
164
165use crate::turn_context::{current_turn_actor_context, scope_actor_context};
166
167use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
168use ai_agents_core::traits::storage::StorageCapability;
169use ai_agents_core::{
170 AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMError, LLMProvider,
171 LLMResponse, LLMToolDefinition, LLMToolRequest, PermissionOutcome, Result, ToolActorContext,
172 ToolApprovalRecord, ToolApprovalStatus, ToolCallSource, ToolCancellationToken, ToolChoice,
173 ToolExecutionContext, ToolExecutionRecord, ToolExecutionRequest, ToolInvoker,
174 ToolPolicyDecisionRecord, ToolResult,
175};
176use ai_agents_disambiguation::{
177 ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
178 ConfirmationParseFuture, DisambiguationConfig, DisambiguationContext, DisambiguationManager,
179 DisambiguationResult,
180};
181use ai_agents_hitl::{
182 ApprovalHandler, ApprovalResolvedOutcome, ApprovalResult, ApprovalTrigger, HITLCheckResult,
183 HITLEngine, RejectAllHandler, TimeoutAction,
184};
185use ai_agents_hooks::{AgentHooks, NoopHooks};
186use ai_agents_llm::LLMRegistry;
187use ai_agents_memory::{
188 CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
189 MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
190};
191use ai_agents_observability::{
192 EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
193 current_observation_context, new_session_id as new_observation_session_id,
194 resolve_language_from_context, with_observation_context, with_observation_purpose,
195};
196use ai_agents_process::{
197 ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
198};
199use ai_agents_reasoning::{
200 CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
201 ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
202 ReflectionMetadata, StepFailureAction,
203};
204use ai_agents_recovery::{
205 ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
206 LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
207};
208use ai_agents_relationships::RelationshipManager;
209use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
210use ai_agents_state::{
211 PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, Transition,
212 TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
213};
214use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
215use ai_agents_tools::{
216 CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
217 QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord, ToolRegistry,
218 ToolSecurityConfig, ToolSecurityEngine,
219};
220
221use super::{
222 Agent, AgentInfo, AgentResponse, AgentStreamEvent, ParallelToolsConfig, StreamChunk,
223 StreamingConfig, ToolCall,
224};
225use crate::optimization::{
226 AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
227 MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
228 RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
229 RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
230 StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
231};
232use crate::spec::StorageConfig;
233
234enum ToolCallOutcome {
236 Continue,
238 TransitionFired,
240 Rejected(AgentResponse),
242}
243
244#[derive(Clone)]
245struct MainToolProtocol {
246 choice: Option<ToolChoice>,
247 tool_ids: Vec<String>,
248 definitions: Vec<LLMToolDefinition>,
249}
250
251struct MainProviderResponse {
252 response: LLMResponse,
253 used_native_tools: bool,
254}
255
256struct CommittedTextResponse<'a> {
260 processed_input: &'a str,
261 input_context: &'a HashMap<String, Value>,
262 answer: String,
263 reasoning_mode: ReasoningMode,
264 auto_detected: bool,
265 iterations: u32,
266 thinking_content: Option<String>,
267 all_tool_calls: Vec<ToolCall>,
268}
269
270struct AgentResponseParts {
274 content: String,
275 all_tool_calls: Vec<ToolCall>,
276 reasoning_mode: ReasoningMode,
277 auto_detected: bool,
278 iterations: u32,
279 thinking: Option<String>,
280 reflection_metadata: Option<ReflectionMetadata>,
281}
282
283type RuntimeStreamTerminalSlot = Arc<RwLock<Option<AgentResponse>>>;
287
288fn new_runtime_stream_terminal_slot() -> RuntimeStreamTerminalSlot {
292 Arc::new(RwLock::new(None))
293}
294
295fn record_runtime_stream_final(slot: &RuntimeStreamTerminalSlot, response: AgentResponse) {
299 *slot.write() = Some(response);
300}
301
302#[derive(Clone, Copy)]
303struct DisambiguationOwnership {
304 epoch: u64,
305 state_generation: Option<u64>,
306}
307
308enum SkillRouteResult {
310 NoMatch,
312 Response { skill_id: String, content: String },
314 NeedsClarification {
316 response: AgentResponse,
317 ownership: Option<DisambiguationOwnership>,
318 },
319}
320
321enum ParallelTransitionSelection {
323 Candidate(TransitionCandidate),
325 NoMatch,
327 ReservationExhausted,
329}
330
331enum PostLoopResult {
333 NoTransition(String),
335 Transitioned(String),
337 NeedsRedispatch,
340}
341
342struct StateTransitionReservation<'a> {
343 reserved: &'a AtomicBool,
344}
345
346impl Drop for StateTransitionReservation<'_> {
347 fn drop(&mut self) {
348 self.reserved.store(false, Ordering::SeqCst);
349 }
350}
351
352struct RootTurnCleanup<'a> {
353 agent: &'a RuntimeAgent,
354}
355
356impl<'a> RootTurnCleanup<'a> {
357 fn new(agent: &'a RuntimeAgent) -> Self {
358 Self { agent }
359 }
360}
361
362impl Drop for RootTurnCleanup<'_> {
363 fn drop(&mut self) {
364 self.agent.end_root_turn();
365 }
366}
367
368#[derive(Debug)]
370struct RuntimeControlState {
371 snapshot_guard: RwLock<()>,
373 version: AtomicU64,
375 emergency_deny: Arc<AtomicBool>,
377 tool_security_override: RwLock<Option<ToolSecurityEngine>>,
379 tool_scope_override: RwLock<Option<Vec<String>>>,
381}
382
383impl Default for RuntimeControlState {
384 fn default() -> Self {
385 Self {
386 snapshot_guard: RwLock::new(()),
387 version: AtomicU64::new(1),
388 emergency_deny: Arc::new(AtomicBool::new(false)),
389 tool_security_override: RwLock::new(None),
390 tool_scope_override: RwLock::new(None),
391 }
392 }
393}
394
395#[derive(Clone)]
397pub struct RuntimeControlHandle {
398 state: Arc<RuntimeControlState>,
399}
400
401impl RuntimeControlHandle {
402 pub fn version(&self) -> u64 {
404 self.state.version.load(Ordering::SeqCst)
405 }
406
407 fn bump(&self) -> u64 {
408 self.state.version.fetch_add(1, Ordering::SeqCst) + 1
409 }
410
411 pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
413 self.try_set_tool_security(config)
414 .expect("invalid tool security configuration")
415 }
416
417 pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
419 config.validate()?;
420 let _guard = self.state.snapshot_guard.write();
421 let generation = self.bump();
422 *self.state.tool_security_override.write() = Some(
423 ToolSecurityEngine::new_with_policy_version(config, generation),
424 );
425 Ok(generation)
426 }
427
428 pub fn clear_tool_security_override(&self) -> u64 {
430 let _guard = self.state.snapshot_guard.write();
431 *self.state.tool_security_override.write() = None;
432 self.bump()
433 }
434
435 pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
437 let _guard = self.state.snapshot_guard.write();
438 *self.state.tool_scope_override.write() = Some(tool_ids);
439 self.bump()
440 }
441
442 pub fn clear_tool_scope_override(&self) -> u64 {
444 let _guard = self.state.snapshot_guard.write();
445 *self.state.tool_scope_override.write() = None;
446 self.bump()
447 }
448
449 pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
451 let _guard = self.state.snapshot_guard.write();
452 self.state.emergency_deny.store(enabled, Ordering::SeqCst);
453 self.bump()
454 }
455
456 pub fn cancel_all(&self) -> u64 {
458 self.set_emergency_deny(true)
459 }
460}
461
462pub struct RuntimeAgent {
463 info: AgentInfo,
464 llm_registry: Arc<LLMRegistry>,
465 memory: Arc<dyn Memory>,
466 tools: Arc<ToolRegistry>,
467 skills: Vec<SkillDefinition>,
468 skill_router: Option<SkillRouter>,
469 skill_executor: Option<SkillExecutor>,
470 base_system_prompt: String,
471 max_iterations: u32,
472 iteration_count: RwLock<u32>,
473 max_context_tokens: u32,
474 memory_token_budget: Option<MemoryTokenBudget>,
475 recovery_manager: RecoveryManager,
476 tool_security: ToolSecurityEngine,
477 process_processor: Option<ProcessProcessor>,
478 message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
479 state_machine: Option<Arc<StateMachine>>,
480 transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
481 context_manager: Arc<ContextManager>,
482 template_renderer: TemplateRenderer,
483 tool_call_history: RwLock<Vec<ToolCallRecord>>,
484 parallel_tools: ParallelToolsConfig,
485 streaming: StreamingConfig,
486 hooks: Arc<dyn AgentHooks>,
487 hitl_engine: Option<HITLEngine>,
488 approval_handler: Arc<dyn ApprovalHandler>,
489 storage_config: StorageConfig,
490 storage: RwLock<Option<Arc<dyn AgentStorage>>>,
491 storage_init: tokio::sync::Mutex<()>,
492 reasoning_config: ReasoningConfig,
493 reflection_config: ReflectionConfig,
494 disambiguation_manager: Option<DisambiguationManager>,
495 disambiguation_epoch: AtomicU64,
497 disambiguation_admission: tokio::sync::RwLock<()>,
499 state_transition_reserved: AtomicBool,
501 persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
503 pending_skill_id: RwLock<Option<String>>,
507 current_plan: RwLock<Option<Plan>>,
508 declared_tool_ids: Option<Vec<String>>,
510 context_initialized: AtomicBool,
512 spawner: Option<Arc<crate::spawner::AgentSpawner>>,
514 spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
516 redispatch_depth: RwLock<u32>,
519 active_turn_context: RwLock<Option<TurnOptimizationContext>>,
521 root_user_message_committed: AtomicBool,
523 actor_id: RwLock<Option<String>>,
525 fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
527 fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
530 actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
532 messages_since_extraction: Arc<RwLock<usize>>,
534 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
536 facts_config: Option<ai_agents_facts::FactsConfig>,
538 session_metadata: RwLock<ai_agents_core::SessionMetadata>,
540 current_session_id: RwLock<Option<String>>,
542 relationship_manager: Option<Arc<RelationshipManager>>,
544 observability_manager: Option<Arc<ObservabilityManager>>,
546 runtime_config: RuntimeConfig,
548 background_maintenance: Arc<BackgroundMaintenanceQueue>,
550 resource_locks: ToolResourceLocks,
552 runtime_control: Arc<RuntimeControlState>,
554 root_turn_gate: RootTurnGate,
556}
557
558impl std::fmt::Debug for RuntimeAgent {
559 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
560 f.debug_struct("RuntimeAgent")
561 .field("info", &self.info)
562 .field("base_system_prompt", &self.base_system_prompt)
563 .field("max_iterations", &self.max_iterations)
564 .field("skills_count", &self.skills.len())
565 .field("max_context_tokens", &self.max_context_tokens)
566 .field("has_state_machine", &self.state_machine.is_some())
567 .field("parallel_tools", &self.parallel_tools)
568 .field("streaming", &self.streaming)
569 .field("has_hooks", &true)
570 .field("has_hitl", &self.hitl_engine.is_some())
571 .field("storage_type", &self.storage_config.storage_type())
572 .field("reasoning_mode", &self.reasoning_config.mode)
573 .field("reflection_enabled", &self.reflection_config.enabled)
574 .field("declared_tool_ids", &self.declared_tool_ids)
575 .field("has_persona", &self.persona_manager.is_some())
576 .field("has_observability", &self.observability_manager.is_some())
577 .finish_non_exhaustive()
578 }
579}
580
581struct ObservabilityClarificationObserver;
582
583impl ClarificationObserver for ObservabilityClarificationObserver {
584 fn observe_question<'a>(
586 &'a self,
587 future: ClarificationQuestionFuture<'a>,
588 ) -> ClarificationQuestionFuture<'a> {
589 Box::pin(async move {
590 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
591 })
592 }
593
594 fn observe_parse<'a>(
596 &'a self,
597 future: ClarificationParseFuture<'a>,
598 ) -> ClarificationParseFuture<'a> {
599 Box::pin(async move {
600 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
601 })
602 }
603
604 fn observe_confirmation_parse<'a>(
606 &'a self,
607 future: ConfirmationParseFuture<'a>,
608 ) -> ConfirmationParseFuture<'a> {
609 Box::pin(async move {
610 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
611 })
612 }
613}
614
615struct ObservabilityProcessStageObserver;
616
617impl ProcessStageObserver for ObservabilityProcessStageObserver {
618 fn observe<'a>(
620 &'a self,
621 hint: ProcessPurposeHint,
622 future: ProcessStageFuture<'a>,
623 ) -> ProcessStageFuture<'a> {
624 Box::pin(async move {
625 with_observation_purpose(observation_purpose_for_process(hint), future).await
626 })
627 }
628}
629
630struct RegistryLLMGetter {
631 registry: Arc<LLMRegistry>,
632}
633
634impl LLMGetter for RegistryLLMGetter {
635 fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
636 self.registry.get(alias).ok()
637 }
638}
639
640impl RuntimeAgent {
641 #[allow(clippy::too_many_arguments)]
643 pub fn new(
644 info: AgentInfo,
645 llm_registry: Arc<LLMRegistry>,
646 memory: Arc<dyn Memory>,
647 tools: Arc<ToolRegistry>,
648 skills: Vec<SkillDefinition>,
649 system_prompt: String,
650 max_iterations: u32,
651 ) -> Self {
652 let (skill_router, skill_executor) = if !skills.is_empty() {
653 let router_llm = llm_registry.router().ok();
654 let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
655 let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
656 (router, Some(executor))
657 } else {
658 (None, None)
659 };
660
661 let context_manager =
662 ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
663
664 Self {
665 info,
666 llm_registry,
667 memory,
668 tools,
669 skills,
670 skill_router,
671 skill_executor,
672 base_system_prompt: system_prompt,
673 max_iterations,
674 iteration_count: RwLock::new(0),
675 max_context_tokens: 128000,
676 memory_token_budget: None,
677 recovery_manager: RecoveryManager::default(),
678 tool_security: ToolSecurityEngine::default(),
679 process_processor: None,
680 message_filters: RwLock::new(HashMap::new()),
681 state_machine: None,
682 transition_evaluator: None,
683 context_manager: Arc::new(context_manager),
684 template_renderer: TemplateRenderer::new(),
685 tool_call_history: RwLock::new(Vec::new()),
686 parallel_tools: ParallelToolsConfig::default(),
687 streaming: StreamingConfig::default(),
688 hooks: Arc::new(NoopHooks),
689 hitl_engine: None,
690 approval_handler: Arc::new(RejectAllHandler::new()),
691 storage_config: StorageConfig::default(),
692 storage: RwLock::new(None),
693 storage_init: tokio::sync::Mutex::new(()),
694 reasoning_config: ReasoningConfig::default(),
695 reflection_config: ReflectionConfig::default(),
696 disambiguation_manager: None,
697 disambiguation_epoch: AtomicU64::new(0),
698 disambiguation_admission: tokio::sync::RwLock::new(()),
699 state_transition_reserved: AtomicBool::new(false),
700 persona_manager: None,
701 pending_skill_id: RwLock::new(None),
702 current_plan: RwLock::new(None),
703 declared_tool_ids: None,
704 context_initialized: AtomicBool::new(false),
705 spawner: None,
706 spawner_registry: None,
707 redispatch_depth: RwLock::new(0),
708 active_turn_context: RwLock::new(None),
709 root_user_message_committed: AtomicBool::new(false),
710 actor_id: RwLock::new(None),
711 fact_store: RwLock::new(None),
712 fact_extractor: RwLock::new(None),
713 actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
714 messages_since_extraction: Arc::new(RwLock::new(0)),
715 actor_memory_config: None,
716 facts_config: None,
717 session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
718 current_session_id: RwLock::new(None),
719 relationship_manager: None,
720 observability_manager: None,
721 runtime_config: RuntimeConfig::default(),
722 background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
723 resource_locks: new_tool_resource_locks(),
724 runtime_control: Arc::new(RuntimeControlState::default()),
725 root_turn_gate: Arc::new(tokio::sync::Mutex::new(())),
726 }
727 }
728
729 pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
730 self.declared_tool_ids = ids;
731 self
732 }
733
734 pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
735 self.storage_config = config;
736 self
737 }
738
739 pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
740 *self.storage.write() = Some(storage);
741 self
742 }
743
744 pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
745 self.resource_locks = locks;
746 self
747 }
748
749 pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
750 self.reasoning_config = config;
751 self
752 }
753
754 pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
755 self.reflection_config = config;
756 self
757 }
758
759 pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
761 self.relationship_manager = Some(manager);
762 self
763 }
764
765 pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
767 self.observability_manager = Some(manager);
768 self
769 }
770
771 pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
773 let max_tasks = config.optimization.post_turn.max_background_tasks;
774 self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
775 self.runtime_config = config;
776 self
777 }
778
779 pub fn runtime_config(&self) -> &RuntimeConfig {
781 &self.runtime_config
782 }
783
784 pub async fn flush_background_tasks(&self) -> Result<()> {
786 self.background_maintenance.flush_all().await
787 }
788
789 pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
791 self.background_maintenance.flush_scope(actor_id).await
792 }
793
794 pub async fn flush_background_tasks_for_purpose(
796 &self,
797 purpose: RuntimeTaskPurpose,
798 ) -> Result<()> {
799 self.background_maintenance.flush_purpose(purpose).await
800 }
801
802 pub async fn flush_background_tasks_for_actor_purpose(
804 &self,
805 actor_id: &str,
806 purpose: RuntimeTaskPurpose,
807 ) -> Result<()> {
808 self.background_maintenance
809 .flush_scope_purpose(actor_id, purpose)
810 .await
811 }
812
813 pub async fn shutdown_background_tasks(&self) -> Result<()> {
815 self.flush_background_tasks().await
816 }
817
818 pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
820 self.observability_manager.clone()
821 }
822
823 async fn export_observability_if_configured(&self) {
825 let Some(manager) = self.observability_manager.as_ref() else {
826 return;
827 };
828 let export = &manager.config().export;
829 if !export.write_report && !export.write_raw_events {
830 return;
831 }
832 if let Err(error) = manager.export().await {
833 warn!(error = %error, "Observability export failed");
834 }
835 }
836
837 pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
839 self.relationship_manager.clone()
840 }
841
842 fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
843 current_turn_actor_context()
844 }
845
846 fn effective_actor_id(&self) -> Option<String> {
847 self.current_turn_actor_context()
848 .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
849 .or_else(|| self.actor_id.read().clone())
850 }
851
852 fn effective_origin_actor_id(&self) -> Option<String> {
853 self.current_turn_actor_context()
854 .and_then(|ctx| ctx.origin_actor_id.clone())
855 .or_else(|| self.actor_id.read().clone())
856 }
857
858 fn record_session_actor_if_needed(&self) {
859 if let Some(actor_id) = self.effective_origin_actor_id() {
860 let mut meta = self.session_metadata.write();
861 meta.actor_id = Some(actor_id.clone());
862 if !meta.actors.iter().any(|a| a == &actor_id) {
863 meta.actors.push(actor_id);
864 }
865 }
866 }
867
868 fn outbound_actor_context(&self) -> crate::TurnActorContext {
869 let mut context = self.current_turn_actor_context().unwrap_or_default();
870 if context.origin_actor_id.is_none() {
871 context.origin_actor_id = self.effective_origin_actor_id();
872 }
873 context.sender_agent_id = Some(self.info.id.clone());
874 context
875 }
876
877 fn observation_session_id(&self) -> Option<String> {
879 let mut current = self.current_session_id.write();
880 if current.is_none() {
881 *current = Some(new_observation_session_id());
882 }
883 current.clone()
884 }
885
886 fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
888 let manager = self.observability_manager.as_ref()?;
889 let context = self.build_context_with_overlays();
890 let language = resolve_language_from_context(manager.config(), &context);
891 let context = current_observation_context()
892 .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
893 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
894 Some(
895 context
896 .with_actor(actor_id.or_else(|| self.effective_actor_id()))
897 .with_session(self.observation_session_id())
898 .with_state(self.current_state())
899 .with_language(Some(language)),
900 )
901 }
902
903 fn current_runtime_observation_context(
905 &self,
906 purpose: ObservationPurpose,
907 ) -> Option<SpanContext> {
908 let manager = self.observability_manager.as_ref()?;
909 let context = self.build_context_with_overlays();
910 let language = resolve_language_from_context(manager.config(), &context);
911 let mut observation = current_observation_context()
912 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
913 observation.agent_id = self.info.id.clone();
914 observation.actor_id = self.effective_actor_id();
915 observation.session_id = self.observation_session_id();
916 observation.state = self.current_state();
917 observation.language = Some(language);
918 observation.purpose = purpose;
919 Some(observation)
920 }
921
922 async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
924 where
925 F: Future<Output = T>,
926 {
927 if let Some(context) = self.current_runtime_observation_context(purpose) {
928 with_observation_context(context, future).await
929 } else {
930 future.await
931 }
932 }
933
934 fn chat_with_actor_context_boxed<'a>(
938 &'a self,
939 input: &'a str,
940 actor_context: crate::TurnActorContext,
941 ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
942 Box::pin(async move {
943 let RootTurnAdmission {
944 guard,
945 identity_stack,
946 } = self.acquire_root_turn().await?;
947 let result = scope_runtime_gate_identity_stack(&identity_stack, async move {
948 let actor_id = actor_context.effective_actor_id().map(str::to_string);
949 let run = async move {
950 scope_actor_context(
951 actor_context,
952 Box::pin(async move { self.run_loop(input).await }),
953 )
954 .await
955 };
956 let result = if let Some(context) = self.build_observation_context(actor_id) {
957 with_observation_context(context, run).await
958 } else {
959 run.await
960 };
961 self.export_observability_if_configured().await;
962 result
963 })
964 .await;
965 drop(guard);
966 result
967 })
968 }
969
970 async fn acquire_root_turn(&self) -> Result<RootTurnAdmission> {
972 let gate_identity = Arc::clone(&self.root_turn_gate);
973 let current_identity_stack = current_runtime_gate_identity_stack();
974 if current_identity_stack
975 .iter()
976 .any(|owned_gate| Arc::ptr_eq(owned_gate, &gate_identity))
977 {
978 return Err(AgentError::Other(format!(
979 "RuntimeAgent '{}' rejected reentrant root turn ownership",
980 self.info.id
981 )));
982 }
983 let guard = Arc::clone(&gate_identity).lock_owned().await;
984 let mut identity_stack = Vec::with_capacity(current_identity_stack.len() + 1);
988 identity_stack.extend(current_identity_stack.iter().cloned());
989 identity_stack.push(gate_identity);
990 Ok(RootTurnAdmission {
991 guard,
992 identity_stack: identity_stack.into(),
993 })
994 }
995
996 pub async fn chat_with_actor_context(
1000 &self,
1001 input: &str,
1002 actor_context: crate::TurnActorContext,
1003 ) -> Result<AgentResponse> {
1004 self.chat_with_actor_context_boxed(input, actor_context)
1005 .await
1006 }
1007
1008 pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
1010 let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
1011 self.chat_with_actor_context(input, actor_context).await
1012 }
1013
1014 pub async fn load_actor_relationship(&self) -> Result<()> {
1016 self.maybe_load_actor_relationship().await;
1017 Ok(())
1018 }
1019
1020 pub async fn update_relationship_dimension(
1022 &self,
1023 dimension: &str,
1024 delta: f64,
1025 reason: Option<&str>,
1026 ) -> Result<ai_agents_relationships::DimensionChange> {
1027 self.update_relationship_dimension_for_perspective(
1028 ai_agents_relationships::RelationshipPerspective::AgentToActor,
1029 dimension,
1030 delta,
1031 reason,
1032 )
1033 .await
1034 }
1035
1036 pub async fn update_relationship_dimension_for_perspective(
1040 &self,
1041 perspective: ai_agents_relationships::RelationshipPerspective,
1042 dimension: &str,
1043 delta: f64,
1044 reason: Option<&str>,
1045 ) -> Result<ai_agents_relationships::DimensionChange> {
1046 let manager = self
1047 .relationship_manager
1048 .as_ref()
1049 .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
1050 let actor_id = self.effective_actor_id().ok_or_else(|| {
1051 AgentError::Config("No actor ID set. Use set_actor_id() first".into())
1052 })?;
1053 let change = manager.update_dimension_for_perspective(
1054 &actor_id,
1055 perspective,
1056 dimension,
1057 delta,
1058 1.0,
1059 reason.unwrap_or("manual relationship update"),
1060 )?;
1061 self.persist_actor_relationship(&actor_id).await?;
1062 info!(
1063 actor_id = %actor_id,
1064 perspective = %change.perspective,
1065 dimension = %change.dimension,
1066 delta = change.delta,
1067 current = change.current,
1068 "relationship updated manually"
1069 );
1070 self.hooks
1071 .on_relationship_change(&actor_id, std::slice::from_ref(&change))
1072 .await;
1073 Ok(change)
1074 }
1075
1076 pub fn reasoning_config(&self) -> &ReasoningConfig {
1077 &self.reasoning_config
1078 }
1079
1080 pub fn reflection_config(&self) -> &ReflectionConfig {
1081 &self.reflection_config
1082 }
1083
1084 pub fn with_facts_config(
1087 mut self,
1088 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1089 facts_config: Option<ai_agents_facts::FactsConfig>,
1090 ) -> Self {
1091 self.actor_memory_config = actor_memory_config;
1092 self.facts_config = facts_config;
1093 self
1094 }
1095
1096 pub fn with_facts(
1099 mut self,
1100 store: Arc<ai_agents_facts::FactStore>,
1101 extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1102 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1103 facts_config: Option<ai_agents_facts::FactsConfig>,
1104 ) -> Self {
1105 *self.fact_store.write() = Some(store);
1106 *self.fact_extractor.write() = extractor;
1107 self.actor_memory_config = actor_memory_config;
1108 self.facts_config = facts_config;
1109 self
1110 }
1111
1112 pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1114 self.fact_store.read().clone()
1115 }
1116
1117 pub fn actor_id(&self) -> Option<String> {
1119 self.actor_id.read().clone()
1120 }
1121
1122 pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1124 *self.actor_id.write() = Some(actor_id.to_string());
1125 {
1126 let mut meta = self.session_metadata.write();
1127 meta.actor_id = Some(actor_id.to_string());
1128 if !meta.actors.iter().any(|a| a == actor_id) {
1129 meta.actors.push(actor_id.to_string());
1130 }
1131 }
1132 Ok(())
1133 }
1134
1135 pub fn clear_actor_id(&self) {
1137 *self.actor_id.write() = None;
1138 self.session_metadata.write().actor_id = None;
1139 }
1140
1141 pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1143 self.set_actor_id(user_id)
1144 }
1145
1146 pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1148 let actor_id = match self.effective_actor_id() {
1149 Some(id) => id,
1150 None => return Ok(()),
1151 };
1152
1153 let store_opt = self.fact_store.read().clone();
1154 if let Some(store) = store_opt {
1155 let facts = store.get_facts(&actor_id).await?;
1156 let count = facts.len();
1157 self.actor_facts_cache
1158 .write()
1159 .insert(actor_id.clone(), facts);
1160 self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1161 tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1162 }
1163
1164 Ok(())
1165 }
1166
1167 async fn maybe_load_actor_memory(&self) {
1169 let Some(actor_id) = self.effective_actor_id() else {
1170 return;
1171 };
1172 if self.actor_facts_cache.read().contains_key(&actor_id) {
1173 return;
1174 }
1175 let _ = self.load_actor_memory().await;
1176 }
1177
1178 async fn pre_turn_session_lifecycle(&self) {
1180 if *self.redispatch_depth.read() > 0 {
1181 return;
1182 }
1183 self.resolve_actor_id_from_context();
1184 self.await_background_before_next_turn().await;
1185 self.record_session_actor_if_needed();
1186 self.maybe_load_actor_memory().await;
1187 self.maybe_load_actor_relationship().await;
1188 *self.messages_since_extraction.write() += 1;
1189 }
1190
1191 async fn post_turn_session_lifecycle(&self) -> Result<()> {
1193 if *self.redispatch_depth.read() > 0 {
1194 return Ok(());
1195 }
1196 *self.messages_since_extraction.write() += 1;
1197 self.run_post_turn_maintenance().await
1198 }
1199
1200 fn begin_root_turn(&self) {
1202 if *self.redispatch_depth.read() == 0 {
1203 let mut guard = self.active_turn_context.write();
1204 if guard.is_none() {
1205 self.root_user_message_committed
1206 .store(false, Ordering::SeqCst);
1207 let max_calls = self
1208 .runtime_config
1209 .optimization
1210 .max_speculative_llm_calls_per_turn;
1211 *guard = Some(TurnOptimizationContext::new(
1212 String::new(),
1213 HashMap::new(),
1214 max_calls,
1215 ));
1216 }
1217 }
1218 }
1219
1220 fn update_active_turn_context(
1221 &self,
1222 processed_input: &str,
1223 input_context: HashMap<String, Value>,
1224 ) {
1225 if *self.redispatch_depth.read() > 0 {
1226 return;
1227 }
1228 let max_calls = self
1229 .runtime_config
1230 .optimization
1231 .max_speculative_llm_calls_per_turn;
1232 let mut guard = self.active_turn_context.write();
1233 match guard.as_mut() {
1234 Some(context) => {
1235 context.processed_input = processed_input.to_string();
1236 context.input_context = input_context;
1237 context.max_speculative_llm_calls = max_calls;
1238 }
1239 None => {
1240 *guard = Some(TurnOptimizationContext::new(
1241 processed_input,
1242 input_context,
1243 max_calls,
1244 ));
1245 }
1246 }
1247 }
1248
1249 async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1251 if *self.redispatch_depth.read() > 0 {
1252 return Ok(());
1253 }
1254 if !self
1255 .root_user_message_committed
1256 .swap(true, Ordering::SeqCst)
1257 {
1258 self.memory
1259 .add_message(ChatMessage::user(processed_input))
1260 .await?;
1261 if let Some(context) = self.active_turn_context.write().as_mut() {
1262 context.mark_user_message_committed();
1263 }
1264 }
1265 Ok(())
1266 }
1267
1268 fn end_root_turn(&self) {
1270 if *self.redispatch_depth.read() == 0 {
1271 self.root_user_message_committed
1272 .store(false, Ordering::SeqCst);
1273 *self.active_turn_context.write() = None;
1274 }
1275 }
1276
1277 fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1278 self.begin_root_turn();
1279 let mut guard = self.active_turn_context.write();
1280 let Some(context) = guard.as_mut() else {
1281 return false;
1282 };
1283 context.reserve_speculative_llm_call_for(kind)
1284 }
1285
1286 fn branch_context_preview(&self) -> String {
1287 let context = self.build_context_with_overlays();
1288 let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1289 const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1290 if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1291 value = value
1292 .chars()
1293 .take(MAX_CONTEXT_PREVIEW_CHARS)
1294 .collect::<String>();
1295 value.push_str("...");
1296 }
1297 value
1298 }
1299
1300 async fn await_background_before_next_turn(&self) {
1302 let optimization = &self.runtime_config.optimization;
1303 if !optimization.enabled {
1304 return;
1305 }
1306 let actor_id = self.effective_actor_id();
1307 let post = &optimization.post_turn;
1308 self.await_background_task(
1309 post.facts.await_before_next_turn,
1310 RuntimeTaskPurpose::PostTurnFacts,
1311 actor_id.as_deref(),
1312 "facts",
1313 )
1314 .await;
1315 self.await_background_task(
1316 post.relationships.await_before_next_turn,
1317 RuntimeTaskPurpose::PostTurnRelationship,
1318 actor_id.as_deref(),
1319 "relationships",
1320 )
1321 .await;
1322 }
1323
1324 async fn await_background_task(
1325 &self,
1326 policy: AwaitBeforeNextTurn,
1327 purpose: RuntimeTaskPurpose,
1328 actor_id: Option<&str>,
1329 label: &str,
1330 ) {
1331 match policy {
1332 AwaitBeforeNextTurn::Never => {}
1333 AwaitBeforeNextTurn::Always => {
1334 if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1335 warn!(label = label, error = %error, "background maintenance flush failed");
1336 }
1337 }
1338 AwaitBeforeNextTurn::SameActor => {
1339 if let Some(actor_id) = actor_id
1340 && let Err(error) = self
1341 .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1342 .await
1343 {
1344 warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1345 }
1346 }
1347 }
1348 }
1349
1350 async fn run_post_turn_maintenance(&self) -> Result<()> {
1352 let optimization = &self.runtime_config.optimization;
1353 if !optimization.enabled {
1354 self.auto_extract_facts().await;
1355 self.auto_update_relationship().await;
1356 return Ok(());
1357 }
1358
1359 let facts_mode = effective_maintenance_mode(
1360 optimization.post_turn.facts.mode,
1361 optimization.parallel_post_turn_memory,
1362 );
1363 let relationships_mode = effective_maintenance_mode(
1364 optimization.post_turn.relationships.mode,
1365 optimization.parallel_post_turn_memory,
1366 );
1367
1368 match (facts_mode, relationships_mode) {
1369 (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1370 self.auto_extract_facts().await;
1371 self.auto_update_relationship().await;
1372 }
1373 (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1374 let facts = self.auto_extract_facts();
1375 let relationships = self.auto_update_relationship();
1376 tokio::join!(facts, relationships);
1377 }
1378 (MaintenanceMode::Background, MaintenanceMode::Background) => {
1379 self.schedule_facts_background().await?;
1380 self.schedule_relationship_background().await?;
1381 }
1382 (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1383 | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1384 self.schedule_facts_background().await?;
1385 self.auto_update_relationship().await;
1386 }
1387 (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1388 | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1389 self.auto_extract_facts().await;
1390 self.schedule_relationship_background().await?;
1391 }
1392 _ => {
1393 self.auto_extract_facts().await;
1394 self.auto_update_relationship().await;
1395 }
1396 }
1397 Ok(())
1398 }
1399
1400 async fn schedule_facts_background(&self) -> Result<()> {
1401 let policy = self.runtime_config.optimization.post_turn.facts.clone();
1402 let should_extract = self
1403 .facts_config
1404 .as_ref()
1405 .map(|c| c.enabled && c.auto_extract)
1406 .unwrap_or(false);
1407 if !should_extract {
1408 return Ok(());
1409 }
1410 let msgs_since = *self.messages_since_extraction.read();
1411 if msgs_since < 2 {
1412 return Ok(());
1413 }
1414 let Some(actor_id) = self.effective_actor_id() else {
1415 self.record_skipped_maintenance(
1416 "facts",
1417 ObservationPurpose::FactsExtraction,
1418 "missing_actor",
1419 Some(&policy),
1420 );
1421 return Ok(());
1422 };
1423 let Some(extractor) = self.fact_extractor.read().clone() else {
1424 return Ok(());
1425 };
1426 let messages = match self.memory.get_messages(None).await {
1427 Ok(messages) => messages,
1428 Err(error) => {
1429 warn!(error = %error, "failed to snapshot messages for fact extraction");
1430 return Ok(());
1431 }
1432 };
1433 let recent: Vec<_> = messages
1434 .iter()
1435 .rev()
1436 .take(msgs_since)
1437 .rev()
1438 .cloned()
1439 .collect();
1440 if recent.is_empty() {
1441 return Ok(());
1442 }
1443 let existing = self
1444 .actor_facts_cache
1445 .read()
1446 .get(&actor_id)
1447 .cloned()
1448 .unwrap_or_default();
1449 let categories = self
1450 .facts_config
1451 .as_ref()
1452 .map(|c| c.custom_categories.clone())
1453 .unwrap_or_default();
1454 let store = self.fact_store.read().clone();
1455 let cache = Arc::clone(&self.actor_facts_cache);
1456 let counter = Arc::clone(&self.messages_since_extraction);
1457 let hooks = Arc::clone(&self.hooks);
1458 let agent_id = self.info.id.clone();
1459 let observation = current_observation_context();
1460 let key = MaintenanceSequenceKey::actor(
1461 agent_id,
1462 actor_id.clone(),
1463 RuntimeTaskPurpose::PostTurnFacts,
1464 );
1465 let actor_for_task = actor_id.clone();
1466 let task = async move {
1467 let run = async move {
1468 let facts = extractor
1469 .extract(&recent, &existing, Some(&actor_for_task), &categories)
1470 .await?;
1471 if !facts.is_empty() {
1472 if let Some(store) = store {
1473 let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1474 cache.write().insert(actor_for_task.clone(), authoritative);
1475 } else {
1476 cache
1477 .write()
1478 .entry(actor_for_task.clone())
1479 .or_default()
1480 .extend(facts.clone());
1481 }
1482 {
1483 let mut count = counter.write();
1484 if *count <= msgs_since {
1485 *count = 0;
1486 } else {
1487 *count -= msgs_since;
1488 }
1489 }
1490 hooks.on_facts_extracted(&actor_for_task, &facts).await;
1491 }
1492 Ok(())
1493 };
1494 if let Some(context) = observation {
1495 with_observation_context(
1496 context.with_purpose(ObservationPurpose::FactsExtraction),
1497 run,
1498 )
1499 .await
1500 } else {
1501 run.await
1502 }
1503 };
1504 self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1505 .await
1506 }
1507
1508 async fn schedule_relationship_background(&self) -> Result<()> {
1509 let policy = self
1510 .runtime_config
1511 .optimization
1512 .post_turn
1513 .relationships
1514 .clone();
1515 let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1516 return Ok(());
1517 };
1518 let Some(actor_id) = self.effective_actor_id() else {
1519 self.record_skipped_maintenance(
1520 "relationships",
1521 ObservationPurpose::RelationshipUpdate,
1522 "missing_actor",
1523 Some(&policy),
1524 );
1525 return Ok(());
1526 };
1527 let recent_messages = manager.config().auto_update.recent_messages;
1528 let messages = match self.memory.get_messages(Some(recent_messages)).await {
1529 Ok(messages) => messages,
1530 Err(error) => {
1531 warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1532 return Ok(());
1533 }
1534 };
1535 let storage = self.storage.read().clone();
1536 let hooks = Arc::clone(&self.hooks);
1537 let agent_id = self.info.id.clone();
1538 let observation = current_observation_context();
1539 let key = MaintenanceSequenceKey::actor(
1540 agent_id.clone(),
1541 actor_id.clone(),
1542 RuntimeTaskPurpose::PostTurnRelationship,
1543 );
1544 let actor_for_task = actor_id.clone();
1545 let task = async move {
1546 let run = async move {
1547 if manager.config().auto_update.enabled {
1548 let update = manager.auto_update(&actor_for_task, &messages).await?;
1549 if !update.changes.is_empty() {
1550 hooks
1551 .on_relationship_change(&actor_for_task, &update.changes)
1552 .await;
1553 }
1554 if let Some(ref event) = update.event {
1555 hooks.on_notable_event(&actor_for_task, event).await;
1556 }
1557 }
1558 if manager.config().persistence.enabled
1559 && let (Some(storage), Some(value)) =
1560 (storage, manager.relationship_as_value(&actor_for_task)?)
1561 {
1562 storage
1563 .save_relationship(&agent_id, &actor_for_task, &value)
1564 .await?;
1565 }
1566 Ok(())
1567 };
1568 if let Some(context) = observation {
1569 with_observation_context(
1570 context.with_purpose(ObservationPurpose::RelationshipUpdate),
1571 run,
1572 )
1573 .await
1574 } else {
1575 run.await
1576 }
1577 };
1578 self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1579 .await
1580 }
1581
1582 async fn spawn_or_handle_background<F>(
1584 &self,
1585 key: Option<MaintenanceSequenceKey>,
1586 task: F,
1587 label: &'static str,
1588 policy: &crate::optimization::config::MaintenanceTaskPolicy,
1589 ) -> Result<()>
1590 where
1591 F: Future<Output = Result<()>> + Send + 'static,
1592 {
1593 if self.background_maintenance.is_full() {
1594 match self
1595 .runtime_config
1596 .optimization
1597 .post_turn
1598 .on_background_overflow
1599 {
1600 BackgroundOverflowPolicy::RunInline => {
1601 record_background_maintenance_event(
1602 self.observability_manager.as_ref(),
1603 label,
1604 EventStatus::Success,
1605 0,
1606 "inline_overflow",
1607 None,
1608 Some(policy),
1609 );
1610 let start = Instant::now();
1611 match task.await {
1612 Ok(()) => record_background_maintenance_event(
1613 self.observability_manager.as_ref(),
1614 label,
1615 EventStatus::Success,
1616 start.elapsed().as_millis() as u64,
1617 "inline_completed",
1618 None,
1619 Some(policy),
1620 ),
1621 Err(error) => {
1622 warn!(label = label, error = %error, "inline maintenance fallback failed");
1623 record_background_maintenance_event(
1624 self.observability_manager.as_ref(),
1625 label,
1626 EventStatus::Error,
1627 start.elapsed().as_millis() as u64,
1628 "inline_failed",
1629 Some(error.to_string()),
1630 Some(policy),
1631 );
1632 return Err(error);
1633 }
1634 }
1635 }
1636 BackgroundOverflowPolicy::Drop => {
1637 self.record_skipped_maintenance(
1638 label,
1639 ObservationPurpose::Other(label.to_string()),
1640 "queue_full",
1641 Some(policy),
1642 );
1643 }
1644 BackgroundOverflowPolicy::Error => {
1645 record_background_maintenance_event(
1646 self.observability_manager.as_ref(),
1647 label,
1648 EventStatus::Error,
1649 0,
1650 "queue_full",
1651 None,
1652 Some(policy),
1653 );
1654 warn!(label = label, "background maintenance queue full");
1655 return Err(AgentError::Other(format!(
1656 "background maintenance queue is full for {}",
1657 label
1658 )));
1659 }
1660 }
1661 return Ok(());
1662 }
1663
1664 record_background_maintenance_event(
1665 self.observability_manager.as_ref(),
1666 label,
1667 EventStatus::Success,
1668 0,
1669 "scheduled",
1670 None,
1671 Some(policy),
1672 );
1673 let manager = self.observability_manager.clone();
1674 let policy_for_task = policy.clone();
1675 let observed_task = async move {
1676 let start = Instant::now();
1677 let result = task.await;
1678 match &result {
1679 Ok(()) => record_background_maintenance_event(
1680 manager.as_ref(),
1681 label,
1682 EventStatus::Success,
1683 start.elapsed().as_millis() as u64,
1684 "completed",
1685 None,
1686 Some(&policy_for_task),
1687 ),
1688 Err(error) => record_background_maintenance_event(
1689 manager.as_ref(),
1690 label,
1691 EventStatus::Error,
1692 start.elapsed().as_millis() as u64,
1693 "failed",
1694 Some(error.to_string()),
1695 Some(&policy_for_task),
1696 ),
1697 }
1698 result
1699 };
1700
1701 if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1702 record_background_maintenance_event(
1703 self.observability_manager.as_ref(),
1704 label,
1705 EventStatus::Error,
1706 0,
1707 "spawn_failed",
1708 Some(error.to_string()),
1709 Some(policy),
1710 );
1711 warn!(label = label, error = %error, "background maintenance spawn failed");
1712 return Err(error);
1713 }
1714 Ok(())
1715 }
1716
1717 fn record_skipped_maintenance(
1719 &self,
1720 label: &str,
1721 purpose: ObservationPurpose,
1722 reason: &str,
1723 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1724 ) {
1725 if let Some(manager) = self.observability_manager.as_ref() {
1726 let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1727 tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1728 manager.record_lifecycle_event(
1729 EventType::MemoryOperation {
1730 operation: format!("{}_maintenance", label),
1731 },
1732 purpose,
1733 EventStatus::Skipped,
1734 0,
1735 tags,
1736 None,
1737 );
1738 }
1739 }
1740
1741 pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1743 let Some(actor_id) = self.effective_actor_id() else {
1744 return Vec::new();
1745 };
1746 self.actor_facts_cache
1747 .read()
1748 .get(&actor_id)
1749 .cloned()
1750 .unwrap_or_default()
1751 }
1752
1753 pub fn relationship_memory_text(&self) -> Option<String> {
1755 self.format_relationship_for_context().map(|(_, text)| text)
1756 }
1757
1758 pub async fn extract_facts(
1760 &self,
1761 last_n: usize,
1762 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1763 self.extract_facts_with_source(last_n, "manual").await
1764 }
1765
1766 async fn extract_facts_with_source(
1767 &self,
1768 last_n: usize,
1769 source: &'static str,
1770 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1771 let extractor = match self.fact_extractor.read().clone() {
1772 Some(e) => e,
1773 None => return Ok(vec![]),
1774 };
1775
1776 let messages = self.memory.get_messages(None).await?;
1777 let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1778
1779 if recent.is_empty() {
1780 return Ok(vec![]);
1781 }
1782
1783 let actor_id = self.effective_actor_id();
1784 let existing = actor_id
1785 .as_ref()
1786 .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1787 .unwrap_or_default();
1788
1789 let categories = self
1790 .facts_config
1791 .as_ref()
1792 .map(|c| c.custom_categories.clone())
1793 .unwrap_or_default();
1794
1795 let facts = self
1796 .observe_purpose(
1797 ObservationPurpose::FactsExtraction,
1798 extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1799 )
1800 .await?;
1801
1802 if !facts.is_empty() {
1804 let fact_store_opt = self.fact_store.read().clone();
1805 let mut stored_total = 0usize;
1806 let mut cache_updated = false;
1807 if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1808 let authoritative = store.add_facts(aid, facts.clone()).await?;
1810 stored_total = authoritative.len();
1811 self.actor_facts_cache
1812 .write()
1813 .insert(aid.clone(), authoritative);
1814 cache_updated = true;
1815 } else if let Some(aid) = &actor_id {
1816 let mut cache = self.actor_facts_cache.write();
1817 let entry = cache.entry(aid.clone()).or_default();
1818 entry.extend(facts.clone());
1819 stored_total = entry.len();
1820 cache_updated = true;
1821 }
1822
1823 info!(
1824 actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1825 source = source,
1826 requested_messages = last_n,
1827 message_count = recent.len(),
1828 extracted_count = facts.len(),
1829 cache_updated = cache_updated,
1830 stored_total = stored_total,
1831 "facts extracted"
1832 );
1833
1834 if let Some(ref aid) = actor_id {
1835 self.hooks.on_facts_extracted(aid, &facts).await;
1836 }
1837 }
1838
1839 Ok(facts)
1840 }
1841
1842 fn resolve_actor_id_from_context(&self) {
1845 if self
1846 .current_turn_actor_context()
1847 .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1848 .is_some()
1849 {
1850 return;
1851 }
1852
1853 if let Some(ref am_config) = self.actor_memory_config
1854 && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1855 && let Some(ref path) = am_config.identification.context_path
1856 {
1857 let val = self
1859 .context_manager
1860 .get_path(path)
1861 .or_else(|| self.context_manager.get(path));
1862 if let Some(val) = val
1863 && let Some(id_str) = val.as_str()
1864 {
1865 let current = self.actor_id.read().clone();
1866 if current.as_deref() != Some(id_str) {
1867 *self.actor_id.write() = Some(id_str.to_string());
1868 let mut meta = self.session_metadata.write();
1869 meta.actor_id = Some(id_str.to_string());
1870 if !meta.actors.iter().any(|a| a == id_str) {
1871 meta.actors.push(id_str.to_string());
1872 }
1873 }
1874 }
1875 }
1876 }
1877
1878 fn format_actor_facts_for_context(&self) -> String {
1880 let should_inject = self
1882 .facts_config
1883 .as_ref()
1884 .map(|c| c.inject_in_context)
1885 .unwrap_or(true);
1886 if !should_inject {
1887 return String::new();
1888 }
1889
1890 let Some(actor_id) = self.effective_actor_id() else {
1891 return String::new();
1892 };
1893
1894 let facts = self
1895 .actor_facts_cache
1896 .read()
1897 .get(&actor_id)
1898 .cloned()
1899 .unwrap_or_default();
1900 if facts.is_empty() {
1901 return String::new();
1902 }
1903
1904 let am_config = self.actor_memory_config.as_ref();
1905 let facts_budget = self
1908 .memory_token_budget
1909 .as_ref()
1910 .map(|b| b.allocation.facts as usize)
1911 .filter(|n| *n > 0);
1912 let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1913 let max_tokens = facts_budget.unwrap_or(default_max);
1914
1915 let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1917 if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1918 return String::new();
1919 }
1920 if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1921 && !cfg.injection.categories.is_empty()
1922 {
1923 facts
1924 .iter()
1925 .filter(|f| {
1926 cfg.injection
1927 .categories
1928 .iter()
1929 .any(|c| f.category.to_string() == *c)
1930 })
1931 .cloned()
1932 .collect()
1933 } else {
1934 facts.clone()
1935 }
1936 } else {
1937 facts.clone()
1938 };
1939
1940 if filtered.is_empty() {
1941 return String::new();
1942 }
1943
1944 if let Some(store) = self.fact_store.read().clone() {
1945 store.format_for_context(&filtered, max_tokens)
1946 } else {
1947 String::new()
1948 }
1949 }
1950
1951 fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
1952 let context = self.build_context_with_overlays();
1953 let mut root = Value::Object(context.into_iter().collect());
1954 for (path, value) in staged {
1955 if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
1956 root = updated;
1957 }
1958 }
1959 match root {
1960 Value::Object(obj) => obj.into_iter().collect(),
1961 _ => HashMap::new(),
1962 }
1963 }
1964
1965 fn build_context_with_overlays(&self) -> HashMap<String, Value> {
1966 let mut context = self.context_manager.get_all();
1967 let mut root = Value::Object(context.clone().into_iter().collect());
1968
1969 if let Some(turn_ctx) = self.current_turn_actor_context() {
1970 if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
1971 && let Ok(updated) = ai_agents_core::set_dot_path(
1972 root.clone(),
1973 "interaction.origin_actor_id",
1974 serde_json::json!(origin_actor_id),
1975 )
1976 {
1977 root = updated;
1978 }
1979 if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
1980 && let Ok(updated) = ai_agents_core::set_dot_path(
1981 root.clone(),
1982 "interaction.sender_agent_id",
1983 serde_json::json!(sender_agent_id),
1984 )
1985 {
1986 root = updated;
1987 }
1988 }
1989
1990 if let Some(ref actor_id) = self.effective_actor_id()
1991 && let Ok(updated) = ai_agents_core::set_dot_path(
1992 root.clone(),
1993 "interaction.actor_id",
1994 serde_json::json!(actor_id),
1995 )
1996 {
1997 root = updated;
1998 }
1999
2000 if let Some(manager) = self.relationship_manager.as_ref()
2001 && let Some(actor_id) = self.effective_actor_id()
2002 && let Some(value) = manager.to_context_value(&actor_id)
2003 && let Ok(updated) = ai_agents_core::set_dot_path(
2004 root.clone(),
2005 &manager.config().injection.context_path,
2006 value,
2007 )
2008 {
2009 root = updated;
2010 }
2011
2012 if let Value::Object(obj) = root {
2013 context = obj.into_iter().collect();
2014 }
2015
2016 context
2017 }
2018
2019 fn resolve_actor_name_from_context(&self) -> Option<String> {
2020 for path in ["actor.name", "user.name", "player.name", "customer.name"] {
2021 if let Some(value) = self.context_manager.get_path(path)
2022 && let Some(name) = value.as_str()
2023 {
2024 return Some(name.to_string());
2025 }
2026 }
2027 None
2028 }
2029
2030 async fn maybe_load_actor_relationship(&self) {
2031 let Some(manager) = self.relationship_manager.as_ref() else {
2032 return;
2033 };
2034 let Some(actor_id) = self.effective_actor_id() else {
2035 return;
2036 };
2037
2038 let mut should_fire_loaded = false;
2039 if manager.get(&actor_id).is_none() {
2040 let mut loaded = false;
2041 if manager.config().persistence.enabled {
2042 let storage = self.storage.read().clone();
2043 if let Some(storage) = storage {
2044 match storage.load_relationship(&self.info.id, &actor_id).await {
2045 Ok(Some(value)) => match manager.insert_from_value(value) {
2046 Ok(_) => loaded = true,
2047 Err(e) => {
2048 warn!(actor = %actor_id, error = %e, "failed to restore relationship")
2049 }
2050 },
2051 Ok(None) => {}
2052 Err(e) => {
2053 warn!(actor = %actor_id, error = %e, "failed to load relationship")
2054 }
2055 }
2056 }
2057 }
2058
2059 if !loaded {
2060 manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
2061 }
2062 should_fire_loaded = true;
2063 }
2064
2065 let actor_name = self.resolve_actor_name_from_context();
2066 let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
2067 if should_fire_loaded {
2068 self.hooks
2069 .on_relationship_loaded(&actor_id, &relationship)
2070 .await;
2071 }
2072 }
2073
2074 fn format_relationship_for_context(&self) -> Option<(String, String)> {
2075 let manager = self.relationship_manager.as_ref()?;
2076 if !manager.config().injection.enabled {
2077 return None;
2078 }
2079 let actor_id = self.effective_actor_id()?;
2080 let relationship = manager.get(&actor_id)?;
2081 let local_cap = manager.config().injection.max_tokens;
2082 let global_cap = self
2083 .memory_token_budget
2084 .as_ref()
2085 .map(|b| b.allocation.relationships as usize)
2086 .filter(|n| *n > 0);
2087 let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
2088 let text = ai_agents_relationships::format_relationship(
2089 &relationship,
2090 &manager.config().injection.format,
2091 max_tokens,
2092 );
2093 if text.is_empty() {
2094 None
2095 } else {
2096 Some((manager.config().injection.prompt_variable.clone(), text))
2097 }
2098 }
2099
2100 async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2101 let Some(manager) = self.relationship_manager.as_ref() else {
2102 return Ok(());
2103 };
2104 if !manager.config().persistence.enabled {
2105 return Ok(());
2106 }
2107 let storage = self.storage.read().clone();
2108 let Some(storage) = storage else {
2109 return Ok(());
2110 };
2111 if let Some(value) = manager.relationship_as_value(actor_id)? {
2112 storage
2113 .save_relationship(&self.info.id, actor_id, &value)
2114 .await?;
2115 }
2116 Ok(())
2117 }
2118
2119 async fn auto_update_relationship(&self) {
2120 let Some(manager) = self.relationship_manager.as_ref() else {
2121 return;
2122 };
2123 let Some(actor_id) = self.effective_actor_id() else {
2124 return;
2125 };
2126 if !manager.config().auto_update.enabled {
2127 let _ = self.persist_actor_relationship(&actor_id).await;
2128 return;
2129 }
2130
2131 let recent_messages = manager.config().auto_update.recent_messages;
2132 let messages = match self.memory.get_messages(Some(recent_messages)).await {
2133 Ok(messages) => messages,
2134 Err(e) => {
2135 warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2136 return;
2137 }
2138 };
2139
2140 match self
2141 .observe_purpose(
2142 ObservationPurpose::RelationshipUpdate,
2143 manager.auto_update(&actor_id, &messages),
2144 )
2145 .await
2146 {
2147 Ok(update) => {
2148 if !update.changes.is_empty() {
2149 self.hooks
2150 .on_relationship_change(&actor_id, &update.changes)
2151 .await;
2152 }
2153 if let Some(ref event) = update.event {
2154 self.hooks.on_notable_event(&actor_id, event).await;
2155 }
2156 let persisted = match self.persist_actor_relationship(&actor_id).await {
2157 Ok(()) => true,
2158 Err(e) => {
2159 warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2160 false
2161 }
2162 };
2163 if !update.changes.is_empty() || update.event.is_some() {
2164 let changed_dimensions: Vec<String> = update
2165 .changes
2166 .iter()
2167 .map(|change| format!("{}:{}", change.perspective, change.dimension))
2168 .collect();
2169 info!(
2170 actor_id = %actor_id,
2171 change_count = update.changes.len(),
2172 changed_dimensions = ?changed_dimensions,
2173 event_present = update.event.is_some(),
2174 persisted = persisted,
2175 "relationship updated"
2176 );
2177 } else {
2178 debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2179 }
2180 }
2181 Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2182 }
2183 }
2184
2185 async fn auto_extract_facts(&self) {
2187 let should_extract = self
2188 .facts_config
2189 .as_ref()
2190 .map(|c| c.enabled && c.auto_extract)
2191 .unwrap_or(false);
2192
2193 if !should_extract {
2194 debug!("fact extraction skipped because auto extraction is disabled");
2195 return;
2196 }
2197
2198 let msgs_since = *self.messages_since_extraction.read();
2199 if msgs_since < 2 {
2200 debug!(
2201 messages_since_extraction = msgs_since,
2202 "fact extraction skipped until threshold is reached"
2203 );
2204 return;
2205 }
2206
2207 match self.extract_facts_with_source(msgs_since, "auto").await {
2208 Ok(facts) => {
2209 if !facts.is_empty() {
2210 *self.messages_since_extraction.write() = 0;
2211 } else {
2212 debug!("fact extraction ran but found no new facts");
2213 }
2214 }
2215 Err(e) => {
2216 warn!("fact extraction failed: {}", e);
2217 }
2218 }
2219 }
2220
2221 pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2222 self.persona_manager = Some(manager);
2223 self
2224 }
2225
2226 pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2227 self.persona_manager.as_ref()
2228 }
2229
2230 pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2231 if config.is_enabled() {
2232 let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2233 .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2234 self.disambiguation_manager = Some(manager);
2235 }
2236 self
2237 }
2238
2239 pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2240 self.disambiguation_manager.as_ref()
2241 }
2242
2243 pub fn has_disambiguation(&self) -> bool {
2244 self.disambiguation_manager
2245 .as_ref()
2246 .is_some_and(|m| m.is_enabled())
2247 }
2248
2249 pub async fn init_storage(&self) -> Result<()> {
2250 let _guard = self.storage_init.lock().await;
2254 let mut storage = self.storage.read().clone();
2255 if storage.is_none() && !self.storage_config.is_none() {
2256 let storage_config = self.convert_storage_config();
2257 storage = create_storage(&storage_config).await?;
2258 *self.storage.write() = storage.clone();
2259 }
2260
2261 self.validate_storage_requirements(storage.as_deref())?;
2262 self.complete_facts_init().await;
2263 Ok(())
2264 }
2265
2266 fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2267 let facts_required = self
2268 .facts_config
2269 .as_ref()
2270 .is_some_and(|config| config.enabled)
2271 || self
2272 .actor_memory_config
2273 .as_ref()
2274 .is_some_and(|config| config.enabled);
2275 let relationships_required = self
2276 .relationship_manager
2277 .as_ref()
2278 .is_some_and(|manager| manager.config().persistence.enabled);
2279
2280 let Some(storage) = storage else {
2281 let mut requirements = Vec::new();
2282 if facts_required {
2283 requirements.push("actor facts or actor memory");
2284 }
2285 if relationships_required {
2286 requirements.push("persistent relationships");
2287 }
2288 if requirements.is_empty() {
2289 return Ok(());
2290 }
2291 return Err(AgentError::Config(format!(
2292 "Storage is required for enabled {} but none is configured or injected",
2293 requirements.join(" and ")
2294 )));
2295 };
2296
2297 if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2301 return Err(AgentError::UnsupportedStorageCapability(
2302 StorageCapability::ActorFacts,
2303 ));
2304 }
2305 if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2306 return Err(AgentError::UnsupportedStorageCapability(
2307 StorageCapability::ActorRelationships,
2308 ));
2309 }
2310 Ok(())
2311 }
2312
2313 async fn complete_facts_init(&self) {
2316 if self.fact_store.read().is_some() {
2317 return;
2318 }
2319 let storage = match self.storage.read().clone() {
2320 Some(s) => s,
2321 None => return,
2322 };
2323
2324 let facts_enabled = self
2325 .facts_config
2326 .as_ref()
2327 .map(|f| f.enabled)
2328 .unwrap_or(false);
2329 let actor_memory_enabled = self
2330 .actor_memory_config
2331 .as_ref()
2332 .map(|a| a.enabled)
2333 .unwrap_or(false);
2334
2335 if !facts_enabled && !actor_memory_enabled {
2336 return;
2337 }
2338
2339 let fc = self.facts_config.clone().unwrap_or_default();
2340 let store = Arc::new(ai_agents_facts::FactStore::new(
2341 storage,
2342 self.info.id.clone(),
2343 fc.clone(),
2344 ));
2345
2346 let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2347 let extractor_llm = fc
2348 .extractor_llm
2349 .as_ref()
2350 .and_then(|alias| self.llm_registry.get(alias).ok())
2351 .or_else(|| self.llm_registry.router().ok())
2352 .or_else(|| self.llm_registry.default().ok());
2353 extractor_llm.map(|llm| {
2354 Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2355 as Arc<dyn ai_agents_facts::FactExtractor>
2356 })
2357 } else {
2358 None
2359 };
2360
2361 *self.fact_store.write() = Some(store);
2362 *self.fact_extractor.write() = extractor;
2363 debug!(
2364 agent = %self.info.id,
2365 facts_enabled,
2366 actor_memory_enabled,
2367 "facts storage initialized"
2368 );
2369 }
2370
2371 fn convert_storage_config(&self) -> StorageStorageConfig {
2372 crate::spec::storage::to_storage_config(&self.storage_config)
2373 }
2374
2375 pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2376 self.storage.read().clone()
2377 }
2378
2379 pub fn storage_config(&self) -> &StorageConfig {
2380 &self.storage_config
2381 }
2382
2383 pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2385 self.spawner.as_ref()
2386 }
2387
2388 pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2390 self.spawner_registry.as_ref()
2391 }
2392
2393 pub fn has_spawner(&self) -> bool {
2394 self.spawner_registry.is_some()
2395 }
2396
2397 pub fn with_spawner_handles(
2398 mut self,
2399 spawner: Arc<crate::spawner::AgentSpawner>,
2400 registry: Arc<crate::spawner::AgentRegistry>,
2401 ) -> Self {
2402 self.spawner = Some(spawner);
2403 self.spawner_registry = Some(registry);
2404 self
2405 }
2406
2407 pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2408 self.hooks = hooks;
2409 self
2410 }
2411
2412 pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2413 self.parallel_tools = config;
2414 self
2415 }
2416
2417 pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2418 self.streaming = config;
2419 self
2420 }
2421
2422 pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2423 self.hitl_engine = Some(engine);
2424 self.approval_handler = handler;
2425 self
2426 }
2427
2428 pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2429 self.max_context_tokens = tokens;
2430 self
2431 }
2432
2433 pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2434 self.memory_token_budget = Some(budget);
2435 self
2436 }
2437
2438 pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2439 self.recovery_manager = manager;
2440 self
2441 }
2442
2443 pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2444 self.tool_security = engine;
2445 self
2446 }
2447
2448 pub fn runtime_control(&self) -> RuntimeControlHandle {
2450 RuntimeControlHandle {
2451 state: Arc::clone(&self.runtime_control),
2452 }
2453 }
2454
2455 pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2457 self.tools.set_question_handler(handler);
2458 }
2459
2460 pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2462 self.tools.set_diagnostics_provider(provider);
2463 }
2464
2465 pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2467 self.tools.set_command_runner(runner);
2468 }
2469
2470 pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2472 self.tools.set_web_search_provider(provider);
2473 }
2474
2475 pub fn todos(&self) -> Vec<TodoItem> {
2477 self.tools.todos()
2478 }
2479
2480 fn active_tool_security(&self) -> ToolSecurityEngine {
2482 self.runtime_control
2483 .tool_security_override
2484 .read()
2485 .clone()
2486 .unwrap_or_else(|| self.tool_security.clone())
2487 }
2488
2489 fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2491 let _guard = self.runtime_control.snapshot_guard.read();
2492 RuntimeSafetySnapshot {
2493 version: self.runtime_control.version.load(Ordering::SeqCst),
2494 emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2495 tool_security: self
2496 .runtime_control
2497 .tool_security_override
2498 .read()
2499 .clone()
2500 .unwrap_or_else(|| self.tool_security.clone()),
2501 tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2502 }
2503 }
2504
2505 fn admit_tool_execution(
2507 &self,
2508 expected_runtime_version: u64,
2509 expected_policy_version: u64,
2510 expected_state_generation: Option<u64>,
2511 canonical_id: &str,
2512 ) -> SecurityCheckResult {
2513 let _guard = self.runtime_control.snapshot_guard.read();
2514 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2515 return SecurityCheckResult::Block {
2516 reason: "runtime emergency deny is enabled".to_string(),
2517 };
2518 }
2519 let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2520 let security_engine = self
2521 .runtime_control
2522 .tool_security_override
2523 .read()
2524 .clone()
2525 .unwrap_or_else(|| self.tool_security.clone());
2526 if runtime_version != expected_runtime_version
2527 || security_engine.policy_version() != expected_policy_version
2528 {
2529 return SecurityCheckResult::Block {
2530 reason: "runtime safety controls changed before admission".to_string(),
2531 };
2532 }
2533 let current_state_generation = self
2534 .state_machine
2535 .as_ref()
2536 .map(|state_machine| state_machine.generation());
2537 if current_state_generation != expected_state_generation {
2538 return SecurityCheckResult::Block {
2539 reason: "state scope changed before admission".to_string(),
2540 };
2541 }
2542 security_engine.admit_tool_execution(canonical_id)
2543 }
2544
2545 pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2546 let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2547 self.process_processor = Some(processor);
2548 self
2549 }
2550
2551 pub fn with_state_machine(
2552 mut self,
2553 state_machine: Arc<StateMachine>,
2554 evaluator: Arc<dyn TransitionEvaluator>,
2555 ) -> Self {
2556 self.state_machine = Some(state_machine);
2557 self.transition_evaluator = Some(evaluator);
2558 self
2559 }
2560
2561 pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2562 self.context_manager = manager;
2563 self
2564 }
2565
2566 pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2567 self.message_filters.write().insert(name.into(), filter);
2568 }
2569
2570 pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2571 self.context_manager.update(key, value)
2572 }
2573
2574 pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2575 self.context_manager.update(path, value)
2576 }
2577
2578 pub fn get_context(&self) -> HashMap<String, Value> {
2579 self.build_context_with_overlays()
2580 }
2581
2582 pub fn remove_context(&self, key: &str) -> Option<Value> {
2583 self.context_manager.remove(key)
2584 }
2585
2586 pub async fn refresh_context(&self, key: &str) -> Result<()> {
2587 self.context_manager.refresh(key).await
2588 }
2589
2590 pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2591 self.context_manager.register_provider(name, provider);
2592 }
2593
2594 pub fn current_state(&self) -> Option<String> {
2595 self.state_machine.as_ref().map(|sm| sm.current())
2596 }
2597
2598 async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2600 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2601 let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2602 return;
2603 };
2604 if disambiguator.has_pending_confirmation().await {
2605 disambiguator.clear_pending().await;
2606 *self.pending_skill_id.write() = None;
2607 info!(
2608 confirmation_event = "invalidated",
2609 invalidation_reason = reason,
2610 "Runtime invalidated pending confirmation"
2611 );
2612 }
2613 }
2614
2615 async fn admit_disambiguation_redispatch(
2617 &self,
2618 expected_epoch: u64,
2619 expected_state_generation: Option<u64>,
2620 ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2621 let admission = self.disambiguation_admission.read().await;
2622 let state_generation = self
2623 .state_machine
2624 .as_ref()
2625 .map(|state_machine| state_machine.generation());
2626 if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2627 || state_generation != expected_state_generation
2628 {
2629 return Err(AgentError::Other(
2630 "Disambiguation ownership changed before redispatch admission".to_string(),
2631 ));
2632 }
2633 Ok(admission)
2634 }
2635
2636 fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2638 self.state_transition_reserved
2639 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2640 .ok()
2641 .map(|_| StateTransitionReservation {
2642 reserved: &self.state_transition_reserved,
2643 })
2644 }
2645
2646 async fn admit_optional_disambiguation_ownership(
2648 &self,
2649 ownership: Option<DisambiguationOwnership>,
2650 ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2651 match ownership {
2652 Some(ownership) => self
2653 .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2654 .await
2655 .map(Some),
2656 None => Ok(None),
2657 }
2658 }
2659
2660 pub async fn transition_to(&self, state: &str) -> Result<()> {
2662 let Some(ref sm) = self.state_machine else {
2663 return Ok(());
2664 };
2665 let claim_admission = self.disambiguation_admission.write().await;
2666 let reservation = self.reserve_state_transition().ok_or_else(|| {
2667 AgentError::Other("Another state transition is already in progress".to_string())
2668 })?;
2669 let from_state = sm.current();
2670 let expected_state_generation = sm.generation();
2671 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2672 let history_before = sm.history();
2673 drop(claim_admission);
2674
2675 self.execute_state_exit_actions(&from_state).await;
2676
2677 let admission = self.disambiguation_admission.write().await;
2678 if sm.current() != from_state
2679 || sm.generation() != expected_state_generation
2680 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2681 {
2682 return Err(AgentError::Other(
2683 "State ownership changed during manual transition preparation".to_string(),
2684 ));
2685 }
2686 sm.transition_to(state, "manual transition")?;
2687 self.invalidate_pending_confirmation("state_transition")
2688 .await;
2689 let entered = sm.current();
2690 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2691 drop(admission);
2692
2693 self.execute_state_enter_actions(&entered, is_reentry).await;
2694 drop(reservation);
2695 info!(to = %entered, "Manual state transition");
2696 Ok(())
2697 }
2698
2699 pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2700 self.state_machine
2701 .as_ref()
2702 .map(|sm| sm.history())
2703 .unwrap_or_default()
2704 }
2705
2706 pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2708 self.session_metadata.read().clone()
2709 }
2710
2711 pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2714 let allowed = self
2715 .actor_memory_config
2716 .as_ref()
2717 .map(|c| c.privacy.allow_deletion)
2718 .unwrap_or(true);
2719 if !allowed {
2720 return Err(AgentError::Config(
2721 "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2722 ));
2723 }
2724 let storage = self.storage.read().clone();
2725 if let Some(storage) = storage {
2726 if !storage.supports(StorageCapability::ActorDataDeletion) {
2730 return Err(AgentError::UnsupportedStorageCapability(
2731 StorageCapability::ActorDataDeletion,
2732 ));
2733 }
2734 storage.delete_actor_data(&self.info.id, actor_id).await?;
2735 } else {
2736 let store = { self.fact_store.read().clone() };
2740 if let Some(store) = store {
2741 store.delete_actor_data(actor_id).await?;
2742 }
2743 }
2744 if let Some(manager) = self.relationship_manager.as_ref() {
2745 manager.remove(actor_id);
2746 }
2747 self.actor_facts_cache.write().remove(actor_id);
2748 Ok(())
2749 }
2750
2751 pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2753 *self.session_metadata.write() = meta;
2754 }
2755
2756 pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2758 let storage = self.storage.read().clone();
2759 match storage {
2760 Some(s) => {
2761 let count = s.cleanup_expired().await?;
2762 if count > 0 {
2763 self.hooks.on_sessions_expired(count).await;
2764 }
2765 Ok(count)
2766 }
2767 None => Err(AgentError::Config(
2768 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2769 )),
2770 }
2771 }
2772
2773 pub async fn list_sessions_filtered(
2775 &self,
2776 filter: &ai_agents_core::SessionFilter,
2777 ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2778 let storage = self.storage.read().clone();
2779 match storage {
2780 Some(s) => s.list_sessions_filtered(filter).await,
2781 None => Err(AgentError::Config(
2782 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2783 )),
2784 }
2785 }
2786
2787 pub async fn save_state(&self) -> Result<AgentSnapshot> {
2788 let memory_snapshot = self.memory.snapshot().await?;
2789 let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2790 let context_snapshot = self.context_manager.snapshot();
2791
2792 let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2793 .with_memory(memory_snapshot)
2794 .with_context(context_snapshot)
2795 .with_state_machine(
2796 state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2797 current_state: String::new(),
2798 previous_state: None,
2799 turn_count: 0,
2800 no_transition_count: 0,
2801 history: vec![],
2802 }),
2803 );
2804
2805 if let Some(ref persona) = self.persona_manager {
2806 snapshot.persona = Some(persona.snapshot_as_value()?);
2807 }
2808
2809 if let Some(ref relationships) = self.relationship_manager {
2810 snapshot.relationships = Some(relationships.snapshot_as_value()?);
2811 }
2812
2813 Ok(snapshot)
2814 }
2815
2816 pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2818 let mut snapshot = self.save_state().await?;
2819 if let Some(ref registry) = self.spawner_registry {
2820 let entries = registry.list_with_specs();
2821 if !entries.is_empty() {
2822 snapshot = snapshot.with_spawned_agents(entries);
2823 }
2824 }
2825 Ok(snapshot)
2826 }
2827
2828 pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2830 let _admission = self.disambiguation_admission.write().await;
2831 if self.state_transition_reserved.load(Ordering::SeqCst) {
2832 return Err(AgentError::Other(
2833 "Cannot restore state while a state transition is in progress".to_string(),
2834 ));
2835 }
2836 self.invalidate_pending_confirmation("state_restore").await;
2837 *self.pending_skill_id.write() = None;
2838 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2839 disambiguator.clear_pending().await;
2840 }
2841 self.memory.restore(snapshot.memory).await?;
2842
2843 if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2844 && !sm_snapshot.current_state.is_empty()
2845 {
2846 sm.restore(sm_snapshot)?;
2847 }
2848
2849 self.context_manager.restore(snapshot.context);
2850
2851 if let (Some(persona_value), Some(persona_manager)) =
2852 (snapshot.persona, &self.persona_manager)
2853 {
2854 persona_manager.restore_from_value(persona_value)?;
2855 }
2856
2857 if let (Some(relationship_value), Some(relationship_manager)) =
2858 (snapshot.relationships, &self.relationship_manager)
2859 {
2860 relationship_manager.restore_from_value(relationship_value)?;
2861 }
2862
2863 info!(agent_id = %snapshot.agent_id, "State restored");
2864 Ok(())
2865 }
2866
2867 pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2868 let snapshot = self.save_state().await?;
2869 storage.save(session_id, &snapshot).await
2870 }
2871
2872 async fn load_session_restore(
2873 storage: &dyn AgentStorage,
2874 session_id: &str,
2875 ) -> Result<Option<StoredSessionRestore>> {
2876 let Some(snapshot) = storage.load(session_id).await? else {
2877 return Ok(None);
2878 };
2879 let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2883 storage.load_metadata(session_id).await?
2884 } else {
2885 None
2886 };
2887 Ok(Some(StoredSessionRestore { snapshot, metadata }))
2888 }
2889
2890 async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2891 Ok(RuntimeSessionRestorePoint {
2892 snapshot: self.save_state().await?,
2893 metadata: self.session_metadata(),
2894 actor_id: self.actor_id(),
2895 session_id: self.current_session_id.read().clone(),
2896 })
2897 }
2898
2899 async fn apply_session_restore_unchecked(
2900 &self,
2901 session_id: &str,
2902 stored: StoredSessionRestore,
2903 ) -> Result<()> {
2904 self.restore_state(stored.snapshot).await?;
2905 let metadata = stored.metadata.unwrap_or_default();
2906 if let Some(actor_id) = metadata.actor_id.as_deref() {
2907 self.set_actor_id(actor_id)?;
2908 } else {
2909 self.clear_actor_id();
2910 }
2911 self.set_session_metadata(metadata);
2912 *self.current_session_id.write() = Some(session_id.to_string());
2913 Ok(())
2914 }
2915
2916 async fn restore_session_restore_point(
2917 &self,
2918 restore_point: &RuntimeSessionRestorePoint,
2919 ) -> Result<()> {
2920 self.restore_state(restore_point.snapshot.clone()).await?;
2921 if let Some(actor_id) = restore_point.actor_id.as_deref() {
2922 self.set_actor_id(actor_id)?;
2923 } else {
2924 self.clear_actor_id();
2925 }
2926 self.set_session_metadata(restore_point.metadata.clone());
2927 *self.current_session_id.write() = restore_point.session_id.clone();
2928 Ok(())
2929 }
2930
2931 async fn apply_session_restore(
2932 &self,
2933 session_id: &str,
2934 stored: StoredSessionRestore,
2935 ) -> Result<()> {
2936 let before = self.capture_session_restore_point().await?;
2937 if let Err(error) = self
2938 .apply_session_restore_unchecked(session_id, stored)
2939 .await
2940 {
2941 return match self.restore_session_restore_point(&before).await {
2942 Ok(()) => Err(error),
2943 Err(rollback_error) => Err(AgentError::Other(format!(
2944 "Session restore failed: {error}; rollback failed: {rollback_error}"
2945 ))),
2946 };
2947 }
2948 Ok(())
2949 }
2950
2951 async fn rollback_session_restore_set(
2952 parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
2953 children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
2954 ) -> Vec<String> {
2955 let mut errors = Vec::new();
2956 if let Some((agent, restore_point)) = parent
2957 && let Err(error) = agent.restore_session_restore_point(restore_point).await
2958 {
2959 errors.push(format!("parent: {error}"));
2960 }
2961 for (id, agent, restore_point) in children {
2962 if let Err(error) = agent.restore_session_restore_point(restore_point).await {
2963 errors.push(format!("child '{id}': {error}"));
2964 }
2965 }
2966 errors
2967 }
2968
2969 fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
2970 if rollback_errors.is_empty() {
2971 AgentError::Other(format!(
2972 "Session restore failed: {error}; runtime state was rolled back"
2973 ))
2974 } else {
2975 AgentError::Other(format!(
2976 "Session restore failed: {error}; rollback also failed for {}",
2977 rollback_errors.join(", ")
2978 ))
2979 }
2980 }
2981
2982 pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
2983 let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
2984 return Ok(false);
2985 };
2986 self.apply_session_restore(session_id, stored).await?;
2987 Ok(true)
2988 }
2989
2990 pub async fn save_session(&self, session_id: &str) -> Result<()> {
2991 let storage = self.storage.read().clone();
2992 match storage {
2993 Some(s) => {
2994 let is_new = {
2996 let cur = self.current_session_id.read().clone();
2997 cur.as_deref() != Some(session_id)
2998 };
2999 if is_new {
3000 *self.current_session_id.write() = Some(session_id.to_string());
3001 self.hooks.on_session_created(session_id).await;
3002 }
3003
3004 {
3006 let now = chrono::Utc::now();
3007 let msg_count = self
3008 .memory
3009 .get_messages(None)
3010 .await
3011 .map(|v| v.len())
3012 .unwrap_or(0);
3013 let mut meta = self.session_metadata.write();
3014 meta.last_active = now;
3015 meta.message_count = msg_count;
3016 if meta.actor_id.is_none() {
3017 meta.actor_id = self.actor_id.read().clone();
3018 }
3019 }
3020
3021 let snapshot = self.save_state().await?;
3022 if s.supports(StorageCapability::SessionMetadata) {
3026 let metadata = self.session_metadata.read().clone();
3027 s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
3028 .await
3029 } else {
3030 s.save(session_id, &snapshot).await
3031 }
3032 }
3033 None => Err(AgentError::Config(
3034 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3035 )),
3036 }
3037 }
3038
3039 pub async fn load_session(&self, session_id: &str) -> Result<bool> {
3040 let storage = self.storage.read().clone();
3041 match storage {
3042 Some(storage) => self.load_from(storage.as_ref(), session_id).await,
3043 None => Err(AgentError::Config(
3044 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3045 )),
3046 }
3047 }
3048
3049 pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
3051 self.init_storage().await?;
3052 let storage = self.storage.read().clone().ok_or_else(|| {
3053 AgentError::Config(
3054 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3055 )
3056 })?;
3057 let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
3058 .await?
3059 .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
3060 let manifest = target_parent
3061 .snapshot
3062 .spawned_agents
3063 .clone()
3064 .unwrap_or_default();
3065
3066 let registry = self.spawner_registry.as_ref().cloned();
3067 let spawner = if manifest.is_empty() {
3068 self.spawner.as_ref().cloned()
3069 } else {
3070 Some(self.spawner.as_ref().cloned().ok_or_else(|| {
3071 AgentError::Config(
3072 "Saved session contains child agents but this runtime has no spawner".into(),
3073 )
3074 })?)
3075 };
3076 let registry = if manifest.is_empty() {
3077 registry
3078 } else {
3079 Some(registry.ok_or_else(|| {
3080 AgentError::Config(
3081 "Saved session contains child agents but this runtime has no registry".into(),
3082 )
3083 })?)
3084 };
3085
3086 let mut target_ids = HashSet::with_capacity(manifest.len());
3087 let mut prepared = Vec::with_capacity(manifest.len());
3088 for entry in manifest {
3089 if !target_ids.insert(entry.id.clone()) {
3090 return Err(AgentError::InvalidSpec(format!(
3091 "Saved child manifest contains duplicate ID: {}",
3092 entry.id
3093 )));
3094 }
3095 let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
3096 spawner
3097 .as_ref()
3098 .expect("non-empty manifests require a spawner")
3099 .validate_explicit_child(&entry.id, &spec)?;
3100 prepared.push((entry.id, spec));
3101 }
3102
3103 let current_ids = registry
3104 .as_ref()
3105 .map(|registry| {
3106 registry
3107 .list()
3108 .into_iter()
3109 .map(|info| info.id)
3110 .collect::<HashSet<_>>()
3111 })
3112 .unwrap_or_default();
3113 let removal_count = current_ids.difference(&target_ids).count();
3114 let additions = prepared
3115 .iter()
3116 .filter(|(id, _)| !current_ids.contains(id))
3117 .cloned()
3118 .collect::<Vec<_>>();
3119
3120 let mut existing = Vec::new();
3121 if let Some(registry) = registry.as_ref() {
3122 for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3123 let agent = registry.get(id).ok_or_else(|| {
3124 AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3125 })?;
3126 let child_storage = agent.storage().ok_or_else(|| {
3127 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3128 })?;
3129 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3130 .await?
3131 .ok_or_else(|| {
3132 AgentError::Persistence(format!(
3133 "Child '{id}' has no saved session '{session_id}'"
3134 ))
3135 })?;
3136 existing.push((id.clone(), agent, stored));
3137 }
3138 }
3139
3140 let mut staged = Vec::with_capacity(additions.len());
3141 if !additions.is_empty() {
3142 let spawner = spawner
3143 .as_ref()
3144 .expect("restored additions require a spawner");
3145 let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3146 for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3147 let spawned = spawner
3148 .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3149 .await?;
3150 let child_storage = spawned.agent.storage().ok_or_else(|| {
3151 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3152 })?;
3153 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3154 .await?
3155 .ok_or_else(|| {
3156 AgentError::Persistence(format!(
3157 "Child '{id}' has no saved session '{session_id}'"
3158 ))
3159 })?;
3160 staged.push((spawned, stored));
3161 }
3162 } else if let Some(spawner) = spawner.as_ref() {
3163 spawner.reserve_restore_capacity(0, removal_count)?;
3164 }
3165
3166 let parent_before = self.capture_session_restore_point().await?;
3167 let mut existing_before = Vec::with_capacity(existing.len());
3168 for (id, agent, _) in &existing {
3169 existing_before.push((
3170 id.clone(),
3171 Arc::clone(agent),
3172 agent.capture_session_restore_point().await?,
3173 ));
3174 }
3175
3176 for (_, agent, stored) in &existing {
3180 if let Err(error) = agent
3181 .apply_session_restore_unchecked(session_id, stored.clone())
3182 .await
3183 {
3184 drop(staged);
3185 let rollback_errors =
3186 Self::rollback_session_restore_set(None, &existing_before).await;
3187 return Err(Self::restore_failure(error, rollback_errors));
3188 }
3189 }
3190 for (spawned, stored) in &staged {
3191 if let Err(error) = spawned
3192 .agent
3193 .apply_session_restore_unchecked(session_id, stored.clone())
3194 .await
3195 {
3196 drop(staged);
3197 let rollback_errors =
3198 Self::rollback_session_restore_set(None, &existing_before).await;
3199 return Err(Self::restore_failure(error, rollback_errors));
3200 }
3201 }
3202 if let Err(error) = self
3203 .apply_session_restore_unchecked(session_id, target_parent)
3204 .await
3205 {
3206 drop(staged);
3207 let rollback_errors =
3208 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3209 .await;
3210 return Err(Self::restore_failure(error, rollback_errors));
3211 }
3212
3213 if let Some(registry) = registry.as_ref()
3214 && let Err(error) = registry
3215 .reconcile(
3216 &target_ids,
3217 staged.into_iter().map(|(spawned, _)| spawned).collect(),
3218 )
3219 .await
3220 {
3221 let rollback_errors =
3222 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3223 .await;
3224 return Err(Self::restore_failure(error, rollback_errors));
3225 }
3226
3227 Ok(target_ids.len())
3228 }
3229
3230 pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3231 let storage = self.storage.read().clone();
3232 match storage {
3233 Some(s) => s.delete(session_id).await,
3234 None => Err(AgentError::Config(
3235 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3236 )),
3237 }
3238 }
3239
3240 pub async fn list_sessions(&self) -> Result<Vec<String>> {
3241 let storage = self.storage.read().clone();
3242 match storage {
3243 Some(s) => s.list_sessions().await,
3244 None => Err(AgentError::Config(
3245 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3246 )),
3247 }
3248 }
3249
3250 fn estimate_tokens(&self, text: &str) -> u32 {
3251 (text.len() as f32 / 4.0).ceil() as u32
3252 }
3253
3254 fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3255 messages
3256 .iter()
3257 .map(|m| self.estimate_tokens(&m.content))
3258 .sum()
3259 }
3260
3261 fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3262 if messages.len() <= keep_recent + 1 {
3263 return;
3264 }
3265 let system_msg = messages.remove(0);
3266 let to_remove = messages.len().saturating_sub(keep_recent);
3267 messages.drain(..to_remove);
3268 messages.insert(0, system_msg);
3269 }
3270
3271 fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3272 match config {
3273 FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3274 FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3275 FilterConfig::SkipPattern { skip_if_contains } => {
3276 Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3277 }
3278 FilterConfig::Custom { name } => {
3279 let filters = self.message_filters.read();
3280 filters
3281 .get(name)
3282 .cloned()
3283 .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3284 }
3285 }
3286 }
3287
3288 async fn summarize_context(
3289 &self,
3290 messages: &mut Vec<ChatMessage>,
3291 summarizer_llm: Option<&str>,
3292 max_summary_tokens: u32,
3293 custom_prompt: Option<&str>,
3294 keep_recent: usize,
3295 filter: Option<&FilterConfig>,
3296 ) -> Result<()> {
3297 let system_msg = messages.remove(0);
3298
3299 let to_summarize_count = messages.len().saturating_sub(keep_recent);
3300 if to_summarize_count == 0 {
3301 messages.insert(0, system_msg);
3302 return Ok(());
3303 }
3304
3305 let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3306 let mut to_summarize = std::mem::take(messages);
3307
3308 if let Some(filter_config) = filter {
3309 let filter = self.get_filter(filter_config);
3310 to_summarize = filter.filter(to_summarize);
3311 }
3312
3313 if to_summarize.is_empty() {
3314 *messages = recent_msgs;
3315 messages.insert(0, system_msg);
3316 return Ok(());
3317 }
3318
3319 let conversation_text = to_summarize
3320 .iter()
3321 .map(|m| format!("{:?}: {}", m.role, m.content))
3322 .collect::<Vec<_>>()
3323 .join("\n");
3324
3325 let default_prompt = format!(
3326 "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3327 max_summary_tokens, conversation_text
3328 );
3329
3330 let summary_prompt = custom_prompt
3331 .map(|p| format!("{}\n\n{}", p, conversation_text))
3332 .unwrap_or(default_prompt);
3333
3334 let summarizer = if let Some(alias) = summarizer_llm {
3335 self.llm_registry
3336 .get(alias)
3337 .map_err(|e| AgentError::Config(e.to_string()))?
3338 } else {
3339 self.llm_registry
3340 .router()
3341 .or_else(|_| self.llm_registry.default())
3342 .map_err(|e| AgentError::Config(e.to_string()))?
3343 };
3344
3345 let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3346 let response = self
3347 .observe_purpose(
3348 ObservationPurpose::Summarization,
3349 summarizer.complete(&summary_msgs, None),
3350 )
3351 .await?;
3352
3353 let summary_message = ChatMessage::system(format!(
3354 "[Previous conversation summary]\n{}",
3355 response.content
3356 ));
3357
3358 *messages = vec![system_msg, summary_message];
3359 messages.extend(recent_msgs);
3360
3361 debug!(
3362 summarized_count = to_summarize_count,
3363 kept_recent = keep_recent,
3364 "Context summarized"
3365 );
3366
3367 Ok(())
3368 }
3369
3370 fn render_system_prompt(&self) -> Result<String> {
3371 let mut context = self.build_context_with_overlays();
3372
3373 let facts_text = self.format_actor_facts_for_context();
3375 if !facts_text.is_empty() {
3376 context.insert(
3377 "actor_facts".to_string(),
3378 serde_json::Value::String(facts_text),
3379 );
3380 }
3381
3382 if let Some((key, text)) = self.format_relationship_for_context() {
3383 context.insert(key, serde_json::Value::String(text));
3384 }
3385
3386 self.template_renderer
3387 .render(&self.base_system_prompt, &context)
3388 }
3389
3390 fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3392 let mut seen = HashSet::new();
3393 ids.iter()
3394 .filter_map(|id| self.tools.canonical_id(id))
3395 .filter(|canonical_id| seen.insert(canonical_id.clone()))
3396 .collect()
3397 }
3398
3399 fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3401 let Some(declared) = self.declared_tool_ids.as_deref() else {
3402 return Vec::new();
3403 };
3404 let mut effective = self.canonical_unique_tool_ids(declared);
3405 if let Some(scope) = scope_override {
3406 let scope: HashSet<String> =
3407 self.canonical_unique_tool_ids(scope).into_iter().collect();
3408 effective.retain(|canonical_id| scope.contains(canonical_id));
3409 }
3410 effective
3411 }
3412
3413 async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3415 Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3416 }
3417
3418 async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3420 let scope_override = self.runtime_control.tool_scope_override.read().clone();
3421 self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3422 .await
3423 }
3424
3425 async fn get_available_tool_ids_snapshot_for_scope(
3427 &self,
3428 scope_override: Option<&[String]>,
3429 ) -> Result<AvailableToolIdsSnapshot> {
3430 let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3431 let (state_generation, state_scopes) = self
3432 .state_machine
3433 .as_ref()
3434 .map(|state_machine| {
3435 let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3436 (Some(generation), scopes)
3437 })
3438 .unwrap_or((None, Vec::new()));
3439
3440 if available.is_empty() || state_scopes.is_empty() {
3441 return Ok(AvailableToolIdsSnapshot {
3442 tool_ids: available,
3443 state_generation,
3444 });
3445 }
3446
3447 let eval_ctx = self.build_evaluation_context().await?;
3448 let llm_getter = RegistryLLMGetter {
3449 registry: self.llm_registry.clone(),
3450 };
3451 let evaluator = ConditionEvaluator::new(llm_getter);
3452
3453 for state_scope in state_scopes {
3454 if state_scope.is_empty() {
3455 available.clear();
3456 break;
3457 }
3458
3459 let mut allowed = HashSet::new();
3460 for tool_ref in &state_scope {
3461 let tool_id = tool_ref.id();
3462 let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3463 continue;
3464 };
3465 let condition_matches = if let Some(condition) = tool_ref.condition() {
3466 match evaluator.evaluate(condition, &eval_ctx).await {
3467 Ok(matches) => matches,
3468 Err(error) => {
3469 warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3470 false
3471 }
3472 }
3473 } else {
3474 true
3475 };
3476 if condition_matches {
3477 allowed.insert(canonical_id);
3478 } else {
3479 debug!(tool = tool_id, "Tool condition not met, skipping");
3480 }
3481 }
3482 available.retain(|canonical_id| allowed.contains(canonical_id));
3483 if available.is_empty() {
3484 break;
3485 }
3486 }
3487
3488 Ok(AvailableToolIdsSnapshot {
3489 tool_ids: available,
3490 state_generation,
3491 })
3492 }
3493
3494 async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3495 let context = self.build_context_with_overlays();
3496 let messages = self.memory.get_messages(Some(10)).await?;
3497 let tool_history = self.tool_call_history.read().clone();
3498
3499 let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3500 (Some(sm.current()), sm.turn_count(), sm.previous())
3501 } else {
3502 (None, 0, None)
3503 };
3504
3505 Ok(EvaluationContext::default()
3506 .with_context(context)
3507 .with_state(state_name, turn_count, previous_state)
3508 .with_called_tools(tool_history)
3509 .with_messages(messages))
3510 }
3511
3512 fn record_tool_call(&self, tool_id: &str, result: Value) {
3513 self.tool_call_history.write().push(ToolCallRecord {
3514 tool_id: tool_id.to_string(),
3515 result,
3516 timestamp: chrono::Utc::now(),
3517 });
3518 }
3519
3520 async fn get_effective_system_prompt_with_persona_hooks(
3521 &self,
3522 fire_persona_hooks: bool,
3523 include_tool_prompt: bool,
3524 ) -> Result<String> {
3525 let rendered_base = self.render_system_prompt()?;
3526
3527 let persona_prefix = if let Some(ref persona) = self.persona_manager {
3528 let context = self.build_context_with_overlays();
3529 if fire_persona_hooks {
3530 let render_result = persona.render_prompt(&context)?;
3531 for content in &render_result.newly_revealed {
3532 self.hooks.on_secret_revealed(content).await;
3533 }
3534 render_result.prompt
3535 } else {
3536 persona.render_prompt_preview(&context)?
3537 }
3538 } else {
3539 String::new()
3540 };
3541
3542 if let Some(ref sm) = self.state_machine
3543 && let Some(state_def) = sm.current_definition()
3544 {
3545 let state_prompt = if let Some(ref prompt) = state_def.prompt {
3546 let context = self.build_context_with_overlays();
3547 self.template_renderer.render_with_state(
3548 prompt,
3549 &context,
3550 &sm.current(),
3551 sm.previous().as_deref(),
3552 sm.turn_count(),
3553 state_def.max_turns,
3554 )?
3555 } else {
3556 String::new()
3557 };
3558
3559 let combined = match state_def.prompt_mode {
3560 PromptMode::Append => {
3561 if state_prompt.is_empty() {
3562 rendered_base
3563 } else {
3564 format!(
3565 "{}\n\n[Current State: {}]\n{}",
3566 rendered_base,
3567 sm.current(),
3568 state_prompt
3569 )
3570 }
3571 }
3572 PromptMode::Replace => {
3573 if state_prompt.is_empty() {
3574 rendered_base
3575 } else {
3576 state_prompt
3577 }
3578 }
3579 PromptMode::Prepend => {
3580 if state_prompt.is_empty() {
3581 rendered_base
3582 } else {
3583 format!("{}\n\n{}", state_prompt, rendered_base)
3584 }
3585 }
3586 };
3587
3588 let with_persona = if persona_prefix.is_empty() {
3590 combined
3591 } else {
3592 format!("{}\n\n{}", persona_prefix, combined)
3593 };
3594
3595 if include_tool_prompt {
3596 let available_tool_ids = self.get_available_tool_ids().await?;
3597 if !available_tool_ids.is_empty() {
3598 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3599 &available_tool_ids,
3600 None,
3601 self.parallel_tools.enabled,
3602 self.runtime_config.tool_schema_prompt_mode,
3603 );
3604 if !tools_prompt.is_empty() {
3605 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3606 }
3607 }
3608 }
3609 return Ok(with_persona);
3610 }
3611
3612 let with_persona = if persona_prefix.is_empty() {
3614 rendered_base
3615 } else {
3616 format!("{}\n\n{}", persona_prefix, rendered_base)
3617 };
3618
3619 if include_tool_prompt {
3620 let available_tool_ids = self.get_available_tool_ids().await?;
3621 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3622 &available_tool_ids,
3623 None,
3624 self.parallel_tools.enabled,
3625 self.runtime_config.tool_schema_prompt_mode,
3626 );
3627 if !tools_prompt.is_empty() {
3628 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3629 }
3630 }
3631 Ok(with_persona)
3632 }
3633
3634 fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3635 if let Some(ref sm) = self.state_machine
3636 && let Some(state_def) = sm.current_definition()
3637 && let Some(ref llm_alias) = state_def.llm
3638 {
3639 return self
3640 .llm_registry
3641 .get(llm_alias)
3642 .map_err(|e| AgentError::Config(e.to_string()));
3643 }
3644 self.llm_registry
3645 .default()
3646 .map_err(|e| AgentError::Config(e.to_string()))
3647 }
3648
3649 fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3650 if let Some(ref sm) = self.state_machine
3651 && let Some(state_def) = sm.current_definition()
3652 && let Some(ref state_reasoning) = state_def.reasoning
3653 {
3654 return state_reasoning.clone();
3655 }
3656 self.reasoning_config.clone()
3657 }
3658
3659 fn get_effective_reflection_config(&self) -> ReflectionConfig {
3660 if let Some(ref sm) = self.state_machine
3661 && let Some(state_def) = sm.current_definition()
3662 && let Some(ref state_reflection) = state_def.reflection
3663 {
3664 return state_reflection.clone();
3665 }
3666 self.reflection_config.clone()
3667 }
3668
3669 fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3670 skill
3671 .reasoning
3672 .clone()
3673 .unwrap_or_else(|| self.get_effective_reasoning_config())
3674 }
3675
3676 fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3677 skill
3678 .reflection
3679 .clone()
3680 .unwrap_or_else(|| self.get_effective_reflection_config())
3681 }
3682
3683 async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3684 let recent_messages: Vec<String> = self
3685 .memory
3686 .get_messages(Some(5))
3687 .await?
3688 .iter()
3689 .rev()
3690 .map(|m| format!("{:?}: {}", m.role, m.content))
3691 .collect();
3692
3693 let current_state = self.current_state().map(|s| s.to_string());
3694
3695 let state_prompt: Option<String> = self
3698 .state_machine
3699 .as_ref()
3700 .and_then(|sm| sm.current_definition())
3701 .and_then(|def| def.prompt.clone());
3702
3703 let available_tools: Vec<String> = self
3704 .get_available_tool_ids()
3705 .await
3706 .unwrap_or_else(|_| self.tools.list_ids());
3707
3708 let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3709
3710 let mut user_context = self.build_context_with_overlays();
3711 user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3712 if let Some(state_generation) = self
3713 .state_machine
3714 .as_ref()
3715 .map(|state_machine| state_machine.generation())
3716 {
3717 user_context.insert(
3718 DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3719 serde_json::json!(state_generation),
3720 );
3721 }
3722
3723 let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3725 sm.current_definition()
3726 .map(|def| {
3727 def.transitions
3728 .iter()
3729 .filter_map(|t| t.intent.clone())
3730 .collect()
3731 })
3732 .unwrap_or_default()
3733 } else {
3734 Vec::new()
3735 };
3736
3737 Ok(DisambiguationContext::from_agent_state(
3738 recent_messages,
3739 current_state,
3740 state_prompt,
3741 available_tools,
3742 available_skills,
3743 available_intents,
3744 user_context,
3745 ))
3746 }
3747
3748 fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3749 if let Some(ref sm) = self.state_machine
3750 && let Some(state_def) = sm.current_definition()
3751 {
3752 let parent_def = sm.get_parent_definition();
3753 let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3754 if !effective_skills.is_empty() {
3755 return self
3756 .skills
3757 .iter()
3758 .filter(|s| effective_skills.contains(&&s.id))
3759 .collect();
3760 }
3761 }
3762 self.skills.iter().collect()
3763 }
3764
3765 async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3766 self.build_messages_internal(true, None, true).await
3767 }
3768
3769 async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3770 self.build_messages_internal(false, Some(user_message), true)
3771 .await
3772 }
3773
3774 async fn build_messages_internal(
3775 &self,
3776 fire_persona_hooks: bool,
3777 ephemeral_user_message: Option<&str>,
3778 include_tool_prompt: bool,
3779 ) -> Result<Vec<ChatMessage>> {
3780 let system_prompt = self
3781 .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3782 .await?;
3783 let mut messages = vec![ChatMessage::system(&system_prompt)];
3784
3785 let context = self.memory.get_context().await?;
3786 let history = if let Some(ref budget) = self.memory_token_budget {
3787 context.to_llm_messages_with_allocation(&budget.allocation)
3788 } else {
3789 context.to_llm_messages()
3790 };
3791 messages.extend(history);
3792 if let Some(user_message) = ephemeral_user_message {
3793 messages.push(ChatMessage::user(user_message));
3794 }
3795
3796 let total_tokens = self.estimate_total_tokens(&messages);
3797
3798 if total_tokens > self.max_context_tokens {
3799 debug!(
3800 total = total_tokens,
3801 limit = self.max_context_tokens,
3802 "Context overflow"
3803 );
3804
3805 match &self.recovery_manager.config().llm.on_context_overflow {
3806 ContextOverflowAction::Error => {
3807 return Err(AgentError::LLM(format!(
3808 "Context overflow: {} tokens > {} limit",
3809 total_tokens, self.max_context_tokens
3810 )));
3811 }
3812 ContextOverflowAction::Truncate { keep_recent } => {
3813 self.truncate_context(&mut messages, *keep_recent);
3814 }
3815 ContextOverflowAction::Summarize {
3816 summarizer_llm,
3817 max_summary_tokens,
3818 custom_prompt,
3819 keep_recent,
3820 filter,
3821 } => {
3822 self.summarize_context(
3823 &mut messages,
3824 summarizer_llm.as_deref(),
3825 *max_summary_tokens,
3826 custom_prompt.as_deref(),
3827 *keep_recent,
3828 filter.as_ref(),
3829 )
3830 .await?;
3831 }
3832 }
3833 }
3834
3835 Ok(messages)
3836 }
3837
3838 async fn main_tool_protocol(
3839 &self,
3840 llm: &dyn LLMProvider,
3841 ephemeral_new_turn: bool,
3842 ) -> Result<MainToolProtocol> {
3843 let mut choice = llm.configured_tool_choice();
3844 if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3845 return Ok(MainToolProtocol {
3846 choice,
3847 tool_ids: Vec::new(),
3848 definitions: Vec::new(),
3849 });
3850 }
3851
3852 let mut tool_ids = self.get_available_tool_ids().await?;
3853 tool_ids.sort();
3854 tool_ids.dedup();
3855 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3856 let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3857 AgentError::Config(format!(
3858 "specific tool choice '{expected}' is not registered"
3859 ))
3860 })?;
3861 if canonical != *expected {
3862 return Err(AgentError::Config(format!(
3863 "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3864 )));
3865 }
3866 if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3867 return Err(AgentError::Config(format!(
3868 "specific tool choice '{expected}' is outside the effective tool grant"
3869 )));
3870 }
3871 }
3872 if matches!(
3873 choice.as_ref(),
3874 Some(ToolChoice::Required | ToolChoice::Specific(_))
3875 ) && tool_ids.is_empty()
3876 {
3877 return Err(AgentError::Config(
3878 "required tool choice has no tool inside the effective grant".to_string(),
3879 ));
3880 }
3881 if !ephemeral_new_turn
3882 && let Some(configured_choice) = choice.as_ref()
3883 && matches!(
3884 configured_choice,
3885 ToolChoice::Required | ToolChoice::Specific(_)
3886 )
3887 && self
3888 .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3889 .await?
3890 {
3891 choice = Some(ToolChoice::Auto);
3892 }
3893 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3894 tool_ids.retain(|tool_id| tool_id == expected);
3895 }
3896
3897 let definitions = tool_ids
3898 .iter()
3899 .map(|tool_id| {
3900 let tool = self.tools.get(tool_id).ok_or_else(|| {
3901 AgentError::Config(format!(
3902 "effective tool '{tool_id}' disappeared before provider exposure"
3903 ))
3904 })?;
3905 Ok(LLMToolDefinition {
3906 name: tool_id.clone(),
3907 description: tool.description().to_string(),
3908 input_schema: tool.input_schema(),
3909 })
3910 })
3911 .collect::<Result<Vec<_>>>()?;
3912
3913 Ok(MainToolProtocol {
3917 choice,
3918 tool_ids,
3919 definitions,
3920 })
3921 }
3922
3923 async fn tool_choice_satisfied_in_current_turn(
3924 &self,
3925 choice: &ToolChoice,
3926 effective_tool_ids: &[String],
3927 ) -> Result<bool> {
3928 let messages = self.memory.get_messages(None).await?;
3929 let mut saw_tool_result = false;
3930 for message in messages.iter().rev() {
3931 match message.role {
3932 ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
3933 saw_tool_result = true;
3934 }
3935 ai_agents_core::Role::Assistant if saw_tool_result => {
3936 let Some(calls) = self.parse_tool_calls(&message.content) else {
3937 continue;
3938 };
3939 let calls_are_effective = !calls.is_empty()
3940 && calls.iter().all(|call| {
3941 self.tools
3942 .canonical_id(&call.name)
3943 .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
3944 });
3945 return Ok(calls_are_effective
3946 && match choice {
3947 ToolChoice::Required => true,
3948 ToolChoice::Specific(expected) => calls.iter().all(|call| {
3949 self.tools.canonical_id(&call.name).as_deref()
3950 == Some(expected.as_str())
3951 }),
3952 _ => false,
3953 });
3954 }
3955 ai_agents_core::Role::User => return Ok(false),
3956 _ => {}
3957 }
3958 }
3959 Ok(false)
3960 }
3961
3962 fn provider_can_use_native_tools(
3963 &self,
3964 llm: &dyn LLMProvider,
3965 protocol: &MainToolProtocol,
3966 ) -> bool {
3967 let Some(choice) = protocol.choice.as_ref() else {
3968 return false;
3969 };
3970 if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
3971 return false;
3972 }
3973 llm.supports_tool_choice(choice)
3974 && protocol.definitions.iter().all(|definition| {
3975 !definition.name.is_empty()
3976 && definition.name.len() <= 64
3977 && definition
3978 .name
3979 .bytes()
3980 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
3981 })
3982 }
3983
3984 fn prompt_messages_for_tool_protocol(
3985 &self,
3986 messages: &[ChatMessage],
3987 protocol: &MainToolProtocol,
3988 corrective: bool,
3989 ) -> Vec<ChatMessage> {
3990 let mut messages = messages.to_vec();
3991 let Some(choice) = protocol.choice.as_ref() else {
3992 return messages;
3993 };
3994 if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
3995 return messages;
3996 }
3997
3998 let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
3999 &protocol.tool_ids,
4000 None,
4001 self.parallel_tools.enabled,
4002 self.runtime_config.tool_schema_prompt_mode,
4003 );
4004 match choice {
4005 ToolChoice::Required => tool_prompt.push_str(
4006 "\n\nYou must call at least one listed tool before giving a final answer.",
4007 ),
4008 ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
4009 "\n\nYou must call the '{tool_id}' tool before giving a final answer."
4010 )),
4011 ToolChoice::Auto => {}
4012 ToolChoice::None => return messages,
4013 _ => return messages,
4014 }
4015 if let Some(system) = messages
4016 .iter_mut()
4017 .find(|message| message.role == ai_agents_core::Role::System)
4018 {
4019 system.content.push_str("\n\n");
4020 system.content.push_str(&tool_prompt);
4021 } else {
4022 messages.insert(0, ChatMessage::system(tool_prompt));
4023 }
4024 if corrective {
4025 let instruction = match choice {
4026 ToolChoice::Required => {
4027 "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
4028 }
4029 ToolChoice::Specific(tool_id) => {
4030 messages.push(ChatMessage::user(format!(
4031 "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
4032 )));
4033 return messages;
4034 }
4035 _ => return messages,
4036 };
4037 messages.push(ChatMessage::user(instruction));
4038 }
4039 messages
4040 }
4041
4042 async fn invoke_main_provider(
4043 &self,
4044 llm: Arc<dyn LLMProvider>,
4045 messages: &[ChatMessage],
4046 protocol: &MainToolProtocol,
4047 corrective: bool,
4048 ) -> std::result::Result<MainProviderResponse, LLMError> {
4049 let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
4050 let response = if use_native {
4051 let request = LLMToolRequest {
4052 tools: protocol.definitions.clone(),
4053 choice: protocol
4054 .choice
4055 .clone()
4056 .expect("native tool requests require an explicit choice"),
4057 };
4058 self.observe_purpose(
4059 ObservationPurpose::MainResponse,
4060 llm.complete_with_tools(messages, None, &request),
4061 )
4062 .await?
4063 } else {
4064 let prompt_messages =
4065 self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
4066 self.observe_purpose(
4067 ObservationPurpose::MainResponse,
4068 llm.complete(&prompt_messages, None),
4069 )
4070 .await?
4071 };
4072 Ok(MainProviderResponse {
4073 response,
4074 used_native_tools: use_native,
4075 })
4076 }
4077
4078 async fn complete_main_attempt_with_recovery(
4079 &self,
4080 llm: Arc<dyn LLMProvider>,
4081 messages: &[ChatMessage],
4082 protocol: &MainToolProtocol,
4083 corrective: bool,
4084 ) -> Result<MainProviderResponse> {
4085 let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
4086 self.recovery_manager
4087 .with_retry("llm_call", None, || {
4088 let llm = Arc::clone(&llm);
4089 async move {
4090 self.invoke_main_provider(llm, messages, protocol, corrective)
4091 .await
4092 .map_err(|error| error.classify())
4093 }
4094 })
4095 .await
4096 .map_err(|error| AgentError::LLM(error.to_string()))
4097 } else {
4098 self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4099 .await
4100 .map_err(|error| AgentError::LLM(error.to_string()))
4101 };
4102
4103 match primary_result {
4104 Ok(response) => Ok(response),
4105 Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4106 LLMFailureAction::FallbackLlm { fallback_llm } => {
4107 let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4108 AgentError::Config(format!(
4109 "Fallback LLM '{fallback_llm}' not found: {error}"
4110 ))
4111 })?;
4112 self.invoke_main_provider(fallback, messages, protocol, corrective)
4113 .await
4114 .map_err(|error| AgentError::LLM(error.to_string()))
4115 }
4116 LLMFailureAction::FallbackResponse { message } => {
4117 if matches!(
4118 protocol.choice.as_ref(),
4119 Some(ToolChoice::Required | ToolChoice::Specific(_))
4120 ) {
4121 Err(AgentError::LLM(format!(
4122 "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4123 )))
4124 } else {
4125 Ok(MainProviderResponse {
4126 response: LLMResponse::new(message.clone(), FinishReason::Stop),
4127 used_native_tools: false,
4128 })
4129 }
4130 }
4131 LLMFailureAction::Error => Err(primary_error),
4132 },
4133 }
4134 }
4135
4136 fn normalize_main_provider_response(
4137 &self,
4138 mut response: LLMResponse,
4139 protocol: &MainToolProtocol,
4140 ) -> Result<(LLMResponse, bool)> {
4141 let native_calls = response
4142 .tool_calls()
4143 .map_err(|error| AgentError::LLM(error.to_string()))?;
4144 let calls = match native_calls {
4145 Some(calls) => {
4146 let markers = calls
4147 .iter()
4148 .map(|call| {
4149 serde_json::json!({
4150 "_ai_agents_native_tool_call": true,
4151 "id": call.id,
4152 "tool": call.name,
4153 "arguments": call.arguments,
4154 })
4155 })
4156 .collect::<Vec<_>>();
4157 response.content = if markers.len() == 1 {
4158 markers[0].to_string()
4159 } else {
4160 serde_json::Value::Array(markers).to_string()
4161 };
4162 Some(calls)
4163 }
4164 None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4165 self.parse_tool_calls(response.content.trim())
4166 }
4167 None => None,
4168 };
4169
4170 if protocol.choice.is_some()
4171 && let Some(calls) = calls.as_ref()
4172 && calls.iter().any(|call| {
4173 self.tools
4174 .canonical_id(&call.name)
4175 .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4176 })
4177 {
4178 return Err(AgentError::LLM(
4179 "Provider returned a tool call outside the effective grant".to_string(),
4180 ));
4181 }
4182
4183 let compliant = match protocol.choice.as_ref() {
4184 Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4185 Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4186 !calls.is_empty()
4187 && calls.iter().all(|call| {
4188 self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4189 })
4190 }),
4191 _ => true,
4192 };
4193 Ok((response, compliant))
4194 }
4195
4196 async fn complete_main_llm_with_recovery(
4197 &self,
4198 llm: Arc<dyn LLMProvider>,
4199 messages: &[ChatMessage],
4200 protocol: &MainToolProtocol,
4201 ) -> Result<LLMResponse> {
4202 let first = self
4203 .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4204 .await?;
4205 let (response, compliant) =
4206 self.normalize_main_provider_response(first.response, protocol)?;
4207 if compliant {
4208 return Ok(response);
4209 }
4210 if first.used_native_tools {
4211 return Err(AgentError::LLM(
4212 "Provider returned no compliant native call for required tool choice".to_string(),
4213 ));
4214 }
4215
4216 let corrected = self
4217 .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4218 .await?;
4219 let (response, compliant) =
4220 self.normalize_main_provider_response(corrected.response, protocol)?;
4221 if compliant {
4222 return Ok(response);
4223 }
4224 Err(AgentError::LLM(
4225 "Provider returned no compliant tool call after one corrective retry".to_string(),
4226 ))
4227 }
4228
4229 fn is_native_tool_call_content(content: &str) -> bool {
4230 let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4231 return false;
4232 };
4233 match value {
4234 serde_json::Value::Array(values) => {
4235 !values.is_empty()
4236 && values.iter().all(|value| {
4237 value
4238 .get("_ai_agents_native_tool_call")
4239 .and_then(|marker| marker.as_bool())
4240 == Some(true)
4241 })
4242 }
4243 serde_json::Value::Object(map) => {
4244 map.get("_ai_agents_native_tool_call")
4245 .and_then(|marker| marker.as_bool())
4246 == Some(true)
4247 }
4248 _ => false,
4249 }
4250 }
4251
4252 fn tool_result_message(
4253 tool_call: &ToolCall,
4254 output: &str,
4255 native_tool_call: bool,
4256 ) -> ChatMessage {
4257 if !native_tool_call {
4258 return ChatMessage::function(&tool_call.name, output);
4259 }
4260 let output = serde_json::from_str::<serde_json::Value>(output)
4261 .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4262 ChatMessage::function(
4263 &tool_call.name,
4264 serde_json::json!({
4265 "_ai_agents_native_tool_result": true,
4266 "id": tool_call.id,
4267 "tool": tool_call.name,
4268 "output": output,
4269 })
4270 .to_string(),
4271 )
4272 }
4273
4274 fn parse_main_tool_calls(
4275 &self,
4276 content: &str,
4277 protocol: &MainToolProtocol,
4278 ) -> Option<Vec<ToolCall>> {
4279 if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4280 None
4281 } else {
4282 self.parse_tool_calls(content)
4283 }
4284 }
4285
4286 fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4287 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4289 if let Some(arr) = parsed.as_array() {
4291 let calls: Vec<ToolCall> = arr
4292 .iter()
4293 .filter_map(|v| self.extract_tool_call_from_value(v))
4294 .collect();
4295 if !calls.is_empty() {
4296 return Some(calls);
4297 }
4298 }
4299 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4301 return Some(vec![tool_call]);
4302 }
4303 }
4304
4305 if let Some(json_str) = self.extract_json_from_content(content)
4307 && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4308 {
4309 if let Some(arr) = parsed.as_array() {
4311 let calls: Vec<ToolCall> = arr
4312 .iter()
4313 .filter_map(|v| self.extract_tool_call_from_value(v))
4314 .collect();
4315 if !calls.is_empty() {
4316 return Some(calls);
4317 }
4318 }
4319 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4321 return Some(vec![tool_call]);
4322 }
4323 }
4324
4325 None
4326 }
4327
4328 fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4329 if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4330 let arguments = parsed
4331 .get("arguments")
4332 .cloned()
4333 .unwrap_or(serde_json::json!({}));
4334 return Some(ToolCall {
4335 id: parsed
4336 .get("id")
4337 .and_then(|value| value.as_str())
4338 .filter(|id| !id.is_empty())
4339 .map(str::to_string)
4340 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4341 name: tool_name.to_string(),
4342 arguments,
4343 });
4344 }
4345 None
4346 }
4347
4348 fn extract_json_from_content(&self, content: &str) -> Option<String> {
4350 if let Some(result) = self.extract_json_array_from_content(content) {
4352 return Some(result);
4353 }
4354 self.extract_json_object_from_content(content)
4355 }
4356
4357 fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4359 let start = content.find('[')?;
4360 let content_from_start = &content[start..];
4361
4362 let mut depth = 0;
4363 let mut end = 0;
4364 for (i, ch) in content_from_start.char_indices() {
4365 match ch {
4366 '[' => depth += 1,
4367 ']' => {
4368 depth -= 1;
4369 if depth == 0 {
4370 end = i + 1;
4371 break;
4372 }
4373 }
4374 _ => {}
4375 }
4376 }
4377
4378 if end > 0 {
4379 let json_str = &content_from_start[..end];
4380 if json_str.contains("\"tool\"") {
4382 return Some(json_str.to_string());
4383 }
4384 }
4385
4386 None
4387 }
4388
4389 fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4391 let start = content.find('{')?;
4392 let content_from_start = &content[start..];
4393
4394 let mut depth = 0;
4396 let mut end = 0;
4397 for (i, ch) in content_from_start.char_indices() {
4398 match ch {
4399 '{' => depth += 1,
4400 '}' => {
4401 depth -= 1;
4402 if depth == 0 {
4403 end = i + 1;
4404 break;
4405 }
4406 }
4407 _ => {}
4408 }
4409 }
4410
4411 if end > 0 {
4412 let json_str = &content_from_start[..end];
4413 if json_str.contains("\"tool\"") {
4415 return Some(json_str.to_string());
4416 }
4417 }
4418
4419 None
4420 }
4421
4422 #[allow(clippy::too_many_arguments)]
4426 fn record_from_parts(
4427 &self,
4428 request: &ToolExecutionRequest,
4429 canonical_id: String,
4430 executed_arguments: Value,
4431 started_at: chrono::DateTime<chrono::Utc>,
4432 start: Instant,
4433 executed: bool,
4434 success: bool,
4435 output: String,
4436 metadata: HashMap<String, Value>,
4437 policy: ToolPolicyDecisionRecord,
4438 approval: Option<ToolApprovalRecord>,
4439 timed_out: bool,
4440 output_truncated: bool,
4441 ) -> ToolExecutionRecord {
4442 let versions = ToolDecisionVersions {
4443 policy: self.active_tool_security().policy_version(),
4444 registry: self.tools.version(),
4445 runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4446 state: self
4447 .state_machine
4448 .as_ref()
4449 .map(|state_machine| state_machine.generation()),
4450 };
4451 self.record_from_parts_at(
4452 request,
4453 canonical_id,
4454 executed_arguments,
4455 started_at,
4456 start,
4457 executed,
4458 success,
4459 output,
4460 metadata,
4461 policy,
4462 approval,
4463 timed_out,
4464 output_truncated,
4465 versions,
4466 )
4467 }
4468
4469 #[allow(clippy::too_many_arguments)]
4471 fn record_from_parts_at(
4472 &self,
4473 request: &ToolExecutionRequest,
4474 canonical_id: String,
4475 executed_arguments: Value,
4476 started_at: chrono::DateTime<chrono::Utc>,
4477 start: Instant,
4478 executed: bool,
4479 success: bool,
4480 output: String,
4481 metadata: HashMap<String, Value>,
4482 policy: ToolPolicyDecisionRecord,
4483 approval: Option<ToolApprovalRecord>,
4484 timed_out: bool,
4485 output_truncated: bool,
4486 versions: ToolDecisionVersions,
4487 ) -> ToolExecutionRecord {
4488 ToolExecutionRecord {
4489 call_id: request.call_id.clone(),
4490 requested_name: request.requested_name.clone(),
4491 canonical_id,
4492 source: request.source.clone(),
4493 arguments: request.arguments.clone(),
4494 executed_arguments,
4495 policy_version: versions.policy,
4496 registry_version: versions.registry,
4497 runtime_config_version: versions.runtime_control,
4498 executed,
4499 success,
4500 output,
4501 metadata,
4502 policy,
4503 approval,
4504 started_at,
4505 duration_ms: start.elapsed().as_millis() as u64,
4506 timed_out,
4507 cancelled: false,
4508 cancellation_reason: None,
4509 output_truncated,
4510 }
4511 }
4512
4513 async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4515 let result = ToolResult {
4516 success: record.success,
4517 output: record.model_output_string(),
4518 metadata: if record.metadata.is_empty() {
4519 None
4520 } else {
4521 Some(record.metadata.clone())
4522 },
4523 };
4524 self.hooks
4525 .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4526 .await;
4527 self.hooks.on_tool_execution_record(record).await;
4528 self.record_tool_call(&record.canonical_id, record.model_output_value());
4529 if !record.success {
4530 self.hooks
4531 .on_error(&AgentError::Tool(record.output.clone()))
4532 .await;
4533 }
4534 }
4535
4536 async fn finish_tool_record_after_resource_guards(
4538 &self,
4539 resource_guards: ToolResourceGuards,
4540 record: &ToolExecutionRecord,
4541 ) {
4542 drop(resource_guards);
4543 self.finish_tool_record(record).await;
4544 }
4545
4546 async fn execute_resolved_tool_once(
4548 &self,
4549 tool: Arc<dyn ai_agents_core::Tool>,
4550 args: Value,
4551 mut ctx: ToolExecutionContext,
4552 timeout_ms: u64,
4553 ) -> Result<(ToolResult, bool, bool, bool)> {
4554 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4555 return Ok((
4556 ToolResult::error("Tool execution cancelled by runtime control"),
4557 false,
4558 true,
4559 false,
4560 ));
4561 }
4562 ctx.deadline = Some(chrono::Utc::now() + chrono::Duration::milliseconds(timeout_ms as i64));
4567 let invoked = Arc::new(AtomicBool::new(false));
4571 let invoked_by_future = Arc::clone(&invoked);
4572 let actor_context = current_turn_actor_context();
4573 let future = async move {
4574 invoked_by_future.store(true, Ordering::SeqCst);
4575 if let Some(actor_context) = actor_context {
4576 scope_actor_context(actor_context, tool.execute(args, ctx)).await
4577 } else {
4578 tool.execute(args, ctx).await
4579 }
4580 };
4581 tokio::pin!(future);
4582 let timeout = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms));
4583 tokio::pin!(timeout);
4584 let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4585
4586 loop {
4587 tokio::select! {
4588 result = &mut future => return Ok((result, false, false, true)),
4589 _ = &mut timeout => {
4590 return Ok((
4591 ToolResult::error("Tool execution timed out"),
4592 true,
4593 false,
4594 invoked.load(Ordering::SeqCst),
4595 ));
4596 }
4597 _ = cancel_tick.tick() => {
4598 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4599 return Ok((
4600 ToolResult::error("Tool execution cancelled by runtime control"),
4601 false,
4602 true,
4603 invoked.load(Ordering::SeqCst),
4604 ));
4605 }
4606 }
4607 }
4608 }
4609 }
4610
4611 fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4613 let Some(max_chars) = max_chars else {
4614 return (output, false);
4615 };
4616 let mut chars = output.chars();
4617 let truncated: String = chars.by_ref().take(max_chars).collect();
4618 if chars.next().is_some() {
4619 (truncated, true)
4620 } else {
4621 (output, false)
4622 }
4623 }
4624
4625 async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4627 let locks = {
4628 let mut table = self.resource_locks.write();
4629 table.retain(|_, lock| lock.strong_count() > 0);
4630 keys.iter()
4631 .map(|key| {
4632 if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4633 lock
4634 } else {
4635 let lock = Arc::new(tokio::sync::Mutex::new(()));
4636 table.insert(key.clone(), Arc::downgrade(&lock));
4637 lock
4638 }
4639 })
4640 .collect::<Vec<_>>()
4641 };
4642 let mut resource_guards = ToolResourceGuards {
4643 guards: Vec::with_capacity(locks.len()),
4644 locks: Arc::clone(&self.resource_locks),
4645 };
4646 let mut locks = locks.into_iter();
4647 while let Some(lock) = locks.next() {
4648 let mut lock = Box::pin(lock.lock_owned());
4649 loop {
4650 tokio::select! {
4651 guard = &mut lock => {
4652 resource_guards.guards.push(guard);
4653 break;
4654 }
4655 _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4656 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4657 drop(lock);
4658 drop(locks);
4659 drop(resource_guards);
4660 return None;
4661 }
4662 }
4663 }
4664 }
4665 }
4666 Some(resource_guards)
4667 }
4668
4669 async fn run_tool_with_retries(
4673 &self,
4674 canonical_id: &str,
4675 tool: Arc<dyn ai_agents_core::Tool>,
4676 args: Value,
4677 ctx: ToolExecutionContext,
4678 timeout_ms: u64,
4679 max_retries: u32,
4680 ) -> Result<(ToolResult, bool, bool, bool)> {
4681 let max_retries = if ctx.classification.safely_retryable {
4682 max_retries
4683 } else {
4684 0
4685 };
4686 let mut attempts = 0;
4687 let mut invoked = false;
4688 loop {
4689 let (result, timed_out, cancelled, attempt_invoked) = self
4690 .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout_ms)
4691 .await?;
4692 invoked |= attempt_invoked;
4693 if result.success || timed_out || cancelled || attempts >= max_retries {
4694 return Ok((result, timed_out, cancelled, invoked));
4695 }
4696 attempts += 1;
4697 warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4698 }
4699 }
4700
4701 fn host_tool_unavailability(&self, canonical_id: &str) -> Option<(&'static str, &'static str)> {
4703 match canonical_id {
4704 "command" if !self.tools.command_runner_available() => Some((
4705 "Command runner is unavailable",
4706 "command runner is unavailable",
4707 )),
4708 "diagnostics" if !self.tools.diagnostics_available() => Some((
4709 "Diagnostics provider is unavailable",
4710 "diagnostics provider is unavailable",
4711 )),
4712 "web_search" if !self.tools.web_search_available() => Some((
4713 "Web search provider is unavailable",
4714 "web search provider is unavailable",
4715 )),
4716 _ => None,
4717 }
4718 }
4719
4720 fn execute_tool_record(
4722 &self,
4723 request: ToolExecutionRequest,
4724 ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4725 Box::pin(self.execute_tool_record_inner(request))
4726 }
4727
4728 async fn execute_tool_record_inner(
4732 &self,
4733 request: ToolExecutionRequest,
4734 ) -> Result<ToolExecutionRecord> {
4735 let started_at = chrono::Utc::now();
4736 let start = Instant::now();
4737 info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4738
4739 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4740 let record = self.record_from_parts(
4741 &request,
4742 request.requested_name.clone(),
4743 request.arguments.clone(),
4744 started_at,
4745 start,
4746 false,
4747 false,
4748 "Tool execution is disabled by runtime control".to_string(),
4749 HashMap::new(),
4750 ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4751 None,
4752 false,
4753 false,
4754 );
4755 self.finish_tool_record(&record).await;
4756 return Ok(record);
4757 }
4758
4759 let Some(resolved) = self.tools.resolve(&request.requested_name) else {
4760 let record = self.record_from_parts(
4761 &request,
4762 request.requested_name.clone(),
4763 request.arguments.clone(),
4764 started_at,
4765 start,
4766 false,
4767 false,
4768 format!("Tool '{}' is unavailable", request.requested_name),
4769 HashMap::new(),
4770 ToolPolicyDecisionRecord::unavailable(format!(
4771 "Tool '{}' is not registered",
4772 request.requested_name
4773 )),
4774 None,
4775 false,
4776 false,
4777 );
4778 self.finish_tool_record(&record).await;
4779 return Ok(record);
4780 };
4781
4782 let canonical_id = resolved.identity.canonical_id.clone();
4783
4784 let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4785 if !initial_scope_snapshot
4786 .tool_ids
4787 .iter()
4788 .any(|id| id == &canonical_id)
4789 {
4790 let record = self.record_from_parts(
4791 &request,
4792 canonical_id.clone(),
4793 request.arguments.clone(),
4794 started_at,
4795 start,
4796 false,
4797 false,
4798 format!(
4799 "Tool '{}' is not available in the current scope",
4800 canonical_id
4801 ),
4802 HashMap::new(),
4803 ToolPolicyDecisionRecord::deny(format!(
4804 "Tool '{}' is not granted by the current top-level and state tool scope",
4805 canonical_id
4806 )),
4807 None,
4808 false,
4809 false,
4810 );
4811 self.finish_tool_record(&record).await;
4812 return Ok(record);
4813 }
4814
4815 let approval_control_snapshot = self.runtime_safety_snapshot();
4816 let security_engine = approval_control_snapshot.tool_security.clone();
4817 let bindings = resolved.tool.policy_bindings();
4818 let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4819 &canonical_id,
4820 &request.arguments,
4821 &bindings,
4822 );
4823 self.hooks
4824 .on_tool_start(&canonical_id, &executed_arguments)
4825 .await;
4826
4827 let mut metadata = HashMap::new();
4828 let safety = resolved.tool.safety_metadata();
4829 let classification = resolved.tool.classify_call(&executed_arguments);
4830 let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
4831 metadata.insert(
4832 "classification".to_string(),
4833 serde_json::to_value(&classification).unwrap_or(Value::Null),
4834 );
4835 metadata.insert(
4836 "effective_limits".to_string(),
4837 serde_json::to_value(&limits).unwrap_or(Value::Null),
4838 );
4839 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
4840 if !policy_snapshot.is_null() {
4841 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
4842 }
4843
4844 let mut approval_record = Some(ToolApprovalRecord {
4845 status: ToolApprovalStatus::NotRequired,
4846 reason: None,
4847 modified_arguments: None,
4848 });
4849
4850 let mut security_result = security_engine
4851 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
4852 .await?;
4853 if (security_result.is_allowed()
4858 || matches!(
4859 &security_result,
4860 SecurityCheckResult::RequireConfirmation { .. }
4861 ))
4862 && let Some((output, reason)) = self.host_tool_unavailability(&canonical_id)
4863 {
4864 let record = self.record_from_parts(
4865 &request,
4866 canonical_id,
4867 executed_arguments,
4868 started_at,
4869 start,
4870 false,
4871 false,
4872 output.to_string(),
4873 metadata,
4874 ToolPolicyDecisionRecord::unavailable(reason),
4875 Some(ToolApprovalRecord {
4876 status: ToolApprovalStatus::Unavailable,
4877 reason: Some(reason.to_string()),
4878 modified_arguments: None,
4879 }),
4880 false,
4881 false,
4882 );
4883 self.finish_tool_record(&record).await;
4884 return Ok(record);
4885 }
4886 match &security_result {
4887 SecurityCheckResult::Allow => {}
4888 SecurityCheckResult::Warn { message } => {
4889 warn!(tool = %canonical_id, message = %message, "Tool security warning");
4890 }
4891 SecurityCheckResult::Block { reason } => {
4892 let record = self.record_from_parts(
4893 &request,
4894 canonical_id,
4895 executed_arguments,
4896 started_at,
4897 start,
4898 false,
4899 false,
4900 format!("Denied: {}", reason),
4901 metadata,
4902 ToolPolicyDecisionRecord::deny(reason.clone()),
4903 approval_record,
4904 false,
4905 false,
4906 );
4907 self.finish_tool_record(&record).await;
4908 return Ok(record);
4909 }
4910 SecurityCheckResult::Unavailable { reason } => {
4911 let record = self.record_from_parts(
4912 &request,
4913 canonical_id,
4914 executed_arguments,
4915 started_at,
4916 start,
4917 false,
4918 false,
4919 format!("Unavailable: {}", reason),
4920 metadata,
4921 ToolPolicyDecisionRecord::unavailable(reason.clone()),
4922 approval_record,
4923 false,
4924 false,
4925 );
4926 self.finish_tool_record(&record).await;
4927 return Ok(record);
4928 }
4929 SecurityCheckResult::RequireConfirmation { message } => {
4930 if self.hitl_engine.is_none() {
4931 approval_record = Some(ToolApprovalRecord {
4932 status: ToolApprovalStatus::Unavailable,
4933 reason: Some("No HITL engine configured".to_string()),
4934 modified_arguments: None,
4935 });
4936 let record = self.record_from_parts(
4937 &request,
4938 canonical_id,
4939 executed_arguments,
4940 started_at,
4941 start,
4942 false,
4943 false,
4944 format!("Approval unavailable: {}", message),
4945 metadata,
4946 ToolPolicyDecisionRecord::approval(message.clone()),
4947 approval_record,
4948 false,
4949 false,
4950 );
4951 self.finish_tool_record(&record).await;
4952 return Ok(record);
4953 }
4954
4955 let check_result = HITLCheckResult::required(
4956 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4957 HashMap::new(),
4958 message.clone(),
4959 None,
4960 );
4961 match self.request_hitl_approval(check_result).await? {
4962 ApprovalResult::Approved => {
4963 merge_approved_record(&mut approval_record);
4964 }
4965 ApprovalResult::Modified { changes } => {
4966 if let Some(obj) = executed_arguments.as_object_mut() {
4967 for (key, value) in changes {
4968 obj.insert(key, value);
4969 }
4970 }
4971 security_result = security_engine
4972 .validate_tool_execution_with_bindings(
4973 &canonical_id,
4974 &executed_arguments,
4975 &bindings,
4976 )
4977 .await?;
4978 if !matches!(
4979 security_result,
4980 SecurityCheckResult::Allow
4981 | SecurityCheckResult::Warn { .. }
4982 | SecurityCheckResult::RequireConfirmation { .. }
4983 ) {
4984 let reason = security_result
4985 .reason()
4986 .unwrap_or("modified arguments failed policy")
4987 .to_string();
4988 let record = self.record_from_parts(
4989 &request,
4990 canonical_id,
4991 executed_arguments.clone(),
4992 started_at,
4993 start,
4994 false,
4995 false,
4996 reason.clone(),
4997 metadata,
4998 ToolPolicyDecisionRecord::deny(reason),
4999 Some(ToolApprovalRecord {
5000 status: ToolApprovalStatus::Modified,
5001 reason: None,
5002 modified_arguments: Some(executed_arguments),
5003 }),
5004 false,
5005 false,
5006 );
5007 self.finish_tool_record(&record).await;
5008 return Ok(record);
5009 }
5010 approval_record = Some(ToolApprovalRecord {
5011 status: ToolApprovalStatus::Modified,
5012 reason: None,
5013 modified_arguments: Some(executed_arguments.clone()),
5014 });
5015 }
5016 ApprovalResult::Rejected { reason } => {
5017 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5018 approval_record = Some(ToolApprovalRecord {
5019 status: ToolApprovalStatus::Rejected,
5020 reason: Some(reason.clone()),
5021 modified_arguments: None,
5022 });
5023 let record = self.record_from_parts(
5024 &request,
5025 canonical_id,
5026 executed_arguments,
5027 started_at,
5028 start,
5029 false,
5030 false,
5031 format!("Approval rejected: {}", reason),
5032 metadata,
5033 ToolPolicyDecisionRecord::approval(reason),
5034 approval_record,
5035 false,
5036 false,
5037 );
5038 self.finish_tool_record(&record).await;
5039 return Ok(record);
5040 }
5041 ApprovalResult::Timeout => {
5042 approval_record = Some(ToolApprovalRecord {
5043 status: ToolApprovalStatus::Timeout,
5044 reason: Some("approval timeout".to_string()),
5045 modified_arguments: None,
5046 });
5047 let record = self.record_from_parts(
5048 &request,
5049 canonical_id,
5050 executed_arguments,
5051 started_at,
5052 start,
5053 false,
5054 false,
5055 "Approval timed out".to_string(),
5056 metadata,
5057 ToolPolicyDecisionRecord::approval("approval timeout"),
5058 approval_record,
5059 false,
5060 false,
5061 );
5062 self.finish_tool_record(&record).await;
5063 return Ok(record);
5064 }
5065 }
5066 }
5067 }
5068
5069 if approval_record
5070 .as_ref()
5071 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
5072 && let Some(message) =
5073 security_engine.classification_approval_message(&canonical_id, &classification)
5074 {
5075 if self.hitl_engine.is_none() {
5076 approval_record = Some(ToolApprovalRecord {
5077 status: ToolApprovalStatus::Unavailable,
5078 reason: Some("No HITL engine configured".to_string()),
5079 modified_arguments: None,
5080 });
5081 let record = self.record_from_parts(
5082 &request,
5083 canonical_id,
5084 executed_arguments,
5085 started_at,
5086 start,
5087 false,
5088 false,
5089 format!("Approval unavailable: {}", message),
5090 metadata,
5091 ToolPolicyDecisionRecord::approval(message),
5092 approval_record,
5093 false,
5094 false,
5095 );
5096 self.finish_tool_record(&record).await;
5097 return Ok(record);
5098 }
5099 let check_result = HITLCheckResult::required(
5100 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
5101 HashMap::new(),
5102 message.clone(),
5103 None,
5104 );
5105 match self.request_hitl_approval(check_result).await? {
5106 ApprovalResult::Approved => {
5107 merge_approved_record(&mut approval_record);
5108 }
5109 ApprovalResult::Modified { changes } => {
5110 if let Some(obj) = executed_arguments.as_object_mut() {
5111 for (key, value) in changes {
5112 obj.insert(key, value);
5113 }
5114 }
5115 let modified_security = security_engine
5116 .validate_tool_execution_with_bindings(
5117 &canonical_id,
5118 &executed_arguments,
5119 &bindings,
5120 )
5121 .await?;
5122 if !matches!(
5123 modified_security,
5124 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5125 ) {
5126 let reason = modified_security
5127 .reason()
5128 .unwrap_or("modified arguments failed policy")
5129 .to_string();
5130 let record = self.record_from_parts(
5131 &request,
5132 canonical_id,
5133 executed_arguments.clone(),
5134 started_at,
5135 start,
5136 false,
5137 false,
5138 reason.clone(),
5139 metadata,
5140 ToolPolicyDecisionRecord::deny(reason),
5141 Some(ToolApprovalRecord {
5142 status: ToolApprovalStatus::Modified,
5143 reason: None,
5144 modified_arguments: Some(executed_arguments),
5145 }),
5146 false,
5147 false,
5148 );
5149 self.finish_tool_record(&record).await;
5150 return Ok(record);
5151 }
5152 approval_record = Some(ToolApprovalRecord {
5153 status: ToolApprovalStatus::Modified,
5154 reason: None,
5155 modified_arguments: Some(executed_arguments.clone()),
5156 });
5157 }
5158 ApprovalResult::Rejected { reason } => {
5159 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5160 let record = self.record_from_parts(
5161 &request,
5162 canonical_id,
5163 executed_arguments,
5164 started_at,
5165 start,
5166 false,
5167 false,
5168 format!("Approval rejected: {}", reason),
5169 metadata,
5170 ToolPolicyDecisionRecord::approval(reason.clone()),
5171 Some(ToolApprovalRecord {
5172 status: ToolApprovalStatus::Rejected,
5173 reason: Some(reason),
5174 modified_arguments: None,
5175 }),
5176 false,
5177 false,
5178 );
5179 self.finish_tool_record(&record).await;
5180 return Ok(record);
5181 }
5182 ApprovalResult::Timeout => {
5183 let record = self.record_from_parts(
5184 &request,
5185 canonical_id,
5186 executed_arguments,
5187 started_at,
5188 start,
5189 false,
5190 false,
5191 "Approval timed out".to_string(),
5192 metadata,
5193 ToolPolicyDecisionRecord::approval("approval timeout"),
5194 Some(ToolApprovalRecord {
5195 status: ToolApprovalStatus::Timeout,
5196 reason: Some("approval timeout".to_string()),
5197 modified_arguments: None,
5198 }),
5199 false,
5200 false,
5201 );
5202 self.finish_tool_record(&record).await;
5203 return Ok(record);
5204 }
5205 }
5206 }
5207
5208 let hitl_lang_ctx = self.build_hitl_language_context();
5209 if let Some(ref hitl_engine) = self.hitl_engine {
5210 let check_result = self
5211 .observe_purpose(
5212 ObservationPurpose::HitlLocalization,
5213 hitl_engine.check_tool_with_localization(
5214 &canonical_id,
5215 &executed_arguments,
5216 &hitl_lang_ctx,
5217 self.approval_handler.as_ref(),
5218 Some(&self.llm_registry),
5219 ),
5220 )
5221 .await?;
5222 if check_result.is_required() {
5223 match self.request_hitl_approval(check_result).await? {
5224 ApprovalResult::Approved => {
5225 merge_approved_record(&mut approval_record);
5226 }
5227 ApprovalResult::Modified { changes } => {
5228 if let Some(obj) = executed_arguments.as_object_mut() {
5229 for (key, value) in changes {
5230 obj.insert(key, value);
5231 }
5232 }
5233 let modified_security = security_engine
5234 .validate_tool_execution_with_bindings(
5235 &canonical_id,
5236 &executed_arguments,
5237 &bindings,
5238 )
5239 .await?;
5240 if !matches!(
5241 modified_security,
5242 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5243 ) {
5244 let reason = modified_security
5245 .reason()
5246 .unwrap_or("modified arguments failed policy")
5247 .to_string();
5248 let record = self.record_from_parts(
5249 &request,
5250 canonical_id,
5251 executed_arguments.clone(),
5252 started_at,
5253 start,
5254 false,
5255 false,
5256 reason.clone(),
5257 metadata,
5258 ToolPolicyDecisionRecord::deny(reason),
5259 Some(ToolApprovalRecord {
5260 status: ToolApprovalStatus::Modified,
5261 reason: None,
5262 modified_arguments: Some(executed_arguments),
5263 }),
5264 false,
5265 false,
5266 );
5267 self.finish_tool_record(&record).await;
5268 return Ok(record);
5269 }
5270 approval_record = Some(ToolApprovalRecord {
5271 status: ToolApprovalStatus::Modified,
5272 reason: None,
5273 modified_arguments: Some(executed_arguments.clone()),
5274 });
5275 }
5276 ApprovalResult::Rejected { reason } => {
5277 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5278 let record = self.record_from_parts(
5279 &request,
5280 canonical_id,
5281 executed_arguments,
5282 started_at,
5283 start,
5284 false,
5285 false,
5286 format!("Approval rejected: {}", reason),
5287 metadata,
5288 ToolPolicyDecisionRecord::approval(reason.clone()),
5289 Some(ToolApprovalRecord {
5290 status: ToolApprovalStatus::Rejected,
5291 reason: Some(reason),
5292 modified_arguments: None,
5293 }),
5294 false,
5295 false,
5296 );
5297 self.finish_tool_record(&record).await;
5298 return Ok(record);
5299 }
5300 ApprovalResult::Timeout => {
5301 let record = self.record_from_parts(
5302 &request,
5303 canonical_id,
5304 executed_arguments,
5305 started_at,
5306 start,
5307 false,
5308 false,
5309 "Approval timed out".to_string(),
5310 metadata,
5311 ToolPolicyDecisionRecord::approval("approval timeout"),
5312 Some(ToolApprovalRecord {
5313 status: ToolApprovalStatus::Timeout,
5314 reason: Some("approval timeout".to_string()),
5315 modified_arguments: None,
5316 }),
5317 false,
5318 false,
5319 );
5320 self.finish_tool_record(&record).await;
5321 return Ok(record);
5322 }
5323 }
5324 }
5325
5326 let condition_check = self
5327 .observe_purpose(
5328 ObservationPurpose::HitlLocalization,
5329 hitl_engine.check_conditions_with_localization(
5330 &executed_arguments,
5331 &hitl_lang_ctx,
5332 self.approval_handler.as_ref(),
5333 Some(&self.llm_registry),
5334 ),
5335 )
5336 .await?;
5337 if condition_check.is_required() {
5338 match self.request_hitl_approval(condition_check).await? {
5339 ApprovalResult::Approved => {
5340 merge_approved_record(&mut approval_record);
5341 }
5342 ApprovalResult::Modified { changes } => {
5343 if let Some(obj) = executed_arguments.as_object_mut() {
5344 for (key, value) in changes {
5345 obj.insert(key, value);
5346 }
5347 }
5348 let modified_security = security_engine
5349 .validate_tool_execution_with_bindings(
5350 &canonical_id,
5351 &executed_arguments,
5352 &bindings,
5353 )
5354 .await?;
5355 if !matches!(
5356 modified_security,
5357 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5358 ) {
5359 let reason = modified_security
5360 .reason()
5361 .unwrap_or("modified arguments failed policy")
5362 .to_string();
5363 let record = self.record_from_parts(
5364 &request,
5365 canonical_id,
5366 executed_arguments,
5367 started_at,
5368 start,
5369 false,
5370 false,
5371 reason.clone(),
5372 metadata,
5373 ToolPolicyDecisionRecord::deny(reason),
5374 approval_record,
5375 false,
5376 false,
5377 );
5378 self.finish_tool_record(&record).await;
5379 return Ok(record);
5380 }
5381 approval_record = Some(ToolApprovalRecord {
5382 status: ToolApprovalStatus::Modified,
5383 reason: None,
5384 modified_arguments: Some(executed_arguments.clone()),
5385 });
5386 }
5387 ApprovalResult::Rejected { reason } => {
5388 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5389 let record = self.record_from_parts(
5390 &request,
5391 canonical_id,
5392 executed_arguments,
5393 started_at,
5394 start,
5395 false,
5396 false,
5397 format!("Approval rejected: {}", reason),
5398 metadata,
5399 ToolPolicyDecisionRecord::approval(reason.clone()),
5400 Some(ToolApprovalRecord {
5401 status: ToolApprovalStatus::Rejected,
5402 reason: Some(reason),
5403 modified_arguments: None,
5404 }),
5405 false,
5406 false,
5407 );
5408 self.finish_tool_record(&record).await;
5409 return Ok(record);
5410 }
5411 ApprovalResult::Timeout => {
5412 let record = self.record_from_parts(
5413 &request,
5414 canonical_id,
5415 executed_arguments,
5416 started_at,
5417 start,
5418 false,
5419 false,
5420 "Approval timed out".to_string(),
5421 metadata,
5422 ToolPolicyDecisionRecord::approval("approval timeout"),
5423 Some(ToolApprovalRecord {
5424 status: ToolApprovalStatus::Timeout,
5425 reason: Some("approval timeout".to_string()),
5426 modified_arguments: None,
5427 }),
5428 false,
5429 false,
5430 );
5431 self.finish_tool_record(&record).await;
5432 return Ok(record);
5433 }
5434 }
5435 }
5436 }
5437
5438 executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5443 &canonical_id,
5444 &executed_arguments,
5445 &bindings,
5446 );
5447 if let Some(record) = approval_record.as_mut()
5448 && matches!(record.status, ToolApprovalStatus::Modified)
5449 {
5450 record.modified_arguments = Some(executed_arguments.clone());
5451 }
5452 let binding_security_result = security_engine
5453 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5454 .await?;
5455 let approval_confirmation_required = matches!(
5456 binding_security_result,
5457 SecurityCheckResult::RequireConfirmation { .. }
5458 ) || security_engine
5459 .classification_approval_message(
5460 &canonical_id,
5461 &resolved.tool.classify_call(&executed_arguments),
5462 )
5463 .is_some();
5464 let approval_binding = approval_record.as_ref().and_then(|record| {
5465 matches!(
5466 record.status,
5467 ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5468 )
5469 .then(|| ToolApprovalBinding {
5470 canonical_id: canonical_id.clone(),
5471 arguments: executed_arguments.clone(),
5472 confirmation_required: approval_confirmation_required,
5473 policy_version: security_engine.policy_version(),
5474 runtime_control_version: approval_control_snapshot.version,
5475 state_generation: initial_scope_snapshot.state_generation,
5476 reviewed_tool: Arc::clone(&resolved.tool),
5477 })
5478 });
5479
5480 let control_snapshot = self.runtime_safety_snapshot();
5485 let resolved = self.tools.resolve(&request.requested_name);
5486 let registry_version = self.tools.version();
5487 let mut versions = ToolDecisionVersions {
5488 policy: control_snapshot.tool_security.policy_version(),
5489 registry: registry_version,
5490 runtime_control: control_snapshot.version,
5491 state: None,
5492 };
5493 metadata.insert(
5494 "runtime_scope_snapshot".to_string(),
5495 serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5496 );
5497 let resolved = match resolved {
5498 Some(resolved) => resolved,
5499 None => {
5500 let reason = format!(
5501 "Tool '{}' became unavailable after approval",
5502 request.requested_name
5503 );
5504 let record = self.record_from_parts_at(
5505 &request,
5506 request.requested_name.clone(),
5507 executed_arguments,
5508 started_at,
5509 start,
5510 false,
5511 false,
5512 reason.clone(),
5513 metadata,
5514 ToolPolicyDecisionRecord::unavailable(reason),
5515 approval_record,
5516 false,
5517 false,
5518 versions,
5519 );
5520 self.finish_tool_record(&record).await;
5521 return Ok(record);
5522 }
5523 };
5524
5525 let canonical_id = resolved.identity.canonical_id.clone();
5526 let bindings = resolved.tool.policy_bindings();
5527 let final_arguments = control_snapshot
5528 .tool_security
5529 .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5530 if let Some(record) = approval_record.as_mut()
5531 && matches!(record.status, ToolApprovalStatus::Modified)
5532 {
5533 record.modified_arguments = Some(final_arguments.clone());
5534 }
5535 let classification = resolved.tool.classify_call(&final_arguments);
5536 let safety = resolved.tool.safety_metadata();
5537 let security_engine = control_snapshot.tool_security;
5538 let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
5539 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5540 let resource_lock_keys =
5541 tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5542 metadata.insert(
5543 "classification".to_string(),
5544 serde_json::to_value(&classification).unwrap_or(Value::Null),
5545 );
5546 metadata.insert(
5547 "effective_limits".to_string(),
5548 serde_json::to_value(&limits).unwrap_or(Value::Null),
5549 );
5550 metadata.insert(
5551 "resource_lock_keys".to_string(),
5552 serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5553 );
5554 if policy_snapshot.is_null() {
5555 metadata.remove("policy_snapshot");
5556 } else {
5557 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5558 }
5559
5560 let final_denial = |canonical_id: String,
5561 output: String,
5562 policy: ToolPolicyDecisionRecord,
5563 metadata: HashMap<String, Value>,
5564 decision_versions: ToolDecisionVersions| {
5565 self.record_from_parts_at(
5566 &request,
5567 canonical_id,
5568 final_arguments.clone(),
5569 started_at,
5570 start,
5571 false,
5572 false,
5573 output,
5574 metadata,
5575 policy,
5576 approval_record.clone(),
5577 false,
5578 false,
5579 decision_versions,
5580 )
5581 };
5582
5583 if control_snapshot.emergency_deny {
5584 let reason = "Tool execution is disabled by runtime control".to_string();
5585 let record = final_denial(
5586 canonical_id,
5587 reason.clone(),
5588 ToolPolicyDecisionRecord::deny(reason),
5589 metadata,
5590 versions,
5591 );
5592 self.finish_tool_record(&record).await;
5593 return Ok(record);
5594 }
5595
5596 let available_snapshot = self
5601 .get_available_tool_ids_snapshot_for_scope(
5602 control_snapshot.tool_scope_override.as_deref(),
5603 )
5604 .await?;
5605 versions.state = available_snapshot.state_generation;
5606 metadata.insert(
5607 "available_tool_ids_snapshot".to_string(),
5608 serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5609 );
5610 metadata.insert(
5611 "state_generation_snapshot".to_string(),
5612 serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5613 );
5614 if !available_snapshot
5615 .tool_ids
5616 .iter()
5617 .any(|tool_id| tool_id == &canonical_id)
5618 {
5619 let reason = format!(
5620 "Tool '{}' is not available in the final runtime scope",
5621 canonical_id
5622 );
5623 let record = final_denial(
5624 canonical_id,
5625 reason.clone(),
5626 ToolPolicyDecisionRecord::deny(reason),
5627 metadata,
5628 versions,
5629 );
5630 self.finish_tool_record(&record).await;
5631 return Ok(record);
5632 }
5633
5634 let final_security_result = security_engine
5639 .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5640 .await?;
5641 match &final_security_result {
5642 SecurityCheckResult::Block { reason } => {
5643 let record = final_denial(
5644 canonical_id,
5645 format!("Denied: {}", reason),
5646 ToolPolicyDecisionRecord::deny(reason.clone()),
5647 metadata,
5648 versions,
5649 );
5650 self.finish_tool_record(&record).await;
5651 return Ok(record);
5652 }
5653 SecurityCheckResult::Unavailable { reason } => {
5654 let record = final_denial(
5655 canonical_id,
5656 format!("Unavailable: {}", reason),
5657 ToolPolicyDecisionRecord::unavailable(reason.clone()),
5658 metadata,
5659 versions,
5660 );
5661 self.finish_tool_record(&record).await;
5662 return Ok(record);
5663 }
5664 SecurityCheckResult::Warn { message } => {
5665 warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5666 }
5667 SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5668 }
5669 let final_confirmation_required = matches!(
5670 final_security_result,
5671 SecurityCheckResult::RequireConfirmation { .. }
5672 ) || security_engine
5673 .classification_approval_message(&canonical_id, &classification)
5674 .is_some();
5675 let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5676 binding.is_stale(
5677 &canonical_id,
5678 &final_arguments,
5679 final_confirmation_required,
5680 versions,
5681 &resolved.tool,
5682 )
5683 });
5684 if stale_approval {
5685 let reason = "Approval became stale before final admission".to_string();
5686 let record = final_denial(
5687 canonical_id,
5688 reason.clone(),
5689 ToolPolicyDecisionRecord::deny(reason),
5690 metadata,
5691 versions,
5692 );
5693 self.finish_tool_record(&record).await;
5694 return Ok(record);
5695 }
5696 if final_confirmation_required && approval_binding.is_none() {
5697 let reason = "Final policy requires fresh approval".to_string();
5698 let record = final_denial(
5699 canonical_id,
5700 reason.clone(),
5701 ToolPolicyDecisionRecord::approval(reason),
5702 metadata,
5703 versions,
5704 );
5705 self.finish_tool_record(&record).await;
5706 return Ok(record);
5707 }
5708
5709 if let Some((_, reason)) = self.host_tool_unavailability(&canonical_id) {
5710 let record = final_denial(
5711 canonical_id,
5712 reason.to_string(),
5713 ToolPolicyDecisionRecord::unavailable(reason),
5714 metadata,
5715 versions,
5716 );
5717 self.finish_tool_record(&record).await;
5718 return Ok(record);
5719 }
5720
5721 let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5726 else {
5727 let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5728 let record = final_denial(
5729 canonical_id,
5730 reason.clone(),
5731 ToolPolicyDecisionRecord::deny(reason),
5732 metadata,
5733 versions,
5734 );
5735 self.finish_tool_record(&record).await;
5736 return Ok(record);
5737 };
5738
5739 let admission = self.admit_tool_execution(
5744 versions.runtime_control,
5745 versions.policy,
5746 versions.state,
5747 &canonical_id,
5748 );
5749 if !matches!(admission, SecurityCheckResult::Allow) {
5750 let latest_control = self.runtime_safety_snapshot();
5751 let reason = admission
5752 .reason()
5753 .unwrap_or("tool admission was denied")
5754 .to_string();
5755 let policy = if admission.is_unavailable() {
5756 ToolPolicyDecisionRecord::unavailable(reason.clone())
5757 } else {
5758 ToolPolicyDecisionRecord::deny(reason.clone())
5759 };
5760 let record = self.record_from_parts_at(
5761 &request,
5762 canonical_id,
5763 final_arguments,
5764 started_at,
5765 start,
5766 false,
5767 false,
5768 reason,
5769 metadata,
5770 policy,
5771 approval_record,
5772 false,
5773 false,
5774 ToolDecisionVersions {
5775 policy: latest_control.tool_security.policy_version(),
5776 registry: versions.registry,
5777 runtime_control: latest_control.version,
5778 state: self
5779 .state_machine
5780 .as_ref()
5781 .map(|state_machine| state_machine.generation()),
5782 },
5783 );
5784 self.finish_tool_record_after_resource_guards(resource_guards, &record)
5785 .await;
5786 return Ok(record);
5787 }
5788 let executed_arguments = final_arguments;
5789
5790 let tool_config = self.recovery_manager.get_tool_config(&canonical_id);
5791 let timeout_ms = limits
5792 .timeout_ms
5793 .unwrap_or_else(|| security_engine.get_tool_timeout(&canonical_id));
5794 let turn_actor = current_turn_actor_context();
5795 let actor = ToolActorContext {
5796 actor_id: turn_actor
5797 .as_ref()
5798 .and_then(|context| context.effective_actor_id().map(str::to_string))
5799 .or_else(|| self.actor_id()),
5800 origin_actor_id: turn_actor
5801 .as_ref()
5802 .and_then(|context| context.origin_actor_id.clone()),
5803 sender_agent_id: turn_actor
5804 .as_ref()
5805 .and_then(|context| context.sender_agent_id.clone()),
5806 };
5807 let tool_context = ToolExecutionContext {
5808 requested_name: request.requested_name.clone(),
5809 canonical_id: canonical_id.clone(),
5810 display_name: resolved.identity.display_name.clone(),
5811 provider_id: resolved.identity.provider_id.clone(),
5812 registry_version: versions.registry,
5813 policy_version: versions.policy,
5814 runtime_control_version: versions.runtime_control,
5815 call_id: request.call_id.clone(),
5816 source: request.source.clone(),
5817 actor,
5818 cancellation: ToolCancellationToken::new(
5819 Arc::clone(&self.runtime_control.emergency_deny),
5820 Some("runtime control cancellation".to_string()),
5821 ),
5822 started_at,
5823 deadline: None,
5824 permission: ToolPolicyDecisionRecord::allow(),
5825 approval: approval_record.clone(),
5826 classification: classification.clone(),
5827 safety,
5828 limits: limits.clone(),
5829 policy_snapshot,
5830 custom_config: security_engine.custom_config(&canonical_id),
5831 };
5832 let (mut result, timed_out, cancelled, invoked) = self
5833 .run_tool_with_retries(
5834 &canonical_id,
5835 resolved.tool.clone(),
5836 executed_arguments.clone(),
5837 tool_context,
5838 timeout_ms,
5839 tool_config.max_retries,
5840 )
5841 .await?;
5842
5843 let fallback_tool = if !result.success {
5844 match &tool_config.on_failure {
5845 ToolFailureAction::Skip => {
5846 result = ToolResult::ok(format!(
5847 "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
5848 canonical_id
5849 ));
5850 None
5851 }
5852 ToolFailureAction::Fallback { fallback_tool } => Some(fallback_tool.clone()),
5853 ToolFailureAction::ReportError => None,
5854 }
5855 } else {
5856 None
5857 };
5858
5859 let output_cap = limits.max_output_chars;
5860 let (output, output_truncated) =
5861 Self::truncate_tool_output(result.output.clone(), output_cap);
5862 if let Some(result_metadata) = result.metadata {
5863 metadata.extend(result_metadata);
5864 }
5865 let mut record = self.record_from_parts_at(
5866 &request,
5867 canonical_id,
5868 executed_arguments,
5869 started_at,
5870 start,
5871 invoked,
5872 result.success,
5873 output,
5874 metadata,
5875 ToolPolicyDecisionRecord::allow(),
5876 approval_record,
5877 timed_out,
5878 output_truncated,
5879 versions,
5880 );
5881 record.cancelled = cancelled;
5882 if cancelled {
5883 record.cancellation_reason = Some("runtime control cancellation".to_string());
5884 }
5885 if let Some(fallback_tool) = fallback_tool {
5886 let fallback_arguments = record.executed_arguments.clone();
5887 let original_tool = record.canonical_id.clone();
5888 self.finish_tool_record_after_resource_guards(resource_guards, &record)
5892 .await;
5893 let fallback_request = ToolExecutionRequest::new(
5894 request.call_id.clone(),
5895 fallback_tool,
5896 fallback_arguments,
5897 ToolCallSource::Fallback { original_tool },
5898 );
5899 return Box::pin(self.execute_tool_record(fallback_request)).await;
5900 }
5901 self.finish_tool_record_after_resource_guards(resource_guards, &record)
5902 .await;
5903 Ok(record)
5904 }
5905
5906 #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
5907 async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
5908 let record = self
5909 .execute_tool_record(ToolExecutionRequest::new(
5910 tool_call.id.clone(),
5911 tool_call.name.clone(),
5912 tool_call.arguments.clone(),
5913 ToolCallSource::Model,
5914 ))
5915 .await?;
5916 if record.success {
5917 Ok(record.model_output_string())
5918 } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
5919 Err(AgentError::HITLRejected(record.model_output_string()))
5920 } else {
5921 Err(AgentError::Tool(record.model_output_string()))
5922 }
5923 }
5924
5925 async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
5931 let Some(ref router) = self.skill_router else {
5932 return Ok(None);
5933 };
5934 let available_skills = self.get_available_skills();
5935 if available_skills.is_empty() {
5936 return Ok(None);
5937 }
5938 let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
5939 let Some(skill_id) = self
5940 .observe_purpose(
5941 ObservationPurpose::SkillRouting,
5942 router.select_skill_filtered(input, &skill_ids),
5943 )
5944 .await?
5945 else {
5946 return Ok(None);
5947 };
5948 let skill = router
5949 .get_skill(&skill_id)
5950 .cloned()
5951 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
5952 info!(skill_id = %skill_id, "Skill selected");
5953 Ok(Some(SkillCandidate::new(skill_id, skill)))
5954 }
5955
5956 async fn commit_skill_candidate_route_result(
5961 &self,
5962 candidate: SkillCandidate,
5963 input: &str,
5964 ) -> Result<SkillRouteResult> {
5965 let skill_id = candidate.skill_id;
5966 let skill = candidate.skill;
5967 let expected_state_generation = self
5968 .state_machine
5969 .as_ref()
5970 .map(|state_machine| state_machine.generation());
5971 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
5972 if let Some(ref skill_disambig) = skill.disambiguation
5973 && skill_disambig.enabled.unwrap_or(false)
5974 && let Some(ref disambiguator) = self.disambiguation_manager
5975 {
5976 let context = self.build_disambiguation_context().await?;
5977 let state_override = self
5978 .state_machine
5979 .as_ref()
5980 .and_then(|sm| sm.current_definition())
5981 .and_then(|def| def.disambiguation.clone());
5982
5983 let disambiguation_result = self
5984 .observe_purpose(
5985 ObservationPurpose::DisambiguationDetection,
5986 disambiguator.process_input_with_override(
5987 input,
5988 &context,
5989 state_override.as_ref(),
5990 Some(skill_disambig),
5991 ),
5992 )
5993 .await?;
5994 let current_state_generation = self
5995 .state_machine
5996 .as_ref()
5997 .map(|state_machine| state_machine.generation());
5998 if current_state_generation != expected_state_generation
5999 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6000 {
6001 disambiguator.clear_pending().await;
6002 *self.pending_skill_id.write() = None;
6003 return Err(AgentError::Other(
6004 "State or reset ownership changed during skill disambiguation".to_string(),
6005 ));
6006 }
6007 match disambiguation_result {
6008 DisambiguationResult::Clear => {
6009 debug!(skill_id = %skill_id, "Skill disambiguation: clear");
6010 }
6011 DisambiguationResult::NeedsClarification {
6012 question,
6013 detection,
6014 } => {
6015 let admission = self
6016 .admit_disambiguation_redispatch(
6017 expected_disambiguation_epoch,
6018 expected_state_generation,
6019 )
6020 .await?;
6021 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
6022 info!(
6023 skill_id = %skill_id,
6024 ambiguity_type = ?detection.ambiguity_type,
6025 confidence = detection.confidence,
6026 "Skill requires clarification before execution"
6027 );
6028 *self.pending_skill_id.write() = Some(skill_id.clone());
6029 let response = AgentResponse::new(&question.question).with_metadata(
6030 "disambiguation",
6031 serde_json::json!({
6032 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
6033 "skill_id": skill_id,
6034 "options": question.options,
6035 "clarifying": question.clarifying,
6036 "detection": {
6037 "type": detection.ambiguity_type,
6038 "confidence": detection.confidence,
6039 "what_is_unclear": detection.what_is_unclear,
6040 }
6041 }),
6042 );
6043 drop(admission);
6044 return Ok(SkillRouteResult::NeedsClarification {
6045 response,
6046 ownership: Some(DisambiguationOwnership {
6047 epoch: expected_disambiguation_epoch,
6048 state_generation: expected_state_generation,
6049 }),
6050 });
6051 }
6052 DisambiguationResult::Clarified { enriched_input, .. } => {
6053 info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
6054 let admission = self
6055 .admit_disambiguation_redispatch(
6056 expected_disambiguation_epoch,
6057 expected_state_generation,
6058 )
6059 .await?;
6060 drop(admission);
6061 let content = self.execute_skill(&skill, &enriched_input).await?;
6062 return Ok(SkillRouteResult::Response { skill_id, content });
6063 }
6064 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
6065 info!(skill_id = %skill_id, "Skill disambiguation best guess");
6066 let admission = self
6067 .admit_disambiguation_redispatch(
6068 expected_disambiguation_epoch,
6069 expected_state_generation,
6070 )
6071 .await?;
6072 drop(admission);
6073 let content = self.execute_skill(&skill, &enriched_input).await?;
6074 return Ok(SkillRouteResult::Response { skill_id, content });
6075 }
6076 DisambiguationResult::GiveUp { reason } => {
6077 warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
6078 let apology = self
6079 .generate_localized_apology(
6080 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
6081 &reason,
6082 )
6083 .await
6084 .unwrap_or_else(|_| {
6085 format!("I'm sorry, I couldn't understand your request: {}", reason)
6086 });
6087 return Ok(SkillRouteResult::NeedsClarification {
6088 response: AgentResponse::new(&apology),
6089 ownership: None,
6090 });
6091 }
6092 DisambiguationResult::Escalate { reason } => {
6093 info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6094 let apology = self
6095 .generate_localized_apology(
6096 "Explain briefly that you're transferring the user to a human agent for help.",
6097 &reason,
6098 )
6099 .await
6100 .unwrap_or_else(|_| {
6101 format!("I need human assistance to help with your request: {}", reason)
6102 });
6103 return Ok(SkillRouteResult::NeedsClarification {
6104 response: AgentResponse::new(&apology),
6105 ownership: None,
6106 });
6107 }
6108 DisambiguationResult::Abandoned { .. } => {
6109 debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6110 return Ok(SkillRouteResult::NoMatch);
6111 }
6112 }
6113 }
6114 let admission = self
6115 .admit_disambiguation_redispatch(
6116 expected_disambiguation_epoch,
6117 expected_state_generation,
6118 )
6119 .await?;
6120 drop(admission);
6121 let content = self.execute_skill(&skill, input).await?;
6122 Ok(SkillRouteResult::Response { skill_id, content })
6123 }
6124
6125 async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6127 if let Some(candidate) = self.select_skill_candidate(input).await? {
6128 self.commit_skill_candidate_route_result(candidate, input)
6129 .await
6130 } else {
6131 Ok(SkillRouteResult::NoMatch)
6132 }
6133 }
6134
6135 async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6137 if let Some(ref executor) = self.skill_executor {
6138 let skill_reasoning = self.get_skill_reasoning_config(skill);
6139 let skill_reflection = self.get_skill_reflection_config(skill);
6140
6141 debug!(
6142 skill_id = %skill.id,
6143 reasoning_mode = ?skill_reasoning.mode,
6144 reflection_enabled = ?skill_reflection.enabled,
6145 "Skill reasoning/reflection config"
6146 );
6147
6148 let response = self
6149 .observe_purpose(
6150 ObservationPurpose::SkillPrompt,
6151 executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6152 )
6153 .await?;
6154
6155 if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6156 let should_reflect = self
6157 .should_reflect_with_config(input, &response, &skill_reflection)
6158 .await?;
6159 if should_reflect {
6160 let evaluated = self
6161 .evaluate_and_retry_with_config(input, response, &skill_reflection)
6162 .await?;
6163 return Ok(evaluated);
6164 }
6165 }
6166
6167 return Ok(response);
6168 }
6169 Err(AgentError::Skill(
6170 "No skill executor configured".to_string(),
6171 ))
6172 }
6173
6174 async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6177 let skill = self
6178 .skill_router
6179 .as_ref()
6180 .and_then(|r| r.get_skill(skill_id).cloned())
6181 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6182 self.execute_skill(&skill, input).await
6183 }
6184
6185 async fn should_reflect_with_config(
6186 &self,
6187 input: &str,
6188 response: &str,
6189 config: &ReflectionConfig,
6190 ) -> Result<bool> {
6191 if !config.requires_evaluation() {
6192 return Ok(false);
6193 }
6194
6195 if config.is_enabled() {
6196 return Ok(true);
6197 }
6198
6199 let evaluator_llm = config
6200 .evaluator_llm
6201 .as_ref()
6202 .and_then(|alias| self.llm_registry.get(alias).ok())
6203 .or_else(|| self.llm_registry.router().ok())
6204 .or_else(|| self.llm_registry.default().ok());
6205
6206 let Some(llm) = evaluator_llm else {
6207 return Ok(false);
6208 };
6209
6210 let response_preview: String = response.chars().take(500).collect();
6211 let prompt = format!(
6212 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6213
6214User query: "{}"
6215Response: "{}"
6216
6217Answer YES or NO only."#,
6218 input, response_preview
6219 );
6220
6221 let messages = vec![ChatMessage::user(&prompt)];
6222 let result = self
6223 .observe_purpose(
6224 ObservationPurpose::ReflectionDecision,
6225 llm.complete(&messages, None),
6226 )
6227 .await;
6228
6229 match result {
6230 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6231 Err(_) => Ok(false),
6232 }
6233 }
6234
6235 async fn evaluate_and_retry_with_config(
6236 &self,
6237 input: &str,
6238 mut response: String,
6239 config: &ReflectionConfig,
6240 ) -> Result<String> {
6241 let llm = self.get_state_llm()?;
6242 let mut attempts = 0u32;
6243 let max_retries = config.max_retries;
6244
6245 loop {
6246 let evaluation = self
6247 .evaluate_response_with_config(input, &response, config)
6248 .await?;
6249
6250 if evaluation.passed || attempts >= max_retries {
6251 info!(
6252 passed = evaluation.passed,
6253 confidence = evaluation.confidence,
6254 attempts = attempts + 1,
6255 "Skill reflection evaluation complete"
6256 );
6257 return Ok(response);
6258 }
6259
6260 debug!(
6261 attempt = attempts + 1,
6262 failed_criteria = evaluation.failed_criteria().count(),
6263 "Skill response did not meet criteria, retrying"
6264 );
6265
6266 let feedback: Vec<String> = evaluation
6267 .failed_criteria()
6268 .map(|c| format!("- {}", c.criterion))
6269 .collect();
6270
6271 let retry_prompt = format!(
6272 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6273 feedback.join("\n"),
6274 input
6275 );
6276
6277 let messages = vec![ChatMessage::user(&retry_prompt)];
6278 let retry_response = self
6279 .observe_purpose(
6280 ObservationPurpose::ReflectionEvaluation,
6281 llm.complete(&messages, None),
6282 )
6283 .await
6284 .map_err(|e| AgentError::LLM(e.to_string()))?;
6285
6286 response = retry_response.content.trim().to_string();
6287 attempts += 1;
6288 }
6289 }
6290
6291 async fn evaluate_response_with_config(
6292 &self,
6293 input: &str,
6294 response: &str,
6295 config: &ReflectionConfig,
6296 ) -> Result<EvaluationResult> {
6297 let evaluator_llm = config
6298 .evaluator_llm
6299 .as_ref()
6300 .and_then(|alias| self.llm_registry.get(alias).ok())
6301 .or_else(|| self.llm_registry.router().ok())
6302 .or_else(|| self.llm_registry.default().ok())
6303 .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6304
6305 let criteria = &config.criteria;
6306 let criteria_list = criteria
6307 .iter()
6308 .enumerate()
6309 .map(|(i, c)| format!("{}. {}", i + 1, c))
6310 .collect::<Vec<_>>()
6311 .join("\n");
6312
6313 let prompt = format!(
6314 r#"Evaluate this response against the criteria.
6315
6316User query: "{}"
6317
6318Response to evaluate: "{}"
6319
6320Criteria:
6321{}
6322
6323For each criterion, respond with:
6324- criterion number
6325- PASS or FAIL
6326- brief reason
6327
6328Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6329
6330Format:
63311. PASS/FAIL - reason
63322. PASS/FAIL - reason
6333...
6334CONFIDENCE: 0.X
6335OVERALL: PASS/FAIL"#,
6336 input, response, criteria_list
6337 );
6338
6339 let messages = vec![ChatMessage::user(&prompt)];
6340 let eval_response = self
6341 .observe_purpose(
6342 ObservationPurpose::ReflectionEvaluation,
6343 evaluator_llm.complete(&messages, None),
6344 )
6345 .await
6346 .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6347
6348 let content = eval_response.content.to_uppercase();
6349 let llm_pass = content.contains("OVERALL: PASS");
6350
6351 let confidence = content
6352 .lines()
6353 .find(|l| l.contains("CONFIDENCE:"))
6354 .and_then(|l| {
6355 l.split(':')
6356 .nth(1)
6357 .and_then(|v| v.trim().parse::<f32>().ok())
6358 })
6359 .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6360
6361 let overall_pass = llm_pass && confidence >= config.pass_threshold;
6364
6365 let mut criteria_results = Vec::new();
6366 for (i, criterion) in criteria.iter().enumerate() {
6367 let line_marker = format!("{}.", i + 1);
6368 let passed = eval_response
6369 .content
6370 .lines()
6371 .find(|l| l.contains(&line_marker))
6372 .map(|l| l.to_uppercase().contains("PASS"))
6373 .unwrap_or(overall_pass);
6374
6375 if passed {
6376 criteria_results.push(CriterionResult::pass(criterion));
6377 } else {
6378 criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6379 }
6380 }
6381
6382 Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6383 }
6384
6385 async fn process_input(&self, input: &str) -> Result<ProcessData> {
6387 if let Some(processor) = self.get_state_process_processor() {
6388 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6389 return self
6390 .observe_purpose(purpose, processor.process_input(input))
6391 .await;
6392 }
6393 if let Some(ref processor) = self.process_processor {
6394 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6395 self.observe_purpose(purpose, processor.process_input(input))
6396 .await
6397 } else {
6398 Ok(ProcessData::new(input))
6399 }
6400 }
6401
6402 async fn process_output(
6404 &self,
6405 output: &str,
6406 input_context: &std::collections::HashMap<String, serde_json::Value>,
6407 ) -> Result<ProcessData> {
6408 if let Some(processor) = self.get_state_process_processor() {
6409 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6410 return self
6411 .observe_purpose(purpose, processor.process_output(output, input_context))
6412 .await;
6413 }
6414 if let Some(ref processor) = self.process_processor {
6415 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6416 self.observe_purpose(purpose, processor.process_output(output, input_context))
6417 .await
6418 } else {
6419 Ok(ProcessData::new(output))
6420 }
6421 }
6422
6423 fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6425 let sm = self.state_machine.as_ref()?;
6426 let def = sm.current_definition()?;
6427 let config = def.process.as_ref()?;
6428 let mut processor = ProcessProcessor::new(config.clone());
6429 if let Some(ref registry) = Some(self.llm_registry.clone()) {
6430 processor = processor.with_llm_registry(registry.clone());
6431 }
6432 processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6433 Some(processor)
6434 }
6435
6436 async fn check_turn_timeout(&self) -> Result<()> {
6438 let Some(ref sm) = self.state_machine else {
6439 return Ok(());
6440 };
6441 let Some(timeout_state) = sm.check_timeout() else {
6442 return Ok(());
6443 };
6444 let claim_admission = self.disambiguation_admission.write().await;
6445 if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6446 return Ok(());
6447 }
6448 let Some(reservation) = self.reserve_state_transition() else {
6449 return Ok(());
6450 };
6451 let from_state = sm.current();
6452 let expected_state_generation = sm.generation();
6453 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6454 let history_before = sm.history();
6455 drop(claim_admission);
6456
6457 self.execute_state_exit_actions(&from_state).await;
6458
6459 let admission = self.disambiguation_admission.write().await;
6460 if sm.current() != from_state
6461 || sm.generation() != expected_state_generation
6462 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6463 || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6464 {
6465 return Ok(());
6466 }
6467 sm.transition_to(&timeout_state, "max_turns exceeded")?;
6468 self.invalidate_pending_confirmation("state_timeout").await;
6469 let entered = sm.current();
6470 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6471 drop(admission);
6472
6473 self.execute_state_enter_actions(&entered, is_reentry).await;
6474 drop(reservation);
6475 info!(to = %entered, "Timeout transition");
6476 Ok(())
6477 }
6478
6479 fn increment_turn(&self) {
6480 if let Some(ref sm) = self.state_machine {
6481 sm.increment_turn();
6482 }
6483 }
6484
6485 fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6486 let sm = self.state_machine.as_ref()?;
6487 let current = sm.current();
6488 let transitions: Vec<_> = sm
6489 .auto_transitions()
6490 .into_iter()
6491 .filter(|t| match t.cooldown_turns {
6492 Some(cd) if cd > 0 => {
6493 let resolved = sm.config().resolve_full_path(¤t, &t.to);
6494 !sm.is_on_cooldown(&resolved, cd)
6495 }
6496 _ => true,
6497 })
6498 .collect();
6499 Some((transitions, current))
6500 }
6501
6502 fn transition_reason(transition: &Transition) -> String {
6503 if transition.when.is_empty() {
6504 "guard condition met".to_string()
6505 } else {
6506 transition.when.clone()
6507 }
6508 }
6509
6510 fn build_transition_context(
6512 &self,
6513 user_message: &str,
6514 response: &str,
6515 current_state: &str,
6516 staged: Option<&HashMap<String, Value>>,
6517 ) -> TransitionContext {
6518 let context_map = staged
6519 .map(|writes| self.build_context_with_staged(writes))
6520 .unwrap_or_else(|| self.build_context_with_overlays());
6521 TransitionContext::new(user_message, response, current_state).with_context(context_map)
6522 }
6523
6524 async fn select_transition_candidate(
6526 &self,
6527 user_message: &str,
6528 response: &str,
6529 ) -> Result<Option<TransitionCandidate>> {
6530 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6531 return Ok(None);
6532 };
6533 let transitions: Vec<Transition> = transitions
6534 .into_iter()
6535 .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6536 .collect();
6537 if transitions.is_empty() {
6538 return Ok(None);
6539 }
6540 let Some(evaluator) = self.transition_evaluator.as_ref() else {
6541 return Ok(None);
6542 };
6543 let context = self.build_transition_context(user_message, response, ¤t_state, None);
6544 let selected = self
6545 .observe_purpose(
6546 ObservationPurpose::StateTransitionEvaluation,
6547 evaluator.select_transition(&transitions, &context),
6548 )
6549 .await?;
6550 Ok(selected.map(|index| {
6551 let transition = transitions[index].clone();
6552 TransitionCandidate::new(
6553 current_state,
6554 transition.clone(),
6555 Self::transition_reason(&transition),
6556 )
6557 }))
6558 }
6559
6560 fn select_deterministic_transition_candidate(
6562 &self,
6563 user_message: &str,
6564 current_state: &str,
6565 transitions: &[Transition],
6566 staged: &HashMap<String, Value>,
6567 ) -> Option<TransitionCandidate> {
6568 let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6569
6570 for transition in transitions {
6571 if let Some(guard) = transition.guard.as_ref()
6572 && evaluate_guard(guard, &context)
6573 {
6574 return Some(TransitionCandidate::new(
6575 current_state,
6576 transition.clone(),
6577 Self::transition_reason(transition),
6578 ));
6579 }
6580 }
6581
6582 let resolved_intent = context
6583 .context
6584 .get("resolved_intent")
6585 .and_then(Value::as_str)
6586 .filter(|value| !value.is_empty());
6587 if let Some(resolved_intent) = resolved_intent {
6588 for transition in transitions {
6589 if transition.intent.as_deref() == Some(resolved_intent) {
6590 return Some(TransitionCandidate::new(
6591 current_state,
6592 transition.clone(),
6593 Self::transition_reason(transition),
6594 ));
6595 }
6596 }
6597 }
6598
6599 None
6600 }
6601
6602 async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6604 self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6605 .await
6606 }
6607
6608 async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6610 let approved = self.check_state_hitl(Some(from_state), target).await?;
6611 if !approved {
6612 info!(to = %target, "State transition rejected by HITL");
6613 }
6614 Ok(approved)
6615 }
6616
6617 async fn apply_transition_target(
6619 &self,
6620 from_state: &str,
6621 target: &str,
6622 reason: &str,
6623 staged: Option<&HashMap<String, Value>>,
6624 ) -> Result<bool> {
6625 let Some(ref sm) = self.state_machine else {
6626 return Ok(false);
6627 };
6628 let claim_admission = self.disambiguation_admission.write().await;
6629 if sm.current() != from_state {
6630 return Ok(false);
6631 }
6632 let Some(reservation) = self.reserve_state_transition() else {
6633 return Ok(false);
6634 };
6635 let expected_state_generation = sm.generation();
6636 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6637 let history_before = sm.history();
6638 drop(claim_admission);
6639
6640 self.execute_state_exit_actions(from_state).await;
6641
6642 let admission = self.disambiguation_admission.write().await;
6643 if sm.current() != from_state
6644 || sm.generation() != expected_state_generation
6645 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6646 {
6647 return Ok(false);
6648 }
6649 sm.transition_to(target, reason)?;
6650 self.invalidate_pending_confirmation("state_transition")
6651 .await;
6652 sm.reset_no_transition();
6653 if let Some(staged) = staged {
6654 self.commit_staged_context_writes(staged);
6655 }
6656 let entered = sm.current();
6657 let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6658 drop(admission);
6659
6660 self.execute_state_enter_actions(&entered, is_reentry).await;
6661 drop(reservation);
6662 self.hooks
6663 .on_state_transition(Some(from_state), &entered, reason)
6664 .await;
6665 info!(from = %from_state, to = %entered, "State transition");
6666 Ok(true)
6667 }
6668
6669 async fn commit_transition_target(
6671 &self,
6672 from_state: &str,
6673 target: &str,
6674 reason: &str,
6675 ) -> Result<bool> {
6676 if !self.approve_transition_target(from_state, target).await? {
6677 return Ok(false);
6678 }
6679 self.apply_transition_target(from_state, target, reason, None)
6680 .await
6681 }
6682
6683 async fn apply_pre_response_transition_candidate(
6685 &self,
6686 candidate: &TransitionCandidate,
6687 staged: &HashMap<String, Value>,
6688 processed_input: &str,
6689 ) -> Result<bool> {
6690 self.commit_root_user_message(processed_input).await?;
6691 self.apply_transition_target(
6692 &candidate.from_state,
6693 candidate.target(),
6694 &candidate.reason,
6695 Some(staged),
6696 )
6697 .await
6698 }
6699
6700 async fn commit_pre_response_transition_candidate(
6702 &self,
6703 candidate: &TransitionCandidate,
6704 staged: &HashMap<String, Value>,
6705 processed_input: &str,
6706 ) -> Result<bool> {
6707 if !self
6708 .approve_transition_target(&candidate.from_state, candidate.target())
6709 .await?
6710 {
6711 return Ok(false);
6712 }
6713 self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6714 .await
6715 }
6716
6717 async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6719 let Some(ref sm) = self.state_machine else {
6720 return Ok(false);
6721 };
6722 sm.increment_no_transition();
6723 let Some(fallback) = sm.check_fallback() else {
6724 return Ok(false);
6725 };
6726 self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6727 .await
6728 }
6729
6730 async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6732 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6733 return Ok(false);
6734 };
6735 if transitions.is_empty() {
6736 return Ok(false);
6737 }
6738 if let Some(candidate) = self
6739 .select_transition_candidate(user_message, response)
6740 .await?
6741 {
6742 return self.commit_transition_candidate(&candidate).await;
6743 }
6744 self.handle_transition_miss(¤t_state).await
6745 }
6746
6747 async fn try_pre_response_transition(
6749 &self,
6750 processed_input: &str,
6751 ) -> Result<Option<AgentResponse>> {
6752 let optimization = &self.runtime_config.optimization;
6753 if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6754 return Ok(None);
6755 }
6756 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6757 return Ok(None);
6758 };
6759 let eligible: Vec<Transition> = transitions
6760 .into_iter()
6761 .filter(|transition| !transition.requires_response)
6762 .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
6763 .collect();
6764 if eligible.is_empty() {
6765 return Ok(None);
6766 }
6767
6768 let empty_staged = HashMap::new();
6769 let mut extracted_staged: Option<HashMap<String, Value>> = None;
6770 let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
6771
6772 for transition in &eligible {
6773 let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
6774 let staged_for_eval = if use_extractors {
6775 if extracted_staged.is_none() {
6776 extracted_staged =
6777 Some(self.run_context_extractors_staged(processed_input).await);
6778 }
6779 extracted_staged.as_ref().unwrap_or(&empty_staged)
6780 } else {
6781 &empty_staged
6782 };
6783
6784 if let Some(candidate) = self.select_deterministic_transition_candidate(
6785 processed_input,
6786 ¤t_state,
6787 std::slice::from_ref(transition),
6788 staged_for_eval,
6789 ) {
6790 let staged_for_commit = if use_extractors {
6791 staged_for_eval.clone()
6792 } else {
6793 HashMap::new()
6794 };
6795 selected = Some((candidate, staged_for_commit));
6796 break;
6797 }
6798 }
6799
6800 let Some((candidate, staged)) = selected else {
6801 return Ok(None);
6802 };
6803
6804 if !self
6805 .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
6806 .await?
6807 {
6808 return Ok(None);
6809 }
6810 self.redispatch_current_state(processed_input)
6811 .await
6812 .map(Some)
6813 }
6814
6815 async fn try_speculative_branches(
6820 &self,
6821 processed_input: &str,
6822 input_context: &HashMap<String, Value>,
6823 ) -> Result<Option<AgentResponse>> {
6824 let optimization = &self.runtime_config.optimization;
6825 if !optimization.enabled {
6826 return Ok(None);
6827 }
6828
6829 let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
6830 if !matches!(
6831 effective_reasoning_mode,
6832 ReasoningMode::None | ReasoningMode::Auto
6833 ) {
6834 return Ok(None);
6835 }
6836
6837 let mut transition_enabled =
6838 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
6839 let mut skill_enabled = optimization.speculative_skill_routing
6840 && self.skill_router.is_some()
6841 && self.pending_skill_id.read().is_none();
6842 let mut reasoning_enabled = optimization.speculative_reasoning_auto
6843 && matches!(effective_reasoning_mode, ReasoningMode::Auto);
6844
6845 if matches!(effective_reasoning_mode, ReasoningMode::Auto)
6846 && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
6847 {
6848 return Ok(None);
6849 }
6850
6851 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6852 return Ok(None);
6853 }
6854
6855 let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
6856 let mut speculative_call_slots = optimization
6857 .max_speculative_llm_calls_per_turn
6858 .saturating_sub(1);
6859 if reasoning_enabled {
6860 if optional_slots == 0 || speculative_call_slots == 0 {
6861 return Ok(None);
6862 }
6863 optional_slots -= 1;
6864 speculative_call_slots -= 1;
6865 }
6866 if transition_enabled {
6867 if optional_slots == 0 {
6868 transition_enabled = false;
6869 } else {
6870 optional_slots -= 1;
6871 }
6872 }
6873 if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
6874 skill_enabled = false;
6875 }
6876
6877 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6878 return Ok(None);
6879 }
6880
6881 let main_kind = if transition_enabled {
6882 RuntimeOptimizationKind::ParallelStateTransition
6883 } else if skill_enabled {
6884 RuntimeOptimizationKind::SpeculativeSkillRouting
6885 } else {
6886 RuntimeOptimizationKind::SpeculativeReasoningAuto
6887 };
6888 if !self.reserve_active_speculative_llm_call(main_kind) {
6889 return Ok(None);
6890 }
6891
6892 let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
6893 let main_branch = RuntimeBranch::new(
6894 RuntimeTaskPurpose::MainResponse,
6895 main_kind,
6896 RuntimeTaskPriority::Normal,
6897 RuntimeCommitBehavior::FinalResponse,
6898 );
6899 let transition_branch = RuntimeBranch::new(
6900 RuntimeTaskPurpose::StateTransition,
6901 RuntimeOptimizationKind::ParallelStateTransition,
6902 RuntimeTaskPriority::Critical,
6903 RuntimeCommitBehavior::TransitionDecision,
6904 );
6905 let skill_branch = RuntimeBranch::new(
6906 RuntimeTaskPurpose::SkillRouting,
6907 RuntimeOptimizationKind::SpeculativeSkillRouting,
6908 RuntimeTaskPriority::High,
6909 RuntimeCommitBehavior::SkillSelection,
6910 );
6911 let reasoning_branch = RuntimeBranch::new(
6912 RuntimeTaskPurpose::ReasoningJudge,
6913 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6914 RuntimeTaskPriority::Normal,
6915 RuntimeCommitBehavior::ReasoningDecision,
6916 );
6917 let main_id = main_branch.branch_id();
6918 let transition_id = transition_branch.branch_id();
6919 let skill_id = skill_branch.branch_id();
6920 let reasoning_id = reasoning_branch.branch_id();
6921
6922 let main_id_for_future = main_id.clone();
6923 if !branch_set.schedule(
6924 main_branch,
6925 Box::pin(async move {
6926 match crate::optimization::observability::with_branch_observation(
6927 &main_id_for_future,
6928 main_kind,
6929 RuntimeCommitBehavior::FinalResponse,
6930 self.generate_main_response_draft(processed_input, &ReasoningMode::None),
6931 )
6932 .await
6933 {
6934 Ok(draft) => RuntimeBranchResult::MainDraft(draft),
6935 Err(error) => RuntimeBranchResult::Failed(error),
6936 }
6937 }),
6938 ) {
6939 return Ok(None);
6940 }
6941
6942 if transition_enabled {
6943 let transition_id_for_future = transition_id.clone();
6944 if !branch_set.schedule(
6945 transition_branch,
6946 Box::pin(async move {
6947 match crate::optimization::observability::with_branch_observation(
6948 &transition_id_for_future,
6949 RuntimeOptimizationKind::ParallelStateTransition,
6950 RuntimeCommitBehavior::TransitionDecision,
6951 self.select_parallel_transition_candidate(processed_input),
6952 )
6953 .await
6954 {
6955 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
6956 RuntimeBranchResult::Transition(Some(candidate))
6957 }
6958 Ok(ParallelTransitionSelection::NoMatch) => {
6959 RuntimeBranchResult::Transition(None)
6960 }
6961 Ok(ParallelTransitionSelection::ReservationExhausted) => {
6962 RuntimeBranchResult::Cancelled
6963 }
6964 Err(error) => RuntimeBranchResult::Failed(error),
6965 }
6966 }),
6967 ) {
6968 transition_enabled = false;
6969 }
6970 }
6971
6972 if skill_enabled {
6973 let skill_id_for_future = skill_id.clone();
6974 if !branch_set.schedule(
6975 skill_branch,
6976 Box::pin(async move {
6977 if !self.reserve_active_speculative_llm_call(
6978 RuntimeOptimizationKind::SpeculativeSkillRouting,
6979 ) {
6980 return RuntimeBranchResult::Cancelled;
6981 }
6982 match crate::optimization::observability::with_branch_observation(
6983 &skill_id_for_future,
6984 RuntimeOptimizationKind::SpeculativeSkillRouting,
6985 RuntimeCommitBehavior::SkillSelection,
6986 self.select_skill_candidate(processed_input),
6987 )
6988 .await
6989 {
6990 Ok(candidate) => RuntimeBranchResult::Skill(candidate),
6991 Err(error) => RuntimeBranchResult::Failed(error),
6992 }
6993 }),
6994 ) {
6995 skill_enabled = false;
6996 }
6997 }
6998
6999 if reasoning_enabled {
7000 let reasoning_id_for_future = reasoning_id.clone();
7001 if !branch_set.schedule(
7002 reasoning_branch,
7003 Box::pin(async move {
7004 if !self.reserve_active_speculative_llm_call(
7005 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7006 ) {
7007 return RuntimeBranchResult::Cancelled;
7008 }
7009 match crate::optimization::observability::with_branch_observation(
7010 &reasoning_id_for_future,
7011 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7012 RuntimeCommitBehavior::ReasoningDecision,
7013 self.determine_reasoning_mode_strict(processed_input),
7014 )
7015 .await
7016 {
7017 Ok(mode) => RuntimeBranchResult::Reasoning(mode),
7018 Err(error) => RuntimeBranchResult::Failed(error),
7019 }
7020 }),
7021 ) {
7022 reasoning_enabled = false;
7023 }
7024 }
7025
7026 if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
7027 self.finalize_pending_branches(branch_set.cancel_pending());
7028 return Ok(None);
7029 }
7030
7031 if !transition_enabled && !skill_enabled && !reasoning_enabled {
7032 self.finalize_pending_branches(branch_set.cancel_pending());
7033 return Ok(None);
7034 }
7035
7036 let mut main_pending = true;
7037 let mut skill_pending = skill_enabled;
7038 let mut reasoning_pending = reasoning_enabled;
7039 let mut transition_finalized = !transition_enabled;
7040 let mut skill_finalized = !skill_enabled;
7041 let mut reasoning_finalized = !reasoning_enabled;
7042 let mut main_result: Option<Result<MainResponseDraft>> = None;
7043 let mut transition_candidate: Option<TransitionCandidate> = None;
7044 let mut skill_candidate: Option<SkillCandidate> = None;
7045 let mut reasoning_decision: Option<ReasoningMode> = None;
7046 let mut transition_fallback_required = false;
7047 let mut skill_fallback_required = false;
7048 let mut reasoning_fallback_required = false;
7049
7050 loop {
7051 if let Some(candidate) = transition_candidate.take() {
7052 if self
7053 .approve_transition_target(&candidate.from_state, candidate.target())
7054 .await?
7055 {
7056 self.finalize_pending_branches(branch_set.cancel_pending());
7058 if !main_pending {
7059 self.finalize_branch_loss(
7060 &main_id,
7061 main_kind,
7062 RuntimeCommitBehavior::FinalResponse,
7063 false,
7064 main_result.as_ref().map(|result| result.is_err()),
7065 );
7066 }
7067 if skill_enabled && !skill_pending {
7068 self.finalize_branch_loss(
7069 &skill_id,
7070 RuntimeOptimizationKind::SpeculativeSkillRouting,
7071 RuntimeCommitBehavior::SkillSelection,
7072 false,
7073 Some(false),
7074 );
7075 }
7076 if reasoning_enabled && !reasoning_pending {
7077 self.finalize_branch_loss(
7078 &reasoning_id,
7079 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7080 RuntimeCommitBehavior::ReasoningDecision,
7081 false,
7082 Some(false),
7083 );
7084 }
7085 if !self
7086 .apply_pre_response_transition_candidate(
7087 &candidate,
7088 &HashMap::new(),
7089 processed_input,
7090 )
7091 .await?
7092 {
7093 self.finalize_optional_branch(
7094 &transition_id,
7095 RuntimeOptimizationKind::ParallelStateTransition,
7096 RuntimeCommitBehavior::TransitionDecision,
7097 "discarded",
7098 false,
7099 );
7100 return Ok(None);
7101 }
7102 self.finalize_optional_branch(
7103 &transition_id,
7104 RuntimeOptimizationKind::ParallelStateTransition,
7105 RuntimeCommitBehavior::TransitionDecision,
7106 "committed",
7107 true,
7108 );
7109 return self
7110 .redispatch_current_state(processed_input)
7111 .await
7112 .map(Some);
7113 }
7114 self.finalize_optional_branch(
7115 &transition_id,
7116 RuntimeOptimizationKind::ParallelStateTransition,
7117 RuntimeCommitBehavior::TransitionDecision,
7118 "discarded",
7119 false,
7120 );
7121 transition_finalized = true;
7122 }
7123
7124 if transition_finalized && skill_candidate.is_some() {
7125 let candidate = skill_candidate.take().unwrap();
7126 self.finalize_optional_branch(
7127 &skill_id,
7128 RuntimeOptimizationKind::SpeculativeSkillRouting,
7129 RuntimeCommitBehavior::SkillSelection,
7130 "committed",
7131 true,
7132 );
7133 if !main_pending {
7134 self.finalize_branch_loss(
7135 &main_id,
7136 main_kind,
7137 RuntimeCommitBehavior::FinalResponse,
7138 false,
7139 main_result.as_ref().map(|result| result.is_err()),
7140 );
7141 }
7142 if reasoning_enabled && !reasoning_pending {
7143 self.finalize_branch_loss(
7144 &reasoning_id,
7145 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7146 RuntimeCommitBehavior::ReasoningDecision,
7147 false,
7148 Some(false),
7149 );
7150 }
7151 self.finalize_pending_branches(branch_set.cancel_pending());
7152 self.commit_root_user_message(processed_input).await?;
7153 return match self
7154 .commit_skill_candidate_route_result(candidate, processed_input)
7155 .await?
7156 {
7157 SkillRouteResult::Response { skill_id, content } => self
7158 .handle_skill_response(processed_input, &skill_id, content, input_context)
7159 .await
7160 .map(Some),
7161 SkillRouteResult::NeedsClarification {
7162 response,
7163 ownership,
7164 } => {
7165 let admission = self
7166 .admit_optional_disambiguation_ownership(ownership)
7167 .await?;
7168 if response
7169 .metadata
7170 .as_ref()
7171 .and_then(|m| m.get("disambiguation"))
7172 .and_then(|d| d.get("status"))
7173 .and_then(|s| s.as_str())
7174 == Some("awaiting_clarification")
7175 {
7176 self.memory
7177 .add_message(ChatMessage::assistant(&response.content))
7178 .await?;
7179 }
7180 drop(admission);
7181 self.finish_turn_if_root(&response).await?;
7182 Ok(Some(response))
7183 }
7184 SkillRouteResult::NoMatch => Ok(None),
7185 };
7186 }
7187
7188 if transition_finalized
7189 && skill_finalized
7190 && let Some(reasoning_mode) = reasoning_decision.take()
7191 {
7192 if !matches!(reasoning_mode, ReasoningMode::None) {
7193 self.finalize_optional_branch(
7194 &reasoning_id,
7195 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7196 RuntimeCommitBehavior::ReasoningDecision,
7197 "committed",
7198 true,
7199 );
7200 if !main_pending {
7201 self.finalize_branch_loss(
7202 &main_id,
7203 main_kind,
7204 RuntimeCommitBehavior::FinalResponse,
7205 false,
7206 main_result.as_ref().map(|result| result.is_err()),
7207 );
7208 }
7209 self.finalize_pending_branches(branch_set.cancel_pending());
7210 self.commit_root_user_message(processed_input).await?;
7211 return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7212 self.handle_plan_and_execute(processed_input, input_context, true)
7213 .await
7214 .map(Some)
7215 } else {
7216 self.run_committed_response_loop_with_reasoning(
7217 processed_input,
7218 input_context,
7219 reasoning_mode,
7220 true,
7221 )
7222 .await
7223 .map(Some)
7224 };
7225 }
7226 self.finalize_optional_branch(
7227 &reasoning_id,
7228 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7229 RuntimeCommitBehavior::ReasoningDecision,
7230 "committed",
7231 true,
7232 );
7233 reasoning_finalized = true;
7234 }
7235
7236 if transition_finalized && skill_finalized && reasoning_finalized {
7237 if transition_fallback_required
7238 || skill_fallback_required
7239 || reasoning_fallback_required
7240 {
7241 if !main_pending {
7242 self.finalize_branch_loss(
7243 &main_id,
7244 main_kind,
7245 RuntimeCommitBehavior::FinalResponse,
7246 false,
7247 main_result.as_ref().map(|result| result.is_err()),
7248 );
7249 }
7250 self.finalize_pending_branches(branch_set.cancel_pending());
7251 return Ok(None);
7252 }
7253
7254 if let Some(result) = main_result.take() {
7255 let draft = match result {
7256 Ok(draft) => draft,
7257 Err(error) => {
7258 self.finalize_optional_branch(
7259 &main_id,
7260 main_kind,
7261 RuntimeCommitBehavior::FinalResponse,
7262 "failed",
7263 false,
7264 );
7265 self.finalize_pending_branches(branch_set.cancel_pending());
7266 return Err(error);
7267 }
7268 };
7269 self.finalize_optional_branch(
7270 &main_id,
7271 main_kind,
7272 RuntimeCommitBehavior::FinalResponse,
7273 "committed",
7274 true,
7275 );
7276 self.finalize_pending_branches(branch_set.cancel_pending());
7277 return self
7278 .commit_main_response_draft(
7279 processed_input,
7280 input_context,
7281 draft,
7282 ReasoningMode::None,
7283 reasoning_enabled,
7284 )
7285 .await
7286 .map(Some);
7287 }
7288 }
7289
7290 if branch_set.is_empty() {
7291 return Ok(None);
7292 }
7293
7294 let Some(outcome) = branch_set.next_completed().await else {
7295 return Ok(None);
7296 };
7297 let branch_id = outcome.branch.branch_id();
7298 match outcome.result {
7299 RuntimeBranchResult::MainDraft(draft) => {
7300 main_pending = false;
7301 main_result = Some(Ok(draft));
7302 }
7303 RuntimeBranchResult::Transition(candidate) => {
7304 if let Some(candidate) = candidate {
7305 transition_candidate = Some(candidate);
7306 } else {
7307 self.finalize_optional_branch(
7308 &transition_id,
7309 RuntimeOptimizationKind::ParallelStateTransition,
7310 RuntimeCommitBehavior::TransitionDecision,
7311 "discarded",
7312 false,
7313 );
7314 transition_finalized = true;
7315 }
7316 }
7317 RuntimeBranchResult::Skill(candidate) => {
7318 skill_pending = false;
7319 if let Some(candidate) = candidate {
7320 skill_candidate = Some(candidate);
7321 } else {
7322 self.finalize_optional_branch(
7323 &skill_id,
7324 RuntimeOptimizationKind::SpeculativeSkillRouting,
7325 RuntimeCommitBehavior::SkillSelection,
7326 "discarded",
7327 false,
7328 );
7329 skill_finalized = true;
7330 }
7331 }
7332 RuntimeBranchResult::Reasoning(mode) => {
7333 reasoning_pending = false;
7334 reasoning_decision = Some(mode);
7335 }
7336 RuntimeBranchResult::Failed(error) => {
7337 if branch_id == main_id {
7338 main_pending = false;
7339 main_result = Some(Err(error));
7340 } else if branch_id == transition_id {
7341 self.finalize_optional_branch(
7342 &transition_id,
7343 RuntimeOptimizationKind::ParallelStateTransition,
7344 RuntimeCommitBehavior::TransitionDecision,
7345 "failed",
7346 false,
7347 );
7348 transition_finalized = true;
7349 } else if branch_id == skill_id {
7350 skill_pending = false;
7351 self.finalize_optional_branch(
7352 &skill_id,
7353 RuntimeOptimizationKind::SpeculativeSkillRouting,
7354 RuntimeCommitBehavior::SkillSelection,
7355 "failed",
7356 false,
7357 );
7358 skill_finalized = true;
7359 } else if branch_id == reasoning_id {
7360 reasoning_pending = false;
7361 self.finalize_optional_branch(
7362 &reasoning_id,
7363 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7364 RuntimeCommitBehavior::ReasoningDecision,
7365 "failed",
7366 false,
7367 );
7368 reasoning_finalized = true;
7369 }
7370 }
7371 RuntimeBranchResult::Cancelled => {
7372 self.finalize_optional_branch(
7373 &branch_id,
7374 outcome.branch.optimization,
7375 outcome.branch.commit_behavior,
7376 "cancelled",
7377 false,
7378 );
7379 if branch_id == main_id {
7380 main_pending = false;
7381 main_result =
7382 Some(Err(AgentError::Other("main branch cancelled".to_string())));
7383 } else if branch_id == transition_id {
7384 transition_finalized = true;
7385 transition_fallback_required = true;
7386 } else if branch_id == skill_id {
7387 skill_pending = false;
7388 skill_finalized = true;
7389 skill_fallback_required = true;
7390 } else if branch_id == reasoning_id {
7391 reasoning_pending = false;
7392 reasoning_finalized = true;
7393 reasoning_fallback_required = true;
7394 }
7395 }
7396 }
7397 }
7398 }
7399
7400 fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7401 for branch in branches {
7402 self.finalize_optional_branch(
7403 &branch.branch_id(),
7404 branch.optimization,
7405 branch.commit_behavior,
7406 "cancelled",
7407 false,
7408 );
7409 }
7410 }
7411
7412 fn finalize_branch_loss(
7417 &self,
7418 branch_id: &str,
7419 optimization: RuntimeOptimizationKind,
7420 commit_behavior: RuntimeCommitBehavior,
7421 pending: bool,
7422 completed_failed: Option<bool>,
7423 ) {
7424 let status = if pending {
7425 "cancelled"
7426 } else if completed_failed.unwrap_or(false) {
7427 "failed"
7428 } else {
7429 "discarded"
7430 };
7431 self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7432 }
7433
7434 fn finalize_optional_branch(
7439 &self,
7440 branch_id: &str,
7441 optimization: RuntimeOptimizationKind,
7442 commit_behavior: RuntimeCommitBehavior,
7443 status: &str,
7444 winner: bool,
7445 ) {
7446 crate::optimization::observability::finalize_branch(
7447 self.observability_manager.as_ref(),
7448 branch_id,
7449 status,
7450 winner,
7451 optimization,
7452 commit_behavior,
7453 );
7454 }
7455
7456 fn has_parallel_transition_candidates(&self) -> bool {
7461 self.transitions_available_for_commit()
7462 .map(|(transitions, _)| {
7463 transitions
7464 .iter()
7465 .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7466 })
7467 .unwrap_or(false)
7468 }
7469
7470 async fn select_parallel_transition_candidate(
7475 &self,
7476 processed_input: &str,
7477 ) -> Result<ParallelTransitionSelection> {
7478 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7479 return Ok(ParallelTransitionSelection::NoMatch);
7480 };
7481 let parallel: Vec<Transition> = transitions
7482 .into_iter()
7483 .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7484 .filter(|transition| !transition.requires_response)
7485 .collect();
7486 if parallel.is_empty() {
7487 return Ok(ParallelTransitionSelection::NoMatch);
7488 }
7489 let empty_staged = HashMap::new();
7490 if let Some(candidate) = self.select_deterministic_transition_candidate(
7491 processed_input,
7492 ¤t_state,
7493 ¶llel,
7494 &empty_staged,
7495 ) {
7496 return Ok(ParallelTransitionSelection::Candidate(candidate));
7497 }
7498 let when_transitions: Vec<(usize, &Transition)> = parallel
7499 .iter()
7500 .enumerate()
7501 .filter(|(_, transition)| !transition.when.trim().is_empty())
7502 .collect();
7503 if when_transitions.is_empty() {
7504 return Ok(ParallelTransitionSelection::NoMatch);
7505 }
7506 let llm = self
7507 .llm_registry
7508 .router()
7509 .or_else(|_| self.llm_registry.default())
7510 .map_err(|e| AgentError::Config(e.to_string()))?;
7511 let conditions = when_transitions
7512 .iter()
7513 .enumerate()
7514 .map(|(display_idx, (_, transition))| {
7515 format!("{}. {}", display_idx + 1, transition.when)
7516 })
7517 .collect::<Vec<_>>()
7518 .join("\n");
7519 if !self
7520 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7521 {
7522 return Ok(ParallelTransitionSelection::ReservationExhausted);
7523 }
7524 let context_preview = self.branch_context_preview();
7525 let prompt = format!(
7526 "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-{}).",
7527 current_state,
7528 processed_input,
7529 context_preview,
7530 conditions,
7531 when_transitions.len()
7532 );
7533 let response = self
7534 .observe_purpose(
7535 ObservationPurpose::StateTransitionEvaluation,
7536 llm.complete(&[ChatMessage::user(prompt)], None),
7537 )
7538 .await
7539 .map_err(|e| AgentError::LLM(e.to_string()))?;
7540 let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7541 if choice == 0 || choice > when_transitions.len() {
7542 return Ok(ParallelTransitionSelection::NoMatch);
7543 }
7544 let transition = when_transitions[choice - 1].1.clone();
7545 Ok(ParallelTransitionSelection::Candidate(
7546 TransitionCandidate::new(
7547 current_state,
7548 transition.clone(),
7549 Self::transition_reason(&transition),
7550 ),
7551 ))
7552 }
7553
7554 async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7556 const MAX_REDISPATCH_DEPTH: u32 = 3;
7557 let current_depth = *self.redispatch_depth.read();
7558 if current_depth >= MAX_REDISPATCH_DEPTH {
7559 warn!(depth = current_depth, "Re-dispatch depth limit reached");
7560 let response = AgentResponse::new("");
7561 self.finish_turn_if_root(&response).await?;
7562 return Ok(response);
7563 }
7564 *self.redispatch_depth.write() += 1;
7565 if let Some(context) = self.active_turn_context.write().as_mut() {
7566 context.enter_redispatch();
7567 }
7568 let result = Box::pin(self.run_loop_internal(processed_input)).await;
7569 *self.redispatch_depth.write() -= 1;
7570 if let Some(context) = self.active_turn_context.write().as_mut() {
7571 context.exit_redispatch();
7572 }
7573 let response = result?;
7574 self.finish_turn_if_root(&response).await?;
7575 Ok(response)
7576 }
7577
7578 async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7580 if *self.redispatch_depth.read() == 0 {
7581 self.post_turn_session_lifecycle().await?;
7582 if let Some(context) = self.active_turn_context.write().as_mut() {
7583 context.mark_post_turn_lifecycle_completed();
7584 }
7585 self.hooks.on_response(response).await;
7586 self.end_root_turn();
7587 }
7588 Ok(())
7589 }
7590
7591 async fn execute_state_exit_actions(&self, state_path: &str) {
7593 if let Some(ref sm) = self.state_machine
7594 && let Some(def) = sm.get_definition(state_path)
7595 && !def.on_exit.is_empty()
7596 {
7597 debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7598 self.execute_state_actions(&def.on_exit).await;
7599 }
7600 }
7601
7602 fn state_was_previously_entered(
7604 state_path: &str,
7605 from_state: &str,
7606 history_before: &[StateTransitionEvent],
7607 ) -> bool {
7608 state_path == from_state
7609 || history_before
7610 .iter()
7611 .any(|event| event.from == state_path || event.to == state_path)
7612 }
7613
7614 async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7616 if let Some(ref sm) = self.state_machine
7617 && let Some(def) = sm.get_definition(state_path)
7618 {
7619 if is_reentry && !def.on_reenter.is_empty() {
7620 debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7621 self.execute_state_actions(&def.on_reenter).await;
7622 } else if !def.on_enter.is_empty() {
7623 debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7624 self.execute_state_actions(&def.on_enter).await;
7625 }
7626 }
7627 }
7628
7629 async fn execute_state_actions(&self, actions: &[StateAction]) {
7631 for (action_index, action) in actions.iter().enumerate() {
7632 match action {
7633 StateAction::Tool { tool, args } => {
7634 let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7635 let args_value = self.render_action_args(&raw_args);
7636 let state = self.state_machine.as_ref().map(|sm| sm.current());
7637 let request = ToolExecutionRequest::new(
7638 uuid::Uuid::new_v4().to_string(),
7639 tool.clone(),
7640 args_value,
7641 ToolCallSource::StateAction {
7642 state,
7643 action_index,
7644 },
7645 );
7646 match self.execute_tool_record(request).await {
7647 Ok(record) if record.success => {
7648 debug!(tool = %record.canonical_id, "State action: tool executed");
7649 let _ = self.context_manager.set(
7650 "last_tool_result",
7651 serde_json::Value::String(record.model_output_string()),
7652 );
7653 let _ = self.context_manager.set(
7654 "last_tool_record",
7655 serde_json::to_value(record).unwrap_or(Value::Null),
7656 );
7657 }
7658 Ok(record) => {
7659 warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7660 }
7661 Err(e) => {
7662 warn!(tool = %tool, error = %e, "State action: tool failed")
7663 }
7664 }
7665 }
7666 StateAction::Skill { skill } => {
7667 if let Some(ref executor) = self.skill_executor {
7668 if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7669 match executor
7670 .execute_with_invoker(def, "", serde_json::json!({}), self)
7671 .await
7672 {
7673 Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7674 Err(e) => {
7675 warn!(skill = %skill, error = %e, "State action: skill failed")
7676 }
7677 }
7678 } else {
7679 warn!(skill = %skill, "State action: skill not found");
7680 }
7681 }
7682 }
7683 StateAction::SetContext { set_context } => {
7684 for (key, value) in set_context {
7685 if let Err(e) = self.context_manager.set(key, value.clone()) {
7686 warn!(key = %key, error = %e, "State action: set_context failed");
7687 } else {
7688 debug!(key = %key, "State action: context set");
7689 }
7690 }
7691 }
7692 StateAction::Prompt {
7693 prompt,
7694 llm,
7695 store_as,
7696 } => {
7697 let llm_result = if let Some(alias) = llm {
7698 self.llm_registry.get(alias)
7699 } else {
7700 self.llm_registry.default()
7701 };
7702 match llm_result {
7703 Ok(llm_provider) => {
7704 let context = self.build_context_with_overlays();
7706 let rendered_prompt = self
7707 .template_renderer
7708 .render(prompt, &context)
7709 .unwrap_or_else(|_| prompt.clone());
7710 let recent =
7711 self.memory.get_messages(Some(5)).await.unwrap_or_default();
7712 let mut messages: Vec<ChatMessage> = recent;
7713 messages.push(ChatMessage::user(&rendered_prompt));
7714 match self
7715 .observe_purpose(
7716 ObservationPurpose::StateAction,
7717 llm_provider.complete(&messages, None),
7718 )
7719 .await
7720 {
7721 Ok(response) => {
7722 if let Some(key) = store_as {
7723 let _ = self
7724 .context_manager
7725 .set(key, Value::String(response.content));
7726 debug!(key = %key, "State action: prompt result stored");
7727 }
7728 }
7729 Err(e) => {
7730 warn!(error = %e, "State action: prompt LLM call failed");
7731 }
7732 }
7733 }
7734 Err(e) => {
7735 warn!(error = %e, "State action: LLM not found for prompt");
7736 }
7737 }
7738 }
7739 }
7740 }
7741 }
7742
7743 async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7744 let extractors = match &self.state_machine {
7745 Some(sm) => match sm.current_definition() {
7746 Some(def) if !def.extract.is_empty() => def.extract.clone(),
7747 _ => return HashMap::new(),
7748 },
7749 None => return HashMap::new(),
7750 };
7751
7752 let mut staged = HashMap::new();
7753 for extractor in &extractors {
7754 let prompt = if let Some(ref custom) = extractor.llm_extract {
7755 format!(
7756 "User message:\n\"{}\"\n\nInstruction:\n{}",
7757 user_message, custom
7758 )
7759 } else if let Some(ref desc) = extractor.description {
7760 format!(
7761 "From the following message, extract: {}\n\n\
7762 Message: \"{}\"\n\n\
7763 If the information is present, return ONLY the extracted value.\n\
7764 If NOT present, return exactly: __NONE__",
7765 desc, user_message
7766 )
7767 } else {
7768 continue;
7769 };
7770
7771 let llm = match self
7772 .llm_registry
7773 .get(&extractor.llm)
7774 .or_else(|_| self.llm_registry.get("router"))
7775 .or_else(|_| self.llm_registry.get("default"))
7776 {
7777 Ok(llm) => llm,
7778 Err(e) => {
7779 warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
7780 continue;
7781 }
7782 };
7783
7784 let messages = vec![ChatMessage::user(&prompt)];
7785 match self
7786 .observe_purpose(
7787 ObservationPurpose::ContextExtraction,
7788 llm.complete(&messages, None),
7789 )
7790 .await
7791 {
7792 Ok(response) => {
7793 let value = response.content.trim().to_string();
7794 if value != "__NONE__" && !value.is_empty() {
7795 staged.insert(
7796 extractor.key.clone(),
7797 serde_json::Value::String(value.clone()),
7798 );
7799 debug!(key = %extractor.key, value = %value, "Context extracted");
7800 } else if extractor.required {
7801 warn!(key = %extractor.key, "Required extraction returned no value");
7802 }
7803 }
7804 Err(e) => {
7805 warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
7806 }
7807 }
7808 }
7809 staged
7810 }
7811
7812 fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
7813 for (key, value) in staged {
7814 if let Err(error) = self.context_manager.update(key, value.clone()) {
7815 warn!(key = %key, error = %error, "staged context write failed");
7816 }
7817 }
7818 }
7819
7820 async fn run_context_extractors(&self, user_message: &str) {
7822 let staged = self.run_context_extractors_staged(user_message).await;
7823 self.commit_staged_context_writes(&staged);
7824 }
7825
7826 async fn check_memory_compression(&self) -> Result<()> {
7827 if self.memory.needs_compression() {
7828 let result = self.memory.compress(None).await?;
7829 if let CompressResult::Compressed {
7830 messages_summarized,
7831 new_summary_length,
7832 tokens_saved,
7833 } = result
7834 {
7835 let event = MemoryCompressEvent::new(
7836 messages_summarized,
7837 tokens_saved,
7838 new_summary_length as u32,
7839 );
7840 self.hooks.on_memory_compress(&event).await;
7841 debug!(
7842 messages = messages_summarized,
7843 tokens_saved = tokens_saved,
7844 "Memory compressed"
7845 );
7846 }
7847 }
7848
7849 self.handle_memory_overflow().await?;
7851 self.check_memory_budget().await;
7852
7853 Ok(())
7854 }
7855
7856 async fn check_memory_budget(&self) {
7857 let Some(ref budget) = self.memory_token_budget else {
7858 return;
7859 };
7860
7861 let context = match self.memory.get_context().await {
7862 Ok(ctx) => ctx,
7863 Err(_) => return,
7864 };
7865
7866 let used_tokens = context.estimated_tokens();
7868 if budget.is_over_warn_threshold(used_tokens) {
7869 let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
7870 self.hooks.on_memory_budget_warning(&event).await;
7871 debug!(
7872 used = used_tokens,
7873 total = budget.total,
7874 percent = event.usage_percent,
7875 "Memory budget warning"
7876 );
7877 }
7878
7879 if let Some(ref summary) = context.summary {
7881 let summary_tokens = ai_agents_memory::estimate_tokens(summary);
7882 let summary_budget = budget.allocation.summary;
7883 if summary_budget > 0 {
7884 let warn_threshold =
7885 (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7886 if summary_tokens >= warn_threshold {
7887 let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
7888 self.hooks.on_memory_budget_warning(&event).await;
7889 }
7890 }
7891 }
7892
7893 let recent_tokens: u32 = context
7895 .messages
7896 .iter()
7897 .map(ai_agents_memory::estimate_message_tokens)
7898 .sum();
7899 let recent_budget = budget.allocation.recent_messages;
7900 if recent_budget > 0 {
7901 let warn_threshold =
7902 (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7903 if recent_tokens >= warn_threshold {
7904 let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
7905 self.hooks.on_memory_budget_warning(&event).await;
7906 }
7907 }
7908
7909 let relationship_budget = budget.allocation.relationships;
7910 if relationship_budget > 0 {
7911 let relationship_tokens = self
7912 .relationship_memory_text()
7913 .map(|text| ai_agents_memory::estimate_tokens(&text))
7914 .unwrap_or(0);
7915 let warn_threshold =
7916 (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7917 if relationship_tokens >= warn_threshold {
7918 let event = MemoryBudgetEvent::new(
7919 "relationships",
7920 relationship_tokens,
7921 relationship_budget,
7922 );
7923 self.hooks.on_memory_budget_warning(&event).await;
7924 }
7925 }
7926 }
7927
7928 async fn handle_memory_overflow(&self) -> Result<()> {
7929 let Some(ref budget) = self.memory_token_budget else {
7930 return Ok(());
7931 };
7932
7933 let context = self.memory.get_context().await?;
7934 let used_tokens = context.estimated_tokens();
7935
7936 if used_tokens <= budget.total {
7937 return Ok(());
7938 }
7939
7940 match budget.overflow_strategy {
7941 OverflowStrategy::TruncateOldest => {
7942 let tokens_to_free = used_tokens - budget.total;
7943 let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
7944 if messages_to_evict > 0 {
7945 self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
7946 .await?;
7947 }
7948 }
7949 OverflowStrategy::SummarizeMore => {
7950 let max_attempts = context.total_messages.max(1);
7951 for _ in 0..max_attempts {
7952 match self.memory.compress(None).await? {
7953 CompressResult::Compressed {
7954 messages_summarized,
7955 ..
7956 } if messages_summarized > 0 => {
7957 let context = self.memory.get_context().await?;
7958 if context.estimated_tokens() <= budget.total {
7959 return Ok(());
7960 }
7961 }
7962 _ => break,
7963 }
7964 }
7965 let context = self.memory.get_context().await?;
7966 let used_tokens = context.estimated_tokens();
7967 if used_tokens > budget.total {
7968 return Err(AgentError::MemoryBudgetExceeded {
7969 used: used_tokens,
7970 budget: budget.total,
7971 });
7972 }
7973 }
7974 OverflowStrategy::Error => {
7975 return Err(AgentError::MemoryBudgetExceeded {
7976 used: used_tokens,
7977 budget: budget.total,
7978 });
7979 }
7980 }
7981 Ok(())
7982 }
7983
7984 fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
7985 ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
7987 }
7988
7989 async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
7990 let evicted = self.memory.evict_oldest(count).await?;
7991 if !evicted.is_empty() {
7992 let event = MemoryEvictEvent {
7993 reason,
7994 messages_evicted: evicted.len(),
7995 importance_scores: vec![],
7996 };
7997 self.hooks.on_memory_evict(&event).await;
7998 debug!(count = evicted.len(), "Messages evicted from memory");
7999 }
8000 Ok(())
8001 }
8002
8003 #[instrument(skip(self, input), fields(agent = %self.info.name))]
8004 async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
8005 match self.determine_reasoning_mode_strict(input).await {
8006 Ok(mode) => Ok(mode),
8007 Err(_) => Ok(ReasoningMode::None),
8008 }
8009 }
8010
8011 async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
8012 let effective_config = self.get_effective_reasoning_config();
8013
8014 if !matches!(effective_config.mode, ReasoningMode::Auto) {
8015 return Ok(effective_config.mode.clone());
8016 }
8017
8018 let judge_llm = effective_config
8019 .judge_llm
8020 .as_ref()
8021 .and_then(|alias| self.llm_registry.get(alias).ok())
8022 .or_else(|| self.llm_registry.router().ok())
8023 .or_else(|| self.llm_registry.default().ok());
8024
8025 let Some(llm) = judge_llm else {
8026 return Ok(ReasoningMode::None);
8027 };
8028
8029 let prompt = format!(
8030 r#"Analyze this user request and determine the appropriate reasoning mode.
8031
8032User request: "{}"
8033
8034Choose ONE of these modes:
8035- none: Simple queries, greetings, direct answers (fastest)
8036- cot: Complex analysis, multi-step reasoning, math problems
8037- react: Tasks requiring multiple tool calls with observation
8038- plan_and_execute: Complex multi-step tasks requiring coordination
8039
8040Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
8041 input
8042 );
8043
8044 let messages = vec![ChatMessage::user(&prompt)];
8045 let response = self
8046 .observe_purpose(
8047 ObservationPurpose::ReflectionDecision,
8048 llm.complete(&messages, None),
8049 )
8050 .await
8051 .map_err(|e| AgentError::LLM(e.to_string()))?;
8052
8053 let mode_str = response.content.trim().to_lowercase();
8054 Ok(match mode_str.as_str() {
8055 "cot" => ReasoningMode::CoT,
8056 "react" => ReasoningMode::React,
8057 "plan_and_execute" => ReasoningMode::PlanAndExecute,
8058 _ => ReasoningMode::None,
8059 })
8060 }
8061
8062 async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
8063 let effective_config = self.get_effective_reflection_config();
8064
8065 if !effective_config.requires_evaluation() {
8066 return Ok(false);
8067 }
8068
8069 if effective_config.is_enabled() {
8070 return Ok(true);
8071 }
8072
8073 let evaluator_llm = effective_config
8074 .evaluator_llm
8075 .as_ref()
8076 .and_then(|alias| self.llm_registry.get(alias).ok())
8077 .or_else(|| self.llm_registry.router().ok())
8078 .or_else(|| self.llm_registry.default().ok());
8079
8080 let Some(llm) = evaluator_llm else {
8081 return Ok(false);
8082 };
8083
8084 let response_preview: String = response.chars().take(500).collect();
8085 let prompt = format!(
8086 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8087
8088User query: "{}"
8089Response: "{}"
8090
8091Answer YES or NO only."#,
8092 input, response_preview
8093 );
8094
8095 let messages = vec![ChatMessage::user(&prompt)];
8096 let result = self
8097 .observe_purpose(
8098 ObservationPurpose::ReflectionDecision,
8099 llm.complete(&messages, None),
8100 )
8101 .await;
8102
8103 match result {
8104 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8105 Err(_) => Ok(false),
8106 }
8107 }
8108
8109 fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8110 format!(
8111 "{}\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>",
8112 base_prompt
8113 )
8114 }
8115
8116 fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8117 format!(
8118 "{}\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>",
8119 base_prompt
8120 )
8121 }
8122
8123 async fn generate_plan(&self, input: &str) -> Result<Plan> {
8124 let effective = self.get_effective_reasoning_config();
8125 let planning_config = effective.get_planning();
8126
8127 let planner_llm = planning_config
8128 .and_then(|c| c.planner_llm.as_ref())
8129 .and_then(|alias| self.llm_registry.get(alias).ok())
8130 .or_else(|| self.llm_registry.router().ok())
8131 .or_else(|| self.llm_registry.default().ok())
8132 .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8133
8134 let mut available_tool_ids: Vec<String> = self
8135 .get_available_tool_ids()
8136 .await
8137 .unwrap_or_else(|_| self.tools.list_ids());
8138 let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8139
8140 if let Some(config) = planning_config {
8142 if !config.available.tools.is_all() {
8143 available_tool_ids.retain(|t| config.available.tools.allows(t));
8144 }
8145 if !config.available.skills.is_all() {
8146 available_skills.retain(|s| config.available.skills.allows(s));
8147 }
8148 }
8149
8150 let tool_descriptions: Vec<String> = available_tool_ids
8153 .iter()
8154 .filter_map(|id| {
8155 self.tools.get(id).map(|tool| {
8156 let schema = tool.input_schema();
8157 let args_desc = schema
8158 .get("properties")
8159 .and_then(|p| serde_json::to_string(p).ok())
8160 .unwrap_or_else(|| "{}".to_string());
8161 format!(
8162 "- {} ({}): {}\n Arguments: {}",
8163 id,
8164 tool.name(),
8165 tool.description(),
8166 args_desc
8167 )
8168 })
8169 })
8170 .collect();
8171
8172 let tools_section = if tool_descriptions.is_empty() {
8173 "Available tools: none".to_string()
8174 } else {
8175 format!("Available tools:\n{}", tool_descriptions.join("\n"))
8176 };
8177
8178 let skills_section = if available_skills.is_empty() {
8179 "Available skills: none".to_string()
8180 } else {
8181 format!("Available skills: {}", available_skills.join(", "))
8182 };
8183
8184 let prompt = format!(
8185 r#"Create a step-by-step plan to accomplish this goal.
8186
8187Goal: "{}"
8188
8189{}
8190
8191{}
8192
8193Create a plan with clear steps. For each step, specify:
8194- description: What this step accomplishes
8195- action_type: "tool", "skill", "think", or "respond"
8196- action_target: The tool/skill id (if applicable)
8197- args: The arguments object matching the tool's schema (if action_type is "tool")
8198- dependencies: List of step IDs this depends on (empty if none)
8199
8200Respond in JSON format:
8201{{
8202 "steps": [
8203 {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8204 {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8205 ]
8206}}"#,
8207 input, tools_section, skills_section,
8208 );
8209
8210 let messages = vec![ChatMessage::user(&prompt)];
8211 let response = self
8212 .observe_purpose(
8213 ObservationPurpose::PlanGeneration,
8214 planner_llm.complete(&messages, None),
8215 )
8216 .await
8217 .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8218
8219 let mut plan = Plan::new(input);
8220
8221 if let Some(json_start) = response.content.find('{')
8222 && let Some(json_end) = response.content.rfind('}')
8223 {
8224 let json_str = &response.content[json_start..=json_end];
8225 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8226 && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8227 {
8228 for step_value in steps {
8229 let id = step_value
8230 .get("id")
8231 .and_then(|v| v.as_str())
8232 .unwrap_or("step");
8233 let desc = step_value
8234 .get("description")
8235 .and_then(|v| v.as_str())
8236 .unwrap_or("");
8237 let action_type = step_value
8238 .get("action_type")
8239 .and_then(|v| v.as_str())
8240 .unwrap_or("think");
8241 let action_target = step_value
8242 .get("action_target")
8243 .and_then(|v| v.as_str())
8244 .unwrap_or("");
8245 let args = step_value
8246 .get("args")
8247 .cloned()
8248 .unwrap_or(serde_json::json!({}));
8249 let deps: Vec<String> = step_value
8250 .get("dependencies")
8251 .and_then(|v| v.as_array())
8252 .map(|arr| {
8253 arr.iter()
8254 .filter_map(|v| v.as_str().map(String::from))
8255 .collect()
8256 })
8257 .unwrap_or_default();
8258
8259 let action = match action_type {
8260 "tool" => PlanAction::tool(action_target, args),
8261 "skill" => PlanAction::skill(action_target),
8262 "respond" => PlanAction::respond(action_target),
8263 _ => PlanAction::think(desc),
8264 };
8265
8266 let step = PlanStep::new(desc, action)
8267 .with_id(id)
8268 .with_dependencies(deps);
8269 plan.add_step(step);
8270 }
8271 }
8272 }
8273
8274 if plan.steps.is_empty() {
8275 plan.add_step(PlanStep::new(
8276 "Process the request",
8277 PlanAction::think(input),
8278 ));
8279 plan.add_step(PlanStep::new(
8280 "Provide response",
8281 PlanAction::respond("Answer based on analysis"),
8282 ));
8283 }
8284
8285 Ok(plan)
8286 }
8287
8288 async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8289 let llm = self.get_state_llm()?;
8290 let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8291 let effective = self.get_effective_reasoning_config();
8292 let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8293
8294 plan.status = PlanStatus::InProgress;
8295
8296 for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8297 let step = &plan.steps[step_idx];
8298
8299 let deps_satisfied = step.dependencies.iter().all(|dep| {
8300 plan.steps
8301 .iter()
8302 .find(|s| &s.id == dep)
8303 .map(|s| s.status.is_completed())
8304 .unwrap_or(false)
8305 });
8306
8307 if !deps_satisfied {
8308 continue;
8309 }
8310
8311 plan.steps[step_idx].mark_running();
8312
8313 let result = match &plan.steps[step_idx].action {
8314 PlanAction::Tool { tool, args } => {
8315 let has_dep_results = plan.steps[step_idx]
8321 .dependencies
8322 .iter()
8323 .any(|dep| results.contains_key(dep));
8324
8325 let final_args = if has_dep_results {
8326 let dep_context: String = plan.steps[step_idx]
8327 .dependencies
8328 .iter()
8329 .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8330 .collect::<Vec<_>>()
8331 .join("\n");
8332
8333 let tool_schema = self
8334 .tools
8335 .get(tool)
8336 .map(|t| {
8337 let schema = t.input_schema();
8338 let props = schema
8339 .get("properties")
8340 .and_then(|p| serde_json::to_string(p).ok())
8341 .unwrap_or_else(|| "{}".to_string());
8342 format!(
8343 "{}: {}\nArguments schema: {}",
8344 t.id(),
8345 t.description(),
8346 props
8347 )
8348 })
8349 .unwrap_or_default();
8350
8351 let step_desc = &plan.steps[step_idx].description;
8352 let arg_prompt = format!(
8353 "Generate the JSON arguments for a tool call.\n\n\
8354 Tool: {}\n\n\
8355 Task: {}\n\n\
8356 Previous step results:\n{}\n\n\
8357 Planner's draft arguments: {}\n\n\
8358 Produce ONLY a valid JSON object with the correct argument values.\n\
8359 Use actual values from the previous step results, not template references.",
8360 tool_schema,
8361 step_desc,
8362 dep_context,
8363 serde_json::to_string(args).unwrap_or_default()
8364 );
8365 let messages = vec![ChatMessage::user(&arg_prompt)];
8366 match self
8367 .observe_purpose(
8368 ObservationPurpose::PlanStep,
8369 llm.complete(&messages, None),
8370 )
8371 .await
8372 {
8373 Ok(resp) => {
8374 let content = resp.content.trim();
8375 let json_start = content.find('{');
8377 let json_end = content.rfind('}');
8378 if let (Some(start), Some(end)) = (json_start, json_end) {
8379 serde_json::from_str(&content[start..=end])
8380 .unwrap_or_else(|_| args.clone())
8381 } else {
8382 args.clone()
8383 }
8384 }
8385 Err(_) => args.clone(),
8386 }
8387 } else {
8388 args.clone()
8389 };
8390
8391 let request = ToolExecutionRequest::new(
8392 uuid::Uuid::new_v4().to_string(),
8393 tool.clone(),
8394 final_args,
8395 ToolCallSource::Plan {
8396 step_index: step_idx,
8397 },
8398 );
8399 match self.execute_tool_record(request).await {
8400 Ok(record) if record.success => {
8401 serde_json::json!({ "output": record.model_output_string() })
8402 }
8403 Ok(record) => {
8404 plan.steps[step_idx].mark_failed(record.model_output_string());
8405 continue;
8406 }
8407 Err(e) => {
8408 plan.steps[step_idx].mark_failed(e.to_string());
8409 continue;
8410 }
8411 }
8412 }
8413 PlanAction::Skill { skill } => {
8414 if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8415 if let Some(ref executor) = self.skill_executor {
8416 match executor
8417 .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8418 .await
8419 {
8420 Ok(output) => serde_json::json!({ "output": output }),
8421 Err(e) => {
8422 plan.steps[step_idx].mark_failed(e.to_string());
8423 continue;
8424 }
8425 }
8426 } else {
8427 serde_json::json!({ "output": "Skill executor not available" })
8428 }
8429 } else {
8430 plan.steps[step_idx].mark_failed("Skill not found");
8431 continue;
8432 }
8433 }
8434 PlanAction::Think { prompt } => {
8435 let context: String = results
8436 .iter()
8437 .map(|(k, v)| format!("{}: {}", k, v))
8438 .collect::<Vec<_>>()
8439 .join("\n");
8440
8441 let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8442 let messages = vec![ChatMessage::user(&think_prompt)];
8443
8444 match self
8445 .observe_purpose(
8446 ObservationPurpose::PlanStep,
8447 llm.complete(&messages, None),
8448 )
8449 .await
8450 {
8451 Ok(resp) => serde_json::json!({ "output": resp.content }),
8452 Err(e) => {
8453 plan.steps[step_idx].mark_failed(e.to_string());
8454 continue;
8455 }
8456 }
8457 }
8458 PlanAction::Respond { template } => {
8459 let context: String = results
8460 .iter()
8461 .map(|(k, v)| format!("{}: {}", k, v))
8462 .collect::<Vec<_>>()
8463 .join("\n");
8464
8465 let respond_prompt = format!(
8466 "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8467 context, template
8468 );
8469 let messages = vec![ChatMessage::user(&respond_prompt)];
8470
8471 match self
8472 .observe_purpose(
8473 ObservationPurpose::PlanStep,
8474 llm.complete(&messages, None),
8475 )
8476 .await
8477 {
8478 Ok(resp) => serde_json::json!({ "output": resp.content }),
8479 Err(e) => {
8480 plan.steps[step_idx].mark_failed(e.to_string());
8481 continue;
8482 }
8483 }
8484 }
8485 };
8486
8487 results.insert(plan.steps[step_idx].id.clone(), result.clone());
8488 plan.steps[step_idx].mark_completed(Some(result));
8489 }
8490
8491 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8493 if has_failures {
8494 let failed_ids: Vec<String> = plan
8495 .steps
8496 .iter()
8497 .filter(|s| s.status.is_failed())
8498 .map(|s| s.id.clone())
8499 .collect();
8500 plan.status = PlanStatus::Failed {
8501 error: format!("Steps failed: {}", failed_ids.join(", ")),
8502 };
8503 } else {
8504 plan.status = PlanStatus::Completed;
8505 }
8506
8507 let all_outputs: Vec<String> = plan
8509 .steps
8510 .iter()
8511 .filter(|s| s.status.is_completed())
8512 .filter_map(|s| {
8513 s.result
8514 .as_ref()
8515 .and_then(|r| r.get("output"))
8516 .and_then(|o| o.as_str())
8517 .map(|o| format!("{}: {}", s.description, o))
8518 })
8519 .collect();
8520
8521 if all_outputs.is_empty() {
8522 return Ok("Plan execution completed but produced no results.".to_string());
8523 }
8524
8525 if all_outputs.len() == 1 {
8526 return Ok(all_outputs.into_iter().next().unwrap());
8527 }
8528
8529 let context = all_outputs.join("\n\n");
8531 let prompt = format!(
8532 "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8533 plan.goal, context
8534 );
8535 let messages = vec![ChatMessage::user(&prompt)];
8536 match self
8537 .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8538 .await
8539 {
8540 Ok(resp) => Ok(resp.content.trim().to_string()),
8541 Err(_) => Ok(context),
8542 }
8543 }
8544
8545 async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8546 let effective_config = self.get_effective_reflection_config();
8547 self.evaluate_response_with_config(input, response, &effective_config)
8548 .await
8549 }
8550
8551 fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8552 if let Some(start) = content.find("<thinking>")
8553 && let Some(end) = content.find("</thinking>")
8554 {
8555 let thinking = content[start + 10..end].trim().to_string();
8556 let answer = content[end + 11..].trim().to_string();
8557 return (Some(thinking), answer);
8558 }
8559 (None, content.to_string())
8560 }
8561
8562 fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8563 match self.get_effective_reasoning_config().output {
8564 ReasoningOutput::Hidden => answer.to_string(),
8565 ReasoningOutput::Visible => {
8566 if let Some(t) = thinking {
8567 format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8568 } else {
8569 answer.to_string()
8570 }
8571 }
8572 ReasoningOutput::Tagged => {
8573 if let Some(t) = thinking {
8574 format!("<thinking>{}</thinking>\n{}", t, answer)
8575 } else {
8576 answer.to_string()
8577 }
8578 }
8579 }
8580 }
8581
8582 async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8587 self.init_storage().await?;
8591 self.begin_root_turn();
8592 let _root_cleanup = RootTurnCleanup::new(self);
8593 info!(input_len = input.len(), "Starting chat");
8594
8595 self.hooks.on_message_received(input).await;
8596
8597 if !self.context_initialized.swap(true, Ordering::SeqCst) {
8601 self.context_manager.initialize().await?;
8602 debug!("Context manager initialized (defaults, env, builtins)");
8603 }
8604
8605 self.check_turn_timeout().await?;
8606 self.context_manager.refresh_per_turn().await?;
8607
8608 self.clear_disambiguation_context();
8611
8612 if let Some(ref disambiguator) = self.disambiguation_manager {
8614 let disambiguation_context = self.build_disambiguation_context().await?;
8615
8616 let state_override = self
8618 .state_machine
8619 .as_ref()
8620 .and_then(|sm| sm.current_definition())
8621 .and_then(|def| def.disambiguation.clone());
8622
8623 let state_generation = self
8624 .state_machine
8625 .as_ref()
8626 .map(|state_machine| state_machine.generation());
8627 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8628 let mut disambiguation_result = self
8629 .observe_purpose(
8630 ObservationPurpose::DisambiguationDetection,
8631 disambiguator.process_input_with_override(
8632 input,
8633 &disambiguation_context,
8634 state_override.as_ref(),
8635 None,
8636 ),
8637 )
8638 .await?;
8639 let current_state_generation = self
8640 .state_machine
8641 .as_ref()
8642 .map(|state_machine| state_machine.generation());
8643 if current_state_generation != state_generation
8644 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8645 {
8646 disambiguator.clear_pending().await;
8647 *self.pending_skill_id.write() = None;
8648 disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8649 info!(
8650 confirmation_event = "invalidated",
8651 invalidation_reason = "state_generation_changed",
8652 "Disambiguation result invalidated before redispatch"
8653 );
8654 }
8655 match disambiguation_result {
8656 DisambiguationResult::Clear => {
8657 debug!("Input is clear, proceeding normally");
8658 }
8659 DisambiguationResult::NeedsClarification {
8660 question,
8661 detection,
8662 } => {
8663 let admission = self
8664 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8665 .await?;
8666 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8667 info!(
8668 ambiguity_type = ?detection.ambiguity_type,
8669 confidence = detection.confidence,
8670 "Input requires clarification"
8671 );
8672
8673 self.commit_root_user_message(input).await?;
8676 self.memory
8677 .add_message(ChatMessage::assistant(&question.question))
8678 .await?;
8679
8680 let status = if awaiting_confirmation {
8681 "awaiting_confirmation"
8682 } else {
8683 "awaiting_clarification"
8684 };
8685 let response = AgentResponse::new(&question.question).with_metadata(
8686 "disambiguation",
8687 serde_json::json!({
8688 "status": status,
8689 "options": question.options,
8690 "clarifying": question.clarifying,
8691 "detection": {
8692 "type": detection.ambiguity_type,
8693 "confidence": detection.confidence,
8694 "what_is_unclear": detection.what_is_unclear,
8695 }
8696 }),
8697 );
8698 drop(admission);
8699 self.finish_turn_if_root(&response).await?;
8700 return Ok(response);
8701 }
8702 DisambiguationResult::Clarified {
8703 enriched_input,
8704 resolved,
8705 ..
8706 } => {
8707 let admission = match self
8708 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8709 .await
8710 {
8711 Ok(admission) => admission,
8712 Err(error) => {
8713 *self.pending_skill_id.write() = None;
8714 return Err(error);
8715 }
8716 };
8717 info!(
8718 resolved_count = resolved.len(),
8719 enriched = %enriched_input,
8720 "Input clarified, injecting resolved intent into context"
8721 );
8722
8723 for (key, value) in &resolved {
8726 let context_key = format!("disambiguation.{}", key);
8727 let _ = self.context_manager.set(&context_key, value.clone());
8728 }
8729
8730 if let Some(intent) = resolved.get("intent") {
8731 let _ = self.context_manager.set("resolved_intent", intent.clone());
8732 }
8733
8734 let _ = self
8735 .context_manager
8736 .set("disambiguation.resolved", serde_json::Value::Bool(true));
8737
8738 let skill_id = self.pending_skill_id.read().clone();
8742 if let Some(skill_id) = skill_id {
8743 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8744 drop(admission);
8745 return self
8746 .recheck_skill_disambiguation(
8747 &skill_id,
8748 &enriched_input,
8749 disambiguation_epoch,
8750 state_generation,
8751 )
8752 .await;
8753 }
8754
8755 drop(admission);
8756 return self.run_loop_internal(&enriched_input).await;
8757 }
8758 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
8759 info!("Proceeding with best guess interpretation");
8760
8761 let skill_id = self.pending_skill_id.read().clone();
8763 if let Some(skill_id) = skill_id {
8764 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
8765 return self
8766 .recheck_skill_disambiguation(
8767 &skill_id,
8768 &enriched_input,
8769 disambiguation_epoch,
8770 state_generation,
8771 )
8772 .await;
8773 }
8774
8775 return self.run_loop_internal(&enriched_input).await;
8776 }
8777 DisambiguationResult::GiveUp { reason } => {
8778 *self.pending_skill_id.write() = None;
8779 warn!(reason = %reason, "Disambiguation gave up");
8780 let apology = self
8781 .generate_localized_apology(
8782 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
8783 &reason,
8784 )
8785 .await
8786 .unwrap_or_else(|_| {
8787 format!("I'm sorry, I couldn't understand your request: {}", reason)
8788 });
8789 let response = AgentResponse::new(&apology);
8790 self.finish_turn_if_root(&response).await?;
8791 return Ok(response);
8792 }
8793 DisambiguationResult::Escalate { reason } => {
8794 *self.pending_skill_id.write() = None;
8795 info!(reason = %reason, "Escalating to human");
8796 if let Some(ref hitl) = self.hitl_engine {
8797 let trigger =
8798 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
8799 let mut context_map = HashMap::new();
8800 context_map.insert("original_input".to_string(), serde_json::json!(input));
8801 context_map.insert("reason".to_string(), serde_json::json!(&reason));
8802 let check_result = HITLCheckResult::required(
8803 trigger,
8804 context_map,
8805 format!("User request needs human assistance: {}", reason),
8806 Some(hitl.config().default_timeout_seconds),
8807 );
8808 let result = self.request_hitl_approval(check_result).await?;
8809 if matches!(
8810 result,
8811 ApprovalResult::Approved | ApprovalResult::Modified { .. }
8812 ) {
8813 return self.run_loop_internal(input).await;
8814 }
8815 }
8816 let apology = self
8817 .generate_localized_apology(
8818 "Explain briefly that you're transferring the user to a human agent for help.",
8819 &reason,
8820 )
8821 .await
8822 .unwrap_or_else(|_| {
8823 format!("I need human assistance to help with your request: {}", reason)
8824 });
8825 let response = AgentResponse::new(&apology);
8826 self.finish_turn_if_root(&response).await?;
8827 return Ok(response);
8828 }
8829 DisambiguationResult::Abandoned { new_input } => {
8830 *self.pending_skill_id.write() = None;
8831
8832 info!(
8833 has_new_input = new_input.is_some(),
8834 "Clarification abandoned by user"
8835 );
8836
8837 self.commit_root_user_message(input).await?;
8838
8839 match new_input {
8840 Some(fresh_input) => {
8841 return self.run_loop_internal(&fresh_input).await;
8844 }
8845 None => {
8846 let ack = self
8848 .generate_localized_apology(
8849 "The user changed their mind about their previous request. \
8850 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
8851 Do NOT apologize excessively. Be concise.",
8852 "User abandoned clarification",
8853 )
8854 .await
8855 .unwrap_or_else(|_| {
8856 "OK, no problem. What else can I help with?".to_string()
8857 });
8858
8859 self.memory
8860 .add_message(ChatMessage::assistant(&ack))
8861 .await?;
8862
8863 let response = AgentResponse::new(&ack);
8864 self.finish_turn_if_root(&response).await?;
8865 return Ok(response);
8866 }
8867 }
8868 }
8869 }
8870 }
8871
8872 self.run_loop_internal(input).await
8873 }
8874
8875 async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
8877 let llm = self.llm_registry.router().map_err(|e| {
8878 AgentError::LLM(format!(
8879 "Router LLM not available for localized response: {}",
8880 e
8881 ))
8882 })?;
8883
8884 let recent: Vec<String> = self
8885 .memory
8886 .get_messages(Some(3))
8887 .await?
8888 .iter()
8889 .map(|m| m.content.clone())
8890 .collect();
8891
8892 let context_hint = if recent.is_empty() {
8893 String::new()
8894 } else {
8895 format!(
8896 "\nRecent conversation (detect the user's language from this):\n{}\n",
8897 recent.join("\n")
8898 )
8899 };
8900
8901 let prompt = format!(
8902 "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
8903 instruction, reason, context_hint
8904 );
8905
8906 let messages = vec![ChatMessage::user(&prompt)];
8907 let response = self
8908 .observe_purpose(
8909 ObservationPurpose::DisambiguationClarification,
8910 llm.complete(&messages, None),
8911 )
8912 .await
8913 .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
8914
8915 Ok(response.content.trim().to_string())
8916 }
8917
8918 fn render_action_args(&self, args: &Value) -> Value {
8922 let context = self.build_context_with_overlays();
8923 match args {
8924 Value::Object(map) => {
8925 let mut rendered = serde_json::Map::new();
8926 for (k, v) in map {
8927 match v {
8928 Value::String(s) if s.contains("{{") => {
8929 match self.template_renderer.render(s, &context) {
8930 Ok(rendered_str) => {
8931 rendered.insert(k.clone(), Value::String(rendered_str));
8932 }
8933 Err(_) => {
8934 rendered.insert(k.clone(), v.clone());
8935 }
8936 }
8937 }
8938 _ => {
8939 rendered.insert(k.clone(), v.clone());
8940 }
8941 }
8942 }
8943 Value::Object(rendered)
8944 }
8945 _ => args.clone(),
8946 }
8947 }
8948
8949 fn clear_disambiguation_context(&self) {
8951 let _ = self
8952 .context_manager
8953 .set("resolved_intent", serde_json::Value::Null);
8954
8955 let all = self.context_manager.get_all();
8956 for key in all.keys() {
8957 if key.starts_with("disambiguation.") {
8958 let _ = self.context_manager.set(key, serde_json::Value::Null);
8959 }
8960 }
8961 }
8962
8963 async fn recheck_skill_disambiguation(
8969 &self,
8970 skill_id: &str,
8971 enriched_input: &str,
8972 expected_disambiguation_epoch: u64,
8973 expected_state_generation: Option<u64>,
8974 ) -> Result<AgentResponse> {
8975 let skill = self
8976 .skill_router
8977 .as_ref()
8978 .and_then(|r| r.get_skill(skill_id).cloned());
8979
8980 if let Some(ref skill) = skill
8982 && let Some(ref skill_disambig) = skill.disambiguation
8983 && skill_disambig.enabled.unwrap_or(false)
8984 && let Some(ref disambiguator) = self.disambiguation_manager
8985 {
8986 let context = self.build_disambiguation_context().await?;
8987 let state_override = self
8988 .state_machine
8989 .as_ref()
8990 .and_then(|sm| sm.current_definition())
8991 .and_then(|def| def.disambiguation.clone());
8992
8993 let disambiguation_result = self
8994 .observe_purpose(
8995 ObservationPurpose::DisambiguationDetection,
8996 disambiguator.process_input_with_override(
8997 enriched_input,
8998 &context,
8999 state_override.as_ref(),
9000 Some(skill_disambig),
9001 ),
9002 )
9003 .await?;
9004 let current_state_generation = self
9005 .state_machine
9006 .as_ref()
9007 .map(|state_machine| state_machine.generation());
9008 if current_state_generation != expected_state_generation
9009 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
9010 {
9011 disambiguator.clear_pending().await;
9012 *self.pending_skill_id.write() = None;
9013 return Err(AgentError::Other(
9014 "State or reset ownership changed during skill disambiguation recheck"
9015 .to_string(),
9016 ));
9017 }
9018 match disambiguation_result {
9019 DisambiguationResult::Clear => {
9020 debug!(skill_id = %skill_id, "Skill re-check: all fields present");
9021 }
9022 DisambiguationResult::NeedsClarification {
9023 question,
9024 detection,
9025 } => {
9026 let admission = self
9027 .admit_disambiguation_redispatch(
9028 expected_disambiguation_epoch,
9029 expected_state_generation,
9030 )
9031 .await?;
9032 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
9033 info!(
9034 skill_id = %skill_id,
9035 ambiguity_type = ?detection.ambiguity_type,
9036 what_is_unclear = ?detection.what_is_unclear,
9037 "Skill re-check: still missing fields, asking again"
9038 );
9039 self.memory
9043 .add_message(ChatMessage::user(enriched_input))
9044 .await?;
9045 self.memory
9046 .add_message(ChatMessage::assistant(&question.question))
9047 .await?;
9048
9049 let response = AgentResponse::new(&question.question).with_metadata(
9050 "disambiguation",
9051 serde_json::json!({
9052 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
9053 "skill_id": skill_id,
9054 "options": question.options,
9055 "clarifying": question.clarifying,
9056 "detection": {
9057 "type": detection.ambiguity_type,
9058 "confidence": detection.confidence,
9059 "what_is_unclear": detection.what_is_unclear,
9060 }
9061 }),
9062 );
9063 drop(admission);
9064 self.finish_turn_if_root(&response).await?;
9065 return Ok(response);
9066 }
9067 DisambiguationResult::Clarified {
9068 enriched_input: re_enriched,
9069 ..
9070 } => {
9071 debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
9072 let admission = self
9073 .admit_disambiguation_redispatch(
9074 expected_disambiguation_epoch,
9075 expected_state_generation,
9076 )
9077 .await?;
9078 *self.pending_skill_id.write() = None;
9079 drop(admission);
9080 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9081 self.memory
9082 .add_message(ChatMessage::user(&re_enriched))
9083 .await?;
9084 return self
9085 .handle_skill_response(
9086 &re_enriched,
9087 skill_id,
9088 skill_response,
9089 &HashMap::new(),
9090 )
9091 .await;
9092 }
9093 DisambiguationResult::ProceedWithBestGuess {
9094 enriched_input: re_enriched,
9095 } => {
9096 debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9097 let admission = self
9098 .admit_disambiguation_redispatch(
9099 expected_disambiguation_epoch,
9100 expected_state_generation,
9101 )
9102 .await?;
9103 *self.pending_skill_id.write() = None;
9104 drop(admission);
9105 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9106 self.memory
9107 .add_message(ChatMessage::user(&re_enriched))
9108 .await?;
9109 return self
9110 .handle_skill_response(
9111 &re_enriched,
9112 skill_id,
9113 skill_response,
9114 &HashMap::new(),
9115 )
9116 .await;
9117 }
9118 DisambiguationResult::GiveUp { reason } => {
9119 *self.pending_skill_id.write() = None;
9120 let apology = self
9121 .generate_localized_apology(
9122 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9123 &reason,
9124 )
9125 .await
9126 .unwrap_or_else(|_| {
9127 format!("I'm sorry, I couldn't understand your request: {}", reason)
9128 });
9129 let response = AgentResponse::new(&apology);
9130 self.finish_turn_if_root(&response).await?;
9131 return Ok(response);
9132 }
9133 DisambiguationResult::Escalate { reason } => {
9134 *self.pending_skill_id.write() = None;
9135 let apology = self
9136 .generate_localized_apology(
9137 "Explain briefly that you're transferring the user to a human agent for help.",
9138 &reason,
9139 )
9140 .await
9141 .unwrap_or_else(|_| {
9142 format!("I need human assistance to help with your request: {}", reason)
9143 });
9144 let response = AgentResponse::new(&apology);
9145 self.finish_turn_if_root(&response).await?;
9146 return Ok(response);
9147 }
9148 DisambiguationResult::Abandoned { new_input } => {
9149 *self.pending_skill_id.write() = None;
9152 debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9153 if let Some(fresh) = new_input {
9154 return self.run_loop_internal(&fresh).await;
9155 }
9156 let ack = self
9157 .generate_localized_apology(
9158 "The user changed their mind about their previous request. \
9159 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9160 Do NOT apologize excessively. Be concise.",
9161 "User abandoned clarification",
9162 )
9163 .await
9164 .unwrap_or_else(|_| {
9165 "OK, no problem. What else can I help with?".to_string()
9166 });
9167 self.memory
9168 .add_message(ChatMessage::assistant(&ack))
9169 .await?;
9170 let response = AgentResponse::new(&ack);
9171 self.finish_turn_if_root(&response).await?;
9172 return Ok(response);
9173 }
9174 }
9175 }
9176
9177 let admission = self
9179 .admit_disambiguation_redispatch(
9180 expected_disambiguation_epoch,
9181 expected_state_generation,
9182 )
9183 .await?;
9184 *self.pending_skill_id.write() = None;
9185 drop(admission);
9186 let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9187 self.memory
9188 .add_message(ChatMessage::user(enriched_input))
9189 .await?;
9190 self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9191 .await
9192 }
9193
9194 async fn handle_skill_response(
9197 &self,
9198 processed_input: &str,
9199 skill_id: &str,
9200 skill_response: String,
9201 input_context: &HashMap<String, Value>,
9202 ) -> Result<AgentResponse> {
9203 let output_data = self.process_output(&skill_response, input_context).await?;
9204 let final_response = output_data.content;
9205
9206 self.memory
9207 .add_message(ChatMessage::assistant(&final_response))
9208 .await?;
9209
9210 self.check_memory_compression().await?;
9211
9212 self.increment_turn();
9213 self.evaluate_transitions(processed_input, &final_response)
9214 .await?;
9215
9216 let response = AgentResponse::new(final_response)
9217 .with_metadata("skill_id", serde_json::json!(skill_id));
9218 self.finish_turn_if_root(&response).await?;
9219 Ok(response)
9220 }
9221
9222 async fn handle_plan_and_execute(
9225 &self,
9226 processed_input: &str,
9227 input_context: &HashMap<String, Value>,
9228 auto_detected: bool,
9229 ) -> Result<AgentResponse> {
9230 let effective = self.get_effective_reasoning_config();
9231 let plan_reflection = effective
9232 .get_planning()
9233 .map(|c| c.reflection.clone())
9234 .unwrap_or_default();
9235
9236 let max_attempts = if plan_reflection.enabled {
9237 1 + plan_reflection.max_replans
9238 } else {
9239 1
9240 };
9241
9242 let mut plan = self.generate_plan(processed_input).await?;
9243 info!(
9244 plan_id = %plan.id,
9245 steps = plan.steps.len(),
9246 "Plan generated"
9247 );
9248
9249 let mut plan_result = String::new();
9250
9251 for attempt in 0..max_attempts {
9252 *self.current_plan.write() = Some(plan.clone());
9253 plan_result = self.execute_plan(&mut plan).await?;
9254
9255 info!(
9256 plan_status = ?plan.status,
9257 completed_steps = plan.completed_steps().count(),
9258 attempt = attempt + 1,
9259 "Plan execution completed"
9260 );
9261
9262 if !plan_reflection.enabled {
9263 break;
9264 }
9265
9266 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9267 if !has_failures {
9268 break;
9269 }
9270
9271 if attempt + 1 >= max_attempts {
9272 break;
9273 }
9274
9275 match plan_reflection.on_step_failure {
9276 StepFailureAction::Replan => {
9277 info!(attempt = attempt + 1, "Plan had failures, replanning");
9278 plan = self.generate_plan(processed_input).await?;
9279 }
9280 StepFailureAction::Abort => {
9281 warn!("Plan step failed, aborting");
9282 break;
9283 }
9284 StepFailureAction::Skip | StepFailureAction::Continue => {
9285 break;
9286 }
9287 }
9288 }
9289
9290 *self.current_plan.write() = Some(plan);
9291
9292 let output_data = self.process_output(&plan_result, input_context).await?;
9293 let final_content = output_data.content;
9294
9295 self.memory
9296 .add_message(ChatMessage::assistant(&final_content))
9297 .await?;
9298
9299 self.check_memory_compression().await?;
9300 self.increment_turn();
9301 self.evaluate_transitions(processed_input, &final_content)
9302 .await?;
9303
9304 let reasoning_metadata =
9305 ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9306
9307 let response = AgentResponse::new(&final_content).with_metadata(
9308 "reasoning",
9309 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9310 );
9311
9312 self.finish_turn_if_root(&response).await?;
9313 Ok(response)
9314 }
9315
9316 fn inject_reasoning_prompt(
9318 &self,
9319 messages: &mut [ChatMessage],
9320 reasoning_mode: &ReasoningMode,
9321 is_first_iteration: bool,
9322 ) {
9323 if !is_first_iteration {
9324 return;
9325 }
9326 match reasoning_mode {
9327 ReasoningMode::CoT => {
9328 if let Some(msg) = messages.first_mut()
9329 && matches!(msg.role, ai_agents_core::Role::System)
9330 {
9331 msg.content = self.build_cot_system_prompt(&msg.content);
9332 debug!("Applied Chain-of-Thought system prompt");
9333 }
9334 }
9335 ReasoningMode::React => {
9336 if let Some(msg) = messages.first_mut()
9337 && matches!(msg.role, ai_agents_core::Role::System)
9338 {
9339 msg.content = self.build_react_system_prompt(&msg.content);
9340 debug!("Applied ReAct system prompt");
9341 }
9342 }
9343 _ => {}
9344 }
9345 }
9346
9347 async fn generate_main_response_draft(
9352 &self,
9353 processed_input: &str,
9354 reasoning_mode: &ReasoningMode,
9355 ) -> Result<MainResponseDraft> {
9356 let llm = self.get_state_llm()?;
9357 let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9358 let mut messages = self
9359 .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9360 .await?;
9361 self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9362 let response = self
9363 .complete_main_llm_with_recovery(llm, &messages, &protocol)
9364 .await?;
9365 let content = response.content.trim().to_string();
9366 let (thinking, answer) = self.extract_thinking(&content);
9367 if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9368 return Ok(MainResponseDraft::ToolCalls {
9369 raw_content: content,
9370 calls,
9371 thinking,
9372 });
9373 }
9374 Ok(MainResponseDraft::Text {
9375 raw_content: answer,
9376 thinking,
9377 })
9378 }
9379
9380 async fn commit_main_response_draft(
9385 &self,
9386 processed_input: &str,
9387 input_context: &HashMap<String, Value>,
9388 draft: MainResponseDraft,
9389 reasoning_mode: ReasoningMode,
9390 auto_detected: bool,
9391 ) -> Result<AgentResponse> {
9392 self.commit_root_user_message(processed_input).await?;
9393 match draft {
9394 MainResponseDraft::Text {
9395 raw_content,
9396 thinking,
9397 } => {
9398 self.finish_text_response_from_model(CommittedTextResponse {
9399 processed_input,
9400 input_context,
9401 answer: raw_content,
9402 reasoning_mode,
9403 auto_detected,
9404 iterations: 1,
9405 thinking_content: thinking,
9406 all_tool_calls: Vec::new(),
9407 })
9408 .await
9409 }
9410 MainResponseDraft::ToolCalls {
9411 raw_content,
9412 calls,
9413 thinking: _,
9414 } => {
9415 let mut all_tool_calls = Vec::new();
9416 match self
9417 .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9418 .await?
9419 {
9420 ToolCallOutcome::Rejected(response) => {
9421 self.finish_turn_if_root(&response).await?;
9422 Ok(response)
9423 }
9424 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9425 self.continue_after_committed_tool_draft(processed_input)
9426 .await
9427 }
9428 }
9429 }
9430 }
9431 }
9432
9433 async fn continue_after_committed_tool_draft(
9438 &self,
9439 processed_input: &str,
9440 ) -> Result<AgentResponse> {
9441 *self.redispatch_depth.write() += 1;
9442 if let Some(context) = self.active_turn_context.write().as_mut() {
9443 context.enter_redispatch();
9444 }
9445 let result = Box::pin(self.run_loop_internal(processed_input)).await;
9446 *self.redispatch_depth.write() -= 1;
9447 if let Some(context) = self.active_turn_context.write().as_mut() {
9448 context.exit_redispatch();
9449 }
9450 let response = result?;
9451 self.finish_turn_if_root(&response).await?;
9452 Ok(response)
9453 }
9454
9455 async fn finish_text_response_from_model(
9460 &self,
9461 response: CommittedTextResponse<'_>,
9462 ) -> Result<AgentResponse> {
9463 let CommittedTextResponse {
9464 processed_input,
9465 input_context,
9466 answer,
9467 reasoning_mode,
9468 auto_detected,
9469 iterations,
9470 thinking_content,
9471 all_tool_calls,
9472 } = response;
9473 let output_data = self.process_output(&answer, input_context).await?;
9474 let mut final_content = if output_data.metadata.rejected {
9475 output_data
9476 .metadata
9477 .rejection_reason
9478 .unwrap_or_else(|| answer.to_string())
9479 } else {
9480 output_data.content
9481 };
9482 let llm = self.get_state_llm()?;
9483 let reflection_metadata;
9484 (final_content, reflection_metadata) = self
9485 .run_reflection(&*llm, processed_input, final_content)
9486 .await?;
9487 final_content =
9488 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9489 let final_content = {
9490 let result = self
9491 .post_loop_processing(processed_input, final_content)
9492 .await?;
9493 self.apply_post_loop_result(processed_input, result).await?
9494 };
9495 let response = self.build_agent_response(AgentResponseParts {
9496 content: final_content,
9497 all_tool_calls,
9498 reasoning_mode,
9499 auto_detected,
9500 iterations,
9501 thinking: thinking_content,
9502 reflection_metadata,
9503 });
9504 self.finish_turn_if_root(&response).await?;
9505 Ok(response)
9506 }
9507
9508 async fn run_committed_response_loop_with_reasoning(
9513 &self,
9514 processed_input: &str,
9515 input_context: &HashMap<String, Value>,
9516 reasoning_mode: ReasoningMode,
9517 auto_detected: bool,
9518 ) -> Result<AgentResponse> {
9519 self.commit_root_user_message(processed_input).await?;
9520 let llm = self.get_state_llm()?;
9521 let mut iterations = 0u32;
9522 let mut all_tool_calls = Vec::new();
9523 let mut thinking_content = None;
9524 loop {
9525 let effective_max = if reasoning_mode != ReasoningMode::None {
9526 let rc = self.get_effective_reasoning_config();
9527 self.max_iterations.min(rc.max_iterations)
9528 } else {
9529 self.max_iterations
9530 };
9531 if iterations >= effective_max {
9532 return Err(AgentError::Other(format!(
9533 "Max iterations ({}) exceeded",
9534 effective_max
9535 )));
9536 }
9537 iterations += 1;
9538 *self.iteration_count.write() = iterations;
9539 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9540 let mut messages = self
9541 .build_messages_internal(true, None, protocol.choice.is_none())
9542 .await?;
9543 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9544 self.hooks.on_llm_start(&messages).await;
9545 let llm_start = Instant::now();
9546 let response = self
9547 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9548 .await?;
9549 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9550 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9551 let content = response.content.trim();
9552 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9553 match self
9554 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9555 .await?
9556 {
9557 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9558 ToolCallOutcome::Rejected(resp) => {
9559 self.finish_turn_if_root(&resp).await?;
9560 return Ok(resp);
9561 }
9562 }
9563 }
9564 let (extracted_thinking, answer) = self.extract_thinking(content);
9565 if extracted_thinking.is_some() {
9566 thinking_content = extracted_thinking;
9567 }
9568 return self
9569 .finish_text_response_from_model(CommittedTextResponse {
9570 processed_input,
9571 input_context,
9572 answer,
9573 reasoning_mode,
9574 auto_detected,
9575 iterations,
9576 thinking_content,
9577 all_tool_calls,
9578 })
9579 .await;
9580 }
9581 }
9582
9583 async fn handle_tool_calls(
9585 &self,
9586 processed_input: &str,
9587 content: &str,
9588 tool_calls: Vec<ToolCall>,
9589 all_tool_calls: &mut Vec<ToolCall>,
9590 ) -> Result<ToolCallOutcome> {
9591 let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9595 if transition_fired {
9596 self.memory
9597 .add_message(ChatMessage::assistant(
9598 "(Transitioned to new state — tool call handled by workflow)",
9599 ))
9600 .await?;
9601 return Ok(ToolCallOutcome::TransitionFired);
9602 }
9603
9604 self.memory
9606 .add_message(ChatMessage::assistant(content))
9607 .await?;
9608 let native_tool_call = Self::is_native_tool_call_content(content);
9609
9610 let results = self.execute_tools_parallel(&tool_calls).await;
9611
9612 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9613 match result {
9614 Ok(output) => {
9615 self.memory
9616 .add_message(Self::tool_result_message(
9617 tool_call,
9618 &output,
9619 native_tool_call,
9620 ))
9621 .await?;
9622 }
9623 Err(e) => {
9624 if matches!(e, AgentError::HITLRejected(_)) {
9626 self.memory
9627 .add_message(ChatMessage::assistant(format!(
9628 "The operation was rejected by the approver: {}",
9629 e
9630 )))
9631 .await?;
9632 return Ok(ToolCallOutcome::Rejected(AgentResponse {
9634 content: format!("Operation cancelled: {}", e),
9635 metadata: None,
9636 tool_calls: Some(all_tool_calls.clone()),
9637 }));
9638 }
9639 self.memory
9640 .add_message(Self::tool_result_message(
9641 tool_call,
9642 &format!("Error: {}", e),
9643 native_tool_call,
9644 ))
9645 .await?;
9646 }
9647 }
9648 all_tool_calls.push(tool_call.clone());
9649 }
9650 Ok(ToolCallOutcome::Continue)
9651 }
9652
9653 async fn run_reflection(
9655 &self,
9656 llm: &dyn LLMProvider,
9657 processed_input: &str,
9658 mut content: String,
9659 ) -> Result<(String, Option<ReflectionMetadata>)> {
9660 let should_reflect = self.should_reflect(processed_input, &content).await?;
9661 if !should_reflect {
9662 return Ok((content, None));
9663 }
9664
9665 info!("Starting response reflection evaluation");
9666 let mut attempts = 0u32;
9667 let max_retries = self.reflection_config.max_retries;
9668 let mut history: Vec<ReflectionAttempt> = Vec::new();
9669
9670 loop {
9671 let evaluation = self.evaluate_response(processed_input, &content).await?;
9672
9673 if evaluation.passed || attempts >= max_retries {
9674 info!(
9675 passed = evaluation.passed,
9676 confidence = evaluation.confidence,
9677 attempts = attempts + 1,
9678 "Reflection evaluation complete"
9679 );
9680 let reflection_metadata = Some(
9681 ReflectionMetadata::new(evaluation)
9682 .with_attempts(attempts + 1)
9683 .with_history(history),
9684 );
9685 return Ok((content, reflection_metadata));
9686 }
9687
9688 debug!(
9689 attempt = attempts + 1,
9690 failed_criteria = evaluation.failed_criteria().count(),
9691 "Response did not meet criteria, retrying"
9692 );
9693
9694 history.push(
9695 ReflectionAttempt::new(&content, evaluation.clone())
9696 .with_feedback("Response did not meet quality criteria"),
9697 );
9698
9699 let feedback: Vec<String> = evaluation
9700 .failed_criteria()
9701 .map(|c| format!("- {}", c.criterion))
9702 .collect();
9703
9704 let retry_prompt = format!(
9705 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9706 feedback.join("\n")
9707 );
9708
9709 self.memory
9710 .add_message(ChatMessage::user(&retry_prompt))
9711 .await?;
9712
9713 let retry_messages = self.build_messages().await?;
9714 let retry_response = self
9715 .observe_purpose(
9716 ObservationPurpose::ReflectionEvaluation,
9717 llm.complete(&retry_messages, None),
9718 )
9719 .await
9720 .map_err(|e| AgentError::LLM(e.to_string()))?;
9721
9722 content = retry_response.content.trim().to_string();
9723 attempts += 1;
9724 }
9725 }
9726
9727 async fn post_loop_processing(
9730 &self,
9731 processed_input: &str,
9732 content: String,
9733 ) -> Result<PostLoopResult> {
9734 self.increment_turn();
9739
9740 self.run_context_extractors(processed_input).await;
9742
9743 let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9744
9745 if !transitioned {
9746 self.memory
9747 .add_message(ChatMessage::assistant(&content))
9748 .await?;
9749 self.check_memory_compression().await?;
9750 return Ok(PostLoopResult::NoTransition(content));
9751 }
9752
9753 if !self.should_regenerate_after_transition() {
9755 self.memory
9756 .add_message(ChatMessage::assistant(&content))
9757 .await?;
9758 self.check_memory_compression().await?;
9759 return Ok(PostLoopResult::Transitioned(content));
9760 }
9761
9762 if self.needs_redispatch_for_new_state() {
9766 info!("Post-transition NeedsRedispatch: new state requires full dispatch");
9767 return Ok(PostLoopResult::NeedsRedispatch);
9770 }
9771
9772 self.memory
9775 .add_message(ChatMessage::assistant(&content))
9776 .await?;
9777 self.check_memory_compression().await?;
9778
9779 let new_llm = self.get_state_llm()?;
9785 let mut final_content;
9786
9787 for post_iter in 0..self.max_iterations {
9788 let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
9789 let new_messages = self
9790 .build_messages_internal(true, None, protocol.choice.is_none())
9791 .await?;
9792 if post_iter == 0
9793 && let Some(system_msg) = new_messages.first()
9794 && system_msg.role == ai_agents_core::Role::System
9795 {
9796 debug!(
9797 prompt_preview =
9798 &system_msg.content[system_msg.content.len().saturating_sub(200)..],
9799 "Post-transition system prompt (last 200 chars)"
9800 );
9801 }
9802
9803 let new_response = self
9804 .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
9805 .await?;
9806 final_content = new_response.content.trim().to_string();
9807
9808 if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
9811 let native_tool_call = Self::is_native_tool_call_content(&final_content);
9812 debug!(
9813 post_iter = post_iter,
9814 tools = tool_calls.len(),
9815 "Post-transition tool call detected, executing"
9816 );
9817
9818 self.memory
9819 .add_message(ChatMessage::assistant(&final_content))
9820 .await?;
9821
9822 let results = self.execute_tools_parallel(&tool_calls).await;
9823 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9824 match result {
9825 Ok(output) => {
9826 self.memory
9827 .add_message(Self::tool_result_message(
9828 tool_call,
9829 &output,
9830 native_tool_call,
9831 ))
9832 .await?;
9833 }
9834 Err(e) => {
9835 self.memory
9836 .add_message(Self::tool_result_message(
9837 tool_call,
9838 &format!("Error: {}", e),
9839 native_tool_call,
9840 ))
9841 .await?;
9842 }
9843 }
9844 }
9845 continue;
9847 }
9848
9849 self.memory
9851 .add_message(ChatMessage::assistant(&final_content))
9852 .await?;
9853 return Ok(PostLoopResult::Transitioned(final_content));
9854 }
9855
9856 final_content = "Post-transition processing completed.".to_string();
9858 self.memory
9859 .add_message(ChatMessage::assistant(&final_content))
9860 .await?;
9861
9862 Ok(PostLoopResult::Transitioned(final_content))
9863 }
9864
9865 fn should_regenerate_after_transition(&self) -> bool {
9868 if let Some(ref sm) = self.state_machine {
9869 if !sm.config().regenerate_on_transition {
9871 return false;
9872 }
9873 if let Some(def) = sm.current_definition()
9875 && let Some(regen) = def.regenerate_on_enter
9876 {
9877 return regen;
9878 }
9879 }
9880 true
9881 }
9882
9883 fn needs_redispatch_for_new_state(&self) -> bool {
9886 if let Some(ref sm) = self.state_machine
9887 && let Some(def) = sm.current_definition()
9888 {
9889 if def.concurrent.is_some()
9890 || def.group_chat.is_some()
9891 || def.pipeline.is_some()
9892 || def.handoff.is_some()
9893 || def.delegate.is_some()
9894 {
9895 return true;
9896 }
9897 let effective = self.get_effective_reasoning_config();
9899 if !matches!(effective.mode, ReasoningMode::None) {
9900 return true;
9901 }
9902 }
9903 false
9904 }
9905
9906 async fn apply_post_loop_result(
9909 &self,
9910 processed_input: &str,
9911 result: PostLoopResult,
9912 ) -> Result<String> {
9913 match result {
9914 PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
9915 Ok(content)
9916 }
9917 PostLoopResult::NeedsRedispatch => {
9918 const MAX_REDISPATCH_DEPTH: u32 = 3;
9919 let current_depth = *self.redispatch_depth.read();
9920 if current_depth >= MAX_REDISPATCH_DEPTH {
9921 warn!(
9922 depth = current_depth,
9923 "Post-transition re-dispatch depth limit reached, returning empty response"
9924 );
9925 let content = String::new();
9926 self.memory
9927 .add_message(ChatMessage::assistant(&content))
9928 .await?;
9929 return Ok(content);
9930 }
9931 *self.redispatch_depth.write() += 1;
9932 if let Some(context) = self.active_turn_context.write().as_mut() {
9933 context.enter_redispatch();
9934 }
9935 info!(
9936 depth = current_depth + 1,
9937 "Re-dispatching for new state after transition"
9938 );
9939 let resp = Box::pin(self.run_loop_internal(processed_input)).await;
9940 *self.redispatch_depth.write() -= 1;
9941 if let Some(context) = self.active_turn_context.write().as_mut() {
9942 context.exit_redispatch();
9943 }
9944 resp.map(|r| r.content)
9945 }
9946 }
9947 }
9948
9949 fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
9951 let AgentResponseParts {
9952 content,
9953 all_tool_calls,
9954 reasoning_mode,
9955 auto_detected,
9956 iterations,
9957 thinking,
9958 reflection_metadata,
9959 } = parts;
9960 let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
9961 .with_thinking(thinking.clone().unwrap_or_default())
9962 .with_iterations(iterations)
9963 .with_auto_detected(auto_detected);
9964
9965 let mut response = AgentResponse::new(&content);
9966 if !all_tool_calls.is_empty() {
9967 response = response.with_tool_calls(all_tool_calls);
9968 }
9969
9970 if let Some(state) = self.current_state() {
9971 response = response.with_metadata("current_state", serde_json::json!(state));
9972 }
9973
9974 response = response.with_metadata(
9975 "reasoning",
9976 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9977 );
9978
9979 if let Some(ref refl_meta) = reflection_metadata {
9980 response = response.with_metadata(
9981 "reflection",
9982 serde_json::to_value(refl_meta).unwrap_or_default(),
9983 );
9984 }
9985
9986 response
9987 }
9988
9989 async fn handle_delegated_state(
9991 &self,
9992 input: &str,
9993 delegate_id: &str,
9994 state_def: &ai_agents_state::StateDefinition,
9995 ) -> Result<AgentResponse> {
9996 use std::time::Instant;
9997
9998 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
9999 AgentError::Config(format!(
10000 "State delegates to '{}' but no agent registry is configured. \
10001 Add a spawner section with auto_spawn to your YAML.",
10002 delegate_id
10003 ))
10004 })?;
10005
10006 let state_name = self
10007 .state_machine
10008 .as_ref()
10009 .map(|sm| sm.current())
10010 .unwrap_or_else(|| "unknown".to_string());
10011
10012 self.hooks.on_delegate_start(delegate_id, &state_name).await;
10013 let start = Instant::now();
10014
10015 let delegate = registry.get(delegate_id).ok_or_else(|| {
10016 AgentError::Other(format!(
10017 "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
10018 state_name, delegate_id
10019 ))
10020 })?;
10021
10022 let context_mode = state_def.delegate_context.clone().unwrap_or_default();
10024 let effective_input = self
10025 .observe_purpose(
10026 ObservationPurpose::OrchestrationRouting,
10027 crate::orchestration::context::prepare_delegate_input(
10028 input,
10029 &context_mode,
10030 &*self.memory,
10031 self.llm_registry.get("router").ok().as_deref(),
10032 ),
10033 )
10034 .await?;
10035
10036 let response = delegate
10037 .chat_with_actor_context(&effective_input, self.outbound_actor_context())
10038 .await?;
10039
10040 let duration_ms = start.elapsed().as_millis() as u64;
10041 self.hooks
10042 .on_delegate_complete(delegate_id, &state_name, duration_ms)
10043 .await;
10044
10045 let ctx_key = format!("delegation.{}.last_response", delegate_id);
10047 let _ = self.context_manager.set(
10048 &ctx_key,
10049 serde_json::Value::String(response.content.clone()),
10050 );
10051
10052 let _ = self.context_manager.set(
10054 "orchestration",
10055 serde_json::json!({
10056 "type": "delegate",
10057 "agent": delegate_id,
10058 "state": state_name,
10059 "response": response.content,
10060 "duration_ms": duration_ms,
10061 }),
10062 );
10063
10064 self.commit_root_user_message(input).await?;
10065
10066 let post_result = self
10069 .post_loop_processing(
10070 input,
10071 format!("[Delegated to {}]: {}", delegate_id, response.content),
10072 )
10073 .await?;
10074 let final_content = self.apply_post_loop_result(input, post_result).await?;
10075
10076 let mut result = AgentResponse::new(final_content);
10077
10078 let metadata = serde_json::json!({
10079 "orchestration": {
10080 "type": "delegate",
10081 "agent": delegate_id,
10082 "state": state_name,
10083 "response": response.content,
10084 "duration_ms": duration_ms,
10085 }
10086 });
10087 result.metadata = Some(
10088 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10089 metadata,
10090 )
10091 .unwrap_or_default(),
10092 );
10093
10094 self.finish_turn_if_root(&result).await?;
10095 Ok(result)
10096 }
10097
10098 async fn handle_concurrent_state(
10100 &self,
10101 input: &str,
10102 config: &ai_agents_state::ConcurrentStateConfig,
10103 ) -> Result<AgentResponse> {
10104 use std::time::Instant;
10105
10106 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10107 AgentError::Config(
10108 "Concurrent state requires an agent registry. Add a spawner section.".into(),
10109 )
10110 })?;
10111
10112 let context_mode = config.context_mode.clone().unwrap_or_default();
10117 let context_input = self
10118 .observe_purpose(
10119 ObservationPurpose::OrchestrationRouting,
10120 crate::orchestration::context::prepare_delegate_input(
10121 input,
10122 &context_mode,
10123 &*self.memory,
10124 self.llm_registry.get("router").ok().as_deref(),
10125 ),
10126 )
10127 .await?;
10128
10129 let effective_input = if let Some(ref tmpl) = config.input {
10130 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10131 .unwrap_or_else(|_| context_input.clone())
10132 } else {
10133 context_input
10134 };
10135
10136 let start = Instant::now();
10137
10138 let llm_name = config
10139 .aggregation
10140 .synthesizer_llm
10141 .as_deref()
10142 .unwrap_or("router");
10143 let llm_provider = self.llm_registry.get(llm_name).ok();
10144
10145 let vote_parallelism = if self.runtime_config.optimization.enabled
10146 && self
10147 .runtime_config
10148 .optimization
10149 .parallel_orchestration_vote_extraction
10150 {
10151 Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10152 } else {
10153 None
10154 };
10155
10156 let result = self
10157 .observe_purpose(
10158 ObservationPurpose::OrchestrationAggregation,
10159 scope_actor_context(
10160 self.outbound_actor_context(),
10161 crate::orchestration::concurrent(
10162 registry,
10163 &effective_input,
10164 &config.agents,
10165 &config.aggregation,
10166 llm_provider.as_deref(),
10167 config.min_required,
10168 config.timeout_ms,
10169 config.on_partial_failure.clone(),
10170 vote_parallelism,
10171 ),
10172 ),
10173 )
10174 .await?;
10175
10176 let duration_ms = start.elapsed().as_millis() as u64;
10177 let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10178 let strategy = format!("{:?}", config.aggregation.strategy);
10179 self.hooks
10180 .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10181 .await;
10182
10183 let _ = self.context_manager.set(
10185 "concurrent.result",
10186 serde_json::Value::String(result.response.content.clone()),
10187 );
10188
10189 let agents_json: Vec<serde_json::Value> = result
10191 .agent_results
10192 .iter()
10193 .map(|ar| {
10194 serde_json::json!({
10195 "id": ar.agent_id,
10196 "response": ar.response.as_ref().map(|r| r.content.as_str()),
10197 "success": ar.success,
10198 "error": ar.error,
10199 "duration_ms": ar.duration_ms,
10200 })
10201 })
10202 .collect();
10203
10204 let _ = self.context_manager.set(
10206 "orchestration",
10207 serde_json::json!({
10208 "type": "concurrent",
10209 "result": result.response.content,
10210 "strategy": strategy,
10211 "agents": agents_json,
10212 "duration_ms": duration_ms,
10213 }),
10214 );
10215
10216 self.commit_root_user_message(input).await?;
10217
10218 let post_result = self
10219 .post_loop_processing(input, result.response.content.clone())
10220 .await?;
10221 let final_content = self.apply_post_loop_result(input, post_result).await?;
10222
10223 let mut response = AgentResponse::new(final_content);
10224 let metadata = serde_json::json!({
10225 "orchestration": {
10226 "type": "concurrent",
10227 "result": result.response.content,
10228 "strategy": strategy,
10229 "agents": agents_json,
10230 "duration_ms": duration_ms,
10231 }
10232 });
10233 response.metadata = Some(
10234 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10235 metadata,
10236 )
10237 .unwrap_or_default(),
10238 );
10239
10240 self.finish_turn_if_root(&response).await?;
10241 Ok(response)
10242 }
10243
10244 async fn handle_group_chat_state(
10246 &self,
10247 input: &str,
10248 config: &ai_agents_state::GroupChatStateConfig,
10249 ) -> Result<AgentResponse> {
10250 use std::time::Instant;
10251
10252 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10253 AgentError::Config(
10254 "Group chat state requires an agent registry. Add a spawner section.".into(),
10255 )
10256 })?;
10257
10258 let start = Instant::now();
10259
10260 let llm_provider = self.llm_registry.get("router").ok();
10261
10262 let context_mode = config.context_mode.clone().unwrap_or_default();
10264 let context_input = self
10265 .observe_purpose(
10266 ObservationPurpose::OrchestrationRouting,
10267 crate::orchestration::context::prepare_delegate_input(
10268 input,
10269 &context_mode,
10270 &*self.memory,
10271 self.llm_registry.get("router").ok().as_deref(),
10272 ),
10273 )
10274 .await?;
10275
10276 let effective_topic = if let Some(ref tmpl) = config.input {
10278 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10279 .unwrap_or_else(|_| context_input.clone())
10280 } else {
10281 context_input
10282 };
10283
10284 let result = self
10285 .observe_purpose(
10286 ObservationPurpose::OrchestrationConversation,
10287 scope_actor_context(
10288 self.outbound_actor_context(),
10289 crate::orchestration::group_chat(
10290 registry,
10291 &effective_topic,
10292 config,
10293 llm_provider.as_deref(),
10294 Some(&*self.hooks),
10295 ),
10296 ),
10297 )
10298 .await?;
10299
10300 let duration_ms = start.elapsed().as_millis() as u64;
10301
10302 let _ = self.context_manager.set(
10304 "group_chat.conclusion",
10305 serde_json::Value::String(result.response.content.clone()),
10306 );
10307
10308 let transcript_json: Vec<serde_json::Value> = result
10310 .transcript
10311 .iter()
10312 .map(|t| {
10313 serde_json::json!({
10314 "speaker": t.speaker,
10315 "round": t.round,
10316 "content": t.content,
10317 })
10318 })
10319 .collect();
10320
10321 let _ = self.context_manager.set(
10323 "orchestration",
10324 serde_json::json!({
10325 "type": "group_chat",
10326 "conclusion": result.response.content,
10327 "transcript": transcript_json,
10328 "rounds": result.rounds_completed,
10329 "termination": result.termination_reason,
10330 "duration_ms": duration_ms,
10331 }),
10332 );
10333
10334 self.commit_root_user_message(input).await?;
10335
10336 let post_result = self
10337 .post_loop_processing(input, result.response.content.clone())
10338 .await?;
10339 let final_content = self.apply_post_loop_result(input, post_result).await?;
10340
10341 let mut response = AgentResponse::new(final_content);
10342 let metadata = serde_json::json!({
10343 "orchestration": {
10344 "type": "group_chat",
10345 "conclusion": result.response.content,
10346 "transcript": transcript_json,
10347 "rounds": result.rounds_completed,
10348 "termination": result.termination_reason,
10349 "duration_ms": duration_ms,
10350 }
10351 });
10352 response.metadata = Some(
10353 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10354 metadata,
10355 )
10356 .unwrap_or_default(),
10357 );
10358
10359 self.finish_turn_if_root(&response).await?;
10360 Ok(response)
10361 }
10362
10363 async fn handle_pipeline_state(
10365 &self,
10366 input: &str,
10367 config: &ai_agents_state::PipelineStateConfig,
10368 ) -> Result<AgentResponse> {
10369 use std::time::Instant;
10370
10371 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10372 AgentError::Config(
10373 "Pipeline state requires an agent registry. Add a spawner section.".into(),
10374 )
10375 })?;
10376
10377 let start = Instant::now();
10378
10379 let stages: Vec<crate::orchestration::PipelineStage> = config
10380 .stages
10381 .iter()
10382 .map(|entry| {
10383 let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10384 if let Some(tmpl) = entry.input() {
10385 stage = stage.with_input(tmpl);
10386 }
10387 stage
10388 })
10389 .collect();
10390
10391 let context_mode = config.context_mode.clone().unwrap_or_default();
10393 let context_input = self
10394 .observe_purpose(
10395 ObservationPurpose::OrchestrationRouting,
10396 crate::orchestration::context::prepare_delegate_input(
10397 input,
10398 &context_mode,
10399 &*self.memory,
10400 self.llm_registry.get("router").ok().as_deref(),
10401 ),
10402 )
10403 .await?;
10404
10405 let context_values = self.build_context_with_overlays();
10406 let result = self
10407 .observe_purpose(
10408 ObservationPurpose::OrchestrationRouting,
10409 scope_actor_context(
10410 self.outbound_actor_context(),
10411 crate::orchestration::pipeline(
10412 registry,
10413 &context_input,
10414 &stages,
10415 config.timeout_ms,
10416 Some(&*self.hooks),
10417 Some(&context_values),
10418 ),
10419 ),
10420 )
10421 .await?;
10422
10423 let duration_ms = start.elapsed().as_millis() as u64;
10424
10425 let _ = self.context_manager.set(
10427 "pipeline.result",
10428 serde_json::Value::String(result.response.content.clone()),
10429 );
10430
10431 let stages_json: Vec<serde_json::Value> = result
10433 .stage_outputs
10434 .iter()
10435 .map(|s| {
10436 serde_json::json!({
10437 "agent_id": s.agent_id,
10438 "output": s.output,
10439 "duration_ms": s.duration_ms,
10440 "skipped": s.skipped,
10441 })
10442 })
10443 .collect();
10444
10445 let _ = self.context_manager.set(
10447 "orchestration",
10448 serde_json::json!({
10449 "type": "pipeline",
10450 "result": result.response.content,
10451 "stages": stages_json,
10452 "duration_ms": duration_ms,
10453 }),
10454 );
10455
10456 self.commit_root_user_message(input).await?;
10457
10458 let post_result = self
10459 .post_loop_processing(input, result.response.content.clone())
10460 .await?;
10461 let final_content = self.apply_post_loop_result(input, post_result).await?;
10462
10463 let mut response = AgentResponse::new(final_content);
10464 let metadata = serde_json::json!({
10465 "orchestration": {
10466 "type": "pipeline",
10467 "result": result.response.content,
10468 "stages": stages_json,
10469 "duration_ms": duration_ms,
10470 }
10471 });
10472 response.metadata = Some(
10473 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10474 metadata,
10475 )
10476 .unwrap_or_default(),
10477 );
10478
10479 self.finish_turn_if_root(&response).await?;
10480 Ok(response)
10481 }
10482
10483 async fn handle_handoff_state(
10485 &self,
10486 input: &str,
10487 config: &ai_agents_state::HandoffStateConfig,
10488 ) -> Result<AgentResponse> {
10489 use std::time::Instant;
10490
10491 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10492 AgentError::Config(
10493 "Handoff state requires an agent registry. Add a spawner section.".into(),
10494 )
10495 })?;
10496
10497 let llm = self
10498 .llm_registry
10499 .get("router")
10500 .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10501
10502 let start = Instant::now();
10503
10504 let context_mode = config.context_mode.clone().unwrap_or_default();
10506 let context_input = self
10507 .observe_purpose(
10508 ObservationPurpose::OrchestrationRouting,
10509 crate::orchestration::context::prepare_delegate_input(
10510 input,
10511 &context_mode,
10512 &*self.memory,
10513 self.llm_registry.get("router").ok().as_deref(),
10514 ),
10515 )
10516 .await?;
10517
10518 let effective_input = if let Some(ref tmpl) = config.input {
10520 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10521 .unwrap_or_else(|_| context_input.clone())
10522 } else {
10523 context_input
10524 };
10525
10526 let result = self
10527 .observe_purpose(
10528 ObservationPurpose::OrchestrationRouting,
10529 scope_actor_context(
10530 self.outbound_actor_context(),
10531 crate::orchestration::handoff(
10532 registry,
10533 &effective_input,
10534 &config.initial_agent,
10535 &config.available_agents,
10536 config.max_handoffs,
10537 llm.as_ref(),
10538 Some(&*self.hooks),
10539 ),
10540 ),
10541 )
10542 .await?;
10543
10544 let duration_ms = start.elapsed().as_millis() as u64;
10545
10546 let _ = self.context_manager.set(
10548 "handoff.result",
10549 serde_json::Value::String(result.response.content.clone()),
10550 );
10551
10552 let chain_json: Vec<serde_json::Value> = result
10554 .handoff_chain
10555 .iter()
10556 .map(|h| {
10557 serde_json::json!({
10558 "from": h.from_agent,
10559 "to": h.to_agent,
10560 "reason": h.reason,
10561 })
10562 })
10563 .collect();
10564
10565 let _ = self.context_manager.set(
10567 "orchestration",
10568 serde_json::json!({
10569 "type": "handoff",
10570 "result": result.response.content,
10571 "final_agent": result.final_agent,
10572 "handoff_chain": chain_json,
10573 "duration_ms": duration_ms,
10574 }),
10575 );
10576
10577 self.commit_root_user_message(input).await?;
10578
10579 let post_result = self
10580 .post_loop_processing(input, result.response.content.clone())
10581 .await?;
10582 let final_content = self.apply_post_loop_result(input, post_result).await?;
10583
10584 let mut response = AgentResponse::new(final_content);
10585 let metadata = serde_json::json!({
10586 "orchestration": {
10587 "type": "handoff",
10588 "result": result.response.content,
10589 "final_agent": result.final_agent,
10590 "handoff_chain": chain_json,
10591 "duration_ms": duration_ms,
10592 }
10593 });
10594 response.metadata = Some(
10595 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10596 metadata,
10597 )
10598 .unwrap_or_default(),
10599 );
10600
10601 self.finish_turn_if_root(&response).await?;
10602 Ok(response)
10603 }
10604
10605 async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10607 self.begin_root_turn();
10608 self.pre_turn_session_lifecycle().await;
10610
10611 let input_data = self.process_input(input).await?;
10612 self.update_active_turn_context(&input_data.content, input_data.context.clone());
10613
10614 for (key, value) in &input_data.context {
10617 let _ = self.context_manager.set(key, value.clone());
10618 }
10619
10620 if input_data.metadata.rejected {
10621 let reason = input_data
10622 .metadata
10623 .rejection_reason
10624 .unwrap_or_else(|| "Input rejected".to_string());
10625 warn!(reason = %reason, "Input rejected");
10626 let response = AgentResponse::new(reason);
10627 self.finish_turn_if_root(&response).await?;
10628 return Ok(response);
10629 }
10630
10631 let processed_input = &input_data.content;
10632
10633 if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10634 return Ok(response);
10635 }
10636
10637 if let Some(ref sm) = self.state_machine
10639 && let Some(def) = sm.current_definition()
10640 {
10641 if let Some(ref delegate_id) = def.delegate {
10642 return self
10643 .handle_delegated_state(processed_input, delegate_id, &def)
10644 .await;
10645 }
10646 if let Some(ref concurrent_config) = def.concurrent {
10647 return self
10648 .handle_concurrent_state(processed_input, concurrent_config)
10649 .await;
10650 }
10651 if let Some(ref group_chat_config) = def.group_chat {
10652 return self
10653 .handle_group_chat_state(processed_input, group_chat_config)
10654 .await;
10655 }
10656 if let Some(ref pipeline_config) = def.pipeline {
10657 return self
10658 .handle_pipeline_state(processed_input, pipeline_config)
10659 .await;
10660 }
10661 if let Some(ref handoff_config) = def.handoff {
10662 return self
10663 .handle_handoff_state(processed_input, handoff_config)
10664 .await;
10665 }
10666 }
10667
10668 if let Some(response) =
10673 Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10674 {
10675 return Ok(response);
10676 }
10677
10678 match self.try_skill_route(processed_input).await? {
10679 SkillRouteResult::Response { skill_id, content } => {
10680 self.commit_root_user_message(processed_input).await?;
10681 return self
10682 .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10683 .await;
10684 }
10685 SkillRouteResult::NeedsClarification {
10686 response,
10687 ownership,
10688 } => {
10689 let admission = self
10690 .admit_optional_disambiguation_ownership(ownership)
10691 .await?;
10692 self.commit_root_user_message(processed_input).await?;
10693 if let Some(q) = response
10694 .metadata
10695 .as_ref()
10696 .and_then(|m| m.get("disambiguation"))
10697 .and_then(|d| d.get("status"))
10698 .and_then(|s| s.as_str())
10699 && q == "awaiting_clarification"
10700 {
10701 self.memory
10704 .add_message(ChatMessage::assistant(&response.content))
10705 .await?;
10706 }
10707 drop(admission);
10708 self.finish_turn_if_root(&response).await?;
10709 return Ok(response);
10710 }
10711 SkillRouteResult::NoMatch => {} }
10713
10714 let effective_reasoning = self.get_effective_reasoning_config();
10715 let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10716 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10717
10718 info!(
10719 reasoning_mode = ?reasoning_mode,
10720 auto_detected = auto_detected,
10721 reflection_enabled = ?self.reflection_config.enabled,
10722 "Reasoning mode determined"
10723 );
10724
10725 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10726 self.commit_root_user_message(processed_input).await?;
10727 return self
10728 .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10729 .await;
10730 }
10731
10732 self.commit_root_user_message(processed_input).await?;
10733
10734 let mut iterations = 0u32;
10735 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10736 let mut thinking_content: Option<String> = None;
10737
10738 let llm = self.get_state_llm()?;
10739
10740 loop {
10741 let effective_max = if reasoning_mode != ReasoningMode::None {
10743 let rc = self.get_effective_reasoning_config();
10744 self.max_iterations.min(rc.max_iterations)
10745 } else {
10746 self.max_iterations
10747 };
10748
10749 if iterations >= effective_max {
10750 let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10751 self.hooks.on_error(&err).await;
10752 error!(iterations = iterations, "Max iterations exceeded");
10753 return Err(err);
10754 }
10755 iterations += 1;
10756 *self.iteration_count.write() = iterations;
10757
10758 debug!(iteration = iterations, max = effective_max, "LLM call");
10759
10760 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
10761 let mut messages = self
10762 .build_messages_internal(true, None, protocol.choice.is_none())
10763 .await?;
10764 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
10765
10766 self.hooks.on_llm_start(&messages).await;
10767 let llm_start = Instant::now();
10768 let response = self
10769 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
10770 .await?;
10771
10772 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
10773 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
10774
10775 let content = response.content.trim();
10776
10777 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
10778 match self
10779 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
10780 .await?
10781 {
10782 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
10783 ToolCallOutcome::Rejected(resp) => {
10784 self.finish_turn_if_root(&resp).await?;
10785 return Ok(resp);
10786 }
10787 }
10788 }
10789
10790 let (extracted_thinking, answer) = self.extract_thinking(content);
10791 if extracted_thinking.is_some() {
10792 thinking_content = extracted_thinking;
10793 }
10794
10795 let output_data = self.process_output(&answer, &input_data.context).await?;
10796
10797 let mut final_content = if output_data.metadata.rejected {
10798 output_data
10799 .metadata
10800 .rejection_reason
10801 .unwrap_or_else(|| answer.to_string())
10802 } else {
10803 output_data.content
10804 };
10805
10806 let reflection_metadata;
10808 (final_content, reflection_metadata) = self
10809 .run_reflection(&*llm, processed_input, final_content)
10810 .await?;
10811
10812 final_content =
10813 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
10814
10815 let final_content = {
10819 let result = self
10820 .post_loop_processing(processed_input, final_content)
10821 .await?;
10822 self.apply_post_loop_result(processed_input, result).await?
10823 };
10824
10825 let reflected = reflection_metadata.is_some();
10826 let reasoning_mode_debug = format!("{:?}", reasoning_mode);
10827
10828 let response = self.build_agent_response(AgentResponseParts {
10829 content: final_content,
10830 all_tool_calls,
10831 reasoning_mode,
10832 auto_detected,
10833 iterations,
10834 thinking: thinking_content,
10835 reflection_metadata,
10836 });
10837
10838 self.finish_turn_if_root(&response).await?;
10839
10840 let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
10841 info!(
10842 tool_calls = tool_call_count,
10843 response_len = response.content.len(),
10844 reasoning_mode = %reasoning_mode_debug,
10845 reflected = reflected,
10846 "Chat completed"
10847 );
10848 return Ok(response);
10849 }
10850 }
10851
10852 async fn generate_buffered_streaming_draft(
10853 &self,
10854 processed_input: &str,
10855 routing_resolved: Arc<AtomicBool>,
10856 ) -> Result<StreamingDraftResult> {
10857 let llm = self.get_state_llm()?;
10858 if llm.configured_tool_choice().is_some() {
10859 let draft = self
10860 .generate_main_response_draft(processed_input, &ReasoningMode::None)
10861 .await?;
10862 return Ok(StreamingDraftResult::new(draft, Vec::new()));
10863 }
10864 let messages = self.build_messages_for_draft(processed_input).await?;
10865 let mut stream = self
10866 .observe_purpose(
10867 ObservationPurpose::MainResponse,
10868 llm.complete_stream(&messages, None),
10869 )
10870 .await
10871 .map_err(|e| AgentError::LLM(e.to_string()))?;
10872 let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
10873 let mut chunks = Vec::new();
10874 let mut accumulated = String::new();
10875 while let Some(chunk_result) = stream.next().await {
10876 let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
10877 accumulated.push_str(&chunk.delta);
10878 let stream_chunk = StreamChunk::content(chunk.delta);
10879 if routing_resolved.load(Ordering::SeqCst) {
10880 chunks.push(stream_chunk);
10881 } else {
10882 buffer.push(stream_chunk)?;
10883 }
10884 }
10885 chunks.splice(0..0, buffer.drain());
10886 let content = accumulated.trim().to_string();
10887 let draft = if let Some(calls) = self.parse_tool_calls(&content) {
10888 MainResponseDraft::ToolCalls {
10889 raw_content: content,
10890 calls,
10891 thinking: None,
10892 }
10893 } else {
10894 MainResponseDraft::Text {
10895 raw_content: content,
10896 thinking: None,
10897 }
10898 };
10899 Ok(StreamingDraftResult::new(draft, chunks))
10900 }
10901
10902 async fn try_buffered_streaming_branches(
10903 &self,
10904 processed_input: &str,
10905 input_context: &HashMap<String, Value>,
10906 ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
10907 let optimization = &self.runtime_config.optimization;
10908 if !optimization.enabled {
10909 return Ok(None);
10910 }
10911 let transition_enabled =
10912 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
10913 if !transition_enabled {
10914 return Ok(None);
10915 }
10916 let mut branch_scheduler =
10917 TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
10918 if !branch_scheduler.reserve_task() {
10919 return Ok(None);
10920 }
10921 if !self
10922 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
10923 {
10924 branch_scheduler.release_task();
10925 return Ok(None);
10926 }
10927 if !branch_scheduler.reserve_task() {
10928 branch_scheduler.release_task();
10929 return Ok(None);
10930 }
10931 let mut main_branch = RuntimeBranch::new(
10932 RuntimeTaskPurpose::MainResponse,
10933 RuntimeOptimizationKind::BufferedStreamingRouting,
10934 RuntimeTaskPriority::Normal,
10935 RuntimeCommitBehavior::FinalResponse,
10936 );
10937 let mut transition_branch = RuntimeBranch::new(
10938 RuntimeTaskPurpose::StateTransition,
10939 RuntimeOptimizationKind::ParallelStateTransition,
10940 RuntimeTaskPriority::Critical,
10941 RuntimeCommitBehavior::TransitionDecision,
10942 );
10943 let main_id = main_branch.branch_id();
10944 let transition_id = transition_branch.branch_id();
10945 let routing_resolved = Arc::new(AtomicBool::new(false));
10946 let mut main_future =
10947 Box::pin(crate::optimization::observability::with_branch_observation(
10948 &main_id,
10949 RuntimeOptimizationKind::BufferedStreamingRouting,
10950 RuntimeCommitBehavior::FinalResponse,
10951 self.generate_buffered_streaming_draft(
10952 processed_input,
10953 Arc::clone(&routing_resolved),
10954 ),
10955 ));
10956 let mut transition_future =
10957 Box::pin(crate::optimization::observability::with_branch_observation(
10958 &transition_id,
10959 RuntimeOptimizationKind::ParallelStateTransition,
10960 RuntimeCommitBehavior::TransitionDecision,
10961 self.select_parallel_transition_candidate(processed_input),
10962 ));
10963 let mut main_pending = true;
10964 let mut transition_pending = true;
10965 let mut main_result: Option<Result<StreamingDraftResult>> = None;
10966 let mut transition_finalized = false;
10967 let mut transition_candidate: Option<TransitionCandidate> = None;
10968 loop {
10969 if let Some(candidate) = transition_candidate.take() {
10970 if self
10971 .approve_transition_target(&candidate.from_state, candidate.target())
10972 .await?
10973 {
10974 drop(main_future);
10976 drop(transition_future);
10977 self.finalize_branch_loss(
10978 &main_id,
10979 RuntimeOptimizationKind::BufferedStreamingRouting,
10980 RuntimeCommitBehavior::FinalResponse,
10981 main_pending,
10982 main_result.as_ref().map(|result| result.is_err()),
10983 );
10984 if !self
10985 .apply_pre_response_transition_candidate(
10986 &candidate,
10987 &HashMap::new(),
10988 processed_input,
10989 )
10990 .await?
10991 {
10992 self.finalize_optional_branch(
10993 &transition_id,
10994 RuntimeOptimizationKind::ParallelStateTransition,
10995 RuntimeCommitBehavior::TransitionDecision,
10996 "discarded",
10997 false,
10998 );
10999 return Ok(None);
11000 }
11001 self.finalize_optional_branch(
11002 &transition_id,
11003 RuntimeOptimizationKind::ParallelStateTransition,
11004 RuntimeCommitBehavior::TransitionDecision,
11005 "committed",
11006 true,
11007 );
11008 let response = self.redispatch_current_state(processed_input).await?;
11009 return Ok(Some((
11010 response.clone(),
11011 vec![StreamChunk::content(response.content)],
11012 )));
11013 }
11014 self.finalize_optional_branch(
11015 &transition_id,
11016 RuntimeOptimizationKind::ParallelStateTransition,
11017 RuntimeCommitBehavior::TransitionDecision,
11018 "discarded",
11019 false,
11020 );
11021 routing_resolved.store(true, Ordering::SeqCst);
11022 transition_finalized = true;
11023 }
11024 if transition_finalized && let Some(result) = main_result.take() {
11025 let stream_draft = match result {
11026 Ok(stream_draft) => stream_draft,
11027 Err(error) => {
11028 self.finalize_optional_branch(
11029 &main_id,
11030 RuntimeOptimizationKind::BufferedStreamingRouting,
11031 RuntimeCommitBehavior::FinalResponse,
11032 "failed",
11033 false,
11034 );
11035 return Err(error);
11036 }
11037 };
11038 let raw_draft_content = stream_draft.draft.raw_content().to_string();
11039 let buffered_chunks = stream_draft.chunks;
11040 self.finalize_optional_branch(
11041 &main_id,
11042 RuntimeOptimizationKind::BufferedStreamingRouting,
11043 RuntimeCommitBehavior::FinalResponse,
11044 "committed",
11045 true,
11046 );
11047 let response = self
11048 .commit_main_response_draft(
11049 processed_input,
11050 input_context,
11051 stream_draft.draft,
11052 ReasoningMode::None,
11053 false,
11054 )
11055 .await?;
11056 let chunks = if response.content == raw_draft_content {
11057 buffered_chunks
11058 } else {
11059 vec![StreamChunk::content(response.content.clone())]
11060 };
11061 return Ok(Some((response, chunks)));
11062 }
11063 tokio::select! {
11064 result = &mut main_future, if main_pending => {
11065 main_pending = false;
11066 main_branch.transition_to(RuntimeBranchStatus::Completed)?;
11067 main_result = Some(result);
11068 }
11069 result = &mut transition_future, if transition_pending => {
11070 transition_pending = false;
11071 transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
11072 match result {
11073 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
11074 transition_candidate = Some(candidate)
11075 }
11076 Ok(ParallelTransitionSelection::NoMatch) => {
11077 self.finalize_optional_branch(
11078 &transition_id,
11079 RuntimeOptimizationKind::ParallelStateTransition,
11080 RuntimeCommitBehavior::TransitionDecision,
11081 "discarded",
11082 false,
11083 );
11084 routing_resolved.store(true, Ordering::SeqCst);
11085 transition_finalized = true;
11086 }
11087 Ok(ParallelTransitionSelection::ReservationExhausted) => {
11088 self.finalize_optional_branch(
11089 &transition_id,
11090 RuntimeOptimizationKind::ParallelStateTransition,
11091 RuntimeCommitBehavior::TransitionDecision,
11092 "cancelled",
11093 false,
11094 );
11095 routing_resolved.store(true, Ordering::SeqCst);
11096 self.finalize_branch_loss(
11097 &main_id,
11098 RuntimeOptimizationKind::BufferedStreamingRouting,
11099 RuntimeCommitBehavior::FinalResponse,
11100 main_pending,
11101 main_result.as_ref().map(|result| result.is_err()),
11102 );
11103 return Ok(None);
11104 }
11105 Err(_) => {
11106 self.finalize_optional_branch(
11107 &transition_id,
11108 RuntimeOptimizationKind::ParallelStateTransition,
11109 RuntimeCommitBehavior::TransitionDecision,
11110 "failed",
11111 false,
11112 );
11113 routing_resolved.store(true, Ordering::SeqCst);
11114 transition_finalized = true;
11115 }
11116 }
11117 }
11118 }
11119 }
11120 }
11121
11122 fn run_loop_internal_stream<'a>(
11126 &'a self,
11127 input: &'a str,
11128 terminal: RuntimeStreamTerminalSlot,
11129 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11130 let include_tool_events = self.streaming.include_tool_events;
11131 let include_state_events = self.streaming.include_state_events;
11132
11133 Box::pin(async_stream::stream! {
11134 self.begin_root_turn();
11135 self.pre_turn_session_lifecycle().await;
11137
11138 let input_data = match self.process_input(input).await {
11139 Ok(data) => data,
11140 Err(e) => {
11141 yield StreamChunk::error(e.to_string());
11142 return;
11143 }
11144 };
11145 self.update_active_turn_context(&input_data.content, input_data.context.clone());
11146
11147 for (key, value) in &input_data.context {
11149 let _ = self.context_manager.set(key, value.clone());
11150 }
11151
11152 if input_data.metadata.rejected {
11153 let reason = input_data
11154 .metadata
11155 .rejection_reason
11156 .unwrap_or_else(|| "Input rejected".to_string());
11157 warn!(reason = %reason, "Input rejected (stream)");
11158 yield StreamChunk::error(reason);
11159 return;
11160 }
11161
11162 let processed_input = &input_data.content;
11163
11164 if self.runtime_config.optimization.enabled
11165 && matches!(
11166 self.runtime_config.optimization.streaming_policy,
11167 crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11168 )
11169 {
11170 match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11175 Ok(Some((response, chunks))) => {
11176 for chunk in chunks {
11177 yield chunk;
11178 }
11179 record_runtime_stream_final(&terminal, response);
11180 yield StreamChunk::Done {};
11181 return;
11182 }
11183 Ok(None) => {}
11184 Err(e) => {
11185 yield StreamChunk::error(e.to_string());
11186 return;
11187 }
11188 }
11189 }
11190
11191 if self.runtime_config.optimization.enabled
11192 && matches!(
11193 self.runtime_config.optimization.streaming_policy,
11194 crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11195 )
11196 {
11197 match self.try_pre_response_transition(processed_input).await {
11198 Ok(Some(response)) => {
11199 yield StreamChunk::content(&response.content);
11200 record_runtime_stream_final(&terminal, response);
11201 yield StreamChunk::Done {};
11202 return;
11203 }
11204 Ok(None) => {}
11205 Err(e) => {
11206 yield StreamChunk::error(e.to_string());
11207 return;
11208 }
11209 }
11210 }
11211
11212 if let Some(ref sm) = self.state_machine
11214 && let Some(def) = sm.current_definition()
11215 {
11216 let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11217 Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11218 } else if let Some(ref concurrent_config) = def.concurrent {
11219 Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11220 } else if let Some(ref group_chat_config) = def.group_chat {
11221 Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11222 } else if let Some(ref pipeline_config) = def.pipeline {
11223 Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11224 } else if let Some(ref handoff_config) = def.handoff {
11225 Some(self.handle_handoff_state(processed_input, handoff_config).await)
11226 } else {
11227 None
11228 };
11229
11230 if let Some(result) = orchestration_result {
11231 match result {
11232 Ok(response) => {
11233 yield StreamChunk::content(&response.content);
11234 record_runtime_stream_final(&terminal, response);
11235 yield StreamChunk::Done {};
11236 }
11237 Err(e) => {
11238 yield StreamChunk::error(e.to_string());
11239 }
11240 }
11241 return;
11242 }
11243 }
11244
11245 match self.try_skill_route(processed_input).await {
11247 Ok(SkillRouteResult::Response { skill_id, content }) => {
11248 if let Err(e) = self.commit_root_user_message(processed_input).await {
11249 yield StreamChunk::error(e.to_string());
11250 return;
11251 }
11252 match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11253 Ok(resp) => {
11254 yield StreamChunk::content(&resp.content);
11255 record_runtime_stream_final(&terminal, resp);
11256 yield StreamChunk::Done {};
11257 return;
11258 }
11259 Err(e) => {
11260 yield StreamChunk::error(e.to_string());
11261 return;
11262 }
11263 }
11264 }
11265 Ok(SkillRouteResult::NeedsClarification {
11266 response,
11267 ownership,
11268 }) => {
11269 let admission = match self
11270 .admit_optional_disambiguation_ownership(ownership)
11271 .await
11272 {
11273 Ok(admission) => admission,
11274 Err(e) => {
11275 yield StreamChunk::error(e.to_string());
11276 return;
11277 }
11278 };
11279 if let Err(e) = self.commit_root_user_message(processed_input).await {
11280 yield StreamChunk::error(e.to_string());
11281 return;
11282 }
11283 let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11284 drop(admission);
11285 if let Err(e) = self.finish_turn_if_root(&response).await {
11286 yield StreamChunk::error(e.to_string());
11287 return;
11288 }
11289 yield StreamChunk::content(&response.content);
11290 record_runtime_stream_final(&terminal, response);
11291 yield StreamChunk::Done {};
11292 return;
11293 }
11294 Ok(SkillRouteResult::NoMatch) => {} Err(e) => {
11296 yield StreamChunk::error(e.to_string());
11297 return;
11298 }
11299 }
11300
11301 let effective_reasoning = self.get_effective_reasoning_config();
11303 let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11304 Ok(mode) => mode,
11305 Err(e) => {
11306 yield StreamChunk::error(e.to_string());
11307 return;
11308 }
11309 };
11310 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11311
11312 info!(
11313 reasoning_mode = ?reasoning_mode,
11314 auto_detected = auto_detected,
11315 "Reasoning mode determined (stream)"
11316 );
11317
11318 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11320 if let Err(e) = self.commit_root_user_message(processed_input).await {
11321 yield StreamChunk::error(e.to_string());
11322 return;
11323 }
11324 match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11325 Ok(resp) => {
11326 yield StreamChunk::content(&resp.content);
11327 record_runtime_stream_final(&terminal, resp);
11328 yield StreamChunk::Done {};
11329 return;
11330 }
11331 Err(e) => {
11332 yield StreamChunk::error(e.to_string());
11333 return;
11334 }
11335 }
11336 }
11337
11338 if let Err(e) = self.commit_root_user_message(processed_input).await {
11339 yield StreamChunk::error(e.to_string());
11340 return;
11341 }
11342
11343 let llm = match self.get_state_llm() {
11344 Ok(llm) => llm,
11345 Err(e) => {
11346 yield StreamChunk::error(e.to_string());
11347 return;
11348 }
11349 };
11350
11351 let mut iterations = 0u32;
11352 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11353 let mut thinking_content: Option<String> = None;
11354
11355 loop {
11356 let effective_max = if reasoning_mode != ReasoningMode::None {
11358 let rc = self.get_effective_reasoning_config();
11359 self.max_iterations.min(rc.max_iterations)
11360 } else {
11361 self.max_iterations
11362 };
11363
11364 if iterations >= effective_max {
11365 let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11366 let err = AgentError::Other(err_msg.clone());
11367 self.hooks.on_error(&err).await;
11368 error!(iterations = iterations, "Max iterations exceeded (stream)");
11369 yield StreamChunk::error(err_msg);
11370 return;
11371 }
11372 iterations += 1;
11373 *self.iteration_count.write() = iterations;
11374
11375 debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11376
11377 let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11378 Ok(protocol) => protocol,
11379 Err(e) => {
11380 yield StreamChunk::error(e.to_string());
11381 return;
11382 }
11383 };
11384 let mut messages = match self
11385 .build_messages_internal(true, None, protocol.choice.is_none())
11386 .await
11387 {
11388 Ok(m) => m,
11389 Err(e) => {
11390 yield StreamChunk::error(e.to_string());
11391 return;
11392 }
11393 };
11394 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11395
11396 self.hooks.on_llm_start(&messages).await;
11397 let llm_start = Instant::now();
11398
11399 let reflection_active = self
11402 .should_reflect(processed_input, "")
11403 .await
11404 .unwrap_or_default();
11405
11406 let buffered_decision = reflection_active || protocol.choice.is_some();
11407 let content = if buffered_decision {
11408 let response = match self
11412 .complete_main_llm_with_recovery(
11413 Arc::clone(&llm),
11414 &messages,
11415 &protocol,
11416 )
11417 .await
11418 {
11419 Ok(r) => r,
11420 Err(e) => {
11421 yield StreamChunk::error(e.to_string());
11422 return;
11423 }
11424 };
11425 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11426 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11427 response.content.trim().to_string()
11428 } else {
11429 let llm_stream = match self
11431 .observe_purpose(
11432 ObservationPurpose::MainResponse,
11433 llm.complete_stream(&messages, None),
11434 )
11435 .await
11436 {
11437 Ok(s) => s,
11438 Err(e) => {
11439 yield StreamChunk::error(e.to_string());
11440 return;
11441 }
11442 };
11443 let mut accumulated = String::new();
11444 let mut stream_inner = llm_stream;
11445 while let Some(chunk_result) = stream_inner.next().await {
11446 match chunk_result {
11447 Ok(chunk) => {
11448 accumulated.push_str(&chunk.delta);
11449 yield StreamChunk::content(chunk.delta);
11450 }
11451 Err(e) => {
11452 yield StreamChunk::error(e.to_string());
11453 return;
11454 }
11455 }
11456 }
11457 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11458 let llm_response = ai_agents_core::LLMResponse::new(
11460 accumulated.trim(),
11461 ai_agents_core::FinishReason::Stop,
11462 );
11463 self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11464 accumulated.trim().to_string()
11465 };
11466
11467 if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11469 let native_tool_call = Self::is_native_tool_call_content(&content);
11470 let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11473 Ok(v) => v,
11474 Err(e) => {
11475 yield StreamChunk::error(e.to_string());
11476 return;
11477 }
11478 };
11479 if transition_fired {
11480 let _ = self.memory.add_message(ChatMessage::assistant(
11481 "(Transitioned to new state — tool call handled by workflow)",
11482 )).await;
11483
11484 if include_state_events
11485 && let Some(state) = self.current_state()
11486 {
11487 yield StreamChunk::state_transition(None, state);
11488 }
11489 continue;
11490 }
11491
11492 let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11494
11495 let results = self.execute_tools_parallel(&tool_calls).await;
11497
11498 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11499 if include_tool_events {
11500 yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11501 }
11502
11503 match result {
11504 Ok(output) => {
11505 if include_tool_events {
11506 yield StreamChunk::tool_result(
11507 &tool_call.id,
11508 &tool_call.name,
11509 &output,
11510 true,
11511 );
11512 }
11513 let _ = self.memory
11514 .add_message(Self::tool_result_message(
11515 tool_call,
11516 &output,
11517 native_tool_call,
11518 ))
11519 .await;
11520 }
11521 Err(e) => {
11522 if matches!(e, AgentError::HITLRejected(_)) {
11523 let _ = self.memory.add_message(ChatMessage::assistant(
11524 format!("The operation was rejected by the approver: {}", e),
11525 )).await;
11526 let response = AgentResponse {
11527 content: format!("Operation cancelled: {}", e),
11528 metadata: None,
11529 tool_calls: Some(all_tool_calls.clone()),
11530 };
11531 if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11532 yield StreamChunk::error(finalize_error.to_string());
11533 return;
11534 }
11535 let legacy_error = response.content.clone();
11536 record_runtime_stream_final(&terminal, response);
11537 yield StreamChunk::error(legacy_error);
11538 yield StreamChunk::Done {};
11539 return;
11540 }
11541 if include_tool_events {
11542 yield StreamChunk::tool_result(
11543 &tool_call.id,
11544 &tool_call.name,
11545 e.to_string(),
11546 false,
11547 );
11548 }
11549 let _ = self.memory
11550 .add_message(Self::tool_result_message(
11551 tool_call,
11552 &format!("Error: {}", e),
11553 native_tool_call,
11554 ))
11555 .await;
11556 }
11557 }
11558 all_tool_calls.push(tool_call.clone());
11559
11560 if include_tool_events {
11561 yield StreamChunk::tool_end(&tool_call.id);
11562 }
11563 }
11564 continue;
11565 }
11566
11567 let (extracted_thinking, answer) = self.extract_thinking(&content);
11569 if extracted_thinking.is_some() {
11570 thinking_content = extracted_thinking;
11571 }
11572
11573 let output_data = match self.process_output(&answer, &input_data.context).await {
11574 Ok(d) => d,
11575 Err(e) => {
11576 yield StreamChunk::error(e.to_string());
11577 return;
11578 }
11579 };
11580
11581 let final_content = if output_data.metadata.rejected {
11582 output_data
11583 .metadata
11584 .rejection_reason
11585 .unwrap_or_else(|| answer.to_string())
11586 } else {
11587 output_data.content
11588 };
11589
11590 let (final_content, reflection_metadata) = match self
11592 .run_reflection(&*llm, processed_input, final_content)
11593 .await
11594 {
11595 Ok(r) => r,
11596 Err(e) => {
11597 yield StreamChunk::error(e.to_string());
11598 return;
11599 }
11600 };
11601
11602 let final_content = self.format_response_with_thinking(
11603 thinking_content.as_deref(),
11604 &final_content,
11605 );
11606
11607 if buffered_decision {
11609 yield StreamChunk::content(&final_content);
11610 }
11611
11612 let post_result = match self
11616 .post_loop_processing(processed_input, final_content)
11617 .await
11618 {
11619 Ok(r) => r,
11620 Err(e) => {
11621 yield StreamChunk::error(e.to_string());
11622 return;
11623 }
11624 };
11625
11626 let (final_content, transitioned) = match post_result {
11627 PostLoopResult::NoTransition(content) => (content, false),
11628 PostLoopResult::Transitioned(content) => (content, true),
11629 PostLoopResult::NeedsRedispatch => {
11630 const MAX_REDISPATCH_DEPTH: u32 = 3;
11631 let current_depth = *self.redispatch_depth.read();
11632 let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11633 warn!(
11634 depth = current_depth,
11635 "Post-transition re-dispatch depth limit reached (stream)"
11636 );
11637 let c = String::new();
11638 let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11639 c
11640 } else {
11641 *self.redispatch_depth.write() += 1;
11642 if let Some(context) = self.active_turn_context.write().as_mut() {
11643 context.enter_redispatch();
11644 }
11645 info!(
11646 depth = current_depth + 1,
11647 "Re-dispatching for new state after transition (stream)"
11648 );
11649 let result = self.run_loop_internal(processed_input).await;
11650 *self.redispatch_depth.write() -= 1;
11651 if let Some(context) = self.active_turn_context.write().as_mut() {
11652 context.exit_redispatch();
11653 }
11654 match result {
11655 Ok(resp) => resp.content,
11656 Err(e) => {
11657 yield StreamChunk::error(e.to_string());
11658 return;
11659 }
11660 }
11661 };
11662 (content, true)
11663 }
11664 };
11665
11666 if transitioned {
11667 if include_state_events
11668 && let Some(state) = self.current_state()
11669 {
11670 yield StreamChunk::state_transition(None, state);
11671 }
11672 yield StreamChunk::content(&final_content);
11674 }
11675
11676 let final_response = self.build_agent_response(AgentResponseParts {
11678 content: final_content,
11679 all_tool_calls,
11680 reasoning_mode,
11681 auto_detected,
11682 iterations,
11683 thinking: thinking_content,
11684 reflection_metadata,
11685 });
11686 if let Err(e) = self.finish_turn_if_root(&final_response).await {
11687 yield StreamChunk::error(e.to_string());
11688 return;
11689 }
11690
11691 record_runtime_stream_final(&terminal, final_response);
11692 yield StreamChunk::Done {};
11693 return;
11694 }
11695 })
11696 }
11697
11698 fn run_loop_stream<'a>(
11701 &'a self,
11702 input: &'a str,
11703 terminal: RuntimeStreamTerminalSlot,
11704 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11705 Box::pin(async_stream::stream! {
11706 self.begin_root_turn();
11707 let _root_cleanup = RootTurnCleanup::new(self);
11708 self.hooks.on_message_received(input).await;
11709
11710 if !self.context_initialized.swap(true, Ordering::SeqCst) {
11712 if let Err(e) = self.context_manager.initialize().await {
11713 yield StreamChunk::error(e.to_string());
11714 return;
11715 }
11716 debug!("Context manager initialized (defaults, env, builtins)");
11717 }
11718
11719 if let Err(e) = self.check_turn_timeout().await {
11720 yield StreamChunk::error(e.to_string());
11721 return;
11722 }
11723 if let Err(e) = self.context_manager.refresh_per_turn().await {
11724 yield StreamChunk::error(e.to_string());
11725 return;
11726 }
11727
11728 self.clear_disambiguation_context();
11730
11731 if let Some(ref disambiguator) = self.disambiguation_manager {
11733 let disambiguation_context = match self.build_disambiguation_context().await {
11734 Ok(ctx) => ctx,
11735 Err(e) => {
11736 yield StreamChunk::error(e.to_string());
11737 return;
11738 }
11739 };
11740
11741 let state_override = self
11742 .state_machine
11743 .as_ref()
11744 .and_then(|sm| sm.current_definition())
11745 .and_then(|def| def.disambiguation.clone());
11746
11747 let state_generation = self
11748 .state_machine
11749 .as_ref()
11750 .map(|state_machine| state_machine.generation());
11751 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11752 let mut result = match self
11753 .observe_purpose(
11754 ObservationPurpose::DisambiguationDetection,
11755 disambiguator.process_input_with_override(
11756 input,
11757 &disambiguation_context,
11758 state_override.as_ref(),
11759 None,
11760 ),
11761 )
11762 .await
11763 {
11764 Ok(r) => r,
11765 Err(e) => {
11766 yield StreamChunk::error(e.to_string());
11767 return;
11768 }
11769 };
11770 let current_state_generation = self
11771 .state_machine
11772 .as_ref()
11773 .map(|state_machine| state_machine.generation());
11774 if current_state_generation != state_generation
11775 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
11776 {
11777 disambiguator.clear_pending().await;
11778 *self.pending_skill_id.write() = None;
11779 result = DisambiguationResult::Abandoned { new_input: None };
11780 info!(
11781 confirmation_event = "invalidated",
11782 invalidation_reason = "state_generation_changed",
11783 "Streaming disambiguation result invalidated before redispatch"
11784 );
11785 }
11786 match result {
11787 DisambiguationResult::Clear => {
11788 debug!("Input is clear, proceeding normally (stream)");
11789 }
11790 DisambiguationResult::NeedsClarification {
11791 question,
11792 detection,
11793 } => {
11794 let admission = match self
11795 .admit_disambiguation_redispatch(
11796 disambiguation_epoch,
11797 state_generation,
11798 )
11799 .await
11800 {
11801 Ok(admission) => admission,
11802 Err(error) => {
11803 *self.pending_skill_id.write() = None;
11804 yield StreamChunk::error(error.to_string());
11805 return;
11806 }
11807 };
11808 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
11809 info!(
11810 ambiguity_type = ?detection.ambiguity_type,
11811 confidence = detection.confidence,
11812 "Input requires clarification (stream)"
11813 );
11814 if let Err(e) = self.commit_root_user_message(input).await {
11817 yield StreamChunk::error(e.to_string());
11818 return;
11819 }
11820 let _ = self
11821 .memory
11822 .add_message(ChatMessage::assistant(&question.question))
11823 .await;
11824 let status = if awaiting_confirmation {
11825 "awaiting_confirmation"
11826 } else {
11827 "awaiting_clarification"
11828 };
11829 let response = AgentResponse::new(&question.question).with_metadata(
11830 "disambiguation",
11831 serde_json::json!({ "status": status }),
11832 );
11833 drop(admission);
11834 if let Err(e) = self.finish_turn_if_root(&response).await {
11835 yield StreamChunk::error(e.to_string());
11836 return;
11837 }
11838 yield StreamChunk::content(&question.question);
11839 record_runtime_stream_final(&terminal, response);
11840 yield StreamChunk::Done {};
11841 return;
11842 }
11843 DisambiguationResult::Clarified {
11844 enriched_input,
11845 resolved,
11846 ..
11847 } => {
11848 let admission = match self
11849 .admit_disambiguation_redispatch(
11850 disambiguation_epoch,
11851 state_generation,
11852 )
11853 .await
11854 {
11855 Ok(admission) => admission,
11856 Err(error) => {
11857 *self.pending_skill_id.write() = None;
11858 yield StreamChunk::error(error.to_string());
11859 return;
11860 }
11861 };
11862 info!(
11863 resolved_count = resolved.len(),
11864 enriched = %enriched_input,
11865 "Input clarified (stream)"
11866 );
11867 for (key, value) in &resolved {
11868 let context_key = format!("disambiguation.{}", key);
11869 let _ = self.context_manager.set(&context_key, value.clone());
11870 }
11871 if let Some(intent) = resolved.get("intent") {
11872 let _ = self.context_manager.set("resolved_intent", intent.clone());
11873 }
11874 let _ = self
11875 .context_manager
11876 .set("disambiguation.resolved", serde_json::Value::Bool(true));
11877
11878 let skill_id = self.pending_skill_id.read().clone();
11882 if let Some(skill_id) = skill_id {
11883 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
11884 drop(admission);
11885 match self
11886 .recheck_skill_disambiguation(
11887 &skill_id,
11888 &enriched_input,
11889 disambiguation_epoch,
11890 state_generation,
11891 )
11892 .await
11893 {
11894 Ok(resp) => {
11895 yield StreamChunk::content(&resp.content);
11896 record_runtime_stream_final(&terminal, resp);
11897 yield StreamChunk::Done {};
11898 return;
11899 }
11900 Err(e) => {
11901 yield StreamChunk::error(e.to_string());
11902 return;
11903 }
11904 }
11905 }
11906
11907 drop(admission);
11909 let mut inner = self.run_loop_internal_stream(
11910 &enriched_input,
11911 Arc::clone(&terminal),
11912 );
11913 while let Some(chunk) = inner.next().await {
11914 yield chunk;
11915 }
11916 return;
11917 }
11918 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
11919 info!("Proceeding with best guess (stream)");
11920
11921 let skill_id = self.pending_skill_id.read().clone();
11923 if let Some(skill_id) = skill_id {
11924 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
11925 match self
11926 .recheck_skill_disambiguation(
11927 &skill_id,
11928 &enriched_input,
11929 disambiguation_epoch,
11930 state_generation,
11931 )
11932 .await
11933 {
11934 Ok(resp) => {
11935 yield StreamChunk::content(&resp.content);
11936 record_runtime_stream_final(&terminal, resp);
11937 yield StreamChunk::Done {};
11938 return;
11939 }
11940 Err(e) => {
11941 yield StreamChunk::error(e.to_string());
11942 return;
11943 }
11944 }
11945 }
11946
11947 let mut inner = self.run_loop_internal_stream(
11948 &enriched_input,
11949 Arc::clone(&terminal),
11950 );
11951 while let Some(chunk) = inner.next().await {
11952 yield chunk;
11953 }
11954 return;
11955 }
11956 DisambiguationResult::GiveUp { reason } => {
11957 *self.pending_skill_id.write() = None;
11958 warn!(reason = %reason, "Disambiguation gave up (stream)");
11959 let apology = self
11960 .generate_localized_apology(
11961 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
11962 &reason,
11963 )
11964 .await
11965 .unwrap_or_else(|_| {
11966 format!("I'm sorry, I couldn't understand your request: {}", reason)
11967 });
11968 let response = AgentResponse::new(&apology);
11969 if let Err(e) = self.finish_turn_if_root(&response).await {
11970 yield StreamChunk::error(e.to_string());
11971 return;
11972 }
11973 yield StreamChunk::content(&apology);
11974 record_runtime_stream_final(&terminal, response);
11975 yield StreamChunk::Done {};
11976 return;
11977 }
11978 DisambiguationResult::Escalate { reason } => {
11979 *self.pending_skill_id.write() = None;
11980 info!(reason = %reason, "Escalating to human (stream)");
11981 if let Some(ref hitl) = self.hitl_engine {
11982 let trigger =
11983 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
11984 let mut context_map = HashMap::new();
11985 context_map.insert("original_input".to_string(), serde_json::json!(input));
11986 context_map.insert("reason".to_string(), serde_json::json!(&reason));
11987 let check_result = HITLCheckResult::required(
11988 trigger,
11989 context_map,
11990 format!("User request needs human assistance: {}", reason),
11991 Some(hitl.config().default_timeout_seconds),
11992 );
11993 match self.request_hitl_approval(check_result).await {
11994 Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
11995 let mut inner = self.run_loop_internal_stream(
11996 input,
11997 Arc::clone(&terminal),
11998 );
11999 while let Some(chunk) = inner.next().await {
12000 yield chunk;
12001 }
12002 return;
12003 }
12004 Ok(_) => {}
12005 Err(e) => {
12006 yield StreamChunk::error(e.to_string());
12007 return;
12008 }
12009 }
12010 }
12011 let apology = self
12012 .generate_localized_apology(
12013 "Explain briefly that you're transferring the user to a human agent for help.",
12014 &reason,
12015 )
12016 .await
12017 .unwrap_or_else(|_| {
12018 format!("I need human assistance to help with your request: {}", reason)
12019 });
12020 let response = AgentResponse::new(&apology);
12021 if let Err(e) = self.finish_turn_if_root(&response).await {
12022 yield StreamChunk::error(e.to_string());
12023 return;
12024 }
12025 yield StreamChunk::content(&apology);
12026 record_runtime_stream_final(&terminal, response);
12027 yield StreamChunk::Done {};
12028 return;
12029 }
12030 DisambiguationResult::Abandoned { new_input } => {
12031 *self.pending_skill_id.write() = None;
12032
12033 info!(
12034 has_new_input = new_input.is_some(),
12035 "Clarification abandoned by user (stream)"
12036 );
12037
12038 if let Err(e) = self.commit_root_user_message(input).await {
12039 yield StreamChunk::error(e.to_string());
12040 return;
12041 }
12042
12043 match new_input {
12044 Some(fresh_input) => {
12045 let mut inner = self.run_loop_internal_stream(
12047 &fresh_input,
12048 Arc::clone(&terminal),
12049 );
12050 while let Some(chunk) = inner.next().await {
12051 yield chunk;
12052 }
12053 return;
12054 }
12055 None => {
12056 let ack = self
12058 .generate_localized_apology(
12059 "The user changed their mind about their previous request. \
12060 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
12061 Do NOT apologize excessively. Be concise.",
12062 "User abandoned clarification",
12063 )
12064 .await
12065 .unwrap_or_else(|_| {
12066 "OK, no problem. What else can I help with?".to_string()
12067 });
12068
12069 let _ = self
12070 .memory
12071 .add_message(ChatMessage::assistant(&ack))
12072 .await;
12073
12074 let response = AgentResponse::new(&ack);
12075 if let Err(e) = self.finish_turn_if_root(&response).await {
12076 yield StreamChunk::error(e.to_string());
12077 return;
12078 }
12079 yield StreamChunk::content(&ack);
12080 record_runtime_stream_final(&terminal, response);
12081 yield StreamChunk::Done {};
12082 return;
12083 }
12084 }
12085 }
12086 }
12087 }
12088
12089 let mut inner = self.run_loop_internal_stream(input, Arc::clone(&terminal));
12091 while let Some(chunk) = inner.next().await {
12092 yield chunk;
12093 }
12094 })
12095 }
12096
12097 pub fn info(&self) -> AgentInfo {
12098 self.info.clone()
12099 }
12100
12101 pub fn skills(&self) -> &[SkillDefinition] {
12102 &self.skills
12103 }
12104
12105 async fn reset_runtime_state(&self) -> Result<()> {
12107 let _admission = self.disambiguation_admission.write().await;
12108 if self.state_transition_reserved.load(Ordering::SeqCst) {
12109 return Err(AgentError::Other(
12110 "Cannot reset while a state transition is in progress".to_string(),
12111 ));
12112 }
12113 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
12114 *self.pending_skill_id.write() = None;
12115 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
12116 disambiguator.clear_pending().await;
12117 }
12118 self.memory.clear().await?;
12119 *self.iteration_count.write() = 0;
12120 self.tool_call_history.write().clear();
12121 if let Some(ref sm) = self.state_machine {
12122 sm.reset();
12123 }
12124 Ok(())
12125 }
12126
12127 pub async fn reset(&self) -> Result<()> {
12129 self.reset_runtime_state().await
12130 }
12131
12132 pub fn max_context_tokens(&self) -> u32 {
12133 self.max_context_tokens
12134 }
12135
12136 pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12137 &self.llm_registry
12138 }
12139
12140 pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12141 self.state_machine.as_ref()
12142 }
12143
12144 pub fn context_manager(&self) -> &Arc<ContextManager> {
12145 &self.context_manager
12146 }
12147
12148 pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12149 self.tool_call_history.read().clone()
12150 }
12151
12152 pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12153 self.memory_token_budget.as_ref()
12154 }
12155
12156 pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12157 &self.parallel_tools
12158 }
12159
12160 pub fn streaming_config(&self) -> &StreamingConfig {
12161 &self.streaming
12162 }
12163
12164 pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12165 &self.hooks
12166 }
12167
12168 pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12169 self.hitl_engine.as_ref()
12170 }
12171
12172 pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12173 &self.approval_handler
12174 }
12175
12176 fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12178 let mut ctx = HashMap::new();
12179 for key in &["user.language", "input.detected.language", "language"] {
12180 if let Some(val) = self.context_manager.get(key) {
12181 ctx.insert(key.to_string(), val);
12182 }
12183 }
12184 ctx
12185 }
12186
12187 async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12189 let Some(request) = check_result.into_request() else {
12190 return Ok(ApprovalResult::Approved);
12191 };
12192
12193 self.hooks.on_approval_requested(&request).await;
12194
12195 let timeout = request.timeout;
12196
12197 let raw_result = if let Some(duration) = timeout {
12198 match tokio::time::timeout(
12199 duration,
12200 self.approval_handler.request_approval(request.clone()),
12201 )
12202 .await
12203 {
12204 Ok(result) => result,
12205 Err(_) => ApprovalResult::timeout(),
12206 }
12207 } else {
12208 self.approval_handler
12209 .request_approval(request.clone())
12210 .await
12211 };
12212
12213 self.hooks
12214 .on_approval_result(&request.id, &raw_result)
12215 .await;
12216
12217 let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12218 match &raw_result {
12219 ApprovalResult::Approved => (
12220 ApprovalResolvedOutcome::Approved,
12221 Ok(ApprovalResult::Approved),
12222 ),
12223 ApprovalResult::Rejected { reason } => (
12224 ApprovalResolvedOutcome::Rejected {
12225 reason: reason.clone(),
12226 },
12227 Ok(ApprovalResult::Rejected {
12228 reason: reason.clone(),
12229 }),
12230 ),
12231 ApprovalResult::Modified { changes } => (
12232 ApprovalResolvedOutcome::Modified {
12233 changes: changes.clone(),
12234 },
12235 Ok(ApprovalResult::Modified {
12236 changes: changes.clone(),
12237 }),
12238 ),
12239 ApprovalResult::Timeout => {
12240 if let Some(ref engine) = self.hitl_engine {
12241 match engine.config().on_timeout {
12242 TimeoutAction::Approve => (
12243 ApprovalResolvedOutcome::Approved,
12244 Ok(ApprovalResult::Approved),
12245 ),
12246 TimeoutAction::Reject => {
12247 let reason = Some("Timeout".to_string());
12248 (
12249 ApprovalResolvedOutcome::Rejected {
12250 reason: reason.clone(),
12251 },
12252 Ok(ApprovalResult::Rejected { reason }),
12253 )
12254 }
12255 TimeoutAction::Error => {
12256 let message = "HITL approval timeout".to_string();
12257 (
12258 ApprovalResolvedOutcome::Error {
12259 message: message.clone(),
12260 },
12261 Err(AgentError::Other(message)),
12262 )
12263 }
12264 }
12265 } else {
12266 let reason = Some("Timeout (no engine)".to_string());
12267 (
12268 ApprovalResolvedOutcome::Rejected {
12269 reason: reason.clone(),
12270 },
12271 Ok(ApprovalResult::Rejected { reason }),
12272 )
12273 }
12274 }
12275 };
12276
12277 self.hooks
12278 .on_approval_resolved(&request, &raw_result, &outcome)
12279 .await;
12280
12281 effective_result
12282 }
12283
12284 pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12285 if let Some(ref hitl_engine) = self.hitl_engine {
12286 let hitl_lang_ctx = self.build_hitl_language_context();
12287 let check_result = self
12288 .observe_purpose(
12289 ObservationPurpose::HitlLocalization,
12290 hitl_engine.check_state_transition_with_localization(
12291 from,
12292 to,
12293 &hitl_lang_ctx,
12294 self.approval_handler.as_ref(),
12295 Some(&self.llm_registry),
12296 ),
12297 )
12298 .await?;
12299 if check_result.is_required() {
12300 let result = self.request_hitl_approval(check_result).await?;
12301 return Ok(matches!(
12302 result,
12303 ApprovalResult::Approved | ApprovalResult::Modified { .. }
12304 ));
12305 }
12306 }
12307 Ok(true)
12308 }
12309
12310 async fn execute_tools_parallel(
12312 &self,
12313 tool_calls: &[ToolCall],
12314 ) -> Vec<(String, Result<String>)> {
12315 let can_run_parallel = tool_calls.iter().all(|tc| {
12316 self.tools
12317 .resolve(&tc.name)
12318 .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12319 .unwrap_or(false)
12320 });
12321
12322 if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12323 let mut results = Vec::new();
12324 for tc in tool_calls {
12325 let result = self
12326 .observe_purpose(
12327 current_observation_context()
12328 .map(|context| context.purpose)
12329 .unwrap_or_default(),
12330 self.execute_tool_smart(tc),
12331 )
12332 .await;
12333 results.push((tc.id.clone(), result));
12334 }
12335 return results;
12336 }
12337
12338 let chunks: Vec<_> = tool_calls
12339 .chunks(self.parallel_tools.max_parallel)
12340 .collect();
12341
12342 let mut all_results = Vec::new();
12343
12344 for chunk in chunks {
12345 let futures: Vec<_> = chunk
12346 .iter()
12347 .map(|tc| {
12348 let tc = tc.clone();
12349 async move {
12350 let result = self.execute_tool_smart(&tc).await;
12351 (tc.id.clone(), result)
12352 }
12353 })
12354 .collect();
12355
12356 let results = futures::future::join_all(futures).await;
12357 all_results.extend(results);
12358 }
12359
12360 all_results
12361 }
12362
12363 pub async fn chat_stream<'a>(
12367 &'a self,
12368 input: &'a str,
12369 ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12370 let RootTurnAdmission {
12371 guard: root_turn_guard,
12372 identity_stack,
12373 } = self.acquire_root_turn().await?;
12374 scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12378 info!(input_len = input.len(), "Starting streaming chat");
12379 let terminal = new_runtime_stream_terminal_slot();
12380 let inner = self.run_loop_stream(input, terminal);
12381 let observation_context = self.build_observation_context(None);
12382 let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12383 Box::pin(async_stream::stream! {
12384 let mut root_turn_guard = Some(root_turn_guard);
12385 let mut inner = inner;
12386 loop {
12387 let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12388 if let Some(context) = observation_context.as_ref() {
12389 with_observation_context(context.clone(), inner.next()).await
12390 } else {
12391 inner.next().await
12392 }
12393 })
12394 .await;
12395 match next {
12396 Some(StreamChunk::Done {}) => {
12397 while scope_runtime_gate_identity_stack(&identity_stack, async {
12398 if let Some(context) = observation_context.as_ref() {
12399 with_observation_context(context.clone(), inner.next())
12400 .await
12401 .is_some()
12402 } else {
12403 inner.next().await.is_some()
12404 }
12405 })
12406 .await
12407 {}
12408 if observation_context.is_some() {
12409 scope_runtime_gate_identity_stack(
12410 &identity_stack,
12411 self.export_observability_if_configured(),
12412 )
12413 .await;
12414 }
12415 drop(root_turn_guard.take());
12416 yield StreamChunk::Done {};
12417 return;
12418 }
12419 Some(chunk) => yield chunk,
12420 None => {
12421 if observation_context.is_some() {
12422 scope_runtime_gate_identity_stack(
12423 &identity_stack,
12424 self.export_observability_if_configured(),
12425 )
12426 .await;
12427 }
12428 drop(root_turn_guard.take());
12429 return;
12430 }
12431 }
12432 }
12433 });
12434 Ok(stream)
12435 }
12436
12437 pub async fn chat_stream_events<'a>(
12441 &'a self,
12442 input: &'a str,
12443 ) -> Result<Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>>> {
12444 let RootTurnAdmission {
12445 guard: root_turn_guard,
12446 identity_stack,
12447 } = self.acquire_root_turn().await?;
12448 scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12452 info!(input_len = input.len(), "Starting streaming chat events");
12453 let terminal = new_runtime_stream_terminal_slot();
12454 let mut inner = self.run_loop_stream(input, Arc::clone(&terminal));
12455 let observation_context = self.build_observation_context(None);
12456 let stream: Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>> =
12457 Box::pin(async_stream::stream! {
12458 let mut root_turn_guard = Some(root_turn_guard);
12459 loop {
12460 let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12461 if let Some(context) = observation_context.as_ref() {
12462 with_observation_context(context.clone(), inner.next()).await
12463 } else {
12464 inner.next().await
12465 }
12466 })
12467 .await;
12468 match next {
12469 Some(StreamChunk::Done {}) => {
12470 let terminal_event = { terminal.write().take() };
12471 if let Some(response) = terminal_event {
12472 while scope_runtime_gate_identity_stack(&identity_stack, async {
12473 if let Some(context) = observation_context.as_ref() {
12474 with_observation_context(context.clone(), inner.next())
12475 .await
12476 .is_some()
12477 } else {
12478 inner.next().await.is_some()
12479 }
12480 })
12481 .await
12482 {}
12483 if observation_context.is_some() {
12484 scope_runtime_gate_identity_stack(
12485 &identity_stack,
12486 self.export_observability_if_configured(),
12487 )
12488 .await;
12489 }
12490 drop(root_turn_guard.take());
12491 yield AgentStreamEvent::Final(response);
12492 return;
12493 }
12494 }
12495 Some(StreamChunk::Error { message }) => {
12496 let finalized = { terminal.read().is_some() };
12497 if finalized {
12498 continue;
12499 }
12500 while scope_runtime_gate_identity_stack(&identity_stack, async {
12501 if let Some(context) = observation_context.as_ref() {
12502 with_observation_context(context.clone(), inner.next())
12503 .await
12504 .is_some()
12505 } else {
12506 inner.next().await.is_some()
12507 }
12508 })
12509 .await
12510 {}
12511 if observation_context.is_some() {
12512 scope_runtime_gate_identity_stack(
12513 &identity_stack,
12514 self.export_observability_if_configured(),
12515 )
12516 .await;
12517 }
12518 drop(root_turn_guard.take());
12519 yield AgentStreamEvent::Chunk(StreamChunk::Error { message });
12520 return;
12521 }
12522 Some(chunk) => yield AgentStreamEvent::Chunk(chunk),
12523 None => {
12524 if observation_context.is_some() {
12525 scope_runtime_gate_identity_stack(
12526 &identity_stack,
12527 self.export_observability_if_configured(),
12528 )
12529 .await;
12530 }
12531 drop(root_turn_guard.take());
12532 return;
12533 }
12534 }
12535 }
12536 });
12537 Ok(stream)
12538 }
12539}
12540
12541#[async_trait]
12542impl ToolInvoker for RuntimeAgent {
12543 async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12544 self.execute_tool_record(request).await
12545 }
12546}
12547
12548#[async_trait]
12549impl Agent for RuntimeAgent {
12550 async fn chat(&self, input: &str) -> Result<AgentResponse> {
12552 let RootTurnAdmission {
12553 guard,
12554 identity_stack,
12555 } = self.acquire_root_turn().await?;
12556 let result = scope_runtime_gate_identity_stack(&identity_stack, async {
12557 let result = if let Some(context) = self.build_observation_context(None) {
12558 with_observation_context(context, self.run_loop(input)).await
12559 } else {
12560 self.run_loop(input).await
12561 };
12562 self.export_observability_if_configured().await;
12563 result
12564 })
12565 .await;
12566 drop(guard);
12567 result
12568 }
12569
12570 fn info(&self) -> AgentInfo {
12571 self.info.clone()
12572 }
12573
12574 async fn reset(&self) -> Result<()> {
12576 self.reset_runtime_state().await
12577 }
12578}
12579
12580fn background_maintenance_tags(
12590 label: &str,
12591 stage: &str,
12592 reason: Option<&str>,
12593 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12594) -> HashMap<String, String> {
12595 let mut tags = HashMap::new();
12596 tags.insert("runtime.background".to_string(), "true".to_string());
12597 tags.insert("runtime.maintenance".to_string(), label.to_string());
12598 tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12599 if let Some(policy) = policy {
12600 tags.insert(
12601 "runtime.await_before_next_turn".to_string(),
12602 await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12603 );
12604 tags.insert(
12605 "runtime.maintenance_mode".to_string(),
12606 maintenance_mode_label(policy.mode).to_string(),
12607 );
12608 }
12609 if let Some(reason) = reason {
12610 tags.insert("runtime.reason".to_string(), reason.to_string());
12611 }
12612 tags
12613}
12614
12615fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12616 match policy {
12617 AwaitBeforeNextTurn::Never => "never",
12618 AwaitBeforeNextTurn::SameActor => "same_actor",
12619 AwaitBeforeNextTurn::Always => "always",
12620 }
12621}
12622
12623fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12624 match mode {
12625 MaintenanceMode::InlineSerial => "inline_serial",
12626 MaintenanceMode::InlineParallel => "inline_parallel",
12627 MaintenanceMode::Background => "background",
12628 }
12629}
12630
12631fn record_background_maintenance_event(
12633 manager: Option<&Arc<ObservabilityManager>>,
12634 label: &str,
12635 status: EventStatus,
12636 duration_ms: u64,
12637 stage: &str,
12638 reason: Option<String>,
12639 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12640) {
12641 if let Some(manager) = manager {
12642 manager.record_lifecycle_event(
12643 EventType::MemoryOperation {
12644 operation: format!("{}_background_{}", label, stage),
12645 },
12646 ObservationPurpose::Other(format!("{}_maintenance", label)),
12647 status,
12648 duration_ms,
12649 background_maintenance_tags(label, stage, reason.as_deref(), policy),
12650 None,
12651 );
12652 }
12653}
12654
12655fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12656 if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12657 MaintenanceMode::InlineParallel
12658 } else {
12659 mode
12660 }
12661}
12662
12663fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12664 match hint {
12665 ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12666 ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12667 ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12668 ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12669 ObservationPurpose::ProcessTransform
12670 }
12671 }
12672}
12673
12674fn new_tool_resource_locks() -> ToolResourceLocks {
12675 Arc::new(RwLock::new(HashMap::new()))
12676}
12677
12678fn tool_resource_lock_keys(
12683 _canonical_id: &str,
12684 args: &Value,
12685 bindings: &ai_agents_core::ToolPolicyBindings,
12686 classification: &ai_agents_core::ToolCallClassification,
12687) -> Vec<String> {
12688 if classification.concurrency_safe {
12689 return Vec::new();
12690 }
12691
12692 let mut keys = Vec::new();
12693 let mut has_path_resource = false;
12694 for binding in &bindings.path_fields {
12695 let value = value_at_argument_path(args, &binding.field)
12696 .cloned()
12697 .or_else(|| {
12698 binding
12699 .default_path
12700 .as_ref()
12701 .map(|path| Value::String(path.clone()))
12702 });
12703 if let Some(value) = value {
12704 collect_resource_strings(&value, |_| {
12705 has_path_resource = true;
12706 });
12707 }
12708 }
12709 for binding in &bindings.domain_fields {
12710 if let Some(value) = value_at_argument_path(args, &binding.field) {
12711 collect_resource_strings(value, |domain| {
12712 let normalized = if binding.is_url {
12713 normalized_url_resource_key(domain)
12714 } else {
12715 domain.trim().trim_end_matches('.').to_ascii_lowercase()
12716 };
12717 keys.push(format!("domain:{}", normalized));
12718 });
12719 }
12720 }
12721 for binding in &bindings.command_fields {
12722 if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12723 continue;
12724 }
12725 if let Some(value) = value_at_argument_path(args, &binding.field) {
12726 collect_resource_strings(value, |_| {
12727 has_path_resource = true;
12728 });
12729 }
12730 }
12731 if has_path_resource {
12732 keys.push("path-mutation:global".to_string());
12733 }
12734 if keys.is_empty() {
12735 keys.push("side-effect:unbound".to_string());
12736 }
12737 keys.sort();
12738 keys.dedup();
12739 keys
12740}
12741
12742fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12743 let mut current = value;
12744 for segment in field.split('.') {
12745 if segment.is_empty() {
12746 return None;
12747 }
12748 current = current.get(segment)?;
12749 }
12750 Some(current)
12751}
12752
12753fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12754 match value {
12755 Value::String(value) => collect(value),
12756 Value::Array(values) => {
12757 for value in values {
12758 if let Some(value) = value.as_str() {
12759 collect(value);
12760 }
12761 }
12762 }
12763 _ => {}
12764 }
12765}
12766
12767fn normalized_url_resource_key(value: &str) -> String {
12768 let value = value.trim();
12769 let Some((scheme, remainder)) = value.split_once("://") else {
12770 return value.to_ascii_lowercase();
12771 };
12772 let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
12773 let (authority, suffix) = remainder.split_at(authority_end);
12774 format!(
12775 "{}://{}{}",
12776 scheme.to_ascii_lowercase(),
12777 authority.to_ascii_lowercase(),
12778 suffix
12779 )
12780}
12781
12782fn render_concurrent_template(
12783 template: &str,
12784 user_input: &str,
12785 context_values: &std::collections::HashMap<String, serde_json::Value>,
12786) -> Result<String> {
12787 let mut env = minijinja::Environment::new();
12788 env.add_template("concurrent", template)
12789 .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
12790
12791 let mut ctx = std::collections::BTreeMap::new();
12792 ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
12793
12794 let context_obj = minijinja::Value::from_serialize(context_values);
12796 ctx.insert("context".to_string(), context_obj);
12797
12798 let tmpl = env
12799 .get_template("concurrent")
12800 .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
12801
12802 tmpl.render(minijinja::Value::from_serialize(&ctx))
12803 .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
12804}
12805
12806#[cfg(test)]
12807mod tests {
12808 use super::*;
12809 use crate::AgentBuilder;
12810 use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
12811 use ai_agents_llm::mock::MockLLMProvider;
12812 use ai_agents_skills::{SkillDefinition, SkillStep};
12813 use ai_agents_tools::{
12814 CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool,
12815 WebFetchResolver, WebFetchTool, WebFetchTransport, WebFetchTransportRequest,
12816 WebFetchTransportResponse,
12817 };
12818
12819 fn mock_with_response(response: &str) -> MockLLMProvider {
12820 let mut mock = MockLLMProvider::new("test");
12821 mock.set_response(response);
12822 mock
12823 }
12824
12825 fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
12826 let mut mock = MockLLMProvider::new("test");
12827 mock.set_responses(responses.into_iter().map(String::from).collect(), true);
12828 mock
12829 }
12830
12831 fn disambiguation_state_machine(
12833 state_enabled: Option<bool>,
12834 require_confirmation: bool,
12835 ) -> Arc<StateMachine> {
12836 let definition = ai_agents_state::StateDefinition {
12837 prompt: Some("Handle the resolved request.".to_string()),
12838 disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
12839 enabled: state_enabled,
12840 require_confirmation,
12841 ..Default::default()
12842 }),
12843 ..Default::default()
12844 };
12845 let review = ai_agents_state::StateDefinition {
12846 prompt: Some("Review a fresh request.".to_string()),
12847 ..Default::default()
12848 };
12849 Arc::new(
12850 StateMachine::new(ai_agents_state::StateConfig {
12851 initial: "active".to_string(),
12852 states: std::collections::HashMap::from([
12853 ("active".to_string(), definition),
12854 ("review".to_string(), review),
12855 ]),
12856 global_transitions: Vec::new(),
12857 fallback: None,
12858 max_no_transition: None,
12859 regenerate_on_transition: true,
12860 })
12861 .unwrap(),
12862 )
12863 }
12864
12865 fn state_disambiguation_agent(
12867 responses: Vec<&str>,
12868 manager_enabled: bool,
12869 state_enabled: Option<bool>,
12870 require_confirmation: bool,
12871 ) -> (RuntimeAgent, MockLLMProvider) {
12872 state_disambiguation_agent_with_skills(
12873 responses,
12874 manager_enabled,
12875 state_enabled,
12876 require_confirmation,
12877 Vec::new(),
12878 )
12879 }
12880
12881 fn state_disambiguation_agent_with_skills(
12883 responses: Vec<&str>,
12884 manager_enabled: bool,
12885 state_enabled: Option<bool>,
12886 require_confirmation: bool,
12887 skills: Vec<SkillDefinition>,
12888 ) -> (RuntimeAgent, MockLLMProvider) {
12889 let mut mock = MockLLMProvider::new("state-confirmation");
12890 mock.set_responses(responses.into_iter().map(String::from).collect(), false);
12891 let observed = mock.clone();
12892 let agent = AgentBuilder::new()
12893 .system_prompt("Handle requests.")
12894 .llm(Arc::new(mock.clone()))
12895 .llm_alias("router", Arc::new(mock))
12896 .state_machine(disambiguation_state_machine(
12897 state_enabled,
12898 require_confirmation,
12899 ))
12900 .skills(skills)
12901 .build()
12902 .unwrap()
12903 .with_disambiguation(DisambiguationConfig {
12904 enabled: manager_enabled,
12905 ..Default::default()
12906 });
12907 (agent, observed)
12908 }
12909
12910 fn confirmation_skill() -> SkillDefinition {
12912 SkillDefinition {
12913 id: "send_report".to_string(),
12914 description: "Send a report after clarification".to_string(),
12915 trigger: "When the user asks to send a report".to_string(),
12916 steps: vec![SkillStep::Prompt {
12917 prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
12918 llm: None,
12919 }],
12920 reasoning: None,
12921 reflection: None,
12922 disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
12923 enabled: Some(true),
12924 ..Default::default()
12925 }),
12926 }
12927 }
12928
12929 fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
12931 observed
12932 .call_history()
12933 .iter()
12934 .filter(|call| {
12935 call.messages
12936 .iter()
12937 .any(|message| message.content.contains("Execute confirmed report skill"))
12938 })
12939 .count()
12940 }
12941
12942 struct BlockingRuntimeConfirmationObserver {
12943 entered: tokio::sync::Barrier,
12944 release: tokio::sync::Notify,
12945 }
12946
12947 impl BlockingRuntimeConfirmationObserver {
12948 fn new() -> Self {
12949 Self {
12950 entered: tokio::sync::Barrier::new(2),
12951 release: tokio::sync::Notify::new(),
12952 }
12953 }
12954 }
12955
12956 struct ResetOnTransitionHooks {
12957 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
12958 invoked: AtomicBool,
12959 }
12960
12961 #[async_trait]
12962 impl AgentHooks for ResetOnTransitionHooks {
12963 async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
12964 if self.invoked.swap(true, Ordering::SeqCst) {
12965 return;
12966 }
12967 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
12968 if let Some(agent) = agent {
12969 agent.reset().await.unwrap();
12970 }
12971 }
12972 }
12973
12974 impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
12975 fn observe_question<'a>(
12976 &'a self,
12977 future: ClarificationQuestionFuture<'a>,
12978 ) -> ClarificationQuestionFuture<'a> {
12979 future
12980 }
12981
12982 fn observe_parse<'a>(
12983 &'a self,
12984 future: ClarificationParseFuture<'a>,
12985 ) -> ClarificationParseFuture<'a> {
12986 future
12987 }
12988
12989 fn observe_confirmation_parse<'a>(
12990 &'a self,
12991 future: ConfirmationParseFuture<'a>,
12992 ) -> ConfirmationParseFuture<'a> {
12993 Box::pin(async move {
12994 self.entered.wait().await;
12995 self.release.notified().await;
12996 future.await
12997 })
12998 }
12999 }
13000
13001 #[tokio::test]
13002 async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
13003 let (agent, observed) = state_disambiguation_agent(
13004 vec![
13005 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13006 r#"{"question":"What should I send?","options":null}"#,
13007 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13008 r#"{"question":"Should I send the report to Ada?"}"#,
13009 r#"{"status":"confirmed"}"#,
13010 "Request executed.",
13011 ],
13012 true,
13013 None,
13014 true,
13015 );
13016
13017 let clarification = agent.chat("Send it").await.unwrap();
13018 assert_eq!(clarification.content, "What should I send?");
13019 assert_eq!(observed.call_count(), 2);
13020
13021 let confirmation = agent.chat("The report to Ada").await.unwrap();
13022 assert_eq!(confirmation.content, "Should I send the report to Ada?");
13023 assert_eq!(
13024 confirmation
13025 .metadata
13026 .as_ref()
13027 .and_then(|metadata| metadata.get("disambiguation"))
13028 .and_then(|metadata| metadata.get("status"))
13029 .and_then(Value::as_str),
13030 Some("awaiting_confirmation")
13031 );
13032 assert_eq!(observed.call_count(), 4);
13033
13034 let completed = agent.chat("Yes").await.unwrap();
13035 assert_eq!(completed.content, "Request executed.");
13036 assert_eq!(observed.call_count(), 6);
13037 }
13038
13039 #[tokio::test]
13040 async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
13041 let (agent, observed) = state_disambiguation_agent(
13042 vec![
13043 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13044 r#"{"question":"What should I send?","options":null}"#,
13045 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13046 r#"{"question":"Should I send the report to Ada?"}"#,
13047 r#"{"status":"confirmed"}"#,
13048 "Request executed.",
13049 ],
13050 true,
13051 None,
13052 true,
13053 );
13054
13055 let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
13056 let mut clarification = String::new();
13057 while let Some(chunk) = clarification_stream.next().await {
13058 match chunk {
13059 StreamChunk::Content { text } => clarification.push_str(&text),
13060 StreamChunk::Done {} => break,
13061 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13062 _ => {}
13063 }
13064 }
13065 assert_eq!(clarification, "What should I send?");
13066 assert_eq!(observed.call_count(), 2);
13067
13068 let mut confirmation_stream = agent.chat_stream_events("The report to Ada").await.unwrap();
13069 let mut confirmation = None;
13070 while let Some(event) = confirmation_stream.next().await {
13071 match event {
13072 AgentStreamEvent::Final(response) => confirmation = Some(response),
13073 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
13074 panic!("unexpected stream error: {message}")
13075 }
13076 AgentStreamEvent::Chunk(_) => {}
13077 }
13078 }
13079 let confirmation = confirmation.expect("confirmation must finalize");
13080 assert_eq!(confirmation.content, "Should I send the report to Ada?");
13081 assert_eq!(
13082 confirmation
13083 .metadata
13084 .as_ref()
13085 .and_then(|metadata| metadata.get("disambiguation"))
13086 .and_then(|metadata| metadata.get("status"))
13087 .and_then(Value::as_str),
13088 Some("awaiting_confirmation")
13089 );
13090 assert_eq!(observed.call_count(), 4);
13091
13092 let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
13093 let mut completed = String::new();
13094 while let Some(chunk) = completed_stream.next().await {
13095 match chunk {
13096 StreamChunk::Content { text } => completed.push_str(&text),
13097 StreamChunk::Done {} => break,
13098 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13099 _ => {}
13100 }
13101 }
13102 assert_eq!(completed, "Request executed.");
13103 assert_eq!(observed.call_count(), 6);
13104 }
13105
13106 #[tokio::test]
13108 async fn root_turn_gate_serializes_blocking_and_streaming_entry_points() {
13109 let (complete_entered, mut complete_events) = tokio::sync::mpsc::unbounded_channel();
13110 let agent = Arc::new(
13111 AgentBuilder::new()
13112 .system_prompt("Serialize root turns.")
13113 .llm(Arc::new(RootTurnProbeProvider { complete_entered }))
13114 .build()
13115 .unwrap(),
13116 );
13117 let blocking_agent = Arc::clone(&agent);
13118
13119 let legacy_stream = agent.chat_stream("stream owner").await.unwrap();
13120 assert!(agent.root_turn_gate.try_lock().is_err());
13121 let blocking = tokio::spawn(async move { blocking_agent.chat("blocked").await.unwrap() });
13122 assert!(
13123 tokio::time::timeout(std::time::Duration::from_millis(50), complete_events.recv())
13124 .await
13125 .is_err(),
13126 "blocking turn reached the provider while the legacy stream owned the root gate"
13127 );
13128
13129 drop(legacy_stream);
13130 assert_eq!(
13131 tokio::time::timeout(std::time::Duration::from_secs(2), complete_events.recv())
13132 .await
13133 .expect("blocking turn did not enter after stream drop"),
13134 Some(())
13135 );
13136 let response = tokio::time::timeout(std::time::Duration::from_secs(2), blocking)
13137 .await
13138 .expect("blocking turn did not finish after stream drop")
13139 .unwrap();
13140 assert_eq!(response.content, "blocking complete");
13141
13142 let mut event_stream = agent.chat_stream_events("event terminal").await.unwrap();
13143 assert!(agent.root_turn_gate.try_lock().is_err());
13144 let mut saw_final = false;
13145 while let Some(event) = event_stream.next().await {
13146 if matches!(event, AgentStreamEvent::Final(_)) {
13147 saw_final = true;
13148 break;
13149 }
13150 }
13151 assert!(saw_final);
13152 assert!(
13153 agent.root_turn_gate.try_lock().is_ok(),
13154 "authoritative terminal event retained the root gate"
13155 );
13156 }
13157
13158 #[tokio::test]
13160 async fn response_hook_rejects_same_runtime_chat_reentry() {
13161 let hooks = Arc::new(ResponseChatHooks {
13162 target: parking_lot::Mutex::new(None),
13163 invoked: AtomicBool::new(false),
13164 nested_result: parking_lot::Mutex::new(None),
13165 });
13166 let agent = Arc::new(
13167 AgentBuilder::new()
13168 .system_prompt("Reject response hook reentry.")
13169 .llm(Arc::new(mock_with_response("outer response")))
13170 .hooks(hooks.clone())
13171 .build()
13172 .unwrap(),
13173 );
13174 *hooks.target.lock() = Some(Arc::downgrade(&agent));
13175
13176 let response = tokio::time::timeout(
13177 std::time::Duration::from_secs(2),
13178 agent.chat("outer request"),
13179 )
13180 .await
13181 .expect("same-runtime response hook reentry must fail without deadlocking")
13182 .unwrap();
13183
13184 assert_eq!(response.content, "outer response");
13185 let nested_result = hooks
13186 .nested_result
13187 .lock()
13188 .clone()
13189 .expect("response hook must record its nested call");
13190 let error = nested_result.expect_err("same-runtime nested chat must be rejected");
13191 assert!(error.contains("reentrant root turn ownership"));
13192 }
13193
13194 #[tokio::test]
13196 async fn root_turn_gate_allows_nested_runtime_and_rejects_cycles() {
13197 let agent_a = AgentBuilder::new()
13198 .system_prompt("Runtime A.")
13199 .llm(Arc::new(mock_with_response("response A")))
13200 .build()
13201 .unwrap();
13202 let agent_b = AgentBuilder::new()
13203 .system_prompt("Runtime B.")
13204 .llm(Arc::new(mock_with_response("response B")))
13205 .build()
13206 .unwrap();
13207 let RootTurnAdmission {
13208 guard: guard_a,
13209 identity_stack: stack_a,
13210 } = agent_a.acquire_root_turn().await.unwrap();
13211
13212 let cycle_error = scope_runtime_gate_identity_stack(&stack_a, async {
13213 let RootTurnAdmission {
13214 guard: guard_b,
13215 identity_stack: stack_b,
13216 } = agent_b
13217 .acquire_root_turn()
13218 .await
13219 .expect("runtime B must acquire a different gate");
13220 let result =
13221 scope_runtime_gate_identity_stack(&stack_b, agent_a.acquire_root_turn()).await;
13222 drop(guard_b);
13223 match result {
13224 Err(error) => error,
13225 Ok(_) => panic!("runtime A accepted a repeated gate identity"),
13226 }
13227 })
13228 .await;
13229 drop(guard_a);
13230
13231 assert!(
13232 cycle_error
13233 .to_string()
13234 .contains("reentrant root turn ownership")
13235 );
13236 }
13237
13238 #[tokio::test]
13240 async fn concurrent_orchestration_propagates_root_gate_ancestry() {
13241 let registry = Arc::new(crate::spawner::AgentRegistry::new());
13242 let hooks_a = Arc::new(ConcurrentResponseHooks {
13243 registry: Arc::downgrade(®istry),
13244 child_id: "runtime-b".to_string(),
13245 invoked: AtomicBool::new(false),
13246 nested_result: parking_lot::Mutex::new(None),
13247 });
13248 let hooks_b = Arc::new(ResponseChatHooks {
13249 target: parking_lot::Mutex::new(None),
13250 invoked: AtomicBool::new(false),
13251 nested_result: parking_lot::Mutex::new(None),
13252 });
13253 let agent_a = AgentBuilder::new()
13254 .system_prompt("Runtime A dispatches runtime B concurrently.")
13255 .llm(Arc::new(mock_with_response("response A")))
13256 .hooks(hooks_a.clone())
13257 .build()
13258 .unwrap();
13259 let agent_b = AgentBuilder::new()
13260 .system_prompt("Runtime B attempts to re-enter runtime A.")
13261 .llm(Arc::new(mock_with_response("response B")))
13262 .hooks(hooks_b.clone())
13263 .build()
13264 .unwrap();
13265 let spec_a = crate::spec::AgentSpec {
13266 name: "runtime-a".to_string(),
13267 system_prompt: "Runtime A dispatches runtime B concurrently.".to_string(),
13268 ..crate::spec::AgentSpec::default()
13269 };
13270 let spec_b = crate::spec::AgentSpec {
13271 name: "runtime-b".to_string(),
13272 system_prompt: "Runtime B attempts to re-enter runtime A.".to_string(),
13273 ..crate::spec::AgentSpec::default()
13274 };
13275 registry
13276 .register(crate::spawner::SpawnedAgent::from_runtime(
13277 "runtime-a".to_string(),
13278 agent_a,
13279 spec_a,
13280 ))
13281 .await
13282 .unwrap();
13283 registry
13284 .register(crate::spawner::SpawnedAgent::from_runtime(
13285 "runtime-b".to_string(),
13286 agent_b,
13287 spec_b,
13288 ))
13289 .await
13290 .unwrap();
13291 let runtime_a = registry.get("runtime-a").unwrap();
13292 *hooks_b.target.lock() = Some(Arc::downgrade(&runtime_a));
13293
13294 let response = tokio::time::timeout(
13295 std::time::Duration::from_secs(2),
13296 runtime_a.chat("outer concurrent request"),
13297 )
13298 .await
13299 .expect("concurrent orchestration cycle must fail without deadlocking")
13300 .unwrap();
13301
13302 assert_eq!(response.content, "response A");
13303 let child_result = hooks_a
13304 .nested_result
13305 .lock()
13306 .clone()
13307 .expect("runtime A hook must record runtime B completion");
13308 assert_eq!(child_result.unwrap(), "response B");
13309 let cycle_result = hooks_b
13310 .nested_result
13311 .lock()
13312 .clone()
13313 .expect("runtime B hook must record runtime A reentry");
13314 assert!(
13315 cycle_result
13316 .expect_err("runtime A accepted a repeated gate identity")
13317 .contains("reentrant root turn ownership")
13318 );
13319 }
13320
13321 #[tokio::test]
13323 async fn confirmed_skill_route_executes_exactly_once() {
13324 let (agent, observed) = state_disambiguation_agent_with_skills(
13325 vec![
13326 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13327 "send_report",
13328 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13329 r#"{"question":"What should I send?","options":null}"#,
13330 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13331 r#"{"question":"Should I send the report to Ada?"}"#,
13332 r#"{"status":"confirmed"}"#,
13333 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
13334 "Report skill executed.",
13335 ],
13336 true,
13337 None,
13338 true,
13339 vec![confirmation_skill()],
13340 );
13341
13342 let clarification = agent.chat("Send it").await.unwrap();
13343 assert_eq!(clarification.content, "What should I send?");
13344 assert_eq!(confirmation_skill_call_count(&observed), 0);
13345
13346 let confirmation = agent.chat("The report to Ada").await.unwrap();
13347 assert_eq!(confirmation.content, "Should I send the report to Ada?");
13348 assert_eq!(
13349 confirmation
13350 .metadata
13351 .as_ref()
13352 .and_then(|metadata| metadata.get("disambiguation"))
13353 .and_then(|metadata| metadata.get("status"))
13354 .and_then(Value::as_str),
13355 Some("awaiting_confirmation")
13356 );
13357 assert_eq!(confirmation_skill_call_count(&observed), 0);
13358
13359 let completed = agent.chat("Yes").await.unwrap();
13360 assert_eq!(completed.content, "Report skill executed.");
13361 assert_eq!(confirmation_skill_call_count(&observed), 1);
13362 assert!(agent.pending_skill_id.read().is_none());
13363 let messages = agent.memory.get_messages(None).await.unwrap();
13364 assert!(!messages.iter().any(|message| message.content == "Yes"));
13365 }
13366
13367 #[tokio::test]
13369 async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
13370 let (agent, observed) = state_disambiguation_agent_with_skills(
13371 vec![
13372 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13373 "send_report",
13374 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13375 r#"{"question":"What should I send?","options":null}"#,
13376 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13377 r#"{"question":"Should I send the report to Ada?"}"#,
13378 r#"{"status":"confirmed"}"#,
13379 r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
13380 r#"{"question":"When should I send it?","options":null}"#,
13381 ],
13382 true,
13383 None,
13384 true,
13385 vec![confirmation_skill()],
13386 );
13387
13388 agent.chat("Send it").await.unwrap();
13389 agent.chat("The report to Ada").await.unwrap();
13390 let follow_up = agent.chat("Yes").await.unwrap();
13391
13392 assert_eq!(follow_up.content, "When should I send it?");
13393 let metadata = follow_up
13394 .metadata
13395 .as_ref()
13396 .and_then(|metadata| metadata.get("disambiguation"))
13397 .unwrap();
13398 assert_eq!(
13399 metadata.get("status").and_then(Value::as_str),
13400 Some("awaiting_clarification")
13401 );
13402 assert_eq!(
13403 metadata.get("skill_id").and_then(Value::as_str),
13404 Some("send_report")
13405 );
13406 assert!(metadata.get("detection").is_some());
13407 assert_eq!(confirmation_skill_call_count(&observed), 0);
13408 }
13409
13410 #[tokio::test]
13412 async fn rejected_skill_confirmation_never_executes() {
13413 let (agent, observed) = state_disambiguation_agent_with_skills(
13414 vec![
13415 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13416 "send_report",
13417 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13418 r#"{"question":"What should I send?","options":null}"#,
13419 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13420 r#"{"question":"Should I send the report to Ada?"}"#,
13421 r#"{"status":"rejected"}"#,
13422 "Confirmation rejected.",
13423 ],
13424 true,
13425 None,
13426 true,
13427 vec![confirmation_skill()],
13428 );
13429
13430 agent.chat("Send it").await.unwrap();
13431 agent.chat("The report to Ada").await.unwrap();
13432 let rejected = agent.chat("No").await.unwrap();
13433
13434 assert_eq!(rejected.content, "Confirmation rejected.");
13435 assert_eq!(confirmation_skill_call_count(&observed), 0);
13436 assert!(agent.pending_skill_id.read().is_none());
13437 }
13438
13439 #[tokio::test]
13441 async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
13442 let (agent, observed) = state_disambiguation_agent_with_skills(
13443 vec![
13444 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13445 "send_report",
13446 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13447 r#"{"question":"What should I send?","options":null}"#,
13448 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13449 r#"{"question":"Should I send the report to Ada?"}"#,
13450 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13451 "none",
13452 "Fresh response.",
13453 ],
13454 true,
13455 None,
13456 true,
13457 vec![confirmation_skill()],
13458 );
13459
13460 agent.chat("Send it").await.unwrap();
13461 agent.chat("The report to Ada").await.unwrap();
13462 agent.reset().await.unwrap();
13463 assert!(agent.pending_skill_id.read().is_none());
13464 assert!(
13465 !agent
13466 .disambiguation_manager()
13467 .unwrap()
13468 .has_pending_clarification()
13469 .await
13470 );
13471
13472 let mut stream = agent.chat_stream("Yes").await.unwrap();
13473 let mut content = String::new();
13474 while let Some(chunk) = stream.next().await {
13475 match chunk {
13476 StreamChunk::Content { text } => content.push_str(&text),
13477 StreamChunk::Done {} => break,
13478 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13479 _ => {}
13480 }
13481 }
13482
13483 assert_eq!(content, "Fresh response.");
13484 assert_eq!(confirmation_skill_call_count(&observed), 0);
13485 }
13486
13487 #[tokio::test]
13489 async fn trait_reset_clears_pending_skill_confirmation() {
13490 let (agent, _) = state_disambiguation_agent_with_skills(
13491 vec![
13492 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13493 "send_report",
13494 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13495 r#"{"question":"What should I send?","options":null}"#,
13496 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13497 r#"{"question":"Should I send the report to Ada?"}"#,
13498 ],
13499 true,
13500 None,
13501 true,
13502 vec![confirmation_skill()],
13503 );
13504
13505 agent.chat("Send it").await.unwrap();
13506 agent.chat("The report to Ada").await.unwrap();
13507 <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13508
13509 assert!(agent.pending_skill_id.read().is_none());
13510 assert!(
13511 !agent
13512 .disambiguation_manager()
13513 .unwrap()
13514 .has_pending_clarification()
13515 .await
13516 );
13517 }
13518
13519 #[tokio::test]
13521 async fn state_change_invalidates_pending_skill_confirmation() {
13522 let (agent, observed) = state_disambiguation_agent_with_skills(
13523 vec![
13524 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13525 "send_report",
13526 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13527 r#"{"question":"What should I send?","options":null}"#,
13528 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13529 r#"{"question":"Should I send the report to Ada?"}"#,
13530 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13531 "none",
13532 "Fresh response.",
13533 ],
13534 true,
13535 None,
13536 true,
13537 vec![confirmation_skill()],
13538 );
13539
13540 agent.chat("Send it").await.unwrap();
13541 agent.chat("The report to Ada").await.unwrap();
13542 agent.transition_to("review").await.unwrap();
13543 let cancelled = agent.chat("Yes").await.unwrap();
13544
13545 assert_eq!(cancelled.content, "Fresh response.");
13546 assert_eq!(confirmation_skill_call_count(&observed), 0);
13547 assert!(agent.pending_skill_id.read().is_none());
13548 }
13549
13550 #[tokio::test]
13552 async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13553 let (mut agent, observed) = state_disambiguation_agent_with_skills(
13554 vec![
13555 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13556 "send_report",
13557 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13558 r#"{"question":"What should I send?","options":null}"#,
13559 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13560 r#"{"question":"Should I send the report to Ada?"}"#,
13561 r#"{"status":"confirmed"}"#,
13562 "Confirmation cancelled.",
13563 ],
13564 true,
13565 None,
13566 true,
13567 vec![confirmation_skill()],
13568 );
13569 let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13570 let manager = agent
13571 .disambiguation_manager
13572 .take()
13573 .unwrap()
13574 .with_clarification_observer(observer.clone());
13575 agent.disambiguation_manager = Some(manager);
13576 let agent = Arc::new(agent);
13577
13578 agent.chat("Send it").await.unwrap();
13579 agent.chat("The report to Ada").await.unwrap();
13580
13581 let confirming_agent = Arc::clone(&agent);
13582 let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13583 observer.entered.wait().await;
13584 agent.reset().await.unwrap();
13585 observer.release.notify_one();
13586
13587 let response = confirmation.await.unwrap().unwrap();
13588 assert_eq!(response.content, "Confirmation cancelled.");
13589 assert_eq!(confirmation_skill_call_count(&observed), 0);
13590 assert!(agent.pending_skill_id.read().is_none());
13591 }
13592
13593 #[tokio::test]
13595 async fn queued_reset_prevents_stale_confirmation_question_publication() {
13596 let (agent, observed) = state_disambiguation_agent(
13597 vec![
13598 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13599 r#"{"question":"What should I send?","options":null}"#,
13600 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13601 r#"{"question":"Should I send the report to Ada?"}"#,
13602 ],
13603 true,
13604 None,
13605 true,
13606 );
13607 let agent = Arc::new(agent);
13608 agent.chat("Send it").await.unwrap();
13609
13610 let admission = agent.disambiguation_admission.write().await;
13611 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13612 let resetting_agent = Arc::clone(&agent);
13613 let reset = tokio::spawn(async move {
13614 let _ = started_tx.send(());
13615 resetting_agent.reset().await
13616 });
13617 started_rx.await.unwrap();
13618 tokio::task::yield_now().await;
13619
13620 let responding_agent = Arc::clone(&agent);
13621 let response =
13622 tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13623 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13624 while observed.call_count() < 4 {
13625 tokio::task::yield_now().await;
13626 }
13627 })
13628 .await
13629 .expect("clarification processing must reach terminal publication");
13630 drop(admission);
13631
13632 reset.await.unwrap().unwrap();
13633 let error = response.await.unwrap().unwrap_err();
13634 assert!(error.to_string().contains("ownership changed"));
13635 assert!(
13636 !agent
13637 .disambiguation_manager()
13638 .unwrap()
13639 .has_pending_clarification()
13640 .await
13641 );
13642 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13643 }
13644
13645 #[tokio::test]
13647 async fn queued_reset_prevents_stale_skill_clarification_publication() {
13648 let (agent, observed) = state_disambiguation_agent_with_skills(
13649 vec![
13650 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13651 "send_report",
13652 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13653 r#"{"question":"What should I send?","options":null}"#,
13654 ],
13655 true,
13656 None,
13657 true,
13658 vec![confirmation_skill()],
13659 );
13660 let agent = Arc::new(agent);
13661 let admission = agent.disambiguation_admission.write().await;
13662 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13663 let resetting_agent = Arc::clone(&agent);
13664 let reset = tokio::spawn(async move {
13665 let _ = started_tx.send(());
13666 resetting_agent.reset().await
13667 });
13668 started_rx.await.unwrap();
13669 tokio::task::yield_now().await;
13670
13671 let responding_agent = Arc::clone(&agent);
13672 let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13673 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13674 while observed.call_count() < 4 {
13675 tokio::task::yield_now().await;
13676 }
13677 })
13678 .await
13679 .expect("skill clarification must reach terminal publication");
13680 drop(admission);
13681
13682 reset.await.unwrap().unwrap();
13683 let error = response.await.unwrap().unwrap_err();
13684 assert!(error.to_string().contains("ownership changed"));
13685 assert_eq!(confirmation_skill_call_count(&observed), 0);
13686 assert!(agent.pending_skill_id.read().is_none());
13687 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13688 }
13689
13690 #[tokio::test]
13692 async fn transition_hook_can_reset_without_admission_deadlock() {
13693 let hooks = Arc::new(ResetOnTransitionHooks {
13694 agent: parking_lot::Mutex::new(None),
13695 invoked: AtomicBool::new(false),
13696 });
13697 let agent = Arc::new(
13698 AgentBuilder::new()
13699 .system_prompt("Test transition hook reentrancy.")
13700 .llm(Arc::new(mock_with_response("done")))
13701 .state_machine(disambiguation_state_machine(None, false))
13702 .build()
13703 .unwrap()
13704 .with_hooks(hooks.clone()),
13705 );
13706 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13707
13708 let transitioned = tokio::time::timeout(
13709 std::time::Duration::from_secs(2),
13710 agent.apply_transition_target("active", "review", "test transition", None),
13711 )
13712 .await
13713 .expect("transition hook reset must not deadlock")
13714 .unwrap();
13715
13716 assert!(transitioned);
13717 assert!(hooks.invoked.load(Ordering::SeqCst));
13718 assert_eq!(agent.current_state().as_deref(), Some("active"));
13719 }
13720
13721 #[tokio::test]
13723 async fn concurrent_transition_cannot_duplicate_exit_actions() {
13724 let gate = PathMutationGate::new();
13725 let active = ai_agents_state::StateDefinition {
13726 on_exit: vec![StateAction::Tool {
13727 tool: "transition_exit".to_string(),
13728 args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13729 }],
13730 ..Default::default()
13731 };
13732 let state_machine = Arc::new(
13733 StateMachine::new(ai_agents_state::StateConfig {
13734 initial: "active".to_string(),
13735 states: HashMap::from([
13736 ("active".to_string(), active),
13737 (
13738 "review".to_string(),
13739 ai_agents_state::StateDefinition::default(),
13740 ),
13741 ]),
13742 global_transitions: Vec::new(),
13743 fallback: None,
13744 max_no_transition: None,
13745 regenerate_on_transition: true,
13746 })
13747 .unwrap(),
13748 );
13749 let agent = Arc::new(
13750 AgentBuilder::new()
13751 .system_prompt("Test transition reservation.")
13752 .llm(Arc::new(mock_with_response("done")))
13753 .tool(Arc::new(BlockingPathMutationTool {
13754 id: "transition_exit",
13755 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13756 gate: gate.clone(),
13757 }))
13758 .state_machine(state_machine)
13759 .build()
13760 .unwrap(),
13761 );
13762
13763 let first_agent = Arc::clone(&agent);
13764 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13765 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13766 .await
13767 .expect("reserved transition must enter its exit action");
13768
13769 let second = tokio::time::timeout(
13770 std::time::Duration::from_secs(2),
13771 agent.transition_to("review"),
13772 )
13773 .await
13774 .expect("competing transition must fail without waiting for the exit action")
13775 .unwrap_err();
13776 assert!(second.to_string().contains("already in progress"));
13777
13778 gate.release();
13779 first.await.unwrap().unwrap();
13780 assert_eq!(agent.current_state().as_deref(), Some("review"));
13781 }
13782
13783 #[tokio::test]
13785 async fn concurrent_transition_cannot_overtake_enter_actions() {
13786 let gate = PathMutationGate::new();
13787 let review = ai_agents_state::StateDefinition {
13788 on_enter: vec![StateAction::Tool {
13789 tool: "transition_enter".to_string(),
13790 args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
13791 }],
13792 ..Default::default()
13793 };
13794 let state_machine = Arc::new(
13795 StateMachine::new(ai_agents_state::StateConfig {
13796 initial: "active".to_string(),
13797 states: HashMap::from([
13798 (
13799 "active".to_string(),
13800 ai_agents_state::StateDefinition::default(),
13801 ),
13802 ("review".to_string(), review),
13803 ]),
13804 global_transitions: Vec::new(),
13805 fallback: None,
13806 max_no_transition: None,
13807 regenerate_on_transition: true,
13808 })
13809 .unwrap(),
13810 );
13811 let agent = Arc::new(
13812 AgentBuilder::new()
13813 .system_prompt("Test transition lifecycle reservation.")
13814 .llm(Arc::new(mock_with_response("done")))
13815 .tool(Arc::new(BlockingPathMutationTool {
13816 id: "transition_enter",
13817 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13818 gate: gate.clone(),
13819 }))
13820 .state_machine(state_machine)
13821 .build()
13822 .unwrap(),
13823 );
13824
13825 let first_agent = Arc::clone(&agent);
13826 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13827 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13828 .await
13829 .expect("committed transition must enter its destination action");
13830
13831 let second = agent.transition_to("active").await.unwrap_err();
13832 assert!(second.to_string().contains("already in progress"));
13833 assert!(agent.reset().await.is_err());
13834
13835 gate.release();
13836 first.await.unwrap().unwrap();
13837 assert_eq!(agent.current_state().as_deref(), Some("review"));
13838 }
13839
13840 #[tokio::test]
13842 async fn same_state_restore_invalidates_pending_skill_confirmation() {
13843 let (agent, observed) = state_disambiguation_agent_with_skills(
13844 vec![
13845 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13846 "send_report",
13847 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13848 r#"{"question":"What should I send?","options":null}"#,
13849 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13850 r#"{"question":"Should I send the report to Ada?"}"#,
13851 ],
13852 true,
13853 None,
13854 true,
13855 vec![confirmation_skill()],
13856 );
13857
13858 agent.chat("Send it").await.unwrap();
13859 agent.chat("The report to Ada").await.unwrap();
13860 let snapshot = agent.save_state().await.unwrap();
13861 assert_eq!(agent.current_state().as_deref(), Some("active"));
13862
13863 agent.restore_state(snapshot).await.unwrap();
13864
13865 assert_eq!(agent.current_state().as_deref(), Some("active"));
13866 assert!(agent.pending_skill_id.read().is_none());
13867 assert!(
13868 !agent
13869 .disambiguation_manager()
13870 .unwrap()
13871 .has_pending_clarification()
13872 .await
13873 );
13874 assert_eq!(confirmation_skill_call_count(&observed), 0);
13875 }
13876
13877 #[tokio::test]
13879 async fn direct_state_generation_change_invalidates_confirmation() {
13880 let (agent, observed) = state_disambiguation_agent_with_skills(
13881 vec![
13882 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13883 "send_report",
13884 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13885 r#"{"question":"What should I send?","options":null}"#,
13886 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13887 r#"{"question":"Should I send the report to Ada?"}"#,
13888 "Confirmation cancelled.",
13889 ],
13890 true,
13891 None,
13892 true,
13893 vec![confirmation_skill()],
13894 );
13895
13896 agent.chat("Send it").await.unwrap();
13897 agent.chat("The report to Ada").await.unwrap();
13898 let state_machine = agent.state_machine().unwrap();
13899 state_machine
13900 .transition_to("review", "external test")
13901 .unwrap();
13902 state_machine
13903 .transition_to("active", "external test")
13904 .unwrap();
13905
13906 let response = agent.chat("Yes").await.unwrap();
13907
13908 assert_eq!(response.content, "Confirmation cancelled.");
13909 assert_eq!(confirmation_skill_call_count(&observed), 0);
13910 assert!(agent.pending_skill_id.read().is_none());
13911 }
13912
13913 #[tokio::test]
13914 async fn state_confirmation_does_not_add_a_question_for_clear_input() {
13915 let (agent, observed) = state_disambiguation_agent(
13916 vec![
13917 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
13918 "Request executed.",
13919 ],
13920 true,
13921 None,
13922 true,
13923 );
13924
13925 let response = agent.chat("Send the report to Ada").await.unwrap();
13926
13927 assert_eq!(response.content, "Request executed.");
13928 assert_eq!(observed.call_count(), 2);
13929 }
13930
13931 #[tokio::test]
13932 async fn state_override_cannot_activate_a_disabled_top_level_manager() {
13933 let (agent, observed) =
13934 state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
13935
13936 assert!(!agent.has_disambiguation());
13937 let response = agent.chat("Send it").await.unwrap();
13938
13939 assert_eq!(response.content, "Request executed.");
13940 assert_eq!(observed.call_count(), 1);
13941 }
13942
13943 #[tokio::test]
13944 async fn native_required_choice_executes_through_the_shared_tool_path() {
13945 let mut mock = MockLLMProvider::new("native-required");
13946 mock.set_tool_choice(Some(ToolChoice::Required));
13947 mock.add_response(
13948 LLMResponse::new("", FinishReason::ToolCall)
13949 .with_tool_calls(vec![ToolCall {
13950 id: "provider-call-1".to_string(),
13951 name: "calculator".to_string(),
13952 arguments: serde_json::json!({"expression": "2 + 2"}),
13953 }])
13954 .unwrap(),
13955 );
13956 mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
13957 let observed = mock.clone();
13958 let agent = AgentBuilder::new()
13959 .system_prompt("Use the calculator when needed.")
13960 .llm(Arc::new(mock))
13961 .tool(Arc::new(CalculatorTool::new()))
13962 .build()
13963 .unwrap();
13964
13965 let response = agent.chat("What is 2 + 2?").await.unwrap();
13966
13967 assert_eq!(response.content, "The answer is 4.");
13968 assert_eq!(
13969 response.tool_calls.as_ref().unwrap()[0].id,
13970 "provider-call-1"
13971 );
13972 let calls = observed.call_history();
13973 assert_eq!(calls.len(), 2);
13974 assert!(matches!(
13975 calls[0].request.as_ref().map(|request| &request.choice),
13976 Some(ToolChoice::Required)
13977 ));
13978 assert!(matches!(
13979 calls[1].request.as_ref().map(|request| &request.choice),
13980 Some(ToolChoice::Auto)
13981 ));
13982 }
13983
13984 #[tokio::test]
13985 async fn prompt_fallback_uses_one_corrective_retry() {
13986 let mut mock = MockLLMProvider::new("prompt-required");
13987 mock.set_tool_choice(Some(ToolChoice::Required));
13988 mock.set_native_tool_support(false);
13989 mock.set_responses(
13990 vec![
13991 "I can calculate that.".to_string(),
13992 r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
13993 "The answer is 4.".to_string(),
13994 ],
13995 false,
13996 );
13997 let observed = mock.clone();
13998 let agent = AgentBuilder::new()
13999 .system_prompt("Use tools.")
14000 .llm(Arc::new(mock))
14001 .tool(Arc::new(CalculatorTool::new()))
14002 .build()
14003 .unwrap();
14004
14005 let response = agent.chat("What is 2 + 2?").await.unwrap();
14006
14007 assert_eq!(response.content, "The answer is 4.");
14008 assert_eq!(observed.call_count(), 3);
14009 let corrective = &observed.call_history()[1].messages;
14010 assert!(
14011 corrective
14012 .last()
14013 .unwrap()
14014 .content
14015 .contains("previous response")
14016 );
14017 }
14018
14019 #[tokio::test]
14020 async fn prompt_fallback_fails_after_one_noncompliant_retry() {
14021 let mut mock = MockLLMProvider::new("prompt-required-failure");
14022 mock.set_tool_choice(Some(ToolChoice::Required));
14023 mock.set_native_tool_support(false);
14024 mock.set_responses(
14025 vec!["No tool.".to_string(), "Still no tool.".to_string()],
14026 false,
14027 );
14028 let observed = mock.clone();
14029 let agent = AgentBuilder::new()
14030 .system_prompt("Use tools.")
14031 .llm(Arc::new(mock))
14032 .tool(Arc::new(CalculatorTool::new()))
14033 .build()
14034 .unwrap();
14035
14036 let error = agent.chat("What is 2 + 2?").await.unwrap_err();
14037
14038 assert!(error.to_string().contains("one corrective retry"));
14039 assert_eq!(observed.call_count(), 2);
14040 }
14041
14042 #[tokio::test]
14043 async fn specific_choice_cannot_widen_the_effective_grant() {
14044 let mut mock = MockLLMProvider::new("specific-outside-grant");
14045 mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
14046 let observed = mock.clone();
14047 let agent = AgentBuilder::new()
14048 .system_prompt("Use tools.")
14049 .llm(Arc::new(mock))
14050 .tool(Arc::new(CalculatorTool::new()))
14051 .build()
14052 .unwrap();
14053
14054 let error = agent.chat("Generate a value.").await.unwrap_err();
14055
14056 assert!(error.to_string().contains("is not registered"));
14057 assert_eq!(observed.call_count(), 0);
14058 }
14059
14060 #[tokio::test]
14061 async fn none_choice_exposes_no_tool_protocol() {
14062 let mut mock = MockLLMProvider::new("no-tools");
14063 mock.set_tool_choice(Some(ToolChoice::None));
14064 mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
14065 let observed = mock.clone();
14066 let agent = AgentBuilder::new()
14067 .system_prompt("Answer directly.")
14068 .llm(Arc::new(mock))
14069 .tool(Arc::new(CalculatorTool::new()))
14070 .build()
14071 .unwrap();
14072
14073 let response = agent.chat("Hello").await.unwrap();
14074
14075 assert!(response.tool_calls.is_none());
14076 assert_eq!(observed.call_count(), 1);
14077 let call = observed.last_call().unwrap();
14078 assert!(call.request.is_none());
14079 assert!(
14080 call.messages
14081 .iter()
14082 .all(|message| !message.content.contains("Available tools:"))
14083 );
14084 }
14085
14086 struct RuntimeStorage {
14087 capabilities: Box<[StorageCapability]>,
14088 snapshots: RwLock<HashMap<String, AgentSnapshot>>,
14089 metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
14090 metadata_save_calls: AtomicU64,
14091 metadata_load_calls: AtomicU64,
14092 fail_metadata_save: AtomicBool,
14093 fail_metadata_load: AtomicBool,
14094 }
14095
14096 impl RuntimeStorage {
14097 fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
14098 Self {
14099 capabilities: capabilities.into_iter().collect(),
14100 snapshots: RwLock::new(HashMap::new()),
14101 metadata: RwLock::new(HashMap::new()),
14102 metadata_save_calls: AtomicU64::new(0),
14103 metadata_load_calls: AtomicU64::new(0),
14104 fail_metadata_save: AtomicBool::new(false),
14105 fail_metadata_load: AtomicBool::new(false),
14106 }
14107 }
14108 }
14109
14110 #[async_trait]
14111 impl AgentStorage for RuntimeStorage {
14112 fn supports(&self, capability: StorageCapability) -> bool {
14113 self.capabilities.contains(&capability)
14114 }
14115
14116 async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
14117 self.snapshots
14118 .write()
14119 .insert(session_id.to_string(), snapshot.clone());
14120 Ok(())
14121 }
14122
14123 async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
14124 Ok(self.snapshots.read().get(session_id).cloned())
14125 }
14126
14127 async fn delete(&self, session_id: &str) -> Result<()> {
14128 self.snapshots.write().remove(session_id);
14129 Ok(())
14130 }
14131
14132 async fn list_sessions(&self) -> Result<Vec<String>> {
14133 Ok(self.snapshots.read().keys().cloned().collect())
14134 }
14135
14136 async fn save_snapshot_with_metadata(
14137 &self,
14138 session_id: &str,
14139 snapshot: &AgentSnapshot,
14140 metadata: &ai_agents_core::SessionMetadata,
14141 ) -> Result<()> {
14142 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14143 if self.fail_metadata_save.load(Ordering::SeqCst) {
14144 return Err(AgentError::Persistence("metadata save failed".into()));
14145 }
14146 self.snapshots
14147 .write()
14148 .insert(session_id.to_string(), snapshot.clone());
14149 self.metadata
14150 .write()
14151 .insert(session_id.to_string(), metadata.clone());
14152 Ok(())
14153 }
14154
14155 async fn save_metadata(
14156 &self,
14157 session_id: &str,
14158 metadata: &ai_agents_core::SessionMetadata,
14159 ) -> Result<()> {
14160 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14161 if self.fail_metadata_save.load(Ordering::SeqCst) {
14162 return Err(AgentError::Persistence("metadata save failed".into()));
14163 }
14164 self.metadata
14165 .write()
14166 .insert(session_id.to_string(), metadata.clone());
14167 Ok(())
14168 }
14169
14170 async fn load_metadata(
14171 &self,
14172 session_id: &str,
14173 ) -> Result<Option<ai_agents_core::SessionMetadata>> {
14174 self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
14175 if self.fail_metadata_load.load(Ordering::SeqCst) {
14176 return Err(AgentError::Persistence("metadata load failed".into()));
14177 }
14178 Ok(self.metadata.read().get(session_id).cloned())
14179 }
14180 }
14181
14182 fn runtime_storage_agent() -> RuntimeAgent {
14183 AgentBuilder::new()
14184 .system_prompt("Test runtime storage integration.")
14185 .llm(Arc::new(mock_with_response("done")))
14186 .build()
14187 .unwrap()
14188 }
14189
14190 fn restore_spec(id: &str) -> crate::spec::AgentSpec {
14191 crate::spec::AgentSpec {
14192 name: id.to_string(),
14193 system_prompt: format!("Restore child {id}."),
14194 ..crate::spec::AgentSpec::default()
14195 }
14196 }
14197
14198 fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
14199 ai_agents_core::SpawnedAgentEntry {
14200 id: id.to_string(),
14201 name: id.to_string(),
14202 spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
14203 }
14204 }
14205
14206 fn restore_spawner(
14207 storage: Arc<RuntimeStorage>,
14208 max_agents: usize,
14209 ) -> (
14210 Arc<crate::spawner::AgentSpawner>,
14211 Arc<crate::spawner::AgentRegistry>,
14212 ) {
14213 let mut llms = LLMRegistry::new();
14214 llms.register("default", Arc::new(mock_with_response("done")));
14215 (
14216 Arc::new(
14217 crate::spawner::AgentSpawner::new()
14218 .with_shared_llms(llms)
14219 .with_shared_storage(storage)
14220 .with_max_agents(max_agents),
14221 ),
14222 Arc::new(crate::spawner::AgentRegistry::new()),
14223 )
14224 }
14225
14226 async fn save_restore_target(
14227 parent: &RuntimeAgent,
14228 storage: &RuntimeStorage,
14229 session_id: &str,
14230 entries: Vec<ai_agents_core::SpawnedAgentEntry>,
14231 ) {
14232 let mut snapshot = parent.save_state().await.unwrap();
14233 snapshot.spawned_agents = Some(entries);
14234 storage.save(session_id, &snapshot).await.unwrap();
14235 storage
14236 .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
14237 .await
14238 .unwrap();
14239 }
14240
14241 #[tokio::test]
14242 async fn storage_init_requires_storage_for_actor_facts() {
14243 let facts = ai_agents_facts::FactsConfig {
14244 enabled: true,
14245 ..Default::default()
14246 };
14247 let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
14248
14249 let error = agent.init_storage().await.unwrap_err();
14250 assert!(matches!(
14251 error,
14252 AgentError::Config(message)
14253 if message.contains("actor facts or actor memory")
14254 && message.contains("none is configured or injected")
14255 ));
14256 }
14257
14258 #[tokio::test]
14259 async fn storage_init_validates_actor_facts_capability() {
14260 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14261 let actor_memory = ai_agents_facts::ActorMemoryConfig {
14262 enabled: true,
14263 ..Default::default()
14264 };
14265 let agent = runtime_storage_agent()
14266 .with_storage(storage)
14267 .with_facts_config(Some(actor_memory), None);
14268
14269 assert!(matches!(
14270 agent.init_storage().await,
14271 Err(AgentError::UnsupportedStorageCapability(
14272 StorageCapability::ActorFacts
14273 ))
14274 ));
14275 }
14276
14277 #[tokio::test]
14278 async fn blocking_chat_rejects_unsupported_required_storage() {
14279 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14280 let facts = ai_agents_facts::FactsConfig {
14281 enabled: true,
14282 ..Default::default()
14283 };
14284 let agent = runtime_storage_agent()
14285 .with_storage(storage)
14286 .with_facts_config(None, Some(facts));
14287
14288 assert!(matches!(
14289 agent.chat("hello").await,
14290 Err(AgentError::UnsupportedStorageCapability(
14291 StorageCapability::ActorFacts
14292 ))
14293 ));
14294 }
14295
14296 #[tokio::test]
14297 async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
14298 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14299 let config = ai_agents_relationships::RelationshipConfig {
14300 enabled: true,
14301 ..Default::default()
14302 };
14303 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14304 let agent = runtime_storage_agent()
14305 .with_storage(storage)
14306 .with_relationships(manager);
14307
14308 assert!(matches!(
14309 agent.chat_stream("hello").await,
14310 Err(AgentError::UnsupportedStorageCapability(
14311 StorageCapability::ActorRelationships
14312 ))
14313 ));
14314 }
14315
14316 #[tokio::test]
14317 async fn storage_init_completes_facts_for_injected_storage() {
14318 let storage = Arc::new(RuntimeStorage::new([
14319 StorageCapability::Snapshot,
14320 StorageCapability::ActorFacts,
14321 ]));
14322 let facts = ai_agents_facts::FactsConfig {
14323 enabled: true,
14324 ..Default::default()
14325 };
14326 let agent = runtime_storage_agent()
14327 .with_storage(storage)
14328 .with_facts_config(None, Some(facts));
14329
14330 agent.init_storage().await.unwrap();
14331 assert!(agent.fact_store().is_some());
14332 }
14333
14334 #[tokio::test]
14335 async fn storage_init_requires_storage_for_persistent_relationships() {
14336 let config = ai_agents_relationships::RelationshipConfig {
14337 enabled: true,
14338 ..Default::default()
14339 };
14340 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14341 let agent = runtime_storage_agent().with_relationships(manager);
14342
14343 let error = agent.init_storage().await.unwrap_err();
14344 assert!(matches!(
14345 error,
14346 AgentError::Config(message)
14347 if message.contains("persistent relationships")
14348 && message.contains("none is configured or injected")
14349 ));
14350 }
14351
14352 #[tokio::test]
14353 async fn storage_init_validates_persistent_relationships_capability() {
14354 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14355 let config = ai_agents_relationships::RelationshipConfig {
14356 enabled: true,
14357 ..Default::default()
14358 };
14359 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14360 let agent = runtime_storage_agent()
14361 .with_storage(storage)
14362 .with_relationships(manager);
14363
14364 assert!(matches!(
14365 agent.init_storage().await,
14366 Err(AgentError::UnsupportedStorageCapability(
14367 StorageCapability::ActorRelationships
14368 ))
14369 ));
14370 }
14371
14372 #[tokio::test]
14373 async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
14374 let storage = Arc::new(RuntimeStorage::new([
14375 StorageCapability::Snapshot,
14376 StorageCapability::SessionMetadata,
14377 ]));
14378 let agent = runtime_storage_agent().with_storage(storage.clone());
14379 agent.set_actor_id("old-actor").unwrap();
14380 agent.save_session("old").await.unwrap();
14381 storage
14382 .save("target", &agent.save_state().await.unwrap())
14383 .await
14384 .unwrap();
14385 storage
14386 .save_metadata("target", &ai_agents_core::SessionMetadata::default())
14387 .await
14388 .unwrap();
14389
14390 assert!(agent.load_session("target").await.unwrap());
14391
14392 assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
14393 assert_eq!(agent.actor_id(), None);
14394 }
14395
14396 #[tokio::test]
14397 async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
14398 let storage = Arc::new(RuntimeStorage::new([
14399 StorageCapability::Snapshot,
14400 StorageCapability::SessionMetadata,
14401 ]));
14402 let (spawner, registry) = restore_spawner(storage.clone(), 3);
14403 let parent = runtime_storage_agent()
14404 .with_storage(storage.clone())
14405 .with_spawner_handles(Arc::clone(&spawner), Arc::clone(®istry));
14406
14407 for id in ["a", "b"] {
14408 let spawned = spawner
14409 .spawn_with_id(id.to_string(), restore_spec(id))
14410 .await
14411 .unwrap();
14412 spawned.agent.save_session("grow").await.unwrap();
14413 registry.register(spawned).await.unwrap();
14414 }
14415 let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
14416 staged_c
14417 .save("grow", &AgentSnapshot::new("c".into()))
14418 .await
14419 .unwrap();
14420 staged_c
14421 .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
14422 .await
14423 .unwrap();
14424 save_restore_target(
14425 &parent,
14426 storage.as_ref(),
14427 "grow",
14428 vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
14429 )
14430 .await;
14431
14432 assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
14433 assert_eq!(registry.count(), 3);
14434 assert!(registry.contains("c"));
14435 assert_eq!(spawner.spawned_count(), 3);
14436
14437 for id in ["a", "b"] {
14438 registry
14439 .get(id)
14440 .unwrap()
14441 .save_session("shrink")
14442 .await
14443 .unwrap();
14444 }
14445 save_restore_target(
14446 &parent,
14447 storage.as_ref(),
14448 "shrink",
14449 vec![restore_entry("a"), restore_entry("b")],
14450 )
14451 .await;
14452
14453 assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
14454 assert_eq!(registry.count(), 2);
14455 assert!(!registry.contains("c"));
14456 assert_eq!(spawner.spawned_count(), 2);
14457
14458 save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
14459
14460 assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
14461 assert_eq!(registry.count(), 0);
14462 assert_eq!(spawner.spawned_count(), 0);
14463 assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
14464 }
14465
14466 #[tokio::test]
14467 async fn storage_session_metadata_is_called_only_when_advertised() {
14468 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14469 storage.fail_metadata_save.store(true, Ordering::SeqCst);
14470 storage.fail_metadata_load.store(true, Ordering::SeqCst);
14471 let agent = runtime_storage_agent().with_storage(storage.clone());
14472
14473 agent.save_session("session").await.unwrap();
14474 assert!(agent.load_session("session").await.unwrap());
14475 assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
14476 assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
14477 }
14478
14479 #[cfg(feature = "sqlite")]
14480 #[tokio::test]
14481 async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
14482 let directory =
14483 std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
14484 let path = directory.join("sessions.sqlite");
14485 let path_string = path.to_string_lossy().into_owned();
14486 let storage = Arc::new(
14487 ai_agents_storage::SqliteStorage::new(&path_string)
14488 .await
14489 .unwrap(),
14490 );
14491 let agent = runtime_storage_agent().with_storage(storage.clone());
14492 agent.set_session_metadata(ai_agents_core::SessionMetadata {
14493 tags: vec!["initial".into()],
14494 ..Default::default()
14495 });
14496 agent.chat("persist this turn").await.unwrap();
14497 agent.save_session("session").await.unwrap();
14498
14499 agent.set_session_metadata(ai_agents_core::SessionMetadata {
14500 tags: vec!["updated".into()],
14501 ..Default::default()
14502 });
14503 agent.save_session("session").await.unwrap();
14504 assert!(
14505 agent
14506 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14507 tags: Some(vec!["initial".into()]),
14508 ..Default::default()
14509 })
14510 .await
14511 .unwrap()
14512 .is_empty()
14513 );
14514 assert_eq!(
14515 agent
14516 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14517 tags: Some(vec!["updated".into()]),
14518 ..Default::default()
14519 })
14520 .await
14521 .unwrap()
14522 .len(),
14523 1
14524 );
14525 drop(agent);
14526 storage.close().await;
14527 drop(storage);
14528
14529 let reopened_storage = Arc::new(
14530 ai_agents_storage::SqliteStorage::new(&path_string)
14531 .await
14532 .unwrap(),
14533 );
14534 let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14535 assert!(restored.load_session("session").await.unwrap());
14536 assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14537 assert_eq!(
14538 restored.current_session_id.read().as_deref(),
14539 Some("session")
14540 );
14541 assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14542 assert_eq!(
14543 restored
14544 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14545 tags: Some(vec!["updated".into()]),
14546 ..Default::default()
14547 })
14548 .await
14549 .unwrap()
14550 .len(),
14551 1
14552 );
14553
14554 drop(restored);
14555 reopened_storage.close().await;
14556 drop(reopened_storage);
14557 crate::remove_sqlite_test_directory(&directory)
14558 .await
14559 .unwrap();
14560 }
14561
14562 #[tokio::test]
14563 async fn storage_session_metadata_backend_failures_propagate() {
14564 let storage = Arc::new(RuntimeStorage::new([
14565 StorageCapability::Snapshot,
14566 StorageCapability::SessionMetadata,
14567 ]));
14568 let agent = runtime_storage_agent().with_storage(storage.clone());
14569
14570 agent.save_session("session").await.unwrap();
14571 storage
14572 .save("target", &agent.save_state().await.unwrap())
14573 .await
14574 .unwrap();
14575 storage.fail_metadata_load.store(true, Ordering::SeqCst);
14576 assert!(matches!(
14577 agent.load_session("target").await,
14578 Err(AgentError::Persistence(message)) if message == "metadata load failed"
14579 ));
14580 assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14581
14582 storage.fail_metadata_save.store(true, Ordering::SeqCst);
14583 assert!(matches!(
14584 agent.save_session("session").await,
14585 Err(AgentError::Persistence(message)) if message == "metadata save failed"
14586 ));
14587 }
14588
14589 struct ProviderFutureDropSignal {
14590 dropped: Arc<AtomicBool>,
14591 }
14592
14593 impl Drop for ProviderFutureDropSignal {
14594 fn drop(&mut self) {
14595 self.dropped.store(true, Ordering::SeqCst);
14596 }
14597 }
14598
14599 struct BufferedLockingProvider {
14600 lock: Arc<tokio::sync::Mutex<()>>,
14601 stream_started: Arc<tokio::sync::Notify>,
14602 stream_dropped: Arc<AtomicBool>,
14603 committed_after_drop: Arc<AtomicBool>,
14604 }
14605
14606 #[async_trait]
14607 impl LLMProvider for BufferedLockingProvider {
14608 async fn complete(
14609 &self,
14610 _messages: &[ChatMessage],
14611 _config: Option<&LLMConfig>,
14612 ) -> std::result::Result<LLMResponse, LLMError> {
14613 let _guard = self.lock.lock().await;
14614 self.committed_after_drop
14615 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14616 Ok(LLMResponse::new(
14617 "Committed technical response.",
14618 FinishReason::Stop,
14619 ))
14620 }
14621
14622 async fn complete_stream(
14623 &self,
14624 _messages: &[ChatMessage],
14625 _config: Option<&LLMConfig>,
14626 ) -> std::result::Result<
14627 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14628 LLMError,
14629 > {
14630 let _guard = self.lock.lock().await;
14631 let _drop_signal = ProviderFutureDropSignal {
14632 dropped: Arc::clone(&self.stream_dropped),
14633 };
14634 self.stream_started.notify_one();
14635 std::future::pending().await
14636 }
14637
14638 fn provider_name(&self) -> &str {
14639 "buffered-locking"
14640 }
14641
14642 fn supports(&self, _feature: LLMFeature) -> bool {
14643 false
14644 }
14645 }
14646
14647 struct PendingDropStream {
14648 dropped: Arc<AtomicBool>,
14649 dropped_notify: Arc<tokio::sync::Notify>,
14650 }
14651
14652 impl Stream for PendingDropStream {
14653 type Item = std::result::Result<LLMChunk, LLMError>;
14654
14655 fn poll_next(
14656 self: Pin<&mut Self>,
14657 _cx: &mut std::task::Context<'_>,
14658 ) -> std::task::Poll<Option<Self::Item>> {
14659 std::task::Poll::Pending
14660 }
14661 }
14662
14663 impl Drop for PendingDropStream {
14664 fn drop(&mut self) {
14665 self.dropped.store(true, Ordering::SeqCst);
14666 self.dropped_notify.notify_one();
14667 }
14668 }
14669
14670 struct EstablishedStreamProvider {
14671 stream_started: Arc<tokio::sync::Notify>,
14672 stream_dropped: Arc<AtomicBool>,
14673 stream_dropped_notify: Arc<tokio::sync::Notify>,
14674 committed_after_drop: Arc<AtomicBool>,
14675 }
14676
14677 #[async_trait]
14678 impl LLMProvider for EstablishedStreamProvider {
14679 async fn complete(
14680 &self,
14681 _messages: &[ChatMessage],
14682 _config: Option<&LLMConfig>,
14683 ) -> std::result::Result<LLMResponse, LLMError> {
14684 if !self.stream_dropped.load(Ordering::SeqCst) {
14685 self.stream_dropped_notify.notified().await;
14686 }
14687 self.committed_after_drop
14688 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14689 Ok(LLMResponse::new(
14690 "Committed technical response.",
14691 FinishReason::Stop,
14692 ))
14693 }
14694
14695 async fn complete_stream(
14696 &self,
14697 _messages: &[ChatMessage],
14698 _config: Option<&LLMConfig>,
14699 ) -> std::result::Result<
14700 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14701 LLMError,
14702 > {
14703 self.stream_started.notify_one();
14704 Ok(Box::new(PendingDropStream {
14705 dropped: Arc::clone(&self.stream_dropped),
14706 dropped_notify: Arc::clone(&self.stream_dropped_notify),
14707 }))
14708 }
14709
14710 fn provider_name(&self) -> &str {
14711 "established-stream"
14712 }
14713
14714 fn supports(&self, _feature: LLMFeature) -> bool {
14715 false
14716 }
14717 }
14718
14719 struct FirstCallLockingProvider {
14720 lock: Arc<tokio::sync::Mutex<()>>,
14721 first_started: Arc<tokio::sync::Notify>,
14722 first_dropped: Arc<AtomicBool>,
14723 committed_after_drop: Arc<AtomicBool>,
14724 calls: AtomicU64,
14725 }
14726
14727 #[async_trait]
14728 impl LLMProvider for FirstCallLockingProvider {
14729 async fn complete(
14730 &self,
14731 _messages: &[ChatMessage],
14732 _config: Option<&LLMConfig>,
14733 ) -> std::result::Result<LLMResponse, LLMError> {
14734 let _guard = self.lock.lock().await;
14735 let call = self.calls.fetch_add(1, Ordering::SeqCst);
14736 if call == 0 {
14737 let _drop_signal = ProviderFutureDropSignal {
14738 dropped: Arc::clone(&self.first_dropped),
14739 };
14740 self.first_started.notify_one();
14741 return std::future::pending().await;
14742 }
14743 self.committed_after_drop
14744 .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14745 Ok(LLMResponse::new(
14746 "Committed technical response.",
14747 FinishReason::Stop,
14748 ))
14749 }
14750
14751 async fn complete_stream(
14752 &self,
14753 _messages: &[ChatMessage],
14754 _config: Option<&LLMConfig>,
14755 ) -> std::result::Result<
14756 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14757 LLMError,
14758 > {
14759 Err(LLMError::Other(
14760 "streaming is not used in this test".to_string(),
14761 ))
14762 }
14763
14764 fn provider_name(&self) -> &str {
14765 "first-call-locking"
14766 }
14767
14768 fn supports(&self, _feature: LLMFeature) -> bool {
14769 false
14770 }
14771 }
14772
14773 struct RoutingAfterProviderStart {
14774 provider_started: Arc<tokio::sync::Notify>,
14775 }
14776
14777 #[async_trait]
14778 impl LLMProvider for RoutingAfterProviderStart {
14779 async fn complete(
14780 &self,
14781 _messages: &[ChatMessage],
14782 _config: Option<&LLMConfig>,
14783 ) -> std::result::Result<LLMResponse, LLMError> {
14784 self.provider_started.notified().await;
14785 Ok(LLMResponse::new("1", FinishReason::Stop))
14786 }
14787
14788 async fn complete_stream(
14789 &self,
14790 _messages: &[ChatMessage],
14791 _config: Option<&LLMConfig>,
14792 ) -> std::result::Result<
14793 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14794 LLMError,
14795 > {
14796 Err(LLMError::Other(
14797 "streaming is not used in this test".to_string(),
14798 ))
14799 }
14800
14801 fn provider_name(&self) -> &str {
14802 "routing-after-start"
14803 }
14804
14805 fn supports(&self, _feature: LLMFeature) -> bool {
14806 false
14807 }
14808 }
14809
14810 struct ResponseCountingHooks {
14812 responses: Arc<std::sync::atomic::AtomicUsize>,
14813 }
14814
14815 struct RootTurnProbeProvider {
14817 complete_entered: tokio::sync::mpsc::UnboundedSender<()>,
14818 }
14819
14820 struct ResponseChatHooks {
14822 target: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
14823 invoked: AtomicBool,
14824 nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
14825 }
14826
14827 struct ConcurrentResponseHooks {
14829 registry: Weak<crate::spawner::AgentRegistry>,
14830 child_id: String,
14831 invoked: AtomicBool,
14832 nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
14833 }
14834
14835 struct RetryDeadlineTool {
14837 calls: Arc<std::sync::atomic::AtomicUsize>,
14838 deadlines: Arc<parking_lot::Mutex<Vec<chrono::DateTime<chrono::Utc>>>>,
14839 }
14840
14841 struct ToolLifecycleRecordingHooks {
14843 events: parking_lot::Mutex<Vec<String>>,
14844 records: parking_lot::Mutex<Vec<ToolExecutionRecord>>,
14845 }
14846
14847 impl ToolLifecycleRecordingHooks {
14848 fn new() -> Self {
14850 Self {
14851 events: parking_lot::Mutex::new(Vec::new()),
14852 records: parking_lot::Mutex::new(Vec::new()),
14853 }
14854 }
14855
14856 fn events(&self) -> Vec<String> {
14858 self.events.lock().clone()
14859 }
14860
14861 fn records(&self) -> Vec<ToolExecutionRecord> {
14863 self.records.lock().clone()
14864 }
14865 }
14866
14867 struct ContextEchoTool;
14869
14870 #[async_trait]
14871 impl LLMProvider for RootTurnProbeProvider {
14872 async fn complete(
14873 &self,
14874 _messages: &[ChatMessage],
14875 _config: Option<&LLMConfig>,
14876 ) -> std::result::Result<LLMResponse, LLMError> {
14877 let _ = self.complete_entered.send(());
14878 Ok(LLMResponse::new("blocking complete", FinishReason::Stop))
14879 }
14880
14881 async fn complete_stream(
14882 &self,
14883 _messages: &[ChatMessage],
14884 _config: Option<&LLMConfig>,
14885 ) -> std::result::Result<
14886 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14887 LLMError,
14888 > {
14889 Ok(Box::new(futures::stream::iter(vec![Ok(
14890 LLMChunk::final_chunk("stream complete", FinishReason::Stop, None),
14891 )])))
14892 }
14893
14894 fn provider_name(&self) -> &str {
14895 "root-turn-probe"
14896 }
14897
14898 fn supports(&self, feature: LLMFeature) -> bool {
14899 matches!(feature, LLMFeature::Streaming)
14900 }
14901 }
14902
14903 #[async_trait]
14904 impl ai_agents_core::Tool for ContextEchoTool {
14905 fn id(&self) -> &str {
14906 "context_echo"
14907 }
14908
14909 fn name(&self) -> &str {
14910 "Context Echo"
14911 }
14912
14913 fn description(&self) -> &str {
14914 "Returns selected execution context fields."
14915 }
14916
14917 fn input_schema(&self) -> Value {
14918 serde_json::json!({"type": "object"})
14919 }
14920
14921 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14922 ai_agents_core::ToolPolicyBindings {
14923 path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
14924 result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
14925 "max_results",
14926 ai_agents_core::ResultLimitKind::MaxResults,
14927 )],
14928 ..Default::default()
14929 }
14930 }
14931
14932 async fn execute(
14933 &self,
14934 _args: Value,
14935 ctx: ai_agents_core::ToolExecutionContext,
14936 ) -> ToolResult {
14937 ToolResult::ok(
14938 serde_json::json!({
14939 "requested_name": ctx.requested_name,
14940 "canonical_id": ctx.canonical_id,
14941 "display_name": ctx.display_name,
14942 "max_results": ctx.limits.max_results,
14943 "custom_config": ctx.custom_config,
14944 })
14945 .to_string(),
14946 )
14947 }
14948 }
14949
14950 #[async_trait]
14951 impl ai_agents_core::Tool for RetryDeadlineTool {
14952 fn id(&self) -> &str {
14953 "retry_deadline"
14954 }
14955
14956 fn name(&self) -> &str {
14957 "Retry Deadline"
14958 }
14959
14960 fn description(&self) -> &str {
14961 "Records one deadline per retry invocation."
14962 }
14963
14964 fn input_schema(&self) -> Value {
14965 serde_json::json!({"type": "object"})
14966 }
14967
14968 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14969 ai_agents_core::ToolSafetyMetadata::compute()
14970 }
14971
14972 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
14973 let mut classification =
14974 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
14975 classification.timeout_ms = Some(1_000);
14976 classification.safely_retryable = true;
14977 classification
14978 }
14979
14980 async fn execute(
14981 &self,
14982 _args: Value,
14983 ctx: ai_agents_core::ToolExecutionContext,
14984 ) -> ToolResult {
14985 self.deadlines.lock().push(
14986 ctx.deadline
14987 .expect("each invocation must receive a deadline"),
14988 );
14989 let call = self.calls.fetch_add(1, Ordering::SeqCst);
14990 if call == 0 {
14991 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
14992 ToolResult::error("retry")
14993 } else {
14994 ToolResult::ok("done")
14995 }
14996 }
14997 }
14998
14999 struct SlowTool;
15001
15002 struct FlakyWriteTool {
15004 calls: Arc<std::sync::atomic::AtomicUsize>,
15005 }
15006
15007 struct LockedWriteTool {
15009 active: Arc<std::sync::atomic::AtomicUsize>,
15010 max_active: Arc<std::sync::atomic::AtomicUsize>,
15011 }
15012
15013 struct MultiResourceWriteTool {
15014 active: Arc<std::sync::atomic::AtomicUsize>,
15015 max_active: Arc<std::sync::atomic::AtomicUsize>,
15016 }
15017
15018 #[derive(Clone)]
15019 struct PathMutationGate {
15020 entered: Arc<AtomicBool>,
15021 entered_notify: Arc<tokio::sync::Notify>,
15022 release: Arc<tokio::sync::Notify>,
15023 }
15024
15025 impl PathMutationGate {
15026 fn new() -> Self {
15027 Self {
15028 entered: Arc::new(AtomicBool::new(false)),
15029 entered_notify: Arc::new(tokio::sync::Notify::new()),
15030 release: Arc::new(tokio::sync::Notify::new()),
15031 }
15032 }
15033
15034 async fn wait_until_entered(&self) {
15035 if !self.entered.load(Ordering::SeqCst) {
15036 self.entered_notify.notified().await;
15037 }
15038 }
15039
15040 fn release(&self) {
15041 self.release.notify_one();
15042 }
15043 }
15044
15045 struct BlockingPathMutationTool {
15046 id: &'static str,
15047 path_fields: Vec<ai_agents_core::PathPolicyBinding>,
15048 gate: PathMutationGate,
15049 }
15050
15051 struct NoBindingWriteTool {
15052 active: Arc<std::sync::atomic::AtomicUsize>,
15053 max_active: Arc<std::sync::atomic::AtomicUsize>,
15054 }
15055
15056 struct RecoveryTestTool {
15057 id: String,
15058 succeeds: bool,
15059 calls: Arc<std::sync::atomic::AtomicUsize>,
15060 max_output_chars: Option<usize>,
15061 }
15062
15063 struct BlockingApprovalHandler {
15064 entered: Arc<tokio::sync::Barrier>,
15065 release: Arc<tokio::sync::Notify>,
15066 result: ApprovalResult,
15067 }
15068
15069 struct CountingApprovalHandler {
15070 calls: Arc<std::sync::atomic::AtomicUsize>,
15071 }
15072
15073 struct RuntimeWebFetchTransport {
15074 calls: Arc<std::sync::atomic::AtomicUsize>,
15075 }
15076
15077 struct RuntimeWebFetchResolver;
15078
15079 struct ReentrantToolHooks {
15080 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15081 invoked: AtomicBool,
15082 nested_success: AtomicBool,
15083 }
15084
15085 #[async_trait]
15086 impl ai_agents_core::Tool for SlowTool {
15087 fn id(&self) -> &str {
15088 "slow"
15089 }
15090
15091 fn name(&self) -> &str {
15092 "Slow"
15093 }
15094
15095 fn description(&self) -> &str {
15096 "Waits until cancelled or timed out."
15097 }
15098
15099 fn input_schema(&self) -> Value {
15100 serde_json::json!({"type": "object"})
15101 }
15102
15103 async fn execute(
15104 &self,
15105 _args: Value,
15106 _ctx: ai_agents_core::ToolExecutionContext,
15107 ) -> ToolResult {
15108 tokio::time::sleep(std::time::Duration::from_secs(5)).await;
15109 ToolResult::ok("done")
15110 }
15111 }
15112
15113 #[async_trait]
15114 impl ai_agents_core::Tool for FlakyWriteTool {
15115 fn id(&self) -> &str {
15116 "flaky_write"
15117 }
15118
15119 fn name(&self) -> &str {
15120 "Flaky Write"
15121 }
15122
15123 fn description(&self) -> &str {
15124 "Fails on the first write attempt."
15125 }
15126
15127 fn input_schema(&self) -> Value {
15128 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15129 }
15130
15131 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15132 ai_agents_core::ToolPolicyBindings {
15133 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15134 ..Default::default()
15135 }
15136 }
15137
15138 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15139 ai_agents_core::ToolSafetyMetadata {
15140 read_only: false,
15141 concurrency_safe: false,
15142 operation: ai_agents_core::ToolOperationKind::Write,
15143 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15144 requires_network: false,
15145 destructive: false,
15146 open_world: false,
15147 host_dependent: false,
15148 requires_user_interaction: false,
15149 supports_cancellation: true,
15150 default_requires_approval: false,
15151 should_defer_schema: false,
15152 max_output_chars: Some(1024),
15153 max_result_size_chars: Some(1024),
15154 }
15155 }
15156
15157 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15158 let mut classification =
15159 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15160 classification.safely_retryable = false;
15161 classification
15162 }
15163
15164 async fn execute(
15165 &self,
15166 _args: Value,
15167 _ctx: ai_agents_core::ToolExecutionContext,
15168 ) -> ToolResult {
15169 let call = self.calls.fetch_add(1, Ordering::SeqCst);
15170 if call == 0 {
15171 ToolResult::error("first failure")
15172 } else {
15173 ToolResult::ok("second success")
15174 }
15175 }
15176 }
15177
15178 #[async_trait]
15179 impl ai_agents_core::Tool for LockedWriteTool {
15180 fn id(&self) -> &str {
15181 "locked_write"
15182 }
15183
15184 fn name(&self) -> &str {
15185 "Locked Write"
15186 }
15187
15188 fn description(&self) -> &str {
15189 "Tracks concurrent execution on one resource."
15190 }
15191
15192 fn input_schema(&self) -> Value {
15193 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15194 }
15195
15196 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15197 ai_agents_core::ToolPolicyBindings {
15198 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15199 ..Default::default()
15200 }
15201 }
15202
15203 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15204 ai_agents_core::ToolSafetyMetadata {
15205 read_only: false,
15206 concurrency_safe: false,
15207 operation: ai_agents_core::ToolOperationKind::Write,
15208 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15209 requires_network: false,
15210 destructive: false,
15211 open_world: false,
15212 host_dependent: false,
15213 requires_user_interaction: false,
15214 supports_cancellation: true,
15215 default_requires_approval: false,
15216 should_defer_schema: false,
15217 max_output_chars: Some(1024),
15218 max_result_size_chars: Some(1024),
15219 }
15220 }
15221
15222 async fn execute(
15223 &self,
15224 _args: Value,
15225 _ctx: ai_agents_core::ToolExecutionContext,
15226 ) -> ToolResult {
15227 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15228 loop {
15229 let current_max = self.max_active.load(Ordering::SeqCst);
15230 if active <= current_max {
15231 break;
15232 }
15233 if self
15234 .max_active
15235 .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
15236 .is_ok()
15237 {
15238 break;
15239 }
15240 }
15241 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15242 self.active.fetch_sub(1, Ordering::SeqCst);
15243 ToolResult::ok("done")
15244 }
15245 }
15246
15247 #[async_trait]
15248 impl ai_agents_core::Tool for MultiResourceWriteTool {
15249 fn id(&self) -> &str {
15250 "multi_resource_write"
15251 }
15252
15253 fn name(&self) -> &str {
15254 "Multi Resource Write"
15255 }
15256
15257 fn description(&self) -> &str {
15258 "Tracks concurrent execution across source and destination resources."
15259 }
15260
15261 fn input_schema(&self) -> Value {
15262 serde_json::json!({"type": "object"})
15263 }
15264
15265 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15266 ai_agents_core::ToolPolicyBindings {
15267 path_fields: vec![
15268 ai_agents_core::PathPolicyBinding::read_write("source_path"),
15269 ai_agents_core::PathPolicyBinding::write("destination_path"),
15270 ],
15271 ..Default::default()
15272 }
15273 }
15274
15275 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15276 LockedWriteTool {
15277 active: Arc::clone(&self.active),
15278 max_active: Arc::clone(&self.max_active),
15279 }
15280 .safety_metadata()
15281 }
15282
15283 async fn execute(
15284 &self,
15285 _args: Value,
15286 _ctx: ai_agents_core::ToolExecutionContext,
15287 ) -> ToolResult {
15288 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15289 self.max_active.fetch_max(active, Ordering::SeqCst);
15290 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15291 self.active.fetch_sub(1, Ordering::SeqCst);
15292 ToolResult::ok("done")
15293 }
15294 }
15295
15296 #[async_trait]
15297 impl ai_agents_core::Tool for BlockingPathMutationTool {
15298 fn id(&self) -> &str {
15299 self.id
15300 }
15301
15302 fn name(&self) -> &str {
15303 self.id
15304 }
15305
15306 fn description(&self) -> &str {
15307 "Blocks a path mutation until the test releases it."
15308 }
15309
15310 fn input_schema(&self) -> Value {
15311 serde_json::json!({"type": "object"})
15312 }
15313
15314 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15315 ai_agents_core::ToolPolicyBindings {
15316 path_fields: self.path_fields.clone(),
15317 ..Default::default()
15318 }
15319 }
15320
15321 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15322 ai_agents_core::ToolSafetyMetadata {
15323 read_only: false,
15324 concurrency_safe: false,
15325 operation: ai_agents_core::ToolOperationKind::Write,
15326 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15327 requires_network: false,
15328 destructive: false,
15329 open_world: false,
15330 host_dependent: false,
15331 requires_user_interaction: false,
15332 supports_cancellation: true,
15333 default_requires_approval: false,
15334 should_defer_schema: false,
15335 max_output_chars: Some(1024),
15336 max_result_size_chars: Some(1024),
15337 }
15338 }
15339
15340 async fn execute(
15341 &self,
15342 _args: Value,
15343 _ctx: ai_agents_core::ToolExecutionContext,
15344 ) -> ToolResult {
15345 self.gate.entered.store(true, Ordering::SeqCst);
15346 self.gate.entered_notify.notify_one();
15347 self.gate.release.notified().await;
15348 ToolResult::ok("done")
15349 }
15350 }
15351
15352 #[async_trait]
15353 impl ai_agents_core::Tool for NoBindingWriteTool {
15354 fn id(&self) -> &str {
15355 "no_binding_write"
15356 }
15357
15358 fn name(&self) -> &str {
15359 "No Binding Write"
15360 }
15361
15362 fn description(&self) -> &str {
15363 "Tracks concurrent execution without resource bindings."
15364 }
15365
15366 fn input_schema(&self) -> Value {
15367 serde_json::json!({"type": "object"})
15368 }
15369
15370 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15371 LockedWriteTool {
15372 active: Arc::clone(&self.active),
15373 max_active: Arc::clone(&self.max_active),
15374 }
15375 .safety_metadata()
15376 }
15377
15378 async fn execute(
15379 &self,
15380 _args: Value,
15381 _ctx: ai_agents_core::ToolExecutionContext,
15382 ) -> ToolResult {
15383 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15384 self.max_active.fetch_max(active, Ordering::SeqCst);
15385 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15386 self.active.fetch_sub(1, Ordering::SeqCst);
15387 ToolResult::ok("done")
15388 }
15389 }
15390
15391 #[async_trait]
15392 impl ai_agents_core::Tool for RecoveryTestTool {
15393 fn id(&self) -> &str {
15394 &self.id
15395 }
15396
15397 fn name(&self) -> &str {
15398 &self.id
15399 }
15400
15401 fn description(&self) -> &str {
15402 "Records recovery execution and returns a configured result."
15403 }
15404
15405 fn input_schema(&self) -> Value {
15406 serde_json::json!({"type": "object"})
15407 }
15408
15409 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15410 ai_agents_core::ToolPolicyBindings {
15411 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15412 ..Default::default()
15413 }
15414 }
15415
15416 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15418 ai_agents_core::ToolSafetyMetadata {
15419 read_only: false,
15420 concurrency_safe: false,
15421 operation: ai_agents_core::ToolOperationKind::Write,
15422 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15423 requires_network: false,
15424 destructive: false,
15425 open_world: false,
15426 host_dependent: false,
15427 requires_user_interaction: false,
15428 supports_cancellation: true,
15429 default_requires_approval: false,
15430 should_defer_schema: false,
15431 max_output_chars: Some(self.max_output_chars.unwrap_or(1024)),
15432 max_result_size_chars: Some(1024),
15433 }
15434 }
15435
15436 async fn execute(
15438 &self,
15439 _args: Value,
15440 _ctx: ai_agents_core::ToolExecutionContext,
15441 ) -> ToolResult {
15442 self.calls.fetch_add(1, Ordering::SeqCst);
15443 let mut result = if self.succeeds {
15444 ToolResult::ok(format!("{} succeeded", self.id))
15445 } else {
15446 ToolResult::error(format!("{} failed", self.id))
15447 };
15448 result.metadata = Some(HashMap::from([(
15449 "recovery_test_tool".to_string(),
15450 Value::String(self.id.clone()),
15451 )]));
15452 result
15453 }
15454 }
15455
15456 #[async_trait]
15457 impl WebFetchTransport for RuntimeWebFetchTransport {
15458 async fn send(
15460 &self,
15461 _request: WebFetchTransportRequest,
15462 ) -> std::result::Result<WebFetchTransportResponse, String> {
15463 Err("validated addresses are required".to_string())
15464 }
15465
15466 async fn send_validated(
15468 &self,
15469 _request: WebFetchTransportRequest,
15470 _addresses: &[std::net::SocketAddr],
15471 ) -> std::result::Result<WebFetchTransportResponse, String> {
15472 self.calls.fetch_add(1, Ordering::SeqCst);
15473 Ok(WebFetchTransportResponse {
15474 status: 200,
15475 content_type: Some("text/plain".to_string()),
15476 location: None,
15477 body: b"approved".to_vec(),
15478 })
15479 }
15480 }
15481
15482 #[async_trait]
15483 impl WebFetchResolver for RuntimeWebFetchResolver {
15484 async fn resolve(
15486 &self,
15487 _host: &str,
15488 _port: u16,
15489 ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
15490 Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
15491 93, 184, 216, 34,
15492 ))])
15493 }
15494 }
15495
15496 #[async_trait]
15497 impl ApprovalHandler for BlockingApprovalHandler {
15498 async fn request_approval(
15499 &self,
15500 _request: ai_agents_hitl::ApprovalRequest,
15501 ) -> ApprovalResult {
15502 self.entered.wait().await;
15503 self.release.notified().await;
15504 self.result.clone()
15505 }
15506 }
15507
15508 #[async_trait]
15509 impl ApprovalHandler for CountingApprovalHandler {
15510 async fn request_approval(
15511 &self,
15512 _request: ai_agents_hitl::ApprovalRequest,
15513 ) -> ApprovalResult {
15514 self.calls.fetch_add(1, Ordering::SeqCst);
15515 ApprovalResult::Approved
15516 }
15517 }
15518
15519 #[async_trait]
15520 impl AgentHooks for ReentrantToolHooks {
15521 async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
15522 if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
15523 return;
15524 }
15525 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
15526 if let Some(agent) = agent {
15527 let result = agent
15528 .invoke_tool(ToolExecutionRequest::new(
15529 "nested-hook-call",
15530 "reentrant_write",
15531 serde_json::json!({"path": "./hook.txt"}),
15532 ToolCallSource::Manual,
15533 ))
15534 .await;
15535 self.nested_success
15536 .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
15537 }
15538 }
15539 }
15540
15541 #[async_trait]
15542 impl AgentHooks for ResponseCountingHooks {
15543 async fn on_response(&self, _response: &AgentResponse) {
15544 self.responses.fetch_add(1, Ordering::SeqCst);
15545 }
15546 }
15547
15548 #[async_trait]
15549 impl AgentHooks for ResponseChatHooks {
15550 async fn on_response(&self, _response: &AgentResponse) {
15552 if self.invoked.swap(true, Ordering::SeqCst) {
15553 return;
15554 }
15555 let target = self.target.lock().as_ref().and_then(Weak::upgrade);
15556 let result = if let Some(target) = target {
15557 target
15558 .chat("nested response hook call")
15559 .await
15560 .map(|response| response.content)
15561 .map_err(|error| error.to_string())
15562 } else {
15563 Err("response hook target is unavailable".to_string())
15564 };
15565 *self.nested_result.lock() = Some(result);
15566 }
15567 }
15568
15569 #[async_trait]
15570 impl AgentHooks for ConcurrentResponseHooks {
15571 async fn on_response(&self, _response: &AgentResponse) {
15573 if self.invoked.swap(true, Ordering::SeqCst) {
15574 return;
15575 }
15576 let Some(registry) = self.registry.upgrade() else {
15577 *self.nested_result.lock() =
15578 Some(Err("concurrent registry is unavailable".to_string()));
15579 return;
15580 };
15581 let agents = [ai_agents_state::ConcurrentAgentRef::Id(
15582 self.child_id.clone(),
15583 )];
15584 let aggregation = ai_agents_state::AggregationConfig {
15585 strategy: ai_agents_state::AggregationStrategy::FirstWins,
15586 synthesizer_llm: None,
15587 synthesizer_prompt: None,
15588 vote: None,
15589 };
15590 let result = crate::orchestration::concurrent(
15591 ®istry,
15592 "nested concurrent response hook call",
15593 &agents,
15594 &aggregation,
15595 None,
15596 Some(1),
15597 None,
15598 ai_agents_state::PartialFailureAction::Abort,
15599 None,
15600 )
15601 .await
15602 .map(|result| result.response.content)
15603 .map_err(|error| error.to_string());
15604 *self.nested_result.lock() = Some(result);
15605 }
15606 }
15607
15608 #[async_trait]
15609 impl AgentHooks for ToolLifecycleRecordingHooks {
15610 async fn on_tool_start(&self, tool: &str, _args: &Value) {
15611 self.events.lock().push(format!("start:{tool}"));
15612 }
15613
15614 async fn on_tool_complete(&self, tool: &str, result: &ToolResult, _duration_ms: u64) {
15615 self.events
15616 .lock()
15617 .push(format!("complete:{tool}:{}", result.success));
15618 }
15619
15620 async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
15621 self.events.lock().push(format!(
15622 "record:{}:{}",
15623 record.canonical_id, record.executed
15624 ));
15625 self.records.lock().push(record.clone());
15626 }
15627
15628 async fn on_error(&self, _error: &AgentError) {
15630 self.events.lock().push("error".to_string());
15631 }
15632 }
15633
15634 struct ApprovalRecordingHooks {
15635 events: parking_lot::Mutex<Vec<String>>,
15636 }
15637
15638 impl ApprovalRecordingHooks {
15639 fn new() -> Self {
15640 Self {
15641 events: parking_lot::Mutex::new(Vec::new()),
15642 }
15643 }
15644
15645 fn events(&self) -> Vec<String> {
15646 self.events.lock().clone()
15647 }
15648 }
15649
15650 #[async_trait]
15651 impl AgentHooks for ApprovalRecordingHooks {
15652 async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
15653 self.events.lock().push(format!(
15654 "raw:{}:{}",
15655 request_id,
15656 approval_result_name(result)
15657 ));
15658 }
15659
15660 async fn on_approval_resolved(
15661 &self,
15662 request: &ai_agents_hitl::ApprovalRequest,
15663 raw_result: &ApprovalResult,
15664 outcome: &ApprovalResolvedOutcome,
15665 ) {
15666 self.events.lock().push(format!(
15667 "resolved:{}:{}:{}",
15668 request.id,
15669 approval_result_name(raw_result),
15670 approval_outcome_name(outcome)
15671 ));
15672 }
15673 }
15674
15675 fn approval_result_name(result: &ApprovalResult) -> &'static str {
15676 match result {
15677 ApprovalResult::Approved => "approved",
15678 ApprovalResult::Rejected { .. } => "rejected",
15679 ApprovalResult::Modified { .. } => "modified",
15680 ApprovalResult::Timeout => "timeout",
15681 }
15682 }
15683
15684 fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
15685 match outcome {
15686 ApprovalResolvedOutcome::Approved => "approved",
15687 ApprovalResolvedOutcome::Rejected { .. } => "rejected",
15688 ApprovalResolvedOutcome::Modified { .. } => "modified",
15689 ApprovalResolvedOutcome::Error { .. } => "error",
15690 }
15691 }
15692
15693 fn assert_correlated_approval_events(
15694 events: &[String],
15695 raw_status: &str,
15696 outcome_status: &str,
15697 ) {
15698 assert_eq!(events.len(), 2);
15699 let raw: Vec<_> = events[0].split(':').collect();
15700 let resolved: Vec<_> = events[1].split(':').collect();
15701 assert_eq!(raw[0], "raw");
15702 assert_eq!(resolved[0], "resolved");
15703 assert_eq!(raw[1], resolved[1]);
15704 assert_eq!(raw[2], raw_status);
15705 assert_eq!(resolved[2], raw_status);
15706 assert_eq!(resolved[3], outcome_status);
15707 }
15708
15709 fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
15710 let mut security = ToolSecurityConfig {
15711 enabled: true,
15712 fail_closed: true,
15713 ..Default::default()
15714 };
15715 let policy = ai_agents_tools::ToolPolicyConfig {
15716 enabled: policy_enabled,
15717 write_paths: vec![".".to_string()],
15718 require_confirmation: true,
15719 ..Default::default()
15720 };
15721 security.tools.insert("locked_write".to_string(), policy);
15722 security
15723 }
15724
15725 struct MutationTestWorkspace {
15726 root: std::path::PathBuf,
15727 }
15728
15729 impl MutationTestWorkspace {
15730 fn new() -> Self {
15731 let root = std::env::temp_dir().join(format!(
15732 "ai-agents-runtime-mutation-{}",
15733 uuid::Uuid::new_v4()
15734 ));
15735 std::fs::create_dir_all(&root).unwrap();
15736 Self { root }
15737 }
15738 }
15739
15740 impl Drop for MutationTestWorkspace {
15741 fn drop(&mut self) {
15742 let _ = std::fs::remove_dir_all(&self.root);
15743 }
15744 }
15745
15746 async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
15747 tokio::time::timeout(std::time::Duration::from_secs(2), async {
15748 loop {
15749 let strong_count = locks
15750 .read()
15751 .get("path-mutation:global")
15752 .map_or(0, |lock| lock.strong_count());
15753 if strong_count >= minimum {
15754 break;
15755 }
15756 tokio::task::yield_now().await;
15757 }
15758 })
15759 .await
15760 .expect("path mutation call did not reach the shared lock");
15761 }
15762
15763 async fn assert_path_mutation_pair_serialized(
15764 first_id: &'static str,
15765 first_fields: Vec<ai_agents_core::PathPolicyBinding>,
15766 first_args: Value,
15767 second_id: &'static str,
15768 second_fields: Vec<ai_agents_core::PathPolicyBinding>,
15769 second_args: Value,
15770 ) {
15771 let locks = new_tool_resource_locks();
15772 let first_gate = PathMutationGate::new();
15773 let second_gate = PathMutationGate::new();
15774 second_gate.release();
15775 let agent = Arc::new(
15776 AgentBuilder::new()
15777 .system_prompt("Test global path mutation locking.")
15778 .llm(Arc::new(mock_with_response("done")))
15779 .tool(Arc::new(BlockingPathMutationTool {
15780 id: first_id,
15781 path_fields: first_fields,
15782 gate: first_gate.clone(),
15783 }))
15784 .tool(Arc::new(BlockingPathMutationTool {
15785 id: second_id,
15786 path_fields: second_fields,
15787 gate: second_gate.clone(),
15788 }))
15789 .build()
15790 .unwrap()
15791 .with_shared_resource_locks(Arc::clone(&locks)),
15792 );
15793
15794 let first = {
15795 let agent = Arc::clone(&agent);
15796 tokio::spawn(async move {
15797 agent
15798 .invoke_tool(ToolExecutionRequest::new(
15799 format!("{}-first", first_id),
15800 first_id,
15801 first_args,
15802 ToolCallSource::Manual,
15803 ))
15804 .await
15805 .unwrap()
15806 })
15807 };
15808 first_gate.wait_until_entered().await;
15809
15810 let second = {
15811 let agent = Arc::clone(&agent);
15812 tokio::spawn(async move {
15813 agent
15814 .invoke_tool(ToolExecutionRequest::new(
15815 format!("{}-second", second_id),
15816 second_id,
15817 second_args,
15818 ToolCallSource::Manual,
15819 ))
15820 .await
15821 .unwrap()
15822 })
15823 };
15824 wait_for_resource_lock_strong_count(&locks, 2).await;
15825 assert!(!second_gate.entered.load(Ordering::SeqCst));
15826 assert!(!second.is_finished());
15827
15828 first_gate.release();
15829 let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15830 tokio::join!(first, second)
15831 })
15832 .await
15833 .expect("serialized path mutation calls did not finish");
15834 assert!(first.unwrap().success);
15835 assert!(second.unwrap().success);
15836 assert!(second_gate.entered.load(Ordering::SeqCst));
15837 assert!(locks.read().is_empty());
15838 }
15839
15840 #[derive(Clone, Copy)]
15841 enum MutationDenial {
15842 Policy,
15843 Approval,
15844 }
15845
15846 fn mutation_denial_security_config(
15847 tool_id: &str,
15848 workspace: &std::path::Path,
15849 denial: MutationDenial,
15850 ) -> ToolSecurityConfig {
15851 let workspace = workspace.to_string_lossy().into_owned();
15852 let mut policy = ai_agents_tools::ToolPolicyConfig {
15853 read_paths: vec![workspace.clone()],
15854 write_paths: vec![workspace.clone()],
15855 ..Default::default()
15856 };
15857 match denial {
15858 MutationDenial::Policy => policy.blocked_paths = vec![workspace],
15859 MutationDenial::Approval => policy.require_confirmation = true,
15860 }
15861
15862 let mut security = ToolSecurityConfig {
15863 enabled: true,
15864 fail_closed: true,
15865 ..Default::default()
15866 };
15867 security.tools.insert(tool_id.to_string(), policy);
15868 security
15869 }
15870
15871 async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
15872 let workspace = MutationTestWorkspace::new();
15873 let tool_id = tool.id().to_string();
15874 let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
15875 let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
15876 std::fs::write(&preserved, "preserved").unwrap();
15877 let arguments = match tool_id.as_str() {
15878 "copy_path" | "move_path" => serde_json::json!({
15879 "source_path": preserved.to_string_lossy(),
15880 "destination_path": destination.to_string_lossy(),
15881 "dry_run": false
15882 }),
15883 "delete_path" => serde_json::json!({
15884 "path": preserved.to_string_lossy(),
15885 "recursive": false,
15886 "dry_run": false
15887 }),
15888 _ => panic!("unsupported mutation tool: {}", tool_id),
15889 };
15890 let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
15891 let builder = AgentBuilder::new()
15892 .system_prompt("Test mutation denial.")
15893 .llm(Arc::new(mock_with_response("done")))
15894 .tool(tool)
15895 .tool_security(ToolSecurityEngine::new(security));
15896 let builder = match denial {
15897 MutationDenial::Policy => builder,
15898 MutationDenial::Approval => builder
15899 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15900 .approval_handler(Arc::new(RejectAllHandler::new())),
15901 };
15902 let agent = builder.build().unwrap();
15903
15904 let record = agent
15905 .invoke_tool(ToolExecutionRequest::new(
15906 format!("{}-denied", tool_id),
15907 tool_id.clone(),
15908 arguments,
15909 ToolCallSource::Manual,
15910 ))
15911 .await
15912 .unwrap();
15913
15914 assert!(!record.executed, "{} must not be invoked", tool_id);
15915 assert!(!record.success);
15916 match denial {
15917 MutationDenial::Policy => {
15918 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
15919 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15920 &approval.status,
15921 ToolApprovalStatus::NotRequired
15922 )));
15923 }
15924 MutationDenial::Approval => {
15925 assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
15926 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15927 &approval.status,
15928 ToolApprovalStatus::Rejected
15929 )));
15930 }
15931 }
15932 assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
15933 assert!(!destination.exists());
15934 }
15935
15936 fn recovery_manager_with_fallbacks(
15937 fallbacks: impl IntoIterator<Item = (String, String)>,
15938 ) -> RecoveryManager {
15939 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15940
15941 let per_tool = fallbacks
15942 .into_iter()
15943 .map(|(tool, fallback_tool)| {
15944 (
15945 tool,
15946 ToolRetryConfig {
15947 max_retries: 0,
15948 timeout_ms: Some(1_000),
15949 on_failure: ToolFailureAction::Fallback { fallback_tool },
15950 },
15951 )
15952 })
15953 .collect();
15954 RecoveryManager::new(ErrorRecoveryConfig {
15955 tools: ToolRecoveryConfig {
15956 per_tool,
15957 ..Default::default()
15958 },
15959 ..Default::default()
15960 })
15961 }
15962
15963 fn approval_check() -> HITLCheckResult {
15964 HITLCheckResult::required(
15965 ApprovalTrigger::tool("test", serde_json::json!({})),
15966 HashMap::new(),
15967 "Approve?",
15968 None,
15969 )
15970 }
15971
15972 fn agent_with_approval_result(
15973 raw_result: ApprovalResult,
15974 timeout_action: TimeoutAction,
15975 hooks: Arc<ApprovalRecordingHooks>,
15976 ) -> RuntimeAgent {
15977 use ai_agents_hitl::{CallbackHandler, HITLConfig};
15978
15979 let config = HITLConfig {
15980 on_timeout: timeout_action,
15981 ..Default::default()
15982 };
15983 let handler = CallbackHandler::new(move |_| raw_result.clone());
15984 AgentBuilder::new()
15985 .system_prompt("Test HITL hooks.")
15986 .llm(Arc::new(mock_with_response("done")))
15987 .build()
15988 .unwrap()
15989 .with_hooks(hooks)
15990 .with_hitl(HITLEngine::new(config), Arc::new(handler))
15991 }
15992
15993 #[tokio::test]
15994 async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
15995 let cases = vec![
15996 (ApprovalResult::Approved, "approved"),
15997 (
15998 ApprovalResult::Rejected {
15999 reason: Some("denied".to_string()),
16000 },
16001 "rejected",
16002 ),
16003 (
16004 ApprovalResult::Modified {
16005 changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
16006 },
16007 "modified",
16008 ),
16009 ];
16010
16011 for (raw_result, expected) in cases {
16012 let hooks = Arc::new(ApprovalRecordingHooks::new());
16013 let agent =
16014 agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
16015
16016 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16017
16018 assert_eq!(approval_result_name(&result), expected);
16019 assert_correlated_approval_events(&hooks.events(), expected, expected);
16020 }
16021 }
16022
16023 #[tokio::test]
16024 async fn approval_hooks_expose_timeout_policy_decisions() {
16025 for (timeout_action, expected) in [
16026 (TimeoutAction::Approve, "approved"),
16027 (TimeoutAction::Reject, "rejected"),
16028 ] {
16029 let hooks = Arc::new(ApprovalRecordingHooks::new());
16030 let agent =
16031 agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
16032
16033 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16034
16035 assert_eq!(approval_result_name(&result), expected);
16036 assert_correlated_approval_events(&hooks.events(), "timeout", expected);
16037 }
16038 }
16039
16040 #[tokio::test]
16041 async fn timeout_error_fires_correlated_resolved_error_before_returning() {
16042 let hooks = Arc::new(ApprovalRecordingHooks::new());
16043 let agent = agent_with_approval_result(
16044 ApprovalResult::Timeout,
16045 TimeoutAction::Error,
16046 hooks.clone(),
16047 );
16048
16049 let error = agent
16050 .request_hitl_approval(approval_check())
16051 .await
16052 .unwrap_err();
16053
16054 assert!(error.to_string().contains("HITL approval timeout"));
16055 assert_correlated_approval_events(&hooks.events(), "timeout", "error");
16056 }
16057
16058 #[tokio::test]
16060 async fn test_integration_yaml_to_chat_basic() {
16061 let mock = mock_with_response("Hello! How can I help you?");
16062 let agent = AgentBuilder::new()
16063 .system_prompt("You are a test assistant.")
16064 .llm(Arc::new(mock))
16065 .build()
16066 .unwrap();
16067
16068 let response = agent.chat("Hi").await.unwrap();
16069 assert!(!response.content.is_empty());
16070 assert_eq!(response.content, "Hello! How can I help you?");
16071 }
16072
16073 #[tokio::test]
16074 async fn stream_events_emit_one_authoritative_final_without_legacy_done() {
16075 let agent = AgentBuilder::new()
16076 .system_prompt("You are a test assistant.")
16077 .llm(Arc::new(mock_with_response(
16078 "Hello from the final response.",
16079 )))
16080 .build()
16081 .unwrap();
16082
16083 let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16084 let mut final_responses = Vec::new();
16085 let mut legacy_done = 0;
16086 while let Some(event) = stream.next().await {
16087 match event {
16088 AgentStreamEvent::Chunk(StreamChunk::Done {}) => legacy_done += 1,
16089 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16090 panic!("unexpected stream error: {message}")
16091 }
16092 AgentStreamEvent::Final(response) => final_responses.push(response),
16093 AgentStreamEvent::Chunk(_) => {}
16094 }
16095 }
16096
16097 assert_eq!(legacy_done, 0);
16098 assert_eq!(final_responses.len(), 1);
16099 let response = final_responses.pop().unwrap();
16100 assert_eq!(response.content, "Hello from the final response.");
16101 assert!(
16102 response
16103 .metadata
16104 .as_ref()
16105 .is_some_and(|metadata| { metadata.contains_key("reasoning") })
16106 );
16107 }
16108
16109 #[tokio::test]
16110 async fn stream_final_content_includes_output_processing_after_provisional_chunks() {
16111 let yaml = r#"
16112name: ProcessedStreamAgent
16113system_prompt: "Answer directly."
16114process:
16115 output:
16116 - type: format
16117 config:
16118 template: "{{ response }} [finalized]"
16119streaming:
16120 enabled: true
16121"#;
16122 let agent = AgentBuilder::from_yaml(yaml)
16123 .unwrap()
16124 .llm(Arc::new(mock_with_response("provisional answer")))
16125 .auto_configure_features()
16126 .unwrap()
16127 .build()
16128 .unwrap();
16129
16130 let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16131 let mut provisional = String::new();
16132 let mut final_content = None;
16133 while let Some(event) = stream.next().await {
16134 match event {
16135 AgentStreamEvent::Chunk(StreamChunk::Content { text }) => {
16136 provisional.push_str(&text)
16137 }
16138 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16139 panic!("unexpected stream error: {message}")
16140 }
16141 AgentStreamEvent::Final(response) => final_content = Some(response.content),
16142 AgentStreamEvent::Chunk(_) => {}
16143 }
16144 }
16145
16146 assert_eq!(provisional, "provisional answer");
16147 assert_eq!(
16148 final_content.as_deref(),
16149 Some("provisional answer [finalized]")
16150 );
16151 }
16152
16153 #[tokio::test]
16154 async fn stream_events_preserve_tool_progress_and_final_tool_calls() {
16155 let agent = AgentBuilder::new()
16156 .system_prompt("Use the echo tool once, then answer.")
16157 .llm(Arc::new(mock_with_responses(vec![
16158 r#"{"tool":"echo","arguments":{"message":"hello"}}"#,
16159 "Echo completed.",
16160 ])))
16161 .tool(Arc::new(ai_agents_tools::EchoTool::new()))
16162 .build()
16163 .unwrap();
16164
16165 let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
16166 let mut starts = 0;
16167 let mut results = 0;
16168 let mut ends = 0;
16169 let mut final_response = None;
16170 while let Some(event) = stream.next().await {
16171 match event {
16172 AgentStreamEvent::Chunk(StreamChunk::ToolCallStart { name, .. }) => {
16173 assert_eq!(name, "echo");
16174 starts += 1;
16175 }
16176 AgentStreamEvent::Chunk(StreamChunk::ToolResult { name, success, .. }) => {
16177 assert_eq!(name, "echo");
16178 assert!(success);
16179 results += 1;
16180 }
16181 AgentStreamEvent::Chunk(StreamChunk::ToolCallEnd { .. }) => ends += 1,
16182 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16183 panic!("unexpected stream error: {message}")
16184 }
16185 AgentStreamEvent::Final(response) => final_response = Some(response),
16186 AgentStreamEvent::Chunk(_) => {}
16187 }
16188 }
16189
16190 assert_eq!((starts, results, ends), (1, 1, 1));
16191 let response = final_response.expect("tool stream must finalize");
16192 assert_eq!(response.content, "Echo completed.");
16193 assert_eq!(
16194 response.tool_calls.as_ref().map(|calls| calls
16195 .iter()
16196 .map(|call| call.name.as_str())
16197 .collect::<Vec<_>>()),
16198 Some(vec!["echo"])
16199 );
16200 }
16201
16202 #[tokio::test]
16203 async fn legacy_stream_still_emits_one_done_chunk() {
16204 let agent = AgentBuilder::new()
16205 .system_prompt("You are a test assistant.")
16206 .llm(Arc::new(mock_with_response(
16207 "Hello from the legacy stream.",
16208 )))
16209 .build()
16210 .unwrap();
16211
16212 let mut stream = agent.chat_stream("Hi").await.unwrap();
16213 let mut done = 0;
16214 while let Some(chunk) = stream.next().await {
16215 match chunk {
16216 StreamChunk::Done {} => done += 1,
16217 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
16218 _ => {}
16219 }
16220 }
16221
16222 assert_eq!(done, 1);
16223 }
16224
16225 #[tokio::test]
16227 async fn test_integration_multi_turn_conversation() {
16228 let mock = mock_with_responses(vec![
16229 "Hello! I'm your assistant.",
16230 "The weather is sunny today.",
16231 "Goodbye!",
16232 ]);
16233 let agent = AgentBuilder::new()
16234 .system_prompt("You are helpful.")
16235 .llm(Arc::new(mock))
16236 .build()
16237 .unwrap();
16238
16239 let r1 = agent.chat("Hi").await.unwrap();
16240 assert_eq!(r1.content, "Hello! I'm your assistant.");
16241
16242 let r2 = agent.chat("What's the weather?").await.unwrap();
16243 assert_eq!(r2.content, "The weather is sunny today.");
16244
16245 let r3 = agent.chat("Bye").await.unwrap();
16246 assert_eq!(r3.content, "Goodbye!");
16247
16248 let messages = agent.memory.get_messages(None).await.unwrap();
16250 assert_eq!(messages.len(), 6);
16252 }
16253
16254 #[test]
16255 fn later_approval_preserves_modified_evidence() {
16256 let arguments = serde_json::json!({"dry_run": true});
16257 let mut record = Some(ToolApprovalRecord {
16258 status: ToolApprovalStatus::Modified,
16259 reason: None,
16260 modified_arguments: Some(arguments.clone()),
16261 });
16262
16263 merge_approved_record(&mut record);
16264
16265 let record = record.unwrap();
16266 assert!(matches!(record.status, ToolApprovalStatus::Modified));
16267 assert_eq!(record.modified_arguments, Some(arguments));
16268 }
16269
16270 #[test]
16271 fn approval_binding_rejects_replaced_tool_implementation() {
16272 let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16273 let same_tool = Arc::clone(&reviewed_tool);
16274 let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16275 let arguments = serde_json::json!({"path": "."});
16276 let versions = ToolDecisionVersions {
16277 policy: 2,
16278 registry: 3,
16279 runtime_control: 4,
16280 state: Some(5),
16281 };
16282 let binding = ToolApprovalBinding {
16283 canonical_id: "context_echo".to_string(),
16284 arguments: arguments.clone(),
16285 confirmation_required: true,
16286 policy_version: versions.policy,
16287 runtime_control_version: versions.runtime_control,
16288 state_generation: versions.state,
16289 reviewed_tool,
16290 };
16291
16292 assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
16293 assert!(binding.is_stale(
16294 "context_echo",
16295 &arguments,
16296 true,
16297 versions,
16298 &replacement_tool,
16299 ));
16300 }
16301
16302 #[tokio::test]
16303 async fn approved_mutation_to_dry_run_remains_executable() {
16304 use ai_agents_hitl::CallbackHandler;
16305
16306 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16307 changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
16308 });
16309 let agent = AgentBuilder::new()
16310 .system_prompt("Test safer approval modifications.")
16311 .llm(Arc::new(mock_with_response("done")))
16312 .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
16313 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16314 .approval_handler(Arc::new(handler))
16315 .build()
16316 .unwrap();
16317
16318 let record = agent
16319 .invoke_tool(ToolExecutionRequest::new(
16320 "approved-dry-run",
16321 "file_write",
16322 serde_json::json!({
16323 "path": "./approval-dry-run.txt",
16324 "content": "not written"
16325 }),
16326 ToolCallSource::Manual,
16327 ))
16328 .await
16329 .unwrap();
16330
16331 assert!(record.executed);
16332 assert!(record.success);
16333 assert_eq!(record.executed_arguments["dry_run"], true);
16334 assert!(matches!(
16335 record.approval.as_ref().map(|approval| &approval.status),
16336 Some(ToolApprovalStatus::Modified)
16337 ));
16338 let output: Value = serde_json::from_str(&record.output).unwrap();
16339 assert_eq!(output["mutation_performed"], false);
16340 }
16341
16342 #[tokio::test]
16344 async fn shared_executor_approval_reaches_web_fetch_transport() {
16345 use ai_agents_hitl::{CallbackHandler, HITLConfig};
16346 use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
16347
16348 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16349 let tool = WebFetchTool::with_transport_and_resolver(
16350 Arc::new(RuntimeWebFetchTransport {
16351 calls: Arc::clone(&calls),
16352 }),
16353 Arc::new(RuntimeWebFetchResolver),
16354 );
16355 let mut security = ToolSecurityConfig {
16356 enabled: true,
16357 fail_closed: true,
16358 ..Default::default()
16359 };
16360 security.tools.insert(
16361 "web_fetch".to_string(),
16362 ToolPolicyConfig {
16363 domains: DomainPolicyConfig {
16364 requires_approval: vec!["approval.test".to_string()],
16365 ..Default::default()
16366 },
16367 allowed_schemes: vec!["https".to_string()],
16368 allowed_ports: vec![443],
16369 ..Default::default()
16370 },
16371 );
16372 let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
16373 let agent = AgentBuilder::new()
16374 .system_prompt("Test approved web fetch execution.")
16375 .llm(Arc::new(mock_with_response("done")))
16376 .tool(Arc::new(tool))
16377 .tool_security(ToolSecurityEngine::new(security))
16378 .build()
16379 .unwrap()
16380 .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
16381
16382 let record = agent
16383 .invoke_tool(ToolExecutionRequest::new(
16384 "approved-web-fetch",
16385 "web_fetch",
16386 serde_json::json!({
16387 "url": "https://approval.test/page",
16388 "cache_ttl_seconds": 0
16389 }),
16390 ToolCallSource::Manual,
16391 ))
16392 .await
16393 .unwrap();
16394
16395 assert!(record.success);
16396 assert!(
16397 record
16398 .approval
16399 .as_ref()
16400 .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
16401 );
16402 assert_eq!(calls.load(Ordering::SeqCst), 1);
16403 }
16404
16405 #[tokio::test]
16406 async fn context_preserves_requested_and_canonical_identity() {
16407 let mock = mock_with_response("hello");
16408 let mut tools = ai_agents_tools::ToolRegistry::new();
16409 tools.register(Arc::new(ContextEchoTool)).unwrap();
16410
16411 let mut security = ToolSecurityConfig {
16412 enabled: true,
16413 fail_closed: true,
16414 ..Default::default()
16415 };
16416 let mut policy = ai_agents_tools::ToolPolicyConfig {
16417 read_paths: vec![".".to_string()],
16418 max_results: Some(7),
16419 ..Default::default()
16420 };
16421 policy
16422 .config
16423 .insert("backend".to_string(), serde_json::json!("memory"));
16424 security.tools.insert("context_echo".to_string(), policy);
16425
16426 let agent = AgentBuilder::new()
16427 .system_prompt("You are helpful.")
16428 .llm(Arc::new(mock))
16429 .tools(tools)
16430 .tool_security(ToolSecurityEngine::new(security))
16431 .build()
16432 .unwrap();
16433
16434 let record = agent
16435 .invoke_tool(ToolExecutionRequest::new(
16436 "ctx-call",
16437 "Context Echo",
16438 serde_json::json!({"path": ".", "max_results": 99}),
16439 ToolCallSource::Manual,
16440 ))
16441 .await
16442 .unwrap();
16443
16444 assert!(record.success);
16445 assert!(matches!(&record.source, ToolCallSource::Manual));
16446 assert_eq!(record.requested_name, "Context Echo");
16447 assert_eq!(record.canonical_id, "context_echo");
16448 assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
16449 assert_eq!(record.executed_arguments["max_results"], 7);
16450 let output: Value = serde_json::from_str(&record.output).unwrap();
16451 assert_eq!(output["requested_name"], "Context Echo");
16452 assert_eq!(output["canonical_id"], "context_echo");
16453 assert_eq!(output["max_results"], 7);
16454 assert_eq!(output["custom_config"]["backend"], "memory");
16455 assert!(record.metadata.contains_key("effective_limits"));
16456 assert!(record.metadata.contains_key("policy_snapshot"));
16457 }
16458
16459 #[tokio::test]
16460 async fn test_runtime_control_cancels_active_tool_call() {
16461 let mock = mock_with_response("hello");
16462 let agent = Arc::new(
16463 AgentBuilder::new()
16464 .system_prompt("You are helpful.")
16465 .llm(Arc::new(mock))
16466 .tool(Arc::new(SlowTool))
16467 .build()
16468 .unwrap(),
16469 );
16470 let control = agent.runtime_control();
16471 let running_agent = Arc::clone(&agent);
16472 let handle = tokio::spawn(async move {
16473 running_agent
16474 .invoke_tool(ToolExecutionRequest::new(
16475 "slow-call",
16476 "slow",
16477 serde_json::json!({}),
16478 ToolCallSource::Manual,
16479 ))
16480 .await
16481 .unwrap()
16482 });
16483
16484 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
16485 control.cancel_all();
16486 let record = handle.await.unwrap();
16487
16488 assert!(record.executed);
16489 assert!(record.cancelled);
16490 assert!(!record.success);
16491 assert!(record.cancellation_reason.is_some());
16492 }
16493
16494 #[tokio::test]
16495 async fn non_idempotent_tool_calls_are_not_retried() {
16496 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
16497
16498 let mock = mock_with_response("hello");
16499 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16500 let agent = AgentBuilder::new()
16501 .system_prompt("You are helpful.")
16502 .llm(Arc::new(mock))
16503 .tool(Arc::new(FlakyWriteTool {
16504 calls: Arc::clone(&calls),
16505 }))
16506 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
16507 tools: ToolRecoveryConfig {
16508 default: ToolRetryConfig {
16509 max_retries: 2,
16510 ..Default::default()
16511 },
16512 ..Default::default()
16513 },
16514 ..Default::default()
16515 }))
16516 .build()
16517 .unwrap();
16518
16519 let record = agent
16520 .invoke_tool(ToolExecutionRequest::new(
16521 "flaky-call",
16522 "flaky_write",
16523 serde_json::json!({"path": "./tmp.txt"}),
16524 ToolCallSource::Manual,
16525 ))
16526 .await
16527 .unwrap();
16528
16529 assert!(!record.success);
16530 assert_eq!(calls.load(Ordering::SeqCst), 1);
16531 }
16532
16533 #[tokio::test]
16534 async fn safely_retryable_tool_receives_a_fresh_deadline_per_attempt() {
16535 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
16536
16537 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16538 let deadlines = Arc::new(parking_lot::Mutex::new(Vec::new()));
16539 let agent = AgentBuilder::new()
16540 .system_prompt("Test retry deadlines.")
16541 .llm(Arc::new(mock_with_response("done")))
16542 .tool(Arc::new(RetryDeadlineTool {
16543 calls: Arc::clone(&calls),
16544 deadlines: Arc::clone(&deadlines),
16545 }))
16546 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
16547 tools: ToolRecoveryConfig {
16548 per_tool: HashMap::from([(
16549 "retry_deadline".to_string(),
16550 ToolRetryConfig {
16551 max_retries: 1,
16552 ..Default::default()
16553 },
16554 )]),
16555 ..Default::default()
16556 },
16557 ..Default::default()
16558 }))
16559 .build()
16560 .unwrap();
16561
16562 let record = agent
16563 .invoke_tool(ToolExecutionRequest::new(
16564 "retry-deadline-call",
16565 "retry_deadline",
16566 serde_json::json!({}),
16567 ToolCallSource::Manual,
16568 ))
16569 .await
16570 .unwrap();
16571
16572 assert!(record.executed);
16573 assert!(record.success);
16574 assert_eq!(calls.load(Ordering::SeqCst), 2);
16575 let deadlines = deadlines.lock();
16576 assert_eq!(deadlines.len(), 2);
16577 assert!(
16578 deadlines[1] > deadlines[0],
16579 "retry inherited the first invocation deadline"
16580 );
16581 }
16582
16583 #[tokio::test]
16584 async fn side_effecting_tools_are_serialized_per_resource() {
16585 let mock = mock_with_response("hello");
16586 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16587 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16588 let agent = Arc::new(
16589 AgentBuilder::new()
16590 .system_prompt("You are helpful.")
16591 .llm(Arc::new(mock))
16592 .tool(Arc::new(LockedWriteTool {
16593 active: Arc::clone(&active),
16594 max_active: Arc::clone(&max_active),
16595 }))
16596 .build()
16597 .unwrap(),
16598 );
16599
16600 let left = {
16601 let agent = Arc::clone(&agent);
16602 tokio::spawn(async move {
16603 agent
16604 .invoke_tool(ToolExecutionRequest::new(
16605 "lock-1",
16606 "locked_write",
16607 serde_json::json!({"path": "./same.txt"}),
16608 ToolCallSource::Manual,
16609 ))
16610 .await
16611 .unwrap()
16612 })
16613 };
16614 let right = {
16615 let agent = Arc::clone(&agent);
16616 tokio::spawn(async move {
16617 agent
16618 .invoke_tool(ToolExecutionRequest::new(
16619 "lock-2",
16620 "locked_write",
16621 serde_json::json!({"path": "./same.txt"}),
16622 ToolCallSource::Manual,
16623 ))
16624 .await
16625 .unwrap()
16626 })
16627 };
16628
16629 let left = left.await.unwrap();
16630 let right = right.await.unwrap();
16631 assert!(left.success);
16632 assert!(right.success);
16633 assert_eq!(max_active.load(Ordering::SeqCst), 1);
16634 }
16635
16636 #[tokio::test]
16637 async fn path_resources_use_shared_global_lock_and_cleanup() {
16638 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16639 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16640 let bindings = ai_agents_core::ToolPolicyBindings {
16641 path_fields: vec![
16642 ai_agents_core::PathPolicyBinding::read_write("source_path"),
16643 ai_agents_core::PathPolicyBinding::write("destination_path"),
16644 ],
16645 ..Default::default()
16646 };
16647 let classification = ai_agents_core::ToolCallClassification::from_metadata(
16648 &MultiResourceWriteTool {
16649 active: Arc::clone(&active),
16650 max_active: Arc::clone(&max_active),
16651 }
16652 .safety_metadata(),
16653 );
16654 let left_args = serde_json::json!({
16655 "source_path": "./a/../first.txt",
16656 "destination_path": "./second.txt"
16657 });
16658 let right_args = serde_json::json!({
16659 "source_path": "./second.txt",
16660 "destination_path": "./first.txt"
16661 });
16662 let left_keys = tool_resource_lock_keys(
16663 "multi_resource_write",
16664 &left_args,
16665 &bindings,
16666 &classification,
16667 );
16668 let right_keys = tool_resource_lock_keys(
16669 "multi_resource_write",
16670 &right_args,
16671 &bindings,
16672 &classification,
16673 );
16674 assert_eq!(left_keys, right_keys);
16675 assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
16676
16677 let locks = new_tool_resource_locks();
16678 let build_agent = || {
16679 AgentBuilder::new()
16680 .system_prompt("Test shared resource locks.")
16681 .llm(Arc::new(mock_with_response("done")))
16682 .tool(Arc::new(MultiResourceWriteTool {
16683 active: Arc::clone(&active),
16684 max_active: Arc::clone(&max_active),
16685 }))
16686 .build()
16687 .unwrap()
16688 .with_shared_resource_locks(Arc::clone(&locks))
16689 };
16690 let left_agent = Arc::new(build_agent());
16691 let right_agent = Arc::new(build_agent());
16692 let left = tokio::spawn(async move {
16693 left_agent
16694 .invoke_tool(ToolExecutionRequest::new(
16695 "multi-left",
16696 "multi_resource_write",
16697 left_args,
16698 ToolCallSource::Manual,
16699 ))
16700 .await
16701 .unwrap()
16702 });
16703 let right = tokio::spawn(async move {
16704 right_agent
16705 .invoke_tool(ToolExecutionRequest::new(
16706 "multi-right",
16707 "multi_resource_write",
16708 right_args,
16709 ToolCallSource::Manual,
16710 ))
16711 .await
16712 .unwrap()
16713 });
16714 let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
16715 tokio::join!(left, right)
16716 })
16717 .await
16718 .expect("reversed resource acquisition must not deadlock");
16719
16720 assert!(left.unwrap().success);
16721 assert!(right.unwrap().success);
16722 assert_eq!(max_active.load(Ordering::SeqCst), 1);
16723 assert!(locks.read().is_empty());
16724 }
16725
16726 #[tokio::test]
16727 async fn global_path_lock_serializes_copy_destination_with_file_write() {
16728 assert_path_mutation_pair_serialized(
16729 "copy_path",
16730 CopyPathTool::new().policy_bindings().path_fields,
16731 serde_json::json!({
16732 "source_path": "./source.txt",
16733 "destination_path": "./shared.txt"
16734 }),
16735 "file_write",
16736 FileWriteTool::new().policy_bindings().path_fields,
16737 serde_json::json!({"path": "./shared.txt"}),
16738 )
16739 .await;
16740 }
16741
16742 #[tokio::test]
16743 async fn parent_and_spawned_runtime_share_global_path_lock() {
16744 let workspace = MutationTestWorkspace::new();
16745 let destination = workspace.root.join("spawned.txt");
16746 let parent_gate = PathMutationGate::new();
16747 let parent = Arc::new(
16748 AgentBuilder::from_yaml(
16749 r#"
16750name: LockParent
16751system_prompt: parent
16752llm:
16753 default: default
16754tools:
16755 - parent_path_write
16756spawner:
16757 shared_llms: true
16758"#,
16759 )
16760 .unwrap()
16761 .llm(Arc::new(mock_with_response("done")))
16762 .auto_configure_spawner()
16763 .await
16764 .unwrap()
16765 .tool(Arc::new(BlockingPathMutationTool {
16766 id: "parent_path_write",
16767 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16768 gate: parent_gate.clone(),
16769 }))
16770 .build()
16771 .unwrap(),
16772 );
16773
16774 let mut child_spec = crate::spec::AgentSpec {
16775 name: "LockChild".to_string(),
16776 system_prompt: "child".to_string(),
16777 tools: Some(vec![crate::spec::ToolEntry::Simple(
16778 "file_write".to_string(),
16779 )]),
16780 ..Default::default()
16781 };
16782 child_spec.tool_security.enabled = true;
16783 child_spec.tool_security.fail_closed = true;
16784 let file_write_policy = ai_agents_tools::ToolPolicyConfig {
16785 write_paths: vec![workspace.root.to_string_lossy().into_owned()],
16786 allow_without_confirmation: true,
16787 ..Default::default()
16788 };
16789 child_spec
16790 .tool_security
16791 .tools
16792 .insert("file_write".to_string(), file_write_policy);
16793 let spawned = parent
16794 .spawner()
16795 .unwrap()
16796 .spawn_from_spec(child_spec)
16797 .await
16798 .unwrap();
16799 assert!(Arc::ptr_eq(
16800 &parent.resource_locks,
16801 &spawned.agent.resource_locks
16802 ));
16803 assert!(!Arc::ptr_eq(
16804 &parent.runtime_control,
16805 &spawned.agent.runtime_control
16806 ));
16807
16808 let parent_call = {
16809 let parent = Arc::clone(&parent);
16810 let destination = destination.clone();
16811 tokio::spawn(async move {
16812 parent
16813 .invoke_tool(ToolExecutionRequest::new(
16814 "parent-lock-holder",
16815 "parent_path_write",
16816 serde_json::json!({"path": destination}),
16817 ToolCallSource::Manual,
16818 ))
16819 .await
16820 .unwrap()
16821 })
16822 };
16823 parent_gate.wait_until_entered().await;
16824
16825 let child_call = {
16826 let child = Arc::clone(&spawned.agent);
16827 let destination = destination.clone();
16828 tokio::spawn(async move {
16829 child
16830 .invoke_tool(ToolExecutionRequest::new(
16831 "spawned-file-write",
16832 "file_write",
16833 serde_json::json!({
16834 "path": destination,
16835 "content": "spawned",
16836 "dry_run": false
16837 }),
16838 ToolCallSource::Manual,
16839 ))
16840 .await
16841 .unwrap()
16842 })
16843 };
16844 wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
16845 assert!(!child_call.is_finished());
16846
16847 parent_gate.release();
16848 let (parent_record, child_record) =
16849 tokio::time::timeout(std::time::Duration::from_secs(2), async {
16850 tokio::join!(parent_call, child_call)
16851 })
16852 .await
16853 .expect("parent and spawned path mutations did not finish");
16854 assert!(parent_record.unwrap().success);
16855 assert!(child_record.unwrap().success);
16856 assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
16857 assert!(parent.resource_locks.read().is_empty());
16858 }
16859
16860 #[tokio::test]
16861 async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
16862 let locks = new_tool_resource_locks();
16863 let holder_gate = PathMutationGate::new();
16864 let waiter_gate = PathMutationGate::new();
16865 waiter_gate.release();
16866 let holder = Arc::new(
16867 AgentBuilder::new()
16868 .system_prompt("Hold the global path lock.")
16869 .llm(Arc::new(mock_with_response("done")))
16870 .tool(Arc::new(BlockingPathMutationTool {
16871 id: "holder_write",
16872 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16873 gate: holder_gate.clone(),
16874 }))
16875 .build()
16876 .unwrap()
16877 .with_shared_resource_locks(Arc::clone(&locks)),
16878 );
16879 let waiter = Arc::new(
16880 AgentBuilder::new()
16881 .system_prompt("Wait for the global path lock.")
16882 .llm(Arc::new(mock_with_response("done")))
16883 .tool(Arc::new(BlockingPathMutationTool {
16884 id: "waiter_write",
16885 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16886 gate: waiter_gate.clone(),
16887 }))
16888 .build()
16889 .unwrap()
16890 .with_shared_resource_locks(Arc::clone(&locks)),
16891 );
16892
16893 let holder_call = {
16894 let holder = Arc::clone(&holder);
16895 tokio::spawn(async move {
16896 holder
16897 .invoke_tool(ToolExecutionRequest::new(
16898 "holder-call",
16899 "holder_write",
16900 serde_json::json!({"path": "./shared.txt"}),
16901 ToolCallSource::Manual,
16902 ))
16903 .await
16904 .unwrap()
16905 })
16906 };
16907 holder_gate.wait_until_entered().await;
16908
16909 let waiter_call = {
16910 let waiter = Arc::clone(&waiter);
16911 tokio::spawn(async move {
16912 waiter
16913 .invoke_tool(ToolExecutionRequest::new(
16914 "waiter-call",
16915 "waiter_write",
16916 serde_json::json!({"path": "./shared.txt"}),
16917 ToolCallSource::Manual,
16918 ))
16919 .await
16920 .unwrap()
16921 })
16922 };
16923 wait_for_resource_lock_strong_count(&locks, 2).await;
16924 waiter.runtime_control().cancel_all();
16925
16926 let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
16927 .await
16928 .expect("cancelled lock waiter did not finish")
16929 .unwrap();
16930 assert!(!waiter_record.executed);
16931 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
16932 assert_eq!(
16933 locks
16934 .read()
16935 .get("path-mutation:global")
16936 .map_or(0, |lock| lock.strong_count()),
16937 1
16938 );
16939
16940 holder_gate.release();
16941 let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
16942 .await
16943 .expect("lock holder did not finish")
16944 .unwrap();
16945 assert!(holder_record.success);
16946 assert!(locks.read().is_empty());
16947 }
16948
16949 #[tokio::test]
16950 async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
16951 for denial in [MutationDenial::Policy, MutationDenial::Approval] {
16952 let tools: [Arc<dyn Tool>; 3] = [
16953 Arc::new(CopyPathTool::new()),
16954 Arc::new(MovePathTool::new()),
16955 Arc::new(DeletePathTool::new()),
16956 ];
16957 for tool in tools {
16958 assert_path_mutation_denied(tool, denial).await;
16959 }
16960 }
16961 }
16962
16963 #[tokio::test]
16964 async fn policy_denial_keeps_executor_hook_lifecycle_and_record_authority() {
16965 let workspace = MutationTestWorkspace::new();
16966 let target = workspace.root.join("denied.txt");
16967 let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
16968 let agent = AgentBuilder::new()
16969 .system_prompt("Test denied tool hooks.")
16970 .llm(Arc::new(mock_with_response("done")))
16971 .tool(Arc::new(FileWriteTool::new()))
16972 .tool_security(ToolSecurityEngine::new(mutation_denial_security_config(
16973 "file_write",
16974 &workspace.root,
16975 MutationDenial::Policy,
16976 )))
16977 .hooks(hooks.clone())
16978 .build()
16979 .unwrap();
16980
16981 let record = agent
16982 .invoke_tool(ToolExecutionRequest::new(
16983 "denied-hook-call",
16984 "file_write",
16985 serde_json::json!({
16986 "path": target.to_string_lossy(),
16987 "content": "blocked"
16988 }),
16989 ToolCallSource::Manual,
16990 ))
16991 .await
16992 .unwrap();
16993
16994 assert!(!record.executed);
16995 assert!(!record.success);
16996 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
16997 assert_eq!(
16998 hooks.events(),
16999 vec![
17000 "start:file_write",
17001 "complete:file_write:false",
17002 "record:file_write:false",
17003 "error"
17004 ]
17005 );
17006 assert!(!target.exists());
17007 }
17008
17009 #[tokio::test]
17010 async fn approval_argument_changes_are_rechecked_against_final_scope() {
17011 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17012 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17013 let entered = Arc::new(tokio::sync::Barrier::new(2));
17014 let release = Arc::new(tokio::sync::Notify::new());
17015 let handler = Arc::new(BlockingApprovalHandler {
17016 entered: Arc::clone(&entered),
17017 release: Arc::clone(&release),
17018 result: ApprovalResult::Modified {
17019 changes: HashMap::from([(
17020 "path".to_string(),
17021 Value::String("./after-approval.txt".to_string()),
17022 )]),
17023 },
17024 });
17025 let agent = Arc::new(
17026 AgentBuilder::new()
17027 .system_prompt("Test final scope validation.")
17028 .llm(Arc::new(mock_with_response("done")))
17029 .tool(Arc::new(LockedWriteTool {
17030 active: Arc::clone(&active),
17031 max_active: Arc::clone(&max_active),
17032 }))
17033 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17034 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17035 .approval_handler(handler)
17036 .build()
17037 .unwrap(),
17038 );
17039 let control = agent.runtime_control();
17040 let running = Arc::clone(&agent);
17041 let call = tokio::spawn(async move {
17042 running
17043 .invoke_tool(ToolExecutionRequest::new(
17044 "approval-scope",
17045 "locked_write",
17046 serde_json::json!({"path": "./before-approval.txt"}),
17047 ToolCallSource::Manual,
17048 ))
17049 .await
17050 .unwrap()
17051 });
17052 entered.wait().await;
17053 let expected_version = control.set_tool_scope(Vec::new());
17054 release.notify_one();
17055 let record = call.await.unwrap();
17056
17057 assert!(!record.executed);
17058 assert!(!record.success);
17059 assert_eq!(record.runtime_config_version, expected_version);
17060 assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
17061 assert_eq!(max_active.load(Ordering::SeqCst), 0);
17062 assert_eq!(
17063 record.metadata["runtime_scope_snapshot"],
17064 serde_json::json!([])
17065 );
17066 }
17067
17068 #[tokio::test]
17069 async fn approval_is_rechecked_against_final_policy_snapshot() {
17070 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17071 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17072 let entered = Arc::new(tokio::sync::Barrier::new(2));
17073 let release = Arc::new(tokio::sync::Notify::new());
17074 let handler = Arc::new(BlockingApprovalHandler {
17075 entered: Arc::clone(&entered),
17076 release: Arc::clone(&release),
17077 result: ApprovalResult::Approved,
17078 });
17079 let agent = Arc::new(
17080 AgentBuilder::new()
17081 .system_prompt("Test final policy validation.")
17082 .llm(Arc::new(mock_with_response("done")))
17083 .tool(Arc::new(LockedWriteTool {
17084 active: Arc::clone(&active),
17085 max_active: Arc::clone(&max_active),
17086 }))
17087 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17088 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17089 .approval_handler(handler)
17090 .build()
17091 .unwrap(),
17092 );
17093 let control = agent.runtime_control();
17094 let running = Arc::clone(&agent);
17095 let call = tokio::spawn(async move {
17096 running
17097 .invoke_tool(ToolExecutionRequest::new(
17098 "approval-policy",
17099 "locked_write",
17100 serde_json::json!({"path": "./policy.txt"}),
17101 ToolCallSource::Manual,
17102 ))
17103 .await
17104 .unwrap()
17105 });
17106 entered.wait().await;
17107 let expected_version = control.set_tool_security(approval_security_config(false));
17108 release.notify_one();
17109 let record = call.await.unwrap();
17110
17111 assert!(!record.executed);
17112 assert!(!record.success);
17113 assert_eq!(record.runtime_config_version, expected_version);
17114 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17115 assert_eq!(max_active.load(Ordering::SeqCst), 0);
17116 assert!(record.metadata.contains_key("policy_snapshot"));
17117 }
17118
17119 #[test]
17120 fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
17121 let agent = AgentBuilder::new()
17122 .system_prompt("Test runtime policy validation.")
17123 .llm(Arc::new(mock_with_response("done")))
17124 .build()
17125 .unwrap();
17126 let control = agent.runtime_control();
17127 let mut valid = ToolSecurityConfig::default();
17128 valid.tools.insert(
17129 "web_search".to_string(),
17130 ai_agents_tools::ToolPolicyConfig {
17131 max_results: Some(5),
17132 ..Default::default()
17133 },
17134 );
17135 let generation = control.try_set_tool_security(valid).unwrap();
17136
17137 let mut invalid = ToolSecurityConfig::default();
17138 invalid.tools.insert(
17139 "web_search".to_string(),
17140 ai_agents_tools::ToolPolicyConfig {
17141 max_results: Some(0),
17142 ..Default::default()
17143 },
17144 );
17145 let error = control.try_set_tool_security(invalid).unwrap_err();
17146
17147 assert!(
17148 error
17149 .to_string()
17150 .contains("max_results must be greater than 0")
17151 );
17152 assert_eq!(control.version(), generation);
17153 assert_eq!(
17154 control
17155 .state
17156 .tool_security_override
17157 .read()
17158 .as_ref()
17159 .unwrap()
17160 .config()
17161 .tools["web_search"]
17162 .max_results,
17163 Some(5)
17164 );
17165 }
17166
17167 #[tokio::test]
17168 async fn persistent_override_preserves_rate_history_within_generation() {
17169 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17170 let agent = AgentBuilder::new()
17171 .system_prompt("Test persistent policy overrides.")
17172 .llm(Arc::new(mock_with_response("done")))
17173 .tool(Arc::new(RecoveryTestTool {
17174 id: "limited_override".to_string(),
17175 succeeds: true,
17176 calls: Arc::clone(&calls),
17177 max_output_chars: None,
17178 }))
17179 .build()
17180 .unwrap();
17181 let mut security = ToolSecurityConfig {
17182 enabled: true,
17183 fail_closed: true,
17184 ..Default::default()
17185 };
17186 let policy = ai_agents_tools::ToolPolicyConfig {
17187 write_paths: vec![".".to_string()],
17188 rate_limit: Some(1),
17189 ..Default::default()
17190 };
17191 security
17192 .tools
17193 .insert("limited_override".to_string(), policy);
17194 let generation = agent.runtime_control().set_tool_security(security);
17195
17196 let first = agent
17197 .invoke_tool(ToolExecutionRequest::new(
17198 "limited-first",
17199 "limited_override",
17200 serde_json::json!({"path": "./limited.txt"}),
17201 ToolCallSource::Manual,
17202 ))
17203 .await
17204 .unwrap();
17205 let second = agent
17206 .invoke_tool(ToolExecutionRequest::new(
17207 "limited-second",
17208 "limited_override",
17209 serde_json::json!({"path": "./limited.txt"}),
17210 ToolCallSource::Manual,
17211 ))
17212 .await
17213 .unwrap();
17214
17215 assert!(first.success);
17216 assert_eq!(first.policy_version, generation);
17217 assert!(!second.executed);
17218 assert!(second.output.contains("Rate limit exceeded"));
17219 assert_eq!(second.policy_version, generation);
17220 assert_eq!(calls.load(Ordering::SeqCst), 1);
17221 }
17222
17223 #[tokio::test]
17224 async fn concurrent_rate_admission_consumes_capacity_atomically() {
17225 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17226 let tool = Arc::new(RecoveryTestTool {
17227 id: "atomic_rate".to_string(),
17228 succeeds: true,
17229 calls: Arc::clone(&calls),
17230 max_output_chars: None,
17231 });
17232 let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
17233 let bindings = tool.policy_bindings();
17234 let classification = tool.classify_call(&arguments);
17235 let resource_keys =
17236 tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
17237 let mut security = ToolSecurityConfig {
17238 enabled: true,
17239 fail_closed: true,
17240 ..Default::default()
17241 };
17242 let policy = ai_agents_tools::ToolPolicyConfig {
17243 write_paths: vec![".".to_string()],
17244 rate_limit: Some(1),
17245 ..Default::default()
17246 };
17247 security.tools.insert(tool.id().to_string(), policy);
17248 let agent = Arc::new(
17249 AgentBuilder::new()
17250 .system_prompt("Test atomic rate admission.")
17251 .llm(Arc::new(mock_with_response("done")))
17252 .tool(tool)
17253 .tool_security(ToolSecurityEngine::new(security))
17254 .build()
17255 .unwrap(),
17256 );
17257 let held = agent
17258 .acquire_tool_resource_locks(&resource_keys)
17259 .await
17260 .unwrap();
17261 let left = {
17262 let agent = Arc::clone(&agent);
17263 let arguments = arguments.clone();
17264 tokio::spawn(async move {
17265 agent
17266 .invoke_tool(ToolExecutionRequest::new(
17267 "atomic-rate-left",
17268 "atomic_rate",
17269 arguments,
17270 ToolCallSource::Manual,
17271 ))
17272 .await
17273 .unwrap()
17274 })
17275 };
17276 let right = {
17277 let agent = Arc::clone(&agent);
17278 tokio::spawn(async move {
17279 agent
17280 .invoke_tool(ToolExecutionRequest::new(
17281 "atomic-rate-right",
17282 "atomic_rate",
17283 arguments,
17284 ToolCallSource::Manual,
17285 ))
17286 .await
17287 .unwrap()
17288 })
17289 };
17290 tokio::time::sleep(std::time::Duration::from_millis(25)).await;
17291 drop(held);
17292 let (left, right) = tokio::join!(left, right);
17293 let records = [left.unwrap(), right.unwrap()];
17294
17295 assert_eq!(records.iter().filter(|record| record.success).count(), 1);
17296 assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
17297 assert!(
17298 records.iter().any(|record| {
17299 !record.executed && record.output.contains("Rate limit exceeded")
17300 })
17301 );
17302 assert_eq!(calls.load(Ordering::SeqCst), 1);
17303 }
17304
17305 #[tokio::test]
17306 async fn changed_policy_generation_invalidates_pending_approval() {
17307 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17308 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17309 let entered = Arc::new(tokio::sync::Barrier::new(2));
17310 let release = Arc::new(tokio::sync::Notify::new());
17311 let handler = Arc::new(BlockingApprovalHandler {
17312 entered: Arc::clone(&entered),
17313 release: Arc::clone(&release),
17314 result: ApprovalResult::Approved,
17315 });
17316 let agent = Arc::new(
17317 AgentBuilder::new()
17318 .system_prompt("Test stale approval denial.")
17319 .llm(Arc::new(mock_with_response("done")))
17320 .tool(Arc::new(LockedWriteTool {
17321 active: Arc::clone(&active),
17322 max_active: Arc::clone(&max_active),
17323 }))
17324 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17325 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17326 .approval_handler(handler)
17327 .build()
17328 .unwrap(),
17329 );
17330 let running = Arc::clone(&agent);
17331 let call = tokio::spawn(async move {
17332 running
17333 .invoke_tool(ToolExecutionRequest::new(
17334 "stale-approval",
17335 "locked_write",
17336 serde_json::json!({"path": "./stale.txt"}),
17337 ToolCallSource::Manual,
17338 ))
17339 .await
17340 .unwrap()
17341 });
17342 entered.wait().await;
17343 let generation = agent
17344 .runtime_control()
17345 .set_tool_security(approval_security_config(true));
17346 release.notify_one();
17347 let record = call.await.unwrap();
17348
17349 assert!(!record.executed);
17350 assert!(record.output.contains("Approval became stale"));
17351 assert_eq!(record.policy_version, generation);
17352 assert_eq!(max_active.load(Ordering::SeqCst), 0);
17353 }
17354
17355 #[tokio::test]
17356 async fn final_policy_reapplies_argument_caps_after_approval_changes() {
17357 use ai_agents_hitl::CallbackHandler;
17358
17359 let mut security = ToolSecurityConfig {
17360 enabled: true,
17361 fail_closed: true,
17362 ..Default::default()
17363 };
17364 let policy = ai_agents_tools::ToolPolicyConfig {
17365 read_paths: vec![".".to_string()],
17366 max_results: Some(5),
17367 require_confirmation: true,
17368 ..Default::default()
17369 };
17370 security.tools.insert("context_echo".to_string(), policy);
17371 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
17372 changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
17373 });
17374 let agent = AgentBuilder::new()
17375 .system_prompt("Test final argument caps.")
17376 .llm(Arc::new(mock_with_response("done")))
17377 .tool(Arc::new(ContextEchoTool))
17378 .tool_security(ToolSecurityEngine::new(security))
17379 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17380 .approval_handler(Arc::new(handler))
17381 .build()
17382 .unwrap();
17383
17384 let record = agent
17385 .invoke_tool(ToolExecutionRequest::new(
17386 "final-cap",
17387 "context_echo",
17388 serde_json::json!({"path": ".", "max_results": 1}),
17389 ToolCallSource::Manual,
17390 ))
17391 .await
17392 .unwrap();
17393
17394 assert!(record.success);
17395 assert_eq!(record.executed_arguments["max_results"], 5);
17396 assert_eq!(
17397 record.approval.unwrap().modified_arguments.unwrap()["max_results"],
17398 5
17399 );
17400 }
17401
17402 #[tokio::test]
17403 async fn no_binding_writes_use_canonical_fallback_lock() {
17404 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17405 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17406 let agent = Arc::new(
17407 AgentBuilder::new()
17408 .system_prompt("Test fallback resource locks.")
17409 .llm(Arc::new(mock_with_response("done")))
17410 .tool(Arc::new(NoBindingWriteTool {
17411 active: Arc::clone(&active),
17412 max_active: Arc::clone(&max_active),
17413 }))
17414 .build()
17415 .unwrap(),
17416 );
17417 let left = {
17418 let agent = Arc::clone(&agent);
17419 tokio::spawn(async move {
17420 agent
17421 .invoke_tool(ToolExecutionRequest::new(
17422 "no-binding-left",
17423 "no_binding_write",
17424 serde_json::json!({}),
17425 ToolCallSource::Manual,
17426 ))
17427 .await
17428 .unwrap()
17429 })
17430 };
17431 let right = {
17432 let agent = Arc::clone(&agent);
17433 tokio::spawn(async move {
17434 agent
17435 .invoke_tool(ToolExecutionRequest::new(
17436 "no-binding-right",
17437 "no_binding_write",
17438 serde_json::json!({}),
17439 ToolCallSource::Manual,
17440 ))
17441 .await
17442 .unwrap()
17443 })
17444 };
17445 let (left, right) = tokio::join!(left, right);
17446
17447 assert!(left.unwrap().success);
17448 assert!(right.unwrap().success);
17449 assert_eq!(max_active.load(Ordering::SeqCst), 1);
17450 }
17451
17452 #[tokio::test]
17453 async fn parent_and_child_paths_share_a_resource_lock() {
17454 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17455 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17456 let agent = Arc::new(
17457 AgentBuilder::new()
17458 .system_prompt("Test parent child resource locks.")
17459 .llm(Arc::new(mock_with_response("done")))
17460 .tool(Arc::new(LockedWriteTool {
17461 active: Arc::clone(&active),
17462 max_active: Arc::clone(&max_active),
17463 }))
17464 .build()
17465 .unwrap(),
17466 );
17467 let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
17468 let child = format!("{}/child.txt", parent);
17469 let left = {
17470 let agent = Arc::clone(&agent);
17471 tokio::spawn(async move {
17472 agent
17473 .invoke_tool(ToolExecutionRequest::new(
17474 "parent-lock",
17475 "locked_write",
17476 serde_json::json!({"path": parent}),
17477 ToolCallSource::Manual,
17478 ))
17479 .await
17480 .unwrap()
17481 })
17482 };
17483 let right = {
17484 let agent = Arc::clone(&agent);
17485 tokio::spawn(async move {
17486 agent
17487 .invoke_tool(ToolExecutionRequest::new(
17488 "child-lock",
17489 "locked_write",
17490 serde_json::json!({"path": child}),
17491 ToolCallSource::Manual,
17492 ))
17493 .await
17494 .unwrap()
17495 })
17496 };
17497 let (left, right) = tokio::join!(left, right);
17498
17499 assert!(left.unwrap().success);
17500 assert!(right.unwrap().success);
17501 assert_eq!(max_active.load(Ordering::SeqCst), 1);
17502 }
17503
17504 #[tokio::test]
17505 async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
17506 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17507 let hooks = Arc::new(ReentrantToolHooks {
17508 agent: parking_lot::Mutex::new(None),
17509 invoked: AtomicBool::new(false),
17510 nested_success: AtomicBool::new(false),
17511 });
17512 let agent = Arc::new(
17513 AgentBuilder::new()
17514 .system_prompt("Test hook reentrancy.")
17515 .llm(Arc::new(mock_with_response("done")))
17516 .tool(Arc::new(RecoveryTestTool {
17517 id: "reentrant_write".to_string(),
17518 succeeds: true,
17519 calls: Arc::clone(&calls),
17520 max_output_chars: None,
17521 }))
17522 .hooks(hooks.clone())
17523 .build()
17524 .unwrap(),
17525 );
17526 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
17527 let record = tokio::time::timeout(
17528 std::time::Duration::from_secs(2),
17529 agent.invoke_tool(ToolExecutionRequest::new(
17530 "outer-hook-call",
17531 "reentrant_write",
17532 serde_json::json!({"path": "./hook.txt"}),
17533 ToolCallSource::Manual,
17534 )),
17535 )
17536 .await
17537 .expect("tool completion hook must not retain resource guards")
17538 .unwrap();
17539
17540 assert!(record.success);
17541 assert!(hooks.nested_success.load(Ordering::SeqCst));
17542 assert_eq!(calls.load(Ordering::SeqCst), 2);
17543 }
17544
17545 #[tokio::test]
17547 async fn fallback_finalizes_original_record_before_shared_execution() {
17548 let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17549 let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17550 let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
17551 let agent = AgentBuilder::new()
17552 .system_prompt("Test fallback execution.")
17553 .llm(Arc::new(mock_with_response("done")))
17554 .tool(Arc::new(RecoveryTestTool {
17555 id: "primary".to_string(),
17556 succeeds: false,
17557 calls: Arc::clone(&primary_calls),
17558 max_output_chars: None,
17559 }))
17560 .tool(Arc::new(RecoveryTestTool {
17561 id: "fallback".to_string(),
17562 succeeds: true,
17563 calls: Arc::clone(&fallback_calls),
17564 max_output_chars: None,
17565 }))
17566 .recovery_manager(recovery_manager_with_fallbacks([(
17567 "primary".to_string(),
17568 "fallback".to_string(),
17569 )]))
17570 .hooks(hooks.clone())
17571 .build()
17572 .unwrap();
17573 let record = tokio::time::timeout(
17574 std::time::Duration::from_secs(2),
17575 agent.invoke_tool(ToolExecutionRequest::new(
17576 "fallback-call",
17577 "primary",
17578 serde_json::json!({"path": "./shared.txt"}),
17579 ToolCallSource::Manual,
17580 )),
17581 )
17582 .await
17583 .expect("fallback must not retain the primary resource guard")
17584 .unwrap();
17585
17586 assert_eq!(
17587 hooks.events(),
17588 vec![
17589 "start:primary",
17590 "complete:primary:false",
17591 "record:primary:true",
17592 "error",
17593 "start:fallback",
17594 "complete:fallback:true",
17595 "record:fallback:true",
17596 ]
17597 );
17598 let records = hooks.records();
17599 assert_eq!(records.len(), 2);
17600 let original = &records[0];
17601 assert_eq!(original.canonical_id, "primary");
17602 assert!(matches!(original.source, ToolCallSource::Manual));
17603 assert!(original.executed);
17604 assert!(!original.success);
17605
17606 let fallback = &records[1];
17607 assert_eq!(fallback.canonical_id, "fallback");
17608 assert_eq!(fallback.call_id, "fallback-call");
17609 assert!(matches!(
17610 &fallback.source,
17611 ToolCallSource::Fallback { original_tool } if original_tool == "primary"
17612 ));
17613 assert!(fallback.executed);
17614 assert!(fallback.success);
17615 assert_eq!(record.canonical_id, fallback.canonical_id);
17616 assert_eq!(record.output, fallback.output);
17617
17618 let history = agent.tool_call_history();
17619 assert_eq!(
17620 history
17621 .iter()
17622 .map(|entry| entry.tool_id.as_str())
17623 .collect::<Vec<_>>(),
17624 vec!["primary", "fallback"]
17625 );
17626 assert_eq!(history[0].result.get("success"), Some(&Value::Bool(false)));
17627 assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
17628 assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
17629 }
17630
17631 #[tokio::test]
17632 async fn diagnostics_without_provider_records_unavailable_without_execution() {
17633 let mock = mock_with_response("hello");
17634 let yaml = r#"
17635name: DiagnosticsNoProviderAgent
17636system_prompt: "Review diagnostics."
17637tools: [diagnostics]
17638"#;
17639 let agent = AgentBuilder::from_yaml(yaml)
17640 .unwrap()
17641 .llm(Arc::new(mock))
17642 .auto_configure_features()
17643 .unwrap()
17644 .build()
17645 .unwrap();
17646
17647 let record = agent
17648 .invoke_tool(ToolExecutionRequest::new(
17649 "diagnostics-call",
17650 "diagnostics",
17651 serde_json::json!({}),
17652 ToolCallSource::Manual,
17653 ))
17654 .await
17655 .unwrap();
17656
17657 assert!(!record.executed);
17658 assert!(!record.success);
17659 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17660 }
17661
17662 #[tokio::test]
17663 async fn web_search_without_provider_records_unavailable_without_execution() {
17664 let mock = mock_with_response("hello");
17665 let yaml = r#"
17666name: WebSearchNoProviderAgent
17667system_prompt: "You search the web."
17668tools: [web_search]
17669"#;
17670 let agent = AgentBuilder::from_yaml(yaml)
17671 .unwrap()
17672 .llm(Arc::new(mock))
17673 .auto_configure_features()
17674 .unwrap()
17675 .build()
17676 .unwrap();
17677
17678 let record = agent
17679 .invoke_tool(ToolExecutionRequest::new(
17680 "web-search-call",
17681 "web_search",
17682 serde_json::json!({"query": "rust async"}),
17683 ToolCallSource::Manual,
17684 ))
17685 .await
17686 .unwrap();
17687
17688 assert!(!record.executed);
17689 assert!(!record.success);
17690 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17691 }
17692
17693 #[tokio::test]
17694 async fn unavailable_host_tool_does_not_request_approval() {
17695 let approvals = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17696 let handler = Arc::new(CountingApprovalHandler {
17697 calls: Arc::clone(&approvals),
17698 });
17699 let mut security = ToolSecurityConfig {
17700 enabled: true,
17701 fail_closed: true,
17702 ..Default::default()
17703 };
17704 security.tools.insert(
17705 "web_search".to_string(),
17706 ai_agents_tools::ToolPolicyConfig {
17707 enabled: true,
17708 require_confirmation: true,
17709 ..Default::default()
17710 },
17711 );
17712 let yaml = r#"
17713name: UnavailableApprovalAgent
17714system_prompt: "Search only with approval."
17715tools: [web_search]
17716"#;
17717 let agent = AgentBuilder::from_yaml(yaml)
17718 .unwrap()
17719 .llm(Arc::new(mock_with_response("done")))
17720 .auto_configure_features()
17721 .unwrap()
17722 .tool_security(ToolSecurityEngine::new(security))
17723 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17724 .approval_handler(handler)
17725 .build()
17726 .unwrap();
17727
17728 let record = agent
17729 .invoke_tool(ToolExecutionRequest::new(
17730 "unavailable-before-approval",
17731 "web_search",
17732 serde_json::json!({"query": "rust async"}),
17733 ToolCallSource::Manual,
17734 ))
17735 .await
17736 .unwrap();
17737
17738 assert_eq!(approvals.load(Ordering::SeqCst), 0);
17739 assert!(!record.executed);
17740 assert!(!record.success);
17741 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17742 assert!(
17743 record
17744 .approval
17745 .as_ref()
17746 .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Unavailable))
17747 );
17748 }
17749
17750 #[tokio::test]
17751 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
17752 let mock = mock_with_response("hello");
17753 let yaml = r#"
17754name: SpawnerNoGrantAgent
17755system_prompt: "You manage agents."
17756spawner:
17757 max_agents: 2
17758"#;
17759 let agent = AgentBuilder::from_yaml(yaml)
17760 .unwrap()
17761 .llm(Arc::new(mock))
17762 .auto_configure_features()
17763 .unwrap()
17764 .auto_configure_spawner()
17765 .await
17766 .unwrap()
17767 .build()
17768 .unwrap();
17769
17770 let available = agent.get_available_tool_ids().await.unwrap();
17771 assert!(available.is_empty());
17772 }
17773
17774 #[tokio::test]
17775 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
17776 let mock = mock_with_response("hello");
17777 let yaml = r#"
17778name: EmptySpawnerNoGrantAgent
17779system_prompt: "You manage agents."
17780tools: []
17781spawner:
17782 max_agents: 2
17783"#;
17784 let agent = AgentBuilder::from_yaml(yaml)
17785 .unwrap()
17786 .llm(Arc::new(mock))
17787 .auto_configure_features()
17788 .unwrap()
17789 .auto_configure_spawner()
17790 .await
17791 .unwrap()
17792 .build()
17793 .unwrap();
17794
17795 let available = agent.get_available_tool_ids().await.unwrap();
17796 assert!(available.is_empty());
17797 }
17798
17799 #[tokio::test]
17800 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
17801 let mock = mock_with_response("hello");
17802 let yaml = r#"
17803name: ManagementGrantAgent
17804system_prompt: "You manage agents."
17805tools: []
17806spawner:
17807 management_tools: true
17808"#;
17809 let agent = AgentBuilder::from_yaml(yaml)
17810 .unwrap()
17811 .llm(Arc::new(mock))
17812 .auto_configure_features()
17813 .unwrap()
17814 .auto_configure_spawner()
17815 .await
17816 .unwrap()
17817 .build()
17818 .unwrap();
17819
17820 let available = agent.get_available_tool_ids().await.unwrap();
17821 assert_eq!(available.len(), 4);
17822 assert!(available.contains(&"spawn_agent".to_string()));
17823 assert!(available.contains(&"send_agent_message".to_string()));
17824 assert!(available.contains(&"list_agents".to_string()));
17825 assert!(available.contains(&"remove_agent".to_string()));
17826 }
17827
17828 #[tokio::test]
17829 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
17830 let mock = mock_with_response("hello");
17831 let yaml = r#"
17832name: ManagementOmittedToolsGrantAgent
17833system_prompt: "You manage agents."
17834spawner:
17835 management_tools: true
17836"#;
17837 let agent = AgentBuilder::from_yaml(yaml)
17838 .unwrap()
17839 .llm(Arc::new(mock))
17840 .auto_configure_features()
17841 .unwrap()
17842 .auto_configure_spawner()
17843 .await
17844 .unwrap()
17845 .build()
17846 .unwrap();
17847
17848 let available = agent.get_available_tool_ids().await.unwrap();
17849 assert_eq!(available.len(), 4);
17850 assert!(available.contains(&"spawn_agent".to_string()));
17851 assert!(available.contains(&"send_agent_message".to_string()));
17852 assert!(available.contains(&"list_agents".to_string()));
17853 assert!(available.contains(&"remove_agent".to_string()));
17854 }
17855
17856 #[tokio::test]
17857 async fn test_management_tools_selected_grants_only_selected_tools() {
17858 let mock = mock_with_response("hello");
17859 let yaml = r#"
17860name: ManagementSelectedGrantAgent
17861system_prompt: "You manage agents."
17862tools: []
17863spawner:
17864 management_tools:
17865 - spawn_agent
17866 - send_agent_message
17867 - list_agents
17868"#;
17869 let agent = AgentBuilder::from_yaml(yaml)
17870 .unwrap()
17871 .llm(Arc::new(mock))
17872 .auto_configure_features()
17873 .unwrap()
17874 .auto_configure_spawner()
17875 .await
17876 .unwrap()
17877 .build()
17878 .unwrap();
17879
17880 let available = agent.get_available_tool_ids().await.unwrap();
17881 assert_eq!(available.len(), 3);
17882 assert!(available.contains(&"spawn_agent".to_string()));
17883 assert!(available.contains(&"send_agent_message".to_string()));
17884 assert!(available.contains(&"list_agents".to_string()));
17885 assert!(!available.contains(&"remove_agent".to_string()));
17886 }
17887
17888 #[tokio::test]
17889 async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
17890 let mock = mock_with_response("hello");
17891 let yaml = r#"
17892name: OrchestrationGrantAgent
17893system_prompt: "You coordinate agents."
17894llms:
17895 default:
17896 provider: openai
17897 model: gpt-4
17898 router:
17899 provider: openai
17900 model: gpt-4
17901llm:
17902 default: default
17903 router: router
17904tools: []
17905spawner:
17906 orchestration_tools: true
17907"#;
17908 let agent = AgentBuilder::from_yaml(yaml)
17909 .unwrap()
17910 .llm(Arc::new(mock))
17911 .auto_configure_features()
17912 .unwrap()
17913 .auto_configure_spawner()
17914 .await
17915 .unwrap()
17916 .build()
17917 .unwrap();
17918
17919 let available = agent.get_available_tool_ids().await.unwrap();
17920 assert_eq!(available.len(), 5);
17921 assert!(available.contains(&"route_to_agent".to_string()));
17922 assert!(available.contains(&"pipeline_process".to_string()));
17923 assert!(available.contains(&"concurrent_ask".to_string()));
17924 assert!(available.contains(&"group_discussion".to_string()));
17925 assert!(available.contains(&"handoff_conversation".to_string()));
17926 }
17927
17928 #[tokio::test]
17929 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
17930 let mock = mock_with_response("hello");
17931 let yaml = r#"
17932name: PersonaGrantAgent
17933system_prompt: "You can evolve persona."
17934llm:
17935 provider: openai
17936 model: gpt-4
17937tools: []
17938persona:
17939 identity:
17940 name: "Guide"
17941 role: "Helper"
17942 evolution:
17943 enabled: true
17944 allow_llm_evolve: true
17945 mutable_fields:
17946 - traits.personality
17947"#;
17948 let agent = AgentBuilder::from_yaml(yaml)
17949 .unwrap()
17950 .llm(Arc::new(mock))
17951 .build()
17952 .unwrap();
17953
17954 let available = agent.get_available_tool_ids().await.unwrap();
17955 assert_eq!(available, vec!["persona_evolve".to_string()]);
17956 }
17957
17958 #[tokio::test]
17959 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
17960 let mock = mock_with_response("hello");
17961 let yaml = r#"
17962name: PersonaOmittedToolsGrantAgent
17963system_prompt: "You can evolve persona."
17964llm:
17965 provider: openai
17966 model: gpt-4
17967persona:
17968 identity:
17969 name: "Guide"
17970 role: "Helper"
17971 evolution:
17972 enabled: true
17973 allow_llm_evolve: true
17974 mutable_fields:
17975 - traits.personality
17976"#;
17977 let agent = AgentBuilder::from_yaml(yaml)
17978 .unwrap()
17979 .llm(Arc::new(mock))
17980 .build()
17981 .unwrap();
17982
17983 let available = agent.get_available_tool_ids().await.unwrap();
17984 assert_eq!(available, vec!["persona_evolve".to_string()]);
17985 }
17986
17987 #[tokio::test]
17988 async fn test_omitted_yaml_tools_exposes_no_tools() {
17989 let mock = mock_with_response("hello");
17990 let yaml = r#"
17991name: NoToolsAgent
17992system_prompt: "You are helpful."
17993"#;
17994 let agent = AgentBuilder::from_yaml(yaml)
17995 .unwrap()
17996 .llm(Arc::new(mock))
17997 .auto_configure_features()
17998 .unwrap()
17999 .build()
18000 .unwrap();
18001
18002 let available = agent.get_available_tool_ids().await.unwrap();
18003 assert!(available.is_empty());
18004 }
18005
18006 #[tokio::test]
18007 async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
18008 for tools in ["", "tools: []"] {
18009 let yaml = format!(
18010 r#"
18011name: RuntimeScopeNoGrantAgent
18012system_prompt: "No ordinary tools are granted."
18013{tools}
18014"#
18015 );
18016 let agent = AgentBuilder::from_yaml(&yaml)
18017 .unwrap()
18018 .llm(Arc::new(mock_with_response("done")))
18019 .auto_configure_features()
18020 .unwrap()
18021 .build()
18022 .unwrap();
18023
18024 agent
18025 .runtime_control()
18026 .set_tool_scope(vec!["calculator".to_string()]);
18027
18028 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
18029 }
18030 }
18031
18032 #[tokio::test]
18033 async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
18034 let yaml = r#"
18035name: RuntimeScopeWideningAgent
18036system_prompt: "Runtime scope cannot add authority."
18037tools: [calculator]
18038"#;
18039 let agent = AgentBuilder::from_yaml(yaml)
18040 .unwrap()
18041 .llm(Arc::new(mock_with_response("done")))
18042 .auto_configure_features()
18043 .unwrap()
18044 .build()
18045 .unwrap();
18046
18047 agent
18048 .runtime_control()
18049 .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
18050
18051 assert_eq!(
18052 agent.get_available_tool_ids().await.unwrap(),
18053 vec!["calculator".to_string()]
18054 );
18055 }
18056
18057 #[tokio::test]
18058 async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
18059 let yaml = r#"
18060name: RuntimeScopeIntersectionAgent
18061system_prompt: "Use only declared tools."
18062tools: [calculator, datetime]
18063"#;
18064 let agent = AgentBuilder::from_yaml(yaml)
18065 .unwrap()
18066 .llm(Arc::new(mock_with_response("done")))
18067 .auto_configure_features()
18068 .unwrap()
18069 .build()
18070 .unwrap();
18071 let mut aliases = ai_agents_tools::ToolAliases::default();
18072 aliases
18073 .names
18074 .insert("en".to_string(), "calculate_alias".to_string());
18075 agent.tools.set_tool_aliases("calculator", aliases);
18076 let control = agent.runtime_control();
18077
18078 control.set_tool_scope(vec![
18079 "datetime".to_string(),
18080 "calculate_alias".to_string(),
18081 "calculator".to_string(),
18082 "unknown".to_string(),
18083 "datetime".to_string(),
18084 ]);
18085 assert_eq!(
18086 agent.get_available_tool_ids().await.unwrap(),
18087 vec!["calculator".to_string(), "datetime".to_string()]
18088 );
18089
18090 control.set_tool_scope(vec!["datetime".to_string()]);
18091 assert_eq!(
18092 agent.get_available_tool_ids().await.unwrap(),
18093 vec!["datetime".to_string()]
18094 );
18095
18096 control.clear_tool_scope_override();
18097 assert_eq!(
18098 agent.get_available_tool_ids().await.unwrap(),
18099 vec!["calculator".to_string(), "datetime".to_string()]
18100 );
18101 }
18102
18103 #[tokio::test]
18104 async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
18105 let agent = AgentBuilder::new()
18106 .system_prompt("Use registered tools.")
18107 .llm(Arc::new(mock_with_response("done")))
18108 .tool(Arc::new(ContextEchoTool))
18109 .tool(Arc::new(SlowTool))
18110 .build()
18111 .unwrap();
18112
18113 agent.runtime_control().set_tool_scope(vec![
18114 "Context Echo".to_string(),
18115 "context_echo".to_string(),
18116 "unknown".to_string(),
18117 ]);
18118
18119 assert_eq!(
18120 agent.get_available_tool_ids().await.unwrap(),
18121 vec!["context_echo".to_string()]
18122 );
18123 }
18124
18125 #[tokio::test]
18126 async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
18127 let yaml = r#"
18128name: NestedStateScopeAgent
18129system_prompt: "Honor every state scope."
18130tools: [calculator, datetime, echo]
18131states:
18132 initial: root
18133 states:
18134 root:
18135 tools: [calculate_alias, datetime]
18136 initial: middle
18137 states:
18138 middle:
18139 initial: leaf
18140 states:
18141 leaf:
18142 tools: [datetime_alias, echo]
18143"#;
18144 let agent = AgentBuilder::from_yaml(yaml)
18145 .unwrap()
18146 .llm(Arc::new(mock_with_response("done")))
18147 .auto_configure_features()
18148 .unwrap()
18149 .build()
18150 .unwrap();
18151 let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
18152 calculator_aliases
18153 .names
18154 .insert("en".to_string(), "calculate_alias".to_string());
18155 agent
18156 .tools
18157 .set_tool_aliases("calculator", calculator_aliases);
18158 let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
18159 datetime_aliases
18160 .names
18161 .insert("en".to_string(), "datetime_alias".to_string());
18162 agent.tools.set_tool_aliases("datetime", datetime_aliases);
18163 agent.runtime_control().set_tool_scope(vec![
18164 "unknown".to_string(),
18165 "datetime_alias".to_string(),
18166 "calculate_alias".to_string(),
18167 "datetime".to_string(),
18168 ]);
18169
18170 assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
18171 assert_eq!(
18172 agent.get_available_tool_ids().await.unwrap(),
18173 vec!["datetime".to_string()]
18174 );
18175 }
18176
18177 #[tokio::test]
18178 async fn ancestor_empty_state_scope_denies_omitted_descendants() {
18179 let yaml = r#"
18180name: NestedEmptyStateScopeAgent
18181system_prompt: "An empty ancestor scope denies all tools."
18182tools: [calculator]
18183states:
18184 initial: root
18185 states:
18186 root:
18187 tools: []
18188 initial: middle
18189 states:
18190 middle:
18191 initial: leaf
18192 states:
18193 leaf: {}
18194"#;
18195 let agent = AgentBuilder::from_yaml(yaml)
18196 .unwrap()
18197 .llm(Arc::new(mock_with_response("done")))
18198 .auto_configure_features()
18199 .unwrap()
18200 .build()
18201 .unwrap();
18202
18203 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
18204 }
18205
18206 #[tokio::test]
18207 async fn state_change_during_approval_invalidates_the_reviewed_authority() {
18208 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18209 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18210 let entered = Arc::new(tokio::sync::Barrier::new(2));
18211 let release = Arc::new(tokio::sync::Notify::new());
18212 let handler = Arc::new(BlockingApprovalHandler {
18213 entered: Arc::clone(&entered),
18214 release: Arc::clone(&release),
18215 result: ApprovalResult::Approved,
18216 });
18217 let yaml = r#"
18218name: ApprovalStateGenerationAgent
18219system_prompt: "State authority may change during approval."
18220tools: [locked_write]
18221states:
18222 initial: first
18223 states:
18224 first:
18225 tools: [locked_write]
18226 second:
18227 tools: [locked_write]
18228"#;
18229 let agent = Arc::new(
18230 AgentBuilder::from_yaml(yaml)
18231 .unwrap()
18232 .llm(Arc::new(mock_with_response("done")))
18233 .tool(Arc::new(LockedWriteTool {
18234 active: Arc::clone(&active),
18235 max_active: Arc::clone(&max_active),
18236 }))
18237 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
18238 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18239 .approval_handler(handler)
18240 .build()
18241 .unwrap(),
18242 );
18243 let running = Arc::clone(&agent);
18244 let call = tokio::spawn(async move {
18245 running
18246 .invoke_tool(ToolExecutionRequest::new(
18247 "approval-state-generation",
18248 "locked_write",
18249 serde_json::json!({"path": "./state-generation.txt"}),
18250 ToolCallSource::Manual,
18251 ))
18252 .await
18253 .unwrap()
18254 });
18255
18256 entered.wait().await;
18257 agent.transition_to("second").await.unwrap();
18258 release.notify_one();
18259 let record = call.await.unwrap();
18260
18261 assert!(!record.executed);
18262 assert!(record.output.contains("Approval became stale"));
18263 assert_eq!(max_active.load(Ordering::SeqCst), 0);
18264 }
18265
18266 #[tokio::test]
18267 async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
18268 let holder_gate = PathMutationGate::new();
18269 let waiter_gate = PathMutationGate::new();
18270 let yaml = r#"
18271name: LockedStateGenerationAgent
18272system_prompt: "State authority must remain stable through admission."
18273tools: [state_lock_holder, state_lock_waiter]
18274states:
18275 initial: first
18276 states:
18277 first:
18278 tools: [state_lock_holder, state_lock_waiter]
18279 second:
18280 tools: [state_lock_holder, state_lock_waiter]
18281"#;
18282 let agent = Arc::new(
18283 AgentBuilder::from_yaml(yaml)
18284 .unwrap()
18285 .llm(Arc::new(mock_with_response("done")))
18286 .tool(Arc::new(BlockingPathMutationTool {
18287 id: "state_lock_holder",
18288 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
18289 gate: holder_gate.clone(),
18290 }))
18291 .tool(Arc::new(BlockingPathMutationTool {
18292 id: "state_lock_waiter",
18293 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
18294 gate: waiter_gate.clone(),
18295 }))
18296 .build()
18297 .unwrap(),
18298 );
18299 let holder_call = {
18300 let agent = Arc::clone(&agent);
18301 tokio::spawn(async move {
18302 agent
18303 .invoke_tool(ToolExecutionRequest::new(
18304 "state-lock-holder",
18305 "state_lock_holder",
18306 serde_json::json!({"path": "./shared-state-path.txt"}),
18307 ToolCallSource::Manual,
18308 ))
18309 .await
18310 .unwrap()
18311 })
18312 };
18313 holder_gate.wait_until_entered().await;
18314 let waiter_call = {
18315 let agent = Arc::clone(&agent);
18316 tokio::spawn(async move {
18317 agent
18318 .invoke_tool(ToolExecutionRequest::new(
18319 "state-lock-waiter",
18320 "state_lock_waiter",
18321 serde_json::json!({"path": "./shared-state-path.txt"}),
18322 ToolCallSource::Manual,
18323 ))
18324 .await
18325 .unwrap()
18326 })
18327 };
18328
18329 wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
18330 agent.transition_to("second").await.unwrap();
18331 holder_gate.release();
18332 let holder_record = holder_call.await.unwrap();
18333 let waiter_record = waiter_call.await.unwrap();
18334
18335 assert!(holder_record.success);
18336 assert!(!waiter_record.executed);
18337 assert!(
18338 waiter_record
18339 .output
18340 .contains("state scope changed before admission")
18341 );
18342 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
18343 }
18344
18345 #[tokio::test]
18346 async fn test_state_tools_cannot_widen_top_level_grant() {
18347 let mock = mock_with_response("hello");
18348 let yaml = r#"
18349name: NarrowToolsAgent
18350system_prompt: "You are helpful."
18351tools:
18352 - calculator
18353states:
18354 initial: current
18355 states:
18356 current:
18357 tools: [datetime]
18358"#;
18359 let agent = AgentBuilder::from_yaml(yaml)
18360 .unwrap()
18361 .llm(Arc::new(mock))
18362 .auto_configure_features()
18363 .unwrap()
18364 .build()
18365 .unwrap();
18366
18367 let available = agent.get_available_tool_ids().await.unwrap();
18368 assert!(available.is_empty());
18369 }
18370
18371 #[tokio::test]
18373 async fn test_integration_tool_execution() {
18374 let mock = mock_with_responses(vec![
18376 r#"I'll calculate that for you.
18378[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
18379 "The answer is 4.",
18381 ]);
18382 let mut tools = ai_agents_tools::ToolRegistry::new();
18383 tools
18384 .register(Arc::new(ai_agents_tools::CalculatorTool))
18385 .unwrap();
18386
18387 let agent = AgentBuilder::new()
18388 .system_prompt("You are a calculator assistant.")
18389 .llm(Arc::new(mock))
18390 .tools(tools)
18391 .build()
18392 .unwrap();
18393
18394 let response = agent.chat("What is 2+2?").await.unwrap();
18395 assert!(!response.content.is_empty());
18397 }
18398
18399 #[tokio::test]
18400 async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
18401 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18402 let hooks = Arc::new(ResponseCountingHooks {
18403 responses: Arc::clone(&responses),
18404 });
18405 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
18406 let yaml = r#"
18407name: ToolRejectAgent
18408system_prompt: "You use tools when requested."
18409tools:
18410 - echo
18411hitl:
18412 tools:
18413 echo:
18414 require_approval: true
18415 approval_message: "Approve echo?"
18416"#;
18417 let agent = AgentBuilder::from_yaml(yaml)
18418 .unwrap()
18419 .llm(Arc::new(mock))
18420 .auto_configure_features()
18421 .unwrap()
18422 .hooks(hooks)
18423 .build()
18424 .unwrap();
18425
18426 let response = agent.chat("echo hello").await.unwrap();
18427
18428 assert!(
18429 response.content.contains("Operation cancelled"),
18430 "unexpected response: {}",
18431 response.content
18432 );
18433 assert_eq!(responses.load(Ordering::SeqCst), 1);
18434 let messages = agent.memory.get_messages(None).await.unwrap();
18435 assert_eq!(messages.len(), 3);
18436 assert_eq!(messages[0].content, "echo hello");
18437 assert!(messages[1].content.contains("\"tool\":\"echo\""));
18438 assert!(messages[2].content.contains("rejected by the approver"));
18439 }
18440
18441 #[tokio::test]
18442 async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
18443 use futures::StreamExt;
18444
18445 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18446 let hooks = Arc::new(ResponseCountingHooks {
18447 responses: Arc::clone(&responses),
18448 });
18449 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
18450 let yaml = r#"
18451name: ToolRejectStreamingAgent
18452system_prompt: "You use tools when requested."
18453tools:
18454 - echo
18455streaming:
18456 enabled: true
18457hitl:
18458 tools:
18459 echo:
18460 require_approval: true
18461 approval_message: "Approve echo?"
18462"#;
18463 let agent = AgentBuilder::from_yaml(yaml)
18464 .unwrap()
18465 .llm(Arc::new(mock))
18466 .auto_configure_features()
18467 .unwrap()
18468 .hooks(hooks)
18469 .build()
18470 .unwrap();
18471
18472 let mut stream = agent.chat_stream("echo hello").await.unwrap();
18473 let mut terminal_error = String::new();
18474 let mut done = false;
18475 while let Some(chunk) = stream.next().await {
18476 match chunk {
18477 StreamChunk::Error { message } => terminal_error = message,
18478 StreamChunk::Done {} => {
18479 done = true;
18480 break;
18481 }
18482 _ => {}
18483 }
18484 }
18485
18486 assert!(done);
18487 assert!(
18488 terminal_error.contains("Operation cancelled"),
18489 "unexpected terminal error: {}",
18490 terminal_error
18491 );
18492 assert_eq!(responses.load(Ordering::SeqCst), 1);
18493 let messages = agent.memory.get_messages(None).await.unwrap();
18494 assert_eq!(messages.len(), 3);
18495 assert_eq!(messages[0].content, "echo hello");
18496 assert!(messages[1].content.contains("\"tool\":\"echo\""));
18497 assert!(messages[2].content.contains("rejected by the approver"));
18498 }
18499
18500 #[tokio::test]
18501 async fn tool_hitl_rejection_preserves_legacy_error_but_finalizes_event_stream() {
18502 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
18503 let yaml = r#"
18504name: ToolRejectEventAgent
18505system_prompt: "You use tools when requested."
18506tools:
18507 - echo
18508streaming:
18509 enabled: true
18510hitl:
18511 tools:
18512 echo:
18513 require_approval: true
18514 approval_message: "Approve echo?"
18515"#;
18516 let agent = AgentBuilder::from_yaml(yaml)
18517 .unwrap()
18518 .llm(Arc::new(mock))
18519 .auto_configure_features()
18520 .unwrap()
18521 .build()
18522 .unwrap();
18523
18524 let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
18525 let mut error_seen = false;
18526 let mut final_response = None;
18527 while let Some(event) = stream.next().await {
18528 match event {
18529 AgentStreamEvent::Chunk(StreamChunk::Error { .. }) => error_seen = true,
18530 AgentStreamEvent::Final(response) => final_response = Some(response),
18531 AgentStreamEvent::Chunk(_) => {}
18532 }
18533 }
18534
18535 assert!(!error_seen);
18536 assert!(
18537 final_response
18538 .is_some_and(|response| { response.content.contains("Operation cancelled") })
18539 );
18540 }
18541
18542 #[tokio::test]
18543 async fn test_pre_response_guard_transition_skips_old_state_llm() {
18544 let mock = mock_with_response("Billing state response");
18545 let call_counter = mock.clone();
18546 let yaml = r#"
18547name: OptimizedStateAgent
18548system_prompt: "You route before answering."
18549runtime:
18550 optimization:
18551 enabled: true
18552 pre_response_deterministic_transitions: true
18553states:
18554 initial: greeting
18555 states:
18556 greeting:
18557 prompt: "Old state prompt that should be skipped."
18558 transitions:
18559 - to: billing
18560 guard:
18561 context:
18562 topic:
18563 eq: billing
18564 timing: pre_response
18565 billing:
18566 prompt: "Answer from the billing state."
18567"#;
18568 let agent = AgentBuilder::from_yaml(yaml)
18569 .unwrap()
18570 .llm(Arc::new(mock))
18571 .build()
18572 .unwrap();
18573 agent
18574 .set_context("topic", serde_json::json!("billing"))
18575 .unwrap();
18576
18577 let response = agent.chat("I need billing help").await.unwrap();
18578
18579 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18580 assert_eq!(response.content, "Billing state response");
18581 assert_eq!(call_counter.call_count(), 1);
18582 assert_eq!(agent.actor_facts().len(), 0);
18583 }
18584
18585 #[tokio::test]
18586 async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
18587 let mock = mock_with_response("Billing state response");
18588 let call_counter = mock.clone();
18589 let yaml = r#"
18590name: OptimizedStateAgent
18591system_prompt: "You route before answering."
18592runtime:
18593 optimization:
18594 enabled: true
18595 pre_response_deterministic_transitions: true
18596context:
18597 request:
18598 type: runtime
18599 default:
18600 topic: general
18601states:
18602 initial: greeting
18603 states:
18604 greeting:
18605 prompt: "Old state prompt that should be skipped."
18606 transitions:
18607 - to: billing
18608 guard:
18609 context:
18610 request.topic:
18611 eq: billing
18612 timing: pre_response
18613 billing:
18614 prompt: "Answer from the billing state."
18615"#;
18616 let agent = AgentBuilder::from_yaml(yaml)
18617 .unwrap()
18618 .llm(Arc::new(mock))
18619 .build()
18620 .unwrap();
18621 agent
18622 .set_context("request.topic", serde_json::json!("billing"))
18623 .unwrap();
18624
18625 let response = agent.chat("I need billing help").await.unwrap();
18626
18627 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18628 assert_eq!(response.content, "Billing state response");
18629 assert_eq!(call_counter.call_count(), 1);
18630 assert_eq!(
18631 agent.get_context().get("request"),
18632 Some(&serde_json::json!({"topic": "billing"}))
18633 );
18634 }
18635
18636 #[tokio::test]
18637 async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
18638 let mock = mock_with_response("billing");
18639 let yaml = r#"
18640name: OptimizedStateAgent
18641system_prompt: "You route before answering."
18642runtime:
18643 optimization:
18644 enabled: true
18645 pre_response_deterministic_transitions: true
18646hitl:
18647 states:
18648 billing:
18649 on_enter: require_approval
18650 approval_message: "Approve billing route?"
18651states:
18652 initial: greeting
18653 states:
18654 greeting:
18655 prompt: "Old state prompt."
18656 extract:
18657 - key: topic
18658 description: "Support topic"
18659 transitions:
18660 - to: billing
18661 guard:
18662 context:
18663 topic:
18664 eq: billing
18665 timing: pre_response
18666 run_extractors: true
18667 billing:
18668 prompt: "Billing state."
18669"#;
18670 let agent = AgentBuilder::from_yaml(yaml)
18671 .unwrap()
18672 .llm(Arc::new(mock))
18673 .build()
18674 .unwrap();
18675
18676 let response = agent
18677 .try_pre_response_transition("billing please")
18678 .await
18679 .unwrap();
18680
18681 assert!(response.is_none());
18682 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
18683 assert!(!agent.get_context().contains_key("topic"));
18684 assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
18685 }
18686
18687 #[tokio::test]
18688 async fn test_pre_response_extractor_commits_context_on_winning_path() {
18689 let mock = mock_with_responses(vec!["billing", "Billing response"]);
18690 let yaml = r#"
18691name: OptimizedStateAgent
18692system_prompt: "You route before answering."
18693runtime:
18694 optimization:
18695 enabled: true
18696 pre_response_deterministic_transitions: true
18697states:
18698 initial: greeting
18699 states:
18700 greeting:
18701 prompt: "Old state prompt."
18702 extract:
18703 - key: topic
18704 description: "Support topic"
18705 transitions:
18706 - to: billing
18707 guard:
18708 context:
18709 topic:
18710 eq: billing
18711 timing: pre_response
18712 run_extractors: true
18713 billing:
18714 prompt: "Billing state."
18715"#;
18716 let agent = AgentBuilder::from_yaml(yaml)
18717 .unwrap()
18718 .llm(Arc::new(mock))
18719 .build()
18720 .unwrap();
18721
18722 let response = agent.chat("billing please").await.unwrap();
18723
18724 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18725 assert_eq!(response.content, "Billing response");
18726 assert_eq!(
18727 agent.get_context().get("topic"),
18728 Some(&serde_json::json!("billing"))
18729 );
18730 }
18731
18732 #[tokio::test]
18733 async fn test_pre_response_extractor_miss_does_not_mutate_context() {
18734 let mock = mock_with_response("__NONE__");
18735 let yaml = r#"
18736name: OptimizedStateAgent
18737system_prompt: "You route before answering."
18738runtime:
18739 optimization:
18740 enabled: true
18741 pre_response_deterministic_transitions: true
18742states:
18743 initial: greeting
18744 states:
18745 greeting:
18746 prompt: "Old state prompt."
18747 extract:
18748 - key: topic
18749 description: "Support topic"
18750 transitions:
18751 - to: billing
18752 guard:
18753 context:
18754 topic:
18755 eq: billing
18756 timing: pre_response
18757 run_extractors: true
18758 billing:
18759 prompt: "Billing state."
18760"#;
18761 let agent = AgentBuilder::from_yaml(yaml)
18762 .unwrap()
18763 .llm(Arc::new(mock))
18764 .build()
18765 .unwrap();
18766
18767 let response = agent.try_pre_response_transition("hello").await.unwrap();
18768
18769 assert!(response.is_none());
18770 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
18771 assert!(!agent.get_context().contains_key("topic"));
18772 }
18773
18774 #[tokio::test]
18775 async fn test_default_guard_transition_stays_post_response() {
18776 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
18777 let call_counter = mock.clone();
18778 let yaml = r#"
18779name: TimingAgent
18780system_prompt: "You route carefully."
18781runtime:
18782 optimization:
18783 enabled: true
18784 pre_response_deterministic_transitions: true
18785states:
18786 initial: greeting
18787 states:
18788 greeting:
18789 prompt: "Old state prompt."
18790 transitions:
18791 - to: billing
18792 guard:
18793 context:
18794 topic:
18795 eq: billing
18796 billing:
18797 prompt: "Billing state."
18798"#;
18799 let agent = AgentBuilder::from_yaml(yaml)
18800 .unwrap()
18801 .llm(Arc::new(mock))
18802 .build()
18803 .unwrap();
18804 agent
18805 .set_context("topic", serde_json::json!("billing"))
18806 .unwrap();
18807
18808 let response = agent.chat("billing please").await.unwrap();
18809
18810 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18811 assert_eq!(response.content, "Billing response");
18812 assert_eq!(call_counter.call_count(), 2);
18813 }
18814
18815 #[tokio::test]
18816 async fn test_explicit_post_response_guard_transition_stays_post_response() {
18817 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
18818 let call_counter = mock.clone();
18819 let yaml = r#"
18820name: TimingAgent
18821system_prompt: "You route carefully."
18822runtime:
18823 optimization:
18824 enabled: true
18825 pre_response_deterministic_transitions: true
18826states:
18827 initial: greeting
18828 states:
18829 greeting:
18830 prompt: "Old state prompt."
18831 transitions:
18832 - to: billing
18833 guard:
18834 context:
18835 topic:
18836 eq: billing
18837 timing: post_response
18838 billing:
18839 prompt: "Billing state."
18840"#;
18841 let agent = AgentBuilder::from_yaml(yaml)
18842 .unwrap()
18843 .llm(Arc::new(mock))
18844 .build()
18845 .unwrap();
18846 agent
18847 .set_context("topic", serde_json::json!("billing"))
18848 .unwrap();
18849
18850 let response = agent.chat("billing please").await.unwrap();
18851
18852 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18853 assert_eq!(response.content, "Billing response");
18854 assert_eq!(call_counter.call_count(), 2);
18855 }
18856
18857 #[tokio::test]
18858 async fn test_pre_response_extractors_are_transition_scoped() {
18859 let mock = mock_with_responses(vec!["billing", "Billing response"]);
18860 let yaml = r#"
18861name: ScopedExtractorAgent
18862system_prompt: "You route carefully."
18863runtime:
18864 optimization:
18865 enabled: true
18866 pre_response_deterministic_transitions: true
18867states:
18868 initial: greeting
18869 states:
18870 greeting:
18871 prompt: "Old state prompt."
18872 extract:
18873 - key: topic
18874 description: "Support topic"
18875 transitions:
18876 - to: wrong
18877 guard:
18878 context:
18879 topic:
18880 eq: billing
18881 timing: pre_response
18882 - to: billing
18883 guard:
18884 context:
18885 topic:
18886 eq: billing
18887 timing: pre_response
18888 run_extractors: true
18889 wrong:
18890 prompt: "Wrong state."
18891 billing:
18892 prompt: "Billing state."
18893"#;
18894 let agent = AgentBuilder::from_yaml(yaml)
18895 .unwrap()
18896 .llm(Arc::new(mock))
18897 .build()
18898 .unwrap();
18899
18900 let response = agent.chat("billing please").await.unwrap();
18901
18902 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18903 assert_eq!(response.content, "Billing response");
18904 }
18905
18906 #[tokio::test]
18907 async fn test_pre_response_resolved_intent_routes_early() {
18908 let mock = mock_with_response("Billing response");
18909 let yaml = r#"
18910name: IntentAgent
18911system_prompt: "You route carefully."
18912runtime:
18913 optimization:
18914 enabled: true
18915 pre_response_deterministic_transitions: true
18916states:
18917 initial: greeting
18918 states:
18919 greeting:
18920 prompt: "Old state prompt."
18921 transitions:
18922 - to: billing
18923 intent: billing
18924 timing: pre_response
18925 billing:
18926 prompt: "Billing state."
18927"#;
18928 let agent = AgentBuilder::from_yaml(yaml)
18929 .unwrap()
18930 .llm(Arc::new(mock))
18931 .build()
18932 .unwrap();
18933 agent
18934 .set_context("resolved_intent", serde_json::json!("billing"))
18935 .unwrap();
18936
18937 let response = agent
18938 .try_pre_response_transition("I need billing help")
18939 .await
18940 .unwrap()
18941 .unwrap();
18942
18943 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18944 assert_eq!(response.content, "Billing response");
18945 }
18946
18947 #[tokio::test]
18948 async fn test_background_overflow_error_surfaces() {
18949 let mut config = RuntimeConfig::default();
18950 config.optimization.enabled = true;
18951 config.optimization.post_turn.max_background_tasks = 1;
18952 config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
18953 let policy = crate::optimization::MaintenanceTaskPolicy {
18954 mode: MaintenanceMode::Background,
18955 await_before_next_turn: AwaitBeforeNextTurn::Always,
18956 };
18957 let agent = AgentBuilder::new()
18958 .system_prompt("You are helpful.")
18959 .llm(Arc::new(mock_with_response("ok")))
18960 .build()
18961 .unwrap()
18962 .with_runtime_config(config);
18963 agent
18964 .background_maintenance
18965 .spawn(None, async { std::future::pending::<Result<()>>().await })
18966 .unwrap();
18967
18968 let result = agent
18969 .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
18970 .await;
18971
18972 assert!(result.is_err());
18973 }
18974
18975 #[tokio::test]
18976 async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
18977 let default_mock = mock_with_response("Plain draft response");
18978 let router_mock = mock_with_response("cot");
18979 let router_counter = router_mock.clone();
18980 let yaml = r#"
18981name: ReasoningReservationAgent
18982system_prompt: "You answer plainly unless reasoning wins."
18983llm:
18984 default: default
18985 router: router
18986observability:
18987 enabled: true
18988 export:
18989 write_raw_events: true
18990reasoning:
18991 mode: auto
18992 judge_llm: router
18993runtime:
18994 optimization:
18995 enabled: true
18996 max_speculative_llm_calls_per_turn: 1
18997 speculative_reasoning_auto: true
18998 max_parallel_runtime_tasks: 2
18999"#;
19000 let agent = AgentBuilder::from_yaml(yaml)
19001 .unwrap()
19002 .llm_alias("default", Arc::new(default_mock))
19003 .llm_alias("router", Arc::new(router_mock))
19004 .build()
19005 .unwrap();
19006
19007 let response = agent.chat("hello").await.unwrap();
19008
19009 assert_eq!(response.content, "Plain draft response");
19010 assert_eq!(router_counter.call_count(), 1);
19011 let events = agent.observability().unwrap().raw_events();
19012 assert!(!events.iter().any(|event| {
19013 event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
19014 }));
19015 }
19016
19017 #[tokio::test]
19018 async fn test_forced_reasoning_skips_plain_speculative_draft() {
19019 let mock = mock_with_response("Reasoned response");
19020 let yaml = r#"
19021name: ForcedReasoningAgent
19022system_prompt: "You reason before answering."
19023observability:
19024 enabled: true
19025 export:
19026 write_raw_events: true
19027reasoning:
19028 mode: cot
19029runtime:
19030 optimization:
19031 enabled: true
19032 max_speculative_llm_calls_per_turn: 2
19033 speculative_state_transitions: true
19034 max_parallel_runtime_tasks: 2
19035states:
19036 initial: triage
19037 states:
19038 triage:
19039 prompt: "Answer from triage."
19040 transitions:
19041 - to: billing
19042 guard:
19043 context:
19044 route:
19045 eq: billing
19046 timing: parallel
19047 billing:
19048 prompt: "Billing state."
19049"#;
19050 let agent = AgentBuilder::from_yaml(yaml)
19051 .unwrap()
19052 .llm(Arc::new(mock))
19053 .build()
19054 .unwrap();
19055
19056 let response = agent.chat("hello").await.unwrap();
19057
19058 assert_eq!(response.content, "Reasoned response");
19059 let events = agent.observability().unwrap().raw_events();
19060 assert!(
19061 !events
19062 .iter()
19063 .any(|event| event.dimensions.contains_key("branch_status"))
19064 );
19065 }
19066
19067 #[tokio::test]
19068 async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
19069 let default_mock = mock_with_response("Skill committed response");
19070 let router_mock = mock_with_response("helper");
19071 let router_counter = router_mock.clone();
19072 let yaml = r#"
19073name: SkillReservationAgent
19074system_prompt: "Use skills when they match."
19075llm:
19076 default: default
19077 router: router
19078observability:
19079 enabled: true
19080 export:
19081 write_raw_events: true
19082runtime:
19083 optimization:
19084 enabled: true
19085 max_speculative_llm_calls_per_turn: 1
19086 speculative_skill_routing: true
19087 max_parallel_runtime_tasks: 2
19088skills:
19089 - id: helper
19090 description: "Answer helper requests"
19091 trigger: "User asks for helper"
19092 steps:
19093 - prompt: "Answer the helper request: {{ user_input }}"
19094"#;
19095 let agent = AgentBuilder::from_yaml(yaml)
19096 .unwrap()
19097 .llm_alias("default", Arc::new(default_mock))
19098 .llm_alias("router", Arc::new(router_mock))
19099 .build()
19100 .unwrap();
19101
19102 let response = agent.chat("please use helper").await.unwrap();
19103
19104 assert_eq!(response.content, "Skill committed response");
19105 assert_eq!(router_counter.call_count(), 1);
19106 let events = agent.observability().unwrap().raw_events();
19107 assert!(
19108 !events
19109 .iter()
19110 .any(|event| event.dimensions.contains_key("branch_status"))
19111 );
19112 }
19113
19114 #[tokio::test]
19115 async fn test_parallel_transition_low_cap_allows_deterministic_route() {
19116 let mock = mock_with_response("unused");
19117 let call_counter = mock.clone();
19118 let yaml = r#"
19119name: ParallelTransitionLowCapAgent
19120system_prompt: "Route before stale responses when safe."
19121runtime:
19122 optimization:
19123 enabled: true
19124 max_speculative_llm_calls_per_turn: 1
19125 speculative_state_transitions: true
19126 max_parallel_runtime_tasks: 2
19127states:
19128 initial: triage
19129 states:
19130 triage:
19131 prompt: "Triage state."
19132 transitions:
19133 - to: billing
19134 guard:
19135 context:
19136 route:
19137 eq: billing
19138 timing: parallel
19139 billing:
19140 prompt: "Billing state."
19141"#;
19142 let agent = AgentBuilder::from_yaml(yaml)
19143 .unwrap()
19144 .llm(Arc::new(mock))
19145 .build()
19146 .unwrap();
19147 agent
19148 .set_context("route", serde_json::json!("billing"))
19149 .unwrap();
19150 agent.update_active_turn_context("billing help", HashMap::new());
19151 assert!(
19152 agent.reserve_active_speculative_llm_call(
19153 RuntimeOptimizationKind::ParallelStateTransition
19154 )
19155 );
19156
19157 let selection = agent
19158 .select_parallel_transition_candidate("billing help")
19159 .await
19160 .unwrap();
19161 agent.end_root_turn();
19162
19163 match selection {
19164 ParallelTransitionSelection::Candidate(candidate) => {
19165 assert_eq!(candidate.target(), "billing");
19166 }
19167 ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
19168 ParallelTransitionSelection::ReservationExhausted => {
19169 panic!("deterministic route consumed LLM budget")
19170 }
19171 }
19172 assert_eq!(call_counter.call_count(), 0);
19173 }
19174
19175 #[tokio::test]
19176 async fn speculative_transition_drops_loser_before_state_actions() {
19177 let lock = Arc::new(tokio::sync::Mutex::new(()));
19178 let first_started = Arc::new(tokio::sync::Notify::new());
19179 let first_dropped = Arc::new(AtomicBool::new(false));
19180 let committed_after_drop = Arc::new(AtomicBool::new(false));
19181 let default = Arc::new(FirstCallLockingProvider {
19182 lock,
19183 first_started: Arc::clone(&first_started),
19184 first_dropped: Arc::clone(&first_dropped),
19185 committed_after_drop: Arc::clone(&committed_after_drop),
19186 calls: AtomicU64::new(0),
19187 });
19188 let router = Arc::new(RoutingAfterProviderStart {
19189 provider_started: first_started,
19190 });
19191 let yaml = r#"
19192name: SpeculativeCancellationAgent
19193system_prompt: "Route before committed work."
19194llm:
19195 default: default
19196 router: router
19197runtime:
19198 optimization:
19199 enabled: true
19200 max_speculative_llm_calls_per_turn: 2
19201 speculative_state_transitions: true
19202 max_parallel_runtime_tasks: 2
19203states:
19204 initial: triage
19205 states:
19206 triage:
19207 prompt: "Triage state."
19208 transitions:
19209 - to: technical
19210 when: "The request needs technical support"
19211 timing: parallel
19212 technical:
19213 prompt: "Technical state."
19214 on_enter:
19215 - prompt: "Prepare technical context."
19216 llm: default
19217 store_as: preparation
19218"#;
19219 let agent = AgentBuilder::from_yaml(yaml)
19220 .unwrap()
19221 .llm_alias("default", default)
19222 .llm_alias("router", router)
19223 .build()
19224 .unwrap();
19225
19226 let response = tokio::time::timeout(
19227 std::time::Duration::from_secs(2),
19228 agent.chat("I cannot log in because of AUTH-17."),
19229 )
19230 .await
19231 .expect("committed work must not wait on the losing provider future")
19232 .unwrap();
19233
19234 assert_eq!(response.content, "Committed technical response.");
19235 assert_eq!(agent.current_state().as_deref(), Some("technical"));
19236 assert!(first_dropped.load(Ordering::SeqCst));
19237 assert!(committed_after_drop.load(Ordering::SeqCst));
19238 }
19239
19240 #[tokio::test]
19241 async fn buffered_transition_drops_stale_stream_before_redispatch() {
19242 use futures::StreamExt;
19243
19244 let lock = Arc::new(tokio::sync::Mutex::new(()));
19245 let stream_started = Arc::new(tokio::sync::Notify::new());
19246 let stream_dropped = Arc::new(AtomicBool::new(false));
19247 let committed_after_drop = Arc::new(AtomicBool::new(false));
19248 let default = Arc::new(BufferedLockingProvider {
19249 lock,
19250 stream_started: Arc::clone(&stream_started),
19251 stream_dropped: Arc::clone(&stream_dropped),
19252 committed_after_drop: Arc::clone(&committed_after_drop),
19253 });
19254 let router = Arc::new(RoutingAfterProviderStart {
19255 provider_started: stream_started,
19256 });
19257 let yaml = r#"
19258name: BufferedCancellationAgent
19259system_prompt: "Hide stale streamed output."
19260llm:
19261 default: default
19262 router: router
19263streaming:
19264 enabled: true
19265 buffer_size: 8
19266runtime:
19267 optimization:
19268 enabled: true
19269 max_speculative_llm_calls_per_turn: 2
19270 speculative_state_transitions: true
19271 streaming_policy: buffer_until_routing_done
19272 max_parallel_runtime_tasks: 2
19273states:
19274 initial: triage
19275 states:
19276 triage:
19277 prompt: "Triage state."
19278 transitions:
19279 - to: technical
19280 when: "The request needs technical support"
19281 timing: parallel
19282 technical:
19283 prompt: "Technical state."
19284"#;
19285 let agent = AgentBuilder::from_yaml(yaml)
19286 .unwrap()
19287 .llm_alias("default", default)
19288 .llm_alias("router", router)
19289 .build()
19290 .unwrap();
19291
19292 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
19293 let mut stream = agent
19294 .chat_stream("AUTH-17 needs technical help.")
19295 .await
19296 .unwrap();
19297 let mut content = String::new();
19298 while let Some(chunk) = stream.next().await {
19299 match chunk {
19300 StreamChunk::Content { text } => content.push_str(&text),
19301 StreamChunk::Done {} => break,
19302 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
19303 _ => {}
19304 }
19305 }
19306 content
19307 })
19308 .await
19309 .expect("redispatch must not wait on the stale streaming future");
19310
19311 assert_eq!(content, "Committed technical response.");
19312 assert_eq!(agent.current_state().as_deref(), Some("technical"));
19313 assert!(stream_dropped.load(Ordering::SeqCst));
19314 assert!(committed_after_drop.load(Ordering::SeqCst));
19315 }
19316
19317 #[tokio::test]
19318 async fn buffered_transition_drops_established_stream_before_redispatch() {
19319 use futures::StreamExt;
19320
19321 let stream_started = Arc::new(tokio::sync::Notify::new());
19322 let stream_dropped = Arc::new(AtomicBool::new(false));
19323 let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
19324 let committed_after_drop = Arc::new(AtomicBool::new(false));
19325 let default = Arc::new(EstablishedStreamProvider {
19326 stream_started: Arc::clone(&stream_started),
19327 stream_dropped: Arc::clone(&stream_dropped),
19328 stream_dropped_notify,
19329 committed_after_drop: Arc::clone(&committed_after_drop),
19330 });
19331 let router = Arc::new(RoutingAfterProviderStart {
19332 provider_started: stream_started,
19333 });
19334 let yaml = r#"
19335name: EstablishedStreamCancellationAgent
19336system_prompt: "Hide stale streamed output."
19337llm:
19338 default: default
19339 router: router
19340streaming:
19341 enabled: true
19342 buffer_size: 8
19343runtime:
19344 optimization:
19345 enabled: true
19346 max_speculative_llm_calls_per_turn: 2
19347 speculative_state_transitions: true
19348 streaming_policy: buffer_until_routing_done
19349 max_parallel_runtime_tasks: 2
19350states:
19351 initial: triage
19352 states:
19353 triage:
19354 prompt: "Triage state."
19355 transitions:
19356 - to: technical
19357 when: "The request needs technical support"
19358 timing: parallel
19359 technical:
19360 prompt: "Technical state."
19361"#;
19362 let agent = AgentBuilder::from_yaml(yaml)
19363 .unwrap()
19364 .llm_alias("default", default)
19365 .llm_alias("router", router)
19366 .build()
19367 .unwrap();
19368
19369 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
19370 let mut stream = agent
19371 .chat_stream("AUTH-17 needs technical help.")
19372 .await
19373 .unwrap();
19374 let mut content = String::new();
19375 while let Some(chunk) = stream.next().await {
19376 match chunk {
19377 StreamChunk::Content { text } => content.push_str(&text),
19378 StreamChunk::Done {} => break,
19379 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
19380 _ => {}
19381 }
19382 }
19383 content
19384 })
19385 .await
19386 .expect("redispatch must wait for the established stale stream to be dropped");
19387
19388 assert_eq!(content, "Committed technical response.");
19389 assert_eq!(agent.current_state().as_deref(), Some("technical"));
19390 assert!(stream_dropped.load(Ordering::SeqCst));
19391 assert!(committed_after_drop.load(Ordering::SeqCst));
19392 }
19393
19394 #[tokio::test]
19395 async fn test_buffered_streaming_transition_reservation_falls_back() {
19396 use futures::StreamExt;
19397
19398 let mock = mock_with_responses(vec![
19399 "Serial streaming response",
19400 "Serial streaming response",
19401 ]);
19402 let router_mock = mock_with_response("1");
19403 let router_counter = router_mock.clone();
19404 let yaml = r#"
19405name: BufferedReservationFallbackAgent
19406system_prompt: "Stream normally if speculative routing cannot be evaluated."
19407llm:
19408 default: default
19409 router: router
19410observability:
19411 enabled: true
19412 export:
19413 write_raw_events: true
19414streaming:
19415 enabled: true
19416 buffer_size: 8
19417runtime:
19418 optimization:
19419 enabled: true
19420 max_speculative_llm_calls_per_turn: 1
19421 speculative_state_transitions: true
19422 streaming_policy: buffer_until_routing_done
19423 max_parallel_runtime_tasks: 2
19424states:
19425 initial: triage
19426 states:
19427 triage:
19428 prompt: "Triage state."
19429 transitions:
19430 - to: billing
19431 guard:
19432 context:
19433 route:
19434 eq: billing
19435 when: "User asks about billing"
19436 timing: parallel
19437 billing:
19438 prompt: "Billing state."
19439"#;
19440 let agent = AgentBuilder::from_yaml(yaml)
19441 .unwrap()
19442 .llm_alias("default", Arc::new(mock))
19443 .llm_alias("router", Arc::new(router_mock))
19444 .build()
19445 .unwrap();
19446
19447 let mut stream = agent.chat_stream("hello").await.unwrap();
19448 let mut content = String::new();
19449 let mut error = None;
19450 while let Some(chunk) = stream.next().await {
19451 match chunk {
19452 StreamChunk::Content { text } => content.push_str(&text),
19453 StreamChunk::Error { message } => error = Some(message),
19454 StreamChunk::Done {} => break,
19455 _ => {}
19456 }
19457 }
19458
19459 assert_eq!(error, None);
19460 assert_eq!(content, "Serial streaming response");
19461 assert_eq!(router_counter.call_count(), 0);
19462 let events = agent.observability().unwrap().raw_events();
19463 assert!(events.iter().any(|event| {
19464 event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
19465 && event.dimensions.get("commit_behavior")
19466 == Some(&"transition_decision".to_string())
19467 }));
19468 }
19469
19470 #[tokio::test]
19471 async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
19472 let mut mock = mock_with_response("Recovered response");
19473 mock.set_error("boom");
19474 let mut handle = mock.clone();
19475 let agent = AgentBuilder::new()
19476 .system_prompt("You are helpful.")
19477 .llm(Arc::new(mock))
19478 .build()
19479 .unwrap();
19480
19481 assert!(agent.chat("first").await.is_err());
19482 handle.clear_error();
19483 let response = agent.chat("second").await.unwrap();
19484
19485 assert_eq!(response.content, "Recovered response");
19486 let messages = agent.memory.get_messages(None).await.unwrap();
19487 let user_count = messages
19488 .iter()
19489 .filter(|message| message.role == ai_agents_core::Role::User)
19490 .count();
19491 assert_eq!(user_count, 2);
19492 }
19493
19494 #[tokio::test]
19495 async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
19496 use futures::StreamExt;
19497
19498 let mut mock = mock_with_response("Recovered response");
19499 mock.set_error("stream boom");
19500 let mut handle = mock.clone();
19501 let agent = AgentBuilder::new()
19502 .system_prompt("You are helpful.")
19503 .llm(Arc::new(mock))
19504 .build()
19505 .unwrap();
19506
19507 let mut stream = agent.chat_stream("first").await.unwrap();
19508 let mut saw_error = false;
19509 while let Some(chunk) = stream.next().await {
19510 if matches!(chunk, StreamChunk::Error { .. }) {
19511 saw_error = true;
19512 }
19513 }
19514 assert!(saw_error);
19515
19516 handle.clear_error();
19517 let response = agent.chat("second").await.unwrap();
19518
19519 assert_eq!(response.content, "Recovered response");
19520 let messages = agent.memory.get_messages(None).await.unwrap();
19521 let user_count = messages
19522 .iter()
19523 .filter(|message| message.role == ai_agents_core::Role::User)
19524 .count();
19525 assert_eq!(user_count, 2);
19526 }
19527
19528 #[tokio::test]
19529 async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
19530 use futures::StreamExt;
19531
19532 let mut mock = mock_with_response("one two three");
19533 mock.set_latency(10);
19534 let yaml = r#"
19535name: BufferedMissAgent
19536system_prompt: "You stream safely."
19537llm:
19538 default: default
19539streaming:
19540 enabled: true
19541 buffer_size: 1
19542runtime:
19543 optimization:
19544 enabled: true
19545 max_speculative_llm_calls_per_turn: 2
19546 speculative_state_transitions: true
19547 streaming_policy: buffer_until_routing_done
19548 max_parallel_runtime_tasks: 2
19549states:
19550 initial: triage
19551 states:
19552 triage:
19553 prompt: "Answer from triage."
19554 transitions:
19555 - to: billing
19556 guard:
19557 context:
19558 route:
19559 eq: billing
19560 timing: parallel
19561 billing:
19562 prompt: "Billing state."
19563"#;
19564 let agent = AgentBuilder::from_yaml(yaml)
19565 .unwrap()
19566 .llm_alias("default", Arc::new(mock))
19567 .build()
19568 .unwrap();
19569
19570 let mut stream = agent.chat_stream("hello").await.unwrap();
19571 let mut content = String::new();
19572 let mut error = None;
19573 while let Some(chunk) = stream.next().await {
19574 match chunk {
19575 StreamChunk::Content { text } => content.push_str(&text),
19576 StreamChunk::Error { message } => error = Some(message),
19577 StreamChunk::Done {} => break,
19578 _ => {}
19579 }
19580 }
19581
19582 assert_eq!(error, None);
19583 assert_eq!(content, "one two three");
19584 }
19585
19586 #[tokio::test]
19587 async fn test_buffered_streaming_main_failure_finalizes_branch() {
19588 use futures::StreamExt;
19589
19590 let mock = mock_with_response("one two");
19591 let mut router_mock = mock_with_response("0");
19592 router_mock.set_latency(50);
19593 let yaml = r#"
19594name: BufferedFailureAgent
19595system_prompt: "You stream safely."
19596llm:
19597 default: default
19598 router: router
19599observability:
19600 enabled: true
19601 export:
19602 write_raw_events: true
19603streaming:
19604 enabled: true
19605 buffer_size: 1
19606runtime:
19607 optimization:
19608 enabled: true
19609 max_speculative_llm_calls_per_turn: 2
19610 speculative_state_transitions: true
19611 streaming_policy: buffer_until_routing_done
19612 max_parallel_runtime_tasks: 2
19613states:
19614 initial: triage
19615 states:
19616 triage:
19617 prompt: "Ask for the category."
19618 transitions:
19619 - to: billing
19620 when: "User asks about billing"
19621 timing: parallel
19622 billing:
19623 prompt: "Billing state."
19624"#;
19625 let agent = AgentBuilder::from_yaml(yaml)
19626 .unwrap()
19627 .llm_alias("default", Arc::new(mock))
19628 .llm_alias("router", Arc::new(router_mock))
19629 .build()
19630 .unwrap();
19631
19632 let mut stream = agent.chat_stream("hello").await.unwrap();
19633 let mut error = String::new();
19634 while let Some(chunk) = stream.next().await {
19635 if let StreamChunk::Error { message } = chunk {
19636 error = message;
19637 }
19638 }
19639
19640 assert!(
19641 error.contains("stream buffer filled"),
19642 "unexpected stream error: {}",
19643 error
19644 );
19645 let events = agent.observability().unwrap().raw_events();
19646 assert!(events.iter().any(|event| {
19647 event.dimensions.get("branch_status") == Some(&"failed".to_string())
19648 && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
19649 && event.dimensions.get("optimization")
19650 == Some(&"buffered_streaming_routing".to_string())
19651 }));
19652 }
19653
19654 #[tokio::test]
19655 async fn test_streaming_preflight_does_not_emit_old_state_content() {
19656 use futures::StreamExt;
19657
19658 let mock = mock_with_response("Billing streamed response");
19659 let yaml = r#"
19660name: StreamingOptimizedAgent
19661system_prompt: "You route before streaming."
19662runtime:
19663 optimization:
19664 enabled: true
19665 pre_response_deterministic_transitions: true
19666streaming:
19667 enabled: true
19668states:
19669 initial: greeting
19670 states:
19671 greeting:
19672 prompt: "OLD_STATE_SENTINEL"
19673 transitions:
19674 - to: billing
19675 guard:
19676 context:
19677 topic:
19678 eq: billing
19679 timing: pre_response
19680 billing:
19681 prompt: "Billing state."
19682"#;
19683 let agent = AgentBuilder::from_yaml(yaml)
19684 .unwrap()
19685 .llm(Arc::new(mock))
19686 .build()
19687 .unwrap();
19688 agent
19689 .set_context("topic", serde_json::json!("billing"))
19690 .unwrap();
19691
19692 let mut stream = agent.chat_stream("billing please").await.unwrap();
19693 let mut content = String::new();
19694 while let Some(chunk) = stream.next().await {
19695 match chunk {
19696 StreamChunk::Content { text } => content.push_str(&text),
19697 StreamChunk::Error { message } => panic!("stream error: {}", message),
19698 StreamChunk::Done {} => break,
19699 _ => {}
19700 }
19701 }
19702
19703 assert_eq!(agent.current_state().as_deref(), Some("billing"));
19704 assert!(content.contains("Billing streamed response"));
19705 assert!(!content.contains("OLD_STATE_SENTINEL"));
19706 }
19707
19708 #[tokio::test]
19710 async fn test_integration_state_machine_basic() {
19711 let yaml = r#"
19712name: StateAgent
19713system_prompt: "You are a support agent."
19714states:
19715 initial: greeting
19716 states:
19717 greeting:
19718 prompt: "Welcome the user warmly."
19719 transitions:
19720 - to: support
19721 when: "User needs help"
19722 auto: true
19723 support:
19724 prompt: "Help solve the user's problem."
19725"#;
19726 let mock = mock_with_responses(vec![
19727 "Welcome! How can I help?", "1", "I'll help you with that.", ]);
19731 let builder = AgentBuilder::from_yaml(yaml).unwrap();
19732 let agent = builder.llm(Arc::new(mock)).build().unwrap();
19733
19734 assert_eq!(agent.current_state(), Some("greeting".to_string()));
19735 let _ = agent.chat("I need help").await.unwrap();
19736 }
19739
19740 #[tokio::test]
19742 async fn test_integration_state_on_enter_set_context() {
19743 let yaml = r#"
19744name: ActionAgent
19745system_prompt: "You are helpful."
19746states:
19747 initial: step1
19748 states:
19749 step1:
19750 prompt: "Step 1"
19751 on_exit:
19752 - set_context:
19753 step1_exited: true
19754 transitions:
19755 - to: step2
19756 when: "always"
19757 auto: true
19758 step2:
19759 prompt: "Step 2"
19760 on_enter:
19761 - set_context:
19762 step2_entered: true
19763"#;
19764 let mock = mock_with_responses(vec![
19766 "Processing step 1.",
19767 "0", ]);
19769 let builder = AgentBuilder::from_yaml(yaml).unwrap();
19770 let agent = builder.llm(Arc::new(mock)).build().unwrap();
19771
19772 assert_eq!(agent.current_state(), Some("step1".to_string()));
19773
19774 agent.transition_to("step2").await.unwrap();
19776
19777 assert_eq!(agent.current_state(), Some("step2".to_string()));
19778
19779 let ctx = agent.get_context();
19781 assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
19782 assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
19783 }
19784
19785 #[tokio::test]
19786 async fn state_action_tool_preserves_source_in_stored_record() {
19787 let yaml = r#"
19788name: StateActionToolAgent
19789system_prompt: "You are helpful."
19790tools:
19791 - context_echo
19792states:
19793 initial: idle
19794 states:
19795 idle:
19796 prompt: "Idle"
19797 active:
19798 prompt: "Active"
19799 on_enter:
19800 - set_context:
19801 action_started: true
19802 - tool: context_echo
19803 args: {}
19804"#;
19805 let agent = AgentBuilder::from_yaml(yaml)
19806 .unwrap()
19807 .llm(Arc::new(mock_with_response("unused")))
19808 .tool(Arc::new(ContextEchoTool))
19809 .build()
19810 .unwrap();
19811
19812 agent.transition_to("active").await.unwrap();
19813
19814 let record: ToolExecutionRecord = serde_json::from_value(
19815 agent
19816 .get_context()
19817 .get("last_tool_record")
19818 .cloned()
19819 .expect("successful state action must store its execution record"),
19820 )
19821 .unwrap();
19822 assert!(record.executed);
19823 assert!(record.success);
19824 assert_eq!(record.canonical_id, "context_echo");
19825 assert!(matches!(
19826 &record.source,
19827 ToolCallSource::StateAction {
19828 state: Some(state),
19829 action_index: 1,
19830 } if state == "active"
19831 ));
19832 }
19833
19834 #[tokio::test]
19835 async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
19836 let yaml = r#"
19837name: OrdinaryLifecycleAgent
19838system_prompt: "You are helpful."
19839states:
19840 initial: intake
19841 regenerate_on_transition: false
19842 states:
19843 intake:
19844 prompt: "Intake"
19845 transitions:
19846 - to: drafting
19847 guard:
19848 context:
19849 route:
19850 eq: drafting
19851 drafting:
19852 prompt: "Drafting"
19853 on_enter:
19854 - set_context:
19855 draft_version: 1
19856 on_reenter:
19857 - set_context:
19858 draft_version: 2
19859 transitions:
19860 - to: review
19861 guard:
19862 context:
19863 route:
19864 eq: review
19865 review:
19866 prompt: "Review"
19867 on_enter:
19868 - set_context:
19869 review_entry: first
19870 transitions:
19871 - to: drafting
19872 guard:
19873 context:
19874 route:
19875 eq: drafting
19876"#;
19877 let agent = AgentBuilder::from_yaml(yaml)
19878 .unwrap()
19879 .llm(Arc::new(mock_with_responses(vec![
19880 "Intake response",
19881 "Draft response",
19882 "Review response",
19883 ])))
19884 .build()
19885 .unwrap();
19886
19887 agent
19888 .set_context("route", serde_json::json!("drafting"))
19889 .unwrap();
19890 agent.chat("Start a draft").await.unwrap();
19891 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19892 assert_eq!(
19893 agent.get_context().get("draft_version"),
19894 Some(&serde_json::json!(1))
19895 );
19896
19897 agent
19898 .set_context("route", serde_json::json!("review"))
19899 .unwrap();
19900 agent.chat("Review this").await.unwrap();
19901 assert_eq!(agent.current_state().as_deref(), Some("review"));
19902 assert_eq!(
19903 agent.get_context().get("review_entry"),
19904 Some(&serde_json::json!("first"))
19905 );
19906
19907 agent
19908 .set_context("route", serde_json::json!("drafting"))
19909 .unwrap();
19910 agent.chat("Revise this").await.unwrap();
19911 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19912 assert_eq!(
19913 agent.get_context().get("draft_version"),
19914 Some(&serde_json::json!(2))
19915 );
19916 }
19917
19918 #[tokio::test]
19919 async fn test_manual_transition_uses_on_enter_then_on_reenter() {
19920 let yaml = r#"
19921name: ManualLifecycleAgent
19922system_prompt: "You are helpful."
19923states:
19924 initial: intake
19925 states:
19926 intake:
19927 prompt: "Intake"
19928 drafting:
19929 prompt: "Drafting"
19930 on_enter:
19931 - set_context:
19932 draft_version: 1
19933 on_reenter:
19934 - set_context:
19935 draft_version: 2
19936 review:
19937 prompt: "Review"
19938"#;
19939 let agent = AgentBuilder::from_yaml(yaml)
19940 .unwrap()
19941 .llm(Arc::new(mock_with_response("unused")))
19942 .build()
19943 .unwrap();
19944
19945 assert!(!agent.get_context().contains_key("draft_version"));
19946 agent.transition_to("drafting").await.unwrap();
19947 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19948 assert_eq!(
19949 agent.get_context().get("draft_version"),
19950 Some(&serde_json::json!(1))
19951 );
19952
19953 agent.transition_to("review").await.unwrap();
19954 agent.transition_to("drafting").await.unwrap();
19955 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19956 assert_eq!(
19957 agent.get_context().get("draft_version"),
19958 Some(&serde_json::json!(2))
19959 );
19960 }
19961
19962 #[tokio::test]
19963 async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
19964 let yaml = r#"
19965name: TimeoutLifecycleAgent
19966system_prompt: "You are helpful."
19967states:
19968 initial: intake
19969 regenerate_on_transition: false
19970 states:
19971 intake:
19972 prompt: "Intake"
19973 max_turns: 1
19974 timeout_to: drafting
19975 drafting:
19976 prompt: "Drafting"
19977 max_turns: 1
19978 timeout_to: review
19979 on_enter:
19980 - set_context:
19981 draft_version: 1
19982 on_reenter:
19983 - set_context:
19984 draft_version: 2
19985 review:
19986 prompt: "Review"
19987 max_turns: 1
19988 timeout_to: drafting
19989 on_enter:
19990 - set_context:
19991 review_entry: first
19992"#;
19993 let agent = AgentBuilder::from_yaml(yaml)
19994 .unwrap()
19995 .llm(Arc::new(mock_with_responses(vec![
19996 "Intake",
19997 "First draft",
19998 "Review",
19999 "Revised draft",
20000 ])))
20001 .build()
20002 .unwrap();
20003
20004 agent.chat("First turn").await.unwrap();
20005 assert_eq!(agent.current_state().as_deref(), Some("intake"));
20006 assert!(!agent.get_context().contains_key("draft_version"));
20007
20008 agent.chat("Second turn").await.unwrap();
20009 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
20010 assert_eq!(
20011 agent.get_context().get("draft_version"),
20012 Some(&serde_json::json!(1))
20013 );
20014
20015 agent.chat("Third turn").await.unwrap();
20016 assert_eq!(agent.current_state().as_deref(), Some("review"));
20017 assert_eq!(
20018 agent.get_context().get("review_entry"),
20019 Some(&serde_json::json!("first"))
20020 );
20021
20022 agent.chat("Fourth turn").await.unwrap();
20023 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
20024 assert_eq!(
20025 agent.get_context().get("draft_version"),
20026 Some(&serde_json::json!(2))
20027 );
20028 }
20029
20030 #[tokio::test]
20032 async fn test_integration_process_normalize() {
20033 let yaml = r#"
20034name: ProcessAgent
20035system_prompt: "You are helpful."
20036process:
20037 input:
20038 - type: normalize
20039 config:
20040 trim: true
20041 collapse_whitespace: true
20042"#;
20043 let mock = mock_with_response("Got your message.");
20044 let builder = AgentBuilder::from_yaml(yaml).unwrap();
20045 let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
20046
20047 let _ = agent.chat(" hello world ").await.unwrap();
20048
20049 let history = mock.call_history();
20051 assert!(!history.is_empty());
20052 let last_call = history.last().unwrap();
20054 let user_msg = last_call
20055 .messages
20056 .iter()
20057 .find(|m| m.role == ai_agents_core::Role::User)
20058 .unwrap();
20059 assert_eq!(user_msg.content, "hello world");
20060 }
20061
20062 #[tokio::test]
20066 async fn test_integration_memory_compression() {
20067 let yaml = r#"
20068name: MemoryAgent
20069system_prompt: "You are helpful."
20070memory:
20071 type: compacting
20072 max_messages: 100
20073 compress_threshold: 5
20074 max_recent_messages: 3
20075 summarize_batch_size: 2
20076"#;
20077 let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
20079 let mock = mock_with_responses(responses);
20080 let builder = AgentBuilder::from_yaml(yaml).unwrap();
20081 let agent = builder.llm(Arc::new(mock)).build().unwrap();
20082
20083 for i in 0..6 {
20085 let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
20086 }
20087
20088 let messages = agent.memory.get_messages(None).await.unwrap();
20091 assert!(messages.len() <= 12); }
20095
20096 #[tokio::test]
20098 async fn test_integration_multi_llm_registry() {
20099 let mut mock_default = MockLLMProvider::new("default");
20100 mock_default.set_response("Default LLM response.");
20101 let mut mock_router = MockLLMProvider::new("router");
20102 mock_router.set_response("Router response.");
20103
20104 let agent = AgentBuilder::new()
20105 .system_prompt("You are helpful.")
20106 .llm_alias("default", Arc::new(mock_default))
20107 .llm_alias("router", Arc::new(mock_router))
20108 .build()
20109 .unwrap();
20110
20111 let response = agent.chat("Hello").await.unwrap();
20112 assert_eq!(response.content, "Default LLM response.");
20113 }
20114
20115 #[tokio::test]
20117 async fn test_integration_agent_reset() {
20118 let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
20119 let agent = AgentBuilder::new()
20120 .system_prompt("You are helpful.")
20121 .llm(Arc::new(mock))
20122 .build()
20123 .unwrap();
20124
20125 let _ = agent.chat("Hi").await.unwrap();
20126 let messages = agent.memory.get_messages(None).await.unwrap();
20127 assert_eq!(messages.len(), 2); agent.reset().await.unwrap();
20130 let messages = agent.memory.get_messages(None).await.unwrap();
20131 assert_eq!(messages.len(), 0);
20132 }
20133
20134 #[tokio::test]
20136 async fn test_integration_process_validate_reject() {
20137 use ai_agents_process::{ProcessConfig, ProcessProcessor};
20138
20139 let validate_config = ai_agents_process::ValidateStage {
20140 id: Some("length_check".to_string()),
20141 condition: None,
20142 config: ai_agents_process::ValidateConfig {
20143 rules: vec![ai_agents_process::ValidationRule::MinLength {
20144 min_length: 10,
20145 on_fail: ai_agents_process::ValidationAction {
20146 action: ai_agents_process::ValidationActionType::Reject,
20147 message: None,
20148 },
20149 }],
20150 ..Default::default()
20151 },
20152 };
20153 let process_config = ProcessConfig {
20154 input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
20155 ..Default::default()
20156 };
20157 let processor = ProcessProcessor::new(process_config);
20158
20159 let mock = mock_with_response("Should not reach here.");
20160 let agent = AgentBuilder::new()
20161 .system_prompt("You are helpful.")
20162 .llm(Arc::new(mock))
20163 .process_processor(processor)
20164 .build()
20165 .unwrap();
20166
20167 let response = agent.chat("Hi").await.unwrap();
20168 assert!(
20170 response.content.contains("rejected")
20171 || response.content.contains("Input rejected")
20172 || response.content.contains("too short")
20173 || response.content.contains("Too short")
20174 || response.content.len() < 50, "Expected rejection response, got: {}",
20176 response.content
20177 );
20178 }
20179
20180 #[tokio::test]
20182 async fn test_llm_fallback_on_failure() {
20183 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
20184
20185 let mut primary = MockLLMProvider::new("primary");
20186 primary.set_error("Primary LLM is unavailable");
20187
20188 let mut fallback = MockLLMProvider::new("fallback");
20189 fallback.set_response("Fallback response works!");
20190
20191 let agent = AgentBuilder::new()
20192 .system_prompt("You are helpful.")
20193 .llm_alias("default", Arc::new(primary))
20194 .llm_alias("backup", Arc::new(fallback))
20195 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
20196 llm: LLMRecoveryConfig {
20197 on_failure: LLMFailureAction::FallbackLlm {
20198 fallback_llm: "backup".to_string(),
20199 },
20200 ..Default::default()
20201 },
20202 ..Default::default()
20203 }))
20204 .build()
20205 .unwrap();
20206
20207 let response = agent.chat("Hello").await.unwrap();
20208 assert!(
20209 response.content.contains("Fallback response"),
20210 "Expected fallback response, got: {}",
20211 response.content
20212 );
20213 }
20214
20215 #[tokio::test]
20217 async fn test_llm_fallback_response_static_message() {
20218 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
20219
20220 let mut primary = MockLLMProvider::new("primary");
20221 primary.set_error("Primary LLM is unavailable");
20222
20223 let agent = AgentBuilder::new()
20224 .system_prompt("You are helpful.")
20225 .llm(Arc::new(primary))
20226 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
20227 llm: LLMRecoveryConfig {
20228 on_failure: LLMFailureAction::FallbackResponse {
20229 message: "I am temporarily unavailable. Please try again later."
20230 .to_string(),
20231 },
20232 ..Default::default()
20233 },
20234 ..Default::default()
20235 }))
20236 .build()
20237 .unwrap();
20238
20239 let response = agent.chat("Hello").await.unwrap();
20240 assert!(
20241 response.content.contains("temporarily unavailable"),
20242 "Expected static fallback message, got: {}",
20243 response.content
20244 );
20245 }
20246
20247 #[tokio::test]
20249 async fn test_tool_failure_skip() {
20250 use ai_agents_recovery::{
20251 ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
20252 };
20253
20254 let mock = mock_with_responses(vec![
20256 r#"I'll use the nonexistent tool.
20257[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
20258 "The tool was unavailable, but I can still help you.",
20259 ]);
20260
20261 let agent = AgentBuilder::new()
20262 .system_prompt("You are helpful.")
20263 .llm(Arc::new(mock))
20264 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
20265 tools: ToolRecoveryConfig {
20266 default: ToolRetryConfig {
20267 max_retries: 0,
20268 timeout_ms: None,
20269 on_failure: ToolFailureAction::Skip,
20270 },
20271 ..Default::default()
20272 },
20273 ..Default::default()
20274 }))
20275 .build()
20276 .unwrap();
20277
20278 let response = agent.chat("Use the nonexistent tool").await;
20280 assert!(
20281 response.is_ok(),
20282 "Expected Ok with skip policy, got: {:?}",
20283 response
20284 );
20285 }
20286}