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::{Duration, Instant};
11use tracing::{debug, error, info, instrument, warn};
12
13const DISAMBIGUATION_STATE_GENERATION_KEY: &str = "_runtime.disambiguation_state_generation";
14const MAX_TOOL_FALLBACK_HOPS: usize = 16;
15
16pub(crate) type RootTurnGate = Arc<tokio::sync::Mutex<()>>;
18
19pub(crate) type RootTurnGateIdentityStack = Arc<[RootTurnGate]>;
21
22tokio::task_local! {
23 static RUNTIME_GATE_IDENTITY_STACK: RootTurnGateIdentityStack;
24}
25
26pub(crate) fn current_runtime_gate_identity_stack() -> RootTurnGateIdentityStack {
28 RUNTIME_GATE_IDENTITY_STACK
29 .try_with(Arc::clone)
30 .unwrap_or_default()
31}
32
33pub(crate) async fn scope_runtime_gate_identity_stack<F, T>(
35 identity_stack: &RootTurnGateIdentityStack,
36 future: F,
37) -> T
38where
39 F: Future<Output = T>,
40{
41 RUNTIME_GATE_IDENTITY_STACK
42 .scope(Arc::clone(identity_stack), future)
43 .await
44}
45
46pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Weak<tokio::sync::Mutex<()>>>>>;
48
49struct ToolResourceGuards {
53 guards: Vec<tokio::sync::OwnedMutexGuard<()>>,
54 locks: ToolResourceLocks,
55}
56
57struct RootTurnAdmission {
61 guard: tokio::sync::OwnedMutexGuard<()>,
62 identity_stack: RootTurnGateIdentityStack,
63}
64
65#[derive(Clone)]
66struct StoredSessionRestore {
67 snapshot: AgentSnapshot,
68 metadata: Option<ai_agents_core::SessionMetadata>,
69}
70
71struct RuntimeSessionRestorePoint {
72 snapshot: AgentSnapshot,
73 metadata: ai_agents_core::SessionMetadata,
74 actor_id: Option<String>,
75 session_id: Option<String>,
76}
77
78impl Drop for ToolResourceGuards {
79 fn drop(&mut self) {
80 self.guards.clear();
81 self.locks.write().retain(|_, lock| lock.strong_count() > 0);
82 }
83}
84
85#[derive(Clone)]
89struct RuntimeSafetySnapshot {
90 version: u64,
91 emergency_deny: bool,
92 tool_security: ToolSecurityEngine,
93 tool_scope_override: Option<Vec<String>>,
94}
95
96#[derive(Clone, Copy)]
100struct ToolDecisionVersions {
101 policy: u64,
102 registry: u64,
103 runtime_control: u64,
104 state: Option<u64>,
105}
106
107#[derive(Clone, Debug, Default)]
111struct ToolFallbackState {
112 visited_canonical_ids: Vec<String>,
113}
114
115impl ToolFallbackState {
116 fn rejection_reason(&self, canonical_id: &str) -> Option<String> {
120 if self
121 .visited_canonical_ids
122 .iter()
123 .any(|visited| visited == canonical_id)
124 {
125 return Some(format!(
126 "Tool fallback cycle detected at '{canonical_id}' after [{}]",
127 self.visited_canonical_ids.join(" -> ")
128 ));
129 }
130 if self.visited_canonical_ids.len() > MAX_TOOL_FALLBACK_HOPS {
131 return Some(format!(
132 "Tool fallback chain exceeds the maximum of {MAX_TOOL_FALLBACK_HOPS} hops"
133 ));
134 }
135 None
136 }
137
138 fn with_current(mut self, canonical_id: String) -> Self {
142 self.visited_canonical_ids.push(canonical_id);
143 self
144 }
145
146 fn final_rejection_reason(
150 &self,
151 admitted_canonical_id: &str,
152 final_canonical_id: &str,
153 ) -> Option<String> {
154 if admitted_canonical_id == final_canonical_id {
155 return None;
156 }
157 if self
158 .visited_canonical_ids
159 .iter()
160 .any(|visited| visited == final_canonical_id)
161 {
162 return Some(format!(
163 "Tool fallback cycle detected after final resolution changed '{admitted_canonical_id}' to '{final_canonical_id}'"
164 ));
165 }
166 Some(format!(
167 "Tool canonical target changed after initial admission from '{admitted_canonical_id}' to '{final_canonical_id}'"
168 ))
169 }
170}
171
172#[derive(Clone, Copy, Debug)]
176struct ValidatedToolTimeout {
177 timer: Duration,
178 deadline_delta: chrono::Duration,
179}
180
181struct AvailableToolIdsSnapshot {
185 tool_ids: Vec<String>,
186 state_generation: Option<u64>,
187}
188
189#[derive(Clone)]
193struct ToolApprovalBinding {
194 canonical_id: String,
195 arguments: Value,
196 confirmation_required: bool,
197 policy_version: u64,
198 runtime_control_version: u64,
199 state_generation: Option<u64>,
200 reviewed_tool: Arc<dyn ai_agents_core::Tool>,
201}
202
203fn merge_approved_record(record: &mut Option<ToolApprovalRecord>) {
207 if record
208 .as_ref()
209 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::Modified))
210 {
211 return;
212 }
213 *record = Some(ToolApprovalRecord {
214 status: ToolApprovalStatus::Approved,
215 reason: None,
216 modified_arguments: None,
217 });
218}
219
220impl ToolApprovalBinding {
221 fn is_stale(
223 &self,
224 canonical_id: &str,
225 arguments: &Value,
226 confirmation_required: bool,
227 versions: ToolDecisionVersions,
228 resolved_tool: &Arc<dyn ai_agents_core::Tool>,
229 ) -> bool {
230 self.canonical_id != canonical_id
231 || self.arguments != *arguments
232 || self.confirmation_required != confirmation_required
233 || self.policy_version != versions.policy
234 || self.runtime_control_version != versions.runtime_control
235 || self.state_generation != versions.state
236 || !Arc::ptr_eq(&self.reviewed_tool, resolved_tool)
237 }
238}
239
240use crate::turn_context::{current_turn_actor_context, scope_actor_context};
241
242use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
243use ai_agents_core::traits::storage::StorageCapability;
244use ai_agents_core::{
245 AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMError, LLMProvider,
246 LLMResponse, LLMToolDefinition, LLMToolRequest, PermissionOutcome, Result, ToolActorContext,
247 ToolApprovalRecord, ToolApprovalStatus, ToolCallClassification, ToolCallSource,
248 ToolCancellationToken, ToolChoice, ToolExecutionContext, ToolExecutionLimits,
249 ToolExecutionRecord, ToolExecutionRequest, ToolInvoker, ToolPolicyDecisionRecord, ToolResult,
250 ToolSafetyMetadata,
251};
252use ai_agents_disambiguation::{
253 ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
254 ConfirmationParseFuture, DisambiguationConfig, DisambiguationContext, DisambiguationManager,
255 DisambiguationResult,
256};
257use ai_agents_hitl::{
258 ApprovalHandler, ApprovalResolvedOutcome, ApprovalResult, ApprovalTrigger, HITLCheckResult,
259 HITLEngine, RejectAllHandler, TimeoutAction,
260};
261use ai_agents_hooks::{AgentHooks, NoopHooks};
262use ai_agents_llm::LLMRegistry;
263use ai_agents_memory::{
264 CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
265 MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
266};
267use ai_agents_observability::{
268 EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
269 current_observation_context, new_session_id as new_observation_session_id,
270 resolve_language_from_context, with_observation_context, with_observation_purpose,
271};
272use ai_agents_process::{
273 ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
274};
275use ai_agents_reasoning::{
276 CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
277 ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
278 ReflectionMetadata, StepFailureAction,
279};
280use ai_agents_recovery::{
281 ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
282 LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
283};
284use ai_agents_relationships::RelationshipManager;
285use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
286use ai_agents_state::{
287 PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, Transition,
288 TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
289};
290use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
291use ai_agents_tools::{
292 CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
293 MAX_TOOL_TIMEOUT_MS, QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord,
294 ToolRegistry, ToolSecurityConfig, ToolSecurityEngine,
295};
296
297use super::{
298 Agent, AgentInfo, AgentResponse, AgentStreamEvent, ParallelToolsConfig, StreamChunk,
299 StreamingConfig, ToolCall,
300};
301use crate::optimization::{
302 AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
303 MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
304 RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
305 RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
306 StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
307};
308use crate::spec::StorageConfig;
309
310enum ToolCallOutcome {
312 Continue,
314 TransitionFired,
316 Rejected(AgentResponse),
318}
319
320#[derive(Clone)]
321struct MainToolProtocol {
322 choice: Option<ToolChoice>,
323 tool_ids: Vec<String>,
324 definitions: Vec<LLMToolDefinition>,
325}
326
327struct MainProviderResponse {
328 response: LLMResponse,
329 used_native_tools: bool,
330}
331
332struct CommittedTextResponse<'a> {
336 processed_input: &'a str,
337 input_context: &'a HashMap<String, Value>,
338 answer: String,
339 reasoning_mode: ReasoningMode,
340 auto_detected: bool,
341 iterations: u32,
342 thinking_content: Option<String>,
343 all_tool_calls: Vec<ToolCall>,
344}
345
346struct AgentResponseParts {
350 content: String,
351 all_tool_calls: Vec<ToolCall>,
352 reasoning_mode: ReasoningMode,
353 auto_detected: bool,
354 iterations: u32,
355 thinking: Option<String>,
356 reflection_metadata: Option<ReflectionMetadata>,
357}
358
359type RuntimeStreamTerminalSlot = Arc<RwLock<Option<AgentResponse>>>;
363
364fn new_runtime_stream_terminal_slot() -> RuntimeStreamTerminalSlot {
368 Arc::new(RwLock::new(None))
369}
370
371fn record_runtime_stream_final(slot: &RuntimeStreamTerminalSlot, response: AgentResponse) {
375 *slot.write() = Some(response);
376}
377
378#[derive(Clone, Copy)]
379struct DisambiguationOwnership {
380 epoch: u64,
381 state_generation: Option<u64>,
382}
383
384enum SkillRouteResult {
386 NoMatch,
388 Response { skill_id: String, content: String },
390 NeedsClarification {
392 response: AgentResponse,
393 ownership: Option<DisambiguationOwnership>,
394 },
395}
396
397enum ParallelTransitionSelection {
399 Candidate(TransitionCandidate),
401 NoMatch,
403 ReservationExhausted,
405}
406
407enum PostLoopResult {
409 NoTransition(String),
411 Transitioned(String),
413 NeedsRedispatch,
416}
417
418struct StateTransitionReservation<'a> {
419 reserved: &'a AtomicBool,
420}
421
422impl Drop for StateTransitionReservation<'_> {
423 fn drop(&mut self) {
424 self.reserved.store(false, Ordering::SeqCst);
425 }
426}
427
428struct RootTurnCleanup<'a> {
429 agent: &'a RuntimeAgent,
430}
431
432impl<'a> RootTurnCleanup<'a> {
433 fn new(agent: &'a RuntimeAgent) -> Self {
434 Self { agent }
435 }
436}
437
438impl Drop for RootTurnCleanup<'_> {
439 fn drop(&mut self) {
440 self.agent.end_root_turn();
441 }
442}
443
444#[derive(Debug)]
446struct RuntimeControlState {
447 snapshot_guard: RwLock<()>,
449 version: AtomicU64,
451 emergency_deny: Arc<AtomicBool>,
453 tool_security_override: RwLock<Option<ToolSecurityEngine>>,
455 tool_scope_override: RwLock<Option<Vec<String>>>,
457}
458
459impl Default for RuntimeControlState {
460 fn default() -> Self {
461 Self {
462 snapshot_guard: RwLock::new(()),
463 version: AtomicU64::new(1),
464 emergency_deny: Arc::new(AtomicBool::new(false)),
465 tool_security_override: RwLock::new(None),
466 tool_scope_override: RwLock::new(None),
467 }
468 }
469}
470
471#[derive(Clone)]
473pub struct RuntimeControlHandle {
474 state: Arc<RuntimeControlState>,
475}
476
477impl RuntimeControlHandle {
478 pub fn version(&self) -> u64 {
480 self.state.version.load(Ordering::SeqCst)
481 }
482
483 fn bump(&self) -> u64 {
484 self.state.version.fetch_add(1, Ordering::SeqCst) + 1
485 }
486
487 pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
489 self.try_set_tool_security(config)
490 .expect("invalid tool security configuration")
491 }
492
493 pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
495 config.validate()?;
496 let _guard = self.state.snapshot_guard.write();
497 let generation = self.bump();
498 *self.state.tool_security_override.write() = Some(
499 ToolSecurityEngine::new_with_policy_version(config, generation),
500 );
501 Ok(generation)
502 }
503
504 pub fn clear_tool_security_override(&self) -> u64 {
506 let _guard = self.state.snapshot_guard.write();
507 *self.state.tool_security_override.write() = None;
508 self.bump()
509 }
510
511 pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
513 let _guard = self.state.snapshot_guard.write();
514 *self.state.tool_scope_override.write() = Some(tool_ids);
515 self.bump()
516 }
517
518 pub fn clear_tool_scope_override(&self) -> u64 {
520 let _guard = self.state.snapshot_guard.write();
521 *self.state.tool_scope_override.write() = None;
522 self.bump()
523 }
524
525 pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
527 let _guard = self.state.snapshot_guard.write();
528 self.state.emergency_deny.store(enabled, Ordering::SeqCst);
529 self.bump()
530 }
531
532 pub fn cancel_all(&self) -> u64 {
534 self.set_emergency_deny(true)
535 }
536}
537
538pub struct RuntimeAgent {
539 info: AgentInfo,
540 llm_registry: Arc<LLMRegistry>,
541 memory: Arc<dyn Memory>,
542 tools: Arc<ToolRegistry>,
543 skills: Vec<SkillDefinition>,
544 skill_router: Option<SkillRouter>,
545 skill_executor: Option<SkillExecutor>,
546 base_system_prompt: String,
547 max_iterations: u32,
548 iteration_count: RwLock<u32>,
549 max_context_tokens: u32,
550 memory_token_budget: Option<MemoryTokenBudget>,
551 recovery_manager: RecoveryManager,
552 tool_security: ToolSecurityEngine,
553 process_processor: Option<ProcessProcessor>,
554 message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
555 state_machine: Option<Arc<StateMachine>>,
556 transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
557 context_manager: Arc<ContextManager>,
558 template_renderer: TemplateRenderer,
559 tool_call_history: RwLock<Vec<ToolCallRecord>>,
560 parallel_tools: ParallelToolsConfig,
561 streaming: StreamingConfig,
562 hooks: Arc<dyn AgentHooks>,
563 hitl_engine: Option<HITLEngine>,
564 approval_handler: Arc<dyn ApprovalHandler>,
565 storage_config: StorageConfig,
566 storage: RwLock<Option<Arc<dyn AgentStorage>>>,
567 storage_init: tokio::sync::Mutex<()>,
568 reasoning_config: ReasoningConfig,
569 reflection_config: ReflectionConfig,
570 disambiguation_manager: Option<DisambiguationManager>,
571 disambiguation_epoch: AtomicU64,
573 disambiguation_admission: tokio::sync::RwLock<()>,
575 state_transition_reserved: AtomicBool,
577 persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
579 pending_skill_id: RwLock<Option<String>>,
583 current_plan: RwLock<Option<Plan>>,
584 declared_tool_ids: Option<Vec<String>>,
586 context_initialized: AtomicBool,
588 spawner: Option<Arc<crate::spawner::AgentSpawner>>,
590 spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
592 redispatch_depth: RwLock<u32>,
595 active_turn_context: RwLock<Option<TurnOptimizationContext>>,
597 root_user_message_committed: AtomicBool,
599 actor_id: RwLock<Option<String>>,
601 fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
603 fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
606 actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
608 messages_since_extraction: Arc<RwLock<usize>>,
610 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
612 facts_config: Option<ai_agents_facts::FactsConfig>,
614 session_metadata: RwLock<ai_agents_core::SessionMetadata>,
616 current_session_id: RwLock<Option<String>>,
618 relationship_manager: Option<Arc<RelationshipManager>>,
620 observability_manager: Option<Arc<ObservabilityManager>>,
622 runtime_config: RuntimeConfig,
624 background_maintenance: Arc<BackgroundMaintenanceQueue>,
626 resource_locks: ToolResourceLocks,
628 runtime_control: Arc<RuntimeControlState>,
630 root_turn_gate: RootTurnGate,
632}
633
634impl std::fmt::Debug for RuntimeAgent {
635 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
636 f.debug_struct("RuntimeAgent")
637 .field("info", &self.info)
638 .field("base_system_prompt", &self.base_system_prompt)
639 .field("max_iterations", &self.max_iterations)
640 .field("skills_count", &self.skills.len())
641 .field("max_context_tokens", &self.max_context_tokens)
642 .field("has_state_machine", &self.state_machine.is_some())
643 .field("parallel_tools", &self.parallel_tools)
644 .field("streaming", &self.streaming)
645 .field("has_hooks", &true)
646 .field("has_hitl", &self.hitl_engine.is_some())
647 .field("storage_type", &self.storage_config.storage_type())
648 .field("reasoning_mode", &self.reasoning_config.mode)
649 .field("reflection_enabled", &self.reflection_config.enabled)
650 .field("declared_tool_ids", &self.declared_tool_ids)
651 .field("has_persona", &self.persona_manager.is_some())
652 .field("has_observability", &self.observability_manager.is_some())
653 .finish_non_exhaustive()
654 }
655}
656
657struct ObservabilityClarificationObserver;
658
659impl ClarificationObserver for ObservabilityClarificationObserver {
660 fn observe_question<'a>(
662 &'a self,
663 future: ClarificationQuestionFuture<'a>,
664 ) -> ClarificationQuestionFuture<'a> {
665 Box::pin(async move {
666 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
667 })
668 }
669
670 fn observe_parse<'a>(
672 &'a self,
673 future: ClarificationParseFuture<'a>,
674 ) -> ClarificationParseFuture<'a> {
675 Box::pin(async move {
676 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
677 })
678 }
679
680 fn observe_confirmation_parse<'a>(
682 &'a self,
683 future: ConfirmationParseFuture<'a>,
684 ) -> ConfirmationParseFuture<'a> {
685 Box::pin(async move {
686 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
687 })
688 }
689}
690
691struct ObservabilityProcessStageObserver;
692
693impl ProcessStageObserver for ObservabilityProcessStageObserver {
694 fn observe<'a>(
696 &'a self,
697 hint: ProcessPurposeHint,
698 future: ProcessStageFuture<'a>,
699 ) -> ProcessStageFuture<'a> {
700 Box::pin(async move {
701 with_observation_purpose(observation_purpose_for_process(hint), future).await
702 })
703 }
704}
705
706struct RegistryLLMGetter {
707 registry: Arc<LLMRegistry>,
708}
709
710impl LLMGetter for RegistryLLMGetter {
711 fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
712 self.registry.get(alias).ok()
713 }
714}
715
716impl RuntimeAgent {
717 #[allow(clippy::too_many_arguments)]
719 pub fn new(
720 info: AgentInfo,
721 llm_registry: Arc<LLMRegistry>,
722 memory: Arc<dyn Memory>,
723 tools: Arc<ToolRegistry>,
724 skills: Vec<SkillDefinition>,
725 system_prompt: String,
726 max_iterations: u32,
727 ) -> Self {
728 let (skill_router, skill_executor) = if !skills.is_empty() {
729 let router_llm = llm_registry.router().ok();
730 let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
731 let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
732 (router, Some(executor))
733 } else {
734 (None, None)
735 };
736
737 let context_manager =
738 ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
739
740 Self {
741 info,
742 llm_registry,
743 memory,
744 tools,
745 skills,
746 skill_router,
747 skill_executor,
748 base_system_prompt: system_prompt,
749 max_iterations,
750 iteration_count: RwLock::new(0),
751 max_context_tokens: 128000,
752 memory_token_budget: None,
753 recovery_manager: RecoveryManager::default(),
754 tool_security: ToolSecurityEngine::default(),
755 process_processor: None,
756 message_filters: RwLock::new(HashMap::new()),
757 state_machine: None,
758 transition_evaluator: None,
759 context_manager: Arc::new(context_manager),
760 template_renderer: TemplateRenderer::new(),
761 tool_call_history: RwLock::new(Vec::new()),
762 parallel_tools: ParallelToolsConfig::default(),
763 streaming: StreamingConfig::default(),
764 hooks: Arc::new(NoopHooks),
765 hitl_engine: None,
766 approval_handler: Arc::new(RejectAllHandler::new()),
767 storage_config: StorageConfig::default(),
768 storage: RwLock::new(None),
769 storage_init: tokio::sync::Mutex::new(()),
770 reasoning_config: ReasoningConfig::default(),
771 reflection_config: ReflectionConfig::default(),
772 disambiguation_manager: None,
773 disambiguation_epoch: AtomicU64::new(0),
774 disambiguation_admission: tokio::sync::RwLock::new(()),
775 state_transition_reserved: AtomicBool::new(false),
776 persona_manager: None,
777 pending_skill_id: RwLock::new(None),
778 current_plan: RwLock::new(None),
779 declared_tool_ids: None,
780 context_initialized: AtomicBool::new(false),
781 spawner: None,
782 spawner_registry: None,
783 redispatch_depth: RwLock::new(0),
784 active_turn_context: RwLock::new(None),
785 root_user_message_committed: AtomicBool::new(false),
786 actor_id: RwLock::new(None),
787 fact_store: RwLock::new(None),
788 fact_extractor: RwLock::new(None),
789 actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
790 messages_since_extraction: Arc::new(RwLock::new(0)),
791 actor_memory_config: None,
792 facts_config: None,
793 session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
794 current_session_id: RwLock::new(None),
795 relationship_manager: None,
796 observability_manager: None,
797 runtime_config: RuntimeConfig::default(),
798 background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
799 resource_locks: new_tool_resource_locks(),
800 runtime_control: Arc::new(RuntimeControlState::default()),
801 root_turn_gate: Arc::new(tokio::sync::Mutex::new(())),
802 }
803 }
804
805 pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
806 self.declared_tool_ids = ids;
807 self
808 }
809
810 pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
811 self.storage_config = config;
812 self
813 }
814
815 pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
816 *self.storage.write() = Some(storage);
817 self
818 }
819
820 pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
821 self.resource_locks = locks;
822 self
823 }
824
825 pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
826 self.reasoning_config = config;
827 self
828 }
829
830 pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
831 self.reflection_config = config;
832 self
833 }
834
835 pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
837 self.relationship_manager = Some(manager);
838 self
839 }
840
841 pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
843 self.observability_manager = Some(manager);
844 self
845 }
846
847 pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
849 let max_tasks = config.optimization.post_turn.max_background_tasks;
850 self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
851 self.runtime_config = config;
852 self
853 }
854
855 pub fn runtime_config(&self) -> &RuntimeConfig {
857 &self.runtime_config
858 }
859
860 pub async fn flush_background_tasks(&self) -> Result<()> {
862 self.background_maintenance.flush_all().await
863 }
864
865 pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
867 self.background_maintenance.flush_scope(actor_id).await
868 }
869
870 pub async fn flush_background_tasks_for_purpose(
872 &self,
873 purpose: RuntimeTaskPurpose,
874 ) -> Result<()> {
875 self.background_maintenance.flush_purpose(purpose).await
876 }
877
878 pub async fn flush_background_tasks_for_actor_purpose(
880 &self,
881 actor_id: &str,
882 purpose: RuntimeTaskPurpose,
883 ) -> Result<()> {
884 self.background_maintenance
885 .flush_scope_purpose(actor_id, purpose)
886 .await
887 }
888
889 pub async fn shutdown_background_tasks(&self) -> Result<()> {
891 self.flush_background_tasks().await
892 }
893
894 pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
896 self.observability_manager.clone()
897 }
898
899 async fn export_observability_if_configured(&self) {
901 let Some(manager) = self.observability_manager.as_ref() else {
902 return;
903 };
904 let export = &manager.config().export;
905 if !export.write_report && !export.write_raw_events {
906 return;
907 }
908 if let Err(error) = manager.export().await {
909 warn!(error = %error, "Observability export failed");
910 }
911 }
912
913 pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
915 self.relationship_manager.clone()
916 }
917
918 fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
919 current_turn_actor_context()
920 }
921
922 fn effective_actor_id(&self) -> Option<String> {
923 self.current_turn_actor_context()
924 .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
925 .or_else(|| self.actor_id.read().clone())
926 }
927
928 fn effective_origin_actor_id(&self) -> Option<String> {
929 self.current_turn_actor_context()
930 .and_then(|ctx| ctx.origin_actor_id.clone())
931 .or_else(|| self.actor_id.read().clone())
932 }
933
934 fn record_session_actor_if_needed(&self) {
935 if let Some(actor_id) = self.effective_origin_actor_id() {
936 let mut meta = self.session_metadata.write();
937 meta.actor_id = Some(actor_id.clone());
938 if !meta.actors.iter().any(|a| a == &actor_id) {
939 meta.actors.push(actor_id);
940 }
941 }
942 }
943
944 fn outbound_actor_context(&self) -> crate::TurnActorContext {
945 let mut context = self.current_turn_actor_context().unwrap_or_default();
946 if context.origin_actor_id.is_none() {
947 context.origin_actor_id = self.effective_origin_actor_id();
948 }
949 context.sender_agent_id = Some(self.info.id.clone());
950 context
951 }
952
953 fn observation_session_id(&self) -> Option<String> {
955 let mut current = self.current_session_id.write();
956 if current.is_none() {
957 *current = Some(new_observation_session_id());
958 }
959 current.clone()
960 }
961
962 fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
964 let manager = self.observability_manager.as_ref()?;
965 let context = self.build_context_with_overlays();
966 let language = resolve_language_from_context(manager.config(), &context);
967 let context = current_observation_context()
968 .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
969 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
970 Some(
971 context
972 .with_actor(actor_id.or_else(|| self.effective_actor_id()))
973 .with_session(self.observation_session_id())
974 .with_state(self.current_state())
975 .with_language(Some(language)),
976 )
977 }
978
979 fn current_runtime_observation_context(
981 &self,
982 purpose: ObservationPurpose,
983 ) -> Option<SpanContext> {
984 let manager = self.observability_manager.as_ref()?;
985 let context = self.build_context_with_overlays();
986 let language = resolve_language_from_context(manager.config(), &context);
987 let mut observation = current_observation_context()
988 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
989 observation.agent_id = self.info.id.clone();
990 observation.actor_id = self.effective_actor_id();
991 observation.session_id = self.observation_session_id();
992 observation.state = self.current_state();
993 observation.language = Some(language);
994 observation.purpose = purpose;
995 Some(observation)
996 }
997
998 async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
1000 where
1001 F: Future<Output = T>,
1002 {
1003 if let Some(context) = self.current_runtime_observation_context(purpose) {
1004 with_observation_context(context, future).await
1005 } else {
1006 future.await
1007 }
1008 }
1009
1010 fn chat_with_actor_context_boxed<'a>(
1014 &'a self,
1015 input: &'a str,
1016 actor_context: crate::TurnActorContext,
1017 ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
1018 Box::pin(async move {
1019 let RootTurnAdmission {
1020 guard,
1021 identity_stack,
1022 } = self.acquire_root_turn().await?;
1023 let result = scope_runtime_gate_identity_stack(&identity_stack, async move {
1024 let actor_id = actor_context.effective_actor_id().map(str::to_string);
1025 let run = async move {
1026 scope_actor_context(
1027 actor_context,
1028 Box::pin(async move { self.run_loop(input).await }),
1029 )
1030 .await
1031 };
1032 let result = if let Some(context) = self.build_observation_context(actor_id) {
1033 with_observation_context(context, run).await
1034 } else {
1035 run.await
1036 };
1037 self.export_observability_if_configured().await;
1038 result
1039 })
1040 .await;
1041 drop(guard);
1042 result
1043 })
1044 }
1045
1046 async fn acquire_root_turn(&self) -> Result<RootTurnAdmission> {
1048 let gate_identity = Arc::clone(&self.root_turn_gate);
1049 let current_identity_stack = current_runtime_gate_identity_stack();
1050 if current_identity_stack
1051 .iter()
1052 .any(|owned_gate| Arc::ptr_eq(owned_gate, &gate_identity))
1053 {
1054 return Err(AgentError::Other(format!(
1055 "RuntimeAgent '{}' rejected reentrant root turn ownership",
1056 self.info.id
1057 )));
1058 }
1059 let guard = Arc::clone(&gate_identity).lock_owned().await;
1060 let mut identity_stack = Vec::with_capacity(current_identity_stack.len() + 1);
1064 identity_stack.extend(current_identity_stack.iter().cloned());
1065 identity_stack.push(gate_identity);
1066 Ok(RootTurnAdmission {
1067 guard,
1068 identity_stack: identity_stack.into(),
1069 })
1070 }
1071
1072 pub async fn chat_with_actor_context(
1076 &self,
1077 input: &str,
1078 actor_context: crate::TurnActorContext,
1079 ) -> Result<AgentResponse> {
1080 self.chat_with_actor_context_boxed(input, actor_context)
1081 .await
1082 }
1083
1084 pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
1086 let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
1087 self.chat_with_actor_context(input, actor_context).await
1088 }
1089
1090 pub async fn load_actor_relationship(&self) -> Result<()> {
1092 self.maybe_load_actor_relationship().await;
1093 Ok(())
1094 }
1095
1096 pub async fn update_relationship_dimension(
1098 &self,
1099 dimension: &str,
1100 delta: f64,
1101 reason: Option<&str>,
1102 ) -> Result<ai_agents_relationships::DimensionChange> {
1103 self.update_relationship_dimension_for_perspective(
1104 ai_agents_relationships::RelationshipPerspective::AgentToActor,
1105 dimension,
1106 delta,
1107 reason,
1108 )
1109 .await
1110 }
1111
1112 pub async fn update_relationship_dimension_for_perspective(
1116 &self,
1117 perspective: ai_agents_relationships::RelationshipPerspective,
1118 dimension: &str,
1119 delta: f64,
1120 reason: Option<&str>,
1121 ) -> Result<ai_agents_relationships::DimensionChange> {
1122 let manager = self
1123 .relationship_manager
1124 .as_ref()
1125 .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
1126 let actor_id = self.effective_actor_id().ok_or_else(|| {
1127 AgentError::Config("No actor ID set. Use set_actor_id() first".into())
1128 })?;
1129 let change = manager.update_dimension_for_perspective(
1130 &actor_id,
1131 perspective,
1132 dimension,
1133 delta,
1134 1.0,
1135 reason.unwrap_or("manual relationship update"),
1136 )?;
1137 self.persist_actor_relationship(&actor_id).await?;
1138 info!(
1139 actor_id = %actor_id,
1140 perspective = %change.perspective,
1141 dimension = %change.dimension,
1142 delta = change.delta,
1143 current = change.current,
1144 "relationship updated manually"
1145 );
1146 self.hooks
1147 .on_relationship_change(&actor_id, std::slice::from_ref(&change))
1148 .await;
1149 Ok(change)
1150 }
1151
1152 pub fn reasoning_config(&self) -> &ReasoningConfig {
1153 &self.reasoning_config
1154 }
1155
1156 pub fn reflection_config(&self) -> &ReflectionConfig {
1157 &self.reflection_config
1158 }
1159
1160 pub fn with_facts_config(
1163 mut self,
1164 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1165 facts_config: Option<ai_agents_facts::FactsConfig>,
1166 ) -> Self {
1167 self.actor_memory_config = actor_memory_config;
1168 self.facts_config = facts_config;
1169 self
1170 }
1171
1172 pub fn with_facts(
1175 mut self,
1176 store: Arc<ai_agents_facts::FactStore>,
1177 extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1178 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1179 facts_config: Option<ai_agents_facts::FactsConfig>,
1180 ) -> Self {
1181 *self.fact_store.write() = Some(store);
1182 *self.fact_extractor.write() = extractor;
1183 self.actor_memory_config = actor_memory_config;
1184 self.facts_config = facts_config;
1185 self
1186 }
1187
1188 pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1190 self.fact_store.read().clone()
1191 }
1192
1193 pub fn actor_id(&self) -> Option<String> {
1195 self.actor_id.read().clone()
1196 }
1197
1198 pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1200 *self.actor_id.write() = Some(actor_id.to_string());
1201 {
1202 let mut meta = self.session_metadata.write();
1203 meta.actor_id = Some(actor_id.to_string());
1204 if !meta.actors.iter().any(|a| a == actor_id) {
1205 meta.actors.push(actor_id.to_string());
1206 }
1207 }
1208 Ok(())
1209 }
1210
1211 pub fn clear_actor_id(&self) {
1213 *self.actor_id.write() = None;
1214 self.session_metadata.write().actor_id = None;
1215 }
1216
1217 pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1219 self.set_actor_id(user_id)
1220 }
1221
1222 pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1224 let actor_id = match self.effective_actor_id() {
1225 Some(id) => id,
1226 None => return Ok(()),
1227 };
1228
1229 let store_opt = self.fact_store.read().clone();
1230 if let Some(store) = store_opt {
1231 let facts = store.get_facts(&actor_id).await?;
1232 let count = facts.len();
1233 self.actor_facts_cache
1234 .write()
1235 .insert(actor_id.clone(), facts);
1236 self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1237 tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1238 }
1239
1240 Ok(())
1241 }
1242
1243 async fn maybe_load_actor_memory(&self) {
1245 let Some(actor_id) = self.effective_actor_id() else {
1246 return;
1247 };
1248 if self.actor_facts_cache.read().contains_key(&actor_id) {
1249 return;
1250 }
1251 let _ = self.load_actor_memory().await;
1252 }
1253
1254 async fn pre_turn_session_lifecycle(&self) {
1256 if *self.redispatch_depth.read() > 0 {
1257 return;
1258 }
1259 self.resolve_actor_id_from_context();
1260 self.await_background_before_next_turn().await;
1261 self.record_session_actor_if_needed();
1262 self.maybe_load_actor_memory().await;
1263 self.maybe_load_actor_relationship().await;
1264 *self.messages_since_extraction.write() += 1;
1265 }
1266
1267 async fn post_turn_session_lifecycle(&self) -> Result<()> {
1269 if *self.redispatch_depth.read() > 0 {
1270 return Ok(());
1271 }
1272 *self.messages_since_extraction.write() += 1;
1273 self.run_post_turn_maintenance().await
1274 }
1275
1276 fn begin_root_turn(&self) {
1278 if *self.redispatch_depth.read() == 0 {
1279 let mut guard = self.active_turn_context.write();
1280 if guard.is_none() {
1281 self.root_user_message_committed
1282 .store(false, Ordering::SeqCst);
1283 let max_calls = self
1284 .runtime_config
1285 .optimization
1286 .max_speculative_llm_calls_per_turn;
1287 *guard = Some(TurnOptimizationContext::new(
1288 String::new(),
1289 HashMap::new(),
1290 max_calls,
1291 ));
1292 }
1293 }
1294 }
1295
1296 fn update_active_turn_context(
1297 &self,
1298 processed_input: &str,
1299 input_context: HashMap<String, Value>,
1300 ) {
1301 if *self.redispatch_depth.read() > 0 {
1302 return;
1303 }
1304 let max_calls = self
1305 .runtime_config
1306 .optimization
1307 .max_speculative_llm_calls_per_turn;
1308 let mut guard = self.active_turn_context.write();
1309 match guard.as_mut() {
1310 Some(context) => {
1311 context.processed_input = processed_input.to_string();
1312 context.input_context = input_context;
1313 context.max_speculative_llm_calls = max_calls;
1314 }
1315 None => {
1316 *guard = Some(TurnOptimizationContext::new(
1317 processed_input,
1318 input_context,
1319 max_calls,
1320 ));
1321 }
1322 }
1323 }
1324
1325 async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1327 if *self.redispatch_depth.read() > 0 {
1328 return Ok(());
1329 }
1330 if !self
1331 .root_user_message_committed
1332 .swap(true, Ordering::SeqCst)
1333 {
1334 self.memory
1335 .add_message(ChatMessage::user(processed_input))
1336 .await?;
1337 if let Some(context) = self.active_turn_context.write().as_mut() {
1338 context.mark_user_message_committed();
1339 }
1340 }
1341 Ok(())
1342 }
1343
1344 fn end_root_turn(&self) {
1346 if *self.redispatch_depth.read() == 0 {
1347 self.root_user_message_committed
1348 .store(false, Ordering::SeqCst);
1349 *self.active_turn_context.write() = None;
1350 }
1351 }
1352
1353 fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1354 self.begin_root_turn();
1355 let mut guard = self.active_turn_context.write();
1356 let Some(context) = guard.as_mut() else {
1357 return false;
1358 };
1359 context.reserve_speculative_llm_call_for(kind)
1360 }
1361
1362 fn branch_context_preview(&self) -> String {
1363 let context = self.build_context_with_overlays();
1364 let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1365 const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1366 if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1367 value = value
1368 .chars()
1369 .take(MAX_CONTEXT_PREVIEW_CHARS)
1370 .collect::<String>();
1371 value.push_str("...");
1372 }
1373 value
1374 }
1375
1376 async fn await_background_before_next_turn(&self) {
1378 let optimization = &self.runtime_config.optimization;
1379 if !optimization.enabled {
1380 return;
1381 }
1382 let actor_id = self.effective_actor_id();
1383 let post = &optimization.post_turn;
1384 self.await_background_task(
1385 post.facts.await_before_next_turn,
1386 RuntimeTaskPurpose::PostTurnFacts,
1387 actor_id.as_deref(),
1388 "facts",
1389 )
1390 .await;
1391 self.await_background_task(
1392 post.relationships.await_before_next_turn,
1393 RuntimeTaskPurpose::PostTurnRelationship,
1394 actor_id.as_deref(),
1395 "relationships",
1396 )
1397 .await;
1398 }
1399
1400 async fn await_background_task(
1401 &self,
1402 policy: AwaitBeforeNextTurn,
1403 purpose: RuntimeTaskPurpose,
1404 actor_id: Option<&str>,
1405 label: &str,
1406 ) {
1407 match policy {
1408 AwaitBeforeNextTurn::Never => {}
1409 AwaitBeforeNextTurn::Always => {
1410 if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1411 warn!(label = label, error = %error, "background maintenance flush failed");
1412 }
1413 }
1414 AwaitBeforeNextTurn::SameActor => {
1415 if let Some(actor_id) = actor_id
1416 && let Err(error) = self
1417 .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1418 .await
1419 {
1420 warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1421 }
1422 }
1423 }
1424 }
1425
1426 async fn run_post_turn_maintenance(&self) -> Result<()> {
1428 let optimization = &self.runtime_config.optimization;
1429 if !optimization.enabled {
1430 self.auto_extract_facts().await;
1431 self.auto_update_relationship().await;
1432 return Ok(());
1433 }
1434
1435 let facts_mode = effective_maintenance_mode(
1436 optimization.post_turn.facts.mode,
1437 optimization.parallel_post_turn_memory,
1438 );
1439 let relationships_mode = effective_maintenance_mode(
1440 optimization.post_turn.relationships.mode,
1441 optimization.parallel_post_turn_memory,
1442 );
1443
1444 match (facts_mode, relationships_mode) {
1445 (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1446 self.auto_extract_facts().await;
1447 self.auto_update_relationship().await;
1448 }
1449 (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1450 let facts = self.auto_extract_facts();
1451 let relationships = self.auto_update_relationship();
1452 tokio::join!(facts, relationships);
1453 }
1454 (MaintenanceMode::Background, MaintenanceMode::Background) => {
1455 self.schedule_facts_background().await?;
1456 self.schedule_relationship_background().await?;
1457 }
1458 (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1459 | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1460 self.schedule_facts_background().await?;
1461 self.auto_update_relationship().await;
1462 }
1463 (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1464 | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1465 self.auto_extract_facts().await;
1466 self.schedule_relationship_background().await?;
1467 }
1468 _ => {
1469 self.auto_extract_facts().await;
1470 self.auto_update_relationship().await;
1471 }
1472 }
1473 Ok(())
1474 }
1475
1476 async fn schedule_facts_background(&self) -> Result<()> {
1477 let policy = self.runtime_config.optimization.post_turn.facts.clone();
1478 let should_extract = self
1479 .facts_config
1480 .as_ref()
1481 .map(|c| c.enabled && c.auto_extract)
1482 .unwrap_or(false);
1483 if !should_extract {
1484 return Ok(());
1485 }
1486 let msgs_since = *self.messages_since_extraction.read();
1487 if msgs_since < 2 {
1488 return Ok(());
1489 }
1490 let Some(actor_id) = self.effective_actor_id() else {
1491 self.record_skipped_maintenance(
1492 "facts",
1493 ObservationPurpose::FactsExtraction,
1494 "missing_actor",
1495 Some(&policy),
1496 );
1497 return Ok(());
1498 };
1499 let Some(extractor) = self.fact_extractor.read().clone() else {
1500 return Ok(());
1501 };
1502 let messages = match self.memory.get_messages(None).await {
1503 Ok(messages) => messages,
1504 Err(error) => {
1505 warn!(error = %error, "failed to snapshot messages for fact extraction");
1506 return Ok(());
1507 }
1508 };
1509 let recent: Vec<_> = messages
1510 .iter()
1511 .rev()
1512 .take(msgs_since)
1513 .rev()
1514 .cloned()
1515 .collect();
1516 if recent.is_empty() {
1517 return Ok(());
1518 }
1519 let existing = self
1520 .actor_facts_cache
1521 .read()
1522 .get(&actor_id)
1523 .cloned()
1524 .unwrap_or_default();
1525 let categories = self
1526 .facts_config
1527 .as_ref()
1528 .map(|c| c.custom_categories.clone())
1529 .unwrap_or_default();
1530 let store = self.fact_store.read().clone();
1531 let cache = Arc::clone(&self.actor_facts_cache);
1532 let counter = Arc::clone(&self.messages_since_extraction);
1533 let hooks = Arc::clone(&self.hooks);
1534 let agent_id = self.info.id.clone();
1535 let observation = current_observation_context();
1536 let key = MaintenanceSequenceKey::actor(
1537 agent_id,
1538 actor_id.clone(),
1539 RuntimeTaskPurpose::PostTurnFacts,
1540 );
1541 let actor_for_task = actor_id.clone();
1542 let task = async move {
1543 let run = async move {
1544 let facts = extractor
1545 .extract(&recent, &existing, Some(&actor_for_task), &categories)
1546 .await?;
1547 if !facts.is_empty() {
1548 if let Some(store) = store {
1549 let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1550 cache.write().insert(actor_for_task.clone(), authoritative);
1551 } else {
1552 cache
1553 .write()
1554 .entry(actor_for_task.clone())
1555 .or_default()
1556 .extend(facts.clone());
1557 }
1558 {
1559 let mut count = counter.write();
1560 if *count <= msgs_since {
1561 *count = 0;
1562 } else {
1563 *count -= msgs_since;
1564 }
1565 }
1566 hooks.on_facts_extracted(&actor_for_task, &facts).await;
1567 }
1568 Ok(())
1569 };
1570 if let Some(context) = observation {
1571 with_observation_context(
1572 context.with_purpose(ObservationPurpose::FactsExtraction),
1573 run,
1574 )
1575 .await
1576 } else {
1577 run.await
1578 }
1579 };
1580 self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1581 .await
1582 }
1583
1584 async fn schedule_relationship_background(&self) -> Result<()> {
1585 let policy = self
1586 .runtime_config
1587 .optimization
1588 .post_turn
1589 .relationships
1590 .clone();
1591 let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1592 return Ok(());
1593 };
1594 let Some(actor_id) = self.effective_actor_id() else {
1595 self.record_skipped_maintenance(
1596 "relationships",
1597 ObservationPurpose::RelationshipUpdate,
1598 "missing_actor",
1599 Some(&policy),
1600 );
1601 return Ok(());
1602 };
1603 let recent_messages = manager.config().auto_update.recent_messages;
1604 let messages = match self.memory.get_messages(Some(recent_messages)).await {
1605 Ok(messages) => messages,
1606 Err(error) => {
1607 warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1608 return Ok(());
1609 }
1610 };
1611 let storage = self.storage.read().clone();
1612 let hooks = Arc::clone(&self.hooks);
1613 let agent_id = self.info.id.clone();
1614 let observation = current_observation_context();
1615 let key = MaintenanceSequenceKey::actor(
1616 agent_id.clone(),
1617 actor_id.clone(),
1618 RuntimeTaskPurpose::PostTurnRelationship,
1619 );
1620 let actor_for_task = actor_id.clone();
1621 let task = async move {
1622 let run = async move {
1623 if manager.config().auto_update.enabled {
1624 let update = manager.auto_update(&actor_for_task, &messages).await?;
1625 if !update.changes.is_empty() {
1626 hooks
1627 .on_relationship_change(&actor_for_task, &update.changes)
1628 .await;
1629 }
1630 if let Some(ref event) = update.event {
1631 hooks.on_notable_event(&actor_for_task, event).await;
1632 }
1633 }
1634 if manager.config().persistence.enabled
1635 && let (Some(storage), Some(value)) =
1636 (storage, manager.relationship_as_value(&actor_for_task)?)
1637 {
1638 storage
1639 .save_relationship(&agent_id, &actor_for_task, &value)
1640 .await?;
1641 }
1642 Ok(())
1643 };
1644 if let Some(context) = observation {
1645 with_observation_context(
1646 context.with_purpose(ObservationPurpose::RelationshipUpdate),
1647 run,
1648 )
1649 .await
1650 } else {
1651 run.await
1652 }
1653 };
1654 self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1655 .await
1656 }
1657
1658 async fn spawn_or_handle_background<F>(
1660 &self,
1661 key: Option<MaintenanceSequenceKey>,
1662 task: F,
1663 label: &'static str,
1664 policy: &crate::optimization::config::MaintenanceTaskPolicy,
1665 ) -> Result<()>
1666 where
1667 F: Future<Output = Result<()>> + Send + 'static,
1668 {
1669 if self.background_maintenance.is_full() {
1670 match self
1671 .runtime_config
1672 .optimization
1673 .post_turn
1674 .on_background_overflow
1675 {
1676 BackgroundOverflowPolicy::RunInline => {
1677 record_background_maintenance_event(
1678 self.observability_manager.as_ref(),
1679 label,
1680 EventStatus::Success,
1681 0,
1682 "inline_overflow",
1683 None,
1684 Some(policy),
1685 );
1686 let start = Instant::now();
1687 match task.await {
1688 Ok(()) => record_background_maintenance_event(
1689 self.observability_manager.as_ref(),
1690 label,
1691 EventStatus::Success,
1692 start.elapsed().as_millis() as u64,
1693 "inline_completed",
1694 None,
1695 Some(policy),
1696 ),
1697 Err(error) => {
1698 warn!(label = label, error = %error, "inline maintenance fallback failed");
1699 record_background_maintenance_event(
1700 self.observability_manager.as_ref(),
1701 label,
1702 EventStatus::Error,
1703 start.elapsed().as_millis() as u64,
1704 "inline_failed",
1705 Some(error.to_string()),
1706 Some(policy),
1707 );
1708 return Err(error);
1709 }
1710 }
1711 }
1712 BackgroundOverflowPolicy::Drop => {
1713 self.record_skipped_maintenance(
1714 label,
1715 ObservationPurpose::Other(label.to_string()),
1716 "queue_full",
1717 Some(policy),
1718 );
1719 }
1720 BackgroundOverflowPolicy::Error => {
1721 record_background_maintenance_event(
1722 self.observability_manager.as_ref(),
1723 label,
1724 EventStatus::Error,
1725 0,
1726 "queue_full",
1727 None,
1728 Some(policy),
1729 );
1730 warn!(label = label, "background maintenance queue full");
1731 return Err(AgentError::Other(format!(
1732 "background maintenance queue is full for {}",
1733 label
1734 )));
1735 }
1736 }
1737 return Ok(());
1738 }
1739
1740 record_background_maintenance_event(
1741 self.observability_manager.as_ref(),
1742 label,
1743 EventStatus::Success,
1744 0,
1745 "scheduled",
1746 None,
1747 Some(policy),
1748 );
1749 let manager = self.observability_manager.clone();
1750 let policy_for_task = policy.clone();
1751 let observed_task = async move {
1752 let start = Instant::now();
1753 let result = task.await;
1754 match &result {
1755 Ok(()) => record_background_maintenance_event(
1756 manager.as_ref(),
1757 label,
1758 EventStatus::Success,
1759 start.elapsed().as_millis() as u64,
1760 "completed",
1761 None,
1762 Some(&policy_for_task),
1763 ),
1764 Err(error) => record_background_maintenance_event(
1765 manager.as_ref(),
1766 label,
1767 EventStatus::Error,
1768 start.elapsed().as_millis() as u64,
1769 "failed",
1770 Some(error.to_string()),
1771 Some(&policy_for_task),
1772 ),
1773 }
1774 result
1775 };
1776
1777 if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1778 record_background_maintenance_event(
1779 self.observability_manager.as_ref(),
1780 label,
1781 EventStatus::Error,
1782 0,
1783 "spawn_failed",
1784 Some(error.to_string()),
1785 Some(policy),
1786 );
1787 warn!(label = label, error = %error, "background maintenance spawn failed");
1788 return Err(error);
1789 }
1790 Ok(())
1791 }
1792
1793 fn record_skipped_maintenance(
1795 &self,
1796 label: &str,
1797 purpose: ObservationPurpose,
1798 reason: &str,
1799 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1800 ) {
1801 if let Some(manager) = self.observability_manager.as_ref() {
1802 let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1803 tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1804 manager.record_lifecycle_event(
1805 EventType::MemoryOperation {
1806 operation: format!("{}_maintenance", label),
1807 },
1808 purpose,
1809 EventStatus::Skipped,
1810 0,
1811 tags,
1812 None,
1813 );
1814 }
1815 }
1816
1817 pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1819 let Some(actor_id) = self.effective_actor_id() else {
1820 return Vec::new();
1821 };
1822 self.actor_facts_cache
1823 .read()
1824 .get(&actor_id)
1825 .cloned()
1826 .unwrap_or_default()
1827 }
1828
1829 pub fn relationship_memory_text(&self) -> Option<String> {
1831 self.format_relationship_for_context().map(|(_, text)| text)
1832 }
1833
1834 pub async fn extract_facts(
1836 &self,
1837 last_n: usize,
1838 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1839 self.extract_facts_with_source(last_n, "manual").await
1840 }
1841
1842 async fn extract_facts_with_source(
1843 &self,
1844 last_n: usize,
1845 source: &'static str,
1846 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1847 let extractor = match self.fact_extractor.read().clone() {
1848 Some(e) => e,
1849 None => return Ok(vec![]),
1850 };
1851
1852 let messages = self.memory.get_messages(None).await?;
1853 let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1854
1855 if recent.is_empty() {
1856 return Ok(vec![]);
1857 }
1858
1859 let actor_id = self.effective_actor_id();
1860 let existing = actor_id
1861 .as_ref()
1862 .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1863 .unwrap_or_default();
1864
1865 let categories = self
1866 .facts_config
1867 .as_ref()
1868 .map(|c| c.custom_categories.clone())
1869 .unwrap_or_default();
1870
1871 let facts = self
1872 .observe_purpose(
1873 ObservationPurpose::FactsExtraction,
1874 extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1875 )
1876 .await?;
1877
1878 if !facts.is_empty() {
1880 let fact_store_opt = self.fact_store.read().clone();
1881 let mut stored_total = 0usize;
1882 let mut cache_updated = false;
1883 if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1884 let authoritative = store.add_facts(aid, facts.clone()).await?;
1886 stored_total = authoritative.len();
1887 self.actor_facts_cache
1888 .write()
1889 .insert(aid.clone(), authoritative);
1890 cache_updated = true;
1891 } else if let Some(aid) = &actor_id {
1892 let mut cache = self.actor_facts_cache.write();
1893 let entry = cache.entry(aid.clone()).or_default();
1894 entry.extend(facts.clone());
1895 stored_total = entry.len();
1896 cache_updated = true;
1897 }
1898
1899 info!(
1900 actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1901 source = source,
1902 requested_messages = last_n,
1903 message_count = recent.len(),
1904 extracted_count = facts.len(),
1905 cache_updated = cache_updated,
1906 stored_total = stored_total,
1907 "facts extracted"
1908 );
1909
1910 if let Some(ref aid) = actor_id {
1911 self.hooks.on_facts_extracted(aid, &facts).await;
1912 }
1913 }
1914
1915 Ok(facts)
1916 }
1917
1918 fn resolve_actor_id_from_context(&self) {
1921 if self
1922 .current_turn_actor_context()
1923 .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1924 .is_some()
1925 {
1926 return;
1927 }
1928
1929 if let Some(ref am_config) = self.actor_memory_config
1930 && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1931 && let Some(ref path) = am_config.identification.context_path
1932 {
1933 let val = self
1935 .context_manager
1936 .get_path(path)
1937 .or_else(|| self.context_manager.get(path));
1938 if let Some(val) = val
1939 && let Some(id_str) = val.as_str()
1940 {
1941 let current = self.actor_id.read().clone();
1942 if current.as_deref() != Some(id_str) {
1943 *self.actor_id.write() = Some(id_str.to_string());
1944 let mut meta = self.session_metadata.write();
1945 meta.actor_id = Some(id_str.to_string());
1946 if !meta.actors.iter().any(|a| a == id_str) {
1947 meta.actors.push(id_str.to_string());
1948 }
1949 }
1950 }
1951 }
1952 }
1953
1954 fn format_actor_facts_for_context(&self) -> String {
1956 let should_inject = self
1958 .facts_config
1959 .as_ref()
1960 .map(|c| c.inject_in_context)
1961 .unwrap_or(true);
1962 if !should_inject {
1963 return String::new();
1964 }
1965
1966 let Some(actor_id) = self.effective_actor_id() else {
1967 return String::new();
1968 };
1969
1970 let facts = self
1971 .actor_facts_cache
1972 .read()
1973 .get(&actor_id)
1974 .cloned()
1975 .unwrap_or_default();
1976 if facts.is_empty() {
1977 return String::new();
1978 }
1979
1980 let am_config = self.actor_memory_config.as_ref();
1981 let facts_budget = self
1984 .memory_token_budget
1985 .as_ref()
1986 .map(|b| b.allocation.facts as usize)
1987 .filter(|n| *n > 0);
1988 let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1989 let max_tokens = facts_budget.unwrap_or(default_max);
1990
1991 let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1993 if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1994 return String::new();
1995 }
1996 if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1997 && !cfg.injection.categories.is_empty()
1998 {
1999 facts
2000 .iter()
2001 .filter(|f| {
2002 cfg.injection
2003 .categories
2004 .iter()
2005 .any(|c| f.category.to_string() == *c)
2006 })
2007 .cloned()
2008 .collect()
2009 } else {
2010 facts.clone()
2011 }
2012 } else {
2013 facts.clone()
2014 };
2015
2016 if filtered.is_empty() {
2017 return String::new();
2018 }
2019
2020 if let Some(store) = self.fact_store.read().clone() {
2021 store.format_for_context(&filtered, max_tokens)
2022 } else {
2023 String::new()
2024 }
2025 }
2026
2027 fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
2028 let context = self.build_context_with_overlays();
2029 let mut root = Value::Object(context.into_iter().collect());
2030 for (path, value) in staged {
2031 if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
2032 root = updated;
2033 }
2034 }
2035 match root {
2036 Value::Object(obj) => obj.into_iter().collect(),
2037 _ => HashMap::new(),
2038 }
2039 }
2040
2041 fn build_context_with_overlays(&self) -> HashMap<String, Value> {
2042 let mut context = self.context_manager.get_all();
2043 let mut root = Value::Object(context.clone().into_iter().collect());
2044
2045 if let Some(turn_ctx) = self.current_turn_actor_context() {
2046 if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
2047 && let Ok(updated) = ai_agents_core::set_dot_path(
2048 root.clone(),
2049 "interaction.origin_actor_id",
2050 serde_json::json!(origin_actor_id),
2051 )
2052 {
2053 root = updated;
2054 }
2055 if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
2056 && let Ok(updated) = ai_agents_core::set_dot_path(
2057 root.clone(),
2058 "interaction.sender_agent_id",
2059 serde_json::json!(sender_agent_id),
2060 )
2061 {
2062 root = updated;
2063 }
2064 }
2065
2066 if let Some(ref actor_id) = self.effective_actor_id()
2067 && let Ok(updated) = ai_agents_core::set_dot_path(
2068 root.clone(),
2069 "interaction.actor_id",
2070 serde_json::json!(actor_id),
2071 )
2072 {
2073 root = updated;
2074 }
2075
2076 if let Some(manager) = self.relationship_manager.as_ref()
2077 && let Some(actor_id) = self.effective_actor_id()
2078 && let Some(value) = manager.to_context_value(&actor_id)
2079 && let Ok(updated) = ai_agents_core::set_dot_path(
2080 root.clone(),
2081 &manager.config().injection.context_path,
2082 value,
2083 )
2084 {
2085 root = updated;
2086 }
2087
2088 if let Value::Object(obj) = root {
2089 context = obj.into_iter().collect();
2090 }
2091
2092 context
2093 }
2094
2095 fn resolve_actor_name_from_context(&self) -> Option<String> {
2096 for path in ["actor.name", "user.name", "player.name", "customer.name"] {
2097 if let Some(value) = self.context_manager.get_path(path)
2098 && let Some(name) = value.as_str()
2099 {
2100 return Some(name.to_string());
2101 }
2102 }
2103 None
2104 }
2105
2106 async fn maybe_load_actor_relationship(&self) {
2107 let Some(manager) = self.relationship_manager.as_ref() else {
2108 return;
2109 };
2110 let Some(actor_id) = self.effective_actor_id() else {
2111 return;
2112 };
2113
2114 let mut should_fire_loaded = false;
2115 if manager.get(&actor_id).is_none() {
2116 let mut loaded = false;
2117 if manager.config().persistence.enabled {
2118 let storage = self.storage.read().clone();
2119 if let Some(storage) = storage {
2120 match storage.load_relationship(&self.info.id, &actor_id).await {
2121 Ok(Some(value)) => match manager.insert_from_value(value) {
2122 Ok(_) => loaded = true,
2123 Err(e) => {
2124 warn!(actor = %actor_id, error = %e, "failed to restore relationship")
2125 }
2126 },
2127 Ok(None) => {}
2128 Err(e) => {
2129 warn!(actor = %actor_id, error = %e, "failed to load relationship")
2130 }
2131 }
2132 }
2133 }
2134
2135 if !loaded {
2136 manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
2137 }
2138 should_fire_loaded = true;
2139 }
2140
2141 let actor_name = self.resolve_actor_name_from_context();
2142 let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
2143 if should_fire_loaded {
2144 self.hooks
2145 .on_relationship_loaded(&actor_id, &relationship)
2146 .await;
2147 }
2148 }
2149
2150 fn format_relationship_for_context(&self) -> Option<(String, String)> {
2151 let manager = self.relationship_manager.as_ref()?;
2152 if !manager.config().injection.enabled {
2153 return None;
2154 }
2155 let actor_id = self.effective_actor_id()?;
2156 let relationship = manager.get(&actor_id)?;
2157 let local_cap = manager.config().injection.max_tokens;
2158 let global_cap = self
2159 .memory_token_budget
2160 .as_ref()
2161 .map(|b| b.allocation.relationships as usize)
2162 .filter(|n| *n > 0);
2163 let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
2164 let text = ai_agents_relationships::format_relationship(
2165 &relationship,
2166 &manager.config().injection.format,
2167 max_tokens,
2168 );
2169 if text.is_empty() {
2170 None
2171 } else {
2172 Some((manager.config().injection.prompt_variable.clone(), text))
2173 }
2174 }
2175
2176 async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2177 let Some(manager) = self.relationship_manager.as_ref() else {
2178 return Ok(());
2179 };
2180 if !manager.config().persistence.enabled {
2181 return Ok(());
2182 }
2183 let storage = self.storage.read().clone();
2184 let Some(storage) = storage else {
2185 return Ok(());
2186 };
2187 if let Some(value) = manager.relationship_as_value(actor_id)? {
2188 storage
2189 .save_relationship(&self.info.id, actor_id, &value)
2190 .await?;
2191 }
2192 Ok(())
2193 }
2194
2195 async fn auto_update_relationship(&self) {
2196 let Some(manager) = self.relationship_manager.as_ref() else {
2197 return;
2198 };
2199 let Some(actor_id) = self.effective_actor_id() else {
2200 return;
2201 };
2202 if !manager.config().auto_update.enabled {
2203 let _ = self.persist_actor_relationship(&actor_id).await;
2204 return;
2205 }
2206
2207 let recent_messages = manager.config().auto_update.recent_messages;
2208 let messages = match self.memory.get_messages(Some(recent_messages)).await {
2209 Ok(messages) => messages,
2210 Err(e) => {
2211 warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2212 return;
2213 }
2214 };
2215
2216 match self
2217 .observe_purpose(
2218 ObservationPurpose::RelationshipUpdate,
2219 manager.auto_update(&actor_id, &messages),
2220 )
2221 .await
2222 {
2223 Ok(update) => {
2224 if !update.changes.is_empty() {
2225 self.hooks
2226 .on_relationship_change(&actor_id, &update.changes)
2227 .await;
2228 }
2229 if let Some(ref event) = update.event {
2230 self.hooks.on_notable_event(&actor_id, event).await;
2231 }
2232 let persisted = match self.persist_actor_relationship(&actor_id).await {
2233 Ok(()) => true,
2234 Err(e) => {
2235 warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2236 false
2237 }
2238 };
2239 if !update.changes.is_empty() || update.event.is_some() {
2240 let changed_dimensions: Vec<String> = update
2241 .changes
2242 .iter()
2243 .map(|change| format!("{}:{}", change.perspective, change.dimension))
2244 .collect();
2245 info!(
2246 actor_id = %actor_id,
2247 change_count = update.changes.len(),
2248 changed_dimensions = ?changed_dimensions,
2249 event_present = update.event.is_some(),
2250 persisted = persisted,
2251 "relationship updated"
2252 );
2253 } else {
2254 debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2255 }
2256 }
2257 Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2258 }
2259 }
2260
2261 async fn auto_extract_facts(&self) {
2263 let should_extract = self
2264 .facts_config
2265 .as_ref()
2266 .map(|c| c.enabled && c.auto_extract)
2267 .unwrap_or(false);
2268
2269 if !should_extract {
2270 debug!("fact extraction skipped because auto extraction is disabled");
2271 return;
2272 }
2273
2274 let msgs_since = *self.messages_since_extraction.read();
2275 if msgs_since < 2 {
2276 debug!(
2277 messages_since_extraction = msgs_since,
2278 "fact extraction skipped until threshold is reached"
2279 );
2280 return;
2281 }
2282
2283 match self.extract_facts_with_source(msgs_since, "auto").await {
2284 Ok(facts) => {
2285 if !facts.is_empty() {
2286 *self.messages_since_extraction.write() = 0;
2287 } else {
2288 debug!("fact extraction ran but found no new facts");
2289 }
2290 }
2291 Err(e) => {
2292 warn!("fact extraction failed: {}", e);
2293 }
2294 }
2295 }
2296
2297 pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2298 self.persona_manager = Some(manager);
2299 self
2300 }
2301
2302 pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2303 self.persona_manager.as_ref()
2304 }
2305
2306 pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2307 if config.is_enabled() {
2308 let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2309 .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2310 self.disambiguation_manager = Some(manager);
2311 }
2312 self
2313 }
2314
2315 pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2316 self.disambiguation_manager.as_ref()
2317 }
2318
2319 pub fn has_disambiguation(&self) -> bool {
2320 self.disambiguation_manager
2321 .as_ref()
2322 .is_some_and(|m| m.is_enabled())
2323 }
2324
2325 pub async fn init_storage(&self) -> Result<()> {
2326 let _guard = self.storage_init.lock().await;
2330 let mut storage = self.storage.read().clone();
2331 if storage.is_none() && !self.storage_config.is_none() {
2332 let storage_config = self.convert_storage_config();
2333 storage = create_storage(&storage_config).await?;
2334 *self.storage.write() = storage.clone();
2335 }
2336
2337 self.validate_storage_requirements(storage.as_deref())?;
2338 self.complete_facts_init().await;
2339 Ok(())
2340 }
2341
2342 fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2343 let facts_required = self
2344 .facts_config
2345 .as_ref()
2346 .is_some_and(|config| config.enabled)
2347 || self
2348 .actor_memory_config
2349 .as_ref()
2350 .is_some_and(|config| config.enabled);
2351 let relationships_required = self
2352 .relationship_manager
2353 .as_ref()
2354 .is_some_and(|manager| manager.config().persistence.enabled);
2355
2356 let Some(storage) = storage else {
2357 let mut requirements = Vec::new();
2358 if facts_required {
2359 requirements.push("actor facts or actor memory");
2360 }
2361 if relationships_required {
2362 requirements.push("persistent relationships");
2363 }
2364 if requirements.is_empty() {
2365 return Ok(());
2366 }
2367 return Err(AgentError::Config(format!(
2368 "Storage is required for enabled {} but none is configured or injected",
2369 requirements.join(" and ")
2370 )));
2371 };
2372
2373 if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2377 return Err(AgentError::UnsupportedStorageCapability(
2378 StorageCapability::ActorFacts,
2379 ));
2380 }
2381 if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2382 return Err(AgentError::UnsupportedStorageCapability(
2383 StorageCapability::ActorRelationships,
2384 ));
2385 }
2386 Ok(())
2387 }
2388
2389 async fn complete_facts_init(&self) {
2392 if self.fact_store.read().is_some() {
2393 return;
2394 }
2395 let storage = match self.storage.read().clone() {
2396 Some(s) => s,
2397 None => return,
2398 };
2399
2400 let facts_enabled = self
2401 .facts_config
2402 .as_ref()
2403 .map(|f| f.enabled)
2404 .unwrap_or(false);
2405 let actor_memory_enabled = self
2406 .actor_memory_config
2407 .as_ref()
2408 .map(|a| a.enabled)
2409 .unwrap_or(false);
2410
2411 if !facts_enabled && !actor_memory_enabled {
2412 return;
2413 }
2414
2415 let fc = self.facts_config.clone().unwrap_or_default();
2416 let store = Arc::new(ai_agents_facts::FactStore::new(
2417 storage,
2418 self.info.id.clone(),
2419 fc.clone(),
2420 ));
2421
2422 let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2423 let extractor_llm = fc
2424 .extractor_llm
2425 .as_ref()
2426 .and_then(|alias| self.llm_registry.get(alias).ok())
2427 .or_else(|| self.llm_registry.router().ok())
2428 .or_else(|| self.llm_registry.default().ok());
2429 extractor_llm.map(|llm| {
2430 Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2431 as Arc<dyn ai_agents_facts::FactExtractor>
2432 })
2433 } else {
2434 None
2435 };
2436
2437 *self.fact_store.write() = Some(store);
2438 *self.fact_extractor.write() = extractor;
2439 debug!(
2440 agent = %self.info.id,
2441 facts_enabled,
2442 actor_memory_enabled,
2443 "facts storage initialized"
2444 );
2445 }
2446
2447 fn convert_storage_config(&self) -> StorageStorageConfig {
2448 crate::spec::storage::to_storage_config(&self.storage_config)
2449 }
2450
2451 pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2452 self.storage.read().clone()
2453 }
2454
2455 pub fn storage_config(&self) -> &StorageConfig {
2456 &self.storage_config
2457 }
2458
2459 pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2461 self.spawner.as_ref()
2462 }
2463
2464 pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2466 self.spawner_registry.as_ref()
2467 }
2468
2469 pub fn has_spawner(&self) -> bool {
2470 self.spawner_registry.is_some()
2471 }
2472
2473 pub fn with_spawner_handles(
2474 mut self,
2475 spawner: Arc<crate::spawner::AgentSpawner>,
2476 registry: Arc<crate::spawner::AgentRegistry>,
2477 ) -> Self {
2478 self.spawner = Some(spawner);
2479 self.spawner_registry = Some(registry);
2480 self
2481 }
2482
2483 pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2484 self.hooks = hooks;
2485 self
2486 }
2487
2488 pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2489 self.parallel_tools = config;
2490 self
2491 }
2492
2493 pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2494 self.streaming = config;
2495 self
2496 }
2497
2498 pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2499 self.hitl_engine = Some(engine);
2500 self.approval_handler = handler;
2501 self
2502 }
2503
2504 pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2505 self.max_context_tokens = tokens;
2506 self
2507 }
2508
2509 pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2510 self.memory_token_budget = Some(budget);
2511 self
2512 }
2513
2514 pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2515 self.recovery_manager = manager;
2516 self
2517 }
2518
2519 pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2520 self.tool_security = engine;
2521 self
2522 }
2523
2524 pub fn runtime_control(&self) -> RuntimeControlHandle {
2526 RuntimeControlHandle {
2527 state: Arc::clone(&self.runtime_control),
2528 }
2529 }
2530
2531 pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2533 self.tools.set_question_handler(handler);
2534 }
2535
2536 pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2538 self.tools.set_diagnostics_provider(provider);
2539 }
2540
2541 pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2543 self.tools.set_command_runner(runner);
2544 }
2545
2546 pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2548 self.tools.set_web_search_provider(provider);
2549 }
2550
2551 pub fn todos(&self) -> Vec<TodoItem> {
2553 self.tools.todos()
2554 }
2555
2556 fn active_tool_security(&self) -> ToolSecurityEngine {
2558 self.runtime_control
2559 .tool_security_override
2560 .read()
2561 .clone()
2562 .unwrap_or_else(|| self.tool_security.clone())
2563 }
2564
2565 fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2567 let _guard = self.runtime_control.snapshot_guard.read();
2568 RuntimeSafetySnapshot {
2569 version: self.runtime_control.version.load(Ordering::SeqCst),
2570 emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2571 tool_security: self
2572 .runtime_control
2573 .tool_security_override
2574 .read()
2575 .clone()
2576 .unwrap_or_else(|| self.tool_security.clone()),
2577 tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2578 }
2579 }
2580
2581 fn admit_tool_execution(
2583 &self,
2584 expected_runtime_version: u64,
2585 expected_policy_version: u64,
2586 expected_state_generation: Option<u64>,
2587 canonical_id: &str,
2588 ) -> SecurityCheckResult {
2589 let _guard = self.runtime_control.snapshot_guard.read();
2590 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2591 return SecurityCheckResult::Block {
2592 reason: "runtime emergency deny is enabled".to_string(),
2593 };
2594 }
2595 let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2596 let security_engine = self
2597 .runtime_control
2598 .tool_security_override
2599 .read()
2600 .clone()
2601 .unwrap_or_else(|| self.tool_security.clone());
2602 if runtime_version != expected_runtime_version
2603 || security_engine.policy_version() != expected_policy_version
2604 {
2605 return SecurityCheckResult::Block {
2606 reason: "runtime safety controls changed before admission".to_string(),
2607 };
2608 }
2609 let current_state_generation = self
2610 .state_machine
2611 .as_ref()
2612 .map(|state_machine| state_machine.generation());
2613 if current_state_generation != expected_state_generation {
2614 return SecurityCheckResult::Block {
2615 reason: "state scope changed before admission".to_string(),
2616 };
2617 }
2618 security_engine.admit_tool_execution(canonical_id)
2619 }
2620
2621 pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2622 let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2623 self.process_processor = Some(processor);
2624 self
2625 }
2626
2627 pub fn with_state_machine(
2628 mut self,
2629 state_machine: Arc<StateMachine>,
2630 evaluator: Arc<dyn TransitionEvaluator>,
2631 ) -> Self {
2632 self.state_machine = Some(state_machine);
2633 self.transition_evaluator = Some(evaluator);
2634 self
2635 }
2636
2637 pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2638 self.context_manager = manager;
2639 self
2640 }
2641
2642 pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2643 self.message_filters.write().insert(name.into(), filter);
2644 }
2645
2646 pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2647 self.context_manager.update(key, value)
2648 }
2649
2650 pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2651 self.context_manager.update(path, value)
2652 }
2653
2654 pub fn get_context(&self) -> HashMap<String, Value> {
2655 self.build_context_with_overlays()
2656 }
2657
2658 pub fn remove_context(&self, key: &str) -> Option<Value> {
2659 self.context_manager.remove(key)
2660 }
2661
2662 pub async fn refresh_context(&self, key: &str) -> Result<()> {
2663 self.context_manager.refresh(key).await
2664 }
2665
2666 pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2667 self.context_manager.register_provider(name, provider);
2668 }
2669
2670 pub fn current_state(&self) -> Option<String> {
2671 self.state_machine.as_ref().map(|sm| sm.current())
2672 }
2673
2674 async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2676 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2677 let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2678 return;
2679 };
2680 if disambiguator.has_pending_confirmation().await {
2681 disambiguator.clear_pending().await;
2682 *self.pending_skill_id.write() = None;
2683 info!(
2684 confirmation_event = "invalidated",
2685 invalidation_reason = reason,
2686 "Runtime invalidated pending confirmation"
2687 );
2688 }
2689 }
2690
2691 async fn admit_disambiguation_redispatch(
2693 &self,
2694 expected_epoch: u64,
2695 expected_state_generation: Option<u64>,
2696 ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2697 let admission = self.disambiguation_admission.read().await;
2698 let state_generation = self
2699 .state_machine
2700 .as_ref()
2701 .map(|state_machine| state_machine.generation());
2702 if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2703 || state_generation != expected_state_generation
2704 {
2705 return Err(AgentError::Other(
2706 "Disambiguation ownership changed before redispatch admission".to_string(),
2707 ));
2708 }
2709 Ok(admission)
2710 }
2711
2712 fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2714 self.state_transition_reserved
2715 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2716 .ok()
2717 .map(|_| StateTransitionReservation {
2718 reserved: &self.state_transition_reserved,
2719 })
2720 }
2721
2722 async fn admit_optional_disambiguation_ownership(
2724 &self,
2725 ownership: Option<DisambiguationOwnership>,
2726 ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2727 match ownership {
2728 Some(ownership) => self
2729 .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2730 .await
2731 .map(Some),
2732 None => Ok(None),
2733 }
2734 }
2735
2736 pub async fn transition_to(&self, state: &str) -> Result<()> {
2738 let Some(ref sm) = self.state_machine else {
2739 return Ok(());
2740 };
2741 let claim_admission = self.disambiguation_admission.write().await;
2742 let reservation = self.reserve_state_transition().ok_or_else(|| {
2743 AgentError::Other("Another state transition is already in progress".to_string())
2744 })?;
2745 let from_state = sm.current();
2746 let expected_state_generation = sm.generation();
2747 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2748 let history_before = sm.history();
2749 drop(claim_admission);
2750
2751 self.execute_state_exit_actions(&from_state).await;
2752
2753 let admission = self.disambiguation_admission.write().await;
2754 if sm.current() != from_state
2755 || sm.generation() != expected_state_generation
2756 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2757 {
2758 return Err(AgentError::Other(
2759 "State ownership changed during manual transition preparation".to_string(),
2760 ));
2761 }
2762 sm.transition_to(state, "manual transition")?;
2763 self.invalidate_pending_confirmation("state_transition")
2764 .await;
2765 let entered = sm.current();
2766 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2767 drop(admission);
2768
2769 self.execute_state_enter_actions(&entered, is_reentry).await;
2770 drop(reservation);
2771 info!(to = %entered, "Manual state transition");
2772 Ok(())
2773 }
2774
2775 pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2776 self.state_machine
2777 .as_ref()
2778 .map(|sm| sm.history())
2779 .unwrap_or_default()
2780 }
2781
2782 pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2784 self.session_metadata.read().clone()
2785 }
2786
2787 pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2790 let allowed = self
2791 .actor_memory_config
2792 .as_ref()
2793 .map(|c| c.privacy.allow_deletion)
2794 .unwrap_or(true);
2795 if !allowed {
2796 return Err(AgentError::Config(
2797 "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2798 ));
2799 }
2800 let storage = self.storage.read().clone();
2801 if let Some(storage) = storage {
2802 if !storage.supports(StorageCapability::ActorDataDeletion) {
2806 return Err(AgentError::UnsupportedStorageCapability(
2807 StorageCapability::ActorDataDeletion,
2808 ));
2809 }
2810 storage.delete_actor_data(&self.info.id, actor_id).await?;
2811 } else {
2812 let store = { self.fact_store.read().clone() };
2816 if let Some(store) = store {
2817 store.delete_actor_data(actor_id).await?;
2818 }
2819 }
2820 if let Some(manager) = self.relationship_manager.as_ref() {
2821 manager.remove(actor_id);
2822 }
2823 self.actor_facts_cache.write().remove(actor_id);
2824 Ok(())
2825 }
2826
2827 pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2829 *self.session_metadata.write() = meta;
2830 }
2831
2832 pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2834 let storage = self.storage.read().clone();
2835 match storage {
2836 Some(s) => {
2837 let count = s.cleanup_expired().await?;
2838 if count > 0 {
2839 self.hooks.on_sessions_expired(count).await;
2840 }
2841 Ok(count)
2842 }
2843 None => Err(AgentError::Config(
2844 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2845 )),
2846 }
2847 }
2848
2849 pub async fn list_sessions_filtered(
2851 &self,
2852 filter: &ai_agents_core::SessionFilter,
2853 ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2854 let storage = self.storage.read().clone();
2855 match storage {
2856 Some(s) => s.list_sessions_filtered(filter).await,
2857 None => Err(AgentError::Config(
2858 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2859 )),
2860 }
2861 }
2862
2863 pub async fn save_state(&self) -> Result<AgentSnapshot> {
2864 let memory_snapshot = self.memory.snapshot().await?;
2865 let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2866 let context_snapshot = self.context_manager.snapshot();
2867
2868 let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2869 .with_memory(memory_snapshot)
2870 .with_context(context_snapshot)
2871 .with_state_machine(
2872 state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2873 current_state: String::new(),
2874 previous_state: None,
2875 turn_count: 0,
2876 no_transition_count: 0,
2877 history: vec![],
2878 }),
2879 );
2880
2881 if let Some(ref persona) = self.persona_manager {
2882 snapshot.persona = Some(persona.snapshot_as_value()?);
2883 }
2884
2885 if let Some(ref relationships) = self.relationship_manager {
2886 snapshot.relationships = Some(relationships.snapshot_as_value()?);
2887 }
2888
2889 Ok(snapshot)
2890 }
2891
2892 pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2894 let mut snapshot = self.save_state().await?;
2895 if let Some(ref registry) = self.spawner_registry {
2896 let entries = registry.list_with_specs();
2897 if !entries.is_empty() {
2898 snapshot = snapshot.with_spawned_agents(entries);
2899 }
2900 }
2901 Ok(snapshot)
2902 }
2903
2904 pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2906 let _admission = self.disambiguation_admission.write().await;
2907 if self.state_transition_reserved.load(Ordering::SeqCst) {
2908 return Err(AgentError::Other(
2909 "Cannot restore state while a state transition is in progress".to_string(),
2910 ));
2911 }
2912 self.invalidate_pending_confirmation("state_restore").await;
2913 *self.pending_skill_id.write() = None;
2914 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2915 disambiguator.clear_pending().await;
2916 }
2917 self.memory.restore(snapshot.memory).await?;
2918
2919 if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2920 && !sm_snapshot.current_state.is_empty()
2921 {
2922 sm.restore(sm_snapshot)?;
2923 }
2924
2925 self.context_manager.restore(snapshot.context);
2926
2927 if let (Some(persona_value), Some(persona_manager)) =
2928 (snapshot.persona, &self.persona_manager)
2929 {
2930 persona_manager.restore_from_value(persona_value)?;
2931 }
2932
2933 if let (Some(relationship_value), Some(relationship_manager)) =
2934 (snapshot.relationships, &self.relationship_manager)
2935 {
2936 relationship_manager.restore_from_value(relationship_value)?;
2937 }
2938
2939 info!(agent_id = %snapshot.agent_id, "State restored");
2940 Ok(())
2941 }
2942
2943 pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2944 let snapshot = self.save_state().await?;
2945 storage.save(session_id, &snapshot).await
2946 }
2947
2948 async fn load_session_restore(
2949 storage: &dyn AgentStorage,
2950 session_id: &str,
2951 ) -> Result<Option<StoredSessionRestore>> {
2952 let Some(snapshot) = storage.load(session_id).await? else {
2953 return Ok(None);
2954 };
2955 let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2959 storage.load_metadata(session_id).await?
2960 } else {
2961 None
2962 };
2963 Ok(Some(StoredSessionRestore { snapshot, metadata }))
2964 }
2965
2966 async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2967 Ok(RuntimeSessionRestorePoint {
2968 snapshot: self.save_state().await?,
2969 metadata: self.session_metadata(),
2970 actor_id: self.actor_id(),
2971 session_id: self.current_session_id.read().clone(),
2972 })
2973 }
2974
2975 async fn apply_session_restore_unchecked(
2976 &self,
2977 session_id: &str,
2978 stored: StoredSessionRestore,
2979 ) -> Result<()> {
2980 self.restore_state(stored.snapshot).await?;
2981 let metadata = stored.metadata.unwrap_or_default();
2982 if let Some(actor_id) = metadata.actor_id.as_deref() {
2983 self.set_actor_id(actor_id)?;
2984 } else {
2985 self.clear_actor_id();
2986 }
2987 self.set_session_metadata(metadata);
2988 *self.current_session_id.write() = Some(session_id.to_string());
2989 Ok(())
2990 }
2991
2992 async fn restore_session_restore_point(
2993 &self,
2994 restore_point: &RuntimeSessionRestorePoint,
2995 ) -> Result<()> {
2996 self.restore_state(restore_point.snapshot.clone()).await?;
2997 if let Some(actor_id) = restore_point.actor_id.as_deref() {
2998 self.set_actor_id(actor_id)?;
2999 } else {
3000 self.clear_actor_id();
3001 }
3002 self.set_session_metadata(restore_point.metadata.clone());
3003 *self.current_session_id.write() = restore_point.session_id.clone();
3004 Ok(())
3005 }
3006
3007 async fn apply_session_restore(
3008 &self,
3009 session_id: &str,
3010 stored: StoredSessionRestore,
3011 ) -> Result<()> {
3012 let before = self.capture_session_restore_point().await?;
3013 if let Err(error) = self
3014 .apply_session_restore_unchecked(session_id, stored)
3015 .await
3016 {
3017 return match self.restore_session_restore_point(&before).await {
3018 Ok(()) => Err(error),
3019 Err(rollback_error) => Err(AgentError::Other(format!(
3020 "Session restore failed: {error}; rollback failed: {rollback_error}"
3021 ))),
3022 };
3023 }
3024 Ok(())
3025 }
3026
3027 async fn rollback_session_restore_set(
3028 parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
3029 children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
3030 ) -> Vec<String> {
3031 let mut errors = Vec::new();
3032 if let Some((agent, restore_point)) = parent
3033 && let Err(error) = agent.restore_session_restore_point(restore_point).await
3034 {
3035 errors.push(format!("parent: {error}"));
3036 }
3037 for (id, agent, restore_point) in children {
3038 if let Err(error) = agent.restore_session_restore_point(restore_point).await {
3039 errors.push(format!("child '{id}': {error}"));
3040 }
3041 }
3042 errors
3043 }
3044
3045 fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
3046 if rollback_errors.is_empty() {
3047 AgentError::Other(format!(
3048 "Session restore failed: {error}; runtime state was rolled back"
3049 ))
3050 } else {
3051 AgentError::Other(format!(
3052 "Session restore failed: {error}; rollback also failed for {}",
3053 rollback_errors.join(", ")
3054 ))
3055 }
3056 }
3057
3058 pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
3059 let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
3060 return Ok(false);
3061 };
3062 self.apply_session_restore(session_id, stored).await?;
3063 Ok(true)
3064 }
3065
3066 pub async fn save_session(&self, session_id: &str) -> Result<()> {
3067 let storage = self.storage.read().clone();
3068 match storage {
3069 Some(s) => {
3070 let is_new = {
3072 let cur = self.current_session_id.read().clone();
3073 cur.as_deref() != Some(session_id)
3074 };
3075 if is_new {
3076 *self.current_session_id.write() = Some(session_id.to_string());
3077 self.hooks.on_session_created(session_id).await;
3078 }
3079
3080 {
3082 let now = chrono::Utc::now();
3083 let msg_count = self
3084 .memory
3085 .get_messages(None)
3086 .await
3087 .map(|v| v.len())
3088 .unwrap_or(0);
3089 let mut meta = self.session_metadata.write();
3090 meta.last_active = now;
3091 meta.message_count = msg_count;
3092 if meta.actor_id.is_none() {
3093 meta.actor_id = self.actor_id.read().clone();
3094 }
3095 }
3096
3097 let snapshot = self.save_state().await?;
3098 if s.supports(StorageCapability::SessionMetadata) {
3102 let metadata = self.session_metadata.read().clone();
3103 s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
3104 .await
3105 } else {
3106 s.save(session_id, &snapshot).await
3107 }
3108 }
3109 None => Err(AgentError::Config(
3110 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3111 )),
3112 }
3113 }
3114
3115 pub async fn load_session(&self, session_id: &str) -> Result<bool> {
3116 let storage = self.storage.read().clone();
3117 match storage {
3118 Some(storage) => self.load_from(storage.as_ref(), session_id).await,
3119 None => Err(AgentError::Config(
3120 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3121 )),
3122 }
3123 }
3124
3125 pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
3127 self.init_storage().await?;
3128 let storage = self.storage.read().clone().ok_or_else(|| {
3129 AgentError::Config(
3130 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3131 )
3132 })?;
3133 let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
3134 .await?
3135 .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
3136 let manifest = target_parent
3137 .snapshot
3138 .spawned_agents
3139 .clone()
3140 .unwrap_or_default();
3141
3142 let registry = self.spawner_registry.as_ref().cloned();
3143 let spawner = if manifest.is_empty() {
3144 self.spawner.as_ref().cloned()
3145 } else {
3146 Some(self.spawner.as_ref().cloned().ok_or_else(|| {
3147 AgentError::Config(
3148 "Saved session contains child agents but this runtime has no spawner".into(),
3149 )
3150 })?)
3151 };
3152 let registry = if manifest.is_empty() {
3153 registry
3154 } else {
3155 Some(registry.ok_or_else(|| {
3156 AgentError::Config(
3157 "Saved session contains child agents but this runtime has no registry".into(),
3158 )
3159 })?)
3160 };
3161
3162 let mut target_ids = HashSet::with_capacity(manifest.len());
3163 let mut prepared = Vec::with_capacity(manifest.len());
3164 for entry in manifest {
3165 if !target_ids.insert(entry.id.clone()) {
3166 return Err(AgentError::InvalidSpec(format!(
3167 "Saved child manifest contains duplicate ID: {}",
3168 entry.id
3169 )));
3170 }
3171 let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
3172 spawner
3173 .as_ref()
3174 .expect("non-empty manifests require a spawner")
3175 .validate_explicit_child(&entry.id, &spec)?;
3176 prepared.push((entry.id, spec));
3177 }
3178
3179 let current_ids = registry
3180 .as_ref()
3181 .map(|registry| {
3182 registry
3183 .list()
3184 .into_iter()
3185 .map(|info| info.id)
3186 .collect::<HashSet<_>>()
3187 })
3188 .unwrap_or_default();
3189 let removal_count = current_ids.difference(&target_ids).count();
3190 let additions = prepared
3191 .iter()
3192 .filter(|(id, _)| !current_ids.contains(id))
3193 .cloned()
3194 .collect::<Vec<_>>();
3195
3196 let mut existing = Vec::new();
3197 if let Some(registry) = registry.as_ref() {
3198 for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3199 let agent = registry.get(id).ok_or_else(|| {
3200 AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3201 })?;
3202 let child_storage = agent.storage().ok_or_else(|| {
3203 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3204 })?;
3205 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3206 .await?
3207 .ok_or_else(|| {
3208 AgentError::Persistence(format!(
3209 "Child '{id}' has no saved session '{session_id}'"
3210 ))
3211 })?;
3212 existing.push((id.clone(), agent, stored));
3213 }
3214 }
3215
3216 let mut staged = Vec::with_capacity(additions.len());
3217 if !additions.is_empty() {
3218 let spawner = spawner
3219 .as_ref()
3220 .expect("restored additions require a spawner");
3221 let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3222 for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3223 let spawned = spawner
3224 .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3225 .await?;
3226 let child_storage = spawned.agent.storage().ok_or_else(|| {
3227 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3228 })?;
3229 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3230 .await?
3231 .ok_or_else(|| {
3232 AgentError::Persistence(format!(
3233 "Child '{id}' has no saved session '{session_id}'"
3234 ))
3235 })?;
3236 staged.push((spawned, stored));
3237 }
3238 } else if let Some(spawner) = spawner.as_ref() {
3239 spawner.reserve_restore_capacity(0, removal_count)?;
3240 }
3241
3242 let parent_before = self.capture_session_restore_point().await?;
3243 let mut existing_before = Vec::with_capacity(existing.len());
3244 for (id, agent, _) in &existing {
3245 existing_before.push((
3246 id.clone(),
3247 Arc::clone(agent),
3248 agent.capture_session_restore_point().await?,
3249 ));
3250 }
3251
3252 for (_, agent, stored) in &existing {
3256 if let Err(error) = agent
3257 .apply_session_restore_unchecked(session_id, stored.clone())
3258 .await
3259 {
3260 drop(staged);
3261 let rollback_errors =
3262 Self::rollback_session_restore_set(None, &existing_before).await;
3263 return Err(Self::restore_failure(error, rollback_errors));
3264 }
3265 }
3266 for (spawned, stored) in &staged {
3267 if let Err(error) = spawned
3268 .agent
3269 .apply_session_restore_unchecked(session_id, stored.clone())
3270 .await
3271 {
3272 drop(staged);
3273 let rollback_errors =
3274 Self::rollback_session_restore_set(None, &existing_before).await;
3275 return Err(Self::restore_failure(error, rollback_errors));
3276 }
3277 }
3278 if let Err(error) = self
3279 .apply_session_restore_unchecked(session_id, target_parent)
3280 .await
3281 {
3282 drop(staged);
3283 let rollback_errors =
3284 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3285 .await;
3286 return Err(Self::restore_failure(error, rollback_errors));
3287 }
3288
3289 if let Some(registry) = registry.as_ref()
3290 && let Err(error) = registry
3291 .reconcile(
3292 &target_ids,
3293 staged.into_iter().map(|(spawned, _)| spawned).collect(),
3294 )
3295 .await
3296 {
3297 let rollback_errors =
3298 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3299 .await;
3300 return Err(Self::restore_failure(error, rollback_errors));
3301 }
3302
3303 Ok(target_ids.len())
3304 }
3305
3306 pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3307 let storage = self.storage.read().clone();
3308 match storage {
3309 Some(s) => s.delete(session_id).await,
3310 None => Err(AgentError::Config(
3311 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3312 )),
3313 }
3314 }
3315
3316 pub async fn list_sessions(&self) -> Result<Vec<String>> {
3317 let storage = self.storage.read().clone();
3318 match storage {
3319 Some(s) => s.list_sessions().await,
3320 None => Err(AgentError::Config(
3321 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3322 )),
3323 }
3324 }
3325
3326 fn estimate_tokens(&self, text: &str) -> u32 {
3327 (text.len() as f32 / 4.0).ceil() as u32
3328 }
3329
3330 fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3331 messages
3332 .iter()
3333 .map(|m| self.estimate_tokens(&m.content))
3334 .sum()
3335 }
3336
3337 fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3338 if messages.len() <= keep_recent + 1 {
3339 return;
3340 }
3341 let system_msg = messages.remove(0);
3342 let to_remove = messages.len().saturating_sub(keep_recent);
3343 messages.drain(..to_remove);
3344 messages.insert(0, system_msg);
3345 }
3346
3347 fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3348 match config {
3349 FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3350 FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3351 FilterConfig::SkipPattern { skip_if_contains } => {
3352 Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3353 }
3354 FilterConfig::Custom { name } => {
3355 let filters = self.message_filters.read();
3356 filters
3357 .get(name)
3358 .cloned()
3359 .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3360 }
3361 }
3362 }
3363
3364 async fn summarize_context(
3365 &self,
3366 messages: &mut Vec<ChatMessage>,
3367 summarizer_llm: Option<&str>,
3368 max_summary_tokens: u32,
3369 custom_prompt: Option<&str>,
3370 keep_recent: usize,
3371 filter: Option<&FilterConfig>,
3372 ) -> Result<()> {
3373 let system_msg = messages.remove(0);
3374
3375 let to_summarize_count = messages.len().saturating_sub(keep_recent);
3376 if to_summarize_count == 0 {
3377 messages.insert(0, system_msg);
3378 return Ok(());
3379 }
3380
3381 let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3382 let mut to_summarize = std::mem::take(messages);
3383
3384 if let Some(filter_config) = filter {
3385 let filter = self.get_filter(filter_config);
3386 to_summarize = filter.filter(to_summarize);
3387 }
3388
3389 if to_summarize.is_empty() {
3390 *messages = recent_msgs;
3391 messages.insert(0, system_msg);
3392 return Ok(());
3393 }
3394
3395 let conversation_text = to_summarize
3396 .iter()
3397 .map(|m| format!("{:?}: {}", m.role, m.content))
3398 .collect::<Vec<_>>()
3399 .join("\n");
3400
3401 let default_prompt = format!(
3402 "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3403 max_summary_tokens, conversation_text
3404 );
3405
3406 let summary_prompt = custom_prompt
3407 .map(|p| format!("{}\n\n{}", p, conversation_text))
3408 .unwrap_or(default_prompt);
3409
3410 let summarizer = if let Some(alias) = summarizer_llm {
3411 self.llm_registry
3412 .get(alias)
3413 .map_err(|e| AgentError::Config(e.to_string()))?
3414 } else {
3415 self.llm_registry
3416 .router()
3417 .or_else(|_| self.llm_registry.default())
3418 .map_err(|e| AgentError::Config(e.to_string()))?
3419 };
3420
3421 let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3422 let response = self
3423 .observe_purpose(
3424 ObservationPurpose::Summarization,
3425 summarizer.complete(&summary_msgs, None),
3426 )
3427 .await?;
3428
3429 let summary_message = ChatMessage::system(format!(
3430 "[Previous conversation summary]\n{}",
3431 response.content
3432 ));
3433
3434 *messages = vec![system_msg, summary_message];
3435 messages.extend(recent_msgs);
3436
3437 debug!(
3438 summarized_count = to_summarize_count,
3439 kept_recent = keep_recent,
3440 "Context summarized"
3441 );
3442
3443 Ok(())
3444 }
3445
3446 fn render_system_prompt(&self) -> Result<String> {
3447 let mut context = self.build_context_with_overlays();
3448
3449 let facts_text = self.format_actor_facts_for_context();
3451 if !facts_text.is_empty() {
3452 context.insert(
3453 "actor_facts".to_string(),
3454 serde_json::Value::String(facts_text),
3455 );
3456 }
3457
3458 if let Some((key, text)) = self.format_relationship_for_context() {
3459 context.insert(key, serde_json::Value::String(text));
3460 }
3461
3462 self.template_renderer
3463 .render(&self.base_system_prompt, &context)
3464 }
3465
3466 fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3468 let mut seen = HashSet::new();
3469 ids.iter()
3470 .filter_map(|id| self.tools.canonical_id(id))
3471 .filter(|canonical_id| seen.insert(canonical_id.clone()))
3472 .collect()
3473 }
3474
3475 fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3477 let Some(declared) = self.declared_tool_ids.as_deref() else {
3478 return Vec::new();
3479 };
3480 let mut effective = self.canonical_unique_tool_ids(declared);
3481 if let Some(scope) = scope_override {
3482 let scope: HashSet<String> =
3483 self.canonical_unique_tool_ids(scope).into_iter().collect();
3484 effective.retain(|canonical_id| scope.contains(canonical_id));
3485 }
3486 effective
3487 }
3488
3489 async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3491 Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3492 }
3493
3494 async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3496 let scope_override = self.runtime_control.tool_scope_override.read().clone();
3497 self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3498 .await
3499 }
3500
3501 async fn get_available_tool_ids_snapshot_for_scope(
3503 &self,
3504 scope_override: Option<&[String]>,
3505 ) -> Result<AvailableToolIdsSnapshot> {
3506 let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3507 let (state_generation, state_scopes) = self
3508 .state_machine
3509 .as_ref()
3510 .map(|state_machine| {
3511 let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3512 (Some(generation), scopes)
3513 })
3514 .unwrap_or((None, Vec::new()));
3515
3516 if available.is_empty() || state_scopes.is_empty() {
3517 return Ok(AvailableToolIdsSnapshot {
3518 tool_ids: available,
3519 state_generation,
3520 });
3521 }
3522
3523 let eval_ctx = self.build_evaluation_context().await?;
3524 let llm_getter = RegistryLLMGetter {
3525 registry: self.llm_registry.clone(),
3526 };
3527 let evaluator = ConditionEvaluator::new(llm_getter);
3528
3529 for state_scope in state_scopes {
3530 if state_scope.is_empty() {
3531 available.clear();
3532 break;
3533 }
3534
3535 let mut allowed = HashSet::new();
3536 for tool_ref in &state_scope {
3537 let tool_id = tool_ref.id();
3538 let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3539 continue;
3540 };
3541 let condition_matches = if let Some(condition) = tool_ref.condition() {
3542 match evaluator.evaluate(condition, &eval_ctx).await {
3543 Ok(matches) => matches,
3544 Err(error) => {
3545 warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3546 false
3547 }
3548 }
3549 } else {
3550 true
3551 };
3552 if condition_matches {
3553 allowed.insert(canonical_id);
3554 } else {
3555 debug!(tool = tool_id, "Tool condition not met, skipping");
3556 }
3557 }
3558 available.retain(|canonical_id| allowed.contains(canonical_id));
3559 if available.is_empty() {
3560 break;
3561 }
3562 }
3563
3564 Ok(AvailableToolIdsSnapshot {
3565 tool_ids: available,
3566 state_generation,
3567 })
3568 }
3569
3570 async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3571 let context = self.build_context_with_overlays();
3572 let messages = self.memory.get_messages(Some(10)).await?;
3573 let tool_history = self.tool_call_history.read().clone();
3574
3575 let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3576 (Some(sm.current()), sm.turn_count(), sm.previous())
3577 } else {
3578 (None, 0, None)
3579 };
3580
3581 Ok(EvaluationContext::default()
3582 .with_context(context)
3583 .with_state(state_name, turn_count, previous_state)
3584 .with_called_tools(tool_history)
3585 .with_messages(messages))
3586 }
3587
3588 fn record_tool_call(&self, tool_id: &str, result: Value) {
3589 self.tool_call_history.write().push(ToolCallRecord {
3590 tool_id: tool_id.to_string(),
3591 result,
3592 timestamp: chrono::Utc::now(),
3593 });
3594 }
3595
3596 async fn get_effective_system_prompt_with_persona_hooks(
3597 &self,
3598 fire_persona_hooks: bool,
3599 include_tool_prompt: bool,
3600 ) -> Result<String> {
3601 let rendered_base = self.render_system_prompt()?;
3602
3603 let persona_prefix = if let Some(ref persona) = self.persona_manager {
3604 let context = self.build_context_with_overlays();
3605 if fire_persona_hooks {
3606 let render_result = persona.render_prompt(&context)?;
3607 for content in &render_result.newly_revealed {
3608 self.hooks.on_secret_revealed(content).await;
3609 }
3610 render_result.prompt
3611 } else {
3612 persona.render_prompt_preview(&context)?
3613 }
3614 } else {
3615 String::new()
3616 };
3617
3618 if let Some(ref sm) = self.state_machine
3619 && let Some(state_def) = sm.current_definition()
3620 {
3621 let state_prompt = if let Some(ref prompt) = state_def.prompt {
3622 let context = self.build_context_with_overlays();
3623 self.template_renderer.render_with_state(
3624 prompt,
3625 &context,
3626 &sm.current(),
3627 sm.previous().as_deref(),
3628 sm.turn_count(),
3629 state_def.max_turns,
3630 )?
3631 } else {
3632 String::new()
3633 };
3634
3635 let combined = match state_def.prompt_mode {
3636 PromptMode::Append => {
3637 if state_prompt.is_empty() {
3638 rendered_base
3639 } else {
3640 format!(
3641 "{}\n\n[Current State: {}]\n{}",
3642 rendered_base,
3643 sm.current(),
3644 state_prompt
3645 )
3646 }
3647 }
3648 PromptMode::Replace => {
3649 if state_prompt.is_empty() {
3650 rendered_base
3651 } else {
3652 state_prompt
3653 }
3654 }
3655 PromptMode::Prepend => {
3656 if state_prompt.is_empty() {
3657 rendered_base
3658 } else {
3659 format!("{}\n\n{}", state_prompt, rendered_base)
3660 }
3661 }
3662 };
3663
3664 let with_persona = if persona_prefix.is_empty() {
3666 combined
3667 } else {
3668 format!("{}\n\n{}", persona_prefix, combined)
3669 };
3670
3671 if include_tool_prompt {
3672 let available_tool_ids = self.get_available_tool_ids().await?;
3673 if !available_tool_ids.is_empty() {
3674 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3675 &available_tool_ids,
3676 None,
3677 self.parallel_tools.enabled,
3678 self.runtime_config.tool_schema_prompt_mode,
3679 );
3680 if !tools_prompt.is_empty() {
3681 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3682 }
3683 }
3684 }
3685 return Ok(with_persona);
3686 }
3687
3688 let with_persona = if persona_prefix.is_empty() {
3690 rendered_base
3691 } else {
3692 format!("{}\n\n{}", persona_prefix, rendered_base)
3693 };
3694
3695 if include_tool_prompt {
3696 let available_tool_ids = self.get_available_tool_ids().await?;
3697 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3698 &available_tool_ids,
3699 None,
3700 self.parallel_tools.enabled,
3701 self.runtime_config.tool_schema_prompt_mode,
3702 );
3703 if !tools_prompt.is_empty() {
3704 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3705 }
3706 }
3707 Ok(with_persona)
3708 }
3709
3710 fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3711 if let Some(ref sm) = self.state_machine
3712 && let Some(state_def) = sm.current_definition()
3713 && let Some(ref llm_alias) = state_def.llm
3714 {
3715 return self
3716 .llm_registry
3717 .get(llm_alias)
3718 .map_err(|e| AgentError::Config(e.to_string()));
3719 }
3720 self.llm_registry
3721 .default()
3722 .map_err(|e| AgentError::Config(e.to_string()))
3723 }
3724
3725 fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3726 if let Some(ref sm) = self.state_machine
3727 && let Some(state_def) = sm.current_definition()
3728 && let Some(ref state_reasoning) = state_def.reasoning
3729 {
3730 return state_reasoning.clone();
3731 }
3732 self.reasoning_config.clone()
3733 }
3734
3735 fn get_effective_reflection_config(&self) -> ReflectionConfig {
3736 if let Some(ref sm) = self.state_machine
3737 && let Some(state_def) = sm.current_definition()
3738 && let Some(ref state_reflection) = state_def.reflection
3739 {
3740 return state_reflection.clone();
3741 }
3742 self.reflection_config.clone()
3743 }
3744
3745 fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3746 skill
3747 .reasoning
3748 .clone()
3749 .unwrap_or_else(|| self.get_effective_reasoning_config())
3750 }
3751
3752 fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3753 skill
3754 .reflection
3755 .clone()
3756 .unwrap_or_else(|| self.get_effective_reflection_config())
3757 }
3758
3759 async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3760 let recent_messages: Vec<String> = self
3761 .memory
3762 .get_messages(Some(5))
3763 .await?
3764 .iter()
3765 .rev()
3766 .map(|m| format!("{:?}: {}", m.role, m.content))
3767 .collect();
3768
3769 let current_state = self.current_state().map(|s| s.to_string());
3770
3771 let state_prompt: Option<String> = self
3774 .state_machine
3775 .as_ref()
3776 .and_then(|sm| sm.current_definition())
3777 .and_then(|def| def.prompt.clone());
3778
3779 let available_tools: Vec<String> = self
3780 .get_available_tool_ids()
3781 .await
3782 .unwrap_or_else(|_| self.tools.list_ids());
3783
3784 let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3785
3786 let mut user_context = self.build_context_with_overlays();
3787 user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3788 if let Some(state_generation) = self
3789 .state_machine
3790 .as_ref()
3791 .map(|state_machine| state_machine.generation())
3792 {
3793 user_context.insert(
3794 DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3795 serde_json::json!(state_generation),
3796 );
3797 }
3798
3799 let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3801 sm.current_definition()
3802 .map(|def| {
3803 def.transitions
3804 .iter()
3805 .filter_map(|t| t.intent.clone())
3806 .collect()
3807 })
3808 .unwrap_or_default()
3809 } else {
3810 Vec::new()
3811 };
3812
3813 Ok(DisambiguationContext::from_agent_state(
3814 recent_messages,
3815 current_state,
3816 state_prompt,
3817 available_tools,
3818 available_skills,
3819 available_intents,
3820 user_context,
3821 ))
3822 }
3823
3824 fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3825 if let Some(ref sm) = self.state_machine
3826 && let Some(state_def) = sm.current_definition()
3827 {
3828 let parent_def = sm.get_parent_definition();
3829 let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3830 if !effective_skills.is_empty() {
3831 return self
3832 .skills
3833 .iter()
3834 .filter(|s| effective_skills.contains(&&s.id))
3835 .collect();
3836 }
3837 }
3838 self.skills.iter().collect()
3839 }
3840
3841 async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3842 self.build_messages_internal(true, None, true).await
3843 }
3844
3845 async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3846 self.build_messages_internal(false, Some(user_message), true)
3847 .await
3848 }
3849
3850 async fn build_messages_internal(
3851 &self,
3852 fire_persona_hooks: bool,
3853 ephemeral_user_message: Option<&str>,
3854 include_tool_prompt: bool,
3855 ) -> Result<Vec<ChatMessage>> {
3856 let system_prompt = self
3857 .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3858 .await?;
3859 let mut messages = vec![ChatMessage::system(&system_prompt)];
3860
3861 let context = self.memory.get_context().await?;
3862 let history = if let Some(ref budget) = self.memory_token_budget {
3863 context.to_llm_messages_with_allocation(&budget.allocation)
3864 } else {
3865 context.to_llm_messages()
3866 };
3867 messages.extend(history);
3868 if let Some(user_message) = ephemeral_user_message {
3869 messages.push(ChatMessage::user(user_message));
3870 }
3871
3872 let total_tokens = self.estimate_total_tokens(&messages);
3873
3874 if total_tokens > self.max_context_tokens {
3875 debug!(
3876 total = total_tokens,
3877 limit = self.max_context_tokens,
3878 "Context overflow"
3879 );
3880
3881 match &self.recovery_manager.config().llm.on_context_overflow {
3882 ContextOverflowAction::Error => {
3883 return Err(AgentError::LLM(format!(
3884 "Context overflow: {} tokens > {} limit",
3885 total_tokens, self.max_context_tokens
3886 )));
3887 }
3888 ContextOverflowAction::Truncate { keep_recent } => {
3889 self.truncate_context(&mut messages, *keep_recent);
3890 }
3891 ContextOverflowAction::Summarize {
3892 summarizer_llm,
3893 max_summary_tokens,
3894 custom_prompt,
3895 keep_recent,
3896 filter,
3897 } => {
3898 self.summarize_context(
3899 &mut messages,
3900 summarizer_llm.as_deref(),
3901 *max_summary_tokens,
3902 custom_prompt.as_deref(),
3903 *keep_recent,
3904 filter.as_ref(),
3905 )
3906 .await?;
3907 }
3908 }
3909 }
3910
3911 Ok(messages)
3912 }
3913
3914 async fn main_tool_protocol(
3915 &self,
3916 llm: &dyn LLMProvider,
3917 ephemeral_new_turn: bool,
3918 ) -> Result<MainToolProtocol> {
3919 let mut choice = llm.configured_tool_choice();
3920 if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3921 return Ok(MainToolProtocol {
3922 choice,
3923 tool_ids: Vec::new(),
3924 definitions: Vec::new(),
3925 });
3926 }
3927
3928 let mut tool_ids = self.get_available_tool_ids().await?;
3929 tool_ids.sort();
3930 tool_ids.dedup();
3931 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3932 let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3933 AgentError::Config(format!(
3934 "specific tool choice '{expected}' is not registered"
3935 ))
3936 })?;
3937 if canonical != *expected {
3938 return Err(AgentError::Config(format!(
3939 "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3940 )));
3941 }
3942 if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3943 return Err(AgentError::Config(format!(
3944 "specific tool choice '{expected}' is outside the effective tool grant"
3945 )));
3946 }
3947 }
3948 if matches!(
3949 choice.as_ref(),
3950 Some(ToolChoice::Required | ToolChoice::Specific(_))
3951 ) && tool_ids.is_empty()
3952 {
3953 return Err(AgentError::Config(
3954 "required tool choice has no tool inside the effective grant".to_string(),
3955 ));
3956 }
3957 if !ephemeral_new_turn
3958 && let Some(configured_choice) = choice.as_ref()
3959 && matches!(
3960 configured_choice,
3961 ToolChoice::Required | ToolChoice::Specific(_)
3962 )
3963 && self
3964 .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3965 .await?
3966 {
3967 choice = Some(ToolChoice::Auto);
3968 }
3969 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3970 tool_ids.retain(|tool_id| tool_id == expected);
3971 }
3972
3973 let definitions = tool_ids
3974 .iter()
3975 .map(|tool_id| {
3976 let tool = self.tools.get(tool_id).ok_or_else(|| {
3977 AgentError::Config(format!(
3978 "effective tool '{tool_id}' disappeared before provider exposure"
3979 ))
3980 })?;
3981 Ok(LLMToolDefinition {
3982 name: tool_id.clone(),
3983 description: tool.description().to_string(),
3984 input_schema: tool.input_schema(),
3985 })
3986 })
3987 .collect::<Result<Vec<_>>>()?;
3988
3989 Ok(MainToolProtocol {
3993 choice,
3994 tool_ids,
3995 definitions,
3996 })
3997 }
3998
3999 async fn tool_choice_satisfied_in_current_turn(
4000 &self,
4001 choice: &ToolChoice,
4002 effective_tool_ids: &[String],
4003 ) -> Result<bool> {
4004 let messages = self.memory.get_messages(None).await?;
4005 let mut saw_tool_result = false;
4006 for message in messages.iter().rev() {
4007 match message.role {
4008 ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
4009 saw_tool_result = true;
4010 }
4011 ai_agents_core::Role::Assistant if saw_tool_result => {
4012 let Some(calls) = self.parse_tool_calls(&message.content) else {
4013 continue;
4014 };
4015 let calls_are_effective = !calls.is_empty()
4016 && calls.iter().all(|call| {
4017 self.tools
4018 .canonical_id(&call.name)
4019 .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
4020 });
4021 return Ok(calls_are_effective
4022 && match choice {
4023 ToolChoice::Required => true,
4024 ToolChoice::Specific(expected) => calls.iter().all(|call| {
4025 self.tools.canonical_id(&call.name).as_deref()
4026 == Some(expected.as_str())
4027 }),
4028 _ => false,
4029 });
4030 }
4031 ai_agents_core::Role::User => return Ok(false),
4032 _ => {}
4033 }
4034 }
4035 Ok(false)
4036 }
4037
4038 fn provider_can_use_native_tools(
4039 &self,
4040 llm: &dyn LLMProvider,
4041 protocol: &MainToolProtocol,
4042 ) -> bool {
4043 let Some(choice) = protocol.choice.as_ref() else {
4044 return false;
4045 };
4046 if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
4047 return false;
4048 }
4049 llm.supports_tool_choice(choice)
4050 && protocol.definitions.iter().all(|definition| {
4051 !definition.name.is_empty()
4052 && definition.name.len() <= 64
4053 && definition
4054 .name
4055 .bytes()
4056 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
4057 })
4058 }
4059
4060 fn prompt_messages_for_tool_protocol(
4061 &self,
4062 messages: &[ChatMessage],
4063 protocol: &MainToolProtocol,
4064 corrective: bool,
4065 ) -> Vec<ChatMessage> {
4066 let mut messages = messages.to_vec();
4067 let Some(choice) = protocol.choice.as_ref() else {
4068 return messages;
4069 };
4070 if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
4071 return messages;
4072 }
4073
4074 let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
4075 &protocol.tool_ids,
4076 None,
4077 self.parallel_tools.enabled,
4078 self.runtime_config.tool_schema_prompt_mode,
4079 );
4080 match choice {
4081 ToolChoice::Required => tool_prompt.push_str(
4082 "\n\nYou must call at least one listed tool before giving a final answer.",
4083 ),
4084 ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
4085 "\n\nYou must call the '{tool_id}' tool before giving a final answer."
4086 )),
4087 ToolChoice::Auto => {}
4088 ToolChoice::None => return messages,
4089 _ => return messages,
4090 }
4091 if let Some(system) = messages
4092 .iter_mut()
4093 .find(|message| message.role == ai_agents_core::Role::System)
4094 {
4095 system.content.push_str("\n\n");
4096 system.content.push_str(&tool_prompt);
4097 } else {
4098 messages.insert(0, ChatMessage::system(tool_prompt));
4099 }
4100 if corrective {
4101 let instruction = match choice {
4102 ToolChoice::Required => {
4103 "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
4104 }
4105 ToolChoice::Specific(tool_id) => {
4106 messages.push(ChatMessage::user(format!(
4107 "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
4108 )));
4109 return messages;
4110 }
4111 _ => return messages,
4112 };
4113 messages.push(ChatMessage::user(instruction));
4114 }
4115 messages
4116 }
4117
4118 async fn invoke_main_provider(
4119 &self,
4120 llm: Arc<dyn LLMProvider>,
4121 messages: &[ChatMessage],
4122 protocol: &MainToolProtocol,
4123 corrective: bool,
4124 ) -> std::result::Result<MainProviderResponse, LLMError> {
4125 let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
4126 let response = if use_native {
4127 let request = LLMToolRequest {
4128 tools: protocol.definitions.clone(),
4129 choice: protocol
4130 .choice
4131 .clone()
4132 .expect("native tool requests require an explicit choice"),
4133 };
4134 self.observe_purpose(
4135 ObservationPurpose::MainResponse,
4136 llm.complete_with_tools(messages, None, &request),
4137 )
4138 .await?
4139 } else {
4140 let prompt_messages =
4141 self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
4142 self.observe_purpose(
4143 ObservationPurpose::MainResponse,
4144 llm.complete(&prompt_messages, None),
4145 )
4146 .await?
4147 };
4148 Ok(MainProviderResponse {
4149 response,
4150 used_native_tools: use_native,
4151 })
4152 }
4153
4154 async fn complete_main_attempt_with_recovery(
4155 &self,
4156 llm: Arc<dyn LLMProvider>,
4157 messages: &[ChatMessage],
4158 protocol: &MainToolProtocol,
4159 corrective: bool,
4160 ) -> Result<MainProviderResponse> {
4161 let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
4162 self.recovery_manager
4163 .with_retry("llm_call", None, || {
4164 let llm = Arc::clone(&llm);
4165 async move {
4166 self.invoke_main_provider(llm, messages, protocol, corrective)
4167 .await
4168 .map_err(|error| error.classify())
4169 }
4170 })
4171 .await
4172 .map_err(|error| AgentError::LLM(error.to_string()))
4173 } else {
4174 self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4175 .await
4176 .map_err(|error| AgentError::LLM(error.to_string()))
4177 };
4178
4179 match primary_result {
4180 Ok(response) => Ok(response),
4181 Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4182 LLMFailureAction::FallbackLlm { fallback_llm } => {
4183 let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4184 AgentError::Config(format!(
4185 "Fallback LLM '{fallback_llm}' not found: {error}"
4186 ))
4187 })?;
4188 self.invoke_main_provider(fallback, messages, protocol, corrective)
4189 .await
4190 .map_err(|error| AgentError::LLM(error.to_string()))
4191 }
4192 LLMFailureAction::FallbackResponse { message } => {
4193 if matches!(
4194 protocol.choice.as_ref(),
4195 Some(ToolChoice::Required | ToolChoice::Specific(_))
4196 ) {
4197 Err(AgentError::LLM(format!(
4198 "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4199 )))
4200 } else {
4201 Ok(MainProviderResponse {
4202 response: LLMResponse::new(message.clone(), FinishReason::Stop),
4203 used_native_tools: false,
4204 })
4205 }
4206 }
4207 LLMFailureAction::Error => Err(primary_error),
4208 },
4209 }
4210 }
4211
4212 fn normalize_main_provider_response(
4213 &self,
4214 mut response: LLMResponse,
4215 protocol: &MainToolProtocol,
4216 ) -> Result<(LLMResponse, bool)> {
4217 let native_calls = response
4218 .tool_calls()
4219 .map_err(|error| AgentError::LLM(error.to_string()))?;
4220 let calls = match native_calls {
4221 Some(calls) => {
4222 let markers = calls
4223 .iter()
4224 .map(|call| {
4225 serde_json::json!({
4226 "_ai_agents_native_tool_call": true,
4227 "id": call.id,
4228 "tool": call.name,
4229 "arguments": call.arguments,
4230 })
4231 })
4232 .collect::<Vec<_>>();
4233 response.content = if markers.len() == 1 {
4234 markers[0].to_string()
4235 } else {
4236 serde_json::Value::Array(markers).to_string()
4237 };
4238 Some(calls)
4239 }
4240 None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4241 self.parse_tool_calls(response.content.trim())
4242 }
4243 None => None,
4244 };
4245
4246 if protocol.choice.is_some()
4247 && let Some(calls) = calls.as_ref()
4248 && calls.iter().any(|call| {
4249 self.tools
4250 .canonical_id(&call.name)
4251 .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4252 })
4253 {
4254 return Err(AgentError::LLM(
4255 "Provider returned a tool call outside the effective grant".to_string(),
4256 ));
4257 }
4258
4259 let compliant = match protocol.choice.as_ref() {
4260 Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4261 Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4262 !calls.is_empty()
4263 && calls.iter().all(|call| {
4264 self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4265 })
4266 }),
4267 _ => true,
4268 };
4269 Ok((response, compliant))
4270 }
4271
4272 async fn complete_main_llm_with_recovery(
4273 &self,
4274 llm: Arc<dyn LLMProvider>,
4275 messages: &[ChatMessage],
4276 protocol: &MainToolProtocol,
4277 ) -> Result<LLMResponse> {
4278 let first = self
4279 .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4280 .await?;
4281 let (response, compliant) =
4282 self.normalize_main_provider_response(first.response, protocol)?;
4283 if compliant {
4284 return Ok(response);
4285 }
4286 if first.used_native_tools {
4287 return Err(AgentError::LLM(
4288 "Provider returned no compliant native call for required tool choice".to_string(),
4289 ));
4290 }
4291
4292 let corrected = self
4293 .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4294 .await?;
4295 let (response, compliant) =
4296 self.normalize_main_provider_response(corrected.response, protocol)?;
4297 if compliant {
4298 return Ok(response);
4299 }
4300 Err(AgentError::LLM(
4301 "Provider returned no compliant tool call after one corrective retry".to_string(),
4302 ))
4303 }
4304
4305 fn is_native_tool_call_content(content: &str) -> bool {
4306 let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4307 return false;
4308 };
4309 match value {
4310 serde_json::Value::Array(values) => {
4311 !values.is_empty()
4312 && values.iter().all(|value| {
4313 value
4314 .get("_ai_agents_native_tool_call")
4315 .and_then(|marker| marker.as_bool())
4316 == Some(true)
4317 })
4318 }
4319 serde_json::Value::Object(map) => {
4320 map.get("_ai_agents_native_tool_call")
4321 .and_then(|marker| marker.as_bool())
4322 == Some(true)
4323 }
4324 _ => false,
4325 }
4326 }
4327
4328 fn tool_result_message(
4329 tool_call: &ToolCall,
4330 output: &str,
4331 native_tool_call: bool,
4332 ) -> ChatMessage {
4333 if !native_tool_call {
4334 return ChatMessage::function(&tool_call.name, output);
4335 }
4336 let output = serde_json::from_str::<serde_json::Value>(output)
4337 .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4338 ChatMessage::function(
4339 &tool_call.name,
4340 serde_json::json!({
4341 "_ai_agents_native_tool_result": true,
4342 "id": tool_call.id,
4343 "tool": tool_call.name,
4344 "output": output,
4345 })
4346 .to_string(),
4347 )
4348 }
4349
4350 fn parse_main_tool_calls(
4351 &self,
4352 content: &str,
4353 protocol: &MainToolProtocol,
4354 ) -> Option<Vec<ToolCall>> {
4355 if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4356 None
4357 } else {
4358 self.parse_tool_calls(content)
4359 }
4360 }
4361
4362 fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4363 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4365 if let Some(arr) = parsed.as_array() {
4367 let calls: Vec<ToolCall> = arr
4368 .iter()
4369 .filter_map(|v| self.extract_tool_call_from_value(v))
4370 .collect();
4371 if !calls.is_empty() {
4372 return Some(calls);
4373 }
4374 }
4375 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4377 return Some(vec![tool_call]);
4378 }
4379 }
4380
4381 if let Some(json_str) = self.extract_json_from_content(content)
4383 && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4384 {
4385 if let Some(arr) = parsed.as_array() {
4387 let calls: Vec<ToolCall> = arr
4388 .iter()
4389 .filter_map(|v| self.extract_tool_call_from_value(v))
4390 .collect();
4391 if !calls.is_empty() {
4392 return Some(calls);
4393 }
4394 }
4395 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4397 return Some(vec![tool_call]);
4398 }
4399 }
4400
4401 None
4402 }
4403
4404 fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4405 if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4406 let arguments = parsed
4407 .get("arguments")
4408 .cloned()
4409 .unwrap_or(serde_json::json!({}));
4410 return Some(ToolCall {
4411 id: parsed
4412 .get("id")
4413 .and_then(|value| value.as_str())
4414 .filter(|id| !id.is_empty())
4415 .map(str::to_string)
4416 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4417 name: tool_name.to_string(),
4418 arguments,
4419 });
4420 }
4421 None
4422 }
4423
4424 fn extract_json_from_content(&self, content: &str) -> Option<String> {
4426 if let Some(result) = self.extract_json_array_from_content(content) {
4428 return Some(result);
4429 }
4430 self.extract_json_object_from_content(content)
4431 }
4432
4433 fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4435 let start = content.find('[')?;
4436 let content_from_start = &content[start..];
4437
4438 let mut depth = 0;
4439 let mut end = 0;
4440 for (i, ch) in content_from_start.char_indices() {
4441 match ch {
4442 '[' => depth += 1,
4443 ']' => {
4444 depth -= 1;
4445 if depth == 0 {
4446 end = i + 1;
4447 break;
4448 }
4449 }
4450 _ => {}
4451 }
4452 }
4453
4454 if end > 0 {
4455 let json_str = &content_from_start[..end];
4456 if json_str.contains("\"tool\"") {
4458 return Some(json_str.to_string());
4459 }
4460 }
4461
4462 None
4463 }
4464
4465 fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4467 let start = content.find('{')?;
4468 let content_from_start = &content[start..];
4469
4470 let mut depth = 0;
4472 let mut end = 0;
4473 for (i, ch) in content_from_start.char_indices() {
4474 match ch {
4475 '{' => depth += 1,
4476 '}' => {
4477 depth -= 1;
4478 if depth == 0 {
4479 end = i + 1;
4480 break;
4481 }
4482 }
4483 _ => {}
4484 }
4485 }
4486
4487 if end > 0 {
4488 let json_str = &content_from_start[..end];
4489 if json_str.contains("\"tool\"") {
4491 return Some(json_str.to_string());
4492 }
4493 }
4494
4495 None
4496 }
4497
4498 #[allow(clippy::too_many_arguments)]
4502 fn record_from_parts(
4503 &self,
4504 request: &ToolExecutionRequest,
4505 canonical_id: String,
4506 executed_arguments: Value,
4507 started_at: chrono::DateTime<chrono::Utc>,
4508 start: Instant,
4509 executed: bool,
4510 success: bool,
4511 output: String,
4512 metadata: HashMap<String, Value>,
4513 policy: ToolPolicyDecisionRecord,
4514 approval: Option<ToolApprovalRecord>,
4515 timed_out: bool,
4516 output_truncated: bool,
4517 ) -> ToolExecutionRecord {
4518 let versions = ToolDecisionVersions {
4519 policy: self.active_tool_security().policy_version(),
4520 registry: self.tools.version(),
4521 runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4522 state: self
4523 .state_machine
4524 .as_ref()
4525 .map(|state_machine| state_machine.generation()),
4526 };
4527 self.record_from_parts_at(
4528 request,
4529 canonical_id,
4530 executed_arguments,
4531 started_at,
4532 start,
4533 executed,
4534 success,
4535 output,
4536 metadata,
4537 policy,
4538 approval,
4539 timed_out,
4540 output_truncated,
4541 versions,
4542 )
4543 }
4544
4545 #[allow(clippy::too_many_arguments)]
4547 fn record_from_parts_at(
4548 &self,
4549 request: &ToolExecutionRequest,
4550 canonical_id: String,
4551 executed_arguments: Value,
4552 started_at: chrono::DateTime<chrono::Utc>,
4553 start: Instant,
4554 executed: bool,
4555 success: bool,
4556 output: String,
4557 metadata: HashMap<String, Value>,
4558 policy: ToolPolicyDecisionRecord,
4559 approval: Option<ToolApprovalRecord>,
4560 timed_out: bool,
4561 output_truncated: bool,
4562 versions: ToolDecisionVersions,
4563 ) -> ToolExecutionRecord {
4564 ToolExecutionRecord {
4565 call_id: request.call_id.clone(),
4566 requested_name: request.requested_name.clone(),
4567 canonical_id,
4568 source: request.source.clone(),
4569 arguments: request.arguments.clone(),
4570 executed_arguments,
4571 policy_version: versions.policy,
4572 registry_version: versions.registry,
4573 runtime_config_version: versions.runtime_control,
4574 executed,
4575 success,
4576 output,
4577 metadata,
4578 policy,
4579 approval,
4580 started_at,
4581 duration_ms: start.elapsed().as_millis() as u64,
4582 timed_out,
4583 cancelled: false,
4584 cancellation_reason: None,
4585 output_truncated,
4586 }
4587 }
4588
4589 async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4591 let result = ToolResult {
4592 success: record.success,
4593 output: record.model_output_string(),
4594 metadata: if record.metadata.is_empty() {
4595 None
4596 } else {
4597 Some(record.metadata.clone())
4598 },
4599 };
4600 self.hooks
4601 .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4602 .await;
4603 self.hooks.on_tool_execution_record(record).await;
4604 self.record_tool_call(&record.canonical_id, record.model_output_value());
4605 if !record.success {
4606 self.hooks
4607 .on_error(&AgentError::Tool(record.output.clone()))
4608 .await;
4609 }
4610 }
4611
4612 async fn finish_tool_record_after_resource_guards(
4614 &self,
4615 resource_guards: ToolResourceGuards,
4616 record: &ToolExecutionRecord,
4617 ) {
4618 drop(resource_guards);
4619 self.finish_tool_record(record).await;
4620 }
4621
4622 fn validated_tool_timeout(timeout_ms: u64) -> Result<ValidatedToolTimeout> {
4626 if timeout_ms > MAX_TOOL_TIMEOUT_MS {
4627 return Err(AgentError::Config(format!(
4628 "effective tool timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
4629 )));
4630 }
4631 let timer = Duration::from_millis(timeout_ms);
4632 let deadline_delta = chrono::Duration::from_std(timer).map_err(|_| {
4633 AgentError::Config(format!(
4634 "effective tool timeout_ms cannot be represented as a UTC deadline: {timeout_ms}"
4635 ))
4636 })?;
4637 Ok(ValidatedToolTimeout {
4638 timer,
4639 deadline_delta,
4640 })
4641 }
4642
4643 fn effective_tool_limits(
4647 security_engine: &ToolSecurityEngine,
4648 canonical_id: &str,
4649 safety: &ToolSafetyMetadata,
4650 classification: &ToolCallClassification,
4651 recovery_timeout_ms: Option<u64>,
4652 ) -> Result<(ToolExecutionLimits, ValidatedToolTimeout)> {
4653 if let Some(timeout_ms) = classification.timeout_ms {
4654 Self::validated_tool_timeout(timeout_ms)?;
4655 }
4656 if let Some(timeout_ms) = recovery_timeout_ms {
4657 Self::validated_tool_timeout(timeout_ms)?;
4658 }
4659
4660 let mut limits = security_engine.effective_limits(canonical_id, safety, classification);
4661 if let Some(recovery_timeout_ms) = recovery_timeout_ms {
4662 limits.timeout_ms = Some(limits.timeout_ms.map_or(recovery_timeout_ms, |timeout_ms| {
4663 timeout_ms.min(recovery_timeout_ms)
4664 }));
4665 }
4666 let timeout_ms = limits
4667 .timeout_ms
4668 .unwrap_or_else(|| security_engine.get_tool_timeout(canonical_id));
4669 let timeout = Self::validated_tool_timeout(timeout_ms)?;
4670 Ok((limits, timeout))
4671 }
4672
4673 async fn execute_resolved_tool_once(
4675 &self,
4676 tool: Arc<dyn ai_agents_core::Tool>,
4677 args: Value,
4678 mut ctx: ToolExecutionContext,
4679 timeout: ValidatedToolTimeout,
4680 ) -> Result<(ToolResult, bool, bool, bool)> {
4681 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4682 return Ok((
4683 ToolResult::error("Tool execution cancelled by runtime control"),
4684 false,
4685 true,
4686 false,
4687 ));
4688 }
4689 ctx.deadline = Some(
4694 chrono::Utc::now()
4695 .checked_add_signed(timeout.deadline_delta)
4696 .ok_or_else(|| {
4697 AgentError::Config(
4698 "effective tool timeout_ms exceeds the current UTC deadline range"
4699 .to_string(),
4700 )
4701 })?,
4702 );
4703 let invoked = Arc::new(AtomicBool::new(false));
4707 let invoked_by_future = Arc::clone(&invoked);
4708 let actor_context = current_turn_actor_context();
4709 let future = async move {
4710 invoked_by_future.store(true, Ordering::SeqCst);
4711 if let Some(actor_context) = actor_context {
4712 scope_actor_context(actor_context, tool.execute(args, ctx)).await
4713 } else {
4714 tool.execute(args, ctx).await
4715 }
4716 };
4717 tokio::pin!(future);
4718 let timer = tokio::time::sleep(timeout.timer);
4719 tokio::pin!(timer);
4720 let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4721
4722 loop {
4723 tokio::select! {
4724 result = &mut future => return Ok((result, false, false, true)),
4725 _ = &mut timer => {
4726 return Ok((
4727 ToolResult::error("Tool execution timed out"),
4728 true,
4729 false,
4730 invoked.load(Ordering::SeqCst),
4731 ));
4732 }
4733 _ = cancel_tick.tick() => {
4734 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4735 return Ok((
4736 ToolResult::error("Tool execution cancelled by runtime control"),
4737 false,
4738 true,
4739 invoked.load(Ordering::SeqCst),
4740 ));
4741 }
4742 }
4743 }
4744 }
4745 }
4746
4747 fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4749 let Some(max_chars) = max_chars else {
4750 return (output, false);
4751 };
4752 let mut chars = output.chars();
4753 let truncated: String = chars.by_ref().take(max_chars).collect();
4754 if chars.next().is_some() {
4755 (truncated, true)
4756 } else {
4757 (output, false)
4758 }
4759 }
4760
4761 async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4763 let locks = {
4764 let mut table = self.resource_locks.write();
4765 table.retain(|_, lock| lock.strong_count() > 0);
4766 keys.iter()
4767 .map(|key| {
4768 if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4769 lock
4770 } else {
4771 let lock = Arc::new(tokio::sync::Mutex::new(()));
4772 table.insert(key.clone(), Arc::downgrade(&lock));
4773 lock
4774 }
4775 })
4776 .collect::<Vec<_>>()
4777 };
4778 let mut resource_guards = ToolResourceGuards {
4779 guards: Vec::with_capacity(locks.len()),
4780 locks: Arc::clone(&self.resource_locks),
4781 };
4782 let mut locks = locks.into_iter();
4783 while let Some(lock) = locks.next() {
4784 let mut lock = Box::pin(lock.lock_owned());
4785 loop {
4786 tokio::select! {
4787 guard = &mut lock => {
4788 resource_guards.guards.push(guard);
4789 break;
4790 }
4791 _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4792 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4793 drop(lock);
4794 drop(locks);
4795 drop(resource_guards);
4796 return None;
4797 }
4798 }
4799 }
4800 }
4801 }
4802 Some(resource_guards)
4803 }
4804
4805 async fn run_tool_with_retries(
4809 &self,
4810 canonical_id: &str,
4811 tool: Arc<dyn ai_agents_core::Tool>,
4812 args: Value,
4813 ctx: ToolExecutionContext,
4814 timeout: ValidatedToolTimeout,
4815 max_retries: u32,
4816 ) -> Result<(ToolResult, bool, bool, bool)> {
4817 let max_retries = if ctx.classification.safely_retryable {
4818 max_retries
4819 } else {
4820 0
4821 };
4822 let mut attempts = 0;
4823 let mut invoked = false;
4824 loop {
4825 let (result, timed_out, cancelled, attempt_invoked) = self
4826 .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout)
4827 .await?;
4828 invoked |= attempt_invoked;
4829 if result.success || timed_out || cancelled || attempts >= max_retries {
4830 return Ok((result, timed_out, cancelled, invoked));
4831 }
4832 attempts += 1;
4833 warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4834 }
4835 }
4836
4837 fn host_tool_unavailability(&self, canonical_id: &str) -> Option<(&'static str, &'static str)> {
4839 match canonical_id {
4840 "command" if !self.tools.command_runner_available() => Some((
4841 "Command runner is unavailable",
4842 "command runner is unavailable",
4843 )),
4844 "diagnostics" if !self.tools.diagnostics_available() => Some((
4845 "Diagnostics provider is unavailable",
4846 "diagnostics provider is unavailable",
4847 )),
4848 "web_search" if !self.tools.web_search_available() => Some((
4849 "Web search provider is unavailable",
4850 "web search provider is unavailable",
4851 )),
4852 _ => None,
4853 }
4854 }
4855
4856 fn execute_tool_record(
4858 &self,
4859 request: ToolExecutionRequest,
4860 ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4861 Box::pin(self.execute_tool_record_inner(request, ToolFallbackState::default()))
4862 }
4863
4864 async fn execute_tool_record_inner(
4868 &self,
4869 request: ToolExecutionRequest,
4870 fallback_state: ToolFallbackState,
4871 ) -> Result<ToolExecutionRecord> {
4872 let started_at = chrono::Utc::now();
4873 let start = Instant::now();
4874 info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4875
4876 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4877 let record = self.record_from_parts(
4878 &request,
4879 request.requested_name.clone(),
4880 request.arguments.clone(),
4881 started_at,
4882 start,
4883 false,
4884 false,
4885 "Tool execution is disabled by runtime control".to_string(),
4886 HashMap::new(),
4887 ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4888 None,
4889 false,
4890 false,
4891 );
4892 self.finish_tool_record(&record).await;
4893 return Ok(record);
4894 }
4895
4896 let Some(resolved) = self.tools.resolve(&request.requested_name) else {
4897 let record = self.record_from_parts(
4898 &request,
4899 request.requested_name.clone(),
4900 request.arguments.clone(),
4901 started_at,
4902 start,
4903 false,
4904 false,
4905 format!("Tool '{}' is unavailable", request.requested_name),
4906 HashMap::new(),
4907 ToolPolicyDecisionRecord::unavailable(format!(
4908 "Tool '{}' is not registered",
4909 request.requested_name
4910 )),
4911 None,
4912 false,
4913 false,
4914 );
4915 self.finish_tool_record(&record).await;
4916 return Ok(record);
4917 };
4918
4919 let canonical_id = resolved.identity.canonical_id.clone();
4920
4921 let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4922 if !initial_scope_snapshot
4923 .tool_ids
4924 .iter()
4925 .any(|id| id == &canonical_id)
4926 {
4927 let record = self.record_from_parts(
4928 &request,
4929 canonical_id.clone(),
4930 request.arguments.clone(),
4931 started_at,
4932 start,
4933 false,
4934 false,
4935 format!(
4936 "Tool '{}' is not available in the current scope",
4937 canonical_id
4938 ),
4939 HashMap::new(),
4940 ToolPolicyDecisionRecord::deny(format!(
4941 "Tool '{}' is not granted by the current top-level and state tool scope",
4942 canonical_id
4943 )),
4944 None,
4945 false,
4946 false,
4947 );
4948 self.finish_tool_record(&record).await;
4949 return Ok(record);
4950 }
4951
4952 let approval_control_snapshot = self.runtime_safety_snapshot();
4953 let security_engine = approval_control_snapshot.tool_security.clone();
4954 if let Some(reason) = fallback_state.rejection_reason(&canonical_id) {
4955 let mut metadata = HashMap::new();
4956 metadata.insert(
4957 "fallback_chain".to_string(),
4958 serde_json::to_value(&fallback_state.visited_canonical_ids).unwrap_or(Value::Null),
4959 );
4960 let record = self.record_from_parts(
4961 &request,
4962 canonical_id,
4963 request.arguments.clone(),
4964 started_at,
4965 start,
4966 false,
4967 false,
4968 format!("Denied: {reason}"),
4969 metadata,
4970 ToolPolicyDecisionRecord::deny(reason),
4971 None,
4972 false,
4973 false,
4974 );
4975 self.finish_tool_record(&record).await;
4976 return Ok(record);
4977 }
4978 let admitted_canonical_id = canonical_id.clone();
4979 let fallback_state = fallback_state.with_current(canonical_id.clone());
4980 let bindings = resolved.tool.policy_bindings();
4981 let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4982 &canonical_id,
4983 &request.arguments,
4984 &bindings,
4985 );
4986 let mut metadata = HashMap::new();
4987 let safety = resolved.tool.safety_metadata();
4988 let classification = resolved.tool.classify_call(&executed_arguments);
4989 let initial_recovery_timeout_ms = self.recovery_manager.get_tool_timeout(&canonical_id);
4990 let (limits, _) = Self::effective_tool_limits(
4991 &security_engine,
4992 &canonical_id,
4993 &safety,
4994 &classification,
4995 initial_recovery_timeout_ms,
4996 )?;
4997 self.hooks
4998 .on_tool_start(&canonical_id, &executed_arguments)
4999 .await;
5000 metadata.insert(
5001 "classification".to_string(),
5002 serde_json::to_value(&classification).unwrap_or(Value::Null),
5003 );
5004 metadata.insert(
5005 "effective_limits".to_string(),
5006 serde_json::to_value(&limits).unwrap_or(Value::Null),
5007 );
5008 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5009 if !policy_snapshot.is_null() {
5010 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5011 }
5012
5013 let mut approval_record = Some(ToolApprovalRecord {
5014 status: ToolApprovalStatus::NotRequired,
5015 reason: None,
5016 modified_arguments: None,
5017 });
5018
5019 let mut security_result = security_engine
5020 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5021 .await?;
5022 if (security_result.is_allowed()
5027 || matches!(
5028 &security_result,
5029 SecurityCheckResult::RequireConfirmation { .. }
5030 ))
5031 && let Some((output, reason)) = self.host_tool_unavailability(&canonical_id)
5032 {
5033 let record = self.record_from_parts(
5034 &request,
5035 canonical_id,
5036 executed_arguments,
5037 started_at,
5038 start,
5039 false,
5040 false,
5041 output.to_string(),
5042 metadata,
5043 ToolPolicyDecisionRecord::unavailable(reason),
5044 Some(ToolApprovalRecord {
5045 status: ToolApprovalStatus::Unavailable,
5046 reason: Some(reason.to_string()),
5047 modified_arguments: None,
5048 }),
5049 false,
5050 false,
5051 );
5052 self.finish_tool_record(&record).await;
5053 return Ok(record);
5054 }
5055 match &security_result {
5056 SecurityCheckResult::Allow => {}
5057 SecurityCheckResult::Warn { message } => {
5058 warn!(tool = %canonical_id, message = %message, "Tool security warning");
5059 }
5060 SecurityCheckResult::Block { reason } => {
5061 let record = self.record_from_parts(
5062 &request,
5063 canonical_id,
5064 executed_arguments,
5065 started_at,
5066 start,
5067 false,
5068 false,
5069 format!("Denied: {}", reason),
5070 metadata,
5071 ToolPolicyDecisionRecord::deny(reason.clone()),
5072 approval_record,
5073 false,
5074 false,
5075 );
5076 self.finish_tool_record(&record).await;
5077 return Ok(record);
5078 }
5079 SecurityCheckResult::Unavailable { reason } => {
5080 let record = self.record_from_parts(
5081 &request,
5082 canonical_id,
5083 executed_arguments,
5084 started_at,
5085 start,
5086 false,
5087 false,
5088 format!("Unavailable: {}", reason),
5089 metadata,
5090 ToolPolicyDecisionRecord::unavailable(reason.clone()),
5091 approval_record,
5092 false,
5093 false,
5094 );
5095 self.finish_tool_record(&record).await;
5096 return Ok(record);
5097 }
5098 SecurityCheckResult::RequireConfirmation { message } => {
5099 if self.hitl_engine.is_none() {
5100 approval_record = Some(ToolApprovalRecord {
5101 status: ToolApprovalStatus::Unavailable,
5102 reason: Some("No HITL engine configured".to_string()),
5103 modified_arguments: None,
5104 });
5105 let record = self.record_from_parts(
5106 &request,
5107 canonical_id,
5108 executed_arguments,
5109 started_at,
5110 start,
5111 false,
5112 false,
5113 format!("Approval unavailable: {}", message),
5114 metadata,
5115 ToolPolicyDecisionRecord::approval(message.clone()),
5116 approval_record,
5117 false,
5118 false,
5119 );
5120 self.finish_tool_record(&record).await;
5121 return Ok(record);
5122 }
5123
5124 let check_result = HITLCheckResult::required(
5125 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
5126 HashMap::new(),
5127 message.clone(),
5128 None,
5129 );
5130 match self.request_hitl_approval(check_result).await? {
5131 ApprovalResult::Approved => {
5132 merge_approved_record(&mut approval_record);
5133 }
5134 ApprovalResult::Modified { changes } => {
5135 if let Some(obj) = executed_arguments.as_object_mut() {
5136 for (key, value) in changes {
5137 obj.insert(key, value);
5138 }
5139 }
5140 security_result = security_engine
5141 .validate_tool_execution_with_bindings(
5142 &canonical_id,
5143 &executed_arguments,
5144 &bindings,
5145 )
5146 .await?;
5147 if !matches!(
5148 security_result,
5149 SecurityCheckResult::Allow
5150 | SecurityCheckResult::Warn { .. }
5151 | SecurityCheckResult::RequireConfirmation { .. }
5152 ) {
5153 let reason = security_result
5154 .reason()
5155 .unwrap_or("modified arguments failed policy")
5156 .to_string();
5157 let record = self.record_from_parts(
5158 &request,
5159 canonical_id,
5160 executed_arguments.clone(),
5161 started_at,
5162 start,
5163 false,
5164 false,
5165 reason.clone(),
5166 metadata,
5167 ToolPolicyDecisionRecord::deny(reason),
5168 Some(ToolApprovalRecord {
5169 status: ToolApprovalStatus::Modified,
5170 reason: None,
5171 modified_arguments: Some(executed_arguments),
5172 }),
5173 false,
5174 false,
5175 );
5176 self.finish_tool_record(&record).await;
5177 return Ok(record);
5178 }
5179 approval_record = Some(ToolApprovalRecord {
5180 status: ToolApprovalStatus::Modified,
5181 reason: None,
5182 modified_arguments: Some(executed_arguments.clone()),
5183 });
5184 }
5185 ApprovalResult::Rejected { reason } => {
5186 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5187 approval_record = Some(ToolApprovalRecord {
5188 status: ToolApprovalStatus::Rejected,
5189 reason: Some(reason.clone()),
5190 modified_arguments: None,
5191 });
5192 let record = self.record_from_parts(
5193 &request,
5194 canonical_id,
5195 executed_arguments,
5196 started_at,
5197 start,
5198 false,
5199 false,
5200 format!("Approval rejected: {}", reason),
5201 metadata,
5202 ToolPolicyDecisionRecord::approval(reason),
5203 approval_record,
5204 false,
5205 false,
5206 );
5207 self.finish_tool_record(&record).await;
5208 return Ok(record);
5209 }
5210 ApprovalResult::Timeout => {
5211 approval_record = Some(ToolApprovalRecord {
5212 status: ToolApprovalStatus::Timeout,
5213 reason: Some("approval timeout".to_string()),
5214 modified_arguments: None,
5215 });
5216 let record = self.record_from_parts(
5217 &request,
5218 canonical_id,
5219 executed_arguments,
5220 started_at,
5221 start,
5222 false,
5223 false,
5224 "Approval timed out".to_string(),
5225 metadata,
5226 ToolPolicyDecisionRecord::approval("approval timeout"),
5227 approval_record,
5228 false,
5229 false,
5230 );
5231 self.finish_tool_record(&record).await;
5232 return Ok(record);
5233 }
5234 }
5235 }
5236 }
5237
5238 if approval_record
5239 .as_ref()
5240 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
5241 && let Some(message) =
5242 security_engine.classification_approval_message(&canonical_id, &classification)
5243 {
5244 if self.hitl_engine.is_none() {
5245 approval_record = Some(ToolApprovalRecord {
5246 status: ToolApprovalStatus::Unavailable,
5247 reason: Some("No HITL engine configured".to_string()),
5248 modified_arguments: None,
5249 });
5250 let record = self.record_from_parts(
5251 &request,
5252 canonical_id,
5253 executed_arguments,
5254 started_at,
5255 start,
5256 false,
5257 false,
5258 format!("Approval unavailable: {}", message),
5259 metadata,
5260 ToolPolicyDecisionRecord::approval(message),
5261 approval_record,
5262 false,
5263 false,
5264 );
5265 self.finish_tool_record(&record).await;
5266 return Ok(record);
5267 }
5268 let check_result = HITLCheckResult::required(
5269 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
5270 HashMap::new(),
5271 message.clone(),
5272 None,
5273 );
5274 match self.request_hitl_approval(check_result).await? {
5275 ApprovalResult::Approved => {
5276 merge_approved_record(&mut approval_record);
5277 }
5278 ApprovalResult::Modified { changes } => {
5279 if let Some(obj) = executed_arguments.as_object_mut() {
5280 for (key, value) in changes {
5281 obj.insert(key, value);
5282 }
5283 }
5284 let modified_security = security_engine
5285 .validate_tool_execution_with_bindings(
5286 &canonical_id,
5287 &executed_arguments,
5288 &bindings,
5289 )
5290 .await?;
5291 if !matches!(
5292 modified_security,
5293 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5294 ) {
5295 let reason = modified_security
5296 .reason()
5297 .unwrap_or("modified arguments failed policy")
5298 .to_string();
5299 let record = self.record_from_parts(
5300 &request,
5301 canonical_id,
5302 executed_arguments.clone(),
5303 started_at,
5304 start,
5305 false,
5306 false,
5307 reason.clone(),
5308 metadata,
5309 ToolPolicyDecisionRecord::deny(reason),
5310 Some(ToolApprovalRecord {
5311 status: ToolApprovalStatus::Modified,
5312 reason: None,
5313 modified_arguments: Some(executed_arguments),
5314 }),
5315 false,
5316 false,
5317 );
5318 self.finish_tool_record(&record).await;
5319 return Ok(record);
5320 }
5321 approval_record = Some(ToolApprovalRecord {
5322 status: ToolApprovalStatus::Modified,
5323 reason: None,
5324 modified_arguments: Some(executed_arguments.clone()),
5325 });
5326 }
5327 ApprovalResult::Rejected { reason } => {
5328 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5329 let record = self.record_from_parts(
5330 &request,
5331 canonical_id,
5332 executed_arguments,
5333 started_at,
5334 start,
5335 false,
5336 false,
5337 format!("Approval rejected: {}", reason),
5338 metadata,
5339 ToolPolicyDecisionRecord::approval(reason.clone()),
5340 Some(ToolApprovalRecord {
5341 status: ToolApprovalStatus::Rejected,
5342 reason: Some(reason),
5343 modified_arguments: None,
5344 }),
5345 false,
5346 false,
5347 );
5348 self.finish_tool_record(&record).await;
5349 return Ok(record);
5350 }
5351 ApprovalResult::Timeout => {
5352 let record = self.record_from_parts(
5353 &request,
5354 canonical_id,
5355 executed_arguments,
5356 started_at,
5357 start,
5358 false,
5359 false,
5360 "Approval timed out".to_string(),
5361 metadata,
5362 ToolPolicyDecisionRecord::approval("approval timeout"),
5363 Some(ToolApprovalRecord {
5364 status: ToolApprovalStatus::Timeout,
5365 reason: Some("approval timeout".to_string()),
5366 modified_arguments: None,
5367 }),
5368 false,
5369 false,
5370 );
5371 self.finish_tool_record(&record).await;
5372 return Ok(record);
5373 }
5374 }
5375 }
5376
5377 let hitl_lang_ctx = self.build_hitl_language_context();
5378 if let Some(ref hitl_engine) = self.hitl_engine {
5379 let check_result = self
5380 .observe_purpose(
5381 ObservationPurpose::HitlLocalization,
5382 hitl_engine.check_tool_with_localization(
5383 &canonical_id,
5384 &executed_arguments,
5385 &hitl_lang_ctx,
5386 self.approval_handler.as_ref(),
5387 Some(&self.llm_registry),
5388 ),
5389 )
5390 .await?;
5391 if check_result.is_required() {
5392 match self.request_hitl_approval(check_result).await? {
5393 ApprovalResult::Approved => {
5394 merge_approved_record(&mut approval_record);
5395 }
5396 ApprovalResult::Modified { changes } => {
5397 if let Some(obj) = executed_arguments.as_object_mut() {
5398 for (key, value) in changes {
5399 obj.insert(key, value);
5400 }
5401 }
5402 let modified_security = security_engine
5403 .validate_tool_execution_with_bindings(
5404 &canonical_id,
5405 &executed_arguments,
5406 &bindings,
5407 )
5408 .await?;
5409 if !matches!(
5410 modified_security,
5411 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5412 ) {
5413 let reason = modified_security
5414 .reason()
5415 .unwrap_or("modified arguments failed policy")
5416 .to_string();
5417 let record = self.record_from_parts(
5418 &request,
5419 canonical_id,
5420 executed_arguments.clone(),
5421 started_at,
5422 start,
5423 false,
5424 false,
5425 reason.clone(),
5426 metadata,
5427 ToolPolicyDecisionRecord::deny(reason),
5428 Some(ToolApprovalRecord {
5429 status: ToolApprovalStatus::Modified,
5430 reason: None,
5431 modified_arguments: Some(executed_arguments),
5432 }),
5433 false,
5434 false,
5435 );
5436 self.finish_tool_record(&record).await;
5437 return Ok(record);
5438 }
5439 approval_record = Some(ToolApprovalRecord {
5440 status: ToolApprovalStatus::Modified,
5441 reason: None,
5442 modified_arguments: Some(executed_arguments.clone()),
5443 });
5444 }
5445 ApprovalResult::Rejected { reason } => {
5446 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5447 let record = self.record_from_parts(
5448 &request,
5449 canonical_id,
5450 executed_arguments,
5451 started_at,
5452 start,
5453 false,
5454 false,
5455 format!("Approval rejected: {}", reason),
5456 metadata,
5457 ToolPolicyDecisionRecord::approval(reason.clone()),
5458 Some(ToolApprovalRecord {
5459 status: ToolApprovalStatus::Rejected,
5460 reason: Some(reason),
5461 modified_arguments: None,
5462 }),
5463 false,
5464 false,
5465 );
5466 self.finish_tool_record(&record).await;
5467 return Ok(record);
5468 }
5469 ApprovalResult::Timeout => {
5470 let record = self.record_from_parts(
5471 &request,
5472 canonical_id,
5473 executed_arguments,
5474 started_at,
5475 start,
5476 false,
5477 false,
5478 "Approval timed out".to_string(),
5479 metadata,
5480 ToolPolicyDecisionRecord::approval("approval timeout"),
5481 Some(ToolApprovalRecord {
5482 status: ToolApprovalStatus::Timeout,
5483 reason: Some("approval timeout".to_string()),
5484 modified_arguments: None,
5485 }),
5486 false,
5487 false,
5488 );
5489 self.finish_tool_record(&record).await;
5490 return Ok(record);
5491 }
5492 }
5493 }
5494
5495 let condition_check = self
5496 .observe_purpose(
5497 ObservationPurpose::HitlLocalization,
5498 hitl_engine.check_conditions_with_localization(
5499 &executed_arguments,
5500 &hitl_lang_ctx,
5501 self.approval_handler.as_ref(),
5502 Some(&self.llm_registry),
5503 ),
5504 )
5505 .await?;
5506 if condition_check.is_required() {
5507 match self.request_hitl_approval(condition_check).await? {
5508 ApprovalResult::Approved => {
5509 merge_approved_record(&mut approval_record);
5510 }
5511 ApprovalResult::Modified { changes } => {
5512 if let Some(obj) = executed_arguments.as_object_mut() {
5513 for (key, value) in changes {
5514 obj.insert(key, value);
5515 }
5516 }
5517 let modified_security = security_engine
5518 .validate_tool_execution_with_bindings(
5519 &canonical_id,
5520 &executed_arguments,
5521 &bindings,
5522 )
5523 .await?;
5524 if !matches!(
5525 modified_security,
5526 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5527 ) {
5528 let reason = modified_security
5529 .reason()
5530 .unwrap_or("modified arguments failed policy")
5531 .to_string();
5532 let record = self.record_from_parts(
5533 &request,
5534 canonical_id,
5535 executed_arguments,
5536 started_at,
5537 start,
5538 false,
5539 false,
5540 reason.clone(),
5541 metadata,
5542 ToolPolicyDecisionRecord::deny(reason),
5543 approval_record,
5544 false,
5545 false,
5546 );
5547 self.finish_tool_record(&record).await;
5548 return Ok(record);
5549 }
5550 approval_record = Some(ToolApprovalRecord {
5551 status: ToolApprovalStatus::Modified,
5552 reason: None,
5553 modified_arguments: Some(executed_arguments.clone()),
5554 });
5555 }
5556 ApprovalResult::Rejected { reason } => {
5557 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5558 let record = self.record_from_parts(
5559 &request,
5560 canonical_id,
5561 executed_arguments,
5562 started_at,
5563 start,
5564 false,
5565 false,
5566 format!("Approval rejected: {}", reason),
5567 metadata,
5568 ToolPolicyDecisionRecord::approval(reason.clone()),
5569 Some(ToolApprovalRecord {
5570 status: ToolApprovalStatus::Rejected,
5571 reason: Some(reason),
5572 modified_arguments: None,
5573 }),
5574 false,
5575 false,
5576 );
5577 self.finish_tool_record(&record).await;
5578 return Ok(record);
5579 }
5580 ApprovalResult::Timeout => {
5581 let record = self.record_from_parts(
5582 &request,
5583 canonical_id,
5584 executed_arguments,
5585 started_at,
5586 start,
5587 false,
5588 false,
5589 "Approval timed out".to_string(),
5590 metadata,
5591 ToolPolicyDecisionRecord::approval("approval timeout"),
5592 Some(ToolApprovalRecord {
5593 status: ToolApprovalStatus::Timeout,
5594 reason: Some("approval timeout".to_string()),
5595 modified_arguments: None,
5596 }),
5597 false,
5598 false,
5599 );
5600 self.finish_tool_record(&record).await;
5601 return Ok(record);
5602 }
5603 }
5604 }
5605 }
5606
5607 executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5612 &canonical_id,
5613 &executed_arguments,
5614 &bindings,
5615 );
5616 if let Some(record) = approval_record.as_mut()
5617 && matches!(record.status, ToolApprovalStatus::Modified)
5618 {
5619 record.modified_arguments = Some(executed_arguments.clone());
5620 }
5621 let binding_security_result = security_engine
5622 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5623 .await?;
5624 let approval_confirmation_required = matches!(
5625 binding_security_result,
5626 SecurityCheckResult::RequireConfirmation { .. }
5627 ) || security_engine
5628 .classification_approval_message(
5629 &canonical_id,
5630 &resolved.tool.classify_call(&executed_arguments),
5631 )
5632 .is_some();
5633 let approval_binding = approval_record.as_ref().and_then(|record| {
5634 matches!(
5635 record.status,
5636 ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5637 )
5638 .then(|| ToolApprovalBinding {
5639 canonical_id: canonical_id.clone(),
5640 arguments: executed_arguments.clone(),
5641 confirmation_required: approval_confirmation_required,
5642 policy_version: security_engine.policy_version(),
5643 runtime_control_version: approval_control_snapshot.version,
5644 state_generation: initial_scope_snapshot.state_generation,
5645 reviewed_tool: Arc::clone(&resolved.tool),
5646 })
5647 });
5648
5649 let control_snapshot = self.runtime_safety_snapshot();
5654 let resolved = self.tools.resolve(&request.requested_name);
5655 let registry_version = self.tools.version();
5656 let mut versions = ToolDecisionVersions {
5657 policy: control_snapshot.tool_security.policy_version(),
5658 registry: registry_version,
5659 runtime_control: control_snapshot.version,
5660 state: None,
5661 };
5662 metadata.insert(
5663 "runtime_scope_snapshot".to_string(),
5664 serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5665 );
5666 let resolved = match resolved {
5667 Some(resolved) => resolved,
5668 None => {
5669 let reason = format!(
5670 "Tool '{}' became unavailable after approval",
5671 request.requested_name
5672 );
5673 let record = self.record_from_parts_at(
5674 &request,
5675 request.requested_name.clone(),
5676 executed_arguments,
5677 started_at,
5678 start,
5679 false,
5680 false,
5681 reason.clone(),
5682 metadata,
5683 ToolPolicyDecisionRecord::unavailable(reason),
5684 approval_record,
5685 false,
5686 false,
5687 versions,
5688 );
5689 self.finish_tool_record(&record).await;
5690 return Ok(record);
5691 }
5692 };
5693
5694 let canonical_id = resolved.identity.canonical_id.clone();
5695 if let Some(reason) =
5696 fallback_state.final_rejection_reason(&admitted_canonical_id, &canonical_id)
5697 {
5698 metadata.insert(
5702 "fallback_chain".to_string(),
5703 serde_json::to_value(&fallback_state.visited_canonical_ids).unwrap_or(Value::Null),
5704 );
5705 metadata.insert(
5706 "final_resolved_canonical_id".to_string(),
5707 Value::String(canonical_id),
5708 );
5709 let record = self.record_from_parts_at(
5710 &request,
5711 admitted_canonical_id,
5712 executed_arguments,
5713 started_at,
5714 start,
5715 false,
5716 false,
5717 format!("Denied: {reason}"),
5718 metadata,
5719 ToolPolicyDecisionRecord::deny(reason),
5720 approval_record,
5721 false,
5722 false,
5723 versions,
5724 );
5725 self.finish_tool_record(&record).await;
5726 return Ok(record);
5727 }
5728 let bindings = resolved.tool.policy_bindings();
5729 let final_arguments = control_snapshot
5730 .tool_security
5731 .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5732 if let Some(record) = approval_record.as_mut()
5733 && matches!(record.status, ToolApprovalStatus::Modified)
5734 {
5735 record.modified_arguments = Some(final_arguments.clone());
5736 }
5737 let classification = resolved.tool.classify_call(&final_arguments);
5738 let safety = resolved.tool.safety_metadata();
5739 let security_engine = control_snapshot.tool_security;
5740 let tool_config = self.recovery_manager.get_tool_config(&canonical_id).clone();
5741 let recovery_timeout_ms = self.recovery_manager.get_tool_timeout(&canonical_id);
5742 metadata.insert(
5743 "classification".to_string(),
5744 serde_json::to_value(&classification).unwrap_or(Value::Null),
5745 );
5746 let (limits, timeout) = match Self::effective_tool_limits(
5750 &security_engine,
5751 &canonical_id,
5752 &safety,
5753 &classification,
5754 recovery_timeout_ms,
5755 ) {
5756 Ok(effective) => effective,
5757 Err(error) => {
5758 let reason = error.to_string();
5759 metadata.insert(
5760 "configuration_error".to_string(),
5761 Value::String(reason.clone()),
5762 );
5763 let record = self.record_from_parts_at(
5764 &request,
5765 canonical_id,
5766 final_arguments,
5767 started_at,
5768 start,
5769 false,
5770 false,
5771 format!("Denied: {reason}"),
5772 metadata,
5773 ToolPolicyDecisionRecord::deny(reason),
5774 approval_record,
5775 false,
5776 false,
5777 versions,
5778 );
5779 self.finish_tool_record(&record).await;
5780 return Ok(record);
5781 }
5782 };
5783 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5784 let resource_lock_keys =
5785 tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5786 metadata.insert(
5787 "effective_limits".to_string(),
5788 serde_json::to_value(&limits).unwrap_or(Value::Null),
5789 );
5790 metadata.insert(
5791 "resource_lock_keys".to_string(),
5792 serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5793 );
5794 if policy_snapshot.is_null() {
5795 metadata.remove("policy_snapshot");
5796 } else {
5797 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5798 }
5799
5800 let final_denial = |canonical_id: String,
5801 output: String,
5802 policy: ToolPolicyDecisionRecord,
5803 metadata: HashMap<String, Value>,
5804 decision_versions: ToolDecisionVersions| {
5805 self.record_from_parts_at(
5806 &request,
5807 canonical_id,
5808 final_arguments.clone(),
5809 started_at,
5810 start,
5811 false,
5812 false,
5813 output,
5814 metadata,
5815 policy,
5816 approval_record.clone(),
5817 false,
5818 false,
5819 decision_versions,
5820 )
5821 };
5822
5823 if control_snapshot.emergency_deny {
5824 let reason = "Tool execution is disabled by runtime control".to_string();
5825 let record = final_denial(
5826 canonical_id,
5827 reason.clone(),
5828 ToolPolicyDecisionRecord::deny(reason),
5829 metadata,
5830 versions,
5831 );
5832 self.finish_tool_record(&record).await;
5833 return Ok(record);
5834 }
5835
5836 let available_snapshot = self
5841 .get_available_tool_ids_snapshot_for_scope(
5842 control_snapshot.tool_scope_override.as_deref(),
5843 )
5844 .await?;
5845 versions.state = available_snapshot.state_generation;
5846 metadata.insert(
5847 "available_tool_ids_snapshot".to_string(),
5848 serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5849 );
5850 metadata.insert(
5851 "state_generation_snapshot".to_string(),
5852 serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5853 );
5854 if !available_snapshot
5855 .tool_ids
5856 .iter()
5857 .any(|tool_id| tool_id == &canonical_id)
5858 {
5859 let reason = format!(
5860 "Tool '{}' is not available in the final runtime scope",
5861 canonical_id
5862 );
5863 let record = final_denial(
5864 canonical_id,
5865 reason.clone(),
5866 ToolPolicyDecisionRecord::deny(reason),
5867 metadata,
5868 versions,
5869 );
5870 self.finish_tool_record(&record).await;
5871 return Ok(record);
5872 }
5873
5874 let final_security_result = security_engine
5879 .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5880 .await?;
5881 match &final_security_result {
5882 SecurityCheckResult::Block { reason } => {
5883 let record = final_denial(
5884 canonical_id,
5885 format!("Denied: {}", reason),
5886 ToolPolicyDecisionRecord::deny(reason.clone()),
5887 metadata,
5888 versions,
5889 );
5890 self.finish_tool_record(&record).await;
5891 return Ok(record);
5892 }
5893 SecurityCheckResult::Unavailable { reason } => {
5894 let record = final_denial(
5895 canonical_id,
5896 format!("Unavailable: {}", reason),
5897 ToolPolicyDecisionRecord::unavailable(reason.clone()),
5898 metadata,
5899 versions,
5900 );
5901 self.finish_tool_record(&record).await;
5902 return Ok(record);
5903 }
5904 SecurityCheckResult::Warn { message } => {
5905 warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5906 }
5907 SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5908 }
5909 let final_confirmation_required = matches!(
5910 final_security_result,
5911 SecurityCheckResult::RequireConfirmation { .. }
5912 ) || security_engine
5913 .classification_approval_message(&canonical_id, &classification)
5914 .is_some();
5915 let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5916 binding.is_stale(
5917 &canonical_id,
5918 &final_arguments,
5919 final_confirmation_required,
5920 versions,
5921 &resolved.tool,
5922 )
5923 });
5924 if stale_approval {
5925 let reason = "Approval became stale before final admission".to_string();
5926 let record = final_denial(
5927 canonical_id,
5928 reason.clone(),
5929 ToolPolicyDecisionRecord::deny(reason),
5930 metadata,
5931 versions,
5932 );
5933 self.finish_tool_record(&record).await;
5934 return Ok(record);
5935 }
5936 if final_confirmation_required && approval_binding.is_none() {
5937 let reason = "Final policy requires fresh approval".to_string();
5938 let record = final_denial(
5939 canonical_id,
5940 reason.clone(),
5941 ToolPolicyDecisionRecord::approval(reason),
5942 metadata,
5943 versions,
5944 );
5945 self.finish_tool_record(&record).await;
5946 return Ok(record);
5947 }
5948
5949 if let Some((_, reason)) = self.host_tool_unavailability(&canonical_id) {
5950 let record = final_denial(
5951 canonical_id,
5952 reason.to_string(),
5953 ToolPolicyDecisionRecord::unavailable(reason),
5954 metadata,
5955 versions,
5956 );
5957 self.finish_tool_record(&record).await;
5958 return Ok(record);
5959 }
5960
5961 let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5966 else {
5967 let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5971 let mut record = final_denial(
5972 canonical_id,
5973 reason.clone(),
5974 ToolPolicyDecisionRecord::deny(reason),
5975 metadata,
5976 versions,
5977 );
5978 record.cancelled = true;
5979 record.cancellation_reason = Some("runtime control cancellation".to_string());
5980 self.finish_tool_record(&record).await;
5981 return Ok(record);
5982 };
5983
5984 let admission = self.admit_tool_execution(
5989 versions.runtime_control,
5990 versions.policy,
5991 versions.state,
5992 &canonical_id,
5993 );
5994 if !matches!(admission, SecurityCheckResult::Allow) {
5995 let latest_control = self.runtime_safety_snapshot();
5996 let reason = admission
5997 .reason()
5998 .unwrap_or("tool admission was denied")
5999 .to_string();
6000 let policy = if admission.is_unavailable() {
6001 ToolPolicyDecisionRecord::unavailable(reason.clone())
6002 } else {
6003 ToolPolicyDecisionRecord::deny(reason.clone())
6004 };
6005 let record = self.record_from_parts_at(
6006 &request,
6007 canonical_id,
6008 final_arguments,
6009 started_at,
6010 start,
6011 false,
6012 false,
6013 reason,
6014 metadata,
6015 policy,
6016 approval_record,
6017 false,
6018 false,
6019 ToolDecisionVersions {
6020 policy: latest_control.tool_security.policy_version(),
6021 registry: versions.registry,
6022 runtime_control: latest_control.version,
6023 state: self
6024 .state_machine
6025 .as_ref()
6026 .map(|state_machine| state_machine.generation()),
6027 },
6028 );
6029 self.finish_tool_record_after_resource_guards(resource_guards, &record)
6030 .await;
6031 return Ok(record);
6032 }
6033 let executed_arguments = final_arguments;
6034
6035 let turn_actor = current_turn_actor_context();
6036 let actor = ToolActorContext {
6037 actor_id: turn_actor
6038 .as_ref()
6039 .and_then(|context| context.effective_actor_id().map(str::to_string))
6040 .or_else(|| self.actor_id()),
6041 origin_actor_id: turn_actor
6042 .as_ref()
6043 .and_then(|context| context.origin_actor_id.clone()),
6044 sender_agent_id: turn_actor
6045 .as_ref()
6046 .and_then(|context| context.sender_agent_id.clone()),
6047 };
6048 let tool_context = ToolExecutionContext {
6049 requested_name: request.requested_name.clone(),
6050 canonical_id: canonical_id.clone(),
6051 display_name: resolved.identity.display_name.clone(),
6052 provider_id: resolved.identity.provider_id.clone(),
6053 registry_version: versions.registry,
6054 policy_version: versions.policy,
6055 runtime_control_version: versions.runtime_control,
6056 call_id: request.call_id.clone(),
6057 source: request.source.clone(),
6058 actor,
6059 cancellation: ToolCancellationToken::new(
6060 Arc::clone(&self.runtime_control.emergency_deny),
6061 Some("runtime control cancellation".to_string()),
6062 ),
6063 started_at,
6064 deadline: None,
6065 permission: ToolPolicyDecisionRecord::allow(),
6066 approval: approval_record.clone(),
6067 classification: classification.clone(),
6068 safety,
6069 limits: limits.clone(),
6070 policy_snapshot,
6071 custom_config: security_engine.custom_config(&canonical_id),
6072 };
6073 let (mut result, timed_out, cancelled, invoked) = self
6074 .run_tool_with_retries(
6075 &canonical_id,
6076 resolved.tool.clone(),
6077 executed_arguments.clone(),
6078 tool_context,
6079 timeout,
6080 tool_config.max_retries,
6081 )
6082 .await?;
6083
6084 let fallback_tool = if !result.success && !cancelled {
6088 match &tool_config.on_failure {
6089 ToolFailureAction::Skip => {
6090 result = ToolResult::ok(format!(
6091 "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
6092 canonical_id
6093 ));
6094 None
6095 }
6096 ToolFailureAction::Fallback { fallback_tool } => Some(fallback_tool.clone()),
6097 ToolFailureAction::ReportError => None,
6098 }
6099 } else {
6100 None
6101 };
6102
6103 let output_cap = limits.max_output_chars;
6104 let (output, output_truncated) =
6105 Self::truncate_tool_output(result.output.clone(), output_cap);
6106 if let Some(result_metadata) = result.metadata {
6107 metadata.extend(result_metadata);
6108 }
6109 let mut record = self.record_from_parts_at(
6110 &request,
6111 canonical_id,
6112 executed_arguments,
6113 started_at,
6114 start,
6115 invoked,
6116 result.success,
6117 output,
6118 metadata,
6119 ToolPolicyDecisionRecord::allow(),
6120 approval_record,
6121 timed_out,
6122 output_truncated,
6123 versions,
6124 );
6125 record.cancelled = cancelled;
6126 if cancelled {
6127 record.cancellation_reason = Some("runtime control cancellation".to_string());
6128 }
6129 if let Some(fallback_tool) = fallback_tool {
6130 let fallback_arguments = record.executed_arguments.clone();
6131 let original_tool = record.canonical_id.clone();
6132 self.finish_tool_record_after_resource_guards(resource_guards, &record)
6136 .await;
6137 let fallback_request = ToolExecutionRequest::new(
6138 request.call_id.clone(),
6139 fallback_tool,
6140 fallback_arguments,
6141 ToolCallSource::Fallback { original_tool },
6142 );
6143 return Box::pin(self.execute_tool_record_inner(fallback_request, fallback_state))
6144 .await;
6145 }
6146 self.finish_tool_record_after_resource_guards(resource_guards, &record)
6147 .await;
6148 Ok(record)
6149 }
6150
6151 #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
6152 async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
6153 let record = self
6154 .execute_tool_record(ToolExecutionRequest::new(
6155 tool_call.id.clone(),
6156 tool_call.name.clone(),
6157 tool_call.arguments.clone(),
6158 ToolCallSource::Model,
6159 ))
6160 .await?;
6161 if record.success {
6162 Ok(record.model_output_string())
6163 } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
6164 Err(AgentError::HITLRejected(record.model_output_string()))
6165 } else {
6166 Err(AgentError::Tool(record.model_output_string()))
6167 }
6168 }
6169
6170 async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
6176 let Some(ref router) = self.skill_router else {
6177 return Ok(None);
6178 };
6179 let available_skills = self.get_available_skills();
6180 if available_skills.is_empty() {
6181 return Ok(None);
6182 }
6183 let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
6184 let Some(skill_id) = self
6185 .observe_purpose(
6186 ObservationPurpose::SkillRouting,
6187 router.select_skill_filtered(input, &skill_ids),
6188 )
6189 .await?
6190 else {
6191 return Ok(None);
6192 };
6193 let skill = router
6194 .get_skill(&skill_id)
6195 .cloned()
6196 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6197 info!(skill_id = %skill_id, "Skill selected");
6198 Ok(Some(SkillCandidate::new(skill_id, skill)))
6199 }
6200
6201 async fn commit_skill_candidate_route_result(
6206 &self,
6207 candidate: SkillCandidate,
6208 input: &str,
6209 ) -> Result<SkillRouteResult> {
6210 let skill_id = candidate.skill_id;
6211 let skill = candidate.skill;
6212 let expected_state_generation = self
6213 .state_machine
6214 .as_ref()
6215 .map(|state_machine| state_machine.generation());
6216 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6217 if let Some(ref skill_disambig) = skill.disambiguation
6218 && skill_disambig.enabled.unwrap_or(false)
6219 && let Some(ref disambiguator) = self.disambiguation_manager
6220 {
6221 let context = self.build_disambiguation_context().await?;
6222 let state_override = self
6223 .state_machine
6224 .as_ref()
6225 .and_then(|sm| sm.current_definition())
6226 .and_then(|def| def.disambiguation.clone());
6227
6228 let disambiguation_result = self
6229 .observe_purpose(
6230 ObservationPurpose::DisambiguationDetection,
6231 disambiguator.process_input_with_override(
6232 input,
6233 &context,
6234 state_override.as_ref(),
6235 Some(skill_disambig),
6236 ),
6237 )
6238 .await?;
6239 let current_state_generation = self
6240 .state_machine
6241 .as_ref()
6242 .map(|state_machine| state_machine.generation());
6243 if current_state_generation != expected_state_generation
6244 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6245 {
6246 disambiguator.clear_pending().await;
6247 *self.pending_skill_id.write() = None;
6248 return Err(AgentError::Other(
6249 "State or reset ownership changed during skill disambiguation".to_string(),
6250 ));
6251 }
6252 match disambiguation_result {
6253 DisambiguationResult::Clear => {
6254 debug!(skill_id = %skill_id, "Skill disambiguation: clear");
6255 }
6256 DisambiguationResult::NeedsClarification {
6257 question,
6258 detection,
6259 } => {
6260 let admission = self
6261 .admit_disambiguation_redispatch(
6262 expected_disambiguation_epoch,
6263 expected_state_generation,
6264 )
6265 .await?;
6266 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
6267 info!(
6268 skill_id = %skill_id,
6269 ambiguity_type = ?detection.ambiguity_type,
6270 confidence = detection.confidence,
6271 "Skill requires clarification before execution"
6272 );
6273 *self.pending_skill_id.write() = Some(skill_id.clone());
6274 let response = AgentResponse::new(&question.question).with_metadata(
6275 "disambiguation",
6276 serde_json::json!({
6277 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
6278 "skill_id": skill_id,
6279 "options": question.options,
6280 "clarifying": question.clarifying,
6281 "detection": {
6282 "type": detection.ambiguity_type,
6283 "confidence": detection.confidence,
6284 "what_is_unclear": detection.what_is_unclear,
6285 }
6286 }),
6287 );
6288 drop(admission);
6289 return Ok(SkillRouteResult::NeedsClarification {
6290 response,
6291 ownership: Some(DisambiguationOwnership {
6292 epoch: expected_disambiguation_epoch,
6293 state_generation: expected_state_generation,
6294 }),
6295 });
6296 }
6297 DisambiguationResult::Clarified { enriched_input, .. } => {
6298 info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
6299 let admission = self
6300 .admit_disambiguation_redispatch(
6301 expected_disambiguation_epoch,
6302 expected_state_generation,
6303 )
6304 .await?;
6305 drop(admission);
6306 let content = self.execute_skill(&skill, &enriched_input).await?;
6307 return Ok(SkillRouteResult::Response { skill_id, content });
6308 }
6309 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
6310 info!(skill_id = %skill_id, "Skill disambiguation best guess");
6311 let admission = self
6312 .admit_disambiguation_redispatch(
6313 expected_disambiguation_epoch,
6314 expected_state_generation,
6315 )
6316 .await?;
6317 drop(admission);
6318 let content = self.execute_skill(&skill, &enriched_input).await?;
6319 return Ok(SkillRouteResult::Response { skill_id, content });
6320 }
6321 DisambiguationResult::GiveUp { reason } => {
6322 warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
6323 let apology = self
6324 .generate_localized_apology(
6325 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
6326 &reason,
6327 )
6328 .await
6329 .unwrap_or_else(|_| {
6330 format!("I'm sorry, I couldn't understand your request: {}", reason)
6331 });
6332 return Ok(SkillRouteResult::NeedsClarification {
6333 response: AgentResponse::new(&apology),
6334 ownership: None,
6335 });
6336 }
6337 DisambiguationResult::Escalate { reason } => {
6338 info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6339 let apology = self
6340 .generate_localized_apology(
6341 "Explain briefly that you're transferring the user to a human agent for help.",
6342 &reason,
6343 )
6344 .await
6345 .unwrap_or_else(|_| {
6346 format!("I need human assistance to help with your request: {}", reason)
6347 });
6348 return Ok(SkillRouteResult::NeedsClarification {
6349 response: AgentResponse::new(&apology),
6350 ownership: None,
6351 });
6352 }
6353 DisambiguationResult::Abandoned { .. } => {
6354 debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6355 return Ok(SkillRouteResult::NoMatch);
6356 }
6357 }
6358 }
6359 let admission = self
6360 .admit_disambiguation_redispatch(
6361 expected_disambiguation_epoch,
6362 expected_state_generation,
6363 )
6364 .await?;
6365 drop(admission);
6366 let content = self.execute_skill(&skill, input).await?;
6367 Ok(SkillRouteResult::Response { skill_id, content })
6368 }
6369
6370 async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6372 if let Some(candidate) = self.select_skill_candidate(input).await? {
6373 self.commit_skill_candidate_route_result(candidate, input)
6374 .await
6375 } else {
6376 Ok(SkillRouteResult::NoMatch)
6377 }
6378 }
6379
6380 async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6382 if let Some(ref executor) = self.skill_executor {
6383 let skill_reasoning = self.get_skill_reasoning_config(skill);
6384 let skill_reflection = self.get_skill_reflection_config(skill);
6385
6386 debug!(
6387 skill_id = %skill.id,
6388 reasoning_mode = ?skill_reasoning.mode,
6389 reflection_enabled = ?skill_reflection.enabled,
6390 "Skill reasoning/reflection config"
6391 );
6392
6393 let response = self
6394 .observe_purpose(
6395 ObservationPurpose::SkillPrompt,
6396 executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6397 )
6398 .await?;
6399
6400 if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6401 let should_reflect = self
6402 .should_reflect_with_config(input, &response, &skill_reflection)
6403 .await?;
6404 if should_reflect {
6405 let evaluated = self
6406 .evaluate_and_retry_with_config(input, response, &skill_reflection)
6407 .await?;
6408 return Ok(evaluated);
6409 }
6410 }
6411
6412 return Ok(response);
6413 }
6414 Err(AgentError::Skill(
6415 "No skill executor configured".to_string(),
6416 ))
6417 }
6418
6419 async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6422 let skill = self
6423 .skill_router
6424 .as_ref()
6425 .and_then(|r| r.get_skill(skill_id).cloned())
6426 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6427 self.execute_skill(&skill, input).await
6428 }
6429
6430 async fn should_reflect_with_config(
6431 &self,
6432 input: &str,
6433 response: &str,
6434 config: &ReflectionConfig,
6435 ) -> Result<bool> {
6436 if !config.requires_evaluation() {
6437 return Ok(false);
6438 }
6439
6440 if config.is_enabled() {
6441 return Ok(true);
6442 }
6443
6444 let evaluator_llm = config
6445 .evaluator_llm
6446 .as_ref()
6447 .and_then(|alias| self.llm_registry.get(alias).ok())
6448 .or_else(|| self.llm_registry.router().ok())
6449 .or_else(|| self.llm_registry.default().ok());
6450
6451 let Some(llm) = evaluator_llm else {
6452 return Ok(false);
6453 };
6454
6455 let response_preview: String = response.chars().take(500).collect();
6456 let prompt = format!(
6457 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6458
6459User query: "{}"
6460Response: "{}"
6461
6462Answer YES or NO only."#,
6463 input, response_preview
6464 );
6465
6466 let messages = vec![ChatMessage::user(&prompt)];
6467 let result = self
6468 .observe_purpose(
6469 ObservationPurpose::ReflectionDecision,
6470 llm.complete(&messages, None),
6471 )
6472 .await;
6473
6474 match result {
6475 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6476 Err(_) => Ok(false),
6477 }
6478 }
6479
6480 async fn evaluate_and_retry_with_config(
6481 &self,
6482 input: &str,
6483 mut response: String,
6484 config: &ReflectionConfig,
6485 ) -> Result<String> {
6486 let llm = self.get_state_llm()?;
6487 let mut attempts = 0u32;
6488 let max_retries = config.max_retries;
6489
6490 loop {
6491 let evaluation = self
6492 .evaluate_response_with_config(input, &response, config)
6493 .await?;
6494
6495 if evaluation.passed || attempts >= max_retries {
6496 info!(
6497 passed = evaluation.passed,
6498 confidence = evaluation.confidence,
6499 attempts = attempts + 1,
6500 "Skill reflection evaluation complete"
6501 );
6502 return Ok(response);
6503 }
6504
6505 debug!(
6506 attempt = attempts + 1,
6507 failed_criteria = evaluation.failed_criteria().count(),
6508 "Skill response did not meet criteria, retrying"
6509 );
6510
6511 let feedback: Vec<String> = evaluation
6512 .failed_criteria()
6513 .map(|c| format!("- {}", c.criterion))
6514 .collect();
6515
6516 let retry_prompt = format!(
6517 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6518 feedback.join("\n"),
6519 input
6520 );
6521
6522 let messages = vec![ChatMessage::user(&retry_prompt)];
6523 let retry_response = self
6524 .observe_purpose(
6525 ObservationPurpose::ReflectionEvaluation,
6526 llm.complete(&messages, None),
6527 )
6528 .await
6529 .map_err(|e| AgentError::LLM(e.to_string()))?;
6530
6531 response = retry_response.content.trim().to_string();
6532 attempts += 1;
6533 }
6534 }
6535
6536 async fn evaluate_response_with_config(
6537 &self,
6538 input: &str,
6539 response: &str,
6540 config: &ReflectionConfig,
6541 ) -> Result<EvaluationResult> {
6542 let evaluator_llm = config
6543 .evaluator_llm
6544 .as_ref()
6545 .and_then(|alias| self.llm_registry.get(alias).ok())
6546 .or_else(|| self.llm_registry.router().ok())
6547 .or_else(|| self.llm_registry.default().ok())
6548 .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6549
6550 let criteria = &config.criteria;
6551 let criteria_list = criteria
6552 .iter()
6553 .enumerate()
6554 .map(|(i, c)| format!("{}. {}", i + 1, c))
6555 .collect::<Vec<_>>()
6556 .join("\n");
6557
6558 let prompt = format!(
6559 r#"Evaluate this response against the criteria.
6560
6561User query: "{}"
6562
6563Response to evaluate: "{}"
6564
6565Criteria:
6566{}
6567
6568For each criterion, respond with:
6569- criterion number
6570- PASS or FAIL
6571- brief reason
6572
6573Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6574
6575Format:
65761. PASS/FAIL - reason
65772. PASS/FAIL - reason
6578...
6579CONFIDENCE: 0.X
6580OVERALL: PASS/FAIL"#,
6581 input, response, criteria_list
6582 );
6583
6584 let messages = vec![ChatMessage::user(&prompt)];
6585 let eval_response = self
6586 .observe_purpose(
6587 ObservationPurpose::ReflectionEvaluation,
6588 evaluator_llm.complete(&messages, None),
6589 )
6590 .await
6591 .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6592
6593 let content = eval_response.content.to_uppercase();
6594 let llm_pass = content.contains("OVERALL: PASS");
6595
6596 let confidence = content
6597 .lines()
6598 .find(|l| l.contains("CONFIDENCE:"))
6599 .and_then(|l| {
6600 l.split(':')
6601 .nth(1)
6602 .and_then(|v| v.trim().parse::<f32>().ok())
6603 })
6604 .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6605
6606 let overall_pass = llm_pass && confidence >= config.pass_threshold;
6609
6610 let mut criteria_results = Vec::new();
6611 for (i, criterion) in criteria.iter().enumerate() {
6612 let line_marker = format!("{}.", i + 1);
6613 let passed = eval_response
6614 .content
6615 .lines()
6616 .find(|l| l.contains(&line_marker))
6617 .map(|l| l.to_uppercase().contains("PASS"))
6618 .unwrap_or(overall_pass);
6619
6620 if passed {
6621 criteria_results.push(CriterionResult::pass(criterion));
6622 } else {
6623 criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6624 }
6625 }
6626
6627 Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6628 }
6629
6630 async fn process_input(&self, input: &str) -> Result<ProcessData> {
6632 if let Some(processor) = self.get_state_process_processor() {
6633 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6634 return self
6635 .observe_purpose(purpose, processor.process_input(input))
6636 .await;
6637 }
6638 if let Some(ref processor) = self.process_processor {
6639 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6640 self.observe_purpose(purpose, processor.process_input(input))
6641 .await
6642 } else {
6643 Ok(ProcessData::new(input))
6644 }
6645 }
6646
6647 async fn process_output(
6649 &self,
6650 output: &str,
6651 input_context: &std::collections::HashMap<String, serde_json::Value>,
6652 ) -> Result<ProcessData> {
6653 if let Some(processor) = self.get_state_process_processor() {
6654 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6655 return self
6656 .observe_purpose(purpose, processor.process_output(output, input_context))
6657 .await;
6658 }
6659 if let Some(ref processor) = self.process_processor {
6660 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6661 self.observe_purpose(purpose, processor.process_output(output, input_context))
6662 .await
6663 } else {
6664 Ok(ProcessData::new(output))
6665 }
6666 }
6667
6668 fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6670 let sm = self.state_machine.as_ref()?;
6671 let def = sm.current_definition()?;
6672 let config = def.process.as_ref()?;
6673 let mut processor = ProcessProcessor::new(config.clone());
6674 if let Some(ref registry) = Some(self.llm_registry.clone()) {
6675 processor = processor.with_llm_registry(registry.clone());
6676 }
6677 processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6678 Some(processor)
6679 }
6680
6681 async fn check_turn_timeout(&self) -> Result<()> {
6683 let Some(ref sm) = self.state_machine else {
6684 return Ok(());
6685 };
6686 let Some(timeout_state) = sm.check_timeout() else {
6687 return Ok(());
6688 };
6689 let claim_admission = self.disambiguation_admission.write().await;
6690 if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6691 return Ok(());
6692 }
6693 let Some(reservation) = self.reserve_state_transition() else {
6694 return Ok(());
6695 };
6696 let from_state = sm.current();
6697 let expected_state_generation = sm.generation();
6698 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6699 let history_before = sm.history();
6700 drop(claim_admission);
6701
6702 self.execute_state_exit_actions(&from_state).await;
6703
6704 let admission = self.disambiguation_admission.write().await;
6705 if sm.current() != from_state
6706 || sm.generation() != expected_state_generation
6707 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6708 || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6709 {
6710 return Ok(());
6711 }
6712 sm.transition_to(&timeout_state, "max_turns exceeded")?;
6713 self.invalidate_pending_confirmation("state_timeout").await;
6714 let entered = sm.current();
6715 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6716 drop(admission);
6717
6718 self.execute_state_enter_actions(&entered, is_reentry).await;
6719 drop(reservation);
6720 info!(to = %entered, "Timeout transition");
6721 Ok(())
6722 }
6723
6724 fn increment_turn(&self) {
6725 if let Some(ref sm) = self.state_machine {
6726 sm.increment_turn();
6727 }
6728 }
6729
6730 fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6731 let sm = self.state_machine.as_ref()?;
6732 let current = sm.current();
6733 let transitions: Vec<_> = sm
6734 .auto_transitions()
6735 .into_iter()
6736 .filter(|t| match t.cooldown_turns {
6737 Some(cd) if cd > 0 => {
6738 let resolved = sm.config().resolve_full_path(¤t, &t.to);
6739 !sm.is_on_cooldown(&resolved, cd)
6740 }
6741 _ => true,
6742 })
6743 .collect();
6744 Some((transitions, current))
6745 }
6746
6747 fn transition_reason(transition: &Transition) -> String {
6748 if transition.when.is_empty() {
6749 "guard condition met".to_string()
6750 } else {
6751 transition.when.clone()
6752 }
6753 }
6754
6755 fn build_transition_context(
6757 &self,
6758 user_message: &str,
6759 response: &str,
6760 current_state: &str,
6761 staged: Option<&HashMap<String, Value>>,
6762 ) -> TransitionContext {
6763 let context_map = staged
6764 .map(|writes| self.build_context_with_staged(writes))
6765 .unwrap_or_else(|| self.build_context_with_overlays());
6766 TransitionContext::new(user_message, response, current_state).with_context(context_map)
6767 }
6768
6769 async fn select_transition_candidate(
6771 &self,
6772 user_message: &str,
6773 response: &str,
6774 ) -> Result<Option<TransitionCandidate>> {
6775 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6776 return Ok(None);
6777 };
6778 let transitions: Vec<Transition> = transitions
6779 .into_iter()
6780 .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6781 .collect();
6782 if transitions.is_empty() {
6783 return Ok(None);
6784 }
6785 let Some(evaluator) = self.transition_evaluator.as_ref() else {
6786 return Ok(None);
6787 };
6788 let context = self.build_transition_context(user_message, response, ¤t_state, None);
6789 let selected = self
6790 .observe_purpose(
6791 ObservationPurpose::StateTransitionEvaluation,
6792 evaluator.select_transition(&transitions, &context),
6793 )
6794 .await?;
6795 Ok(selected.map(|index| {
6796 let transition = transitions[index].clone();
6797 TransitionCandidate::new(
6798 current_state,
6799 transition.clone(),
6800 Self::transition_reason(&transition),
6801 )
6802 }))
6803 }
6804
6805 fn select_deterministic_transition_candidate(
6807 &self,
6808 user_message: &str,
6809 current_state: &str,
6810 transitions: &[Transition],
6811 staged: &HashMap<String, Value>,
6812 ) -> Option<TransitionCandidate> {
6813 let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6814
6815 for transition in transitions {
6816 if let Some(guard) = transition.guard.as_ref()
6817 && evaluate_guard(guard, &context)
6818 {
6819 return Some(TransitionCandidate::new(
6820 current_state,
6821 transition.clone(),
6822 Self::transition_reason(transition),
6823 ));
6824 }
6825 }
6826
6827 let resolved_intent = context
6828 .context
6829 .get("resolved_intent")
6830 .and_then(Value::as_str)
6831 .filter(|value| !value.is_empty());
6832 if let Some(resolved_intent) = resolved_intent {
6833 for transition in transitions {
6834 if transition.intent.as_deref() == Some(resolved_intent) {
6835 return Some(TransitionCandidate::new(
6836 current_state,
6837 transition.clone(),
6838 Self::transition_reason(transition),
6839 ));
6840 }
6841 }
6842 }
6843
6844 None
6845 }
6846
6847 async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6849 self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6850 .await
6851 }
6852
6853 async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6855 let approved = self.check_state_hitl(Some(from_state), target).await?;
6856 if !approved {
6857 info!(to = %target, "State transition rejected by HITL");
6858 }
6859 Ok(approved)
6860 }
6861
6862 async fn apply_transition_target(
6864 &self,
6865 from_state: &str,
6866 target: &str,
6867 reason: &str,
6868 staged: Option<&HashMap<String, Value>>,
6869 ) -> Result<bool> {
6870 let Some(ref sm) = self.state_machine else {
6871 return Ok(false);
6872 };
6873 let claim_admission = self.disambiguation_admission.write().await;
6874 if sm.current() != from_state {
6875 return Ok(false);
6876 }
6877 let Some(reservation) = self.reserve_state_transition() else {
6878 return Ok(false);
6879 };
6880 let expected_state_generation = sm.generation();
6881 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6882 let history_before = sm.history();
6883 drop(claim_admission);
6884
6885 self.execute_state_exit_actions(from_state).await;
6886
6887 let admission = self.disambiguation_admission.write().await;
6888 if sm.current() != from_state
6889 || sm.generation() != expected_state_generation
6890 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6891 {
6892 return Ok(false);
6893 }
6894 sm.transition_to(target, reason)?;
6895 self.invalidate_pending_confirmation("state_transition")
6896 .await;
6897 sm.reset_no_transition();
6898 if let Some(staged) = staged {
6899 self.commit_staged_context_writes(staged);
6900 }
6901 let entered = sm.current();
6902 let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6903 drop(admission);
6904
6905 self.execute_state_enter_actions(&entered, is_reentry).await;
6906 drop(reservation);
6907 self.hooks
6908 .on_state_transition(Some(from_state), &entered, reason)
6909 .await;
6910 info!(from = %from_state, to = %entered, "State transition");
6911 Ok(true)
6912 }
6913
6914 async fn commit_transition_target(
6916 &self,
6917 from_state: &str,
6918 target: &str,
6919 reason: &str,
6920 ) -> Result<bool> {
6921 if !self.approve_transition_target(from_state, target).await? {
6922 return Ok(false);
6923 }
6924 self.apply_transition_target(from_state, target, reason, None)
6925 .await
6926 }
6927
6928 async fn apply_pre_response_transition_candidate(
6930 &self,
6931 candidate: &TransitionCandidate,
6932 staged: &HashMap<String, Value>,
6933 processed_input: &str,
6934 ) -> Result<bool> {
6935 self.commit_root_user_message(processed_input).await?;
6936 self.apply_transition_target(
6937 &candidate.from_state,
6938 candidate.target(),
6939 &candidate.reason,
6940 Some(staged),
6941 )
6942 .await
6943 }
6944
6945 async fn commit_pre_response_transition_candidate(
6947 &self,
6948 candidate: &TransitionCandidate,
6949 staged: &HashMap<String, Value>,
6950 processed_input: &str,
6951 ) -> Result<bool> {
6952 if !self
6953 .approve_transition_target(&candidate.from_state, candidate.target())
6954 .await?
6955 {
6956 return Ok(false);
6957 }
6958 self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6959 .await
6960 }
6961
6962 async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6964 let Some(ref sm) = self.state_machine else {
6965 return Ok(false);
6966 };
6967 sm.increment_no_transition();
6968 let Some(fallback) = sm.check_fallback() else {
6969 return Ok(false);
6970 };
6971 self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6972 .await
6973 }
6974
6975 async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6977 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6978 return Ok(false);
6979 };
6980 if transitions.is_empty() {
6981 return Ok(false);
6982 }
6983 if let Some(candidate) = self
6984 .select_transition_candidate(user_message, response)
6985 .await?
6986 {
6987 return self.commit_transition_candidate(&candidate).await;
6988 }
6989 self.handle_transition_miss(¤t_state).await
6990 }
6991
6992 async fn try_pre_response_transition(
6994 &self,
6995 processed_input: &str,
6996 ) -> Result<Option<AgentResponse>> {
6997 let optimization = &self.runtime_config.optimization;
6998 if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6999 return Ok(None);
7000 }
7001 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7002 return Ok(None);
7003 };
7004 let eligible: Vec<Transition> = transitions
7005 .into_iter()
7006 .filter(|transition| !transition.requires_response)
7007 .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
7008 .collect();
7009 if eligible.is_empty() {
7010 return Ok(None);
7011 }
7012
7013 let empty_staged = HashMap::new();
7014 let mut extracted_staged: Option<HashMap<String, Value>> = None;
7015 let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
7016
7017 for transition in &eligible {
7018 let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
7019 let staged_for_eval = if use_extractors {
7020 if extracted_staged.is_none() {
7021 extracted_staged =
7022 Some(self.run_context_extractors_staged(processed_input).await);
7023 }
7024 extracted_staged.as_ref().unwrap_or(&empty_staged)
7025 } else {
7026 &empty_staged
7027 };
7028
7029 if let Some(candidate) = self.select_deterministic_transition_candidate(
7030 processed_input,
7031 ¤t_state,
7032 std::slice::from_ref(transition),
7033 staged_for_eval,
7034 ) {
7035 let staged_for_commit = if use_extractors {
7036 staged_for_eval.clone()
7037 } else {
7038 HashMap::new()
7039 };
7040 selected = Some((candidate, staged_for_commit));
7041 break;
7042 }
7043 }
7044
7045 let Some((candidate, staged)) = selected else {
7046 return Ok(None);
7047 };
7048
7049 if !self
7050 .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
7051 .await?
7052 {
7053 return Ok(None);
7054 }
7055 self.redispatch_current_state(processed_input)
7056 .await
7057 .map(Some)
7058 }
7059
7060 async fn try_speculative_branches(
7065 &self,
7066 processed_input: &str,
7067 input_context: &HashMap<String, Value>,
7068 ) -> Result<Option<AgentResponse>> {
7069 let optimization = &self.runtime_config.optimization;
7070 if !optimization.enabled {
7071 return Ok(None);
7072 }
7073
7074 let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
7075 if !matches!(
7076 effective_reasoning_mode,
7077 ReasoningMode::None | ReasoningMode::Auto
7078 ) {
7079 return Ok(None);
7080 }
7081
7082 let mut transition_enabled =
7083 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
7084 let mut skill_enabled = optimization.speculative_skill_routing
7085 && self.skill_router.is_some()
7086 && self.pending_skill_id.read().is_none();
7087 let mut reasoning_enabled = optimization.speculative_reasoning_auto
7088 && matches!(effective_reasoning_mode, ReasoningMode::Auto);
7089
7090 if matches!(effective_reasoning_mode, ReasoningMode::Auto)
7091 && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
7092 {
7093 return Ok(None);
7094 }
7095
7096 if !transition_enabled && !skill_enabled && !reasoning_enabled {
7097 return Ok(None);
7098 }
7099
7100 let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
7101 let mut speculative_call_slots = optimization
7102 .max_speculative_llm_calls_per_turn
7103 .saturating_sub(1);
7104 if reasoning_enabled {
7105 if optional_slots == 0 || speculative_call_slots == 0 {
7106 return Ok(None);
7107 }
7108 optional_slots -= 1;
7109 speculative_call_slots -= 1;
7110 }
7111 if transition_enabled {
7112 if optional_slots == 0 {
7113 transition_enabled = false;
7114 } else {
7115 optional_slots -= 1;
7116 }
7117 }
7118 if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
7119 skill_enabled = false;
7120 }
7121
7122 if !transition_enabled && !skill_enabled && !reasoning_enabled {
7123 return Ok(None);
7124 }
7125
7126 let main_kind = if transition_enabled {
7127 RuntimeOptimizationKind::ParallelStateTransition
7128 } else if skill_enabled {
7129 RuntimeOptimizationKind::SpeculativeSkillRouting
7130 } else {
7131 RuntimeOptimizationKind::SpeculativeReasoningAuto
7132 };
7133 if !self.reserve_active_speculative_llm_call(main_kind) {
7134 return Ok(None);
7135 }
7136
7137 let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
7138 let main_branch = RuntimeBranch::new(
7139 RuntimeTaskPurpose::MainResponse,
7140 main_kind,
7141 RuntimeTaskPriority::Normal,
7142 RuntimeCommitBehavior::FinalResponse,
7143 );
7144 let transition_branch = RuntimeBranch::new(
7145 RuntimeTaskPurpose::StateTransition,
7146 RuntimeOptimizationKind::ParallelStateTransition,
7147 RuntimeTaskPriority::Critical,
7148 RuntimeCommitBehavior::TransitionDecision,
7149 );
7150 let skill_branch = RuntimeBranch::new(
7151 RuntimeTaskPurpose::SkillRouting,
7152 RuntimeOptimizationKind::SpeculativeSkillRouting,
7153 RuntimeTaskPriority::High,
7154 RuntimeCommitBehavior::SkillSelection,
7155 );
7156 let reasoning_branch = RuntimeBranch::new(
7157 RuntimeTaskPurpose::ReasoningJudge,
7158 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7159 RuntimeTaskPriority::Normal,
7160 RuntimeCommitBehavior::ReasoningDecision,
7161 );
7162 let main_id = main_branch.branch_id();
7163 let transition_id = transition_branch.branch_id();
7164 let skill_id = skill_branch.branch_id();
7165 let reasoning_id = reasoning_branch.branch_id();
7166
7167 let main_id_for_future = main_id.clone();
7168 if !branch_set.schedule(
7169 main_branch,
7170 Box::pin(async move {
7171 match crate::optimization::observability::with_branch_observation(
7172 &main_id_for_future,
7173 main_kind,
7174 RuntimeCommitBehavior::FinalResponse,
7175 self.generate_main_response_draft(processed_input, &ReasoningMode::None),
7176 )
7177 .await
7178 {
7179 Ok(draft) => RuntimeBranchResult::MainDraft(draft),
7180 Err(error) => RuntimeBranchResult::Failed(error),
7181 }
7182 }),
7183 ) {
7184 return Ok(None);
7185 }
7186
7187 if transition_enabled {
7188 let transition_id_for_future = transition_id.clone();
7189 if !branch_set.schedule(
7190 transition_branch,
7191 Box::pin(async move {
7192 match crate::optimization::observability::with_branch_observation(
7193 &transition_id_for_future,
7194 RuntimeOptimizationKind::ParallelStateTransition,
7195 RuntimeCommitBehavior::TransitionDecision,
7196 self.select_parallel_transition_candidate(processed_input),
7197 )
7198 .await
7199 {
7200 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
7201 RuntimeBranchResult::Transition(Some(candidate))
7202 }
7203 Ok(ParallelTransitionSelection::NoMatch) => {
7204 RuntimeBranchResult::Transition(None)
7205 }
7206 Ok(ParallelTransitionSelection::ReservationExhausted) => {
7207 RuntimeBranchResult::Cancelled
7208 }
7209 Err(error) => RuntimeBranchResult::Failed(error),
7210 }
7211 }),
7212 ) {
7213 transition_enabled = false;
7214 }
7215 }
7216
7217 if skill_enabled {
7218 let skill_id_for_future = skill_id.clone();
7219 if !branch_set.schedule(
7220 skill_branch,
7221 Box::pin(async move {
7222 if !self.reserve_active_speculative_llm_call(
7223 RuntimeOptimizationKind::SpeculativeSkillRouting,
7224 ) {
7225 return RuntimeBranchResult::Cancelled;
7226 }
7227 match crate::optimization::observability::with_branch_observation(
7228 &skill_id_for_future,
7229 RuntimeOptimizationKind::SpeculativeSkillRouting,
7230 RuntimeCommitBehavior::SkillSelection,
7231 self.select_skill_candidate(processed_input),
7232 )
7233 .await
7234 {
7235 Ok(candidate) => RuntimeBranchResult::Skill(candidate),
7236 Err(error) => RuntimeBranchResult::Failed(error),
7237 }
7238 }),
7239 ) {
7240 skill_enabled = false;
7241 }
7242 }
7243
7244 if reasoning_enabled {
7245 let reasoning_id_for_future = reasoning_id.clone();
7246 if !branch_set.schedule(
7247 reasoning_branch,
7248 Box::pin(async move {
7249 if !self.reserve_active_speculative_llm_call(
7250 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7251 ) {
7252 return RuntimeBranchResult::Cancelled;
7253 }
7254 match crate::optimization::observability::with_branch_observation(
7255 &reasoning_id_for_future,
7256 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7257 RuntimeCommitBehavior::ReasoningDecision,
7258 self.determine_reasoning_mode_strict(processed_input),
7259 )
7260 .await
7261 {
7262 Ok(mode) => RuntimeBranchResult::Reasoning(mode),
7263 Err(error) => RuntimeBranchResult::Failed(error),
7264 }
7265 }),
7266 ) {
7267 reasoning_enabled = false;
7268 }
7269 }
7270
7271 if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
7272 self.finalize_pending_branches(branch_set.cancel_pending());
7273 return Ok(None);
7274 }
7275
7276 if !transition_enabled && !skill_enabled && !reasoning_enabled {
7277 self.finalize_pending_branches(branch_set.cancel_pending());
7278 return Ok(None);
7279 }
7280
7281 let mut main_pending = true;
7282 let mut skill_pending = skill_enabled;
7283 let mut reasoning_pending = reasoning_enabled;
7284 let mut transition_finalized = !transition_enabled;
7285 let mut skill_finalized = !skill_enabled;
7286 let mut reasoning_finalized = !reasoning_enabled;
7287 let mut main_result: Option<Result<MainResponseDraft>> = None;
7288 let mut transition_candidate: Option<TransitionCandidate> = None;
7289 let mut skill_candidate: Option<SkillCandidate> = None;
7290 let mut reasoning_decision: Option<ReasoningMode> = None;
7291 let mut transition_fallback_required = false;
7292 let mut skill_fallback_required = false;
7293 let mut reasoning_fallback_required = false;
7294
7295 loop {
7296 if let Some(candidate) = transition_candidate.take() {
7297 if self
7298 .approve_transition_target(&candidate.from_state, candidate.target())
7299 .await?
7300 {
7301 self.finalize_pending_branches(branch_set.cancel_pending());
7303 if !main_pending {
7304 self.finalize_branch_loss(
7305 &main_id,
7306 main_kind,
7307 RuntimeCommitBehavior::FinalResponse,
7308 false,
7309 main_result.as_ref().map(|result| result.is_err()),
7310 );
7311 }
7312 if skill_enabled && !skill_pending {
7313 self.finalize_branch_loss(
7314 &skill_id,
7315 RuntimeOptimizationKind::SpeculativeSkillRouting,
7316 RuntimeCommitBehavior::SkillSelection,
7317 false,
7318 Some(false),
7319 );
7320 }
7321 if reasoning_enabled && !reasoning_pending {
7322 self.finalize_branch_loss(
7323 &reasoning_id,
7324 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7325 RuntimeCommitBehavior::ReasoningDecision,
7326 false,
7327 Some(false),
7328 );
7329 }
7330 if !self
7331 .apply_pre_response_transition_candidate(
7332 &candidate,
7333 &HashMap::new(),
7334 processed_input,
7335 )
7336 .await?
7337 {
7338 self.finalize_optional_branch(
7339 &transition_id,
7340 RuntimeOptimizationKind::ParallelStateTransition,
7341 RuntimeCommitBehavior::TransitionDecision,
7342 "discarded",
7343 false,
7344 );
7345 return Ok(None);
7346 }
7347 self.finalize_optional_branch(
7348 &transition_id,
7349 RuntimeOptimizationKind::ParallelStateTransition,
7350 RuntimeCommitBehavior::TransitionDecision,
7351 "committed",
7352 true,
7353 );
7354 return self
7355 .redispatch_current_state(processed_input)
7356 .await
7357 .map(Some);
7358 }
7359 self.finalize_optional_branch(
7360 &transition_id,
7361 RuntimeOptimizationKind::ParallelStateTransition,
7362 RuntimeCommitBehavior::TransitionDecision,
7363 "discarded",
7364 false,
7365 );
7366 transition_finalized = true;
7367 }
7368
7369 if transition_finalized && skill_candidate.is_some() {
7370 let candidate = skill_candidate.take().unwrap();
7371 self.finalize_optional_branch(
7372 &skill_id,
7373 RuntimeOptimizationKind::SpeculativeSkillRouting,
7374 RuntimeCommitBehavior::SkillSelection,
7375 "committed",
7376 true,
7377 );
7378 if !main_pending {
7379 self.finalize_branch_loss(
7380 &main_id,
7381 main_kind,
7382 RuntimeCommitBehavior::FinalResponse,
7383 false,
7384 main_result.as_ref().map(|result| result.is_err()),
7385 );
7386 }
7387 if reasoning_enabled && !reasoning_pending {
7388 self.finalize_branch_loss(
7389 &reasoning_id,
7390 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7391 RuntimeCommitBehavior::ReasoningDecision,
7392 false,
7393 Some(false),
7394 );
7395 }
7396 self.finalize_pending_branches(branch_set.cancel_pending());
7397 self.commit_root_user_message(processed_input).await?;
7398 return match self
7399 .commit_skill_candidate_route_result(candidate, processed_input)
7400 .await?
7401 {
7402 SkillRouteResult::Response { skill_id, content } => self
7403 .handle_skill_response(processed_input, &skill_id, content, input_context)
7404 .await
7405 .map(Some),
7406 SkillRouteResult::NeedsClarification {
7407 response,
7408 ownership,
7409 } => {
7410 let admission = self
7411 .admit_optional_disambiguation_ownership(ownership)
7412 .await?;
7413 if response
7414 .metadata
7415 .as_ref()
7416 .and_then(|m| m.get("disambiguation"))
7417 .and_then(|d| d.get("status"))
7418 .and_then(|s| s.as_str())
7419 == Some("awaiting_clarification")
7420 {
7421 self.memory
7422 .add_message(ChatMessage::assistant(&response.content))
7423 .await?;
7424 }
7425 drop(admission);
7426 self.finish_turn_if_root(&response).await?;
7427 Ok(Some(response))
7428 }
7429 SkillRouteResult::NoMatch => Ok(None),
7430 };
7431 }
7432
7433 if transition_finalized
7434 && skill_finalized
7435 && let Some(reasoning_mode) = reasoning_decision.take()
7436 {
7437 if !matches!(reasoning_mode, ReasoningMode::None) {
7438 self.finalize_optional_branch(
7439 &reasoning_id,
7440 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7441 RuntimeCommitBehavior::ReasoningDecision,
7442 "committed",
7443 true,
7444 );
7445 if !main_pending {
7446 self.finalize_branch_loss(
7447 &main_id,
7448 main_kind,
7449 RuntimeCommitBehavior::FinalResponse,
7450 false,
7451 main_result.as_ref().map(|result| result.is_err()),
7452 );
7453 }
7454 self.finalize_pending_branches(branch_set.cancel_pending());
7455 self.commit_root_user_message(processed_input).await?;
7456 return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7457 self.handle_plan_and_execute(processed_input, input_context, true)
7458 .await
7459 .map(Some)
7460 } else {
7461 self.run_committed_response_loop_with_reasoning(
7462 processed_input,
7463 input_context,
7464 reasoning_mode,
7465 true,
7466 )
7467 .await
7468 .map(Some)
7469 };
7470 }
7471 self.finalize_optional_branch(
7472 &reasoning_id,
7473 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7474 RuntimeCommitBehavior::ReasoningDecision,
7475 "committed",
7476 true,
7477 );
7478 reasoning_finalized = true;
7479 }
7480
7481 if transition_finalized && skill_finalized && reasoning_finalized {
7482 if transition_fallback_required
7483 || skill_fallback_required
7484 || reasoning_fallback_required
7485 {
7486 if !main_pending {
7487 self.finalize_branch_loss(
7488 &main_id,
7489 main_kind,
7490 RuntimeCommitBehavior::FinalResponse,
7491 false,
7492 main_result.as_ref().map(|result| result.is_err()),
7493 );
7494 }
7495 self.finalize_pending_branches(branch_set.cancel_pending());
7496 return Ok(None);
7497 }
7498
7499 if let Some(result) = main_result.take() {
7500 let draft = match result {
7501 Ok(draft) => draft,
7502 Err(error) => {
7503 self.finalize_optional_branch(
7504 &main_id,
7505 main_kind,
7506 RuntimeCommitBehavior::FinalResponse,
7507 "failed",
7508 false,
7509 );
7510 self.finalize_pending_branches(branch_set.cancel_pending());
7511 return Err(error);
7512 }
7513 };
7514 self.finalize_optional_branch(
7515 &main_id,
7516 main_kind,
7517 RuntimeCommitBehavior::FinalResponse,
7518 "committed",
7519 true,
7520 );
7521 self.finalize_pending_branches(branch_set.cancel_pending());
7522 return self
7523 .commit_main_response_draft(
7524 processed_input,
7525 input_context,
7526 draft,
7527 ReasoningMode::None,
7528 reasoning_enabled,
7529 )
7530 .await
7531 .map(Some);
7532 }
7533 }
7534
7535 if branch_set.is_empty() {
7536 return Ok(None);
7537 }
7538
7539 let Some(outcome) = branch_set.next_completed().await else {
7540 return Ok(None);
7541 };
7542 let branch_id = outcome.branch.branch_id();
7543 match outcome.result {
7544 RuntimeBranchResult::MainDraft(draft) => {
7545 main_pending = false;
7546 main_result = Some(Ok(draft));
7547 }
7548 RuntimeBranchResult::Transition(candidate) => {
7549 if let Some(candidate) = candidate {
7550 transition_candidate = Some(candidate);
7551 } else {
7552 self.finalize_optional_branch(
7553 &transition_id,
7554 RuntimeOptimizationKind::ParallelStateTransition,
7555 RuntimeCommitBehavior::TransitionDecision,
7556 "discarded",
7557 false,
7558 );
7559 transition_finalized = true;
7560 }
7561 }
7562 RuntimeBranchResult::Skill(candidate) => {
7563 skill_pending = false;
7564 if let Some(candidate) = candidate {
7565 skill_candidate = Some(candidate);
7566 } else {
7567 self.finalize_optional_branch(
7568 &skill_id,
7569 RuntimeOptimizationKind::SpeculativeSkillRouting,
7570 RuntimeCommitBehavior::SkillSelection,
7571 "discarded",
7572 false,
7573 );
7574 skill_finalized = true;
7575 }
7576 }
7577 RuntimeBranchResult::Reasoning(mode) => {
7578 reasoning_pending = false;
7579 reasoning_decision = Some(mode);
7580 }
7581 RuntimeBranchResult::Failed(error) => {
7582 if branch_id == main_id {
7583 main_pending = false;
7584 main_result = Some(Err(error));
7585 } else if branch_id == transition_id {
7586 self.finalize_optional_branch(
7587 &transition_id,
7588 RuntimeOptimizationKind::ParallelStateTransition,
7589 RuntimeCommitBehavior::TransitionDecision,
7590 "failed",
7591 false,
7592 );
7593 transition_finalized = true;
7594 } else if branch_id == skill_id {
7595 skill_pending = false;
7596 self.finalize_optional_branch(
7597 &skill_id,
7598 RuntimeOptimizationKind::SpeculativeSkillRouting,
7599 RuntimeCommitBehavior::SkillSelection,
7600 "failed",
7601 false,
7602 );
7603 skill_finalized = true;
7604 } else if branch_id == reasoning_id {
7605 reasoning_pending = false;
7606 self.finalize_optional_branch(
7607 &reasoning_id,
7608 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7609 RuntimeCommitBehavior::ReasoningDecision,
7610 "failed",
7611 false,
7612 );
7613 reasoning_finalized = true;
7614 }
7615 }
7616 RuntimeBranchResult::Cancelled => {
7617 self.finalize_optional_branch(
7618 &branch_id,
7619 outcome.branch.optimization,
7620 outcome.branch.commit_behavior,
7621 "cancelled",
7622 false,
7623 );
7624 if branch_id == main_id {
7625 main_pending = false;
7626 main_result =
7627 Some(Err(AgentError::Other("main branch cancelled".to_string())));
7628 } else if branch_id == transition_id {
7629 transition_finalized = true;
7630 transition_fallback_required = true;
7631 } else if branch_id == skill_id {
7632 skill_pending = false;
7633 skill_finalized = true;
7634 skill_fallback_required = true;
7635 } else if branch_id == reasoning_id {
7636 reasoning_pending = false;
7637 reasoning_finalized = true;
7638 reasoning_fallback_required = true;
7639 }
7640 }
7641 }
7642 }
7643 }
7644
7645 fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7646 for branch in branches {
7647 self.finalize_optional_branch(
7648 &branch.branch_id(),
7649 branch.optimization,
7650 branch.commit_behavior,
7651 "cancelled",
7652 false,
7653 );
7654 }
7655 }
7656
7657 fn finalize_branch_loss(
7662 &self,
7663 branch_id: &str,
7664 optimization: RuntimeOptimizationKind,
7665 commit_behavior: RuntimeCommitBehavior,
7666 pending: bool,
7667 completed_failed: Option<bool>,
7668 ) {
7669 let status = if pending {
7670 "cancelled"
7671 } else if completed_failed.unwrap_or(false) {
7672 "failed"
7673 } else {
7674 "discarded"
7675 };
7676 self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7677 }
7678
7679 fn finalize_optional_branch(
7684 &self,
7685 branch_id: &str,
7686 optimization: RuntimeOptimizationKind,
7687 commit_behavior: RuntimeCommitBehavior,
7688 status: &str,
7689 winner: bool,
7690 ) {
7691 crate::optimization::observability::finalize_branch(
7692 self.observability_manager.as_ref(),
7693 branch_id,
7694 status,
7695 winner,
7696 optimization,
7697 commit_behavior,
7698 );
7699 }
7700
7701 fn has_parallel_transition_candidates(&self) -> bool {
7706 self.transitions_available_for_commit()
7707 .map(|(transitions, _)| {
7708 transitions
7709 .iter()
7710 .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7711 })
7712 .unwrap_or(false)
7713 }
7714
7715 async fn select_parallel_transition_candidate(
7720 &self,
7721 processed_input: &str,
7722 ) -> Result<ParallelTransitionSelection> {
7723 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7724 return Ok(ParallelTransitionSelection::NoMatch);
7725 };
7726 let parallel: Vec<Transition> = transitions
7727 .into_iter()
7728 .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7729 .filter(|transition| !transition.requires_response)
7730 .collect();
7731 if parallel.is_empty() {
7732 return Ok(ParallelTransitionSelection::NoMatch);
7733 }
7734 let empty_staged = HashMap::new();
7735 if let Some(candidate) = self.select_deterministic_transition_candidate(
7736 processed_input,
7737 ¤t_state,
7738 ¶llel,
7739 &empty_staged,
7740 ) {
7741 return Ok(ParallelTransitionSelection::Candidate(candidate));
7742 }
7743 let when_transitions: Vec<(usize, &Transition)> = parallel
7744 .iter()
7745 .enumerate()
7746 .filter(|(_, transition)| !transition.when.trim().is_empty())
7747 .collect();
7748 if when_transitions.is_empty() {
7749 return Ok(ParallelTransitionSelection::NoMatch);
7750 }
7751 let llm = self
7752 .llm_registry
7753 .router()
7754 .or_else(|_| self.llm_registry.default())
7755 .map_err(|e| AgentError::Config(e.to_string()))?;
7756 let conditions = when_transitions
7757 .iter()
7758 .enumerate()
7759 .map(|(display_idx, (_, transition))| {
7760 format!("{}. {}", display_idx + 1, transition.when)
7761 })
7762 .collect::<Vec<_>>()
7763 .join("\n");
7764 if !self
7765 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7766 {
7767 return Ok(ParallelTransitionSelection::ReservationExhausted);
7768 }
7769 let context_preview = self.branch_context_preview();
7770 let prompt = format!(
7771 "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-{}).",
7772 current_state,
7773 processed_input,
7774 context_preview,
7775 conditions,
7776 when_transitions.len()
7777 );
7778 let response = self
7779 .observe_purpose(
7780 ObservationPurpose::StateTransitionEvaluation,
7781 llm.complete(&[ChatMessage::user(prompt)], None),
7782 )
7783 .await
7784 .map_err(|e| AgentError::LLM(e.to_string()))?;
7785 let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7786 if choice == 0 || choice > when_transitions.len() {
7787 return Ok(ParallelTransitionSelection::NoMatch);
7788 }
7789 let transition = when_transitions[choice - 1].1.clone();
7790 Ok(ParallelTransitionSelection::Candidate(
7791 TransitionCandidate::new(
7792 current_state,
7793 transition.clone(),
7794 Self::transition_reason(&transition),
7795 ),
7796 ))
7797 }
7798
7799 async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7801 const MAX_REDISPATCH_DEPTH: u32 = 3;
7802 let current_depth = *self.redispatch_depth.read();
7803 if current_depth >= MAX_REDISPATCH_DEPTH {
7804 warn!(depth = current_depth, "Re-dispatch depth limit reached");
7805 let response = AgentResponse::new("");
7806 self.finish_turn_if_root(&response).await?;
7807 return Ok(response);
7808 }
7809 *self.redispatch_depth.write() += 1;
7810 if let Some(context) = self.active_turn_context.write().as_mut() {
7811 context.enter_redispatch();
7812 }
7813 let result = Box::pin(self.run_loop_internal(processed_input)).await;
7814 *self.redispatch_depth.write() -= 1;
7815 if let Some(context) = self.active_turn_context.write().as_mut() {
7816 context.exit_redispatch();
7817 }
7818 let response = result?;
7819 self.finish_turn_if_root(&response).await?;
7820 Ok(response)
7821 }
7822
7823 async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7825 if *self.redispatch_depth.read() == 0 {
7826 self.post_turn_session_lifecycle().await?;
7827 if let Some(context) = self.active_turn_context.write().as_mut() {
7828 context.mark_post_turn_lifecycle_completed();
7829 }
7830 self.hooks.on_response(response).await;
7831 self.end_root_turn();
7832 }
7833 Ok(())
7834 }
7835
7836 async fn execute_state_exit_actions(&self, state_path: &str) {
7838 if let Some(ref sm) = self.state_machine
7839 && let Some(def) = sm.get_definition(state_path)
7840 && !def.on_exit.is_empty()
7841 {
7842 debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7843 self.execute_state_actions(&def.on_exit).await;
7844 }
7845 }
7846
7847 fn state_was_previously_entered(
7849 state_path: &str,
7850 from_state: &str,
7851 history_before: &[StateTransitionEvent],
7852 ) -> bool {
7853 state_path == from_state
7854 || history_before
7855 .iter()
7856 .any(|event| event.from == state_path || event.to == state_path)
7857 }
7858
7859 async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7861 if let Some(ref sm) = self.state_machine
7862 && let Some(def) = sm.get_definition(state_path)
7863 {
7864 if is_reentry && !def.on_reenter.is_empty() {
7865 debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7866 self.execute_state_actions(&def.on_reenter).await;
7867 } else if !def.on_enter.is_empty() {
7868 debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7869 self.execute_state_actions(&def.on_enter).await;
7870 }
7871 }
7872 }
7873
7874 async fn execute_state_actions(&self, actions: &[StateAction]) {
7876 for (action_index, action) in actions.iter().enumerate() {
7877 match action {
7878 StateAction::Tool { tool, args } => {
7879 let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7880 let args_value = self.render_action_args(&raw_args);
7881 let state = self.state_machine.as_ref().map(|sm| sm.current());
7882 let request = ToolExecutionRequest::new(
7883 uuid::Uuid::new_v4().to_string(),
7884 tool.clone(),
7885 args_value,
7886 ToolCallSource::StateAction {
7887 state,
7888 action_index,
7889 },
7890 );
7891 match self.execute_tool_record(request).await {
7892 Ok(record) if record.success => {
7893 debug!(tool = %record.canonical_id, "State action: tool executed");
7894 let _ = self.context_manager.set(
7895 "last_tool_result",
7896 serde_json::Value::String(record.model_output_string()),
7897 );
7898 let _ = self.context_manager.set(
7899 "last_tool_record",
7900 serde_json::to_value(record).unwrap_or(Value::Null),
7901 );
7902 }
7903 Ok(record) => {
7904 warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7905 }
7906 Err(e) => {
7907 warn!(tool = %tool, error = %e, "State action: tool failed")
7908 }
7909 }
7910 }
7911 StateAction::Skill { skill } => {
7912 if let Some(ref executor) = self.skill_executor {
7913 if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7914 match executor
7915 .execute_with_invoker(def, "", serde_json::json!({}), self)
7916 .await
7917 {
7918 Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7919 Err(e) => {
7920 warn!(skill = %skill, error = %e, "State action: skill failed")
7921 }
7922 }
7923 } else {
7924 warn!(skill = %skill, "State action: skill not found");
7925 }
7926 }
7927 }
7928 StateAction::SetContext { set_context } => {
7929 for (key, value) in set_context {
7930 if let Err(e) = self.context_manager.set(key, value.clone()) {
7931 warn!(key = %key, error = %e, "State action: set_context failed");
7932 } else {
7933 debug!(key = %key, "State action: context set");
7934 }
7935 }
7936 }
7937 StateAction::Prompt {
7938 prompt,
7939 llm,
7940 store_as,
7941 } => {
7942 let llm_result = if let Some(alias) = llm {
7943 self.llm_registry.get(alias)
7944 } else {
7945 self.llm_registry.default()
7946 };
7947 match llm_result {
7948 Ok(llm_provider) => {
7949 let context = self.build_context_with_overlays();
7951 let rendered_prompt = self
7952 .template_renderer
7953 .render(prompt, &context)
7954 .unwrap_or_else(|_| prompt.clone());
7955 let recent =
7956 self.memory.get_messages(Some(5)).await.unwrap_or_default();
7957 let mut messages: Vec<ChatMessage> = recent;
7958 messages.push(ChatMessage::user(&rendered_prompt));
7959 match self
7960 .observe_purpose(
7961 ObservationPurpose::StateAction,
7962 llm_provider.complete(&messages, None),
7963 )
7964 .await
7965 {
7966 Ok(response) => {
7967 if let Some(key) = store_as {
7968 let _ = self
7969 .context_manager
7970 .set(key, Value::String(response.content));
7971 debug!(key = %key, "State action: prompt result stored");
7972 }
7973 }
7974 Err(e) => {
7975 warn!(error = %e, "State action: prompt LLM call failed");
7976 }
7977 }
7978 }
7979 Err(e) => {
7980 warn!(error = %e, "State action: LLM not found for prompt");
7981 }
7982 }
7983 }
7984 }
7985 }
7986 }
7987
7988 async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7989 let extractors = match &self.state_machine {
7990 Some(sm) => match sm.current_definition() {
7991 Some(def) if !def.extract.is_empty() => def.extract.clone(),
7992 _ => return HashMap::new(),
7993 },
7994 None => return HashMap::new(),
7995 };
7996
7997 let mut staged = HashMap::new();
7998 for extractor in &extractors {
7999 let prompt = if let Some(ref custom) = extractor.llm_extract {
8000 format!(
8001 "User message:\n\"{}\"\n\nInstruction:\n{}",
8002 user_message, custom
8003 )
8004 } else if let Some(ref desc) = extractor.description {
8005 format!(
8006 "From the following message, extract: {}\n\n\
8007 Message: \"{}\"\n\n\
8008 If the information is present, return ONLY the extracted value.\n\
8009 If NOT present, return exactly: __NONE__",
8010 desc, user_message
8011 )
8012 } else {
8013 continue;
8014 };
8015
8016 let llm = match self
8017 .llm_registry
8018 .get(&extractor.llm)
8019 .or_else(|_| self.llm_registry.get("router"))
8020 .or_else(|_| self.llm_registry.get("default"))
8021 {
8022 Ok(llm) => llm,
8023 Err(e) => {
8024 warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
8025 continue;
8026 }
8027 };
8028
8029 let messages = vec![ChatMessage::user(&prompt)];
8030 match self
8031 .observe_purpose(
8032 ObservationPurpose::ContextExtraction,
8033 llm.complete(&messages, None),
8034 )
8035 .await
8036 {
8037 Ok(response) => {
8038 let value = response.content.trim().to_string();
8039 if value != "__NONE__" && !value.is_empty() {
8040 staged.insert(
8041 extractor.key.clone(),
8042 serde_json::Value::String(value.clone()),
8043 );
8044 debug!(key = %extractor.key, value = %value, "Context extracted");
8045 } else if extractor.required {
8046 warn!(key = %extractor.key, "Required extraction returned no value");
8047 }
8048 }
8049 Err(e) => {
8050 warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
8051 }
8052 }
8053 }
8054 staged
8055 }
8056
8057 fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
8058 for (key, value) in staged {
8059 if let Err(error) = self.context_manager.update(key, value.clone()) {
8060 warn!(key = %key, error = %error, "staged context write failed");
8061 }
8062 }
8063 }
8064
8065 async fn run_context_extractors(&self, user_message: &str) {
8067 let staged = self.run_context_extractors_staged(user_message).await;
8068 self.commit_staged_context_writes(&staged);
8069 }
8070
8071 async fn check_memory_compression(&self) -> Result<()> {
8072 if self.memory.needs_compression() {
8073 let result = self.memory.compress(None).await?;
8074 if let CompressResult::Compressed {
8075 messages_summarized,
8076 new_summary_length,
8077 tokens_saved,
8078 } = result
8079 {
8080 let event = MemoryCompressEvent::new(
8081 messages_summarized,
8082 tokens_saved,
8083 new_summary_length as u32,
8084 );
8085 self.hooks.on_memory_compress(&event).await;
8086 debug!(
8087 messages = messages_summarized,
8088 tokens_saved = tokens_saved,
8089 "Memory compressed"
8090 );
8091 }
8092 }
8093
8094 self.handle_memory_overflow().await?;
8096 self.check_memory_budget().await;
8097
8098 Ok(())
8099 }
8100
8101 async fn check_memory_budget(&self) {
8102 let Some(ref budget) = self.memory_token_budget else {
8103 return;
8104 };
8105
8106 let context = match self.memory.get_context().await {
8107 Ok(ctx) => ctx,
8108 Err(_) => return,
8109 };
8110
8111 let used_tokens = context.estimated_tokens();
8113 if budget.is_over_warn_threshold(used_tokens) {
8114 let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
8115 self.hooks.on_memory_budget_warning(&event).await;
8116 debug!(
8117 used = used_tokens,
8118 total = budget.total,
8119 percent = event.usage_percent,
8120 "Memory budget warning"
8121 );
8122 }
8123
8124 if let Some(ref summary) = context.summary {
8126 let summary_tokens = ai_agents_memory::estimate_tokens(summary);
8127 let summary_budget = budget.allocation.summary;
8128 if summary_budget > 0 {
8129 let warn_threshold =
8130 (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
8131 if summary_tokens >= warn_threshold {
8132 let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
8133 self.hooks.on_memory_budget_warning(&event).await;
8134 }
8135 }
8136 }
8137
8138 let recent_tokens: u32 = context
8140 .messages
8141 .iter()
8142 .map(ai_agents_memory::estimate_message_tokens)
8143 .sum();
8144 let recent_budget = budget.allocation.recent_messages;
8145 if recent_budget > 0 {
8146 let warn_threshold =
8147 (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
8148 if recent_tokens >= warn_threshold {
8149 let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
8150 self.hooks.on_memory_budget_warning(&event).await;
8151 }
8152 }
8153
8154 let relationship_budget = budget.allocation.relationships;
8155 if relationship_budget > 0 {
8156 let relationship_tokens = self
8157 .relationship_memory_text()
8158 .map(|text| ai_agents_memory::estimate_tokens(&text))
8159 .unwrap_or(0);
8160 let warn_threshold =
8161 (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
8162 if relationship_tokens >= warn_threshold {
8163 let event = MemoryBudgetEvent::new(
8164 "relationships",
8165 relationship_tokens,
8166 relationship_budget,
8167 );
8168 self.hooks.on_memory_budget_warning(&event).await;
8169 }
8170 }
8171 }
8172
8173 async fn handle_memory_overflow(&self) -> Result<()> {
8174 let Some(ref budget) = self.memory_token_budget else {
8175 return Ok(());
8176 };
8177
8178 let context = self.memory.get_context().await?;
8179 let used_tokens = context.estimated_tokens();
8180
8181 if used_tokens <= budget.total {
8182 return Ok(());
8183 }
8184
8185 match budget.overflow_strategy {
8186 OverflowStrategy::TruncateOldest => {
8187 let tokens_to_free = used_tokens - budget.total;
8188 let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
8189 if messages_to_evict > 0 {
8190 self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
8191 .await?;
8192 }
8193 }
8194 OverflowStrategy::SummarizeMore => {
8195 let max_attempts = context.total_messages.max(1);
8196 for _ in 0..max_attempts {
8197 match self.memory.compress(None).await? {
8198 CompressResult::Compressed {
8199 messages_summarized,
8200 ..
8201 } if messages_summarized > 0 => {
8202 let context = self.memory.get_context().await?;
8203 if context.estimated_tokens() <= budget.total {
8204 return Ok(());
8205 }
8206 }
8207 _ => break,
8208 }
8209 }
8210 let context = self.memory.get_context().await?;
8211 let used_tokens = context.estimated_tokens();
8212 if used_tokens > budget.total {
8213 return Err(AgentError::MemoryBudgetExceeded {
8214 used: used_tokens,
8215 budget: budget.total,
8216 });
8217 }
8218 }
8219 OverflowStrategy::Error => {
8220 return Err(AgentError::MemoryBudgetExceeded {
8221 used: used_tokens,
8222 budget: budget.total,
8223 });
8224 }
8225 }
8226 Ok(())
8227 }
8228
8229 fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
8230 ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
8232 }
8233
8234 async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
8235 let evicted = self.memory.evict_oldest(count).await?;
8236 if !evicted.is_empty() {
8237 let event = MemoryEvictEvent {
8238 reason,
8239 messages_evicted: evicted.len(),
8240 importance_scores: vec![],
8241 };
8242 self.hooks.on_memory_evict(&event).await;
8243 debug!(count = evicted.len(), "Messages evicted from memory");
8244 }
8245 Ok(())
8246 }
8247
8248 #[instrument(skip(self, input), fields(agent = %self.info.name))]
8249 async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
8250 match self.determine_reasoning_mode_strict(input).await {
8251 Ok(mode) => Ok(mode),
8252 Err(_) => Ok(ReasoningMode::None),
8253 }
8254 }
8255
8256 async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
8257 let effective_config = self.get_effective_reasoning_config();
8258
8259 if !matches!(effective_config.mode, ReasoningMode::Auto) {
8260 return Ok(effective_config.mode.clone());
8261 }
8262
8263 let judge_llm = effective_config
8264 .judge_llm
8265 .as_ref()
8266 .and_then(|alias| self.llm_registry.get(alias).ok())
8267 .or_else(|| self.llm_registry.router().ok())
8268 .or_else(|| self.llm_registry.default().ok());
8269
8270 let Some(llm) = judge_llm else {
8271 return Ok(ReasoningMode::None);
8272 };
8273
8274 let prompt = format!(
8275 r#"Analyze this user request and determine the appropriate reasoning mode.
8276
8277User request: "{}"
8278
8279Choose ONE of these modes:
8280- none: Simple queries, greetings, direct answers (fastest)
8281- cot: Complex analysis, multi-step reasoning, math problems
8282- react: Tasks requiring multiple tool calls with observation
8283- plan_and_execute: Complex multi-step tasks requiring coordination
8284
8285Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
8286 input
8287 );
8288
8289 let messages = vec![ChatMessage::user(&prompt)];
8290 let response = self
8291 .observe_purpose(
8292 ObservationPurpose::ReflectionDecision,
8293 llm.complete(&messages, None),
8294 )
8295 .await
8296 .map_err(|e| AgentError::LLM(e.to_string()))?;
8297
8298 let mode_str = response.content.trim().to_lowercase();
8299 Ok(match mode_str.as_str() {
8300 "cot" => ReasoningMode::CoT,
8301 "react" => ReasoningMode::React,
8302 "plan_and_execute" => ReasoningMode::PlanAndExecute,
8303 _ => ReasoningMode::None,
8304 })
8305 }
8306
8307 async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
8308 let effective_config = self.get_effective_reflection_config();
8309
8310 if !effective_config.requires_evaluation() {
8311 return Ok(false);
8312 }
8313
8314 if effective_config.is_enabled() {
8315 return Ok(true);
8316 }
8317
8318 let evaluator_llm = effective_config
8319 .evaluator_llm
8320 .as_ref()
8321 .and_then(|alias| self.llm_registry.get(alias).ok())
8322 .or_else(|| self.llm_registry.router().ok())
8323 .or_else(|| self.llm_registry.default().ok());
8324
8325 let Some(llm) = evaluator_llm else {
8326 return Ok(false);
8327 };
8328
8329 let response_preview: String = response.chars().take(500).collect();
8330 let prompt = format!(
8331 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8332
8333User query: "{}"
8334Response: "{}"
8335
8336Answer YES or NO only."#,
8337 input, response_preview
8338 );
8339
8340 let messages = vec![ChatMessage::user(&prompt)];
8341 let result = self
8342 .observe_purpose(
8343 ObservationPurpose::ReflectionDecision,
8344 llm.complete(&messages, None),
8345 )
8346 .await;
8347
8348 match result {
8349 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8350 Err(_) => Ok(false),
8351 }
8352 }
8353
8354 fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8355 format!(
8356 "{}\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>",
8357 base_prompt
8358 )
8359 }
8360
8361 fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8362 format!(
8363 "{}\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>",
8364 base_prompt
8365 )
8366 }
8367
8368 async fn generate_plan(&self, input: &str) -> Result<Plan> {
8369 let effective = self.get_effective_reasoning_config();
8370 let planning_config = effective.get_planning();
8371
8372 let planner_llm = planning_config
8373 .and_then(|c| c.planner_llm.as_ref())
8374 .and_then(|alias| self.llm_registry.get(alias).ok())
8375 .or_else(|| self.llm_registry.router().ok())
8376 .or_else(|| self.llm_registry.default().ok())
8377 .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8378
8379 let mut available_tool_ids: Vec<String> = self
8380 .get_available_tool_ids()
8381 .await
8382 .unwrap_or_else(|_| self.tools.list_ids());
8383 let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8384
8385 if let Some(config) = planning_config {
8387 if !config.available.tools.is_all() {
8388 available_tool_ids.retain(|t| config.available.tools.allows(t));
8389 }
8390 if !config.available.skills.is_all() {
8391 available_skills.retain(|s| config.available.skills.allows(s));
8392 }
8393 }
8394
8395 let tool_descriptions: Vec<String> = available_tool_ids
8398 .iter()
8399 .filter_map(|id| {
8400 self.tools.get(id).map(|tool| {
8401 let schema = tool.input_schema();
8402 let args_desc = schema
8403 .get("properties")
8404 .and_then(|p| serde_json::to_string(p).ok())
8405 .unwrap_or_else(|| "{}".to_string());
8406 format!(
8407 "- {} ({}): {}\n Arguments: {}",
8408 id,
8409 tool.name(),
8410 tool.description(),
8411 args_desc
8412 )
8413 })
8414 })
8415 .collect();
8416
8417 let tools_section = if tool_descriptions.is_empty() {
8418 "Available tools: none".to_string()
8419 } else {
8420 format!("Available tools:\n{}", tool_descriptions.join("\n"))
8421 };
8422
8423 let skills_section = if available_skills.is_empty() {
8424 "Available skills: none".to_string()
8425 } else {
8426 format!("Available skills: {}", available_skills.join(", "))
8427 };
8428
8429 let prompt = format!(
8430 r#"Create a step-by-step plan to accomplish this goal.
8431
8432Goal: "{}"
8433
8434{}
8435
8436{}
8437
8438Create a plan with clear steps. For each step, specify:
8439- description: What this step accomplishes
8440- action_type: "tool", "skill", "think", or "respond"
8441- action_target: The tool/skill id (if applicable)
8442- args: The arguments object matching the tool's schema (if action_type is "tool")
8443- dependencies: List of step IDs this depends on (empty if none)
8444
8445Respond in JSON format:
8446{{
8447 "steps": [
8448 {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8449 {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8450 ]
8451}}"#,
8452 input, tools_section, skills_section,
8453 );
8454
8455 let messages = vec![ChatMessage::user(&prompt)];
8456 let response = self
8457 .observe_purpose(
8458 ObservationPurpose::PlanGeneration,
8459 planner_llm.complete(&messages, None),
8460 )
8461 .await
8462 .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8463
8464 let mut plan = Plan::new(input);
8465
8466 if let Some(json_start) = response.content.find('{')
8467 && let Some(json_end) = response.content.rfind('}')
8468 {
8469 let json_str = &response.content[json_start..=json_end];
8470 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8471 && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8472 {
8473 for step_value in steps {
8474 let id = step_value
8475 .get("id")
8476 .and_then(|v| v.as_str())
8477 .unwrap_or("step");
8478 let desc = step_value
8479 .get("description")
8480 .and_then(|v| v.as_str())
8481 .unwrap_or("");
8482 let action_type = step_value
8483 .get("action_type")
8484 .and_then(|v| v.as_str())
8485 .unwrap_or("think");
8486 let action_target = step_value
8487 .get("action_target")
8488 .and_then(|v| v.as_str())
8489 .unwrap_or("");
8490 let args = step_value
8491 .get("args")
8492 .cloned()
8493 .unwrap_or(serde_json::json!({}));
8494 let deps: Vec<String> = step_value
8495 .get("dependencies")
8496 .and_then(|v| v.as_array())
8497 .map(|arr| {
8498 arr.iter()
8499 .filter_map(|v| v.as_str().map(String::from))
8500 .collect()
8501 })
8502 .unwrap_or_default();
8503
8504 let action = match action_type {
8505 "tool" => PlanAction::tool(action_target, args),
8506 "skill" => PlanAction::skill(action_target),
8507 "respond" => PlanAction::respond(action_target),
8508 _ => PlanAction::think(desc),
8509 };
8510
8511 let step = PlanStep::new(desc, action)
8512 .with_id(id)
8513 .with_dependencies(deps);
8514 plan.add_step(step);
8515 }
8516 }
8517 }
8518
8519 if plan.steps.is_empty() {
8520 plan.add_step(PlanStep::new(
8521 "Process the request",
8522 PlanAction::think(input),
8523 ));
8524 plan.add_step(PlanStep::new(
8525 "Provide response",
8526 PlanAction::respond("Answer based on analysis"),
8527 ));
8528 }
8529
8530 Ok(plan)
8531 }
8532
8533 async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8534 let llm = self.get_state_llm()?;
8535 let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8536 let effective = self.get_effective_reasoning_config();
8537 let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8538
8539 plan.status = PlanStatus::InProgress;
8540
8541 for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8542 let step = &plan.steps[step_idx];
8543
8544 let deps_satisfied = step.dependencies.iter().all(|dep| {
8545 plan.steps
8546 .iter()
8547 .find(|s| &s.id == dep)
8548 .map(|s| s.status.is_completed())
8549 .unwrap_or(false)
8550 });
8551
8552 if !deps_satisfied {
8553 continue;
8554 }
8555
8556 plan.steps[step_idx].mark_running();
8557
8558 let result = match &plan.steps[step_idx].action {
8559 PlanAction::Tool { tool, args } => {
8560 let has_dep_results = plan.steps[step_idx]
8566 .dependencies
8567 .iter()
8568 .any(|dep| results.contains_key(dep));
8569
8570 let final_args = if has_dep_results {
8571 let dep_context: String = plan.steps[step_idx]
8572 .dependencies
8573 .iter()
8574 .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8575 .collect::<Vec<_>>()
8576 .join("\n");
8577
8578 let tool_schema = self
8579 .tools
8580 .get(tool)
8581 .map(|t| {
8582 let schema = t.input_schema();
8583 let props = schema
8584 .get("properties")
8585 .and_then(|p| serde_json::to_string(p).ok())
8586 .unwrap_or_else(|| "{}".to_string());
8587 format!(
8588 "{}: {}\nArguments schema: {}",
8589 t.id(),
8590 t.description(),
8591 props
8592 )
8593 })
8594 .unwrap_or_default();
8595
8596 let step_desc = &plan.steps[step_idx].description;
8597 let arg_prompt = format!(
8598 "Generate the JSON arguments for a tool call.\n\n\
8599 Tool: {}\n\n\
8600 Task: {}\n\n\
8601 Previous step results:\n{}\n\n\
8602 Planner's draft arguments: {}\n\n\
8603 Produce ONLY a valid JSON object with the correct argument values.\n\
8604 Use actual values from the previous step results, not template references.",
8605 tool_schema,
8606 step_desc,
8607 dep_context,
8608 serde_json::to_string(args).unwrap_or_default()
8609 );
8610 let messages = vec![ChatMessage::user(&arg_prompt)];
8611 match self
8612 .observe_purpose(
8613 ObservationPurpose::PlanStep,
8614 llm.complete(&messages, None),
8615 )
8616 .await
8617 {
8618 Ok(resp) => {
8619 let content = resp.content.trim();
8620 let json_start = content.find('{');
8622 let json_end = content.rfind('}');
8623 if let (Some(start), Some(end)) = (json_start, json_end) {
8624 serde_json::from_str(&content[start..=end])
8625 .unwrap_or_else(|_| args.clone())
8626 } else {
8627 args.clone()
8628 }
8629 }
8630 Err(_) => args.clone(),
8631 }
8632 } else {
8633 args.clone()
8634 };
8635
8636 let request = ToolExecutionRequest::new(
8637 uuid::Uuid::new_v4().to_string(),
8638 tool.clone(),
8639 final_args,
8640 ToolCallSource::Plan {
8641 step_index: step_idx,
8642 },
8643 );
8644 match self.execute_tool_record(request).await {
8645 Ok(record) if record.success => {
8646 serde_json::json!({ "output": record.model_output_string() })
8647 }
8648 Ok(record) => {
8649 plan.steps[step_idx].mark_failed(record.model_output_string());
8650 continue;
8651 }
8652 Err(e) => {
8653 plan.steps[step_idx].mark_failed(e.to_string());
8654 continue;
8655 }
8656 }
8657 }
8658 PlanAction::Skill { skill } => {
8659 if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8660 if let Some(ref executor) = self.skill_executor {
8661 match executor
8662 .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8663 .await
8664 {
8665 Ok(output) => serde_json::json!({ "output": output }),
8666 Err(e) => {
8667 plan.steps[step_idx].mark_failed(e.to_string());
8668 continue;
8669 }
8670 }
8671 } else {
8672 serde_json::json!({ "output": "Skill executor not available" })
8673 }
8674 } else {
8675 plan.steps[step_idx].mark_failed("Skill not found");
8676 continue;
8677 }
8678 }
8679 PlanAction::Think { prompt } => {
8680 let context: String = results
8681 .iter()
8682 .map(|(k, v)| format!("{}: {}", k, v))
8683 .collect::<Vec<_>>()
8684 .join("\n");
8685
8686 let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8687 let messages = vec![ChatMessage::user(&think_prompt)];
8688
8689 match self
8690 .observe_purpose(
8691 ObservationPurpose::PlanStep,
8692 llm.complete(&messages, None),
8693 )
8694 .await
8695 {
8696 Ok(resp) => serde_json::json!({ "output": resp.content }),
8697 Err(e) => {
8698 plan.steps[step_idx].mark_failed(e.to_string());
8699 continue;
8700 }
8701 }
8702 }
8703 PlanAction::Respond { template } => {
8704 let context: String = results
8705 .iter()
8706 .map(|(k, v)| format!("{}: {}", k, v))
8707 .collect::<Vec<_>>()
8708 .join("\n");
8709
8710 let respond_prompt = format!(
8711 "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8712 context, template
8713 );
8714 let messages = vec![ChatMessage::user(&respond_prompt)];
8715
8716 match self
8717 .observe_purpose(
8718 ObservationPurpose::PlanStep,
8719 llm.complete(&messages, None),
8720 )
8721 .await
8722 {
8723 Ok(resp) => serde_json::json!({ "output": resp.content }),
8724 Err(e) => {
8725 plan.steps[step_idx].mark_failed(e.to_string());
8726 continue;
8727 }
8728 }
8729 }
8730 };
8731
8732 results.insert(plan.steps[step_idx].id.clone(), result.clone());
8733 plan.steps[step_idx].mark_completed(Some(result));
8734 }
8735
8736 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8738 if has_failures {
8739 let failed_ids: Vec<String> = plan
8740 .steps
8741 .iter()
8742 .filter(|s| s.status.is_failed())
8743 .map(|s| s.id.clone())
8744 .collect();
8745 plan.status = PlanStatus::Failed {
8746 error: format!("Steps failed: {}", failed_ids.join(", ")),
8747 };
8748 } else {
8749 plan.status = PlanStatus::Completed;
8750 }
8751
8752 let all_outputs: Vec<String> = plan
8754 .steps
8755 .iter()
8756 .filter(|s| s.status.is_completed())
8757 .filter_map(|s| {
8758 s.result
8759 .as_ref()
8760 .and_then(|r| r.get("output"))
8761 .and_then(|o| o.as_str())
8762 .map(|o| format!("{}: {}", s.description, o))
8763 })
8764 .collect();
8765
8766 if all_outputs.is_empty() {
8767 return Ok("Plan execution completed but produced no results.".to_string());
8768 }
8769
8770 if all_outputs.len() == 1 {
8771 return Ok(all_outputs.into_iter().next().unwrap());
8772 }
8773
8774 let context = all_outputs.join("\n\n");
8776 let prompt = format!(
8777 "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8778 plan.goal, context
8779 );
8780 let messages = vec![ChatMessage::user(&prompt)];
8781 match self
8782 .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8783 .await
8784 {
8785 Ok(resp) => Ok(resp.content.trim().to_string()),
8786 Err(_) => Ok(context),
8787 }
8788 }
8789
8790 async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8791 let effective_config = self.get_effective_reflection_config();
8792 self.evaluate_response_with_config(input, response, &effective_config)
8793 .await
8794 }
8795
8796 fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8797 if let Some(start) = content.find("<thinking>")
8798 && let Some(end) = content.find("</thinking>")
8799 {
8800 let thinking = content[start + 10..end].trim().to_string();
8801 let answer = content[end + 11..].trim().to_string();
8802 return (Some(thinking), answer);
8803 }
8804 (None, content.to_string())
8805 }
8806
8807 fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8808 match self.get_effective_reasoning_config().output {
8809 ReasoningOutput::Hidden => answer.to_string(),
8810 ReasoningOutput::Visible => {
8811 if let Some(t) = thinking {
8812 format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8813 } else {
8814 answer.to_string()
8815 }
8816 }
8817 ReasoningOutput::Tagged => {
8818 if let Some(t) = thinking {
8819 format!("<thinking>{}</thinking>\n{}", t, answer)
8820 } else {
8821 answer.to_string()
8822 }
8823 }
8824 }
8825 }
8826
8827 async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8832 self.init_storage().await?;
8836 self.begin_root_turn();
8837 let _root_cleanup = RootTurnCleanup::new(self);
8838 info!(input_len = input.len(), "Starting chat");
8839
8840 self.hooks.on_message_received(input).await;
8841
8842 if !self.context_initialized.swap(true, Ordering::SeqCst) {
8846 self.context_manager.initialize().await?;
8847 debug!("Context manager initialized (defaults, env, builtins)");
8848 }
8849
8850 self.check_turn_timeout().await?;
8851 self.context_manager.refresh_per_turn().await?;
8852
8853 self.clear_disambiguation_context();
8856
8857 if let Some(ref disambiguator) = self.disambiguation_manager {
8859 let disambiguation_context = self.build_disambiguation_context().await?;
8860
8861 let state_override = self
8863 .state_machine
8864 .as_ref()
8865 .and_then(|sm| sm.current_definition())
8866 .and_then(|def| def.disambiguation.clone());
8867
8868 let state_generation = self
8869 .state_machine
8870 .as_ref()
8871 .map(|state_machine| state_machine.generation());
8872 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8873 let mut disambiguation_result = self
8874 .observe_purpose(
8875 ObservationPurpose::DisambiguationDetection,
8876 disambiguator.process_input_with_override(
8877 input,
8878 &disambiguation_context,
8879 state_override.as_ref(),
8880 None,
8881 ),
8882 )
8883 .await?;
8884 let current_state_generation = self
8885 .state_machine
8886 .as_ref()
8887 .map(|state_machine| state_machine.generation());
8888 if current_state_generation != state_generation
8889 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8890 {
8891 disambiguator.clear_pending().await;
8892 *self.pending_skill_id.write() = None;
8893 disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8894 info!(
8895 confirmation_event = "invalidated",
8896 invalidation_reason = "state_generation_changed",
8897 "Disambiguation result invalidated before redispatch"
8898 );
8899 }
8900 match disambiguation_result {
8901 DisambiguationResult::Clear => {
8902 debug!("Input is clear, proceeding normally");
8903 }
8904 DisambiguationResult::NeedsClarification {
8905 question,
8906 detection,
8907 } => {
8908 let admission = self
8909 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8910 .await?;
8911 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8912 info!(
8913 ambiguity_type = ?detection.ambiguity_type,
8914 confidence = detection.confidence,
8915 "Input requires clarification"
8916 );
8917
8918 self.commit_root_user_message(input).await?;
8921 self.memory
8922 .add_message(ChatMessage::assistant(&question.question))
8923 .await?;
8924
8925 let status = if awaiting_confirmation {
8926 "awaiting_confirmation"
8927 } else {
8928 "awaiting_clarification"
8929 };
8930 let response = AgentResponse::new(&question.question).with_metadata(
8931 "disambiguation",
8932 serde_json::json!({
8933 "status": status,
8934 "options": question.options,
8935 "clarifying": question.clarifying,
8936 "detection": {
8937 "type": detection.ambiguity_type,
8938 "confidence": detection.confidence,
8939 "what_is_unclear": detection.what_is_unclear,
8940 }
8941 }),
8942 );
8943 drop(admission);
8944 self.finish_turn_if_root(&response).await?;
8945 return Ok(response);
8946 }
8947 DisambiguationResult::Clarified {
8948 enriched_input,
8949 resolved,
8950 ..
8951 } => {
8952 let admission = match self
8953 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8954 .await
8955 {
8956 Ok(admission) => admission,
8957 Err(error) => {
8958 *self.pending_skill_id.write() = None;
8959 return Err(error);
8960 }
8961 };
8962 info!(
8963 resolved_count = resolved.len(),
8964 enriched = %enriched_input,
8965 "Input clarified, injecting resolved intent into context"
8966 );
8967
8968 for (key, value) in &resolved {
8971 let context_key = format!("disambiguation.{}", key);
8972 let _ = self.context_manager.set(&context_key, value.clone());
8973 }
8974
8975 if let Some(intent) = resolved.get("intent") {
8976 let _ = self.context_manager.set("resolved_intent", intent.clone());
8977 }
8978
8979 let _ = self
8980 .context_manager
8981 .set("disambiguation.resolved", serde_json::Value::Bool(true));
8982
8983 let skill_id = self.pending_skill_id.read().clone();
8987 if let Some(skill_id) = skill_id {
8988 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8989 drop(admission);
8990 return self
8991 .recheck_skill_disambiguation(
8992 &skill_id,
8993 &enriched_input,
8994 disambiguation_epoch,
8995 state_generation,
8996 )
8997 .await;
8998 }
8999
9000 drop(admission);
9001 return self.run_loop_internal(&enriched_input).await;
9002 }
9003 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
9004 info!("Proceeding with best guess interpretation");
9005
9006 let skill_id = self.pending_skill_id.read().clone();
9008 if let Some(skill_id) = skill_id {
9009 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
9010 return self
9011 .recheck_skill_disambiguation(
9012 &skill_id,
9013 &enriched_input,
9014 disambiguation_epoch,
9015 state_generation,
9016 )
9017 .await;
9018 }
9019
9020 return self.run_loop_internal(&enriched_input).await;
9021 }
9022 DisambiguationResult::GiveUp { reason } => {
9023 *self.pending_skill_id.write() = None;
9024 warn!(reason = %reason, "Disambiguation gave up");
9025 let apology = self
9026 .generate_localized_apology(
9027 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9028 &reason,
9029 )
9030 .await
9031 .unwrap_or_else(|_| {
9032 format!("I'm sorry, I couldn't understand your request: {}", reason)
9033 });
9034 let response = AgentResponse::new(&apology);
9035 self.finish_turn_if_root(&response).await?;
9036 return Ok(response);
9037 }
9038 DisambiguationResult::Escalate { reason } => {
9039 *self.pending_skill_id.write() = None;
9040 info!(reason = %reason, "Escalating to human");
9041 if let Some(ref hitl) = self.hitl_engine {
9042 let trigger =
9043 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
9044 let mut context_map = HashMap::new();
9045 context_map.insert("original_input".to_string(), serde_json::json!(input));
9046 context_map.insert("reason".to_string(), serde_json::json!(&reason));
9047 let check_result = HITLCheckResult::required(
9048 trigger,
9049 context_map,
9050 format!("User request needs human assistance: {}", reason),
9051 Some(hitl.config().default_timeout_seconds),
9052 );
9053 let result = self.request_hitl_approval(check_result).await?;
9054 if matches!(
9055 result,
9056 ApprovalResult::Approved | ApprovalResult::Modified { .. }
9057 ) {
9058 return self.run_loop_internal(input).await;
9059 }
9060 }
9061 let apology = self
9062 .generate_localized_apology(
9063 "Explain briefly that you're transferring the user to a human agent for help.",
9064 &reason,
9065 )
9066 .await
9067 .unwrap_or_else(|_| {
9068 format!("I need human assistance to help with your request: {}", reason)
9069 });
9070 let response = AgentResponse::new(&apology);
9071 self.finish_turn_if_root(&response).await?;
9072 return Ok(response);
9073 }
9074 DisambiguationResult::Abandoned { new_input } => {
9075 *self.pending_skill_id.write() = None;
9076
9077 info!(
9078 has_new_input = new_input.is_some(),
9079 "Clarification abandoned by user"
9080 );
9081
9082 self.commit_root_user_message(input).await?;
9083
9084 match new_input {
9085 Some(fresh_input) => {
9086 return self.run_loop_internal(&fresh_input).await;
9089 }
9090 None => {
9091 let ack = self
9093 .generate_localized_apology(
9094 "The user changed their mind about their previous request. \
9095 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9096 Do NOT apologize excessively. Be concise.",
9097 "User abandoned clarification",
9098 )
9099 .await
9100 .unwrap_or_else(|_| {
9101 "OK, no problem. What else can I help with?".to_string()
9102 });
9103
9104 self.memory
9105 .add_message(ChatMessage::assistant(&ack))
9106 .await?;
9107
9108 let response = AgentResponse::new(&ack);
9109 self.finish_turn_if_root(&response).await?;
9110 return Ok(response);
9111 }
9112 }
9113 }
9114 }
9115 }
9116
9117 self.run_loop_internal(input).await
9118 }
9119
9120 async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
9122 let llm = self.llm_registry.router().map_err(|e| {
9123 AgentError::LLM(format!(
9124 "Router LLM not available for localized response: {}",
9125 e
9126 ))
9127 })?;
9128
9129 let recent: Vec<String> = self
9130 .memory
9131 .get_messages(Some(3))
9132 .await?
9133 .iter()
9134 .map(|m| m.content.clone())
9135 .collect();
9136
9137 let context_hint = if recent.is_empty() {
9138 String::new()
9139 } else {
9140 format!(
9141 "\nRecent conversation (detect the user's language from this):\n{}\n",
9142 recent.join("\n")
9143 )
9144 };
9145
9146 let prompt = format!(
9147 "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
9148 instruction, reason, context_hint
9149 );
9150
9151 let messages = vec![ChatMessage::user(&prompt)];
9152 let response = self
9153 .observe_purpose(
9154 ObservationPurpose::DisambiguationClarification,
9155 llm.complete(&messages, None),
9156 )
9157 .await
9158 .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
9159
9160 Ok(response.content.trim().to_string())
9161 }
9162
9163 fn render_action_args(&self, args: &Value) -> Value {
9167 let context = self.build_context_with_overlays();
9168 match args {
9169 Value::Object(map) => {
9170 let mut rendered = serde_json::Map::new();
9171 for (k, v) in map {
9172 match v {
9173 Value::String(s) if s.contains("{{") => {
9174 match self.template_renderer.render(s, &context) {
9175 Ok(rendered_str) => {
9176 rendered.insert(k.clone(), Value::String(rendered_str));
9177 }
9178 Err(_) => {
9179 rendered.insert(k.clone(), v.clone());
9180 }
9181 }
9182 }
9183 _ => {
9184 rendered.insert(k.clone(), v.clone());
9185 }
9186 }
9187 }
9188 Value::Object(rendered)
9189 }
9190 _ => args.clone(),
9191 }
9192 }
9193
9194 fn clear_disambiguation_context(&self) {
9196 let _ = self
9197 .context_manager
9198 .set("resolved_intent", serde_json::Value::Null);
9199
9200 let all = self.context_manager.get_all();
9201 for key in all.keys() {
9202 if key.starts_with("disambiguation.") {
9203 let _ = self.context_manager.set(key, serde_json::Value::Null);
9204 }
9205 }
9206 }
9207
9208 async fn recheck_skill_disambiguation(
9214 &self,
9215 skill_id: &str,
9216 enriched_input: &str,
9217 expected_disambiguation_epoch: u64,
9218 expected_state_generation: Option<u64>,
9219 ) -> Result<AgentResponse> {
9220 let skill = self
9221 .skill_router
9222 .as_ref()
9223 .and_then(|r| r.get_skill(skill_id).cloned());
9224
9225 if let Some(ref skill) = skill
9227 && let Some(ref skill_disambig) = skill.disambiguation
9228 && skill_disambig.enabled.unwrap_or(false)
9229 && let Some(ref disambiguator) = self.disambiguation_manager
9230 {
9231 let context = self.build_disambiguation_context().await?;
9232 let state_override = self
9233 .state_machine
9234 .as_ref()
9235 .and_then(|sm| sm.current_definition())
9236 .and_then(|def| def.disambiguation.clone());
9237
9238 let disambiguation_result = self
9239 .observe_purpose(
9240 ObservationPurpose::DisambiguationDetection,
9241 disambiguator.process_input_with_override(
9242 enriched_input,
9243 &context,
9244 state_override.as_ref(),
9245 Some(skill_disambig),
9246 ),
9247 )
9248 .await?;
9249 let current_state_generation = self
9250 .state_machine
9251 .as_ref()
9252 .map(|state_machine| state_machine.generation());
9253 if current_state_generation != expected_state_generation
9254 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
9255 {
9256 disambiguator.clear_pending().await;
9257 *self.pending_skill_id.write() = None;
9258 return Err(AgentError::Other(
9259 "State or reset ownership changed during skill disambiguation recheck"
9260 .to_string(),
9261 ));
9262 }
9263 match disambiguation_result {
9264 DisambiguationResult::Clear => {
9265 debug!(skill_id = %skill_id, "Skill re-check: all fields present");
9266 }
9267 DisambiguationResult::NeedsClarification {
9268 question,
9269 detection,
9270 } => {
9271 let admission = self
9272 .admit_disambiguation_redispatch(
9273 expected_disambiguation_epoch,
9274 expected_state_generation,
9275 )
9276 .await?;
9277 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
9278 info!(
9279 skill_id = %skill_id,
9280 ambiguity_type = ?detection.ambiguity_type,
9281 what_is_unclear = ?detection.what_is_unclear,
9282 "Skill re-check: still missing fields, asking again"
9283 );
9284 self.memory
9288 .add_message(ChatMessage::user(enriched_input))
9289 .await?;
9290 self.memory
9291 .add_message(ChatMessage::assistant(&question.question))
9292 .await?;
9293
9294 let response = AgentResponse::new(&question.question).with_metadata(
9295 "disambiguation",
9296 serde_json::json!({
9297 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
9298 "skill_id": skill_id,
9299 "options": question.options,
9300 "clarifying": question.clarifying,
9301 "detection": {
9302 "type": detection.ambiguity_type,
9303 "confidence": detection.confidence,
9304 "what_is_unclear": detection.what_is_unclear,
9305 }
9306 }),
9307 );
9308 drop(admission);
9309 self.finish_turn_if_root(&response).await?;
9310 return Ok(response);
9311 }
9312 DisambiguationResult::Clarified {
9313 enriched_input: re_enriched,
9314 ..
9315 } => {
9316 debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
9317 let admission = self
9318 .admit_disambiguation_redispatch(
9319 expected_disambiguation_epoch,
9320 expected_state_generation,
9321 )
9322 .await?;
9323 *self.pending_skill_id.write() = None;
9324 drop(admission);
9325 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9326 self.memory
9327 .add_message(ChatMessage::user(&re_enriched))
9328 .await?;
9329 return self
9330 .handle_skill_response(
9331 &re_enriched,
9332 skill_id,
9333 skill_response,
9334 &HashMap::new(),
9335 )
9336 .await;
9337 }
9338 DisambiguationResult::ProceedWithBestGuess {
9339 enriched_input: re_enriched,
9340 } => {
9341 debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9342 let admission = self
9343 .admit_disambiguation_redispatch(
9344 expected_disambiguation_epoch,
9345 expected_state_generation,
9346 )
9347 .await?;
9348 *self.pending_skill_id.write() = None;
9349 drop(admission);
9350 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9351 self.memory
9352 .add_message(ChatMessage::user(&re_enriched))
9353 .await?;
9354 return self
9355 .handle_skill_response(
9356 &re_enriched,
9357 skill_id,
9358 skill_response,
9359 &HashMap::new(),
9360 )
9361 .await;
9362 }
9363 DisambiguationResult::GiveUp { reason } => {
9364 *self.pending_skill_id.write() = None;
9365 let apology = self
9366 .generate_localized_apology(
9367 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9368 &reason,
9369 )
9370 .await
9371 .unwrap_or_else(|_| {
9372 format!("I'm sorry, I couldn't understand your request: {}", reason)
9373 });
9374 let response = AgentResponse::new(&apology);
9375 self.finish_turn_if_root(&response).await?;
9376 return Ok(response);
9377 }
9378 DisambiguationResult::Escalate { reason } => {
9379 *self.pending_skill_id.write() = None;
9380 let apology = self
9381 .generate_localized_apology(
9382 "Explain briefly that you're transferring the user to a human agent for help.",
9383 &reason,
9384 )
9385 .await
9386 .unwrap_or_else(|_| {
9387 format!("I need human assistance to help with your request: {}", reason)
9388 });
9389 let response = AgentResponse::new(&apology);
9390 self.finish_turn_if_root(&response).await?;
9391 return Ok(response);
9392 }
9393 DisambiguationResult::Abandoned { new_input } => {
9394 *self.pending_skill_id.write() = None;
9397 debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9398 if let Some(fresh) = new_input {
9399 return self.run_loop_internal(&fresh).await;
9400 }
9401 let ack = self
9402 .generate_localized_apology(
9403 "The user changed their mind about their previous request. \
9404 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9405 Do NOT apologize excessively. Be concise.",
9406 "User abandoned clarification",
9407 )
9408 .await
9409 .unwrap_or_else(|_| {
9410 "OK, no problem. What else can I help with?".to_string()
9411 });
9412 self.memory
9413 .add_message(ChatMessage::assistant(&ack))
9414 .await?;
9415 let response = AgentResponse::new(&ack);
9416 self.finish_turn_if_root(&response).await?;
9417 return Ok(response);
9418 }
9419 }
9420 }
9421
9422 let admission = self
9424 .admit_disambiguation_redispatch(
9425 expected_disambiguation_epoch,
9426 expected_state_generation,
9427 )
9428 .await?;
9429 *self.pending_skill_id.write() = None;
9430 drop(admission);
9431 let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9432 self.memory
9433 .add_message(ChatMessage::user(enriched_input))
9434 .await?;
9435 self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9436 .await
9437 }
9438
9439 async fn handle_skill_response(
9442 &self,
9443 processed_input: &str,
9444 skill_id: &str,
9445 skill_response: String,
9446 input_context: &HashMap<String, Value>,
9447 ) -> Result<AgentResponse> {
9448 let output_data = self.process_output(&skill_response, input_context).await?;
9449 let final_response = output_data.content;
9450
9451 self.memory
9452 .add_message(ChatMessage::assistant(&final_response))
9453 .await?;
9454
9455 self.check_memory_compression().await?;
9456
9457 self.increment_turn();
9458 self.evaluate_transitions(processed_input, &final_response)
9459 .await?;
9460
9461 let response = AgentResponse::new(final_response)
9462 .with_metadata("skill_id", serde_json::json!(skill_id));
9463 self.finish_turn_if_root(&response).await?;
9464 Ok(response)
9465 }
9466
9467 async fn handle_plan_and_execute(
9470 &self,
9471 processed_input: &str,
9472 input_context: &HashMap<String, Value>,
9473 auto_detected: bool,
9474 ) -> Result<AgentResponse> {
9475 let effective = self.get_effective_reasoning_config();
9476 let plan_reflection = effective
9477 .get_planning()
9478 .map(|c| c.reflection.clone())
9479 .unwrap_or_default();
9480
9481 let max_attempts = if plan_reflection.enabled {
9482 1 + plan_reflection.max_replans
9483 } else {
9484 1
9485 };
9486
9487 let mut plan = self.generate_plan(processed_input).await?;
9488 info!(
9489 plan_id = %plan.id,
9490 steps = plan.steps.len(),
9491 "Plan generated"
9492 );
9493
9494 let mut plan_result = String::new();
9495
9496 for attempt in 0..max_attempts {
9497 *self.current_plan.write() = Some(plan.clone());
9498 plan_result = self.execute_plan(&mut plan).await?;
9499
9500 info!(
9501 plan_status = ?plan.status,
9502 completed_steps = plan.completed_steps().count(),
9503 attempt = attempt + 1,
9504 "Plan execution completed"
9505 );
9506
9507 if !plan_reflection.enabled {
9508 break;
9509 }
9510
9511 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9512 if !has_failures {
9513 break;
9514 }
9515
9516 if attempt + 1 >= max_attempts {
9517 break;
9518 }
9519
9520 match plan_reflection.on_step_failure {
9521 StepFailureAction::Replan => {
9522 info!(attempt = attempt + 1, "Plan had failures, replanning");
9523 plan = self.generate_plan(processed_input).await?;
9524 }
9525 StepFailureAction::Abort => {
9526 warn!("Plan step failed, aborting");
9527 break;
9528 }
9529 StepFailureAction::Skip | StepFailureAction::Continue => {
9530 break;
9531 }
9532 }
9533 }
9534
9535 *self.current_plan.write() = Some(plan);
9536
9537 let output_data = self.process_output(&plan_result, input_context).await?;
9538 let final_content = output_data.content;
9539
9540 self.memory
9541 .add_message(ChatMessage::assistant(&final_content))
9542 .await?;
9543
9544 self.check_memory_compression().await?;
9545 self.increment_turn();
9546 self.evaluate_transitions(processed_input, &final_content)
9547 .await?;
9548
9549 let reasoning_metadata =
9550 ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9551
9552 let response = AgentResponse::new(&final_content).with_metadata(
9553 "reasoning",
9554 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9555 );
9556
9557 self.finish_turn_if_root(&response).await?;
9558 Ok(response)
9559 }
9560
9561 fn inject_reasoning_prompt(
9563 &self,
9564 messages: &mut [ChatMessage],
9565 reasoning_mode: &ReasoningMode,
9566 is_first_iteration: bool,
9567 ) {
9568 if !is_first_iteration {
9569 return;
9570 }
9571 match reasoning_mode {
9572 ReasoningMode::CoT => {
9573 if let Some(msg) = messages.first_mut()
9574 && matches!(msg.role, ai_agents_core::Role::System)
9575 {
9576 msg.content = self.build_cot_system_prompt(&msg.content);
9577 debug!("Applied Chain-of-Thought system prompt");
9578 }
9579 }
9580 ReasoningMode::React => {
9581 if let Some(msg) = messages.first_mut()
9582 && matches!(msg.role, ai_agents_core::Role::System)
9583 {
9584 msg.content = self.build_react_system_prompt(&msg.content);
9585 debug!("Applied ReAct system prompt");
9586 }
9587 }
9588 _ => {}
9589 }
9590 }
9591
9592 async fn generate_main_response_draft(
9597 &self,
9598 processed_input: &str,
9599 reasoning_mode: &ReasoningMode,
9600 ) -> Result<MainResponseDraft> {
9601 let llm = self.get_state_llm()?;
9602 let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9603 let mut messages = self
9604 .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9605 .await?;
9606 self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9607 let response = self
9608 .complete_main_llm_with_recovery(llm, &messages, &protocol)
9609 .await?;
9610 let content = response.content.trim().to_string();
9611 let (thinking, answer) = self.extract_thinking(&content);
9612 if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9613 return Ok(MainResponseDraft::ToolCalls {
9614 raw_content: content,
9615 calls,
9616 thinking,
9617 });
9618 }
9619 Ok(MainResponseDraft::Text {
9620 raw_content: answer,
9621 thinking,
9622 })
9623 }
9624
9625 async fn commit_main_response_draft(
9630 &self,
9631 processed_input: &str,
9632 input_context: &HashMap<String, Value>,
9633 draft: MainResponseDraft,
9634 reasoning_mode: ReasoningMode,
9635 auto_detected: bool,
9636 ) -> Result<AgentResponse> {
9637 self.commit_root_user_message(processed_input).await?;
9638 match draft {
9639 MainResponseDraft::Text {
9640 raw_content,
9641 thinking,
9642 } => {
9643 self.finish_text_response_from_model(CommittedTextResponse {
9644 processed_input,
9645 input_context,
9646 answer: raw_content,
9647 reasoning_mode,
9648 auto_detected,
9649 iterations: 1,
9650 thinking_content: thinking,
9651 all_tool_calls: Vec::new(),
9652 })
9653 .await
9654 }
9655 MainResponseDraft::ToolCalls {
9656 raw_content,
9657 calls,
9658 thinking: _,
9659 } => {
9660 let mut all_tool_calls = Vec::new();
9661 match self
9662 .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9663 .await?
9664 {
9665 ToolCallOutcome::Rejected(response) => {
9666 self.finish_turn_if_root(&response).await?;
9667 Ok(response)
9668 }
9669 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9670 self.continue_after_committed_tool_draft(processed_input)
9671 .await
9672 }
9673 }
9674 }
9675 }
9676 }
9677
9678 async fn continue_after_committed_tool_draft(
9683 &self,
9684 processed_input: &str,
9685 ) -> Result<AgentResponse> {
9686 *self.redispatch_depth.write() += 1;
9687 if let Some(context) = self.active_turn_context.write().as_mut() {
9688 context.enter_redispatch();
9689 }
9690 let result = Box::pin(self.run_loop_internal(processed_input)).await;
9691 *self.redispatch_depth.write() -= 1;
9692 if let Some(context) = self.active_turn_context.write().as_mut() {
9693 context.exit_redispatch();
9694 }
9695 let response = result?;
9696 self.finish_turn_if_root(&response).await?;
9697 Ok(response)
9698 }
9699
9700 async fn finish_text_response_from_model(
9705 &self,
9706 response: CommittedTextResponse<'_>,
9707 ) -> Result<AgentResponse> {
9708 let CommittedTextResponse {
9709 processed_input,
9710 input_context,
9711 answer,
9712 reasoning_mode,
9713 auto_detected,
9714 iterations,
9715 thinking_content,
9716 all_tool_calls,
9717 } = response;
9718 let output_data = self.process_output(&answer, input_context).await?;
9719 let mut final_content = if output_data.metadata.rejected {
9720 output_data
9721 .metadata
9722 .rejection_reason
9723 .unwrap_or_else(|| answer.to_string())
9724 } else {
9725 output_data.content
9726 };
9727 let llm = self.get_state_llm()?;
9728 let reflection_metadata;
9729 (final_content, reflection_metadata) = self
9730 .run_reflection(&*llm, processed_input, final_content)
9731 .await?;
9732 final_content =
9733 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9734 let final_content = {
9735 let result = self
9736 .post_loop_processing(processed_input, final_content)
9737 .await?;
9738 self.apply_post_loop_result(processed_input, result).await?
9739 };
9740 let response = self.build_agent_response(AgentResponseParts {
9741 content: final_content,
9742 all_tool_calls,
9743 reasoning_mode,
9744 auto_detected,
9745 iterations,
9746 thinking: thinking_content,
9747 reflection_metadata,
9748 });
9749 self.finish_turn_if_root(&response).await?;
9750 Ok(response)
9751 }
9752
9753 async fn run_committed_response_loop_with_reasoning(
9758 &self,
9759 processed_input: &str,
9760 input_context: &HashMap<String, Value>,
9761 reasoning_mode: ReasoningMode,
9762 auto_detected: bool,
9763 ) -> Result<AgentResponse> {
9764 self.commit_root_user_message(processed_input).await?;
9765 let llm = self.get_state_llm()?;
9766 let mut iterations = 0u32;
9767 let mut all_tool_calls = Vec::new();
9768 let mut thinking_content = None;
9769 loop {
9770 let effective_max = if reasoning_mode != ReasoningMode::None {
9771 let rc = self.get_effective_reasoning_config();
9772 self.max_iterations.min(rc.max_iterations)
9773 } else {
9774 self.max_iterations
9775 };
9776 if iterations >= effective_max {
9777 return Err(AgentError::Other(format!(
9778 "Max iterations ({}) exceeded",
9779 effective_max
9780 )));
9781 }
9782 iterations += 1;
9783 *self.iteration_count.write() = iterations;
9784 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9785 let mut messages = self
9786 .build_messages_internal(true, None, protocol.choice.is_none())
9787 .await?;
9788 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9789 self.hooks.on_llm_start(&messages).await;
9790 let llm_start = Instant::now();
9791 let response = self
9792 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9793 .await?;
9794 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9795 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9796 let content = response.content.trim();
9797 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9798 match self
9799 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9800 .await?
9801 {
9802 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9803 ToolCallOutcome::Rejected(resp) => {
9804 self.finish_turn_if_root(&resp).await?;
9805 return Ok(resp);
9806 }
9807 }
9808 }
9809 let (extracted_thinking, answer) = self.extract_thinking(content);
9810 if extracted_thinking.is_some() {
9811 thinking_content = extracted_thinking;
9812 }
9813 return self
9814 .finish_text_response_from_model(CommittedTextResponse {
9815 processed_input,
9816 input_context,
9817 answer,
9818 reasoning_mode,
9819 auto_detected,
9820 iterations,
9821 thinking_content,
9822 all_tool_calls,
9823 })
9824 .await;
9825 }
9826 }
9827
9828 async fn handle_tool_calls(
9830 &self,
9831 processed_input: &str,
9832 content: &str,
9833 tool_calls: Vec<ToolCall>,
9834 all_tool_calls: &mut Vec<ToolCall>,
9835 ) -> Result<ToolCallOutcome> {
9836 let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9840 if transition_fired {
9841 self.memory
9842 .add_message(ChatMessage::assistant(
9843 "(Transitioned to new state — tool call handled by workflow)",
9844 ))
9845 .await?;
9846 return Ok(ToolCallOutcome::TransitionFired);
9847 }
9848
9849 self.memory
9851 .add_message(ChatMessage::assistant(content))
9852 .await?;
9853 let native_tool_call = Self::is_native_tool_call_content(content);
9854
9855 let results = self.execute_tools_parallel(&tool_calls).await;
9856
9857 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9858 match result {
9859 Ok(output) => {
9860 self.memory
9861 .add_message(Self::tool_result_message(
9862 tool_call,
9863 &output,
9864 native_tool_call,
9865 ))
9866 .await?;
9867 }
9868 Err(e) => {
9869 if matches!(e, AgentError::HITLRejected(_)) {
9871 self.memory
9872 .add_message(ChatMessage::assistant(format!(
9873 "The operation was rejected by the approver: {}",
9874 e
9875 )))
9876 .await?;
9877 return Ok(ToolCallOutcome::Rejected(AgentResponse {
9879 content: format!("Operation cancelled: {}", e),
9880 metadata: None,
9881 tool_calls: Some(all_tool_calls.clone()),
9882 }));
9883 }
9884 self.memory
9885 .add_message(Self::tool_result_message(
9886 tool_call,
9887 &format!("Error: {}", e),
9888 native_tool_call,
9889 ))
9890 .await?;
9891 }
9892 }
9893 all_tool_calls.push(tool_call.clone());
9894 }
9895 Ok(ToolCallOutcome::Continue)
9896 }
9897
9898 async fn run_reflection(
9900 &self,
9901 llm: &dyn LLMProvider,
9902 processed_input: &str,
9903 mut content: String,
9904 ) -> Result<(String, Option<ReflectionMetadata>)> {
9905 let should_reflect = self.should_reflect(processed_input, &content).await?;
9906 if !should_reflect {
9907 return Ok((content, None));
9908 }
9909
9910 info!("Starting response reflection evaluation");
9911 let mut attempts = 0u32;
9912 let max_retries = self.reflection_config.max_retries;
9913 let mut history: Vec<ReflectionAttempt> = Vec::new();
9914
9915 loop {
9916 let evaluation = self.evaluate_response(processed_input, &content).await?;
9917
9918 if evaluation.passed || attempts >= max_retries {
9919 info!(
9920 passed = evaluation.passed,
9921 confidence = evaluation.confidence,
9922 attempts = attempts + 1,
9923 "Reflection evaluation complete"
9924 );
9925 let reflection_metadata = Some(
9926 ReflectionMetadata::new(evaluation)
9927 .with_attempts(attempts + 1)
9928 .with_history(history),
9929 );
9930 return Ok((content, reflection_metadata));
9931 }
9932
9933 debug!(
9934 attempt = attempts + 1,
9935 failed_criteria = evaluation.failed_criteria().count(),
9936 "Response did not meet criteria, retrying"
9937 );
9938
9939 history.push(
9940 ReflectionAttempt::new(&content, evaluation.clone())
9941 .with_feedback("Response did not meet quality criteria"),
9942 );
9943
9944 let feedback: Vec<String> = evaluation
9945 .failed_criteria()
9946 .map(|c| format!("- {}", c.criterion))
9947 .collect();
9948
9949 let retry_prompt = format!(
9950 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9951 feedback.join("\n")
9952 );
9953
9954 self.memory
9955 .add_message(ChatMessage::user(&retry_prompt))
9956 .await?;
9957
9958 let retry_messages = self.build_messages().await?;
9959 let retry_response = self
9960 .observe_purpose(
9961 ObservationPurpose::ReflectionEvaluation,
9962 llm.complete(&retry_messages, None),
9963 )
9964 .await
9965 .map_err(|e| AgentError::LLM(e.to_string()))?;
9966
9967 content = retry_response.content.trim().to_string();
9968 attempts += 1;
9969 }
9970 }
9971
9972 async fn post_loop_processing(
9975 &self,
9976 processed_input: &str,
9977 content: String,
9978 ) -> Result<PostLoopResult> {
9979 self.increment_turn();
9984
9985 self.run_context_extractors(processed_input).await;
9987
9988 let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9989
9990 if !transitioned {
9991 self.memory
9992 .add_message(ChatMessage::assistant(&content))
9993 .await?;
9994 self.check_memory_compression().await?;
9995 return Ok(PostLoopResult::NoTransition(content));
9996 }
9997
9998 if !self.should_regenerate_after_transition() {
10000 self.memory
10001 .add_message(ChatMessage::assistant(&content))
10002 .await?;
10003 self.check_memory_compression().await?;
10004 return Ok(PostLoopResult::Transitioned(content));
10005 }
10006
10007 if self.needs_redispatch_for_new_state() {
10011 info!("Post-transition NeedsRedispatch: new state requires full dispatch");
10012 return Ok(PostLoopResult::NeedsRedispatch);
10015 }
10016
10017 self.memory
10020 .add_message(ChatMessage::assistant(&content))
10021 .await?;
10022 self.check_memory_compression().await?;
10023
10024 let new_llm = self.get_state_llm()?;
10030 let mut final_content;
10031
10032 for post_iter in 0..self.max_iterations {
10033 let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
10034 let new_messages = self
10035 .build_messages_internal(true, None, protocol.choice.is_none())
10036 .await?;
10037 if post_iter == 0
10038 && let Some(system_msg) = new_messages.first()
10039 && system_msg.role == ai_agents_core::Role::System
10040 {
10041 debug!(
10042 prompt_preview =
10043 &system_msg.content[system_msg.content.len().saturating_sub(200)..],
10044 "Post-transition system prompt (last 200 chars)"
10045 );
10046 }
10047
10048 let new_response = self
10049 .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
10050 .await?;
10051 final_content = new_response.content.trim().to_string();
10052
10053 if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
10056 let native_tool_call = Self::is_native_tool_call_content(&final_content);
10057 debug!(
10058 post_iter = post_iter,
10059 tools = tool_calls.len(),
10060 "Post-transition tool call detected, executing"
10061 );
10062
10063 self.memory
10064 .add_message(ChatMessage::assistant(&final_content))
10065 .await?;
10066
10067 let results = self.execute_tools_parallel(&tool_calls).await;
10068 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
10069 match result {
10070 Ok(output) => {
10071 self.memory
10072 .add_message(Self::tool_result_message(
10073 tool_call,
10074 &output,
10075 native_tool_call,
10076 ))
10077 .await?;
10078 }
10079 Err(e) => {
10080 self.memory
10081 .add_message(Self::tool_result_message(
10082 tool_call,
10083 &format!("Error: {}", e),
10084 native_tool_call,
10085 ))
10086 .await?;
10087 }
10088 }
10089 }
10090 continue;
10092 }
10093
10094 self.memory
10096 .add_message(ChatMessage::assistant(&final_content))
10097 .await?;
10098 return Ok(PostLoopResult::Transitioned(final_content));
10099 }
10100
10101 final_content = "Post-transition processing completed.".to_string();
10103 self.memory
10104 .add_message(ChatMessage::assistant(&final_content))
10105 .await?;
10106
10107 Ok(PostLoopResult::Transitioned(final_content))
10108 }
10109
10110 fn should_regenerate_after_transition(&self) -> bool {
10113 if let Some(ref sm) = self.state_machine {
10114 if !sm.config().regenerate_on_transition {
10116 return false;
10117 }
10118 if let Some(def) = sm.current_definition()
10120 && let Some(regen) = def.regenerate_on_enter
10121 {
10122 return regen;
10123 }
10124 }
10125 true
10126 }
10127
10128 fn needs_redispatch_for_new_state(&self) -> bool {
10131 if let Some(ref sm) = self.state_machine
10132 && let Some(def) = sm.current_definition()
10133 {
10134 if def.concurrent.is_some()
10135 || def.group_chat.is_some()
10136 || def.pipeline.is_some()
10137 || def.handoff.is_some()
10138 || def.delegate.is_some()
10139 {
10140 return true;
10141 }
10142 let effective = self.get_effective_reasoning_config();
10144 if !matches!(effective.mode, ReasoningMode::None) {
10145 return true;
10146 }
10147 }
10148 false
10149 }
10150
10151 async fn apply_post_loop_result(
10154 &self,
10155 processed_input: &str,
10156 result: PostLoopResult,
10157 ) -> Result<String> {
10158 match result {
10159 PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
10160 Ok(content)
10161 }
10162 PostLoopResult::NeedsRedispatch => {
10163 const MAX_REDISPATCH_DEPTH: u32 = 3;
10164 let current_depth = *self.redispatch_depth.read();
10165 if current_depth >= MAX_REDISPATCH_DEPTH {
10166 warn!(
10167 depth = current_depth,
10168 "Post-transition re-dispatch depth limit reached, returning empty response"
10169 );
10170 let content = String::new();
10171 self.memory
10172 .add_message(ChatMessage::assistant(&content))
10173 .await?;
10174 return Ok(content);
10175 }
10176 *self.redispatch_depth.write() += 1;
10177 if let Some(context) = self.active_turn_context.write().as_mut() {
10178 context.enter_redispatch();
10179 }
10180 info!(
10181 depth = current_depth + 1,
10182 "Re-dispatching for new state after transition"
10183 );
10184 let resp = Box::pin(self.run_loop_internal(processed_input)).await;
10185 *self.redispatch_depth.write() -= 1;
10186 if let Some(context) = self.active_turn_context.write().as_mut() {
10187 context.exit_redispatch();
10188 }
10189 resp.map(|r| r.content)
10190 }
10191 }
10192 }
10193
10194 fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
10196 let AgentResponseParts {
10197 content,
10198 all_tool_calls,
10199 reasoning_mode,
10200 auto_detected,
10201 iterations,
10202 thinking,
10203 reflection_metadata,
10204 } = parts;
10205 let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
10206 .with_thinking(thinking.clone().unwrap_or_default())
10207 .with_iterations(iterations)
10208 .with_auto_detected(auto_detected);
10209
10210 let mut response = AgentResponse::new(&content);
10211 if !all_tool_calls.is_empty() {
10212 response = response.with_tool_calls(all_tool_calls);
10213 }
10214
10215 if let Some(state) = self.current_state() {
10216 response = response.with_metadata("current_state", serde_json::json!(state));
10217 }
10218
10219 response = response.with_metadata(
10220 "reasoning",
10221 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
10222 );
10223
10224 if let Some(ref refl_meta) = reflection_metadata {
10225 response = response.with_metadata(
10226 "reflection",
10227 serde_json::to_value(refl_meta).unwrap_or_default(),
10228 );
10229 }
10230
10231 response
10232 }
10233
10234 async fn handle_delegated_state(
10236 &self,
10237 input: &str,
10238 delegate_id: &str,
10239 state_def: &ai_agents_state::StateDefinition,
10240 ) -> Result<AgentResponse> {
10241 use std::time::Instant;
10242
10243 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10244 AgentError::Config(format!(
10245 "State delegates to '{}' but no agent registry is configured. \
10246 Add a spawner section with auto_spawn to your YAML.",
10247 delegate_id
10248 ))
10249 })?;
10250
10251 let state_name = self
10252 .state_machine
10253 .as_ref()
10254 .map(|sm| sm.current())
10255 .unwrap_or_else(|| "unknown".to_string());
10256
10257 self.hooks.on_delegate_start(delegate_id, &state_name).await;
10258 let start = Instant::now();
10259
10260 let delegate = registry.get(delegate_id).ok_or_else(|| {
10261 AgentError::Other(format!(
10262 "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
10263 state_name, delegate_id
10264 ))
10265 })?;
10266
10267 let context_mode = state_def.delegate_context.clone().unwrap_or_default();
10269 let effective_input = self
10270 .observe_purpose(
10271 ObservationPurpose::OrchestrationRouting,
10272 crate::orchestration::context::prepare_delegate_input(
10273 input,
10274 &context_mode,
10275 &*self.memory,
10276 self.llm_registry.get("router").ok().as_deref(),
10277 ),
10278 )
10279 .await?;
10280
10281 let response = delegate
10282 .chat_with_actor_context(&effective_input, self.outbound_actor_context())
10283 .await?;
10284
10285 let duration_ms = start.elapsed().as_millis() as u64;
10286 self.hooks
10287 .on_delegate_complete(delegate_id, &state_name, duration_ms)
10288 .await;
10289
10290 let ctx_key = format!("delegation.{}.last_response", delegate_id);
10292 let _ = self.context_manager.set(
10293 &ctx_key,
10294 serde_json::Value::String(response.content.clone()),
10295 );
10296
10297 let _ = self.context_manager.set(
10299 "orchestration",
10300 serde_json::json!({
10301 "type": "delegate",
10302 "agent": delegate_id,
10303 "state": state_name,
10304 "response": response.content,
10305 "duration_ms": duration_ms,
10306 }),
10307 );
10308
10309 self.commit_root_user_message(input).await?;
10310
10311 let post_result = self
10314 .post_loop_processing(
10315 input,
10316 format!("[Delegated to {}]: {}", delegate_id, response.content),
10317 )
10318 .await?;
10319 let final_content = self.apply_post_loop_result(input, post_result).await?;
10320
10321 let mut result = AgentResponse::new(final_content);
10322
10323 let metadata = serde_json::json!({
10324 "orchestration": {
10325 "type": "delegate",
10326 "agent": delegate_id,
10327 "state": state_name,
10328 "response": response.content,
10329 "duration_ms": duration_ms,
10330 }
10331 });
10332 result.metadata = Some(
10333 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10334 metadata,
10335 )
10336 .unwrap_or_default(),
10337 );
10338
10339 self.finish_turn_if_root(&result).await?;
10340 Ok(result)
10341 }
10342
10343 async fn handle_concurrent_state(
10345 &self,
10346 input: &str,
10347 config: &ai_agents_state::ConcurrentStateConfig,
10348 ) -> Result<AgentResponse> {
10349 use std::time::Instant;
10350
10351 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10352 AgentError::Config(
10353 "Concurrent state requires an agent registry. Add a spawner section.".into(),
10354 )
10355 })?;
10356
10357 let context_mode = config.context_mode.clone().unwrap_or_default();
10362 let context_input = self
10363 .observe_purpose(
10364 ObservationPurpose::OrchestrationRouting,
10365 crate::orchestration::context::prepare_delegate_input(
10366 input,
10367 &context_mode,
10368 &*self.memory,
10369 self.llm_registry.get("router").ok().as_deref(),
10370 ),
10371 )
10372 .await?;
10373
10374 let effective_input = if let Some(ref tmpl) = config.input {
10375 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10376 .unwrap_or_else(|_| context_input.clone())
10377 } else {
10378 context_input
10379 };
10380
10381 let start = Instant::now();
10382
10383 let llm_name = config
10384 .aggregation
10385 .synthesizer_llm
10386 .as_deref()
10387 .unwrap_or("router");
10388 let llm_provider = self.llm_registry.get(llm_name).ok();
10389
10390 let vote_parallelism = if self.runtime_config.optimization.enabled
10391 && self
10392 .runtime_config
10393 .optimization
10394 .parallel_orchestration_vote_extraction
10395 {
10396 Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10397 } else {
10398 None
10399 };
10400
10401 let result = self
10402 .observe_purpose(
10403 ObservationPurpose::OrchestrationAggregation,
10404 scope_actor_context(
10405 self.outbound_actor_context(),
10406 crate::orchestration::concurrent(
10407 registry,
10408 &effective_input,
10409 &config.agents,
10410 &config.aggregation,
10411 llm_provider.as_deref(),
10412 config.min_required,
10413 config.timeout_ms,
10414 config.on_partial_failure.clone(),
10415 vote_parallelism,
10416 ),
10417 ),
10418 )
10419 .await?;
10420
10421 let duration_ms = start.elapsed().as_millis() as u64;
10422 let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10423 let strategy = format!("{:?}", config.aggregation.strategy);
10424 self.hooks
10425 .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10426 .await;
10427
10428 let _ = self.context_manager.set(
10430 "concurrent.result",
10431 serde_json::Value::String(result.response.content.clone()),
10432 );
10433
10434 let agents_json: Vec<serde_json::Value> = result
10436 .agent_results
10437 .iter()
10438 .map(|ar| {
10439 serde_json::json!({
10440 "id": ar.agent_id,
10441 "response": ar.response.as_ref().map(|r| r.content.as_str()),
10442 "success": ar.success,
10443 "error": ar.error,
10444 "duration_ms": ar.duration_ms,
10445 })
10446 })
10447 .collect();
10448
10449 let _ = self.context_manager.set(
10451 "orchestration",
10452 serde_json::json!({
10453 "type": "concurrent",
10454 "result": result.response.content,
10455 "strategy": strategy,
10456 "agents": agents_json,
10457 "duration_ms": duration_ms,
10458 }),
10459 );
10460
10461 self.commit_root_user_message(input).await?;
10462
10463 let post_result = self
10464 .post_loop_processing(input, result.response.content.clone())
10465 .await?;
10466 let final_content = self.apply_post_loop_result(input, post_result).await?;
10467
10468 let mut response = AgentResponse::new(final_content);
10469 let metadata = serde_json::json!({
10470 "orchestration": {
10471 "type": "concurrent",
10472 "result": result.response.content,
10473 "strategy": strategy,
10474 "agents": agents_json,
10475 "duration_ms": duration_ms,
10476 }
10477 });
10478 response.metadata = Some(
10479 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10480 metadata,
10481 )
10482 .unwrap_or_default(),
10483 );
10484
10485 self.finish_turn_if_root(&response).await?;
10486 Ok(response)
10487 }
10488
10489 async fn handle_group_chat_state(
10491 &self,
10492 input: &str,
10493 config: &ai_agents_state::GroupChatStateConfig,
10494 ) -> Result<AgentResponse> {
10495 use std::time::Instant;
10496
10497 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10498 AgentError::Config(
10499 "Group chat state requires an agent registry. Add a spawner section.".into(),
10500 )
10501 })?;
10502
10503 let start = Instant::now();
10504
10505 let llm_provider = self.llm_registry.get("router").ok();
10506
10507 let context_mode = config.context_mode.clone().unwrap_or_default();
10509 let context_input = self
10510 .observe_purpose(
10511 ObservationPurpose::OrchestrationRouting,
10512 crate::orchestration::context::prepare_delegate_input(
10513 input,
10514 &context_mode,
10515 &*self.memory,
10516 self.llm_registry.get("router").ok().as_deref(),
10517 ),
10518 )
10519 .await?;
10520
10521 let effective_topic = if let Some(ref tmpl) = config.input {
10523 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10524 .unwrap_or_else(|_| context_input.clone())
10525 } else {
10526 context_input
10527 };
10528
10529 let result = self
10530 .observe_purpose(
10531 ObservationPurpose::OrchestrationConversation,
10532 scope_actor_context(
10533 self.outbound_actor_context(),
10534 crate::orchestration::group_chat(
10535 registry,
10536 &effective_topic,
10537 config,
10538 llm_provider.as_deref(),
10539 Some(&*self.hooks),
10540 ),
10541 ),
10542 )
10543 .await?;
10544
10545 let duration_ms = start.elapsed().as_millis() as u64;
10546
10547 let _ = self.context_manager.set(
10549 "group_chat.conclusion",
10550 serde_json::Value::String(result.response.content.clone()),
10551 );
10552
10553 let transcript_json: Vec<serde_json::Value> = result
10555 .transcript
10556 .iter()
10557 .map(|t| {
10558 serde_json::json!({
10559 "speaker": t.speaker,
10560 "round": t.round,
10561 "content": t.content,
10562 })
10563 })
10564 .collect();
10565
10566 let _ = self.context_manager.set(
10568 "orchestration",
10569 serde_json::json!({
10570 "type": "group_chat",
10571 "conclusion": result.response.content,
10572 "transcript": transcript_json,
10573 "rounds": result.rounds_completed,
10574 "termination": result.termination_reason,
10575 "duration_ms": duration_ms,
10576 }),
10577 );
10578
10579 self.commit_root_user_message(input).await?;
10580
10581 let post_result = self
10582 .post_loop_processing(input, result.response.content.clone())
10583 .await?;
10584 let final_content = self.apply_post_loop_result(input, post_result).await?;
10585
10586 let mut response = AgentResponse::new(final_content);
10587 let metadata = serde_json::json!({
10588 "orchestration": {
10589 "type": "group_chat",
10590 "conclusion": result.response.content,
10591 "transcript": transcript_json,
10592 "rounds": result.rounds_completed,
10593 "termination": result.termination_reason,
10594 "duration_ms": duration_ms,
10595 }
10596 });
10597 response.metadata = Some(
10598 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10599 metadata,
10600 )
10601 .unwrap_or_default(),
10602 );
10603
10604 self.finish_turn_if_root(&response).await?;
10605 Ok(response)
10606 }
10607
10608 async fn handle_pipeline_state(
10610 &self,
10611 input: &str,
10612 config: &ai_agents_state::PipelineStateConfig,
10613 ) -> Result<AgentResponse> {
10614 use std::time::Instant;
10615
10616 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10617 AgentError::Config(
10618 "Pipeline state requires an agent registry. Add a spawner section.".into(),
10619 )
10620 })?;
10621
10622 let start = Instant::now();
10623
10624 let stages: Vec<crate::orchestration::PipelineStage> = config
10625 .stages
10626 .iter()
10627 .map(|entry| {
10628 let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10629 if let Some(tmpl) = entry.input() {
10630 stage = stage.with_input(tmpl);
10631 }
10632 stage
10633 })
10634 .collect();
10635
10636 let context_mode = config.context_mode.clone().unwrap_or_default();
10638 let context_input = self
10639 .observe_purpose(
10640 ObservationPurpose::OrchestrationRouting,
10641 crate::orchestration::context::prepare_delegate_input(
10642 input,
10643 &context_mode,
10644 &*self.memory,
10645 self.llm_registry.get("router").ok().as_deref(),
10646 ),
10647 )
10648 .await?;
10649
10650 let context_values = self.build_context_with_overlays();
10651 let result = self
10652 .observe_purpose(
10653 ObservationPurpose::OrchestrationRouting,
10654 scope_actor_context(
10655 self.outbound_actor_context(),
10656 crate::orchestration::pipeline(
10657 registry,
10658 &context_input,
10659 &stages,
10660 config.timeout_ms,
10661 Some(&*self.hooks),
10662 Some(&context_values),
10663 ),
10664 ),
10665 )
10666 .await?;
10667
10668 let duration_ms = start.elapsed().as_millis() as u64;
10669
10670 let _ = self.context_manager.set(
10672 "pipeline.result",
10673 serde_json::Value::String(result.response.content.clone()),
10674 );
10675
10676 let stages_json: Vec<serde_json::Value> = result
10678 .stage_outputs
10679 .iter()
10680 .map(|s| {
10681 serde_json::json!({
10682 "agent_id": s.agent_id,
10683 "output": s.output,
10684 "duration_ms": s.duration_ms,
10685 "skipped": s.skipped,
10686 })
10687 })
10688 .collect();
10689
10690 let _ = self.context_manager.set(
10692 "orchestration",
10693 serde_json::json!({
10694 "type": "pipeline",
10695 "result": result.response.content,
10696 "stages": stages_json,
10697 "duration_ms": duration_ms,
10698 }),
10699 );
10700
10701 self.commit_root_user_message(input).await?;
10702
10703 let post_result = self
10704 .post_loop_processing(input, result.response.content.clone())
10705 .await?;
10706 let final_content = self.apply_post_loop_result(input, post_result).await?;
10707
10708 let mut response = AgentResponse::new(final_content);
10709 let metadata = serde_json::json!({
10710 "orchestration": {
10711 "type": "pipeline",
10712 "result": result.response.content,
10713 "stages": stages_json,
10714 "duration_ms": duration_ms,
10715 }
10716 });
10717 response.metadata = Some(
10718 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10719 metadata,
10720 )
10721 .unwrap_or_default(),
10722 );
10723
10724 self.finish_turn_if_root(&response).await?;
10725 Ok(response)
10726 }
10727
10728 async fn handle_handoff_state(
10730 &self,
10731 input: &str,
10732 config: &ai_agents_state::HandoffStateConfig,
10733 ) -> Result<AgentResponse> {
10734 use std::time::Instant;
10735
10736 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10737 AgentError::Config(
10738 "Handoff state requires an agent registry. Add a spawner section.".into(),
10739 )
10740 })?;
10741
10742 let llm = self
10743 .llm_registry
10744 .get("router")
10745 .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10746
10747 let start = Instant::now();
10748
10749 let context_mode = config.context_mode.clone().unwrap_or_default();
10751 let context_input = self
10752 .observe_purpose(
10753 ObservationPurpose::OrchestrationRouting,
10754 crate::orchestration::context::prepare_delegate_input(
10755 input,
10756 &context_mode,
10757 &*self.memory,
10758 self.llm_registry.get("router").ok().as_deref(),
10759 ),
10760 )
10761 .await?;
10762
10763 let effective_input = if let Some(ref tmpl) = config.input {
10765 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10766 .unwrap_or_else(|_| context_input.clone())
10767 } else {
10768 context_input
10769 };
10770
10771 let result = self
10772 .observe_purpose(
10773 ObservationPurpose::OrchestrationRouting,
10774 scope_actor_context(
10775 self.outbound_actor_context(),
10776 crate::orchestration::handoff(
10777 registry,
10778 &effective_input,
10779 &config.initial_agent,
10780 &config.available_agents,
10781 config.max_handoffs,
10782 llm.as_ref(),
10783 Some(&*self.hooks),
10784 ),
10785 ),
10786 )
10787 .await?;
10788
10789 let duration_ms = start.elapsed().as_millis() as u64;
10790
10791 let _ = self.context_manager.set(
10793 "handoff.result",
10794 serde_json::Value::String(result.response.content.clone()),
10795 );
10796
10797 let chain_json: Vec<serde_json::Value> = result
10799 .handoff_chain
10800 .iter()
10801 .map(|h| {
10802 serde_json::json!({
10803 "from": h.from_agent,
10804 "to": h.to_agent,
10805 "reason": h.reason,
10806 })
10807 })
10808 .collect();
10809
10810 let _ = self.context_manager.set(
10812 "orchestration",
10813 serde_json::json!({
10814 "type": "handoff",
10815 "result": result.response.content,
10816 "final_agent": result.final_agent,
10817 "handoff_chain": chain_json,
10818 "duration_ms": duration_ms,
10819 }),
10820 );
10821
10822 self.commit_root_user_message(input).await?;
10823
10824 let post_result = self
10825 .post_loop_processing(input, result.response.content.clone())
10826 .await?;
10827 let final_content = self.apply_post_loop_result(input, post_result).await?;
10828
10829 let mut response = AgentResponse::new(final_content);
10830 let metadata = serde_json::json!({
10831 "orchestration": {
10832 "type": "handoff",
10833 "result": result.response.content,
10834 "final_agent": result.final_agent,
10835 "handoff_chain": chain_json,
10836 "duration_ms": duration_ms,
10837 }
10838 });
10839 response.metadata = Some(
10840 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10841 metadata,
10842 )
10843 .unwrap_or_default(),
10844 );
10845
10846 self.finish_turn_if_root(&response).await?;
10847 Ok(response)
10848 }
10849
10850 async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10852 self.begin_root_turn();
10853 self.pre_turn_session_lifecycle().await;
10855
10856 let input_data = self.process_input(input).await?;
10857 self.update_active_turn_context(&input_data.content, input_data.context.clone());
10858
10859 for (key, value) in &input_data.context {
10862 let _ = self.context_manager.set(key, value.clone());
10863 }
10864
10865 if input_data.metadata.rejected {
10866 let reason = input_data
10867 .metadata
10868 .rejection_reason
10869 .unwrap_or_else(|| "Input rejected".to_string());
10870 warn!(reason = %reason, "Input rejected");
10871 let response = AgentResponse::new(reason);
10872 self.finish_turn_if_root(&response).await?;
10873 return Ok(response);
10874 }
10875
10876 let processed_input = &input_data.content;
10877
10878 if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10879 return Ok(response);
10880 }
10881
10882 if let Some(ref sm) = self.state_machine
10884 && let Some(def) = sm.current_definition()
10885 {
10886 if let Some(ref delegate_id) = def.delegate {
10887 return self
10888 .handle_delegated_state(processed_input, delegate_id, &def)
10889 .await;
10890 }
10891 if let Some(ref concurrent_config) = def.concurrent {
10892 return self
10893 .handle_concurrent_state(processed_input, concurrent_config)
10894 .await;
10895 }
10896 if let Some(ref group_chat_config) = def.group_chat {
10897 return self
10898 .handle_group_chat_state(processed_input, group_chat_config)
10899 .await;
10900 }
10901 if let Some(ref pipeline_config) = def.pipeline {
10902 return self
10903 .handle_pipeline_state(processed_input, pipeline_config)
10904 .await;
10905 }
10906 if let Some(ref handoff_config) = def.handoff {
10907 return self
10908 .handle_handoff_state(processed_input, handoff_config)
10909 .await;
10910 }
10911 }
10912
10913 if let Some(response) =
10918 Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10919 {
10920 return Ok(response);
10921 }
10922
10923 match self.try_skill_route(processed_input).await? {
10924 SkillRouteResult::Response { skill_id, content } => {
10925 self.commit_root_user_message(processed_input).await?;
10926 return self
10927 .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10928 .await;
10929 }
10930 SkillRouteResult::NeedsClarification {
10931 response,
10932 ownership,
10933 } => {
10934 let admission = self
10935 .admit_optional_disambiguation_ownership(ownership)
10936 .await?;
10937 self.commit_root_user_message(processed_input).await?;
10938 if let Some(q) = response
10939 .metadata
10940 .as_ref()
10941 .and_then(|m| m.get("disambiguation"))
10942 .and_then(|d| d.get("status"))
10943 .and_then(|s| s.as_str())
10944 && q == "awaiting_clarification"
10945 {
10946 self.memory
10949 .add_message(ChatMessage::assistant(&response.content))
10950 .await?;
10951 }
10952 drop(admission);
10953 self.finish_turn_if_root(&response).await?;
10954 return Ok(response);
10955 }
10956 SkillRouteResult::NoMatch => {} }
10958
10959 let effective_reasoning = self.get_effective_reasoning_config();
10960 let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10961 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10962
10963 info!(
10964 reasoning_mode = ?reasoning_mode,
10965 auto_detected = auto_detected,
10966 reflection_enabled = ?self.reflection_config.enabled,
10967 "Reasoning mode determined"
10968 );
10969
10970 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10971 self.commit_root_user_message(processed_input).await?;
10972 return self
10973 .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10974 .await;
10975 }
10976
10977 self.commit_root_user_message(processed_input).await?;
10978
10979 let mut iterations = 0u32;
10980 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10981 let mut thinking_content: Option<String> = None;
10982
10983 let llm = self.get_state_llm()?;
10984
10985 loop {
10986 let effective_max = if reasoning_mode != ReasoningMode::None {
10988 let rc = self.get_effective_reasoning_config();
10989 self.max_iterations.min(rc.max_iterations)
10990 } else {
10991 self.max_iterations
10992 };
10993
10994 if iterations >= effective_max {
10995 let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10996 self.hooks.on_error(&err).await;
10997 error!(iterations = iterations, "Max iterations exceeded");
10998 return Err(err);
10999 }
11000 iterations += 1;
11001 *self.iteration_count.write() = iterations;
11002
11003 debug!(iteration = iterations, max = effective_max, "LLM call");
11004
11005 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
11006 let mut messages = self
11007 .build_messages_internal(true, None, protocol.choice.is_none())
11008 .await?;
11009 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11010
11011 self.hooks.on_llm_start(&messages).await;
11012 let llm_start = Instant::now();
11013 let response = self
11014 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
11015 .await?;
11016
11017 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11018 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11019
11020 let content = response.content.trim();
11021
11022 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
11023 match self
11024 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
11025 .await?
11026 {
11027 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
11028 ToolCallOutcome::Rejected(resp) => {
11029 self.finish_turn_if_root(&resp).await?;
11030 return Ok(resp);
11031 }
11032 }
11033 }
11034
11035 let (extracted_thinking, answer) = self.extract_thinking(content);
11036 if extracted_thinking.is_some() {
11037 thinking_content = extracted_thinking;
11038 }
11039
11040 let output_data = self.process_output(&answer, &input_data.context).await?;
11041
11042 let mut final_content = if output_data.metadata.rejected {
11043 output_data
11044 .metadata
11045 .rejection_reason
11046 .unwrap_or_else(|| answer.to_string())
11047 } else {
11048 output_data.content
11049 };
11050
11051 let reflection_metadata;
11053 (final_content, reflection_metadata) = self
11054 .run_reflection(&*llm, processed_input, final_content)
11055 .await?;
11056
11057 final_content =
11058 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
11059
11060 let final_content = {
11064 let result = self
11065 .post_loop_processing(processed_input, final_content)
11066 .await?;
11067 self.apply_post_loop_result(processed_input, result).await?
11068 };
11069
11070 let reflected = reflection_metadata.is_some();
11071 let reasoning_mode_debug = format!("{:?}", reasoning_mode);
11072
11073 let response = self.build_agent_response(AgentResponseParts {
11074 content: final_content,
11075 all_tool_calls,
11076 reasoning_mode,
11077 auto_detected,
11078 iterations,
11079 thinking: thinking_content,
11080 reflection_metadata,
11081 });
11082
11083 self.finish_turn_if_root(&response).await?;
11084
11085 let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
11086 info!(
11087 tool_calls = tool_call_count,
11088 response_len = response.content.len(),
11089 reasoning_mode = %reasoning_mode_debug,
11090 reflected = reflected,
11091 "Chat completed"
11092 );
11093 return Ok(response);
11094 }
11095 }
11096
11097 async fn generate_buffered_streaming_draft(
11098 &self,
11099 processed_input: &str,
11100 routing_resolved: Arc<AtomicBool>,
11101 ) -> Result<StreamingDraftResult> {
11102 let llm = self.get_state_llm()?;
11103 if llm.configured_tool_choice().is_some() {
11104 let draft = self
11105 .generate_main_response_draft(processed_input, &ReasoningMode::None)
11106 .await?;
11107 return Ok(StreamingDraftResult::new(draft, Vec::new()));
11108 }
11109 let messages = self.build_messages_for_draft(processed_input).await?;
11110 let mut stream = self
11111 .observe_purpose(
11112 ObservationPurpose::MainResponse,
11113 llm.complete_stream(&messages, None),
11114 )
11115 .await
11116 .map_err(|e| AgentError::LLM(e.to_string()))?;
11117 let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
11118 let mut chunks = Vec::new();
11119 let mut accumulated = String::new();
11120 while let Some(chunk_result) = stream.next().await {
11121 let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
11122 accumulated.push_str(&chunk.delta);
11123 let stream_chunk = StreamChunk::content(chunk.delta);
11124 if routing_resolved.load(Ordering::SeqCst) {
11125 chunks.push(stream_chunk);
11126 } else {
11127 buffer.push(stream_chunk)?;
11128 }
11129 }
11130 chunks.splice(0..0, buffer.drain());
11131 let content = accumulated.trim().to_string();
11132 let draft = if let Some(calls) = self.parse_tool_calls(&content) {
11133 MainResponseDraft::ToolCalls {
11134 raw_content: content,
11135 calls,
11136 thinking: None,
11137 }
11138 } else {
11139 MainResponseDraft::Text {
11140 raw_content: content,
11141 thinking: None,
11142 }
11143 };
11144 Ok(StreamingDraftResult::new(draft, chunks))
11145 }
11146
11147 async fn try_buffered_streaming_branches(
11148 &self,
11149 processed_input: &str,
11150 input_context: &HashMap<String, Value>,
11151 ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
11152 let optimization = &self.runtime_config.optimization;
11153 if !optimization.enabled {
11154 return Ok(None);
11155 }
11156 let transition_enabled =
11157 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
11158 if !transition_enabled {
11159 return Ok(None);
11160 }
11161 let mut branch_scheduler =
11162 TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
11163 if !branch_scheduler.reserve_task() {
11164 return Ok(None);
11165 }
11166 if !self
11167 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
11168 {
11169 branch_scheduler.release_task();
11170 return Ok(None);
11171 }
11172 if !branch_scheduler.reserve_task() {
11173 branch_scheduler.release_task();
11174 return Ok(None);
11175 }
11176 let mut main_branch = RuntimeBranch::new(
11177 RuntimeTaskPurpose::MainResponse,
11178 RuntimeOptimizationKind::BufferedStreamingRouting,
11179 RuntimeTaskPriority::Normal,
11180 RuntimeCommitBehavior::FinalResponse,
11181 );
11182 let mut transition_branch = RuntimeBranch::new(
11183 RuntimeTaskPurpose::StateTransition,
11184 RuntimeOptimizationKind::ParallelStateTransition,
11185 RuntimeTaskPriority::Critical,
11186 RuntimeCommitBehavior::TransitionDecision,
11187 );
11188 let main_id = main_branch.branch_id();
11189 let transition_id = transition_branch.branch_id();
11190 let routing_resolved = Arc::new(AtomicBool::new(false));
11191 let mut main_future =
11192 Box::pin(crate::optimization::observability::with_branch_observation(
11193 &main_id,
11194 RuntimeOptimizationKind::BufferedStreamingRouting,
11195 RuntimeCommitBehavior::FinalResponse,
11196 self.generate_buffered_streaming_draft(
11197 processed_input,
11198 Arc::clone(&routing_resolved),
11199 ),
11200 ));
11201 let mut transition_future =
11202 Box::pin(crate::optimization::observability::with_branch_observation(
11203 &transition_id,
11204 RuntimeOptimizationKind::ParallelStateTransition,
11205 RuntimeCommitBehavior::TransitionDecision,
11206 self.select_parallel_transition_candidate(processed_input),
11207 ));
11208 let mut main_pending = true;
11209 let mut transition_pending = true;
11210 let mut main_result: Option<Result<StreamingDraftResult>> = None;
11211 let mut transition_finalized = false;
11212 let mut transition_candidate: Option<TransitionCandidate> = None;
11213 loop {
11214 if let Some(candidate) = transition_candidate.take() {
11215 if self
11216 .approve_transition_target(&candidate.from_state, candidate.target())
11217 .await?
11218 {
11219 drop(main_future);
11221 drop(transition_future);
11222 self.finalize_branch_loss(
11223 &main_id,
11224 RuntimeOptimizationKind::BufferedStreamingRouting,
11225 RuntimeCommitBehavior::FinalResponse,
11226 main_pending,
11227 main_result.as_ref().map(|result| result.is_err()),
11228 );
11229 if !self
11230 .apply_pre_response_transition_candidate(
11231 &candidate,
11232 &HashMap::new(),
11233 processed_input,
11234 )
11235 .await?
11236 {
11237 self.finalize_optional_branch(
11238 &transition_id,
11239 RuntimeOptimizationKind::ParallelStateTransition,
11240 RuntimeCommitBehavior::TransitionDecision,
11241 "discarded",
11242 false,
11243 );
11244 return Ok(None);
11245 }
11246 self.finalize_optional_branch(
11247 &transition_id,
11248 RuntimeOptimizationKind::ParallelStateTransition,
11249 RuntimeCommitBehavior::TransitionDecision,
11250 "committed",
11251 true,
11252 );
11253 let response = self.redispatch_current_state(processed_input).await?;
11254 return Ok(Some((
11255 response.clone(),
11256 vec![StreamChunk::content(response.content)],
11257 )));
11258 }
11259 self.finalize_optional_branch(
11260 &transition_id,
11261 RuntimeOptimizationKind::ParallelStateTransition,
11262 RuntimeCommitBehavior::TransitionDecision,
11263 "discarded",
11264 false,
11265 );
11266 routing_resolved.store(true, Ordering::SeqCst);
11267 transition_finalized = true;
11268 }
11269 if transition_finalized && let Some(result) = main_result.take() {
11270 let stream_draft = match result {
11271 Ok(stream_draft) => stream_draft,
11272 Err(error) => {
11273 self.finalize_optional_branch(
11274 &main_id,
11275 RuntimeOptimizationKind::BufferedStreamingRouting,
11276 RuntimeCommitBehavior::FinalResponse,
11277 "failed",
11278 false,
11279 );
11280 return Err(error);
11281 }
11282 };
11283 let raw_draft_content = stream_draft.draft.raw_content().to_string();
11284 let buffered_chunks = stream_draft.chunks;
11285 self.finalize_optional_branch(
11286 &main_id,
11287 RuntimeOptimizationKind::BufferedStreamingRouting,
11288 RuntimeCommitBehavior::FinalResponse,
11289 "committed",
11290 true,
11291 );
11292 let response = self
11293 .commit_main_response_draft(
11294 processed_input,
11295 input_context,
11296 stream_draft.draft,
11297 ReasoningMode::None,
11298 false,
11299 )
11300 .await?;
11301 let chunks = if response.content == raw_draft_content {
11302 buffered_chunks
11303 } else {
11304 vec![StreamChunk::content(response.content.clone())]
11305 };
11306 return Ok(Some((response, chunks)));
11307 }
11308 tokio::select! {
11309 result = &mut main_future, if main_pending => {
11310 main_pending = false;
11311 main_branch.transition_to(RuntimeBranchStatus::Completed)?;
11312 main_result = Some(result);
11313 }
11314 result = &mut transition_future, if transition_pending => {
11315 transition_pending = false;
11316 transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
11317 match result {
11318 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
11319 transition_candidate = Some(candidate)
11320 }
11321 Ok(ParallelTransitionSelection::NoMatch) => {
11322 self.finalize_optional_branch(
11323 &transition_id,
11324 RuntimeOptimizationKind::ParallelStateTransition,
11325 RuntimeCommitBehavior::TransitionDecision,
11326 "discarded",
11327 false,
11328 );
11329 routing_resolved.store(true, Ordering::SeqCst);
11330 transition_finalized = true;
11331 }
11332 Ok(ParallelTransitionSelection::ReservationExhausted) => {
11333 self.finalize_optional_branch(
11334 &transition_id,
11335 RuntimeOptimizationKind::ParallelStateTransition,
11336 RuntimeCommitBehavior::TransitionDecision,
11337 "cancelled",
11338 false,
11339 );
11340 routing_resolved.store(true, Ordering::SeqCst);
11341 self.finalize_branch_loss(
11342 &main_id,
11343 RuntimeOptimizationKind::BufferedStreamingRouting,
11344 RuntimeCommitBehavior::FinalResponse,
11345 main_pending,
11346 main_result.as_ref().map(|result| result.is_err()),
11347 );
11348 return Ok(None);
11349 }
11350 Err(_) => {
11351 self.finalize_optional_branch(
11352 &transition_id,
11353 RuntimeOptimizationKind::ParallelStateTransition,
11354 RuntimeCommitBehavior::TransitionDecision,
11355 "failed",
11356 false,
11357 );
11358 routing_resolved.store(true, Ordering::SeqCst);
11359 transition_finalized = true;
11360 }
11361 }
11362 }
11363 }
11364 }
11365 }
11366
11367 fn run_loop_internal_stream<'a>(
11371 &'a self,
11372 input: &'a str,
11373 terminal: RuntimeStreamTerminalSlot,
11374 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11375 let include_tool_events = self.streaming.include_tool_events;
11376 let include_state_events = self.streaming.include_state_events;
11377
11378 Box::pin(async_stream::stream! {
11379 self.begin_root_turn();
11380 self.pre_turn_session_lifecycle().await;
11382
11383 let input_data = match self.process_input(input).await {
11384 Ok(data) => data,
11385 Err(e) => {
11386 yield StreamChunk::error(e.to_string());
11387 return;
11388 }
11389 };
11390 self.update_active_turn_context(&input_data.content, input_data.context.clone());
11391
11392 for (key, value) in &input_data.context {
11394 let _ = self.context_manager.set(key, value.clone());
11395 }
11396
11397 if input_data.metadata.rejected {
11398 let reason = input_data
11399 .metadata
11400 .rejection_reason
11401 .unwrap_or_else(|| "Input rejected".to_string());
11402 warn!(reason = %reason, "Input rejected (stream)");
11403 yield StreamChunk::error(reason);
11404 return;
11405 }
11406
11407 let processed_input = &input_data.content;
11408
11409 if self.runtime_config.optimization.enabled
11410 && matches!(
11411 self.runtime_config.optimization.streaming_policy,
11412 crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11413 )
11414 {
11415 match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11420 Ok(Some((response, chunks))) => {
11421 for chunk in chunks {
11422 yield chunk;
11423 }
11424 record_runtime_stream_final(&terminal, response);
11425 yield StreamChunk::Done {};
11426 return;
11427 }
11428 Ok(None) => {}
11429 Err(e) => {
11430 yield StreamChunk::error(e.to_string());
11431 return;
11432 }
11433 }
11434 }
11435
11436 if self.runtime_config.optimization.enabled
11437 && matches!(
11438 self.runtime_config.optimization.streaming_policy,
11439 crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11440 )
11441 {
11442 match self.try_pre_response_transition(processed_input).await {
11443 Ok(Some(response)) => {
11444 yield StreamChunk::content(&response.content);
11445 record_runtime_stream_final(&terminal, response);
11446 yield StreamChunk::Done {};
11447 return;
11448 }
11449 Ok(None) => {}
11450 Err(e) => {
11451 yield StreamChunk::error(e.to_string());
11452 return;
11453 }
11454 }
11455 }
11456
11457 if let Some(ref sm) = self.state_machine
11459 && let Some(def) = sm.current_definition()
11460 {
11461 let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11462 Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11463 } else if let Some(ref concurrent_config) = def.concurrent {
11464 Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11465 } else if let Some(ref group_chat_config) = def.group_chat {
11466 Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11467 } else if let Some(ref pipeline_config) = def.pipeline {
11468 Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11469 } else if let Some(ref handoff_config) = def.handoff {
11470 Some(self.handle_handoff_state(processed_input, handoff_config).await)
11471 } else {
11472 None
11473 };
11474
11475 if let Some(result) = orchestration_result {
11476 match result {
11477 Ok(response) => {
11478 yield StreamChunk::content(&response.content);
11479 record_runtime_stream_final(&terminal, response);
11480 yield StreamChunk::Done {};
11481 }
11482 Err(e) => {
11483 yield StreamChunk::error(e.to_string());
11484 }
11485 }
11486 return;
11487 }
11488 }
11489
11490 match self.try_skill_route(processed_input).await {
11492 Ok(SkillRouteResult::Response { skill_id, content }) => {
11493 if let Err(e) = self.commit_root_user_message(processed_input).await {
11494 yield StreamChunk::error(e.to_string());
11495 return;
11496 }
11497 match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11498 Ok(resp) => {
11499 yield StreamChunk::content(&resp.content);
11500 record_runtime_stream_final(&terminal, resp);
11501 yield StreamChunk::Done {};
11502 return;
11503 }
11504 Err(e) => {
11505 yield StreamChunk::error(e.to_string());
11506 return;
11507 }
11508 }
11509 }
11510 Ok(SkillRouteResult::NeedsClarification {
11511 response,
11512 ownership,
11513 }) => {
11514 let admission = match self
11515 .admit_optional_disambiguation_ownership(ownership)
11516 .await
11517 {
11518 Ok(admission) => admission,
11519 Err(e) => {
11520 yield StreamChunk::error(e.to_string());
11521 return;
11522 }
11523 };
11524 if let Err(e) = self.commit_root_user_message(processed_input).await {
11525 yield StreamChunk::error(e.to_string());
11526 return;
11527 }
11528 let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11529 drop(admission);
11530 if let Err(e) = self.finish_turn_if_root(&response).await {
11531 yield StreamChunk::error(e.to_string());
11532 return;
11533 }
11534 yield StreamChunk::content(&response.content);
11535 record_runtime_stream_final(&terminal, response);
11536 yield StreamChunk::Done {};
11537 return;
11538 }
11539 Ok(SkillRouteResult::NoMatch) => {} Err(e) => {
11541 yield StreamChunk::error(e.to_string());
11542 return;
11543 }
11544 }
11545
11546 let effective_reasoning = self.get_effective_reasoning_config();
11548 let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11549 Ok(mode) => mode,
11550 Err(e) => {
11551 yield StreamChunk::error(e.to_string());
11552 return;
11553 }
11554 };
11555 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11556
11557 info!(
11558 reasoning_mode = ?reasoning_mode,
11559 auto_detected = auto_detected,
11560 "Reasoning mode determined (stream)"
11561 );
11562
11563 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11565 if let Err(e) = self.commit_root_user_message(processed_input).await {
11566 yield StreamChunk::error(e.to_string());
11567 return;
11568 }
11569 match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11570 Ok(resp) => {
11571 yield StreamChunk::content(&resp.content);
11572 record_runtime_stream_final(&terminal, resp);
11573 yield StreamChunk::Done {};
11574 return;
11575 }
11576 Err(e) => {
11577 yield StreamChunk::error(e.to_string());
11578 return;
11579 }
11580 }
11581 }
11582
11583 if let Err(e) = self.commit_root_user_message(processed_input).await {
11584 yield StreamChunk::error(e.to_string());
11585 return;
11586 }
11587
11588 let llm = match self.get_state_llm() {
11589 Ok(llm) => llm,
11590 Err(e) => {
11591 yield StreamChunk::error(e.to_string());
11592 return;
11593 }
11594 };
11595
11596 let mut iterations = 0u32;
11597 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11598 let mut thinking_content: Option<String> = None;
11599
11600 loop {
11601 let effective_max = if reasoning_mode != ReasoningMode::None {
11603 let rc = self.get_effective_reasoning_config();
11604 self.max_iterations.min(rc.max_iterations)
11605 } else {
11606 self.max_iterations
11607 };
11608
11609 if iterations >= effective_max {
11610 let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11611 let err = AgentError::Other(err_msg.clone());
11612 self.hooks.on_error(&err).await;
11613 error!(iterations = iterations, "Max iterations exceeded (stream)");
11614 yield StreamChunk::error(err_msg);
11615 return;
11616 }
11617 iterations += 1;
11618 *self.iteration_count.write() = iterations;
11619
11620 debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11621
11622 let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11623 Ok(protocol) => protocol,
11624 Err(e) => {
11625 yield StreamChunk::error(e.to_string());
11626 return;
11627 }
11628 };
11629 let mut messages = match self
11630 .build_messages_internal(true, None, protocol.choice.is_none())
11631 .await
11632 {
11633 Ok(m) => m,
11634 Err(e) => {
11635 yield StreamChunk::error(e.to_string());
11636 return;
11637 }
11638 };
11639 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11640
11641 self.hooks.on_llm_start(&messages).await;
11642 let llm_start = Instant::now();
11643
11644 let reflection_active = self
11647 .should_reflect(processed_input, "")
11648 .await
11649 .unwrap_or_default();
11650
11651 let buffered_decision = reflection_active || protocol.choice.is_some();
11652 let content = if buffered_decision {
11653 let response = match self
11657 .complete_main_llm_with_recovery(
11658 Arc::clone(&llm),
11659 &messages,
11660 &protocol,
11661 )
11662 .await
11663 {
11664 Ok(r) => r,
11665 Err(e) => {
11666 yield StreamChunk::error(e.to_string());
11667 return;
11668 }
11669 };
11670 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11671 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11672 response.content.trim().to_string()
11673 } else {
11674 let llm_stream = match self
11676 .observe_purpose(
11677 ObservationPurpose::MainResponse,
11678 llm.complete_stream(&messages, None),
11679 )
11680 .await
11681 {
11682 Ok(s) => s,
11683 Err(e) => {
11684 yield StreamChunk::error(e.to_string());
11685 return;
11686 }
11687 };
11688 let mut accumulated = String::new();
11689 let mut stream_inner = llm_stream;
11690 while let Some(chunk_result) = stream_inner.next().await {
11691 match chunk_result {
11692 Ok(chunk) => {
11693 accumulated.push_str(&chunk.delta);
11694 yield StreamChunk::content(chunk.delta);
11695 }
11696 Err(e) => {
11697 yield StreamChunk::error(e.to_string());
11698 return;
11699 }
11700 }
11701 }
11702 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11703 let llm_response = ai_agents_core::LLMResponse::new(
11705 accumulated.trim(),
11706 ai_agents_core::FinishReason::Stop,
11707 );
11708 self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11709 accumulated.trim().to_string()
11710 };
11711
11712 if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11714 let native_tool_call = Self::is_native_tool_call_content(&content);
11715 let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11718 Ok(v) => v,
11719 Err(e) => {
11720 yield StreamChunk::error(e.to_string());
11721 return;
11722 }
11723 };
11724 if transition_fired {
11725 let _ = self.memory.add_message(ChatMessage::assistant(
11726 "(Transitioned to new state — tool call handled by workflow)",
11727 )).await;
11728
11729 if include_state_events
11730 && let Some(state) = self.current_state()
11731 {
11732 yield StreamChunk::state_transition(None, state);
11733 }
11734 continue;
11735 }
11736
11737 let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11739
11740 let results = self.execute_tools_parallel(&tool_calls).await;
11742
11743 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11744 if include_tool_events {
11745 yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11746 }
11747
11748 match result {
11749 Ok(output) => {
11750 if include_tool_events {
11751 yield StreamChunk::tool_result(
11752 &tool_call.id,
11753 &tool_call.name,
11754 &output,
11755 true,
11756 );
11757 }
11758 let _ = self.memory
11759 .add_message(Self::tool_result_message(
11760 tool_call,
11761 &output,
11762 native_tool_call,
11763 ))
11764 .await;
11765 }
11766 Err(e) => {
11767 if matches!(e, AgentError::HITLRejected(_)) {
11768 let _ = self.memory.add_message(ChatMessage::assistant(
11769 format!("The operation was rejected by the approver: {}", e),
11770 )).await;
11771 let response = AgentResponse {
11772 content: format!("Operation cancelled: {}", e),
11773 metadata: None,
11774 tool_calls: Some(all_tool_calls.clone()),
11775 };
11776 if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11777 yield StreamChunk::error(finalize_error.to_string());
11778 return;
11779 }
11780 let legacy_error = response.content.clone();
11781 record_runtime_stream_final(&terminal, response);
11782 yield StreamChunk::error(legacy_error);
11783 yield StreamChunk::Done {};
11784 return;
11785 }
11786 if include_tool_events {
11787 yield StreamChunk::tool_result(
11788 &tool_call.id,
11789 &tool_call.name,
11790 e.to_string(),
11791 false,
11792 );
11793 }
11794 let _ = self.memory
11795 .add_message(Self::tool_result_message(
11796 tool_call,
11797 &format!("Error: {}", e),
11798 native_tool_call,
11799 ))
11800 .await;
11801 }
11802 }
11803 all_tool_calls.push(tool_call.clone());
11804
11805 if include_tool_events {
11806 yield StreamChunk::tool_end(&tool_call.id);
11807 }
11808 }
11809 continue;
11810 }
11811
11812 let (extracted_thinking, answer) = self.extract_thinking(&content);
11814 if extracted_thinking.is_some() {
11815 thinking_content = extracted_thinking;
11816 }
11817
11818 let output_data = match self.process_output(&answer, &input_data.context).await {
11819 Ok(d) => d,
11820 Err(e) => {
11821 yield StreamChunk::error(e.to_string());
11822 return;
11823 }
11824 };
11825
11826 let final_content = if output_data.metadata.rejected {
11827 output_data
11828 .metadata
11829 .rejection_reason
11830 .unwrap_or_else(|| answer.to_string())
11831 } else {
11832 output_data.content
11833 };
11834
11835 let (final_content, reflection_metadata) = match self
11837 .run_reflection(&*llm, processed_input, final_content)
11838 .await
11839 {
11840 Ok(r) => r,
11841 Err(e) => {
11842 yield StreamChunk::error(e.to_string());
11843 return;
11844 }
11845 };
11846
11847 let final_content = self.format_response_with_thinking(
11848 thinking_content.as_deref(),
11849 &final_content,
11850 );
11851
11852 if buffered_decision {
11854 yield StreamChunk::content(&final_content);
11855 }
11856
11857 let post_result = match self
11861 .post_loop_processing(processed_input, final_content)
11862 .await
11863 {
11864 Ok(r) => r,
11865 Err(e) => {
11866 yield StreamChunk::error(e.to_string());
11867 return;
11868 }
11869 };
11870
11871 let (final_content, transitioned) = match post_result {
11872 PostLoopResult::NoTransition(content) => (content, false),
11873 PostLoopResult::Transitioned(content) => (content, true),
11874 PostLoopResult::NeedsRedispatch => {
11875 const MAX_REDISPATCH_DEPTH: u32 = 3;
11876 let current_depth = *self.redispatch_depth.read();
11877 let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11878 warn!(
11879 depth = current_depth,
11880 "Post-transition re-dispatch depth limit reached (stream)"
11881 );
11882 let c = String::new();
11883 let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11884 c
11885 } else {
11886 *self.redispatch_depth.write() += 1;
11887 if let Some(context) = self.active_turn_context.write().as_mut() {
11888 context.enter_redispatch();
11889 }
11890 info!(
11891 depth = current_depth + 1,
11892 "Re-dispatching for new state after transition (stream)"
11893 );
11894 let result = self.run_loop_internal(processed_input).await;
11895 *self.redispatch_depth.write() -= 1;
11896 if let Some(context) = self.active_turn_context.write().as_mut() {
11897 context.exit_redispatch();
11898 }
11899 match result {
11900 Ok(resp) => resp.content,
11901 Err(e) => {
11902 yield StreamChunk::error(e.to_string());
11903 return;
11904 }
11905 }
11906 };
11907 (content, true)
11908 }
11909 };
11910
11911 if transitioned {
11912 if include_state_events
11913 && let Some(state) = self.current_state()
11914 {
11915 yield StreamChunk::state_transition(None, state);
11916 }
11917 yield StreamChunk::content(&final_content);
11919 }
11920
11921 let final_response = self.build_agent_response(AgentResponseParts {
11923 content: final_content,
11924 all_tool_calls,
11925 reasoning_mode,
11926 auto_detected,
11927 iterations,
11928 thinking: thinking_content,
11929 reflection_metadata,
11930 });
11931 if let Err(e) = self.finish_turn_if_root(&final_response).await {
11932 yield StreamChunk::error(e.to_string());
11933 return;
11934 }
11935
11936 record_runtime_stream_final(&terminal, final_response);
11937 yield StreamChunk::Done {};
11938 return;
11939 }
11940 })
11941 }
11942
11943 fn run_loop_stream<'a>(
11946 &'a self,
11947 input: &'a str,
11948 terminal: RuntimeStreamTerminalSlot,
11949 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11950 Box::pin(async_stream::stream! {
11951 self.begin_root_turn();
11952 let _root_cleanup = RootTurnCleanup::new(self);
11953 self.hooks.on_message_received(input).await;
11954
11955 if !self.context_initialized.swap(true, Ordering::SeqCst) {
11957 if let Err(e) = self.context_manager.initialize().await {
11958 yield StreamChunk::error(e.to_string());
11959 return;
11960 }
11961 debug!("Context manager initialized (defaults, env, builtins)");
11962 }
11963
11964 if let Err(e) = self.check_turn_timeout().await {
11965 yield StreamChunk::error(e.to_string());
11966 return;
11967 }
11968 if let Err(e) = self.context_manager.refresh_per_turn().await {
11969 yield StreamChunk::error(e.to_string());
11970 return;
11971 }
11972
11973 self.clear_disambiguation_context();
11975
11976 if let Some(ref disambiguator) = self.disambiguation_manager {
11978 let disambiguation_context = match self.build_disambiguation_context().await {
11979 Ok(ctx) => ctx,
11980 Err(e) => {
11981 yield StreamChunk::error(e.to_string());
11982 return;
11983 }
11984 };
11985
11986 let state_override = self
11987 .state_machine
11988 .as_ref()
11989 .and_then(|sm| sm.current_definition())
11990 .and_then(|def| def.disambiguation.clone());
11991
11992 let state_generation = self
11993 .state_machine
11994 .as_ref()
11995 .map(|state_machine| state_machine.generation());
11996 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11997 let mut result = match self
11998 .observe_purpose(
11999 ObservationPurpose::DisambiguationDetection,
12000 disambiguator.process_input_with_override(
12001 input,
12002 &disambiguation_context,
12003 state_override.as_ref(),
12004 None,
12005 ),
12006 )
12007 .await
12008 {
12009 Ok(r) => r,
12010 Err(e) => {
12011 yield StreamChunk::error(e.to_string());
12012 return;
12013 }
12014 };
12015 let current_state_generation = self
12016 .state_machine
12017 .as_ref()
12018 .map(|state_machine| state_machine.generation());
12019 if current_state_generation != state_generation
12020 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
12021 {
12022 disambiguator.clear_pending().await;
12023 *self.pending_skill_id.write() = None;
12024 result = DisambiguationResult::Abandoned { new_input: None };
12025 info!(
12026 confirmation_event = "invalidated",
12027 invalidation_reason = "state_generation_changed",
12028 "Streaming disambiguation result invalidated before redispatch"
12029 );
12030 }
12031 match result {
12032 DisambiguationResult::Clear => {
12033 debug!("Input is clear, proceeding normally (stream)");
12034 }
12035 DisambiguationResult::NeedsClarification {
12036 question,
12037 detection,
12038 } => {
12039 let admission = match self
12040 .admit_disambiguation_redispatch(
12041 disambiguation_epoch,
12042 state_generation,
12043 )
12044 .await
12045 {
12046 Ok(admission) => admission,
12047 Err(error) => {
12048 *self.pending_skill_id.write() = None;
12049 yield StreamChunk::error(error.to_string());
12050 return;
12051 }
12052 };
12053 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
12054 info!(
12055 ambiguity_type = ?detection.ambiguity_type,
12056 confidence = detection.confidence,
12057 "Input requires clarification (stream)"
12058 );
12059 if let Err(e) = self.commit_root_user_message(input).await {
12062 yield StreamChunk::error(e.to_string());
12063 return;
12064 }
12065 let _ = self
12066 .memory
12067 .add_message(ChatMessage::assistant(&question.question))
12068 .await;
12069 let status = if awaiting_confirmation {
12070 "awaiting_confirmation"
12071 } else {
12072 "awaiting_clarification"
12073 };
12074 let response = AgentResponse::new(&question.question).with_metadata(
12075 "disambiguation",
12076 serde_json::json!({ "status": status }),
12077 );
12078 drop(admission);
12079 if let Err(e) = self.finish_turn_if_root(&response).await {
12080 yield StreamChunk::error(e.to_string());
12081 return;
12082 }
12083 yield StreamChunk::content(&question.question);
12084 record_runtime_stream_final(&terminal, response);
12085 yield StreamChunk::Done {};
12086 return;
12087 }
12088 DisambiguationResult::Clarified {
12089 enriched_input,
12090 resolved,
12091 ..
12092 } => {
12093 let admission = match self
12094 .admit_disambiguation_redispatch(
12095 disambiguation_epoch,
12096 state_generation,
12097 )
12098 .await
12099 {
12100 Ok(admission) => admission,
12101 Err(error) => {
12102 *self.pending_skill_id.write() = None;
12103 yield StreamChunk::error(error.to_string());
12104 return;
12105 }
12106 };
12107 info!(
12108 resolved_count = resolved.len(),
12109 enriched = %enriched_input,
12110 "Input clarified (stream)"
12111 );
12112 for (key, value) in &resolved {
12113 let context_key = format!("disambiguation.{}", key);
12114 let _ = self.context_manager.set(&context_key, value.clone());
12115 }
12116 if let Some(intent) = resolved.get("intent") {
12117 let _ = self.context_manager.set("resolved_intent", intent.clone());
12118 }
12119 let _ = self
12120 .context_manager
12121 .set("disambiguation.resolved", serde_json::Value::Bool(true));
12122
12123 let skill_id = self.pending_skill_id.read().clone();
12127 if let Some(skill_id) = skill_id {
12128 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
12129 drop(admission);
12130 match self
12131 .recheck_skill_disambiguation(
12132 &skill_id,
12133 &enriched_input,
12134 disambiguation_epoch,
12135 state_generation,
12136 )
12137 .await
12138 {
12139 Ok(resp) => {
12140 yield StreamChunk::content(&resp.content);
12141 record_runtime_stream_final(&terminal, resp);
12142 yield StreamChunk::Done {};
12143 return;
12144 }
12145 Err(e) => {
12146 yield StreamChunk::error(e.to_string());
12147 return;
12148 }
12149 }
12150 }
12151
12152 drop(admission);
12154 let mut inner = self.run_loop_internal_stream(
12155 &enriched_input,
12156 Arc::clone(&terminal),
12157 );
12158 while let Some(chunk) = inner.next().await {
12159 yield chunk;
12160 }
12161 return;
12162 }
12163 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
12164 info!("Proceeding with best guess (stream)");
12165
12166 let skill_id = self.pending_skill_id.read().clone();
12168 if let Some(skill_id) = skill_id {
12169 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
12170 match self
12171 .recheck_skill_disambiguation(
12172 &skill_id,
12173 &enriched_input,
12174 disambiguation_epoch,
12175 state_generation,
12176 )
12177 .await
12178 {
12179 Ok(resp) => {
12180 yield StreamChunk::content(&resp.content);
12181 record_runtime_stream_final(&terminal, resp);
12182 yield StreamChunk::Done {};
12183 return;
12184 }
12185 Err(e) => {
12186 yield StreamChunk::error(e.to_string());
12187 return;
12188 }
12189 }
12190 }
12191
12192 let mut inner = self.run_loop_internal_stream(
12193 &enriched_input,
12194 Arc::clone(&terminal),
12195 );
12196 while let Some(chunk) = inner.next().await {
12197 yield chunk;
12198 }
12199 return;
12200 }
12201 DisambiguationResult::GiveUp { reason } => {
12202 *self.pending_skill_id.write() = None;
12203 warn!(reason = %reason, "Disambiguation gave up (stream)");
12204 let apology = self
12205 .generate_localized_apology(
12206 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
12207 &reason,
12208 )
12209 .await
12210 .unwrap_or_else(|_| {
12211 format!("I'm sorry, I couldn't understand your request: {}", reason)
12212 });
12213 let response = AgentResponse::new(&apology);
12214 if let Err(e) = self.finish_turn_if_root(&response).await {
12215 yield StreamChunk::error(e.to_string());
12216 return;
12217 }
12218 yield StreamChunk::content(&apology);
12219 record_runtime_stream_final(&terminal, response);
12220 yield StreamChunk::Done {};
12221 return;
12222 }
12223 DisambiguationResult::Escalate { reason } => {
12224 *self.pending_skill_id.write() = None;
12225 info!(reason = %reason, "Escalating to human (stream)");
12226 if let Some(ref hitl) = self.hitl_engine {
12227 let trigger =
12228 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
12229 let mut context_map = HashMap::new();
12230 context_map.insert("original_input".to_string(), serde_json::json!(input));
12231 context_map.insert("reason".to_string(), serde_json::json!(&reason));
12232 let check_result = HITLCheckResult::required(
12233 trigger,
12234 context_map,
12235 format!("User request needs human assistance: {}", reason),
12236 Some(hitl.config().default_timeout_seconds),
12237 );
12238 match self.request_hitl_approval(check_result).await {
12239 Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
12240 let mut inner = self.run_loop_internal_stream(
12241 input,
12242 Arc::clone(&terminal),
12243 );
12244 while let Some(chunk) = inner.next().await {
12245 yield chunk;
12246 }
12247 return;
12248 }
12249 Ok(_) => {}
12250 Err(e) => {
12251 yield StreamChunk::error(e.to_string());
12252 return;
12253 }
12254 }
12255 }
12256 let apology = self
12257 .generate_localized_apology(
12258 "Explain briefly that you're transferring the user to a human agent for help.",
12259 &reason,
12260 )
12261 .await
12262 .unwrap_or_else(|_| {
12263 format!("I need human assistance to help with your request: {}", reason)
12264 });
12265 let response = AgentResponse::new(&apology);
12266 if let Err(e) = self.finish_turn_if_root(&response).await {
12267 yield StreamChunk::error(e.to_string());
12268 return;
12269 }
12270 yield StreamChunk::content(&apology);
12271 record_runtime_stream_final(&terminal, response);
12272 yield StreamChunk::Done {};
12273 return;
12274 }
12275 DisambiguationResult::Abandoned { new_input } => {
12276 *self.pending_skill_id.write() = None;
12277
12278 info!(
12279 has_new_input = new_input.is_some(),
12280 "Clarification abandoned by user (stream)"
12281 );
12282
12283 if let Err(e) = self.commit_root_user_message(input).await {
12284 yield StreamChunk::error(e.to_string());
12285 return;
12286 }
12287
12288 match new_input {
12289 Some(fresh_input) => {
12290 let mut inner = self.run_loop_internal_stream(
12292 &fresh_input,
12293 Arc::clone(&terminal),
12294 );
12295 while let Some(chunk) = inner.next().await {
12296 yield chunk;
12297 }
12298 return;
12299 }
12300 None => {
12301 let ack = self
12303 .generate_localized_apology(
12304 "The user changed their mind about their previous request. \
12305 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
12306 Do NOT apologize excessively. Be concise.",
12307 "User abandoned clarification",
12308 )
12309 .await
12310 .unwrap_or_else(|_| {
12311 "OK, no problem. What else can I help with?".to_string()
12312 });
12313
12314 let _ = self
12315 .memory
12316 .add_message(ChatMessage::assistant(&ack))
12317 .await;
12318
12319 let response = AgentResponse::new(&ack);
12320 if let Err(e) = self.finish_turn_if_root(&response).await {
12321 yield StreamChunk::error(e.to_string());
12322 return;
12323 }
12324 yield StreamChunk::content(&ack);
12325 record_runtime_stream_final(&terminal, response);
12326 yield StreamChunk::Done {};
12327 return;
12328 }
12329 }
12330 }
12331 }
12332 }
12333
12334 let mut inner = self.run_loop_internal_stream(input, Arc::clone(&terminal));
12336 while let Some(chunk) = inner.next().await {
12337 yield chunk;
12338 }
12339 })
12340 }
12341
12342 pub fn info(&self) -> AgentInfo {
12343 self.info.clone()
12344 }
12345
12346 pub fn skills(&self) -> &[SkillDefinition] {
12347 &self.skills
12348 }
12349
12350 async fn reset_runtime_state(&self) -> Result<()> {
12352 let _admission = self.disambiguation_admission.write().await;
12353 if self.state_transition_reserved.load(Ordering::SeqCst) {
12354 return Err(AgentError::Other(
12355 "Cannot reset while a state transition is in progress".to_string(),
12356 ));
12357 }
12358 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
12359 *self.pending_skill_id.write() = None;
12360 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
12361 disambiguator.clear_pending().await;
12362 }
12363 self.memory.clear().await?;
12364 *self.iteration_count.write() = 0;
12365 self.tool_call_history.write().clear();
12366 if let Some(ref sm) = self.state_machine {
12367 sm.reset();
12368 }
12369 Ok(())
12370 }
12371
12372 pub async fn reset(&self) -> Result<()> {
12374 self.reset_runtime_state().await
12375 }
12376
12377 pub fn max_context_tokens(&self) -> u32 {
12378 self.max_context_tokens
12379 }
12380
12381 pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12382 &self.llm_registry
12383 }
12384
12385 pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12386 self.state_machine.as_ref()
12387 }
12388
12389 pub fn context_manager(&self) -> &Arc<ContextManager> {
12390 &self.context_manager
12391 }
12392
12393 pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12394 self.tool_call_history.read().clone()
12395 }
12396
12397 pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12398 self.memory_token_budget.as_ref()
12399 }
12400
12401 pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12402 &self.parallel_tools
12403 }
12404
12405 pub fn streaming_config(&self) -> &StreamingConfig {
12406 &self.streaming
12407 }
12408
12409 pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12410 &self.hooks
12411 }
12412
12413 pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12414 self.hitl_engine.as_ref()
12415 }
12416
12417 pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12418 &self.approval_handler
12419 }
12420
12421 fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12423 let mut ctx = HashMap::new();
12424 for key in &["user.language", "input.detected.language", "language"] {
12425 if let Some(val) = self.context_manager.get(key) {
12426 ctx.insert(key.to_string(), val);
12427 }
12428 }
12429 ctx
12430 }
12431
12432 async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12434 let Some(request) = check_result.into_request() else {
12435 return Ok(ApprovalResult::Approved);
12436 };
12437
12438 self.hooks.on_approval_requested(&request).await;
12439
12440 let timeout = request.timeout;
12441
12442 let raw_result = if let Some(duration) = timeout {
12443 match tokio::time::timeout(
12444 duration,
12445 self.approval_handler.request_approval(request.clone()),
12446 )
12447 .await
12448 {
12449 Ok(result) => result,
12450 Err(_) => ApprovalResult::timeout(),
12451 }
12452 } else {
12453 self.approval_handler
12454 .request_approval(request.clone())
12455 .await
12456 };
12457
12458 self.hooks
12459 .on_approval_result(&request.id, &raw_result)
12460 .await;
12461
12462 let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12463 match &raw_result {
12464 ApprovalResult::Approved => (
12465 ApprovalResolvedOutcome::Approved,
12466 Ok(ApprovalResult::Approved),
12467 ),
12468 ApprovalResult::Rejected { reason } => (
12469 ApprovalResolvedOutcome::Rejected {
12470 reason: reason.clone(),
12471 },
12472 Ok(ApprovalResult::Rejected {
12473 reason: reason.clone(),
12474 }),
12475 ),
12476 ApprovalResult::Modified { changes } => (
12477 ApprovalResolvedOutcome::Modified {
12478 changes: changes.clone(),
12479 },
12480 Ok(ApprovalResult::Modified {
12481 changes: changes.clone(),
12482 }),
12483 ),
12484 ApprovalResult::Timeout => {
12485 if let Some(ref engine) = self.hitl_engine {
12486 match engine.config().on_timeout {
12487 TimeoutAction::Approve => (
12488 ApprovalResolvedOutcome::Approved,
12489 Ok(ApprovalResult::Approved),
12490 ),
12491 TimeoutAction::Reject => {
12492 let reason = Some("Timeout".to_string());
12493 (
12494 ApprovalResolvedOutcome::Rejected {
12495 reason: reason.clone(),
12496 },
12497 Ok(ApprovalResult::Rejected { reason }),
12498 )
12499 }
12500 TimeoutAction::Error => {
12501 let message = "HITL approval timeout".to_string();
12502 (
12503 ApprovalResolvedOutcome::Error {
12504 message: message.clone(),
12505 },
12506 Err(AgentError::Other(message)),
12507 )
12508 }
12509 }
12510 } else {
12511 let reason = Some("Timeout (no engine)".to_string());
12512 (
12513 ApprovalResolvedOutcome::Rejected {
12514 reason: reason.clone(),
12515 },
12516 Ok(ApprovalResult::Rejected { reason }),
12517 )
12518 }
12519 }
12520 };
12521
12522 self.hooks
12523 .on_approval_resolved(&request, &raw_result, &outcome)
12524 .await;
12525
12526 effective_result
12527 }
12528
12529 pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12530 if let Some(ref hitl_engine) = self.hitl_engine {
12531 let hitl_lang_ctx = self.build_hitl_language_context();
12532 let check_result = self
12533 .observe_purpose(
12534 ObservationPurpose::HitlLocalization,
12535 hitl_engine.check_state_transition_with_localization(
12536 from,
12537 to,
12538 &hitl_lang_ctx,
12539 self.approval_handler.as_ref(),
12540 Some(&self.llm_registry),
12541 ),
12542 )
12543 .await?;
12544 if check_result.is_required() {
12545 let result = self.request_hitl_approval(check_result).await?;
12546 return Ok(matches!(
12547 result,
12548 ApprovalResult::Approved | ApprovalResult::Modified { .. }
12549 ));
12550 }
12551 }
12552 Ok(true)
12553 }
12554
12555 async fn execute_tools_parallel(
12557 &self,
12558 tool_calls: &[ToolCall],
12559 ) -> Vec<(String, Result<String>)> {
12560 let can_run_parallel = tool_calls.iter().all(|tc| {
12561 self.tools
12562 .resolve(&tc.name)
12563 .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12564 .unwrap_or(false)
12565 });
12566
12567 if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12568 let mut results = Vec::new();
12569 for tc in tool_calls {
12570 let result = self
12571 .observe_purpose(
12572 current_observation_context()
12573 .map(|context| context.purpose)
12574 .unwrap_or_default(),
12575 self.execute_tool_smart(tc),
12576 )
12577 .await;
12578 results.push((tc.id.clone(), result));
12579 }
12580 return results;
12581 }
12582
12583 let chunks: Vec<_> = tool_calls
12584 .chunks(self.parallel_tools.max_parallel)
12585 .collect();
12586
12587 let mut all_results = Vec::new();
12588
12589 for chunk in chunks {
12590 let futures: Vec<_> = chunk
12591 .iter()
12592 .map(|tc| {
12593 let tc = tc.clone();
12594 async move {
12595 let result = self.execute_tool_smart(&tc).await;
12596 (tc.id.clone(), result)
12597 }
12598 })
12599 .collect();
12600
12601 let results = futures::future::join_all(futures).await;
12602 all_results.extend(results);
12603 }
12604
12605 all_results
12606 }
12607
12608 pub async fn chat_stream<'a>(
12612 &'a self,
12613 input: &'a str,
12614 ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12615 let RootTurnAdmission {
12616 guard: root_turn_guard,
12617 identity_stack,
12618 } = self.acquire_root_turn().await?;
12619 scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12623 info!(input_len = input.len(), "Starting streaming chat");
12624 let terminal = new_runtime_stream_terminal_slot();
12625 let inner = self.run_loop_stream(input, terminal);
12626 let observation_context = self.build_observation_context(None);
12627 let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12628 Box::pin(async_stream::stream! {
12629 let mut root_turn_guard = Some(root_turn_guard);
12630 let mut inner = inner;
12631 loop {
12632 let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12633 if let Some(context) = observation_context.as_ref() {
12634 with_observation_context(context.clone(), inner.next()).await
12635 } else {
12636 inner.next().await
12637 }
12638 })
12639 .await;
12640 match next {
12641 Some(StreamChunk::Done {}) => {
12642 while scope_runtime_gate_identity_stack(&identity_stack, async {
12643 if let Some(context) = observation_context.as_ref() {
12644 with_observation_context(context.clone(), inner.next())
12645 .await
12646 .is_some()
12647 } else {
12648 inner.next().await.is_some()
12649 }
12650 })
12651 .await
12652 {}
12653 if observation_context.is_some() {
12654 scope_runtime_gate_identity_stack(
12655 &identity_stack,
12656 self.export_observability_if_configured(),
12657 )
12658 .await;
12659 }
12660 drop(root_turn_guard.take());
12661 yield StreamChunk::Done {};
12662 return;
12663 }
12664 Some(chunk) => yield chunk,
12665 None => {
12666 if observation_context.is_some() {
12667 scope_runtime_gate_identity_stack(
12668 &identity_stack,
12669 self.export_observability_if_configured(),
12670 )
12671 .await;
12672 }
12673 drop(root_turn_guard.take());
12674 return;
12675 }
12676 }
12677 }
12678 });
12679 Ok(stream)
12680 }
12681
12682 pub async fn chat_stream_events<'a>(
12686 &'a self,
12687 input: &'a str,
12688 ) -> Result<Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>>> {
12689 let RootTurnAdmission {
12690 guard: root_turn_guard,
12691 identity_stack,
12692 } = self.acquire_root_turn().await?;
12693 scope_runtime_gate_identity_stack(&identity_stack, self.init_storage()).await?;
12697 info!(input_len = input.len(), "Starting streaming chat events");
12698 let terminal = new_runtime_stream_terminal_slot();
12699 let mut inner = self.run_loop_stream(input, Arc::clone(&terminal));
12700 let observation_context = self.build_observation_context(None);
12701 let stream: Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>> =
12702 Box::pin(async_stream::stream! {
12703 let mut root_turn_guard = Some(root_turn_guard);
12704 loop {
12705 let next = scope_runtime_gate_identity_stack(&identity_stack, async {
12706 if let Some(context) = observation_context.as_ref() {
12707 with_observation_context(context.clone(), inner.next()).await
12708 } else {
12709 inner.next().await
12710 }
12711 })
12712 .await;
12713 match next {
12714 Some(StreamChunk::Done {}) => {
12715 let terminal_event = { terminal.write().take() };
12716 if let Some(response) = terminal_event {
12717 while scope_runtime_gate_identity_stack(&identity_stack, async {
12718 if let Some(context) = observation_context.as_ref() {
12719 with_observation_context(context.clone(), inner.next())
12720 .await
12721 .is_some()
12722 } else {
12723 inner.next().await.is_some()
12724 }
12725 })
12726 .await
12727 {}
12728 if observation_context.is_some() {
12729 scope_runtime_gate_identity_stack(
12730 &identity_stack,
12731 self.export_observability_if_configured(),
12732 )
12733 .await;
12734 }
12735 drop(root_turn_guard.take());
12736 yield AgentStreamEvent::Final(response);
12737 return;
12738 }
12739 }
12740 Some(StreamChunk::Error { message }) => {
12741 let finalized = { terminal.read().is_some() };
12742 if finalized {
12743 continue;
12744 }
12745 while scope_runtime_gate_identity_stack(&identity_stack, async {
12746 if let Some(context) = observation_context.as_ref() {
12747 with_observation_context(context.clone(), inner.next())
12748 .await
12749 .is_some()
12750 } else {
12751 inner.next().await.is_some()
12752 }
12753 })
12754 .await
12755 {}
12756 if observation_context.is_some() {
12757 scope_runtime_gate_identity_stack(
12758 &identity_stack,
12759 self.export_observability_if_configured(),
12760 )
12761 .await;
12762 }
12763 drop(root_turn_guard.take());
12764 yield AgentStreamEvent::Chunk(StreamChunk::Error { message });
12765 return;
12766 }
12767 Some(chunk) => yield AgentStreamEvent::Chunk(chunk),
12768 None => {
12769 if observation_context.is_some() {
12770 scope_runtime_gate_identity_stack(
12771 &identity_stack,
12772 self.export_observability_if_configured(),
12773 )
12774 .await;
12775 }
12776 drop(root_turn_guard.take());
12777 return;
12778 }
12779 }
12780 }
12781 });
12782 Ok(stream)
12783 }
12784}
12785
12786#[async_trait]
12787impl ToolInvoker for RuntimeAgent {
12788 async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12789 self.execute_tool_record(request).await
12790 }
12791}
12792
12793#[async_trait]
12794impl Agent for RuntimeAgent {
12795 async fn chat(&self, input: &str) -> Result<AgentResponse> {
12797 let RootTurnAdmission {
12798 guard,
12799 identity_stack,
12800 } = self.acquire_root_turn().await?;
12801 let result = scope_runtime_gate_identity_stack(&identity_stack, async {
12802 let result = if let Some(context) = self.build_observation_context(None) {
12803 with_observation_context(context, self.run_loop(input)).await
12804 } else {
12805 self.run_loop(input).await
12806 };
12807 self.export_observability_if_configured().await;
12808 result
12809 })
12810 .await;
12811 drop(guard);
12812 result
12813 }
12814
12815 fn info(&self) -> AgentInfo {
12816 self.info.clone()
12817 }
12818
12819 async fn reset(&self) -> Result<()> {
12821 self.reset_runtime_state().await
12822 }
12823}
12824
12825fn background_maintenance_tags(
12835 label: &str,
12836 stage: &str,
12837 reason: Option<&str>,
12838 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12839) -> HashMap<String, String> {
12840 let mut tags = HashMap::new();
12841 tags.insert("runtime.background".to_string(), "true".to_string());
12842 tags.insert("runtime.maintenance".to_string(), label.to_string());
12843 tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12844 if let Some(policy) = policy {
12845 tags.insert(
12846 "runtime.await_before_next_turn".to_string(),
12847 await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12848 );
12849 tags.insert(
12850 "runtime.maintenance_mode".to_string(),
12851 maintenance_mode_label(policy.mode).to_string(),
12852 );
12853 }
12854 if let Some(reason) = reason {
12855 tags.insert("runtime.reason".to_string(), reason.to_string());
12856 }
12857 tags
12858}
12859
12860fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12861 match policy {
12862 AwaitBeforeNextTurn::Never => "never",
12863 AwaitBeforeNextTurn::SameActor => "same_actor",
12864 AwaitBeforeNextTurn::Always => "always",
12865 }
12866}
12867
12868fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12869 match mode {
12870 MaintenanceMode::InlineSerial => "inline_serial",
12871 MaintenanceMode::InlineParallel => "inline_parallel",
12872 MaintenanceMode::Background => "background",
12873 }
12874}
12875
12876fn record_background_maintenance_event(
12878 manager: Option<&Arc<ObservabilityManager>>,
12879 label: &str,
12880 status: EventStatus,
12881 duration_ms: u64,
12882 stage: &str,
12883 reason: Option<String>,
12884 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12885) {
12886 if let Some(manager) = manager {
12887 manager.record_lifecycle_event(
12888 EventType::MemoryOperation {
12889 operation: format!("{}_background_{}", label, stage),
12890 },
12891 ObservationPurpose::Other(format!("{}_maintenance", label)),
12892 status,
12893 duration_ms,
12894 background_maintenance_tags(label, stage, reason.as_deref(), policy),
12895 None,
12896 );
12897 }
12898}
12899
12900fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12901 if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12902 MaintenanceMode::InlineParallel
12903 } else {
12904 mode
12905 }
12906}
12907
12908fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12909 match hint {
12910 ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12911 ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12912 ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12913 ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12914 ObservationPurpose::ProcessTransform
12915 }
12916 }
12917}
12918
12919fn new_tool_resource_locks() -> ToolResourceLocks {
12920 Arc::new(RwLock::new(HashMap::new()))
12921}
12922
12923fn tool_resource_lock_keys(
12928 _canonical_id: &str,
12929 args: &Value,
12930 bindings: &ai_agents_core::ToolPolicyBindings,
12931 classification: &ai_agents_core::ToolCallClassification,
12932) -> Vec<String> {
12933 if classification.concurrency_safe {
12934 return Vec::new();
12935 }
12936
12937 let mut keys = Vec::new();
12938 let mut has_path_resource = false;
12939 for binding in &bindings.path_fields {
12940 let value = value_at_argument_path(args, &binding.field)
12941 .cloned()
12942 .or_else(|| {
12943 binding
12944 .default_path
12945 .as_ref()
12946 .map(|path| Value::String(path.clone()))
12947 });
12948 if let Some(value) = value {
12949 collect_resource_strings(&value, |_| {
12950 has_path_resource = true;
12951 });
12952 }
12953 }
12954 for binding in &bindings.domain_fields {
12955 if let Some(value) = value_at_argument_path(args, &binding.field) {
12956 collect_resource_strings(value, |domain| {
12957 let normalized = if binding.is_url {
12958 normalized_url_resource_key(domain)
12959 } else {
12960 domain.trim().trim_end_matches('.').to_ascii_lowercase()
12961 };
12962 keys.push(format!("domain:{}", normalized));
12963 });
12964 }
12965 }
12966 for binding in &bindings.command_fields {
12967 if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12968 continue;
12969 }
12970 if let Some(value) = value_at_argument_path(args, &binding.field) {
12971 collect_resource_strings(value, |_| {
12972 has_path_resource = true;
12973 });
12974 }
12975 }
12976 if has_path_resource {
12977 keys.push("path-mutation:global".to_string());
12978 }
12979 if keys.is_empty() {
12980 keys.push("side-effect:unbound".to_string());
12981 }
12982 keys.sort();
12983 keys.dedup();
12984 keys
12985}
12986
12987fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12988 let mut current = value;
12989 for segment in field.split('.') {
12990 if segment.is_empty() {
12991 return None;
12992 }
12993 current = current.get(segment)?;
12994 }
12995 Some(current)
12996}
12997
12998fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12999 match value {
13000 Value::String(value) => collect(value),
13001 Value::Array(values) => {
13002 for value in values {
13003 if let Some(value) = value.as_str() {
13004 collect(value);
13005 }
13006 }
13007 }
13008 _ => {}
13009 }
13010}
13011
13012fn normalized_url_resource_key(value: &str) -> String {
13013 let value = value.trim();
13014 let Some((scheme, remainder)) = value.split_once("://") else {
13015 return value.to_ascii_lowercase();
13016 };
13017 let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
13018 let (authority, suffix) = remainder.split_at(authority_end);
13019 format!(
13020 "{}://{}{}",
13021 scheme.to_ascii_lowercase(),
13022 authority.to_ascii_lowercase(),
13023 suffix
13024 )
13025}
13026
13027fn render_concurrent_template(
13028 template: &str,
13029 user_input: &str,
13030 context_values: &std::collections::HashMap<String, serde_json::Value>,
13031) -> Result<String> {
13032 let mut env = minijinja::Environment::new();
13033 env.add_template("concurrent", template)
13034 .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
13035
13036 let mut ctx = std::collections::BTreeMap::new();
13037 ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
13038
13039 let context_obj = minijinja::Value::from_serialize(context_values);
13041 ctx.insert("context".to_string(), context_obj);
13042
13043 let tmpl = env
13044 .get_template("concurrent")
13045 .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
13046
13047 tmpl.render(minijinja::Value::from_serialize(&ctx))
13048 .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
13049}
13050
13051#[cfg(test)]
13052mod tests {
13053 use super::*;
13054 use crate::AgentBuilder;
13055 use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
13056 use ai_agents_llm::mock::MockLLMProvider;
13057 use ai_agents_skills::{SkillDefinition, SkillStep};
13058 use ai_agents_tools::{
13059 CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool, ToolAliases,
13060 ToolDescriptor, ToolProvider, ToolProviderError, ToolProviderType, WebFetchResolver,
13061 WebFetchTool, WebFetchTransport, WebFetchTransportRequest, WebFetchTransportResponse,
13062 };
13063
13064 fn mock_with_response(response: &str) -> MockLLMProvider {
13065 let mut mock = MockLLMProvider::new("test");
13066 mock.set_response(response);
13067 mock
13068 }
13069
13070 fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
13071 let mut mock = MockLLMProvider::new("test");
13072 mock.set_responses(responses.into_iter().map(String::from).collect(), true);
13073 mock
13074 }
13075
13076 fn disambiguation_state_machine(
13078 state_enabled: Option<bool>,
13079 require_confirmation: bool,
13080 ) -> Arc<StateMachine> {
13081 let definition = ai_agents_state::StateDefinition {
13082 prompt: Some("Handle the resolved request.".to_string()),
13083 disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
13084 enabled: state_enabled,
13085 require_confirmation,
13086 ..Default::default()
13087 }),
13088 ..Default::default()
13089 };
13090 let review = ai_agents_state::StateDefinition {
13091 prompt: Some("Review a fresh request.".to_string()),
13092 ..Default::default()
13093 };
13094 Arc::new(
13095 StateMachine::new(ai_agents_state::StateConfig {
13096 initial: "active".to_string(),
13097 states: std::collections::HashMap::from([
13098 ("active".to_string(), definition),
13099 ("review".to_string(), review),
13100 ]),
13101 global_transitions: Vec::new(),
13102 fallback: None,
13103 max_no_transition: None,
13104 regenerate_on_transition: true,
13105 })
13106 .unwrap(),
13107 )
13108 }
13109
13110 fn state_disambiguation_agent(
13112 responses: Vec<&str>,
13113 manager_enabled: bool,
13114 state_enabled: Option<bool>,
13115 require_confirmation: bool,
13116 ) -> (RuntimeAgent, MockLLMProvider) {
13117 state_disambiguation_agent_with_skills(
13118 responses,
13119 manager_enabled,
13120 state_enabled,
13121 require_confirmation,
13122 Vec::new(),
13123 )
13124 }
13125
13126 fn state_disambiguation_agent_with_skills(
13128 responses: Vec<&str>,
13129 manager_enabled: bool,
13130 state_enabled: Option<bool>,
13131 require_confirmation: bool,
13132 skills: Vec<SkillDefinition>,
13133 ) -> (RuntimeAgent, MockLLMProvider) {
13134 let mut mock = MockLLMProvider::new("state-confirmation");
13135 mock.set_responses(responses.into_iter().map(String::from).collect(), false);
13136 let observed = mock.clone();
13137 let agent = AgentBuilder::new()
13138 .system_prompt("Handle requests.")
13139 .llm(Arc::new(mock.clone()))
13140 .llm_alias("router", Arc::new(mock))
13141 .state_machine(disambiguation_state_machine(
13142 state_enabled,
13143 require_confirmation,
13144 ))
13145 .skills(skills)
13146 .build()
13147 .unwrap()
13148 .with_disambiguation(DisambiguationConfig {
13149 enabled: manager_enabled,
13150 ..Default::default()
13151 });
13152 (agent, observed)
13153 }
13154
13155 fn confirmation_skill() -> SkillDefinition {
13157 SkillDefinition {
13158 id: "send_report".to_string(),
13159 description: "Send a report after clarification".to_string(),
13160 trigger: "When the user asks to send a report".to_string(),
13161 steps: vec![SkillStep::Prompt {
13162 prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
13163 llm: None,
13164 }],
13165 reasoning: None,
13166 reflection: None,
13167 disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
13168 enabled: Some(true),
13169 ..Default::default()
13170 }),
13171 }
13172 }
13173
13174 fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
13176 observed
13177 .call_history()
13178 .iter()
13179 .filter(|call| {
13180 call.messages
13181 .iter()
13182 .any(|message| message.content.contains("Execute confirmed report skill"))
13183 })
13184 .count()
13185 }
13186
13187 struct BlockingRuntimeConfirmationObserver {
13188 entered: tokio::sync::Barrier,
13189 release: tokio::sync::Notify,
13190 }
13191
13192 impl BlockingRuntimeConfirmationObserver {
13193 fn new() -> Self {
13194 Self {
13195 entered: tokio::sync::Barrier::new(2),
13196 release: tokio::sync::Notify::new(),
13197 }
13198 }
13199 }
13200
13201 struct ResetOnTransitionHooks {
13202 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
13203 invoked: AtomicBool,
13204 }
13205
13206 #[async_trait]
13207 impl AgentHooks for ResetOnTransitionHooks {
13208 async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
13209 if self.invoked.swap(true, Ordering::SeqCst) {
13210 return;
13211 }
13212 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
13213 if let Some(agent) = agent {
13214 agent.reset().await.unwrap();
13215 }
13216 }
13217 }
13218
13219 impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
13220 fn observe_question<'a>(
13221 &'a self,
13222 future: ClarificationQuestionFuture<'a>,
13223 ) -> ClarificationQuestionFuture<'a> {
13224 future
13225 }
13226
13227 fn observe_parse<'a>(
13228 &'a self,
13229 future: ClarificationParseFuture<'a>,
13230 ) -> ClarificationParseFuture<'a> {
13231 future
13232 }
13233
13234 fn observe_confirmation_parse<'a>(
13235 &'a self,
13236 future: ConfirmationParseFuture<'a>,
13237 ) -> ConfirmationParseFuture<'a> {
13238 Box::pin(async move {
13239 self.entered.wait().await;
13240 self.release.notified().await;
13241 future.await
13242 })
13243 }
13244 }
13245
13246 #[tokio::test]
13247 async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
13248 let (agent, observed) = state_disambiguation_agent(
13249 vec![
13250 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13251 r#"{"question":"What should I send?","options":null}"#,
13252 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13253 r#"{"question":"Should I send the report to Ada?"}"#,
13254 r#"{"status":"confirmed"}"#,
13255 "Request executed.",
13256 ],
13257 true,
13258 None,
13259 true,
13260 );
13261
13262 let clarification = agent.chat("Send it").await.unwrap();
13263 assert_eq!(clarification.content, "What should I send?");
13264 assert_eq!(observed.call_count(), 2);
13265
13266 let confirmation = agent.chat("The report to Ada").await.unwrap();
13267 assert_eq!(confirmation.content, "Should I send the report to Ada?");
13268 assert_eq!(
13269 confirmation
13270 .metadata
13271 .as_ref()
13272 .and_then(|metadata| metadata.get("disambiguation"))
13273 .and_then(|metadata| metadata.get("status"))
13274 .and_then(Value::as_str),
13275 Some("awaiting_confirmation")
13276 );
13277 assert_eq!(observed.call_count(), 4);
13278
13279 let completed = agent.chat("Yes").await.unwrap();
13280 assert_eq!(completed.content, "Request executed.");
13281 assert_eq!(observed.call_count(), 6);
13282 }
13283
13284 #[tokio::test]
13285 async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
13286 let (agent, observed) = state_disambiguation_agent(
13287 vec![
13288 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13289 r#"{"question":"What should I send?","options":null}"#,
13290 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13291 r#"{"question":"Should I send the report to Ada?"}"#,
13292 r#"{"status":"confirmed"}"#,
13293 "Request executed.",
13294 ],
13295 true,
13296 None,
13297 true,
13298 );
13299
13300 let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
13301 let mut clarification = String::new();
13302 while let Some(chunk) = clarification_stream.next().await {
13303 match chunk {
13304 StreamChunk::Content { text } => clarification.push_str(&text),
13305 StreamChunk::Done {} => break,
13306 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13307 _ => {}
13308 }
13309 }
13310 assert_eq!(clarification, "What should I send?");
13311 assert_eq!(observed.call_count(), 2);
13312
13313 let mut confirmation_stream = agent.chat_stream_events("The report to Ada").await.unwrap();
13314 let mut confirmation = None;
13315 while let Some(event) = confirmation_stream.next().await {
13316 match event {
13317 AgentStreamEvent::Final(response) => confirmation = Some(response),
13318 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
13319 panic!("unexpected stream error: {message}")
13320 }
13321 AgentStreamEvent::Chunk(_) => {}
13322 }
13323 }
13324 let confirmation = confirmation.expect("confirmation must finalize");
13325 assert_eq!(confirmation.content, "Should I send the report to Ada?");
13326 assert_eq!(
13327 confirmation
13328 .metadata
13329 .as_ref()
13330 .and_then(|metadata| metadata.get("disambiguation"))
13331 .and_then(|metadata| metadata.get("status"))
13332 .and_then(Value::as_str),
13333 Some("awaiting_confirmation")
13334 );
13335 assert_eq!(observed.call_count(), 4);
13336
13337 let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
13338 let mut completed = String::new();
13339 while let Some(chunk) = completed_stream.next().await {
13340 match chunk {
13341 StreamChunk::Content { text } => completed.push_str(&text),
13342 StreamChunk::Done {} => break,
13343 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13344 _ => {}
13345 }
13346 }
13347 assert_eq!(completed, "Request executed.");
13348 assert_eq!(observed.call_count(), 6);
13349 }
13350
13351 #[tokio::test]
13353 async fn root_turn_gate_serializes_blocking_and_streaming_entry_points() {
13354 let (complete_entered, mut complete_events) = tokio::sync::mpsc::unbounded_channel();
13355 let agent = Arc::new(
13356 AgentBuilder::new()
13357 .system_prompt("Serialize root turns.")
13358 .llm(Arc::new(RootTurnProbeProvider { complete_entered }))
13359 .build()
13360 .unwrap(),
13361 );
13362 let blocking_agent = Arc::clone(&agent);
13363
13364 let legacy_stream = agent.chat_stream("stream owner").await.unwrap();
13365 assert!(agent.root_turn_gate.try_lock().is_err());
13366 let blocking = tokio::spawn(async move { blocking_agent.chat("blocked").await.unwrap() });
13367 assert!(
13368 tokio::time::timeout(std::time::Duration::from_millis(50), complete_events.recv())
13369 .await
13370 .is_err(),
13371 "blocking turn reached the provider while the legacy stream owned the root gate"
13372 );
13373
13374 drop(legacy_stream);
13375 assert_eq!(
13376 tokio::time::timeout(std::time::Duration::from_secs(2), complete_events.recv())
13377 .await
13378 .expect("blocking turn did not enter after stream drop"),
13379 Some(())
13380 );
13381 let response = tokio::time::timeout(std::time::Duration::from_secs(2), blocking)
13382 .await
13383 .expect("blocking turn did not finish after stream drop")
13384 .unwrap();
13385 assert_eq!(response.content, "blocking complete");
13386
13387 let mut event_stream = agent.chat_stream_events("event terminal").await.unwrap();
13388 assert!(agent.root_turn_gate.try_lock().is_err());
13389 let mut saw_final = false;
13390 while let Some(event) = event_stream.next().await {
13391 if matches!(event, AgentStreamEvent::Final(_)) {
13392 saw_final = true;
13393 break;
13394 }
13395 }
13396 assert!(saw_final);
13397 assert!(
13398 agent.root_turn_gate.try_lock().is_ok(),
13399 "authoritative terminal event retained the root gate"
13400 );
13401 }
13402
13403 #[tokio::test]
13405 async fn response_hook_rejects_same_runtime_chat_reentry() {
13406 let hooks = Arc::new(ResponseChatHooks {
13407 target: parking_lot::Mutex::new(None),
13408 invoked: AtomicBool::new(false),
13409 nested_result: parking_lot::Mutex::new(None),
13410 });
13411 let agent = Arc::new(
13412 AgentBuilder::new()
13413 .system_prompt("Reject response hook reentry.")
13414 .llm(Arc::new(mock_with_response("outer response")))
13415 .hooks(hooks.clone())
13416 .build()
13417 .unwrap(),
13418 );
13419 *hooks.target.lock() = Some(Arc::downgrade(&agent));
13420
13421 let response = tokio::time::timeout(
13422 std::time::Duration::from_secs(2),
13423 agent.chat("outer request"),
13424 )
13425 .await
13426 .expect("same-runtime response hook reentry must fail without deadlocking")
13427 .unwrap();
13428
13429 assert_eq!(response.content, "outer response");
13430 let nested_result = hooks
13431 .nested_result
13432 .lock()
13433 .clone()
13434 .expect("response hook must record its nested call");
13435 let error = nested_result.expect_err("same-runtime nested chat must be rejected");
13436 assert!(error.contains("reentrant root turn ownership"));
13437 }
13438
13439 #[tokio::test]
13441 async fn root_turn_gate_allows_nested_runtime_and_rejects_cycles() {
13442 let agent_a = AgentBuilder::new()
13443 .system_prompt("Runtime A.")
13444 .llm(Arc::new(mock_with_response("response A")))
13445 .build()
13446 .unwrap();
13447 let agent_b = AgentBuilder::new()
13448 .system_prompt("Runtime B.")
13449 .llm(Arc::new(mock_with_response("response B")))
13450 .build()
13451 .unwrap();
13452 let RootTurnAdmission {
13453 guard: guard_a,
13454 identity_stack: stack_a,
13455 } = agent_a.acquire_root_turn().await.unwrap();
13456
13457 let cycle_error = scope_runtime_gate_identity_stack(&stack_a, async {
13458 let RootTurnAdmission {
13459 guard: guard_b,
13460 identity_stack: stack_b,
13461 } = agent_b
13462 .acquire_root_turn()
13463 .await
13464 .expect("runtime B must acquire a different gate");
13465 let result =
13466 scope_runtime_gate_identity_stack(&stack_b, agent_a.acquire_root_turn()).await;
13467 drop(guard_b);
13468 match result {
13469 Err(error) => error,
13470 Ok(_) => panic!("runtime A accepted a repeated gate identity"),
13471 }
13472 })
13473 .await;
13474 drop(guard_a);
13475
13476 assert!(
13477 cycle_error
13478 .to_string()
13479 .contains("reentrant root turn ownership")
13480 );
13481 }
13482
13483 #[tokio::test]
13485 async fn concurrent_orchestration_propagates_root_gate_ancestry() {
13486 let registry = Arc::new(crate::spawner::AgentRegistry::new());
13487 let hooks_a = Arc::new(ConcurrentResponseHooks {
13488 registry: Arc::downgrade(®istry),
13489 child_id: "runtime-b".to_string(),
13490 invoked: AtomicBool::new(false),
13491 nested_result: parking_lot::Mutex::new(None),
13492 });
13493 let hooks_b = Arc::new(ResponseChatHooks {
13494 target: parking_lot::Mutex::new(None),
13495 invoked: AtomicBool::new(false),
13496 nested_result: parking_lot::Mutex::new(None),
13497 });
13498 let agent_a = AgentBuilder::new()
13499 .system_prompt("Runtime A dispatches runtime B concurrently.")
13500 .llm(Arc::new(mock_with_response("response A")))
13501 .hooks(hooks_a.clone())
13502 .build()
13503 .unwrap();
13504 let agent_b = AgentBuilder::new()
13505 .system_prompt("Runtime B attempts to re-enter runtime A.")
13506 .llm(Arc::new(mock_with_response("response B")))
13507 .hooks(hooks_b.clone())
13508 .build()
13509 .unwrap();
13510 let spec_a = crate::spec::AgentSpec {
13511 name: "runtime-a".to_string(),
13512 system_prompt: "Runtime A dispatches runtime B concurrently.".to_string(),
13513 ..crate::spec::AgentSpec::default()
13514 };
13515 let spec_b = crate::spec::AgentSpec {
13516 name: "runtime-b".to_string(),
13517 system_prompt: "Runtime B attempts to re-enter runtime A.".to_string(),
13518 ..crate::spec::AgentSpec::default()
13519 };
13520 registry
13521 .register(crate::spawner::SpawnedAgent::from_runtime(
13522 "runtime-a".to_string(),
13523 agent_a,
13524 spec_a,
13525 ))
13526 .await
13527 .unwrap();
13528 registry
13529 .register(crate::spawner::SpawnedAgent::from_runtime(
13530 "runtime-b".to_string(),
13531 agent_b,
13532 spec_b,
13533 ))
13534 .await
13535 .unwrap();
13536 let runtime_a = registry.get("runtime-a").unwrap();
13537 *hooks_b.target.lock() = Some(Arc::downgrade(&runtime_a));
13538
13539 let response = tokio::time::timeout(
13540 std::time::Duration::from_secs(2),
13541 runtime_a.chat("outer concurrent request"),
13542 )
13543 .await
13544 .expect("concurrent orchestration cycle must fail without deadlocking")
13545 .unwrap();
13546
13547 assert_eq!(response.content, "response A");
13548 let child_result = hooks_a
13549 .nested_result
13550 .lock()
13551 .clone()
13552 .expect("runtime A hook must record runtime B completion");
13553 assert_eq!(child_result.unwrap(), "response B");
13554 let cycle_result = hooks_b
13555 .nested_result
13556 .lock()
13557 .clone()
13558 .expect("runtime B hook must record runtime A reentry");
13559 assert!(
13560 cycle_result
13561 .expect_err("runtime A accepted a repeated gate identity")
13562 .contains("reentrant root turn ownership")
13563 );
13564 }
13565
13566 #[tokio::test]
13568 async fn confirmed_skill_route_executes_exactly_once() {
13569 let (agent, observed) = state_disambiguation_agent_with_skills(
13570 vec![
13571 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13572 "send_report",
13573 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13574 r#"{"question":"What should I send?","options":null}"#,
13575 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13576 r#"{"question":"Should I send the report to Ada?"}"#,
13577 r#"{"status":"confirmed"}"#,
13578 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
13579 "Report skill executed.",
13580 ],
13581 true,
13582 None,
13583 true,
13584 vec![confirmation_skill()],
13585 );
13586
13587 let clarification = agent.chat("Send it").await.unwrap();
13588 assert_eq!(clarification.content, "What should I send?");
13589 assert_eq!(confirmation_skill_call_count(&observed), 0);
13590
13591 let confirmation = agent.chat("The report to Ada").await.unwrap();
13592 assert_eq!(confirmation.content, "Should I send the report to Ada?");
13593 assert_eq!(
13594 confirmation
13595 .metadata
13596 .as_ref()
13597 .and_then(|metadata| metadata.get("disambiguation"))
13598 .and_then(|metadata| metadata.get("status"))
13599 .and_then(Value::as_str),
13600 Some("awaiting_confirmation")
13601 );
13602 assert_eq!(confirmation_skill_call_count(&observed), 0);
13603
13604 let completed = agent.chat("Yes").await.unwrap();
13605 assert_eq!(completed.content, "Report skill executed.");
13606 assert_eq!(confirmation_skill_call_count(&observed), 1);
13607 assert!(agent.pending_skill_id.read().is_none());
13608 let messages = agent.memory.get_messages(None).await.unwrap();
13609 assert!(!messages.iter().any(|message| message.content == "Yes"));
13610 }
13611
13612 #[tokio::test]
13614 async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
13615 let (agent, observed) = state_disambiguation_agent_with_skills(
13616 vec![
13617 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13618 "send_report",
13619 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13620 r#"{"question":"What should I send?","options":null}"#,
13621 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13622 r#"{"question":"Should I send the report to Ada?"}"#,
13623 r#"{"status":"confirmed"}"#,
13624 r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
13625 r#"{"question":"When should I send it?","options":null}"#,
13626 ],
13627 true,
13628 None,
13629 true,
13630 vec![confirmation_skill()],
13631 );
13632
13633 agent.chat("Send it").await.unwrap();
13634 agent.chat("The report to Ada").await.unwrap();
13635 let follow_up = agent.chat("Yes").await.unwrap();
13636
13637 assert_eq!(follow_up.content, "When should I send it?");
13638 let metadata = follow_up
13639 .metadata
13640 .as_ref()
13641 .and_then(|metadata| metadata.get("disambiguation"))
13642 .unwrap();
13643 assert_eq!(
13644 metadata.get("status").and_then(Value::as_str),
13645 Some("awaiting_clarification")
13646 );
13647 assert_eq!(
13648 metadata.get("skill_id").and_then(Value::as_str),
13649 Some("send_report")
13650 );
13651 assert!(metadata.get("detection").is_some());
13652 assert_eq!(confirmation_skill_call_count(&observed), 0);
13653 }
13654
13655 #[tokio::test]
13657 async fn rejected_skill_confirmation_never_executes() {
13658 let (agent, observed) = state_disambiguation_agent_with_skills(
13659 vec![
13660 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13661 "send_report",
13662 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13663 r#"{"question":"What should I send?","options":null}"#,
13664 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13665 r#"{"question":"Should I send the report to Ada?"}"#,
13666 r#"{"status":"rejected"}"#,
13667 "Confirmation rejected.",
13668 ],
13669 true,
13670 None,
13671 true,
13672 vec![confirmation_skill()],
13673 );
13674
13675 agent.chat("Send it").await.unwrap();
13676 agent.chat("The report to Ada").await.unwrap();
13677 let rejected = agent.chat("No").await.unwrap();
13678
13679 assert_eq!(rejected.content, "Confirmation rejected.");
13680 assert_eq!(confirmation_skill_call_count(&observed), 0);
13681 assert!(agent.pending_skill_id.read().is_none());
13682 }
13683
13684 #[tokio::test]
13686 async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
13687 let (agent, observed) = state_disambiguation_agent_with_skills(
13688 vec![
13689 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13690 "send_report",
13691 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13692 r#"{"question":"What should I send?","options":null}"#,
13693 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13694 r#"{"question":"Should I send the report to Ada?"}"#,
13695 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13696 "none",
13697 "Fresh response.",
13698 ],
13699 true,
13700 None,
13701 true,
13702 vec![confirmation_skill()],
13703 );
13704
13705 agent.chat("Send it").await.unwrap();
13706 agent.chat("The report to Ada").await.unwrap();
13707 agent.reset().await.unwrap();
13708 assert!(agent.pending_skill_id.read().is_none());
13709 assert!(
13710 !agent
13711 .disambiguation_manager()
13712 .unwrap()
13713 .has_pending_clarification()
13714 .await
13715 );
13716
13717 let mut stream = agent.chat_stream("Yes").await.unwrap();
13718 let mut content = String::new();
13719 while let Some(chunk) = stream.next().await {
13720 match chunk {
13721 StreamChunk::Content { text } => content.push_str(&text),
13722 StreamChunk::Done {} => break,
13723 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13724 _ => {}
13725 }
13726 }
13727
13728 assert_eq!(content, "Fresh response.");
13729 assert_eq!(confirmation_skill_call_count(&observed), 0);
13730 }
13731
13732 #[tokio::test]
13734 async fn trait_reset_clears_pending_skill_confirmation() {
13735 let (agent, _) = state_disambiguation_agent_with_skills(
13736 vec![
13737 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13738 "send_report",
13739 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13740 r#"{"question":"What should I send?","options":null}"#,
13741 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13742 r#"{"question":"Should I send the report to Ada?"}"#,
13743 ],
13744 true,
13745 None,
13746 true,
13747 vec![confirmation_skill()],
13748 );
13749
13750 agent.chat("Send it").await.unwrap();
13751 agent.chat("The report to Ada").await.unwrap();
13752 <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13753
13754 assert!(agent.pending_skill_id.read().is_none());
13755 assert!(
13756 !agent
13757 .disambiguation_manager()
13758 .unwrap()
13759 .has_pending_clarification()
13760 .await
13761 );
13762 }
13763
13764 #[tokio::test]
13766 async fn state_change_invalidates_pending_skill_confirmation() {
13767 let (agent, observed) = state_disambiguation_agent_with_skills(
13768 vec![
13769 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13770 "send_report",
13771 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13772 r#"{"question":"What should I send?","options":null}"#,
13773 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13774 r#"{"question":"Should I send the report to Ada?"}"#,
13775 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13776 "none",
13777 "Fresh response.",
13778 ],
13779 true,
13780 None,
13781 true,
13782 vec![confirmation_skill()],
13783 );
13784
13785 agent.chat("Send it").await.unwrap();
13786 agent.chat("The report to Ada").await.unwrap();
13787 agent.transition_to("review").await.unwrap();
13788 let cancelled = agent.chat("Yes").await.unwrap();
13789
13790 assert_eq!(cancelled.content, "Fresh response.");
13791 assert_eq!(confirmation_skill_call_count(&observed), 0);
13792 assert!(agent.pending_skill_id.read().is_none());
13793 }
13794
13795 #[tokio::test]
13797 async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13798 let (mut agent, observed) = state_disambiguation_agent_with_skills(
13799 vec![
13800 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13801 "send_report",
13802 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13803 r#"{"question":"What should I send?","options":null}"#,
13804 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13805 r#"{"question":"Should I send the report to Ada?"}"#,
13806 r#"{"status":"confirmed"}"#,
13807 "Confirmation cancelled.",
13808 ],
13809 true,
13810 None,
13811 true,
13812 vec![confirmation_skill()],
13813 );
13814 let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13815 let manager = agent
13816 .disambiguation_manager
13817 .take()
13818 .unwrap()
13819 .with_clarification_observer(observer.clone());
13820 agent.disambiguation_manager = Some(manager);
13821 let agent = Arc::new(agent);
13822
13823 agent.chat("Send it").await.unwrap();
13824 agent.chat("The report to Ada").await.unwrap();
13825
13826 let confirming_agent = Arc::clone(&agent);
13827 let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13828 observer.entered.wait().await;
13829 agent.reset().await.unwrap();
13830 observer.release.notify_one();
13831
13832 let response = confirmation.await.unwrap().unwrap();
13833 assert_eq!(response.content, "Confirmation cancelled.");
13834 assert_eq!(confirmation_skill_call_count(&observed), 0);
13835 assert!(agent.pending_skill_id.read().is_none());
13836 }
13837
13838 #[tokio::test]
13840 async fn queued_reset_prevents_stale_confirmation_question_publication() {
13841 let (agent, observed) = state_disambiguation_agent(
13842 vec![
13843 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13844 r#"{"question":"What should I send?","options":null}"#,
13845 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13846 r#"{"question":"Should I send the report to Ada?"}"#,
13847 ],
13848 true,
13849 None,
13850 true,
13851 );
13852 let agent = Arc::new(agent);
13853 agent.chat("Send it").await.unwrap();
13854
13855 let admission = agent.disambiguation_admission.write().await;
13856 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13857 let resetting_agent = Arc::clone(&agent);
13858 let reset = tokio::spawn(async move {
13859 let _ = started_tx.send(());
13860 resetting_agent.reset().await
13861 });
13862 started_rx.await.unwrap();
13863 tokio::task::yield_now().await;
13864
13865 let responding_agent = Arc::clone(&agent);
13866 let response =
13867 tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13868 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13869 while observed.call_count() < 4 {
13870 tokio::task::yield_now().await;
13871 }
13872 })
13873 .await
13874 .expect("clarification processing must reach terminal publication");
13875 drop(admission);
13876
13877 reset.await.unwrap().unwrap();
13878 let error = response.await.unwrap().unwrap_err();
13879 assert!(error.to_string().contains("ownership changed"));
13880 assert!(
13881 !agent
13882 .disambiguation_manager()
13883 .unwrap()
13884 .has_pending_clarification()
13885 .await
13886 );
13887 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13888 }
13889
13890 #[tokio::test]
13892 async fn queued_reset_prevents_stale_skill_clarification_publication() {
13893 let (agent, observed) = state_disambiguation_agent_with_skills(
13894 vec![
13895 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13896 "send_report",
13897 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13898 r#"{"question":"What should I send?","options":null}"#,
13899 ],
13900 true,
13901 None,
13902 true,
13903 vec![confirmation_skill()],
13904 );
13905 let agent = Arc::new(agent);
13906 let admission = agent.disambiguation_admission.write().await;
13907 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13908 let resetting_agent = Arc::clone(&agent);
13909 let reset = tokio::spawn(async move {
13910 let _ = started_tx.send(());
13911 resetting_agent.reset().await
13912 });
13913 started_rx.await.unwrap();
13914 tokio::task::yield_now().await;
13915
13916 let responding_agent = Arc::clone(&agent);
13917 let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13918 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13919 while observed.call_count() < 4 {
13920 tokio::task::yield_now().await;
13921 }
13922 })
13923 .await
13924 .expect("skill clarification must reach terminal publication");
13925 drop(admission);
13926
13927 reset.await.unwrap().unwrap();
13928 let error = response.await.unwrap().unwrap_err();
13929 assert!(error.to_string().contains("ownership changed"));
13930 assert_eq!(confirmation_skill_call_count(&observed), 0);
13931 assert!(agent.pending_skill_id.read().is_none());
13932 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13933 }
13934
13935 #[tokio::test]
13937 async fn transition_hook_can_reset_without_admission_deadlock() {
13938 let hooks = Arc::new(ResetOnTransitionHooks {
13939 agent: parking_lot::Mutex::new(None),
13940 invoked: AtomicBool::new(false),
13941 });
13942 let agent = Arc::new(
13943 AgentBuilder::new()
13944 .system_prompt("Test transition hook reentrancy.")
13945 .llm(Arc::new(mock_with_response("done")))
13946 .state_machine(disambiguation_state_machine(None, false))
13947 .build()
13948 .unwrap()
13949 .with_hooks(hooks.clone()),
13950 );
13951 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13952
13953 let transitioned = tokio::time::timeout(
13954 std::time::Duration::from_secs(2),
13955 agent.apply_transition_target("active", "review", "test transition", None),
13956 )
13957 .await
13958 .expect("transition hook reset must not deadlock")
13959 .unwrap();
13960
13961 assert!(transitioned);
13962 assert!(hooks.invoked.load(Ordering::SeqCst));
13963 assert_eq!(agent.current_state().as_deref(), Some("active"));
13964 }
13965
13966 #[tokio::test]
13968 async fn concurrent_transition_cannot_duplicate_exit_actions() {
13969 let gate = PathMutationGate::new();
13970 let active = ai_agents_state::StateDefinition {
13971 on_exit: vec![StateAction::Tool {
13972 tool: "transition_exit".to_string(),
13973 args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13974 }],
13975 ..Default::default()
13976 };
13977 let state_machine = Arc::new(
13978 StateMachine::new(ai_agents_state::StateConfig {
13979 initial: "active".to_string(),
13980 states: HashMap::from([
13981 ("active".to_string(), active),
13982 (
13983 "review".to_string(),
13984 ai_agents_state::StateDefinition::default(),
13985 ),
13986 ]),
13987 global_transitions: Vec::new(),
13988 fallback: None,
13989 max_no_transition: None,
13990 regenerate_on_transition: true,
13991 })
13992 .unwrap(),
13993 );
13994 let agent = Arc::new(
13995 AgentBuilder::new()
13996 .system_prompt("Test transition reservation.")
13997 .llm(Arc::new(mock_with_response("done")))
13998 .tool(Arc::new(BlockingPathMutationTool {
13999 id: "transition_exit",
14000 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14001 gate: gate.clone(),
14002 }))
14003 .state_machine(state_machine)
14004 .build()
14005 .unwrap(),
14006 );
14007
14008 let first_agent = Arc::clone(&agent);
14009 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
14010 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
14011 .await
14012 .expect("reserved transition must enter its exit action");
14013
14014 let second = tokio::time::timeout(
14015 std::time::Duration::from_secs(2),
14016 agent.transition_to("review"),
14017 )
14018 .await
14019 .expect("competing transition must fail without waiting for the exit action")
14020 .unwrap_err();
14021 assert!(second.to_string().contains("already in progress"));
14022
14023 gate.release();
14024 first.await.unwrap().unwrap();
14025 assert_eq!(agent.current_state().as_deref(), Some("review"));
14026 }
14027
14028 #[tokio::test]
14030 async fn concurrent_transition_cannot_overtake_enter_actions() {
14031 let gate = PathMutationGate::new();
14032 let review = ai_agents_state::StateDefinition {
14033 on_enter: vec![StateAction::Tool {
14034 tool: "transition_enter".to_string(),
14035 args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
14036 }],
14037 ..Default::default()
14038 };
14039 let state_machine = Arc::new(
14040 StateMachine::new(ai_agents_state::StateConfig {
14041 initial: "active".to_string(),
14042 states: HashMap::from([
14043 (
14044 "active".to_string(),
14045 ai_agents_state::StateDefinition::default(),
14046 ),
14047 ("review".to_string(), review),
14048 ]),
14049 global_transitions: Vec::new(),
14050 fallback: None,
14051 max_no_transition: None,
14052 regenerate_on_transition: true,
14053 })
14054 .unwrap(),
14055 );
14056 let agent = Arc::new(
14057 AgentBuilder::new()
14058 .system_prompt("Test transition lifecycle reservation.")
14059 .llm(Arc::new(mock_with_response("done")))
14060 .tool(Arc::new(BlockingPathMutationTool {
14061 id: "transition_enter",
14062 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14063 gate: gate.clone(),
14064 }))
14065 .state_machine(state_machine)
14066 .build()
14067 .unwrap(),
14068 );
14069
14070 let first_agent = Arc::clone(&agent);
14071 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
14072 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
14073 .await
14074 .expect("committed transition must enter its destination action");
14075
14076 let second = agent.transition_to("active").await.unwrap_err();
14077 assert!(second.to_string().contains("already in progress"));
14078 assert!(agent.reset().await.is_err());
14079
14080 gate.release();
14081 first.await.unwrap().unwrap();
14082 assert_eq!(agent.current_state().as_deref(), Some("review"));
14083 }
14084
14085 #[tokio::test]
14087 async fn same_state_restore_invalidates_pending_skill_confirmation() {
14088 let (agent, observed) = state_disambiguation_agent_with_skills(
14089 vec![
14090 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
14091 "send_report",
14092 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
14093 r#"{"question":"What should I send?","options":null}"#,
14094 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
14095 r#"{"question":"Should I send the report to Ada?"}"#,
14096 ],
14097 true,
14098 None,
14099 true,
14100 vec![confirmation_skill()],
14101 );
14102
14103 agent.chat("Send it").await.unwrap();
14104 agent.chat("The report to Ada").await.unwrap();
14105 let snapshot = agent.save_state().await.unwrap();
14106 assert_eq!(agent.current_state().as_deref(), Some("active"));
14107
14108 agent.restore_state(snapshot).await.unwrap();
14109
14110 assert_eq!(agent.current_state().as_deref(), Some("active"));
14111 assert!(agent.pending_skill_id.read().is_none());
14112 assert!(
14113 !agent
14114 .disambiguation_manager()
14115 .unwrap()
14116 .has_pending_clarification()
14117 .await
14118 );
14119 assert_eq!(confirmation_skill_call_count(&observed), 0);
14120 }
14121
14122 #[tokio::test]
14124 async fn direct_state_generation_change_invalidates_confirmation() {
14125 let (agent, observed) = state_disambiguation_agent_with_skills(
14126 vec![
14127 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
14128 "send_report",
14129 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
14130 r#"{"question":"What should I send?","options":null}"#,
14131 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
14132 r#"{"question":"Should I send the report to Ada?"}"#,
14133 "Confirmation cancelled.",
14134 ],
14135 true,
14136 None,
14137 true,
14138 vec![confirmation_skill()],
14139 );
14140
14141 agent.chat("Send it").await.unwrap();
14142 agent.chat("The report to Ada").await.unwrap();
14143 let state_machine = agent.state_machine().unwrap();
14144 state_machine
14145 .transition_to("review", "external test")
14146 .unwrap();
14147 state_machine
14148 .transition_to("active", "external test")
14149 .unwrap();
14150
14151 let response = agent.chat("Yes").await.unwrap();
14152
14153 assert_eq!(response.content, "Confirmation cancelled.");
14154 assert_eq!(confirmation_skill_call_count(&observed), 0);
14155 assert!(agent.pending_skill_id.read().is_none());
14156 }
14157
14158 #[tokio::test]
14159 async fn state_confirmation_does_not_add_a_question_for_clear_input() {
14160 let (agent, observed) = state_disambiguation_agent(
14161 vec![
14162 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
14163 "Request executed.",
14164 ],
14165 true,
14166 None,
14167 true,
14168 );
14169
14170 let response = agent.chat("Send the report to Ada").await.unwrap();
14171
14172 assert_eq!(response.content, "Request executed.");
14173 assert_eq!(observed.call_count(), 2);
14174 }
14175
14176 #[tokio::test]
14177 async fn state_override_cannot_activate_a_disabled_top_level_manager() {
14178 let (agent, observed) =
14179 state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
14180
14181 assert!(!agent.has_disambiguation());
14182 let response = agent.chat("Send it").await.unwrap();
14183
14184 assert_eq!(response.content, "Request executed.");
14185 assert_eq!(observed.call_count(), 1);
14186 }
14187
14188 #[tokio::test]
14189 async fn native_required_choice_executes_through_the_shared_tool_path() {
14190 let mut mock = MockLLMProvider::new("native-required");
14191 mock.set_tool_choice(Some(ToolChoice::Required));
14192 mock.add_response(
14193 LLMResponse::new("", FinishReason::ToolCall)
14194 .with_tool_calls(vec![ToolCall {
14195 id: "provider-call-1".to_string(),
14196 name: "calculator".to_string(),
14197 arguments: serde_json::json!({"expression": "2 + 2"}),
14198 }])
14199 .unwrap(),
14200 );
14201 mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
14202 let observed = mock.clone();
14203 let agent = AgentBuilder::new()
14204 .system_prompt("Use the calculator when needed.")
14205 .llm(Arc::new(mock))
14206 .tool(Arc::new(CalculatorTool::new()))
14207 .build()
14208 .unwrap();
14209
14210 let response = agent.chat("What is 2 + 2?").await.unwrap();
14211
14212 assert_eq!(response.content, "The answer is 4.");
14213 assert_eq!(
14214 response.tool_calls.as_ref().unwrap()[0].id,
14215 "provider-call-1"
14216 );
14217 let calls = observed.call_history();
14218 assert_eq!(calls.len(), 2);
14219 assert!(matches!(
14220 calls[0].request.as_ref().map(|request| &request.choice),
14221 Some(ToolChoice::Required)
14222 ));
14223 assert!(matches!(
14224 calls[1].request.as_ref().map(|request| &request.choice),
14225 Some(ToolChoice::Auto)
14226 ));
14227 }
14228
14229 #[tokio::test]
14230 async fn prompt_fallback_uses_one_corrective_retry() {
14231 let mut mock = MockLLMProvider::new("prompt-required");
14232 mock.set_tool_choice(Some(ToolChoice::Required));
14233 mock.set_native_tool_support(false);
14234 mock.set_responses(
14235 vec![
14236 "I can calculate that.".to_string(),
14237 r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
14238 "The answer is 4.".to_string(),
14239 ],
14240 false,
14241 );
14242 let observed = mock.clone();
14243 let agent = AgentBuilder::new()
14244 .system_prompt("Use tools.")
14245 .llm(Arc::new(mock))
14246 .tool(Arc::new(CalculatorTool::new()))
14247 .build()
14248 .unwrap();
14249
14250 let response = agent.chat("What is 2 + 2?").await.unwrap();
14251
14252 assert_eq!(response.content, "The answer is 4.");
14253 assert_eq!(observed.call_count(), 3);
14254 let corrective = &observed.call_history()[1].messages;
14255 assert!(
14256 corrective
14257 .last()
14258 .unwrap()
14259 .content
14260 .contains("previous response")
14261 );
14262 }
14263
14264 #[tokio::test]
14265 async fn prompt_fallback_fails_after_one_noncompliant_retry() {
14266 let mut mock = MockLLMProvider::new("prompt-required-failure");
14267 mock.set_tool_choice(Some(ToolChoice::Required));
14268 mock.set_native_tool_support(false);
14269 mock.set_responses(
14270 vec!["No tool.".to_string(), "Still no tool.".to_string()],
14271 false,
14272 );
14273 let observed = mock.clone();
14274 let agent = AgentBuilder::new()
14275 .system_prompt("Use tools.")
14276 .llm(Arc::new(mock))
14277 .tool(Arc::new(CalculatorTool::new()))
14278 .build()
14279 .unwrap();
14280
14281 let error = agent.chat("What is 2 + 2?").await.unwrap_err();
14282
14283 assert!(error.to_string().contains("one corrective retry"));
14284 assert_eq!(observed.call_count(), 2);
14285 }
14286
14287 #[tokio::test]
14288 async fn specific_choice_cannot_widen_the_effective_grant() {
14289 let mut mock = MockLLMProvider::new("specific-outside-grant");
14290 mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
14291 let observed = mock.clone();
14292 let agent = AgentBuilder::new()
14293 .system_prompt("Use tools.")
14294 .llm(Arc::new(mock))
14295 .tool(Arc::new(CalculatorTool::new()))
14296 .build()
14297 .unwrap();
14298
14299 let error = agent.chat("Generate a value.").await.unwrap_err();
14300
14301 assert!(error.to_string().contains("is not registered"));
14302 assert_eq!(observed.call_count(), 0);
14303 }
14304
14305 #[tokio::test]
14306 async fn none_choice_exposes_no_tool_protocol() {
14307 let mut mock = MockLLMProvider::new("no-tools");
14308 mock.set_tool_choice(Some(ToolChoice::None));
14309 mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
14310 let observed = mock.clone();
14311 let agent = AgentBuilder::new()
14312 .system_prompt("Answer directly.")
14313 .llm(Arc::new(mock))
14314 .tool(Arc::new(CalculatorTool::new()))
14315 .build()
14316 .unwrap();
14317
14318 let response = agent.chat("Hello").await.unwrap();
14319
14320 assert!(response.tool_calls.is_none());
14321 assert_eq!(observed.call_count(), 1);
14322 let call = observed.last_call().unwrap();
14323 assert!(call.request.is_none());
14324 assert!(
14325 call.messages
14326 .iter()
14327 .all(|message| !message.content.contains("Available tools:"))
14328 );
14329 }
14330
14331 struct RuntimeStorage {
14332 capabilities: Box<[StorageCapability]>,
14333 snapshots: RwLock<HashMap<String, AgentSnapshot>>,
14334 metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
14335 metadata_save_calls: AtomicU64,
14336 metadata_load_calls: AtomicU64,
14337 fail_metadata_save: AtomicBool,
14338 fail_metadata_load: AtomicBool,
14339 }
14340
14341 impl RuntimeStorage {
14342 fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
14343 Self {
14344 capabilities: capabilities.into_iter().collect(),
14345 snapshots: RwLock::new(HashMap::new()),
14346 metadata: RwLock::new(HashMap::new()),
14347 metadata_save_calls: AtomicU64::new(0),
14348 metadata_load_calls: AtomicU64::new(0),
14349 fail_metadata_save: AtomicBool::new(false),
14350 fail_metadata_load: AtomicBool::new(false),
14351 }
14352 }
14353 }
14354
14355 #[async_trait]
14356 impl AgentStorage for RuntimeStorage {
14357 fn supports(&self, capability: StorageCapability) -> bool {
14358 self.capabilities.contains(&capability)
14359 }
14360
14361 async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
14362 self.snapshots
14363 .write()
14364 .insert(session_id.to_string(), snapshot.clone());
14365 Ok(())
14366 }
14367
14368 async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
14369 Ok(self.snapshots.read().get(session_id).cloned())
14370 }
14371
14372 async fn delete(&self, session_id: &str) -> Result<()> {
14373 self.snapshots.write().remove(session_id);
14374 Ok(())
14375 }
14376
14377 async fn list_sessions(&self) -> Result<Vec<String>> {
14378 Ok(self.snapshots.read().keys().cloned().collect())
14379 }
14380
14381 async fn save_snapshot_with_metadata(
14382 &self,
14383 session_id: &str,
14384 snapshot: &AgentSnapshot,
14385 metadata: &ai_agents_core::SessionMetadata,
14386 ) -> Result<()> {
14387 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14388 if self.fail_metadata_save.load(Ordering::SeqCst) {
14389 return Err(AgentError::Persistence("metadata save failed".into()));
14390 }
14391 self.snapshots
14392 .write()
14393 .insert(session_id.to_string(), snapshot.clone());
14394 self.metadata
14395 .write()
14396 .insert(session_id.to_string(), metadata.clone());
14397 Ok(())
14398 }
14399
14400 async fn save_metadata(
14401 &self,
14402 session_id: &str,
14403 metadata: &ai_agents_core::SessionMetadata,
14404 ) -> Result<()> {
14405 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
14406 if self.fail_metadata_save.load(Ordering::SeqCst) {
14407 return Err(AgentError::Persistence("metadata save failed".into()));
14408 }
14409 self.metadata
14410 .write()
14411 .insert(session_id.to_string(), metadata.clone());
14412 Ok(())
14413 }
14414
14415 async fn load_metadata(
14416 &self,
14417 session_id: &str,
14418 ) -> Result<Option<ai_agents_core::SessionMetadata>> {
14419 self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
14420 if self.fail_metadata_load.load(Ordering::SeqCst) {
14421 return Err(AgentError::Persistence("metadata load failed".into()));
14422 }
14423 Ok(self.metadata.read().get(session_id).cloned())
14424 }
14425 }
14426
14427 fn runtime_storage_agent() -> RuntimeAgent {
14428 AgentBuilder::new()
14429 .system_prompt("Test runtime storage integration.")
14430 .llm(Arc::new(mock_with_response("done")))
14431 .build()
14432 .unwrap()
14433 }
14434
14435 fn restore_spec(id: &str) -> crate::spec::AgentSpec {
14436 crate::spec::AgentSpec {
14437 name: id.to_string(),
14438 system_prompt: format!("Restore child {id}."),
14439 ..crate::spec::AgentSpec::default()
14440 }
14441 }
14442
14443 fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
14444 ai_agents_core::SpawnedAgentEntry {
14445 id: id.to_string(),
14446 name: id.to_string(),
14447 spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
14448 }
14449 }
14450
14451 fn restore_spawner(
14452 storage: Arc<RuntimeStorage>,
14453 max_agents: usize,
14454 ) -> (
14455 Arc<crate::spawner::AgentSpawner>,
14456 Arc<crate::spawner::AgentRegistry>,
14457 ) {
14458 let mut llms = LLMRegistry::new();
14459 llms.register("default", Arc::new(mock_with_response("done")));
14460 (
14461 Arc::new(
14462 crate::spawner::AgentSpawner::new()
14463 .with_shared_llms(llms)
14464 .with_shared_storage(storage)
14465 .with_max_agents(max_agents),
14466 ),
14467 Arc::new(crate::spawner::AgentRegistry::new()),
14468 )
14469 }
14470
14471 async fn save_restore_target(
14472 parent: &RuntimeAgent,
14473 storage: &RuntimeStorage,
14474 session_id: &str,
14475 entries: Vec<ai_agents_core::SpawnedAgentEntry>,
14476 ) {
14477 let mut snapshot = parent.save_state().await.unwrap();
14478 snapshot.spawned_agents = Some(entries);
14479 storage.save(session_id, &snapshot).await.unwrap();
14480 storage
14481 .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
14482 .await
14483 .unwrap();
14484 }
14485
14486 #[tokio::test]
14487 async fn storage_init_requires_storage_for_actor_facts() {
14488 let facts = ai_agents_facts::FactsConfig {
14489 enabled: true,
14490 ..Default::default()
14491 };
14492 let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
14493
14494 let error = agent.init_storage().await.unwrap_err();
14495 assert!(matches!(
14496 error,
14497 AgentError::Config(message)
14498 if message.contains("actor facts or actor memory")
14499 && message.contains("none is configured or injected")
14500 ));
14501 }
14502
14503 #[tokio::test]
14504 async fn storage_init_validates_actor_facts_capability() {
14505 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14506 let actor_memory = ai_agents_facts::ActorMemoryConfig {
14507 enabled: true,
14508 ..Default::default()
14509 };
14510 let agent = runtime_storage_agent()
14511 .with_storage(storage)
14512 .with_facts_config(Some(actor_memory), None);
14513
14514 assert!(matches!(
14515 agent.init_storage().await,
14516 Err(AgentError::UnsupportedStorageCapability(
14517 StorageCapability::ActorFacts
14518 ))
14519 ));
14520 }
14521
14522 #[tokio::test]
14523 async fn blocking_chat_rejects_unsupported_required_storage() {
14524 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14525 let facts = ai_agents_facts::FactsConfig {
14526 enabled: true,
14527 ..Default::default()
14528 };
14529 let agent = runtime_storage_agent()
14530 .with_storage(storage)
14531 .with_facts_config(None, Some(facts));
14532
14533 assert!(matches!(
14534 agent.chat("hello").await,
14535 Err(AgentError::UnsupportedStorageCapability(
14536 StorageCapability::ActorFacts
14537 ))
14538 ));
14539 }
14540
14541 #[tokio::test]
14542 async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
14543 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14544 let config = ai_agents_relationships::RelationshipConfig {
14545 enabled: true,
14546 ..Default::default()
14547 };
14548 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14549 let agent = runtime_storage_agent()
14550 .with_storage(storage)
14551 .with_relationships(manager);
14552
14553 assert!(matches!(
14554 agent.chat_stream("hello").await,
14555 Err(AgentError::UnsupportedStorageCapability(
14556 StorageCapability::ActorRelationships
14557 ))
14558 ));
14559 }
14560
14561 #[tokio::test]
14562 async fn storage_init_completes_facts_for_injected_storage() {
14563 let storage = Arc::new(RuntimeStorage::new([
14564 StorageCapability::Snapshot,
14565 StorageCapability::ActorFacts,
14566 ]));
14567 let facts = ai_agents_facts::FactsConfig {
14568 enabled: true,
14569 ..Default::default()
14570 };
14571 let agent = runtime_storage_agent()
14572 .with_storage(storage)
14573 .with_facts_config(None, Some(facts));
14574
14575 agent.init_storage().await.unwrap();
14576 assert!(agent.fact_store().is_some());
14577 }
14578
14579 #[tokio::test]
14580 async fn storage_init_requires_storage_for_persistent_relationships() {
14581 let config = ai_agents_relationships::RelationshipConfig {
14582 enabled: true,
14583 ..Default::default()
14584 };
14585 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14586 let agent = runtime_storage_agent().with_relationships(manager);
14587
14588 let error = agent.init_storage().await.unwrap_err();
14589 assert!(matches!(
14590 error,
14591 AgentError::Config(message)
14592 if message.contains("persistent relationships")
14593 && message.contains("none is configured or injected")
14594 ));
14595 }
14596
14597 #[tokio::test]
14598 async fn storage_init_validates_persistent_relationships_capability() {
14599 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14600 let config = ai_agents_relationships::RelationshipConfig {
14601 enabled: true,
14602 ..Default::default()
14603 };
14604 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
14605 let agent = runtime_storage_agent()
14606 .with_storage(storage)
14607 .with_relationships(manager);
14608
14609 assert!(matches!(
14610 agent.init_storage().await,
14611 Err(AgentError::UnsupportedStorageCapability(
14612 StorageCapability::ActorRelationships
14613 ))
14614 ));
14615 }
14616
14617 #[tokio::test]
14618 async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
14619 let storage = Arc::new(RuntimeStorage::new([
14620 StorageCapability::Snapshot,
14621 StorageCapability::SessionMetadata,
14622 ]));
14623 let agent = runtime_storage_agent().with_storage(storage.clone());
14624 agent.set_actor_id("old-actor").unwrap();
14625 agent.save_session("old").await.unwrap();
14626 storage
14627 .save("target", &agent.save_state().await.unwrap())
14628 .await
14629 .unwrap();
14630 storage
14631 .save_metadata("target", &ai_agents_core::SessionMetadata::default())
14632 .await
14633 .unwrap();
14634
14635 assert!(agent.load_session("target").await.unwrap());
14636
14637 assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
14638 assert_eq!(agent.actor_id(), None);
14639 }
14640
14641 #[tokio::test]
14642 async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
14643 let storage = Arc::new(RuntimeStorage::new([
14644 StorageCapability::Snapshot,
14645 StorageCapability::SessionMetadata,
14646 ]));
14647 let (spawner, registry) = restore_spawner(storage.clone(), 3);
14648 let parent = runtime_storage_agent()
14649 .with_storage(storage.clone())
14650 .with_spawner_handles(Arc::clone(&spawner), Arc::clone(®istry));
14651
14652 for id in ["a", "b"] {
14653 let spawned = spawner
14654 .spawn_with_id(id.to_string(), restore_spec(id))
14655 .await
14656 .unwrap();
14657 spawned.agent.save_session("grow").await.unwrap();
14658 registry.register(spawned).await.unwrap();
14659 }
14660 let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
14661 staged_c
14662 .save("grow", &AgentSnapshot::new("c".into()))
14663 .await
14664 .unwrap();
14665 staged_c
14666 .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
14667 .await
14668 .unwrap();
14669 save_restore_target(
14670 &parent,
14671 storage.as_ref(),
14672 "grow",
14673 vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
14674 )
14675 .await;
14676
14677 assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
14678 assert_eq!(registry.count(), 3);
14679 assert!(registry.contains("c"));
14680 assert_eq!(spawner.spawned_count(), 3);
14681
14682 for id in ["a", "b"] {
14683 registry
14684 .get(id)
14685 .unwrap()
14686 .save_session("shrink")
14687 .await
14688 .unwrap();
14689 }
14690 save_restore_target(
14691 &parent,
14692 storage.as_ref(),
14693 "shrink",
14694 vec![restore_entry("a"), restore_entry("b")],
14695 )
14696 .await;
14697
14698 assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
14699 assert_eq!(registry.count(), 2);
14700 assert!(!registry.contains("c"));
14701 assert_eq!(spawner.spawned_count(), 2);
14702
14703 save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
14704
14705 assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
14706 assert_eq!(registry.count(), 0);
14707 assert_eq!(spawner.spawned_count(), 0);
14708 assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
14709 }
14710
14711 #[tokio::test]
14712 async fn storage_session_metadata_is_called_only_when_advertised() {
14713 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14714 storage.fail_metadata_save.store(true, Ordering::SeqCst);
14715 storage.fail_metadata_load.store(true, Ordering::SeqCst);
14716 let agent = runtime_storage_agent().with_storage(storage.clone());
14717
14718 agent.save_session("session").await.unwrap();
14719 assert!(agent.load_session("session").await.unwrap());
14720 assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
14721 assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
14722 }
14723
14724 #[cfg(feature = "sqlite")]
14725 #[tokio::test]
14726 async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
14727 let directory =
14728 std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
14729 let path = directory.join("sessions.sqlite");
14730 let path_string = path.to_string_lossy().into_owned();
14731 let storage = Arc::new(
14732 ai_agents_storage::SqliteStorage::new(&path_string)
14733 .await
14734 .unwrap(),
14735 );
14736 let agent = runtime_storage_agent().with_storage(storage.clone());
14737 agent.set_session_metadata(ai_agents_core::SessionMetadata {
14738 tags: vec!["initial".into()],
14739 ..Default::default()
14740 });
14741 agent.chat("persist this turn").await.unwrap();
14742 agent.save_session("session").await.unwrap();
14743
14744 agent.set_session_metadata(ai_agents_core::SessionMetadata {
14745 tags: vec!["updated".into()],
14746 ..Default::default()
14747 });
14748 agent.save_session("session").await.unwrap();
14749 assert!(
14750 agent
14751 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14752 tags: Some(vec!["initial".into()]),
14753 ..Default::default()
14754 })
14755 .await
14756 .unwrap()
14757 .is_empty()
14758 );
14759 assert_eq!(
14760 agent
14761 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14762 tags: Some(vec!["updated".into()]),
14763 ..Default::default()
14764 })
14765 .await
14766 .unwrap()
14767 .len(),
14768 1
14769 );
14770 drop(agent);
14771 storage.close().await;
14772 drop(storage);
14773
14774 let reopened_storage = Arc::new(
14775 ai_agents_storage::SqliteStorage::new(&path_string)
14776 .await
14777 .unwrap(),
14778 );
14779 let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14780 assert!(restored.load_session("session").await.unwrap());
14781 assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14782 assert_eq!(
14783 restored.current_session_id.read().as_deref(),
14784 Some("session")
14785 );
14786 assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14787 assert_eq!(
14788 restored
14789 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14790 tags: Some(vec!["updated".into()]),
14791 ..Default::default()
14792 })
14793 .await
14794 .unwrap()
14795 .len(),
14796 1
14797 );
14798
14799 drop(restored);
14800 reopened_storage.close().await;
14801 drop(reopened_storage);
14802 crate::remove_sqlite_test_directory(&directory)
14803 .await
14804 .unwrap();
14805 }
14806
14807 #[tokio::test]
14808 async fn storage_session_metadata_backend_failures_propagate() {
14809 let storage = Arc::new(RuntimeStorage::new([
14810 StorageCapability::Snapshot,
14811 StorageCapability::SessionMetadata,
14812 ]));
14813 let agent = runtime_storage_agent().with_storage(storage.clone());
14814
14815 agent.save_session("session").await.unwrap();
14816 storage
14817 .save("target", &agent.save_state().await.unwrap())
14818 .await
14819 .unwrap();
14820 storage.fail_metadata_load.store(true, Ordering::SeqCst);
14821 assert!(matches!(
14822 agent.load_session("target").await,
14823 Err(AgentError::Persistence(message)) if message == "metadata load failed"
14824 ));
14825 assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14826
14827 storage.fail_metadata_save.store(true, Ordering::SeqCst);
14828 assert!(matches!(
14829 agent.save_session("session").await,
14830 Err(AgentError::Persistence(message)) if message == "metadata save failed"
14831 ));
14832 }
14833
14834 struct ProviderFutureDropSignal {
14835 dropped: Arc<AtomicBool>,
14836 }
14837
14838 impl Drop for ProviderFutureDropSignal {
14839 fn drop(&mut self) {
14840 self.dropped.store(true, Ordering::SeqCst);
14841 }
14842 }
14843
14844 struct BufferedLockingProvider {
14845 lock: Arc<tokio::sync::Mutex<()>>,
14846 stream_started: Arc<tokio::sync::Notify>,
14847 stream_dropped: Arc<AtomicBool>,
14848 committed_after_drop: Arc<AtomicBool>,
14849 }
14850
14851 #[async_trait]
14852 impl LLMProvider for BufferedLockingProvider {
14853 async fn complete(
14854 &self,
14855 _messages: &[ChatMessage],
14856 _config: Option<&LLMConfig>,
14857 ) -> std::result::Result<LLMResponse, LLMError> {
14858 let _guard = self.lock.lock().await;
14859 self.committed_after_drop
14860 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14861 Ok(LLMResponse::new(
14862 "Committed technical response.",
14863 FinishReason::Stop,
14864 ))
14865 }
14866
14867 async fn complete_stream(
14868 &self,
14869 _messages: &[ChatMessage],
14870 _config: Option<&LLMConfig>,
14871 ) -> std::result::Result<
14872 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14873 LLMError,
14874 > {
14875 let _guard = self.lock.lock().await;
14876 let _drop_signal = ProviderFutureDropSignal {
14877 dropped: Arc::clone(&self.stream_dropped),
14878 };
14879 self.stream_started.notify_one();
14880 std::future::pending().await
14881 }
14882
14883 fn provider_name(&self) -> &str {
14884 "buffered-locking"
14885 }
14886
14887 fn supports(&self, _feature: LLMFeature) -> bool {
14888 false
14889 }
14890 }
14891
14892 struct PendingDropStream {
14893 dropped: Arc<AtomicBool>,
14894 dropped_notify: Arc<tokio::sync::Notify>,
14895 }
14896
14897 impl Stream for PendingDropStream {
14898 type Item = std::result::Result<LLMChunk, LLMError>;
14899
14900 fn poll_next(
14901 self: Pin<&mut Self>,
14902 _cx: &mut std::task::Context<'_>,
14903 ) -> std::task::Poll<Option<Self::Item>> {
14904 std::task::Poll::Pending
14905 }
14906 }
14907
14908 impl Drop for PendingDropStream {
14909 fn drop(&mut self) {
14910 self.dropped.store(true, Ordering::SeqCst);
14911 self.dropped_notify.notify_one();
14912 }
14913 }
14914
14915 struct EstablishedStreamProvider {
14916 stream_started: Arc<tokio::sync::Notify>,
14917 stream_dropped: Arc<AtomicBool>,
14918 stream_dropped_notify: Arc<tokio::sync::Notify>,
14919 committed_after_drop: Arc<AtomicBool>,
14920 }
14921
14922 #[async_trait]
14923 impl LLMProvider for EstablishedStreamProvider {
14924 async fn complete(
14925 &self,
14926 _messages: &[ChatMessage],
14927 _config: Option<&LLMConfig>,
14928 ) -> std::result::Result<LLMResponse, LLMError> {
14929 if !self.stream_dropped.load(Ordering::SeqCst) {
14930 self.stream_dropped_notify.notified().await;
14931 }
14932 self.committed_after_drop
14933 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14934 Ok(LLMResponse::new(
14935 "Committed technical response.",
14936 FinishReason::Stop,
14937 ))
14938 }
14939
14940 async fn complete_stream(
14941 &self,
14942 _messages: &[ChatMessage],
14943 _config: Option<&LLMConfig>,
14944 ) -> std::result::Result<
14945 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14946 LLMError,
14947 > {
14948 self.stream_started.notify_one();
14949 Ok(Box::new(PendingDropStream {
14950 dropped: Arc::clone(&self.stream_dropped),
14951 dropped_notify: Arc::clone(&self.stream_dropped_notify),
14952 }))
14953 }
14954
14955 fn provider_name(&self) -> &str {
14956 "established-stream"
14957 }
14958
14959 fn supports(&self, _feature: LLMFeature) -> bool {
14960 false
14961 }
14962 }
14963
14964 struct FirstCallLockingProvider {
14965 lock: Arc<tokio::sync::Mutex<()>>,
14966 first_started: Arc<tokio::sync::Notify>,
14967 first_dropped: Arc<AtomicBool>,
14968 committed_after_drop: Arc<AtomicBool>,
14969 calls: AtomicU64,
14970 }
14971
14972 #[async_trait]
14973 impl LLMProvider for FirstCallLockingProvider {
14974 async fn complete(
14975 &self,
14976 _messages: &[ChatMessage],
14977 _config: Option<&LLMConfig>,
14978 ) -> std::result::Result<LLMResponse, LLMError> {
14979 let _guard = self.lock.lock().await;
14980 let call = self.calls.fetch_add(1, Ordering::SeqCst);
14981 if call == 0 {
14982 let _drop_signal = ProviderFutureDropSignal {
14983 dropped: Arc::clone(&self.first_dropped),
14984 };
14985 self.first_started.notify_one();
14986 return std::future::pending().await;
14987 }
14988 self.committed_after_drop
14989 .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14990 Ok(LLMResponse::new(
14991 "Committed technical response.",
14992 FinishReason::Stop,
14993 ))
14994 }
14995
14996 async fn complete_stream(
14997 &self,
14998 _messages: &[ChatMessage],
14999 _config: Option<&LLMConfig>,
15000 ) -> std::result::Result<
15001 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
15002 LLMError,
15003 > {
15004 Err(LLMError::Other(
15005 "streaming is not used in this test".to_string(),
15006 ))
15007 }
15008
15009 fn provider_name(&self) -> &str {
15010 "first-call-locking"
15011 }
15012
15013 fn supports(&self, _feature: LLMFeature) -> bool {
15014 false
15015 }
15016 }
15017
15018 struct RoutingAfterProviderStart {
15019 provider_started: Arc<tokio::sync::Notify>,
15020 }
15021
15022 #[async_trait]
15023 impl LLMProvider for RoutingAfterProviderStart {
15024 async fn complete(
15025 &self,
15026 _messages: &[ChatMessage],
15027 _config: Option<&LLMConfig>,
15028 ) -> std::result::Result<LLMResponse, LLMError> {
15029 self.provider_started.notified().await;
15030 Ok(LLMResponse::new("1", FinishReason::Stop))
15031 }
15032
15033 async fn complete_stream(
15034 &self,
15035 _messages: &[ChatMessage],
15036 _config: Option<&LLMConfig>,
15037 ) -> std::result::Result<
15038 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
15039 LLMError,
15040 > {
15041 Err(LLMError::Other(
15042 "streaming is not used in this test".to_string(),
15043 ))
15044 }
15045
15046 fn provider_name(&self) -> &str {
15047 "routing-after-start"
15048 }
15049
15050 fn supports(&self, _feature: LLMFeature) -> bool {
15051 false
15052 }
15053 }
15054
15055 struct ResponseCountingHooks {
15057 responses: Arc<std::sync::atomic::AtomicUsize>,
15058 }
15059
15060 struct RootTurnProbeProvider {
15062 complete_entered: tokio::sync::mpsc::UnboundedSender<()>,
15063 }
15064
15065 struct ResponseChatHooks {
15067 target: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15068 invoked: AtomicBool,
15069 nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
15070 }
15071
15072 struct ConcurrentResponseHooks {
15074 registry: Weak<crate::spawner::AgentRegistry>,
15075 child_id: String,
15076 invoked: AtomicBool,
15077 nested_result: parking_lot::Mutex<Option<std::result::Result<String, String>>>,
15078 }
15079
15080 struct RetryDeadlineTool {
15082 calls: Arc<std::sync::atomic::AtomicUsize>,
15083 deadlines: Arc<parking_lot::Mutex<Vec<chrono::DateTime<chrono::Utc>>>>,
15084 remaining_ms: Arc<parking_lot::Mutex<Vec<i64>>>,
15085 }
15086
15087 struct ToolLifecycleRecordingHooks {
15089 events: parking_lot::Mutex<Vec<String>>,
15090 records: parking_lot::Mutex<Vec<ToolExecutionRecord>>,
15091 }
15092
15093 impl ToolLifecycleRecordingHooks {
15094 fn new() -> Self {
15096 Self {
15097 events: parking_lot::Mutex::new(Vec::new()),
15098 records: parking_lot::Mutex::new(Vec::new()),
15099 }
15100 }
15101
15102 fn events(&self) -> Vec<String> {
15104 self.events.lock().clone()
15105 }
15106
15107 fn records(&self) -> Vec<ToolExecutionRecord> {
15109 self.records.lock().clone()
15110 }
15111 }
15112
15113 struct ContextEchoTool;
15115
15116 #[async_trait]
15117 impl LLMProvider for RootTurnProbeProvider {
15118 async fn complete(
15119 &self,
15120 _messages: &[ChatMessage],
15121 _config: Option<&LLMConfig>,
15122 ) -> std::result::Result<LLMResponse, LLMError> {
15123 let _ = self.complete_entered.send(());
15124 Ok(LLMResponse::new("blocking complete", FinishReason::Stop))
15125 }
15126
15127 async fn complete_stream(
15128 &self,
15129 _messages: &[ChatMessage],
15130 _config: Option<&LLMConfig>,
15131 ) -> std::result::Result<
15132 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
15133 LLMError,
15134 > {
15135 Ok(Box::new(futures::stream::iter(vec![Ok(
15136 LLMChunk::final_chunk("stream complete", FinishReason::Stop, None),
15137 )])))
15138 }
15139
15140 fn provider_name(&self) -> &str {
15141 "root-turn-probe"
15142 }
15143
15144 fn supports(&self, feature: LLMFeature) -> bool {
15145 matches!(feature, LLMFeature::Streaming)
15146 }
15147 }
15148
15149 #[async_trait]
15150 impl ai_agents_core::Tool for ContextEchoTool {
15151 fn id(&self) -> &str {
15152 "context_echo"
15153 }
15154
15155 fn name(&self) -> &str {
15156 "Context Echo"
15157 }
15158
15159 fn description(&self) -> &str {
15160 "Returns selected execution context fields."
15161 }
15162
15163 fn input_schema(&self) -> Value {
15164 serde_json::json!({"type": "object"})
15165 }
15166
15167 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15168 ai_agents_core::ToolPolicyBindings {
15169 path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
15170 result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
15171 "max_results",
15172 ai_agents_core::ResultLimitKind::MaxResults,
15173 )],
15174 ..Default::default()
15175 }
15176 }
15177
15178 async fn execute(
15179 &self,
15180 _args: Value,
15181 ctx: ai_agents_core::ToolExecutionContext,
15182 ) -> ToolResult {
15183 ToolResult::ok(
15184 serde_json::json!({
15185 "requested_name": ctx.requested_name,
15186 "canonical_id": ctx.canonical_id,
15187 "display_name": ctx.display_name,
15188 "max_results": ctx.limits.max_results,
15189 "custom_config": ctx.custom_config,
15190 })
15191 .to_string(),
15192 )
15193 }
15194 }
15195
15196 #[async_trait]
15197 impl ai_agents_core::Tool for RetryDeadlineTool {
15198 fn id(&self) -> &str {
15199 "retry_deadline"
15200 }
15201
15202 fn name(&self) -> &str {
15203 "Retry Deadline"
15204 }
15205
15206 fn description(&self) -> &str {
15207 "Records one deadline per retry invocation."
15208 }
15209
15210 fn input_schema(&self) -> Value {
15211 serde_json::json!({"type": "object"})
15212 }
15213
15214 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15215 ai_agents_core::ToolSafetyMetadata::compute()
15216 }
15217
15218 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15219 let mut classification =
15220 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15221 classification.timeout_ms = Some(1_000);
15222 classification.safely_retryable = true;
15223 classification
15224 }
15225
15226 async fn execute(
15228 &self,
15229 _args: Value,
15230 ctx: ai_agents_core::ToolExecutionContext,
15231 ) -> ToolResult {
15232 let deadline = ctx
15233 .deadline
15234 .expect("each invocation must receive a deadline");
15235 self.remaining_ms.lock().push(
15236 deadline
15237 .signed_duration_since(chrono::Utc::now())
15238 .num_milliseconds(),
15239 );
15240 self.deadlines.lock().push(deadline);
15241 let call = self.calls.fetch_add(1, Ordering::SeqCst);
15242 if call == 0 {
15243 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
15244 ToolResult::error("retry")
15245 } else {
15246 ToolResult::ok("done")
15247 }
15248 }
15249 }
15250
15251 struct ClassifiedTimeoutTool {
15253 id: &'static str,
15254 calls: Arc<std::sync::atomic::AtomicUsize>,
15255 timeout_ms: u64,
15256 sleep_ms: u64,
15257 requires_approval: bool,
15258 remaining_ms: Arc<parking_lot::Mutex<Vec<i64>>>,
15259 }
15260
15261 struct ApprovalModifiedTimeoutTool {
15263 calls: Arc<std::sync::atomic::AtomicUsize>,
15264 }
15265
15266 struct SlowTool;
15268
15269 struct FlakyWriteTool {
15271 calls: Arc<std::sync::atomic::AtomicUsize>,
15272 }
15273
15274 struct LockedWriteTool {
15276 active: Arc<std::sync::atomic::AtomicUsize>,
15277 max_active: Arc<std::sync::atomic::AtomicUsize>,
15278 }
15279
15280 struct MultiResourceWriteTool {
15281 active: Arc<std::sync::atomic::AtomicUsize>,
15282 max_active: Arc<std::sync::atomic::AtomicUsize>,
15283 }
15284
15285 #[derive(Clone)]
15286 struct PathMutationGate {
15287 entered: Arc<AtomicBool>,
15288 entered_notify: Arc<tokio::sync::Notify>,
15289 release: Arc<tokio::sync::Notify>,
15290 }
15291
15292 impl PathMutationGate {
15293 fn new() -> Self {
15294 Self {
15295 entered: Arc::new(AtomicBool::new(false)),
15296 entered_notify: Arc::new(tokio::sync::Notify::new()),
15297 release: Arc::new(tokio::sync::Notify::new()),
15298 }
15299 }
15300
15301 async fn wait_until_entered(&self) {
15302 if !self.entered.load(Ordering::SeqCst) {
15303 self.entered_notify.notified().await;
15304 }
15305 }
15306
15307 fn release(&self) {
15308 self.release.notify_one();
15309 }
15310 }
15311
15312 struct BlockingPathMutationTool {
15313 id: &'static str,
15314 path_fields: Vec<ai_agents_core::PathPolicyBinding>,
15315 gate: PathMutationGate,
15316 }
15317
15318 struct NoBindingWriteTool {
15319 active: Arc<std::sync::atomic::AtomicUsize>,
15320 max_active: Arc<std::sync::atomic::AtomicUsize>,
15321 }
15322
15323 struct RecoveryTestTool {
15324 id: String,
15325 succeeds: bool,
15326 calls: Arc<std::sync::atomic::AtomicUsize>,
15327 max_output_chars: Option<usize>,
15328 }
15329
15330 struct BlockingApprovalHandler {
15331 entered: Arc<tokio::sync::Barrier>,
15332 release: Arc<tokio::sync::Notify>,
15333 result: ApprovalResult,
15334 }
15335
15336 struct CountingApprovalHandler {
15337 calls: Arc<std::sync::atomic::AtomicUsize>,
15338 }
15339
15340 struct DriftingFallbackProvider {
15342 refreshed: AtomicBool,
15343 primary_calls: Arc<std::sync::atomic::AtomicUsize>,
15344 secondary_calls: Arc<std::sync::atomic::AtomicUsize>,
15345 }
15346
15347 struct RefreshFallbackProviderHooks {
15349 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15350 lifecycle: Arc<ToolLifecycleRecordingHooks>,
15351 }
15352
15353 struct RuntimeWebFetchTransport {
15354 calls: Arc<std::sync::atomic::AtomicUsize>,
15355 }
15356
15357 struct RuntimeWebFetchResolver;
15358
15359 struct ReentrantToolHooks {
15360 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
15361 invoked: AtomicBool,
15362 nested_success: AtomicBool,
15363 }
15364
15365 #[async_trait]
15366 impl ai_agents_core::Tool for ClassifiedTimeoutTool {
15367 fn id(&self) -> &str {
15369 self.id
15370 }
15371
15372 fn name(&self) -> &str {
15374 "Classified Timeout"
15375 }
15376
15377 fn description(&self) -> &str {
15379 "Records and waits under one call-level timeout."
15380 }
15381
15382 fn input_schema(&self) -> Value {
15384 serde_json::json!({"type": "object"})
15385 }
15386
15387 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15389 let mut classification =
15390 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15391 classification.timeout_ms = Some(self.timeout_ms);
15392 classification.requires_approval = self.requires_approval;
15393 classification
15394 }
15395
15396 async fn execute(
15398 &self,
15399 _args: Value,
15400 ctx: ai_agents_core::ToolExecutionContext,
15401 ) -> ToolResult {
15402 self.calls.fetch_add(1, Ordering::SeqCst);
15403 let deadline = ctx
15404 .deadline
15405 .expect("each invocation must receive a deadline");
15406 self.remaining_ms.lock().push(
15407 deadline
15408 .signed_duration_since(chrono::Utc::now())
15409 .num_milliseconds(),
15410 );
15411 tokio::time::sleep(Duration::from_millis(self.sleep_ms)).await;
15412 ToolResult::ok("done")
15413 }
15414 }
15415
15416 #[async_trait]
15417 impl ai_agents_core::Tool for ApprovalModifiedTimeoutTool {
15418 fn id(&self) -> &str {
15420 "approval_modified_timeout"
15421 }
15422
15423 fn name(&self) -> &str {
15425 "Approval Modified Timeout"
15426 }
15427
15428 fn description(&self) -> &str {
15430 "Becomes invalid only after approval modifies its arguments."
15431 }
15432
15433 fn input_schema(&self) -> Value {
15435 serde_json::json!({"type": "object"})
15436 }
15437
15438 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15440 ai_agents_core::ToolPolicyBindings {
15441 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15442 ..Default::default()
15443 }
15444 }
15445
15446 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15448 ai_agents_core::ToolSafetyMetadata {
15449 read_only: false,
15450 concurrency_safe: false,
15451 operation: ai_agents_core::ToolOperationKind::Write,
15452 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15453 requires_network: false,
15454 destructive: false,
15455 open_world: false,
15456 host_dependent: false,
15457 requires_user_interaction: false,
15458 supports_cancellation: true,
15459 default_requires_approval: true,
15460 should_defer_schema: false,
15461 max_output_chars: Some(1024),
15462 max_result_size_chars: Some(1024),
15463 }
15464 }
15465
15466 fn classify_call(&self, args: &Value) -> ai_agents_core::ToolCallClassification {
15468 let mut classification =
15469 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15470 classification.timeout_ms = Some(if args["invalid_timeout"].as_bool() == Some(true) {
15471 u64::MAX
15472 } else {
15473 1_000
15474 });
15475 classification
15476 }
15477
15478 async fn execute(
15480 &self,
15481 _args: Value,
15482 _ctx: ai_agents_core::ToolExecutionContext,
15483 ) -> ToolResult {
15484 self.calls.fetch_add(1, Ordering::SeqCst);
15485 ToolResult::ok("unexpected")
15486 }
15487 }
15488
15489 #[async_trait]
15490 impl ai_agents_core::Tool for SlowTool {
15491 fn id(&self) -> &str {
15492 "slow"
15493 }
15494
15495 fn name(&self) -> &str {
15496 "Slow"
15497 }
15498
15499 fn description(&self) -> &str {
15500 "Waits until cancelled or timed out."
15501 }
15502
15503 fn input_schema(&self) -> Value {
15504 serde_json::json!({"type": "object"})
15505 }
15506
15507 async fn execute(
15508 &self,
15509 _args: Value,
15510 _ctx: ai_agents_core::ToolExecutionContext,
15511 ) -> ToolResult {
15512 tokio::time::sleep(std::time::Duration::from_secs(5)).await;
15513 ToolResult::ok("done")
15514 }
15515 }
15516
15517 #[async_trait]
15518 impl ai_agents_core::Tool for FlakyWriteTool {
15519 fn id(&self) -> &str {
15520 "flaky_write"
15521 }
15522
15523 fn name(&self) -> &str {
15524 "Flaky Write"
15525 }
15526
15527 fn description(&self) -> &str {
15528 "Fails on the first write attempt."
15529 }
15530
15531 fn input_schema(&self) -> Value {
15532 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15533 }
15534
15535 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15536 ai_agents_core::ToolPolicyBindings {
15537 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15538 ..Default::default()
15539 }
15540 }
15541
15542 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15543 ai_agents_core::ToolSafetyMetadata {
15544 read_only: false,
15545 concurrency_safe: false,
15546 operation: ai_agents_core::ToolOperationKind::Write,
15547 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15548 requires_network: false,
15549 destructive: false,
15550 open_world: false,
15551 host_dependent: false,
15552 requires_user_interaction: false,
15553 supports_cancellation: true,
15554 default_requires_approval: false,
15555 should_defer_schema: false,
15556 max_output_chars: Some(1024),
15557 max_result_size_chars: Some(1024),
15558 }
15559 }
15560
15561 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
15562 let mut classification =
15563 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
15564 classification.safely_retryable = false;
15565 classification
15566 }
15567
15568 async fn execute(
15569 &self,
15570 _args: Value,
15571 _ctx: ai_agents_core::ToolExecutionContext,
15572 ) -> ToolResult {
15573 let call = self.calls.fetch_add(1, Ordering::SeqCst);
15574 if call == 0 {
15575 ToolResult::error("first failure")
15576 } else {
15577 ToolResult::ok("second success")
15578 }
15579 }
15580 }
15581
15582 #[async_trait]
15583 impl ai_agents_core::Tool for LockedWriteTool {
15584 fn id(&self) -> &str {
15585 "locked_write"
15586 }
15587
15588 fn name(&self) -> &str {
15589 "Locked Write"
15590 }
15591
15592 fn description(&self) -> &str {
15593 "Tracks concurrent execution on one resource."
15594 }
15595
15596 fn input_schema(&self) -> Value {
15597 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
15598 }
15599
15600 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15601 ai_agents_core::ToolPolicyBindings {
15602 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15603 ..Default::default()
15604 }
15605 }
15606
15607 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15608 ai_agents_core::ToolSafetyMetadata {
15609 read_only: false,
15610 concurrency_safe: false,
15611 operation: ai_agents_core::ToolOperationKind::Write,
15612 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15613 requires_network: false,
15614 destructive: false,
15615 open_world: false,
15616 host_dependent: false,
15617 requires_user_interaction: false,
15618 supports_cancellation: true,
15619 default_requires_approval: false,
15620 should_defer_schema: false,
15621 max_output_chars: Some(1024),
15622 max_result_size_chars: Some(1024),
15623 }
15624 }
15625
15626 async fn execute(
15627 &self,
15628 _args: Value,
15629 _ctx: ai_agents_core::ToolExecutionContext,
15630 ) -> ToolResult {
15631 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15632 loop {
15633 let current_max = self.max_active.load(Ordering::SeqCst);
15634 if active <= current_max {
15635 break;
15636 }
15637 if self
15638 .max_active
15639 .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
15640 .is_ok()
15641 {
15642 break;
15643 }
15644 }
15645 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15646 self.active.fetch_sub(1, Ordering::SeqCst);
15647 ToolResult::ok("done")
15648 }
15649 }
15650
15651 #[async_trait]
15652 impl ai_agents_core::Tool for MultiResourceWriteTool {
15653 fn id(&self) -> &str {
15654 "multi_resource_write"
15655 }
15656
15657 fn name(&self) -> &str {
15658 "Multi Resource Write"
15659 }
15660
15661 fn description(&self) -> &str {
15662 "Tracks concurrent execution across source and destination resources."
15663 }
15664
15665 fn input_schema(&self) -> Value {
15666 serde_json::json!({"type": "object"})
15667 }
15668
15669 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15670 ai_agents_core::ToolPolicyBindings {
15671 path_fields: vec![
15672 ai_agents_core::PathPolicyBinding::read_write("source_path"),
15673 ai_agents_core::PathPolicyBinding::write("destination_path"),
15674 ],
15675 ..Default::default()
15676 }
15677 }
15678
15679 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15680 LockedWriteTool {
15681 active: Arc::clone(&self.active),
15682 max_active: Arc::clone(&self.max_active),
15683 }
15684 .safety_metadata()
15685 }
15686
15687 async fn execute(
15688 &self,
15689 _args: Value,
15690 _ctx: ai_agents_core::ToolExecutionContext,
15691 ) -> ToolResult {
15692 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15693 self.max_active.fetch_max(active, Ordering::SeqCst);
15694 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15695 self.active.fetch_sub(1, Ordering::SeqCst);
15696 ToolResult::ok("done")
15697 }
15698 }
15699
15700 #[async_trait]
15701 impl ai_agents_core::Tool for BlockingPathMutationTool {
15702 fn id(&self) -> &str {
15703 self.id
15704 }
15705
15706 fn name(&self) -> &str {
15707 self.id
15708 }
15709
15710 fn description(&self) -> &str {
15711 "Blocks a path mutation until the test releases it."
15712 }
15713
15714 fn input_schema(&self) -> Value {
15715 serde_json::json!({"type": "object"})
15716 }
15717
15718 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15719 ai_agents_core::ToolPolicyBindings {
15720 path_fields: self.path_fields.clone(),
15721 ..Default::default()
15722 }
15723 }
15724
15725 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15726 ai_agents_core::ToolSafetyMetadata {
15727 read_only: false,
15728 concurrency_safe: false,
15729 operation: ai_agents_core::ToolOperationKind::Write,
15730 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15731 requires_network: false,
15732 destructive: false,
15733 open_world: false,
15734 host_dependent: false,
15735 requires_user_interaction: false,
15736 supports_cancellation: true,
15737 default_requires_approval: false,
15738 should_defer_schema: false,
15739 max_output_chars: Some(1024),
15740 max_result_size_chars: Some(1024),
15741 }
15742 }
15743
15744 async fn execute(
15745 &self,
15746 _args: Value,
15747 _ctx: ai_agents_core::ToolExecutionContext,
15748 ) -> ToolResult {
15749 self.gate.entered.store(true, Ordering::SeqCst);
15750 self.gate.entered_notify.notify_one();
15751 self.gate.release.notified().await;
15752 ToolResult::ok("done")
15753 }
15754 }
15755
15756 #[async_trait]
15757 impl ai_agents_core::Tool for NoBindingWriteTool {
15758 fn id(&self) -> &str {
15759 "no_binding_write"
15760 }
15761
15762 fn name(&self) -> &str {
15763 "No Binding Write"
15764 }
15765
15766 fn description(&self) -> &str {
15767 "Tracks concurrent execution without resource bindings."
15768 }
15769
15770 fn input_schema(&self) -> Value {
15771 serde_json::json!({"type": "object"})
15772 }
15773
15774 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15775 LockedWriteTool {
15776 active: Arc::clone(&self.active),
15777 max_active: Arc::clone(&self.max_active),
15778 }
15779 .safety_metadata()
15780 }
15781
15782 async fn execute(
15783 &self,
15784 _args: Value,
15785 _ctx: ai_agents_core::ToolExecutionContext,
15786 ) -> ToolResult {
15787 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
15788 self.max_active.fetch_max(active, Ordering::SeqCst);
15789 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
15790 self.active.fetch_sub(1, Ordering::SeqCst);
15791 ToolResult::ok("done")
15792 }
15793 }
15794
15795 #[async_trait]
15796 impl ai_agents_core::Tool for RecoveryTestTool {
15797 fn id(&self) -> &str {
15798 &self.id
15799 }
15800
15801 fn name(&self) -> &str {
15802 &self.id
15803 }
15804
15805 fn description(&self) -> &str {
15806 "Records recovery execution and returns a configured result."
15807 }
15808
15809 fn input_schema(&self) -> Value {
15810 serde_json::json!({"type": "object"})
15811 }
15812
15813 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
15814 ai_agents_core::ToolPolicyBindings {
15815 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15816 ..Default::default()
15817 }
15818 }
15819
15820 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
15822 ai_agents_core::ToolSafetyMetadata {
15823 read_only: false,
15824 concurrency_safe: false,
15825 operation: ai_agents_core::ToolOperationKind::Write,
15826 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
15827 requires_network: false,
15828 destructive: false,
15829 open_world: false,
15830 host_dependent: false,
15831 requires_user_interaction: false,
15832 supports_cancellation: true,
15833 default_requires_approval: false,
15834 should_defer_schema: false,
15835 max_output_chars: Some(self.max_output_chars.unwrap_or(1024)),
15836 max_result_size_chars: Some(1024),
15837 }
15838 }
15839
15840 async fn execute(
15842 &self,
15843 _args: Value,
15844 _ctx: ai_agents_core::ToolExecutionContext,
15845 ) -> ToolResult {
15846 self.calls.fetch_add(1, Ordering::SeqCst);
15847 let mut result = if self.succeeds {
15848 ToolResult::ok(format!("{} succeeded", self.id))
15849 } else {
15850 ToolResult::error(format!("{} failed", self.id))
15851 };
15852 result.metadata = Some(HashMap::from([(
15853 "recovery_test_tool".to_string(),
15854 Value::String(self.id.clone()),
15855 )]));
15856 result
15857 }
15858 }
15859
15860 #[async_trait]
15861 impl WebFetchTransport for RuntimeWebFetchTransport {
15862 async fn send(
15864 &self,
15865 _request: WebFetchTransportRequest,
15866 ) -> std::result::Result<WebFetchTransportResponse, String> {
15867 Err("validated addresses are required".to_string())
15868 }
15869
15870 async fn send_validated(
15872 &self,
15873 _request: WebFetchTransportRequest,
15874 _addresses: &[std::net::SocketAddr],
15875 ) -> std::result::Result<WebFetchTransportResponse, String> {
15876 self.calls.fetch_add(1, Ordering::SeqCst);
15877 Ok(WebFetchTransportResponse {
15878 status: 200,
15879 content_type: Some("text/plain".to_string()),
15880 location: None,
15881 body: b"approved".to_vec(),
15882 })
15883 }
15884 }
15885
15886 #[async_trait]
15887 impl WebFetchResolver for RuntimeWebFetchResolver {
15888 async fn resolve(
15890 &self,
15891 _host: &str,
15892 _port: u16,
15893 ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
15894 Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
15895 93, 184, 216, 34,
15896 ))])
15897 }
15898 }
15899
15900 #[async_trait]
15901 impl ToolProvider for DriftingFallbackProvider {
15902 fn id(&self) -> &str {
15904 "drifting_fallback"
15905 }
15906
15907 fn name(&self) -> &str {
15909 "Drifting Fallback"
15910 }
15911
15912 fn provider_type(&self) -> ToolProviderType {
15914 ToolProviderType::Custom
15915 }
15916
15917 async fn list_tools(&self) -> Vec<ToolDescriptor> {
15919 let alias = ToolAliases::new().with_name("en", "fallback alias");
15920 let mut primary = ToolDescriptor::new(
15921 "primary",
15922 "Primary",
15923 "Fails before fallback.",
15924 serde_json::json!({"type": "object"}),
15925 );
15926 let mut secondary = ToolDescriptor::new(
15927 "secondary",
15928 "Secondary",
15929 "Must not execute after final canonical drift.",
15930 serde_json::json!({"type": "object"}),
15931 );
15932 if self.refreshed.load(Ordering::SeqCst) {
15933 primary = primary.with_aliases(alias);
15934 } else {
15935 secondary = secondary.with_aliases(alias);
15936 }
15937 vec![primary, secondary]
15938 }
15939
15940 async fn get_tool(&self, tool_id: &str) -> Option<Arc<dyn Tool>> {
15942 let calls = match tool_id {
15943 "primary" => Arc::clone(&self.primary_calls),
15944 "secondary" => Arc::clone(&self.secondary_calls),
15945 _ => return None,
15946 };
15947 Some(Arc::new(RecoveryTestTool {
15948 id: tool_id.to_string(),
15949 succeeds: false,
15950 calls,
15951 max_output_chars: None,
15952 }))
15953 }
15954
15955 fn supports_refresh(&self) -> bool {
15957 true
15958 }
15959
15960 async fn refresh(&self) -> std::result::Result<(), ToolProviderError> {
15962 self.refreshed.store(true, Ordering::SeqCst);
15963 Ok(())
15964 }
15965 }
15966
15967 #[async_trait]
15968 impl AgentHooks for RefreshFallbackProviderHooks {
15969 async fn on_tool_start(&self, tool: &str, args: &Value) {
15971 self.lifecycle.on_tool_start(tool, args).await;
15972 if tool != "secondary" {
15973 return;
15974 }
15975 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
15976 if let Some(agent) = agent {
15977 agent
15978 .tools
15979 .refresh_provider("drifting_fallback")
15980 .await
15981 .unwrap();
15982 }
15983 }
15984
15985 async fn on_tool_complete(&self, tool: &str, result: &ToolResult, duration_ms: u64) {
15986 self.lifecycle
15987 .on_tool_complete(tool, result, duration_ms)
15988 .await;
15989 }
15990
15991 async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
15992 self.lifecycle.on_tool_execution_record(record).await;
15993 }
15994
15995 async fn on_error(&self, error: &AgentError) {
15996 self.lifecycle.on_error(error).await;
15997 }
15998 }
15999
16000 #[async_trait]
16001 impl ApprovalHandler for BlockingApprovalHandler {
16002 async fn request_approval(
16003 &self,
16004 _request: ai_agents_hitl::ApprovalRequest,
16005 ) -> ApprovalResult {
16006 self.entered.wait().await;
16007 self.release.notified().await;
16008 self.result.clone()
16009 }
16010 }
16011
16012 #[async_trait]
16013 impl ApprovalHandler for CountingApprovalHandler {
16014 async fn request_approval(
16015 &self,
16016 _request: ai_agents_hitl::ApprovalRequest,
16017 ) -> ApprovalResult {
16018 self.calls.fetch_add(1, Ordering::SeqCst);
16019 ApprovalResult::Approved
16020 }
16021 }
16022
16023 #[async_trait]
16024 impl AgentHooks for ReentrantToolHooks {
16025 async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
16026 if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
16027 return;
16028 }
16029 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
16030 if let Some(agent) = agent {
16031 let result = agent
16032 .invoke_tool(ToolExecutionRequest::new(
16033 "nested-hook-call",
16034 "reentrant_write",
16035 serde_json::json!({"path": "./hook.txt"}),
16036 ToolCallSource::Manual,
16037 ))
16038 .await;
16039 self.nested_success
16040 .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
16041 }
16042 }
16043 }
16044
16045 #[async_trait]
16046 impl AgentHooks for ResponseCountingHooks {
16047 async fn on_response(&self, _response: &AgentResponse) {
16048 self.responses.fetch_add(1, Ordering::SeqCst);
16049 }
16050 }
16051
16052 #[async_trait]
16053 impl AgentHooks for ResponseChatHooks {
16054 async fn on_response(&self, _response: &AgentResponse) {
16056 if self.invoked.swap(true, Ordering::SeqCst) {
16057 return;
16058 }
16059 let target = self.target.lock().as_ref().and_then(Weak::upgrade);
16060 let result = if let Some(target) = target {
16061 target
16062 .chat("nested response hook call")
16063 .await
16064 .map(|response| response.content)
16065 .map_err(|error| error.to_string())
16066 } else {
16067 Err("response hook target is unavailable".to_string())
16068 };
16069 *self.nested_result.lock() = Some(result);
16070 }
16071 }
16072
16073 #[async_trait]
16074 impl AgentHooks for ConcurrentResponseHooks {
16075 async fn on_response(&self, _response: &AgentResponse) {
16077 if self.invoked.swap(true, Ordering::SeqCst) {
16078 return;
16079 }
16080 let Some(registry) = self.registry.upgrade() else {
16081 *self.nested_result.lock() =
16082 Some(Err("concurrent registry is unavailable".to_string()));
16083 return;
16084 };
16085 let agents = [ai_agents_state::ConcurrentAgentRef::Id(
16086 self.child_id.clone(),
16087 )];
16088 let aggregation = ai_agents_state::AggregationConfig {
16089 strategy: ai_agents_state::AggregationStrategy::FirstWins,
16090 synthesizer_llm: None,
16091 synthesizer_prompt: None,
16092 vote: None,
16093 };
16094 let result = crate::orchestration::concurrent(
16095 ®istry,
16096 "nested concurrent response hook call",
16097 &agents,
16098 &aggregation,
16099 None,
16100 Some(1),
16101 None,
16102 ai_agents_state::PartialFailureAction::Abort,
16103 None,
16104 )
16105 .await
16106 .map(|result| result.response.content)
16107 .map_err(|error| error.to_string());
16108 *self.nested_result.lock() = Some(result);
16109 }
16110 }
16111
16112 #[async_trait]
16113 impl AgentHooks for ToolLifecycleRecordingHooks {
16114 async fn on_tool_start(&self, tool: &str, _args: &Value) {
16115 self.events.lock().push(format!("start:{tool}"));
16116 }
16117
16118 async fn on_tool_complete(&self, tool: &str, result: &ToolResult, _duration_ms: u64) {
16119 self.events
16120 .lock()
16121 .push(format!("complete:{tool}:{}", result.success));
16122 }
16123
16124 async fn on_tool_execution_record(&self, record: &ToolExecutionRecord) {
16125 self.events.lock().push(format!(
16126 "record:{}:{}",
16127 record.canonical_id, record.executed
16128 ));
16129 self.records.lock().push(record.clone());
16130 }
16131
16132 async fn on_error(&self, _error: &AgentError) {
16134 self.events.lock().push("error".to_string());
16135 }
16136 }
16137
16138 struct ApprovalRecordingHooks {
16139 events: parking_lot::Mutex<Vec<String>>,
16140 }
16141
16142 impl ApprovalRecordingHooks {
16143 fn new() -> Self {
16144 Self {
16145 events: parking_lot::Mutex::new(Vec::new()),
16146 }
16147 }
16148
16149 fn events(&self) -> Vec<String> {
16150 self.events.lock().clone()
16151 }
16152 }
16153
16154 #[async_trait]
16155 impl AgentHooks for ApprovalRecordingHooks {
16156 async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
16157 self.events.lock().push(format!(
16158 "raw:{}:{}",
16159 request_id,
16160 approval_result_name(result)
16161 ));
16162 }
16163
16164 async fn on_approval_resolved(
16165 &self,
16166 request: &ai_agents_hitl::ApprovalRequest,
16167 raw_result: &ApprovalResult,
16168 outcome: &ApprovalResolvedOutcome,
16169 ) {
16170 self.events.lock().push(format!(
16171 "resolved:{}:{}:{}",
16172 request.id,
16173 approval_result_name(raw_result),
16174 approval_outcome_name(outcome)
16175 ));
16176 }
16177 }
16178
16179 fn approval_result_name(result: &ApprovalResult) -> &'static str {
16180 match result {
16181 ApprovalResult::Approved => "approved",
16182 ApprovalResult::Rejected { .. } => "rejected",
16183 ApprovalResult::Modified { .. } => "modified",
16184 ApprovalResult::Timeout => "timeout",
16185 }
16186 }
16187
16188 fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
16189 match outcome {
16190 ApprovalResolvedOutcome::Approved => "approved",
16191 ApprovalResolvedOutcome::Rejected { .. } => "rejected",
16192 ApprovalResolvedOutcome::Modified { .. } => "modified",
16193 ApprovalResolvedOutcome::Error { .. } => "error",
16194 }
16195 }
16196
16197 fn assert_correlated_approval_events(
16198 events: &[String],
16199 raw_status: &str,
16200 outcome_status: &str,
16201 ) {
16202 assert_eq!(events.len(), 2);
16203 let raw: Vec<_> = events[0].split(':').collect();
16204 let resolved: Vec<_> = events[1].split(':').collect();
16205 assert_eq!(raw[0], "raw");
16206 assert_eq!(resolved[0], "resolved");
16207 assert_eq!(raw[1], resolved[1]);
16208 assert_eq!(raw[2], raw_status);
16209 assert_eq!(resolved[2], raw_status);
16210 assert_eq!(resolved[3], outcome_status);
16211 }
16212
16213 fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
16214 let mut security = ToolSecurityConfig {
16215 enabled: true,
16216 fail_closed: true,
16217 ..Default::default()
16218 };
16219 let policy = ai_agents_tools::ToolPolicyConfig {
16220 enabled: policy_enabled,
16221 write_paths: vec![".".to_string()],
16222 require_confirmation: true,
16223 ..Default::default()
16224 };
16225 security.tools.insert("locked_write".to_string(), policy);
16226 security
16227 }
16228
16229 struct MutationTestWorkspace {
16230 root: std::path::PathBuf,
16231 }
16232
16233 impl MutationTestWorkspace {
16234 fn new() -> Self {
16235 let root = std::env::temp_dir().join(format!(
16236 "ai-agents-runtime-mutation-{}",
16237 uuid::Uuid::new_v4()
16238 ));
16239 std::fs::create_dir_all(&root).unwrap();
16240 Self { root }
16241 }
16242 }
16243
16244 impl Drop for MutationTestWorkspace {
16245 fn drop(&mut self) {
16246 let _ = std::fs::remove_dir_all(&self.root);
16247 }
16248 }
16249
16250 async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
16251 tokio::time::timeout(std::time::Duration::from_secs(2), async {
16252 loop {
16253 let strong_count = locks
16254 .read()
16255 .get("path-mutation:global")
16256 .map_or(0, |lock| lock.strong_count());
16257 if strong_count >= minimum {
16258 break;
16259 }
16260 tokio::task::yield_now().await;
16261 }
16262 })
16263 .await
16264 .expect("path mutation call did not reach the shared lock");
16265 }
16266
16267 async fn assert_path_mutation_pair_serialized(
16268 first_id: &'static str,
16269 first_fields: Vec<ai_agents_core::PathPolicyBinding>,
16270 first_args: Value,
16271 second_id: &'static str,
16272 second_fields: Vec<ai_agents_core::PathPolicyBinding>,
16273 second_args: Value,
16274 ) {
16275 let locks = new_tool_resource_locks();
16276 let first_gate = PathMutationGate::new();
16277 let second_gate = PathMutationGate::new();
16278 second_gate.release();
16279 let agent = Arc::new(
16280 AgentBuilder::new()
16281 .system_prompt("Test global path mutation locking.")
16282 .llm(Arc::new(mock_with_response("done")))
16283 .tool(Arc::new(BlockingPathMutationTool {
16284 id: first_id,
16285 path_fields: first_fields,
16286 gate: first_gate.clone(),
16287 }))
16288 .tool(Arc::new(BlockingPathMutationTool {
16289 id: second_id,
16290 path_fields: second_fields,
16291 gate: second_gate.clone(),
16292 }))
16293 .build()
16294 .unwrap()
16295 .with_shared_resource_locks(Arc::clone(&locks)),
16296 );
16297
16298 let first = {
16299 let agent = Arc::clone(&agent);
16300 tokio::spawn(async move {
16301 agent
16302 .invoke_tool(ToolExecutionRequest::new(
16303 format!("{}-first", first_id),
16304 first_id,
16305 first_args,
16306 ToolCallSource::Manual,
16307 ))
16308 .await
16309 .unwrap()
16310 })
16311 };
16312 first_gate.wait_until_entered().await;
16313
16314 let second = {
16315 let agent = Arc::clone(&agent);
16316 tokio::spawn(async move {
16317 agent
16318 .invoke_tool(ToolExecutionRequest::new(
16319 format!("{}-second", second_id),
16320 second_id,
16321 second_args,
16322 ToolCallSource::Manual,
16323 ))
16324 .await
16325 .unwrap()
16326 })
16327 };
16328 wait_for_resource_lock_strong_count(&locks, 2).await;
16329 assert!(!second_gate.entered.load(Ordering::SeqCst));
16330 assert!(!second.is_finished());
16331
16332 first_gate.release();
16333 let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
16334 tokio::join!(first, second)
16335 })
16336 .await
16337 .expect("serialized path mutation calls did not finish");
16338 assert!(first.unwrap().success);
16339 assert!(second.unwrap().success);
16340 assert!(second_gate.entered.load(Ordering::SeqCst));
16341 assert!(locks.read().is_empty());
16342 }
16343
16344 #[derive(Clone, Copy)]
16345 enum MutationDenial {
16346 Policy,
16347 Approval,
16348 }
16349
16350 fn mutation_denial_security_config(
16351 tool_id: &str,
16352 workspace: &std::path::Path,
16353 denial: MutationDenial,
16354 ) -> ToolSecurityConfig {
16355 let workspace = workspace.to_string_lossy().into_owned();
16356 let mut policy = ai_agents_tools::ToolPolicyConfig {
16357 read_paths: vec![workspace.clone()],
16358 write_paths: vec![workspace.clone()],
16359 ..Default::default()
16360 };
16361 match denial {
16362 MutationDenial::Policy => policy.blocked_paths = vec![workspace],
16363 MutationDenial::Approval => policy.require_confirmation = true,
16364 }
16365
16366 let mut security = ToolSecurityConfig {
16367 enabled: true,
16368 fail_closed: true,
16369 ..Default::default()
16370 };
16371 security.tools.insert(tool_id.to_string(), policy);
16372 security
16373 }
16374
16375 async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
16376 let workspace = MutationTestWorkspace::new();
16377 let tool_id = tool.id().to_string();
16378 let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
16379 let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
16380 std::fs::write(&preserved, "preserved").unwrap();
16381 let arguments = match tool_id.as_str() {
16382 "copy_path" | "move_path" => serde_json::json!({
16383 "source_path": preserved.to_string_lossy(),
16384 "destination_path": destination.to_string_lossy(),
16385 "dry_run": false
16386 }),
16387 "delete_path" => serde_json::json!({
16388 "path": preserved.to_string_lossy(),
16389 "recursive": false,
16390 "dry_run": false
16391 }),
16392 _ => panic!("unsupported mutation tool: {}", tool_id),
16393 };
16394 let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
16395 let builder = AgentBuilder::new()
16396 .system_prompt("Test mutation denial.")
16397 .llm(Arc::new(mock_with_response("done")))
16398 .tool(tool)
16399 .tool_security(ToolSecurityEngine::new(security));
16400 let builder = match denial {
16401 MutationDenial::Policy => builder,
16402 MutationDenial::Approval => builder
16403 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16404 .approval_handler(Arc::new(RejectAllHandler::new())),
16405 };
16406 let agent = builder.build().unwrap();
16407
16408 let record = agent
16409 .invoke_tool(ToolExecutionRequest::new(
16410 format!("{}-denied", tool_id),
16411 tool_id.clone(),
16412 arguments,
16413 ToolCallSource::Manual,
16414 ))
16415 .await
16416 .unwrap();
16417
16418 assert!(!record.executed, "{} must not be invoked", tool_id);
16419 assert!(!record.success);
16420 match denial {
16421 MutationDenial::Policy => {
16422 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
16423 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
16424 &approval.status,
16425 ToolApprovalStatus::NotRequired
16426 )));
16427 }
16428 MutationDenial::Approval => {
16429 assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
16430 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
16431 &approval.status,
16432 ToolApprovalStatus::Rejected
16433 )));
16434 }
16435 }
16436 assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
16437 assert!(!destination.exists());
16438 }
16439
16440 fn recovery_manager_with_fallbacks(
16441 fallbacks: impl IntoIterator<Item = (String, String)>,
16442 ) -> RecoveryManager {
16443 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
16444
16445 let per_tool = fallbacks
16446 .into_iter()
16447 .map(|(tool, fallback_tool)| {
16448 (
16449 tool,
16450 ToolRetryConfig {
16451 max_retries: 0,
16452 timeout_ms: Some(1_000),
16453 on_failure: ToolFailureAction::Fallback { fallback_tool },
16454 },
16455 )
16456 })
16457 .collect();
16458 RecoveryManager::new(ErrorRecoveryConfig {
16459 tools: ToolRecoveryConfig {
16460 per_tool,
16461 ..Default::default()
16462 },
16463 ..Default::default()
16464 })
16465 }
16466
16467 fn approval_check() -> HITLCheckResult {
16468 HITLCheckResult::required(
16469 ApprovalTrigger::tool("test", serde_json::json!({})),
16470 HashMap::new(),
16471 "Approve?",
16472 None,
16473 )
16474 }
16475
16476 fn agent_with_approval_result(
16477 raw_result: ApprovalResult,
16478 timeout_action: TimeoutAction,
16479 hooks: Arc<ApprovalRecordingHooks>,
16480 ) -> RuntimeAgent {
16481 use ai_agents_hitl::{CallbackHandler, HITLConfig};
16482
16483 let config = HITLConfig {
16484 on_timeout: timeout_action,
16485 ..Default::default()
16486 };
16487 let handler = CallbackHandler::new(move |_| raw_result.clone());
16488 AgentBuilder::new()
16489 .system_prompt("Test HITL hooks.")
16490 .llm(Arc::new(mock_with_response("done")))
16491 .build()
16492 .unwrap()
16493 .with_hooks(hooks)
16494 .with_hitl(HITLEngine::new(config), Arc::new(handler))
16495 }
16496
16497 #[tokio::test]
16498 async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
16499 let cases = vec![
16500 (ApprovalResult::Approved, "approved"),
16501 (
16502 ApprovalResult::Rejected {
16503 reason: Some("denied".to_string()),
16504 },
16505 "rejected",
16506 ),
16507 (
16508 ApprovalResult::Modified {
16509 changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
16510 },
16511 "modified",
16512 ),
16513 ];
16514
16515 for (raw_result, expected) in cases {
16516 let hooks = Arc::new(ApprovalRecordingHooks::new());
16517 let agent =
16518 agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
16519
16520 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16521
16522 assert_eq!(approval_result_name(&result), expected);
16523 assert_correlated_approval_events(&hooks.events(), expected, expected);
16524 }
16525 }
16526
16527 #[tokio::test]
16528 async fn approval_hooks_expose_timeout_policy_decisions() {
16529 for (timeout_action, expected) in [
16530 (TimeoutAction::Approve, "approved"),
16531 (TimeoutAction::Reject, "rejected"),
16532 ] {
16533 let hooks = Arc::new(ApprovalRecordingHooks::new());
16534 let agent =
16535 agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
16536
16537 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
16538
16539 assert_eq!(approval_result_name(&result), expected);
16540 assert_correlated_approval_events(&hooks.events(), "timeout", expected);
16541 }
16542 }
16543
16544 #[tokio::test]
16545 async fn timeout_error_fires_correlated_resolved_error_before_returning() {
16546 let hooks = Arc::new(ApprovalRecordingHooks::new());
16547 let agent = agent_with_approval_result(
16548 ApprovalResult::Timeout,
16549 TimeoutAction::Error,
16550 hooks.clone(),
16551 );
16552
16553 let error = agent
16554 .request_hitl_approval(approval_check())
16555 .await
16556 .unwrap_err();
16557
16558 assert!(error.to_string().contains("HITL approval timeout"));
16559 assert_correlated_approval_events(&hooks.events(), "timeout", "error");
16560 }
16561
16562 #[tokio::test]
16564 async fn test_integration_yaml_to_chat_basic() {
16565 let mock = mock_with_response("Hello! How can I help you?");
16566 let agent = AgentBuilder::new()
16567 .system_prompt("You are a test assistant.")
16568 .llm(Arc::new(mock))
16569 .build()
16570 .unwrap();
16571
16572 let response = agent.chat("Hi").await.unwrap();
16573 assert!(!response.content.is_empty());
16574 assert_eq!(response.content, "Hello! How can I help you?");
16575 }
16576
16577 #[tokio::test]
16578 async fn stream_events_emit_one_authoritative_final_without_legacy_done() {
16579 let agent = AgentBuilder::new()
16580 .system_prompt("You are a test assistant.")
16581 .llm(Arc::new(mock_with_response(
16582 "Hello from the final response.",
16583 )))
16584 .build()
16585 .unwrap();
16586
16587 let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16588 let mut final_responses = Vec::new();
16589 let mut legacy_done = 0;
16590 while let Some(event) = stream.next().await {
16591 match event {
16592 AgentStreamEvent::Chunk(StreamChunk::Done {}) => legacy_done += 1,
16593 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16594 panic!("unexpected stream error: {message}")
16595 }
16596 AgentStreamEvent::Final(response) => final_responses.push(response),
16597 AgentStreamEvent::Chunk(_) => {}
16598 }
16599 }
16600
16601 assert_eq!(legacy_done, 0);
16602 assert_eq!(final_responses.len(), 1);
16603 let response = final_responses.pop().unwrap();
16604 assert_eq!(response.content, "Hello from the final response.");
16605 assert!(
16606 response
16607 .metadata
16608 .as_ref()
16609 .is_some_and(|metadata| { metadata.contains_key("reasoning") })
16610 );
16611 }
16612
16613 #[tokio::test]
16614 async fn stream_final_content_includes_output_processing_after_provisional_chunks() {
16615 let yaml = r#"
16616name: ProcessedStreamAgent
16617system_prompt: "Answer directly."
16618process:
16619 output:
16620 - type: format
16621 config:
16622 template: "{{ response }} [finalized]"
16623streaming:
16624 enabled: true
16625"#;
16626 let agent = AgentBuilder::from_yaml(yaml)
16627 .unwrap()
16628 .llm(Arc::new(mock_with_response("provisional answer")))
16629 .auto_configure_features()
16630 .unwrap()
16631 .build()
16632 .unwrap();
16633
16634 let mut stream = agent.chat_stream_events("Hi").await.unwrap();
16635 let mut provisional = String::new();
16636 let mut final_content = None;
16637 while let Some(event) = stream.next().await {
16638 match event {
16639 AgentStreamEvent::Chunk(StreamChunk::Content { text }) => {
16640 provisional.push_str(&text)
16641 }
16642 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16643 panic!("unexpected stream error: {message}")
16644 }
16645 AgentStreamEvent::Final(response) => final_content = Some(response.content),
16646 AgentStreamEvent::Chunk(_) => {}
16647 }
16648 }
16649
16650 assert_eq!(provisional, "provisional answer");
16651 assert_eq!(
16652 final_content.as_deref(),
16653 Some("provisional answer [finalized]")
16654 );
16655 }
16656
16657 #[tokio::test]
16658 async fn stream_events_preserve_tool_progress_and_final_tool_calls() {
16659 let agent = AgentBuilder::new()
16660 .system_prompt("Use the echo tool once, then answer.")
16661 .llm(Arc::new(mock_with_responses(vec![
16662 r#"{"tool":"echo","arguments":{"message":"hello"}}"#,
16663 "Echo completed.",
16664 ])))
16665 .tool(Arc::new(ai_agents_tools::EchoTool::new()))
16666 .build()
16667 .unwrap();
16668
16669 let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
16670 let mut starts = 0;
16671 let mut results = 0;
16672 let mut ends = 0;
16673 let mut final_response = None;
16674 while let Some(event) = stream.next().await {
16675 match event {
16676 AgentStreamEvent::Chunk(StreamChunk::ToolCallStart { name, .. }) => {
16677 assert_eq!(name, "echo");
16678 starts += 1;
16679 }
16680 AgentStreamEvent::Chunk(StreamChunk::ToolResult { name, success, .. }) => {
16681 assert_eq!(name, "echo");
16682 assert!(success);
16683 results += 1;
16684 }
16685 AgentStreamEvent::Chunk(StreamChunk::ToolCallEnd { .. }) => ends += 1,
16686 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
16687 panic!("unexpected stream error: {message}")
16688 }
16689 AgentStreamEvent::Final(response) => final_response = Some(response),
16690 AgentStreamEvent::Chunk(_) => {}
16691 }
16692 }
16693
16694 assert_eq!((starts, results, ends), (1, 1, 1));
16695 let response = final_response.expect("tool stream must finalize");
16696 assert_eq!(response.content, "Echo completed.");
16697 assert_eq!(
16698 response.tool_calls.as_ref().map(|calls| calls
16699 .iter()
16700 .map(|call| call.name.as_str())
16701 .collect::<Vec<_>>()),
16702 Some(vec!["echo"])
16703 );
16704 }
16705
16706 #[tokio::test]
16707 async fn legacy_stream_still_emits_one_done_chunk() {
16708 let agent = AgentBuilder::new()
16709 .system_prompt("You are a test assistant.")
16710 .llm(Arc::new(mock_with_response(
16711 "Hello from the legacy stream.",
16712 )))
16713 .build()
16714 .unwrap();
16715
16716 let mut stream = agent.chat_stream("Hi").await.unwrap();
16717 let mut done = 0;
16718 while let Some(chunk) = stream.next().await {
16719 match chunk {
16720 StreamChunk::Done {} => done += 1,
16721 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
16722 _ => {}
16723 }
16724 }
16725
16726 assert_eq!(done, 1);
16727 }
16728
16729 #[tokio::test]
16731 async fn test_integration_multi_turn_conversation() {
16732 let mock = mock_with_responses(vec![
16733 "Hello! I'm your assistant.",
16734 "The weather is sunny today.",
16735 "Goodbye!",
16736 ]);
16737 let agent = AgentBuilder::new()
16738 .system_prompt("You are helpful.")
16739 .llm(Arc::new(mock))
16740 .build()
16741 .unwrap();
16742
16743 let r1 = agent.chat("Hi").await.unwrap();
16744 assert_eq!(r1.content, "Hello! I'm your assistant.");
16745
16746 let r2 = agent.chat("What's the weather?").await.unwrap();
16747 assert_eq!(r2.content, "The weather is sunny today.");
16748
16749 let r3 = agent.chat("Bye").await.unwrap();
16750 assert_eq!(r3.content, "Goodbye!");
16751
16752 let messages = agent.memory.get_messages(None).await.unwrap();
16754 assert_eq!(messages.len(), 6);
16756 }
16757
16758 #[test]
16759 fn later_approval_preserves_modified_evidence() {
16760 let arguments = serde_json::json!({"dry_run": true});
16761 let mut record = Some(ToolApprovalRecord {
16762 status: ToolApprovalStatus::Modified,
16763 reason: None,
16764 modified_arguments: Some(arguments.clone()),
16765 });
16766
16767 merge_approved_record(&mut record);
16768
16769 let record = record.unwrap();
16770 assert!(matches!(record.status, ToolApprovalStatus::Modified));
16771 assert_eq!(record.modified_arguments, Some(arguments));
16772 }
16773
16774 #[test]
16775 fn approval_binding_rejects_replaced_tool_implementation() {
16776 let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16777 let same_tool = Arc::clone(&reviewed_tool);
16778 let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
16779 let arguments = serde_json::json!({"path": "."});
16780 let versions = ToolDecisionVersions {
16781 policy: 2,
16782 registry: 3,
16783 runtime_control: 4,
16784 state: Some(5),
16785 };
16786 let binding = ToolApprovalBinding {
16787 canonical_id: "context_echo".to_string(),
16788 arguments: arguments.clone(),
16789 confirmation_required: true,
16790 policy_version: versions.policy,
16791 runtime_control_version: versions.runtime_control,
16792 state_generation: versions.state,
16793 reviewed_tool,
16794 };
16795
16796 assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
16797 assert!(binding.is_stale(
16798 "context_echo",
16799 &arguments,
16800 true,
16801 versions,
16802 &replacement_tool,
16803 ));
16804 }
16805
16806 #[tokio::test]
16807 async fn approved_mutation_to_dry_run_remains_executable() {
16808 use ai_agents_hitl::CallbackHandler;
16809
16810 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16811 changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
16812 });
16813 let agent = AgentBuilder::new()
16814 .system_prompt("Test safer approval modifications.")
16815 .llm(Arc::new(mock_with_response("done")))
16816 .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
16817 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16818 .approval_handler(Arc::new(handler))
16819 .build()
16820 .unwrap();
16821
16822 let record = agent
16823 .invoke_tool(ToolExecutionRequest::new(
16824 "approved-dry-run",
16825 "file_write",
16826 serde_json::json!({
16827 "path": "./approval-dry-run.txt",
16828 "content": "not written"
16829 }),
16830 ToolCallSource::Manual,
16831 ))
16832 .await
16833 .unwrap();
16834
16835 assert!(record.executed);
16836 assert!(record.success);
16837 assert_eq!(record.executed_arguments["dry_run"], true);
16838 assert!(matches!(
16839 record.approval.as_ref().map(|approval| &approval.status),
16840 Some(ToolApprovalStatus::Modified)
16841 ));
16842 let output: Value = serde_json::from_str(&record.output).unwrap();
16843 assert_eq!(output["mutation_performed"], false);
16844 }
16845
16846 #[tokio::test]
16848 async fn shared_executor_approval_reaches_web_fetch_transport() {
16849 use ai_agents_hitl::{CallbackHandler, HITLConfig};
16850 use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
16851
16852 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16853 let tool = WebFetchTool::with_transport_and_resolver(
16854 Arc::new(RuntimeWebFetchTransport {
16855 calls: Arc::clone(&calls),
16856 }),
16857 Arc::new(RuntimeWebFetchResolver),
16858 );
16859 let mut security = ToolSecurityConfig {
16860 enabled: true,
16861 fail_closed: true,
16862 ..Default::default()
16863 };
16864 security.tools.insert(
16865 "web_fetch".to_string(),
16866 ToolPolicyConfig {
16867 domains: DomainPolicyConfig {
16868 requires_approval: vec!["approval.test".to_string()],
16869 ..Default::default()
16870 },
16871 allowed_schemes: vec!["https".to_string()],
16872 allowed_ports: vec![443],
16873 ..Default::default()
16874 },
16875 );
16876 let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
16877 let agent = AgentBuilder::new()
16878 .system_prompt("Test approved web fetch execution.")
16879 .llm(Arc::new(mock_with_response("done")))
16880 .tool(Arc::new(tool))
16881 .tool_security(ToolSecurityEngine::new(security))
16882 .build()
16883 .unwrap()
16884 .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
16885
16886 let record = agent
16887 .invoke_tool(ToolExecutionRequest::new(
16888 "approved-web-fetch",
16889 "web_fetch",
16890 serde_json::json!({
16891 "url": "https://approval.test/page",
16892 "cache_ttl_seconds": 0
16893 }),
16894 ToolCallSource::Manual,
16895 ))
16896 .await
16897 .unwrap();
16898
16899 assert!(record.success);
16900 assert!(
16901 record
16902 .approval
16903 .as_ref()
16904 .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
16905 );
16906 assert_eq!(calls.load(Ordering::SeqCst), 1);
16907 }
16908
16909 #[tokio::test]
16910 async fn context_preserves_requested_and_canonical_identity() {
16911 let mock = mock_with_response("hello");
16912 let mut tools = ai_agents_tools::ToolRegistry::new();
16913 tools.register(Arc::new(ContextEchoTool)).unwrap();
16914
16915 let mut security = ToolSecurityConfig {
16916 enabled: true,
16917 fail_closed: true,
16918 ..Default::default()
16919 };
16920 let mut policy = ai_agents_tools::ToolPolicyConfig {
16921 read_paths: vec![".".to_string()],
16922 max_results: Some(7),
16923 ..Default::default()
16924 };
16925 policy
16926 .config
16927 .insert("backend".to_string(), serde_json::json!("memory"));
16928 security.tools.insert("context_echo".to_string(), policy);
16929
16930 let agent = AgentBuilder::new()
16931 .system_prompt("You are helpful.")
16932 .llm(Arc::new(mock))
16933 .tools(tools)
16934 .tool_security(ToolSecurityEngine::new(security))
16935 .build()
16936 .unwrap();
16937
16938 let record = agent
16939 .invoke_tool(ToolExecutionRequest::new(
16940 "ctx-call",
16941 "Context Echo",
16942 serde_json::json!({"path": ".", "max_results": 99}),
16943 ToolCallSource::Manual,
16944 ))
16945 .await
16946 .unwrap();
16947
16948 assert!(record.success);
16949 assert!(matches!(&record.source, ToolCallSource::Manual));
16950 assert_eq!(record.requested_name, "Context Echo");
16951 assert_eq!(record.canonical_id, "context_echo");
16952 assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
16953 assert_eq!(record.executed_arguments["max_results"], 7);
16954 let output: Value = serde_json::from_str(&record.output).unwrap();
16955 assert_eq!(output["requested_name"], "Context Echo");
16956 assert_eq!(output["canonical_id"], "context_echo");
16957 assert_eq!(output["max_results"], 7);
16958 assert_eq!(output["custom_config"]["backend"], "memory");
16959 assert!(record.metadata.contains_key("effective_limits"));
16960 assert!(record.metadata.contains_key("policy_snapshot"));
16961 }
16962
16963 #[tokio::test]
16964 async fn test_runtime_control_cancels_active_tool_call() {
16965 let mock = mock_with_response("hello");
16966 let agent = Arc::new(
16967 AgentBuilder::new()
16968 .system_prompt("You are helpful.")
16969 .llm(Arc::new(mock))
16970 .tool(Arc::new(SlowTool))
16971 .build()
16972 .unwrap(),
16973 );
16974 let control = agent.runtime_control();
16975 let running_agent = Arc::clone(&agent);
16976 let handle = tokio::spawn(async move {
16977 running_agent
16978 .invoke_tool(ToolExecutionRequest::new(
16979 "slow-call",
16980 "slow",
16981 serde_json::json!({}),
16982 ToolCallSource::Manual,
16983 ))
16984 .await
16985 .unwrap()
16986 });
16987
16988 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
16989 control.cancel_all();
16990 let record = handle.await.unwrap();
16991
16992 assert!(record.executed);
16993 assert!(record.cancelled);
16994 assert!(!record.success);
16995 assert!(record.cancellation_reason.is_some());
16996 }
16997
16998 #[tokio::test]
17000 async fn cancelled_tool_does_not_enter_fallback() {
17001 let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17002 let agent = Arc::new(
17003 AgentBuilder::new()
17004 .system_prompt("Test cancellation before fallback.")
17005 .llm(Arc::new(mock_with_response("done")))
17006 .tool(Arc::new(SlowTool))
17007 .tool(Arc::new(RecoveryTestTool {
17008 id: "fallback".to_string(),
17009 succeeds: true,
17010 calls: Arc::clone(&fallback_calls),
17011 max_output_chars: None,
17012 }))
17013 .recovery_manager(recovery_manager_with_fallbacks([(
17014 "slow".to_string(),
17015 "fallback".to_string(),
17016 )]))
17017 .build()
17018 .unwrap(),
17019 );
17020 let control = agent.runtime_control();
17021 let running_agent = Arc::clone(&agent);
17022 let handle = tokio::spawn(async move {
17023 running_agent
17024 .invoke_tool(ToolExecutionRequest::new(
17025 "cancelled-fallback-call",
17026 "slow",
17027 serde_json::json!({}),
17028 ToolCallSource::Manual,
17029 ))
17030 .await
17031 .unwrap()
17032 });
17033
17034 tokio::time::sleep(Duration::from_millis(100)).await;
17035 control.cancel_all();
17036 let record = handle.await.unwrap();
17037
17038 assert!(record.executed);
17039 assert!(record.cancelled);
17040 assert!(!record.success);
17041 assert_eq!(record.canonical_id, "slow");
17042 assert_eq!(fallback_calls.load(Ordering::SeqCst), 0);
17043 assert_eq!(agent.tool_call_history().len(), 1);
17044 }
17045
17046 #[tokio::test]
17047 async fn non_idempotent_tool_calls_are_not_retried() {
17048 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17049
17050 let mock = mock_with_response("hello");
17051 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17052 let agent = AgentBuilder::new()
17053 .system_prompt("You are helpful.")
17054 .llm(Arc::new(mock))
17055 .tool(Arc::new(FlakyWriteTool {
17056 calls: Arc::clone(&calls),
17057 }))
17058 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17059 tools: ToolRecoveryConfig {
17060 default: ToolRetryConfig {
17061 max_retries: 2,
17062 ..Default::default()
17063 },
17064 ..Default::default()
17065 },
17066 ..Default::default()
17067 }))
17068 .build()
17069 .unwrap();
17070
17071 let record = agent
17072 .invoke_tool(ToolExecutionRequest::new(
17073 "flaky-call",
17074 "flaky_write",
17075 serde_json::json!({"path": "./tmp.txt"}),
17076 ToolCallSource::Manual,
17077 ))
17078 .await
17079 .unwrap();
17080
17081 assert!(!record.success);
17082 assert_eq!(calls.load(Ordering::SeqCst), 1);
17083 }
17084
17085 #[tokio::test]
17086 async fn safely_retryable_tool_receives_a_fresh_deadline_per_attempt() {
17087 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17088
17089 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17090 let deadlines = Arc::new(parking_lot::Mutex::new(Vec::new()));
17091 let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17092 let agent = AgentBuilder::new()
17093 .system_prompt("Test retry deadlines.")
17094 .llm(Arc::new(mock_with_response("done")))
17095 .tool(Arc::new(RetryDeadlineTool {
17096 calls: Arc::clone(&calls),
17097 deadlines: Arc::clone(&deadlines),
17098 remaining_ms: Arc::clone(&remaining_ms),
17099 }))
17100 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17101 tools: ToolRecoveryConfig {
17102 per_tool: HashMap::from([(
17103 "retry_deadline".to_string(),
17104 ToolRetryConfig {
17105 max_retries: 1,
17106 ..Default::default()
17107 },
17108 )]),
17109 ..Default::default()
17110 },
17111 ..Default::default()
17112 }))
17113 .build()
17114 .unwrap();
17115
17116 let record = agent
17117 .invoke_tool(ToolExecutionRequest::new(
17118 "retry-deadline-call",
17119 "retry_deadline",
17120 serde_json::json!({}),
17121 ToolCallSource::Manual,
17122 ))
17123 .await
17124 .unwrap();
17125
17126 assert!(record.executed);
17127 assert!(record.success);
17128 assert_eq!(calls.load(Ordering::SeqCst), 2);
17129 let deadlines = deadlines.lock();
17130 assert_eq!(deadlines.len(), 2);
17131 assert!(
17132 deadlines[1] > deadlines[0],
17133 "retry inherited the first invocation deadline"
17134 );
17135 let remaining_ms = remaining_ms.lock();
17136 assert_eq!(remaining_ms.len(), 2);
17137 assert!(
17138 remaining_ms
17139 .iter()
17140 .all(|remaining| (800..=1_000).contains(remaining))
17141 );
17142 }
17143
17144 #[tokio::test]
17146 async fn call_classification_timeout_controls_deadline_and_timer() {
17147 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17148 let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17149 let agent = AgentBuilder::new()
17150 .system_prompt("Test call-level timeout.")
17151 .llm(Arc::new(mock_with_response("done")))
17152 .tool(Arc::new(ClassifiedTimeoutTool {
17153 id: "classified_timeout",
17154 calls: Arc::clone(&calls),
17155 timeout_ms: 100,
17156 sleep_ms: 150,
17157 requires_approval: false,
17158 remaining_ms: Arc::clone(&remaining_ms),
17159 }))
17160 .build()
17161 .unwrap();
17162
17163 let started = Instant::now();
17164 let record = agent
17165 .invoke_tool(ToolExecutionRequest::new(
17166 "classified-timeout-call",
17167 "classified_timeout",
17168 serde_json::json!({}),
17169 ToolCallSource::Manual,
17170 ))
17171 .await
17172 .unwrap();
17173
17174 assert!(record.executed);
17175 assert!(record.timed_out);
17176 assert!(!record.success);
17177 assert_eq!(calls.load(Ordering::SeqCst), 1);
17178 assert!(started.elapsed() < Duration::from_secs(1));
17179 let remaining_ms = remaining_ms.lock();
17180 assert_eq!(remaining_ms.len(), 1);
17181 assert!((1..=100).contains(&remaining_ms[0]));
17182 }
17183
17184 #[tokio::test]
17186 async fn recovery_timeout_only_lowers_call_and_policy_timeouts() {
17187 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17188
17189 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17190 let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17191 let agent = AgentBuilder::new()
17192 .system_prompt("Test recovery timeout.")
17193 .llm(Arc::new(mock_with_response("done")))
17194 .tool(Arc::new(ClassifiedTimeoutTool {
17195 id: "recovery_timeout",
17196 calls: Arc::clone(&calls),
17197 timeout_ms: 1_000,
17198 sleep_ms: 150,
17199 requires_approval: false,
17200 remaining_ms: Arc::clone(&remaining_ms),
17201 }))
17202 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17203 tools: ToolRecoveryConfig {
17204 per_tool: HashMap::from([(
17205 "recovery_timeout".to_string(),
17206 ToolRetryConfig {
17207 timeout_ms: Some(100),
17208 ..Default::default()
17209 },
17210 )]),
17211 ..Default::default()
17212 },
17213 ..Default::default()
17214 }))
17215 .build()
17216 .unwrap();
17217
17218 let started = Instant::now();
17219 let record = agent
17220 .invoke_tool(ToolExecutionRequest::new(
17221 "recovery-timeout-call",
17222 "recovery_timeout",
17223 serde_json::json!({}),
17224 ToolCallSource::Manual,
17225 ))
17226 .await
17227 .unwrap();
17228
17229 assert!(record.executed);
17230 assert!(record.timed_out);
17231 assert!(!record.success);
17232 assert_eq!(calls.load(Ordering::SeqCst), 1);
17233 assert!(started.elapsed() < Duration::from_secs(1));
17234 assert_eq!(record.metadata["effective_limits"]["timeout_ms"], 100);
17235 let remaining_ms = remaining_ms.lock();
17236 assert_eq!(remaining_ms.len(), 1);
17237 assert!((1..=100).contains(&remaining_ms[0]));
17238 }
17239
17240 #[tokio::test]
17242 async fn recovery_default_timeout_controls_deadline_and_timer() {
17243 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
17244
17245 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17246 let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17247 let agent = AgentBuilder::new()
17248 .system_prompt("Test default recovery timeout.")
17249 .llm(Arc::new(mock_with_response("done")))
17250 .tool(Arc::new(ClassifiedTimeoutTool {
17251 id: "default_recovery_timeout",
17252 calls: Arc::clone(&calls),
17253 timeout_ms: 1_000,
17254 sleep_ms: 150,
17255 requires_approval: false,
17256 remaining_ms: Arc::clone(&remaining_ms),
17257 }))
17258 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
17259 tools: ToolRecoveryConfig {
17260 default: ToolRetryConfig {
17261 timeout_ms: Some(100),
17262 ..Default::default()
17263 },
17264 ..Default::default()
17265 },
17266 ..Default::default()
17267 }))
17268 .build()
17269 .unwrap();
17270
17271 let started = Instant::now();
17272 let record = agent
17273 .invoke_tool(ToolExecutionRequest::new(
17274 "default-recovery-timeout-call",
17275 "default_recovery_timeout",
17276 serde_json::json!({}),
17277 ToolCallSource::Manual,
17278 ))
17279 .await
17280 .unwrap();
17281
17282 assert!(record.executed);
17283 assert!(record.timed_out);
17284 assert!(!record.success);
17285 assert_eq!(calls.load(Ordering::SeqCst), 1);
17286 assert!(started.elapsed() < Duration::from_secs(1));
17287 assert_eq!(record.metadata["effective_limits"]["timeout_ms"], 100);
17288 let remaining_ms = remaining_ms.lock();
17289 assert_eq!(remaining_ms.len(), 1);
17290 assert!((1..=100).contains(&remaining_ms[0]));
17291 }
17292
17293 #[test]
17295 fn recovery_timeout_cannot_widen_security_baseline() {
17296 let security_engine = ToolSecurityEngine::new(ToolSecurityConfig {
17297 default_timeout_ms: 100,
17298 ..Default::default()
17299 });
17300 let safety = ToolSafetyMetadata::compute();
17301 let mut classification = ToolCallClassification::from_metadata(&safety);
17302 classification.timeout_ms = Some(500);
17303
17304 let (limits, timeout) = RuntimeAgent::effective_tool_limits(
17305 &security_engine,
17306 "recovery_cannot_widen",
17307 &safety,
17308 &classification,
17309 Some(1_000),
17310 )
17311 .unwrap();
17312
17313 assert_eq!(limits.timeout_ms, Some(100));
17314 assert_eq!(timeout.timer, Duration::from_millis(100));
17315 }
17316
17317 #[tokio::test]
17319 async fn invalid_call_timeout_stops_before_approval_or_tool_invocation() {
17320 let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17321 let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17322 let remaining_ms = Arc::new(parking_lot::Mutex::new(Vec::new()));
17323 let mut security = ToolSecurityConfig {
17324 enabled: true,
17325 ..Default::default()
17326 };
17327 security.tools.insert(
17328 "invalid_call_timeout".to_string(),
17329 ai_agents_tools::ToolPolicyConfig {
17330 require_confirmation: true,
17331 ..Default::default()
17332 },
17333 );
17334 let agent = AgentBuilder::new()
17335 .system_prompt("Test invalid call timeout.")
17336 .llm(Arc::new(mock_with_response("done")))
17337 .tool(Arc::new(ClassifiedTimeoutTool {
17338 id: "invalid_call_timeout",
17339 calls: Arc::clone(&tool_calls),
17340 timeout_ms: u64::MAX,
17341 sleep_ms: 0,
17342 requires_approval: false,
17343 remaining_ms,
17344 }))
17345 .tool_security(ToolSecurityEngine::new(security))
17346 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17347 .approval_handler(Arc::new(CountingApprovalHandler {
17348 calls: Arc::clone(&approval_calls),
17349 }))
17350 .build()
17351 .unwrap();
17352
17353 let error = agent
17354 .invoke_tool(ToolExecutionRequest::new(
17355 "invalid-call-timeout",
17356 "invalid_call_timeout",
17357 serde_json::json!({}),
17358 ToolCallSource::Manual,
17359 ))
17360 .await
17361 .unwrap_err();
17362
17363 assert!(error.to_string().contains(
17364 "effective tool timeout_ms must be no greater than 3153600000000000 milliseconds"
17365 ));
17366 assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
17367 assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
17368 }
17369
17370 #[tokio::test]
17372 async fn invalid_modified_call_timeout_stops_before_lock_or_invocation() {
17373 use ai_agents_hitl::CallbackHandler;
17374
17375 let blocker_gate = PathMutationGate::new();
17376 let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17377 let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
17378 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
17379 changes: HashMap::from([("invalid_timeout".to_string(), Value::Bool(true))]),
17380 });
17381 let agent = Arc::new(
17382 AgentBuilder::new()
17383 .system_prompt("Test final call timeout validation.")
17384 .llm(Arc::new(mock_with_response("done")))
17385 .tool(Arc::new(BlockingPathMutationTool {
17386 id: "timeout_lock_blocker",
17387 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17388 gate: blocker_gate.clone(),
17389 }))
17390 .tool(Arc::new(ApprovalModifiedTimeoutTool {
17391 calls: Arc::clone(&tool_calls),
17392 }))
17393 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17394 .approval_handler(Arc::new(handler))
17395 .hooks(hooks.clone())
17396 .build()
17397 .unwrap(),
17398 );
17399 let blocking_agent = Arc::clone(&agent);
17400 let blocker = tokio::spawn(async move {
17401 blocking_agent
17402 .invoke_tool(ToolExecutionRequest::new(
17403 "timeout-lock-blocker",
17404 "timeout_lock_blocker",
17405 serde_json::json!({"path": "./shared-timeout.txt"}),
17406 ToolCallSource::Manual,
17407 ))
17408 .await
17409 .unwrap()
17410 });
17411 blocker_gate.wait_until_entered().await;
17412
17413 let record = tokio::time::timeout(
17414 Duration::from_millis(500),
17415 agent.invoke_tool(ToolExecutionRequest::new(
17416 "invalid-modified-timeout",
17417 "approval_modified_timeout",
17418 serde_json::json!({
17419 "path": "./shared-timeout.txt",
17420 "invalid_timeout": false
17421 }),
17422 ToolCallSource::Manual,
17423 )),
17424 )
17425 .await
17426 .expect("final timeout validation must not wait for the held path lock")
17427 .unwrap();
17428
17429 blocker_gate.release();
17430 assert!(blocker.await.unwrap().success);
17431 assert!(!record.executed);
17432 assert!(!record.success);
17433 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
17434 assert!(record.output.contains(
17435 "effective tool timeout_ms must be no greater than 3153600000000000 milliseconds"
17436 ));
17437 assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
17438 let invalid_request_events = hooks
17439 .events()
17440 .into_iter()
17441 .filter(|event| event.contains("approval_modified_timeout") || event == "error")
17442 .collect::<Vec<_>>();
17443 assert_eq!(
17444 invalid_request_events,
17445 vec![
17446 "start:approval_modified_timeout",
17447 "complete:approval_modified_timeout:false",
17448 "record:approval_modified_timeout:false",
17449 "error"
17450 ]
17451 );
17452 }
17453
17454 #[tokio::test]
17455 async fn side_effecting_tools_are_serialized_per_resource() {
17456 let mock = mock_with_response("hello");
17457 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17458 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17459 let agent = Arc::new(
17460 AgentBuilder::new()
17461 .system_prompt("You are helpful.")
17462 .llm(Arc::new(mock))
17463 .tool(Arc::new(LockedWriteTool {
17464 active: Arc::clone(&active),
17465 max_active: Arc::clone(&max_active),
17466 }))
17467 .build()
17468 .unwrap(),
17469 );
17470
17471 let left = {
17472 let agent = Arc::clone(&agent);
17473 tokio::spawn(async move {
17474 agent
17475 .invoke_tool(ToolExecutionRequest::new(
17476 "lock-1",
17477 "locked_write",
17478 serde_json::json!({"path": "./same.txt"}),
17479 ToolCallSource::Manual,
17480 ))
17481 .await
17482 .unwrap()
17483 })
17484 };
17485 let right = {
17486 let agent = Arc::clone(&agent);
17487 tokio::spawn(async move {
17488 agent
17489 .invoke_tool(ToolExecutionRequest::new(
17490 "lock-2",
17491 "locked_write",
17492 serde_json::json!({"path": "./same.txt"}),
17493 ToolCallSource::Manual,
17494 ))
17495 .await
17496 .unwrap()
17497 })
17498 };
17499
17500 let left = left.await.unwrap();
17501 let right = right.await.unwrap();
17502 assert!(left.success);
17503 assert!(right.success);
17504 assert_eq!(max_active.load(Ordering::SeqCst), 1);
17505 }
17506
17507 #[tokio::test]
17508 async fn path_resources_use_shared_global_lock_and_cleanup() {
17509 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17510 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17511 let bindings = ai_agents_core::ToolPolicyBindings {
17512 path_fields: vec![
17513 ai_agents_core::PathPolicyBinding::read_write("source_path"),
17514 ai_agents_core::PathPolicyBinding::write("destination_path"),
17515 ],
17516 ..Default::default()
17517 };
17518 let classification = ai_agents_core::ToolCallClassification::from_metadata(
17519 &MultiResourceWriteTool {
17520 active: Arc::clone(&active),
17521 max_active: Arc::clone(&max_active),
17522 }
17523 .safety_metadata(),
17524 );
17525 let left_args = serde_json::json!({
17526 "source_path": "./a/../first.txt",
17527 "destination_path": "./second.txt"
17528 });
17529 let right_args = serde_json::json!({
17530 "source_path": "./second.txt",
17531 "destination_path": "./first.txt"
17532 });
17533 let left_keys = tool_resource_lock_keys(
17534 "multi_resource_write",
17535 &left_args,
17536 &bindings,
17537 &classification,
17538 );
17539 let right_keys = tool_resource_lock_keys(
17540 "multi_resource_write",
17541 &right_args,
17542 &bindings,
17543 &classification,
17544 );
17545 assert_eq!(left_keys, right_keys);
17546 assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
17547
17548 let locks = new_tool_resource_locks();
17549 let build_agent = || {
17550 AgentBuilder::new()
17551 .system_prompt("Test shared resource locks.")
17552 .llm(Arc::new(mock_with_response("done")))
17553 .tool(Arc::new(MultiResourceWriteTool {
17554 active: Arc::clone(&active),
17555 max_active: Arc::clone(&max_active),
17556 }))
17557 .build()
17558 .unwrap()
17559 .with_shared_resource_locks(Arc::clone(&locks))
17560 };
17561 let left_agent = Arc::new(build_agent());
17562 let right_agent = Arc::new(build_agent());
17563 let left = tokio::spawn(async move {
17564 left_agent
17565 .invoke_tool(ToolExecutionRequest::new(
17566 "multi-left",
17567 "multi_resource_write",
17568 left_args,
17569 ToolCallSource::Manual,
17570 ))
17571 .await
17572 .unwrap()
17573 });
17574 let right = tokio::spawn(async move {
17575 right_agent
17576 .invoke_tool(ToolExecutionRequest::new(
17577 "multi-right",
17578 "multi_resource_write",
17579 right_args,
17580 ToolCallSource::Manual,
17581 ))
17582 .await
17583 .unwrap()
17584 });
17585 let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
17586 tokio::join!(left, right)
17587 })
17588 .await
17589 .expect("reversed resource acquisition must not deadlock");
17590
17591 assert!(left.unwrap().success);
17592 assert!(right.unwrap().success);
17593 assert_eq!(max_active.load(Ordering::SeqCst), 1);
17594 assert!(locks.read().is_empty());
17595 }
17596
17597 #[tokio::test]
17598 async fn global_path_lock_serializes_copy_destination_with_file_write() {
17599 assert_path_mutation_pair_serialized(
17600 "copy_path",
17601 CopyPathTool::new().policy_bindings().path_fields,
17602 serde_json::json!({
17603 "source_path": "./source.txt",
17604 "destination_path": "./shared.txt"
17605 }),
17606 "file_write",
17607 FileWriteTool::new().policy_bindings().path_fields,
17608 serde_json::json!({"path": "./shared.txt"}),
17609 )
17610 .await;
17611 }
17612
17613 #[tokio::test]
17614 async fn parent_and_spawned_runtime_share_global_path_lock() {
17615 let workspace = MutationTestWorkspace::new();
17616 let destination = workspace.root.join("spawned.txt");
17617 let parent_gate = PathMutationGate::new();
17618 let parent = Arc::new(
17619 AgentBuilder::from_yaml(
17620 r#"
17621name: LockParent
17622system_prompt: parent
17623llm:
17624 default: default
17625tools:
17626 - parent_path_write
17627spawner:
17628 shared_llms: true
17629"#,
17630 )
17631 .unwrap()
17632 .llm(Arc::new(mock_with_response("done")))
17633 .auto_configure_spawner()
17634 .await
17635 .unwrap()
17636 .tool(Arc::new(BlockingPathMutationTool {
17637 id: "parent_path_write",
17638 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17639 gate: parent_gate.clone(),
17640 }))
17641 .build()
17642 .unwrap(),
17643 );
17644
17645 let mut child_spec = crate::spec::AgentSpec {
17646 name: "LockChild".to_string(),
17647 system_prompt: "child".to_string(),
17648 tools: Some(vec![crate::spec::ToolEntry::Simple(
17649 "file_write".to_string(),
17650 )]),
17651 ..Default::default()
17652 };
17653 child_spec.tool_security.enabled = true;
17654 child_spec.tool_security.fail_closed = true;
17655 let file_write_policy = ai_agents_tools::ToolPolicyConfig {
17656 write_paths: vec![workspace.root.to_string_lossy().into_owned()],
17657 allow_without_confirmation: true,
17658 ..Default::default()
17659 };
17660 child_spec
17661 .tool_security
17662 .tools
17663 .insert("file_write".to_string(), file_write_policy);
17664 let spawned = parent
17665 .spawner()
17666 .unwrap()
17667 .spawn_from_spec(child_spec)
17668 .await
17669 .unwrap();
17670 assert!(Arc::ptr_eq(
17671 &parent.resource_locks,
17672 &spawned.agent.resource_locks
17673 ));
17674 assert!(!Arc::ptr_eq(
17675 &parent.runtime_control,
17676 &spawned.agent.runtime_control
17677 ));
17678
17679 let parent_call = {
17680 let parent = Arc::clone(&parent);
17681 let destination = destination.clone();
17682 tokio::spawn(async move {
17683 parent
17684 .invoke_tool(ToolExecutionRequest::new(
17685 "parent-lock-holder",
17686 "parent_path_write",
17687 serde_json::json!({"path": destination}),
17688 ToolCallSource::Manual,
17689 ))
17690 .await
17691 .unwrap()
17692 })
17693 };
17694 parent_gate.wait_until_entered().await;
17695
17696 let child_call = {
17697 let child = Arc::clone(&spawned.agent);
17698 let destination = destination.clone();
17699 tokio::spawn(async move {
17700 child
17701 .invoke_tool(ToolExecutionRequest::new(
17702 "spawned-file-write",
17703 "file_write",
17704 serde_json::json!({
17705 "path": destination,
17706 "content": "spawned",
17707 "dry_run": false
17708 }),
17709 ToolCallSource::Manual,
17710 ))
17711 .await
17712 .unwrap()
17713 })
17714 };
17715 wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
17716 assert!(!child_call.is_finished());
17717
17718 parent_gate.release();
17719 let (parent_record, child_record) =
17720 tokio::time::timeout(std::time::Duration::from_secs(2), async {
17721 tokio::join!(parent_call, child_call)
17722 })
17723 .await
17724 .expect("parent and spawned path mutations did not finish");
17725 assert!(parent_record.unwrap().success);
17726 assert!(child_record.unwrap().success);
17727 assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
17728 assert!(parent.resource_locks.read().is_empty());
17729 }
17730
17731 #[tokio::test]
17732 async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
17733 let locks = new_tool_resource_locks();
17734 let holder_gate = PathMutationGate::new();
17735 let waiter_gate = PathMutationGate::new();
17736 waiter_gate.release();
17737 let holder = Arc::new(
17738 AgentBuilder::new()
17739 .system_prompt("Hold the global path lock.")
17740 .llm(Arc::new(mock_with_response("done")))
17741 .tool(Arc::new(BlockingPathMutationTool {
17742 id: "holder_write",
17743 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17744 gate: holder_gate.clone(),
17745 }))
17746 .build()
17747 .unwrap()
17748 .with_shared_resource_locks(Arc::clone(&locks)),
17749 );
17750 let waiter = Arc::new(
17751 AgentBuilder::new()
17752 .system_prompt("Wait for the global path lock.")
17753 .llm(Arc::new(mock_with_response("done")))
17754 .tool(Arc::new(BlockingPathMutationTool {
17755 id: "waiter_write",
17756 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17757 gate: waiter_gate.clone(),
17758 }))
17759 .build()
17760 .unwrap()
17761 .with_shared_resource_locks(Arc::clone(&locks)),
17762 );
17763
17764 let holder_call = {
17765 let holder = Arc::clone(&holder);
17766 tokio::spawn(async move {
17767 holder
17768 .invoke_tool(ToolExecutionRequest::new(
17769 "holder-call",
17770 "holder_write",
17771 serde_json::json!({"path": "./shared.txt"}),
17772 ToolCallSource::Manual,
17773 ))
17774 .await
17775 .unwrap()
17776 })
17777 };
17778 holder_gate.wait_until_entered().await;
17779
17780 let waiter_call = {
17781 let waiter = Arc::clone(&waiter);
17782 tokio::spawn(async move {
17783 waiter
17784 .invoke_tool(ToolExecutionRequest::new(
17785 "waiter-call",
17786 "waiter_write",
17787 serde_json::json!({"path": "./shared.txt"}),
17788 ToolCallSource::Manual,
17789 ))
17790 .await
17791 .unwrap()
17792 })
17793 };
17794 wait_for_resource_lock_strong_count(&locks, 2).await;
17795 waiter.runtime_control().cancel_all();
17796
17797 let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
17798 .await
17799 .expect("cancelled lock waiter did not finish")
17800 .unwrap();
17801 assert!(!waiter_record.success);
17802 assert!(!waiter_record.executed);
17803 assert!(waiter_record.cancelled);
17804 assert_eq!(
17805 waiter_record.cancellation_reason.as_deref(),
17806 Some("runtime control cancellation")
17807 );
17808 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
17809 assert_eq!(
17810 locks
17811 .read()
17812 .get("path-mutation:global")
17813 .map_or(0, |lock| lock.strong_count()),
17814 1
17815 );
17816
17817 holder_gate.release();
17818 let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
17819 .await
17820 .expect("lock holder did not finish")
17821 .unwrap();
17822 assert!(holder_record.success);
17823 assert!(locks.read().is_empty());
17824 }
17825
17826 #[tokio::test]
17827 async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
17828 for denial in [MutationDenial::Policy, MutationDenial::Approval] {
17829 let tools: [Arc<dyn Tool>; 3] = [
17830 Arc::new(CopyPathTool::new()),
17831 Arc::new(MovePathTool::new()),
17832 Arc::new(DeletePathTool::new()),
17833 ];
17834 for tool in tools {
17835 assert_path_mutation_denied(tool, denial).await;
17836 }
17837 }
17838 }
17839
17840 #[tokio::test]
17841 async fn policy_denial_keeps_executor_hook_lifecycle_and_record_authority() {
17842 let workspace = MutationTestWorkspace::new();
17843 let target = workspace.root.join("denied.txt");
17844 let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
17845 let agent = AgentBuilder::new()
17846 .system_prompt("Test denied tool hooks.")
17847 .llm(Arc::new(mock_with_response("done")))
17848 .tool(Arc::new(FileWriteTool::new()))
17849 .tool_security(ToolSecurityEngine::new(mutation_denial_security_config(
17850 "file_write",
17851 &workspace.root,
17852 MutationDenial::Policy,
17853 )))
17854 .hooks(hooks.clone())
17855 .build()
17856 .unwrap();
17857
17858 let record = agent
17859 .invoke_tool(ToolExecutionRequest::new(
17860 "denied-hook-call",
17861 "file_write",
17862 serde_json::json!({
17863 "path": target.to_string_lossy(),
17864 "content": "blocked"
17865 }),
17866 ToolCallSource::Manual,
17867 ))
17868 .await
17869 .unwrap();
17870
17871 assert!(!record.executed);
17872 assert!(!record.success);
17873 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
17874 assert_eq!(
17875 hooks.events(),
17876 vec![
17877 "start:file_write",
17878 "complete:file_write:false",
17879 "record:file_write:false",
17880 "error"
17881 ]
17882 );
17883 assert!(!target.exists());
17884 }
17885
17886 #[tokio::test]
17887 async fn approval_argument_changes_are_rechecked_against_final_scope() {
17888 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17889 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17890 let entered = Arc::new(tokio::sync::Barrier::new(2));
17891 let release = Arc::new(tokio::sync::Notify::new());
17892 let handler = Arc::new(BlockingApprovalHandler {
17893 entered: Arc::clone(&entered),
17894 release: Arc::clone(&release),
17895 result: ApprovalResult::Modified {
17896 changes: HashMap::from([(
17897 "path".to_string(),
17898 Value::String("./after-approval.txt".to_string()),
17899 )]),
17900 },
17901 });
17902 let agent = Arc::new(
17903 AgentBuilder::new()
17904 .system_prompt("Test final scope validation.")
17905 .llm(Arc::new(mock_with_response("done")))
17906 .tool(Arc::new(LockedWriteTool {
17907 active: Arc::clone(&active),
17908 max_active: Arc::clone(&max_active),
17909 }))
17910 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17911 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17912 .approval_handler(handler)
17913 .build()
17914 .unwrap(),
17915 );
17916 let control = agent.runtime_control();
17917 let running = Arc::clone(&agent);
17918 let call = tokio::spawn(async move {
17919 running
17920 .invoke_tool(ToolExecutionRequest::new(
17921 "approval-scope",
17922 "locked_write",
17923 serde_json::json!({"path": "./before-approval.txt"}),
17924 ToolCallSource::Manual,
17925 ))
17926 .await
17927 .unwrap()
17928 });
17929 entered.wait().await;
17930 let expected_version = control.set_tool_scope(Vec::new());
17931 release.notify_one();
17932 let record = call.await.unwrap();
17933
17934 assert!(!record.executed);
17935 assert!(!record.success);
17936 assert_eq!(record.runtime_config_version, expected_version);
17937 assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
17938 assert_eq!(max_active.load(Ordering::SeqCst), 0);
17939 assert_eq!(
17940 record.metadata["runtime_scope_snapshot"],
17941 serde_json::json!([])
17942 );
17943 }
17944
17945 #[tokio::test]
17946 async fn approval_is_rechecked_against_final_policy_snapshot() {
17947 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17948 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17949 let entered = Arc::new(tokio::sync::Barrier::new(2));
17950 let release = Arc::new(tokio::sync::Notify::new());
17951 let handler = Arc::new(BlockingApprovalHandler {
17952 entered: Arc::clone(&entered),
17953 release: Arc::clone(&release),
17954 result: ApprovalResult::Approved,
17955 });
17956 let agent = Arc::new(
17957 AgentBuilder::new()
17958 .system_prompt("Test final policy validation.")
17959 .llm(Arc::new(mock_with_response("done")))
17960 .tool(Arc::new(LockedWriteTool {
17961 active: Arc::clone(&active),
17962 max_active: Arc::clone(&max_active),
17963 }))
17964 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17965 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17966 .approval_handler(handler)
17967 .build()
17968 .unwrap(),
17969 );
17970 let control = agent.runtime_control();
17971 let running = Arc::clone(&agent);
17972 let call = tokio::spawn(async move {
17973 running
17974 .invoke_tool(ToolExecutionRequest::new(
17975 "approval-policy",
17976 "locked_write",
17977 serde_json::json!({"path": "./policy.txt"}),
17978 ToolCallSource::Manual,
17979 ))
17980 .await
17981 .unwrap()
17982 });
17983 entered.wait().await;
17984 let expected_version = control.set_tool_security(approval_security_config(false));
17985 release.notify_one();
17986 let record = call.await.unwrap();
17987
17988 assert!(!record.executed);
17989 assert!(!record.success);
17990 assert_eq!(record.runtime_config_version, expected_version);
17991 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
17992 assert_eq!(max_active.load(Ordering::SeqCst), 0);
17993 assert!(record.metadata.contains_key("policy_snapshot"));
17994 }
17995
17996 #[test]
17997 fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
17998 let agent = AgentBuilder::new()
17999 .system_prompt("Test runtime policy validation.")
18000 .llm(Arc::new(mock_with_response("done")))
18001 .build()
18002 .unwrap();
18003 let control = agent.runtime_control();
18004 let mut valid = ToolSecurityConfig::default();
18005 valid.tools.insert(
18006 "web_search".to_string(),
18007 ai_agents_tools::ToolPolicyConfig {
18008 max_results: Some(5),
18009 ..Default::default()
18010 },
18011 );
18012 let generation = control.try_set_tool_security(valid).unwrap();
18013
18014 let mut invalid = ToolSecurityConfig::default();
18015 invalid.tools.insert(
18016 "web_search".to_string(),
18017 ai_agents_tools::ToolPolicyConfig {
18018 max_results: Some(0),
18019 ..Default::default()
18020 },
18021 );
18022 let error = control.try_set_tool_security(invalid).unwrap_err();
18023
18024 assert!(
18025 error
18026 .to_string()
18027 .contains("max_results must be greater than 0")
18028 );
18029 assert_eq!(control.version(), generation);
18030 assert_eq!(
18031 control
18032 .state
18033 .tool_security_override
18034 .read()
18035 .as_ref()
18036 .unwrap()
18037 .config()
18038 .tools["web_search"]
18039 .max_results,
18040 Some(5)
18041 );
18042 }
18043
18044 #[test]
18046 fn invalid_timeout_config_stops_before_approval_or_tool_invocation() {
18047 let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18048 let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18049 let spec = crate::spec::AgentSpec {
18050 tool_security: ToolSecurityConfig {
18051 enabled: true,
18052 default_timeout_ms: u64::MAX,
18053 ..Default::default()
18054 },
18055 ..Default::default()
18056 };
18057
18058 let result = AgentBuilder::from_spec(spec)
18059 .llm(Arc::new(mock_with_response("done")))
18060 .tool(Arc::new(FlakyWriteTool {
18061 calls: Arc::clone(&tool_calls),
18062 }))
18063 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18064 .approval_handler(Arc::new(CountingApprovalHandler {
18065 calls: Arc::clone(&approval_calls),
18066 }))
18067 .build();
18068
18069 assert!(result.is_err());
18070 assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
18071 assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
18072 }
18073
18074 #[test]
18076 fn invalid_recovery_timeout_config_stops_before_approval_or_tool_invocation() {
18077 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
18078
18079 let tool_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18080 let approval_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18081 let spec = crate::spec::AgentSpec {
18082 error_recovery: ErrorRecoveryConfig {
18083 tools: ToolRecoveryConfig {
18084 default: ToolRetryConfig {
18085 timeout_ms: Some(u64::MAX),
18086 ..Default::default()
18087 },
18088 ..Default::default()
18089 },
18090 ..Default::default()
18091 },
18092 ..Default::default()
18093 };
18094
18095 let result = AgentBuilder::from_spec(spec)
18096 .llm(Arc::new(mock_with_response("done")))
18097 .tool(Arc::new(FlakyWriteTool {
18098 calls: Arc::clone(&tool_calls),
18099 }))
18100 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18101 .approval_handler(Arc::new(CountingApprovalHandler {
18102 calls: Arc::clone(&approval_calls),
18103 }))
18104 .build();
18105
18106 assert!(result.is_err());
18107 assert_eq!(approval_calls.load(Ordering::SeqCst), 0);
18108 assert_eq!(tool_calls.load(Ordering::SeqCst), 0);
18109 }
18110
18111 #[test]
18113 fn invalid_timeout_policy_does_not_replace_snapshot_or_generation() {
18114 let agent = AgentBuilder::new()
18115 .system_prompt("Test runtime timeout policy validation.")
18116 .llm(Arc::new(mock_with_response("done")))
18117 .build()
18118 .unwrap();
18119 let control = agent.runtime_control();
18120 let valid = ToolSecurityConfig {
18121 default_timeout_ms: 5_000,
18122 ..Default::default()
18123 };
18124 let generation = control.try_set_tool_security(valid).unwrap();
18125
18126 let invalid = ToolSecurityConfig {
18127 default_timeout_ms: MAX_TOOL_TIMEOUT_MS + 1,
18128 ..Default::default()
18129 };
18130 let error = control.try_set_tool_security(invalid).unwrap_err();
18131
18132 assert!(error.to_string().contains(&format!(
18133 "tool_security.default_timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
18134 )));
18135 assert_eq!(control.version(), generation);
18136 assert_eq!(
18137 control
18138 .state
18139 .tool_security_override
18140 .read()
18141 .as_ref()
18142 .unwrap()
18143 .config()
18144 .default_timeout_ms,
18145 5_000
18146 );
18147 }
18148
18149 #[test]
18151 fn runtime_tool_timeout_conversion_enforces_the_stable_boundary() {
18152 let timeout = RuntimeAgent::validated_tool_timeout(MAX_TOOL_TIMEOUT_MS).unwrap();
18153 assert_eq!(timeout.timer, Duration::from_millis(MAX_TOOL_TIMEOUT_MS));
18154 assert_eq!(
18155 timeout.deadline_delta,
18156 chrono::Duration::milliseconds(MAX_TOOL_TIMEOUT_MS as i64)
18157 );
18158
18159 for timeout_ms in [MAX_TOOL_TIMEOUT_MS + 1, u64::MAX] {
18160 let error = RuntimeAgent::validated_tool_timeout(timeout_ms).unwrap_err();
18161 assert!(error.to_string().contains(&format!(
18162 "effective tool timeout_ms must be no greater than {MAX_TOOL_TIMEOUT_MS} milliseconds"
18163 )));
18164 }
18165 }
18166
18167 #[tokio::test]
18168 async fn persistent_override_preserves_rate_history_within_generation() {
18169 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18170 let agent = AgentBuilder::new()
18171 .system_prompt("Test persistent policy overrides.")
18172 .llm(Arc::new(mock_with_response("done")))
18173 .tool(Arc::new(RecoveryTestTool {
18174 id: "limited_override".to_string(),
18175 succeeds: true,
18176 calls: Arc::clone(&calls),
18177 max_output_chars: None,
18178 }))
18179 .build()
18180 .unwrap();
18181 let mut security = ToolSecurityConfig {
18182 enabled: true,
18183 fail_closed: true,
18184 ..Default::default()
18185 };
18186 let policy = ai_agents_tools::ToolPolicyConfig {
18187 write_paths: vec![".".to_string()],
18188 rate_limit: Some(1),
18189 ..Default::default()
18190 };
18191 security
18192 .tools
18193 .insert("limited_override".to_string(), policy);
18194 let generation = agent.runtime_control().set_tool_security(security);
18195
18196 let first = agent
18197 .invoke_tool(ToolExecutionRequest::new(
18198 "limited-first",
18199 "limited_override",
18200 serde_json::json!({"path": "./limited.txt"}),
18201 ToolCallSource::Manual,
18202 ))
18203 .await
18204 .unwrap();
18205 let second = agent
18206 .invoke_tool(ToolExecutionRequest::new(
18207 "limited-second",
18208 "limited_override",
18209 serde_json::json!({"path": "./limited.txt"}),
18210 ToolCallSource::Manual,
18211 ))
18212 .await
18213 .unwrap();
18214
18215 assert!(first.success);
18216 assert_eq!(first.policy_version, generation);
18217 assert!(!second.executed);
18218 assert!(second.output.contains("Rate limit exceeded"));
18219 assert_eq!(second.policy_version, generation);
18220 assert_eq!(calls.load(Ordering::SeqCst), 1);
18221 }
18222
18223 #[tokio::test]
18224 async fn concurrent_rate_admission_consumes_capacity_atomically() {
18225 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18226 let tool = Arc::new(RecoveryTestTool {
18227 id: "atomic_rate".to_string(),
18228 succeeds: true,
18229 calls: Arc::clone(&calls),
18230 max_output_chars: None,
18231 });
18232 let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
18233 let bindings = tool.policy_bindings();
18234 let classification = tool.classify_call(&arguments);
18235 let resource_keys =
18236 tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
18237 let mut security = ToolSecurityConfig {
18238 enabled: true,
18239 fail_closed: true,
18240 ..Default::default()
18241 };
18242 let policy = ai_agents_tools::ToolPolicyConfig {
18243 write_paths: vec![".".to_string()],
18244 rate_limit: Some(1),
18245 ..Default::default()
18246 };
18247 security.tools.insert(tool.id().to_string(), policy);
18248 let agent = Arc::new(
18249 AgentBuilder::new()
18250 .system_prompt("Test atomic rate admission.")
18251 .llm(Arc::new(mock_with_response("done")))
18252 .tool(tool)
18253 .tool_security(ToolSecurityEngine::new(security))
18254 .build()
18255 .unwrap(),
18256 );
18257 let held = agent
18258 .acquire_tool_resource_locks(&resource_keys)
18259 .await
18260 .unwrap();
18261 let left = {
18262 let agent = Arc::clone(&agent);
18263 let arguments = arguments.clone();
18264 tokio::spawn(async move {
18265 agent
18266 .invoke_tool(ToolExecutionRequest::new(
18267 "atomic-rate-left",
18268 "atomic_rate",
18269 arguments,
18270 ToolCallSource::Manual,
18271 ))
18272 .await
18273 .unwrap()
18274 })
18275 };
18276 let right = {
18277 let agent = Arc::clone(&agent);
18278 tokio::spawn(async move {
18279 agent
18280 .invoke_tool(ToolExecutionRequest::new(
18281 "atomic-rate-right",
18282 "atomic_rate",
18283 arguments,
18284 ToolCallSource::Manual,
18285 ))
18286 .await
18287 .unwrap()
18288 })
18289 };
18290 tokio::time::sleep(std::time::Duration::from_millis(25)).await;
18291 drop(held);
18292 let (left, right) = tokio::join!(left, right);
18293 let records = [left.unwrap(), right.unwrap()];
18294
18295 assert_eq!(records.iter().filter(|record| record.success).count(), 1);
18296 assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
18297 assert!(
18298 records.iter().any(|record| {
18299 !record.executed && record.output.contains("Rate limit exceeded")
18300 })
18301 );
18302 assert_eq!(calls.load(Ordering::SeqCst), 1);
18303 }
18304
18305 #[tokio::test]
18306 async fn changed_policy_generation_invalidates_pending_approval() {
18307 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18308 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18309 let entered = Arc::new(tokio::sync::Barrier::new(2));
18310 let release = Arc::new(tokio::sync::Notify::new());
18311 let handler = Arc::new(BlockingApprovalHandler {
18312 entered: Arc::clone(&entered),
18313 release: Arc::clone(&release),
18314 result: ApprovalResult::Approved,
18315 });
18316 let agent = Arc::new(
18317 AgentBuilder::new()
18318 .system_prompt("Test stale approval denial.")
18319 .llm(Arc::new(mock_with_response("done")))
18320 .tool(Arc::new(LockedWriteTool {
18321 active: Arc::clone(&active),
18322 max_active: Arc::clone(&max_active),
18323 }))
18324 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
18325 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18326 .approval_handler(handler)
18327 .build()
18328 .unwrap(),
18329 );
18330 let running = Arc::clone(&agent);
18331 let call = tokio::spawn(async move {
18332 running
18333 .invoke_tool(ToolExecutionRequest::new(
18334 "stale-approval",
18335 "locked_write",
18336 serde_json::json!({"path": "./stale.txt"}),
18337 ToolCallSource::Manual,
18338 ))
18339 .await
18340 .unwrap()
18341 });
18342 entered.wait().await;
18343 let generation = agent
18344 .runtime_control()
18345 .set_tool_security(approval_security_config(true));
18346 release.notify_one();
18347 let record = call.await.unwrap();
18348
18349 assert!(!record.executed);
18350 assert!(record.output.contains("Approval became stale"));
18351 assert_eq!(record.policy_version, generation);
18352 assert_eq!(max_active.load(Ordering::SeqCst), 0);
18353 }
18354
18355 #[tokio::test]
18356 async fn final_policy_reapplies_argument_caps_after_approval_changes() {
18357 use ai_agents_hitl::CallbackHandler;
18358
18359 let mut security = ToolSecurityConfig {
18360 enabled: true,
18361 fail_closed: true,
18362 ..Default::default()
18363 };
18364 let policy = ai_agents_tools::ToolPolicyConfig {
18365 read_paths: vec![".".to_string()],
18366 max_results: Some(5),
18367 require_confirmation: true,
18368 ..Default::default()
18369 };
18370 security.tools.insert("context_echo".to_string(), policy);
18371 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
18372 changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
18373 });
18374 let agent = AgentBuilder::new()
18375 .system_prompt("Test final argument caps.")
18376 .llm(Arc::new(mock_with_response("done")))
18377 .tool(Arc::new(ContextEchoTool))
18378 .tool_security(ToolSecurityEngine::new(security))
18379 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18380 .approval_handler(Arc::new(handler))
18381 .build()
18382 .unwrap();
18383
18384 let record = agent
18385 .invoke_tool(ToolExecutionRequest::new(
18386 "final-cap",
18387 "context_echo",
18388 serde_json::json!({"path": ".", "max_results": 1}),
18389 ToolCallSource::Manual,
18390 ))
18391 .await
18392 .unwrap();
18393
18394 assert!(record.success);
18395 assert_eq!(record.executed_arguments["max_results"], 5);
18396 assert_eq!(
18397 record.approval.unwrap().modified_arguments.unwrap()["max_results"],
18398 5
18399 );
18400 }
18401
18402 #[tokio::test]
18403 async fn no_binding_writes_use_canonical_fallback_lock() {
18404 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18405 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18406 let agent = Arc::new(
18407 AgentBuilder::new()
18408 .system_prompt("Test fallback resource locks.")
18409 .llm(Arc::new(mock_with_response("done")))
18410 .tool(Arc::new(NoBindingWriteTool {
18411 active: Arc::clone(&active),
18412 max_active: Arc::clone(&max_active),
18413 }))
18414 .build()
18415 .unwrap(),
18416 );
18417 let left = {
18418 let agent = Arc::clone(&agent);
18419 tokio::spawn(async move {
18420 agent
18421 .invoke_tool(ToolExecutionRequest::new(
18422 "no-binding-left",
18423 "no_binding_write",
18424 serde_json::json!({}),
18425 ToolCallSource::Manual,
18426 ))
18427 .await
18428 .unwrap()
18429 })
18430 };
18431 let right = {
18432 let agent = Arc::clone(&agent);
18433 tokio::spawn(async move {
18434 agent
18435 .invoke_tool(ToolExecutionRequest::new(
18436 "no-binding-right",
18437 "no_binding_write",
18438 serde_json::json!({}),
18439 ToolCallSource::Manual,
18440 ))
18441 .await
18442 .unwrap()
18443 })
18444 };
18445 let (left, right) = tokio::join!(left, right);
18446
18447 assert!(left.unwrap().success);
18448 assert!(right.unwrap().success);
18449 assert_eq!(max_active.load(Ordering::SeqCst), 1);
18450 }
18451
18452 #[tokio::test]
18453 async fn parent_and_child_paths_share_a_resource_lock() {
18454 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18455 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18456 let agent = Arc::new(
18457 AgentBuilder::new()
18458 .system_prompt("Test parent child resource locks.")
18459 .llm(Arc::new(mock_with_response("done")))
18460 .tool(Arc::new(LockedWriteTool {
18461 active: Arc::clone(&active),
18462 max_active: Arc::clone(&max_active),
18463 }))
18464 .build()
18465 .unwrap(),
18466 );
18467 let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
18468 let child = format!("{}/child.txt", parent);
18469 let left = {
18470 let agent = Arc::clone(&agent);
18471 tokio::spawn(async move {
18472 agent
18473 .invoke_tool(ToolExecutionRequest::new(
18474 "parent-lock",
18475 "locked_write",
18476 serde_json::json!({"path": parent}),
18477 ToolCallSource::Manual,
18478 ))
18479 .await
18480 .unwrap()
18481 })
18482 };
18483 let right = {
18484 let agent = Arc::clone(&agent);
18485 tokio::spawn(async move {
18486 agent
18487 .invoke_tool(ToolExecutionRequest::new(
18488 "child-lock",
18489 "locked_write",
18490 serde_json::json!({"path": child}),
18491 ToolCallSource::Manual,
18492 ))
18493 .await
18494 .unwrap()
18495 })
18496 };
18497 let (left, right) = tokio::join!(left, right);
18498
18499 assert!(left.unwrap().success);
18500 assert!(right.unwrap().success);
18501 assert_eq!(max_active.load(Ordering::SeqCst), 1);
18502 }
18503
18504 #[tokio::test]
18505 async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
18506 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18507 let hooks = Arc::new(ReentrantToolHooks {
18508 agent: parking_lot::Mutex::new(None),
18509 invoked: AtomicBool::new(false),
18510 nested_success: AtomicBool::new(false),
18511 });
18512 let agent = Arc::new(
18513 AgentBuilder::new()
18514 .system_prompt("Test hook reentrancy.")
18515 .llm(Arc::new(mock_with_response("done")))
18516 .tool(Arc::new(RecoveryTestTool {
18517 id: "reentrant_write".to_string(),
18518 succeeds: true,
18519 calls: Arc::clone(&calls),
18520 max_output_chars: None,
18521 }))
18522 .hooks(hooks.clone())
18523 .build()
18524 .unwrap(),
18525 );
18526 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
18527 let record = tokio::time::timeout(
18528 std::time::Duration::from_secs(2),
18529 agent.invoke_tool(ToolExecutionRequest::new(
18530 "outer-hook-call",
18531 "reentrant_write",
18532 serde_json::json!({"path": "./hook.txt"}),
18533 ToolCallSource::Manual,
18534 )),
18535 )
18536 .await
18537 .expect("tool completion hook must not retain resource guards")
18538 .unwrap();
18539
18540 assert!(record.success);
18541 assert!(hooks.nested_success.load(Ordering::SeqCst));
18542 assert_eq!(calls.load(Ordering::SeqCst), 2);
18543 }
18544
18545 #[tokio::test]
18547 async fn fallback_finalizes_original_record_before_shared_execution() {
18548 let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18549 let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18550 let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
18551 let agent = AgentBuilder::new()
18552 .system_prompt("Test fallback execution.")
18553 .llm(Arc::new(mock_with_response("done")))
18554 .tool(Arc::new(RecoveryTestTool {
18555 id: "primary".to_string(),
18556 succeeds: false,
18557 calls: Arc::clone(&primary_calls),
18558 max_output_chars: None,
18559 }))
18560 .tool(Arc::new(RecoveryTestTool {
18561 id: "fallback".to_string(),
18562 succeeds: true,
18563 calls: Arc::clone(&fallback_calls),
18564 max_output_chars: None,
18565 }))
18566 .recovery_manager(recovery_manager_with_fallbacks([(
18567 "primary".to_string(),
18568 "fallback".to_string(),
18569 )]))
18570 .hooks(hooks.clone())
18571 .build()
18572 .unwrap();
18573 let record = tokio::time::timeout(
18574 std::time::Duration::from_secs(2),
18575 agent.invoke_tool(ToolExecutionRequest::new(
18576 "fallback-call",
18577 "primary",
18578 serde_json::json!({"path": "./shared.txt"}),
18579 ToolCallSource::Manual,
18580 )),
18581 )
18582 .await
18583 .expect("fallback must not retain the primary resource guard")
18584 .unwrap();
18585
18586 assert_eq!(
18587 hooks.events(),
18588 vec![
18589 "start:primary",
18590 "complete:primary:false",
18591 "record:primary:true",
18592 "error",
18593 "start:fallback",
18594 "complete:fallback:true",
18595 "record:fallback:true",
18596 ]
18597 );
18598 let records = hooks.records();
18599 assert_eq!(records.len(), 2);
18600 let original = &records[0];
18601 assert_eq!(original.canonical_id, "primary");
18602 assert!(matches!(original.source, ToolCallSource::Manual));
18603 assert!(original.executed);
18604 assert!(!original.success);
18605
18606 let fallback = &records[1];
18607 assert_eq!(fallback.canonical_id, "fallback");
18608 assert_eq!(fallback.call_id, "fallback-call");
18609 assert!(matches!(
18610 &fallback.source,
18611 ToolCallSource::Fallback { original_tool } if original_tool == "primary"
18612 ));
18613 assert!(fallback.executed);
18614 assert!(fallback.success);
18615 assert_eq!(record.canonical_id, fallback.canonical_id);
18616 assert_eq!(record.output, fallback.output);
18617
18618 let history = agent.tool_call_history();
18619 assert_eq!(
18620 history
18621 .iter()
18622 .map(|entry| entry.tool_id.as_str())
18623 .collect::<Vec<_>>(),
18624 vec!["primary", "fallback"]
18625 );
18626 assert_eq!(history[0].result.get("success"), Some(&Value::Bool(false)));
18627 assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
18628 assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
18629 }
18630
18631 #[tokio::test]
18633 async fn self_fallback_cycle_is_denied_before_reinvocation() {
18634 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18635 let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
18636 let agent = AgentBuilder::new()
18637 .system_prompt("Test self-fallback cycle admission.")
18638 .llm(Arc::new(mock_with_response("done")))
18639 .tool(Arc::new(RecoveryTestTool {
18640 id: "primary".to_string(),
18641 succeeds: false,
18642 calls: Arc::clone(&calls),
18643 max_output_chars: None,
18644 }))
18645 .recovery_manager(recovery_manager_with_fallbacks([(
18646 "primary".to_string(),
18647 "primary".to_string(),
18648 )]))
18649 .hooks(hooks.clone())
18650 .build()
18651 .unwrap();
18652
18653 let record = tokio::time::timeout(
18654 std::time::Duration::from_secs(2),
18655 agent.invoke_tool(ToolExecutionRequest::new(
18656 "self-fallback-call",
18657 "primary",
18658 serde_json::json!({"path": "./shared.txt"}),
18659 ToolCallSource::Manual,
18660 )),
18661 )
18662 .await
18663 .expect("self fallback must terminate without recursive execution")
18664 .unwrap();
18665
18666 assert_eq!(calls.load(Ordering::SeqCst), 1);
18667 assert_eq!(record.canonical_id, "primary");
18668 assert!(!record.executed);
18669 assert!(!record.success);
18670 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
18671 assert!(record.output.contains("fallback cycle"));
18672 assert!(matches!(
18673 record.source,
18674 ToolCallSource::Fallback { ref original_tool } if original_tool == "primary"
18675 ));
18676 assert_eq!(
18677 record.metadata.get("fallback_chain"),
18678 Some(&serde_json::json!(["primary"]))
18679 );
18680 assert_eq!(
18681 hooks.events(),
18682 vec![
18683 "start:primary",
18684 "complete:primary:false",
18685 "record:primary:true",
18686 "error",
18687 "complete:primary:false",
18688 "record:primary:false",
18689 "error",
18690 ]
18691 );
18692 assert_eq!(agent.tool_call_history().len(), 2);
18693 }
18694
18695 #[tokio::test]
18697 async fn alias_mediated_fallback_cycle_is_denied_canonically() {
18698 let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18699 let secondary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18700 let hooks = Arc::new(ToolLifecycleRecordingHooks::new());
18701 let agent = AgentBuilder::new()
18702 .system_prompt("Test canonical fallback cycle admission.")
18703 .llm(Arc::new(mock_with_response("done")))
18704 .tool(Arc::new(RecoveryTestTool {
18705 id: "primary".to_string(),
18706 succeeds: false,
18707 calls: Arc::clone(&primary_calls),
18708 max_output_chars: None,
18709 }))
18710 .tool(Arc::new(RecoveryTestTool {
18711 id: "secondary".to_string(),
18712 succeeds: false,
18713 calls: Arc::clone(&secondary_calls),
18714 max_output_chars: None,
18715 }))
18716 .recovery_manager(recovery_manager_with_fallbacks([
18717 ("primary".to_string(), "secondary".to_string()),
18718 ("secondary".to_string(), "primary alias".to_string()),
18719 ]))
18720 .hooks(hooks.clone())
18721 .build()
18722 .unwrap();
18723 agent.tools.set_tool_aliases(
18724 "primary",
18725 ToolAliases::new().with_name("en", "primary alias"),
18726 );
18727
18728 let record = tokio::time::timeout(
18729 std::time::Duration::from_secs(2),
18730 agent.invoke_tool(ToolExecutionRequest::new(
18731 "alias-fallback-call",
18732 "primary",
18733 serde_json::json!({"path": "./shared.txt"}),
18734 ToolCallSource::Manual,
18735 )),
18736 )
18737 .await
18738 .expect("alias-mediated fallback cycle must terminate")
18739 .unwrap();
18740
18741 assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
18742 assert_eq!(secondary_calls.load(Ordering::SeqCst), 1);
18743 assert_eq!(record.requested_name, "primary alias");
18744 assert_eq!(record.canonical_id, "primary");
18745 assert!(!record.executed);
18746 assert!(record.output.contains("fallback cycle"));
18747 assert_eq!(
18748 record.metadata.get("fallback_chain"),
18749 Some(&serde_json::json!(["primary", "secondary"]))
18750 );
18751 assert_eq!(hooks.records().len(), 3);
18752 assert_eq!(agent.tool_call_history().len(), 3);
18753 }
18754
18755 #[tokio::test]
18757 async fn final_canonical_drift_cannot_bypass_fallback_ancestry() {
18758 let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18759 let secondary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18760 let provider = Arc::new(DriftingFallbackProvider {
18761 refreshed: AtomicBool::new(false),
18762 primary_calls: Arc::clone(&primary_calls),
18763 secondary_calls: Arc::clone(&secondary_calls),
18764 });
18765 let registry = ToolRegistry::new();
18766 registry.register_provider(provider).await.unwrap();
18767 let lifecycle = Arc::new(ToolLifecycleRecordingHooks::new());
18768 let hooks = Arc::new(RefreshFallbackProviderHooks {
18769 agent: parking_lot::Mutex::new(None),
18770 lifecycle: Arc::clone(&lifecycle),
18771 });
18772 let agent = Arc::new(
18773 AgentBuilder::new()
18774 .system_prompt("Test final canonical fallback admission.")
18775 .llm(Arc::new(mock_with_response("done")))
18776 .tools(registry)
18777 .recovery_manager(recovery_manager_with_fallbacks([(
18778 "primary".to_string(),
18779 "fallback alias".to_string(),
18780 )]))
18781 .hooks(hooks.clone())
18782 .build()
18783 .unwrap(),
18784 );
18785 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
18786
18787 let record = agent
18788 .invoke_tool(ToolExecutionRequest::new(
18789 "drifting-fallback-call",
18790 "primary",
18791 serde_json::json!({"path": "./shared.txt"}),
18792 ToolCallSource::Manual,
18793 ))
18794 .await
18795 .unwrap();
18796
18797 assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
18798 assert_eq!(secondary_calls.load(Ordering::SeqCst), 0);
18799 assert_eq!(record.canonical_id, "secondary");
18800 assert!(!record.executed);
18801 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
18802 assert!(record.output.contains("fallback cycle"));
18803 assert_eq!(
18804 record.metadata.get("fallback_chain"),
18805 Some(&serde_json::json!(["primary", "secondary"]))
18806 );
18807 assert_eq!(
18808 record.metadata.get("final_resolved_canonical_id"),
18809 Some(&serde_json::json!("primary"))
18810 );
18811 assert_eq!(
18812 lifecycle.events(),
18813 vec![
18814 "start:primary",
18815 "complete:primary:false",
18816 "record:primary:true",
18817 "error",
18818 "start:secondary",
18819 "complete:secondary:false",
18820 "record:secondary:false",
18821 "error",
18822 ]
18823 );
18824 let records = lifecycle.records();
18825 assert_eq!(records.len(), 2);
18826 assert_eq!(records[1].canonical_id, "secondary");
18827 assert_eq!(
18828 records[1].metadata.get("final_resolved_canonical_id"),
18829 Some(&serde_json::json!("primary"))
18830 );
18831 let history = agent.tool_call_history();
18832 assert_eq!(
18833 history
18834 .iter()
18835 .map(|entry| entry.tool_id.as_str())
18836 .collect::<Vec<_>>(),
18837 vec!["primary", "secondary"]
18838 );
18839 }
18840
18841 #[tokio::test]
18843 async fn acyclic_fallback_chain_is_denied_after_the_hop_limit() {
18844 let tool_count = MAX_TOOL_FALLBACK_HOPS + 2;
18845 let calls = (0..tool_count)
18846 .map(|_| Arc::new(std::sync::atomic::AtomicUsize::new(0)))
18847 .collect::<Vec<_>>();
18848 let mut builder = AgentBuilder::new()
18849 .system_prompt("Test bounded acyclic fallback admission.")
18850 .llm(Arc::new(mock_with_response("done")));
18851 for (index, counter) in calls.iter().enumerate() {
18852 builder = builder.tool(Arc::new(RecoveryTestTool {
18853 id: format!("fallback_{index}"),
18854 succeeds: false,
18855 calls: Arc::clone(counter),
18856 max_output_chars: None,
18857 }));
18858 }
18859 let fallbacks = (0..tool_count - 1).map(|index| {
18860 (
18861 format!("fallback_{index}"),
18862 format!("fallback_{}", index + 1),
18863 )
18864 });
18865 let agent = builder
18866 .recovery_manager(recovery_manager_with_fallbacks(fallbacks))
18867 .build()
18868 .unwrap();
18869
18870 let record = tokio::time::timeout(
18871 std::time::Duration::from_secs(2),
18872 agent.invoke_tool(ToolExecutionRequest::new(
18873 "bounded-fallback-call",
18874 "fallback_0",
18875 serde_json::json!({"path": "./shared.txt"}),
18876 ToolCallSource::Manual,
18877 )),
18878 )
18879 .await
18880 .expect("bounded fallback chain must terminate")
18881 .unwrap();
18882
18883 for counter in calls.iter().take(MAX_TOOL_FALLBACK_HOPS + 1) {
18884 assert_eq!(counter.load(Ordering::SeqCst), 1);
18885 }
18886 assert_eq!(calls[MAX_TOOL_FALLBACK_HOPS + 1].load(Ordering::SeqCst), 0);
18887 assert_eq!(
18888 record.canonical_id,
18889 format!("fallback_{}", MAX_TOOL_FALLBACK_HOPS + 1)
18890 );
18891 assert!(!record.executed);
18892 assert!(record.output.contains("maximum of 16 hops"));
18893 assert_eq!(agent.tool_call_history().len(), tool_count);
18894 }
18895
18896 #[tokio::test]
18897 async fn diagnostics_without_provider_records_unavailable_without_execution() {
18898 let mock = mock_with_response("hello");
18899 let yaml = r#"
18900name: DiagnosticsNoProviderAgent
18901system_prompt: "Review diagnostics."
18902tools: [diagnostics]
18903"#;
18904 let agent = AgentBuilder::from_yaml(yaml)
18905 .unwrap()
18906 .llm(Arc::new(mock))
18907 .auto_configure_features()
18908 .unwrap()
18909 .build()
18910 .unwrap();
18911
18912 let record = agent
18913 .invoke_tool(ToolExecutionRequest::new(
18914 "diagnostics-call",
18915 "diagnostics",
18916 serde_json::json!({}),
18917 ToolCallSource::Manual,
18918 ))
18919 .await
18920 .unwrap();
18921
18922 assert!(!record.executed);
18923 assert!(!record.success);
18924 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
18925 }
18926
18927 #[tokio::test]
18928 async fn web_search_without_provider_records_unavailable_without_execution() {
18929 let mock = mock_with_response("hello");
18930 let yaml = r#"
18931name: WebSearchNoProviderAgent
18932system_prompt: "You search the web."
18933tools: [web_search]
18934"#;
18935 let agent = AgentBuilder::from_yaml(yaml)
18936 .unwrap()
18937 .llm(Arc::new(mock))
18938 .auto_configure_features()
18939 .unwrap()
18940 .build()
18941 .unwrap();
18942
18943 let record = agent
18944 .invoke_tool(ToolExecutionRequest::new(
18945 "web-search-call",
18946 "web_search",
18947 serde_json::json!({"query": "rust async"}),
18948 ToolCallSource::Manual,
18949 ))
18950 .await
18951 .unwrap();
18952
18953 assert!(!record.executed);
18954 assert!(!record.success);
18955 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
18956 }
18957
18958 #[tokio::test]
18959 async fn unavailable_host_tool_does_not_request_approval() {
18960 let approvals = Arc::new(std::sync::atomic::AtomicUsize::new(0));
18961 let handler = Arc::new(CountingApprovalHandler {
18962 calls: Arc::clone(&approvals),
18963 });
18964 let mut security = ToolSecurityConfig {
18965 enabled: true,
18966 fail_closed: true,
18967 ..Default::default()
18968 };
18969 security.tools.insert(
18970 "web_search".to_string(),
18971 ai_agents_tools::ToolPolicyConfig {
18972 enabled: true,
18973 require_confirmation: true,
18974 ..Default::default()
18975 },
18976 );
18977 let yaml = r#"
18978name: UnavailableApprovalAgent
18979system_prompt: "Search only with approval."
18980tools: [web_search]
18981"#;
18982 let agent = AgentBuilder::from_yaml(yaml)
18983 .unwrap()
18984 .llm(Arc::new(mock_with_response("done")))
18985 .auto_configure_features()
18986 .unwrap()
18987 .tool_security(ToolSecurityEngine::new(security))
18988 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
18989 .approval_handler(handler)
18990 .build()
18991 .unwrap();
18992
18993 let record = agent
18994 .invoke_tool(ToolExecutionRequest::new(
18995 "unavailable-before-approval",
18996 "web_search",
18997 serde_json::json!({"query": "rust async"}),
18998 ToolCallSource::Manual,
18999 ))
19000 .await
19001 .unwrap();
19002
19003 assert_eq!(approvals.load(Ordering::SeqCst), 0);
19004 assert!(!record.executed);
19005 assert!(!record.success);
19006 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
19007 assert!(
19008 record
19009 .approval
19010 .as_ref()
19011 .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Unavailable))
19012 );
19013 }
19014
19015 #[tokio::test]
19016 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
19017 let mock = mock_with_response("hello");
19018 let yaml = r#"
19019name: SpawnerNoGrantAgent
19020system_prompt: "You manage agents."
19021spawner:
19022 max_agents: 2
19023"#;
19024 let agent = AgentBuilder::from_yaml(yaml)
19025 .unwrap()
19026 .llm(Arc::new(mock))
19027 .auto_configure_features()
19028 .unwrap()
19029 .auto_configure_spawner()
19030 .await
19031 .unwrap()
19032 .build()
19033 .unwrap();
19034
19035 let available = agent.get_available_tool_ids().await.unwrap();
19036 assert!(available.is_empty());
19037 }
19038
19039 #[tokio::test]
19040 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
19041 let mock = mock_with_response("hello");
19042 let yaml = r#"
19043name: EmptySpawnerNoGrantAgent
19044system_prompt: "You manage agents."
19045tools: []
19046spawner:
19047 max_agents: 2
19048"#;
19049 let agent = AgentBuilder::from_yaml(yaml)
19050 .unwrap()
19051 .llm(Arc::new(mock))
19052 .auto_configure_features()
19053 .unwrap()
19054 .auto_configure_spawner()
19055 .await
19056 .unwrap()
19057 .build()
19058 .unwrap();
19059
19060 let available = agent.get_available_tool_ids().await.unwrap();
19061 assert!(available.is_empty());
19062 }
19063
19064 #[tokio::test]
19065 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
19066 let mock = mock_with_response("hello");
19067 let yaml = r#"
19068name: ManagementGrantAgent
19069system_prompt: "You manage agents."
19070tools: []
19071spawner:
19072 management_tools: true
19073"#;
19074 let agent = AgentBuilder::from_yaml(yaml)
19075 .unwrap()
19076 .llm(Arc::new(mock))
19077 .auto_configure_features()
19078 .unwrap()
19079 .auto_configure_spawner()
19080 .await
19081 .unwrap()
19082 .build()
19083 .unwrap();
19084
19085 let available = agent.get_available_tool_ids().await.unwrap();
19086 assert_eq!(available.len(), 4);
19087 assert!(available.contains(&"spawn_agent".to_string()));
19088 assert!(available.contains(&"send_agent_message".to_string()));
19089 assert!(available.contains(&"list_agents".to_string()));
19090 assert!(available.contains(&"remove_agent".to_string()));
19091 }
19092
19093 #[tokio::test]
19094 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
19095 let mock = mock_with_response("hello");
19096 let yaml = r#"
19097name: ManagementOmittedToolsGrantAgent
19098system_prompt: "You manage agents."
19099spawner:
19100 management_tools: true
19101"#;
19102 let agent = AgentBuilder::from_yaml(yaml)
19103 .unwrap()
19104 .llm(Arc::new(mock))
19105 .auto_configure_features()
19106 .unwrap()
19107 .auto_configure_spawner()
19108 .await
19109 .unwrap()
19110 .build()
19111 .unwrap();
19112
19113 let available = agent.get_available_tool_ids().await.unwrap();
19114 assert_eq!(available.len(), 4);
19115 assert!(available.contains(&"spawn_agent".to_string()));
19116 assert!(available.contains(&"send_agent_message".to_string()));
19117 assert!(available.contains(&"list_agents".to_string()));
19118 assert!(available.contains(&"remove_agent".to_string()));
19119 }
19120
19121 #[tokio::test]
19122 async fn test_management_tools_selected_grants_only_selected_tools() {
19123 let mock = mock_with_response("hello");
19124 let yaml = r#"
19125name: ManagementSelectedGrantAgent
19126system_prompt: "You manage agents."
19127tools: []
19128spawner:
19129 management_tools:
19130 - spawn_agent
19131 - send_agent_message
19132 - list_agents
19133"#;
19134 let agent = AgentBuilder::from_yaml(yaml)
19135 .unwrap()
19136 .llm(Arc::new(mock))
19137 .auto_configure_features()
19138 .unwrap()
19139 .auto_configure_spawner()
19140 .await
19141 .unwrap()
19142 .build()
19143 .unwrap();
19144
19145 let available = agent.get_available_tool_ids().await.unwrap();
19146 assert_eq!(available.len(), 3);
19147 assert!(available.contains(&"spawn_agent".to_string()));
19148 assert!(available.contains(&"send_agent_message".to_string()));
19149 assert!(available.contains(&"list_agents".to_string()));
19150 assert!(!available.contains(&"remove_agent".to_string()));
19151 }
19152
19153 #[tokio::test]
19154 async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
19155 let mock = mock_with_response("hello");
19156 let yaml = r#"
19157name: OrchestrationGrantAgent
19158system_prompt: "You coordinate agents."
19159llms:
19160 default:
19161 provider: openai
19162 model: gpt-4
19163 router:
19164 provider: openai
19165 model: gpt-4
19166llm:
19167 default: default
19168 router: router
19169tools: []
19170spawner:
19171 orchestration_tools: true
19172"#;
19173 let agent = AgentBuilder::from_yaml(yaml)
19174 .unwrap()
19175 .llm(Arc::new(mock))
19176 .auto_configure_features()
19177 .unwrap()
19178 .auto_configure_spawner()
19179 .await
19180 .unwrap()
19181 .build()
19182 .unwrap();
19183
19184 let available = agent.get_available_tool_ids().await.unwrap();
19185 assert_eq!(available.len(), 5);
19186 assert!(available.contains(&"route_to_agent".to_string()));
19187 assert!(available.contains(&"pipeline_process".to_string()));
19188 assert!(available.contains(&"concurrent_ask".to_string()));
19189 assert!(available.contains(&"group_discussion".to_string()));
19190 assert!(available.contains(&"handoff_conversation".to_string()));
19191 }
19192
19193 #[tokio::test]
19194 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
19195 let mock = mock_with_response("hello");
19196 let yaml = r#"
19197name: PersonaGrantAgent
19198system_prompt: "You can evolve persona."
19199llm:
19200 provider: openai
19201 model: gpt-4
19202tools: []
19203persona:
19204 identity:
19205 name: "Guide"
19206 role: "Helper"
19207 evolution:
19208 enabled: true
19209 allow_llm_evolve: true
19210 mutable_fields:
19211 - traits.personality
19212"#;
19213 let agent = AgentBuilder::from_yaml(yaml)
19214 .unwrap()
19215 .llm(Arc::new(mock))
19216 .build()
19217 .unwrap();
19218
19219 let available = agent.get_available_tool_ids().await.unwrap();
19220 assert_eq!(available, vec!["persona_evolve".to_string()]);
19221 }
19222
19223 #[tokio::test]
19224 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
19225 let mock = mock_with_response("hello");
19226 let yaml = r#"
19227name: PersonaOmittedToolsGrantAgent
19228system_prompt: "You can evolve persona."
19229llm:
19230 provider: openai
19231 model: gpt-4
19232persona:
19233 identity:
19234 name: "Guide"
19235 role: "Helper"
19236 evolution:
19237 enabled: true
19238 allow_llm_evolve: true
19239 mutable_fields:
19240 - traits.personality
19241"#;
19242 let agent = AgentBuilder::from_yaml(yaml)
19243 .unwrap()
19244 .llm(Arc::new(mock))
19245 .build()
19246 .unwrap();
19247
19248 let available = agent.get_available_tool_ids().await.unwrap();
19249 assert_eq!(available, vec!["persona_evolve".to_string()]);
19250 }
19251
19252 #[tokio::test]
19253 async fn test_omitted_yaml_tools_exposes_no_tools() {
19254 let mock = mock_with_response("hello");
19255 let yaml = r#"
19256name: NoToolsAgent
19257system_prompt: "You are helpful."
19258"#;
19259 let agent = AgentBuilder::from_yaml(yaml)
19260 .unwrap()
19261 .llm(Arc::new(mock))
19262 .auto_configure_features()
19263 .unwrap()
19264 .build()
19265 .unwrap();
19266
19267 let available = agent.get_available_tool_ids().await.unwrap();
19268 assert!(available.is_empty());
19269 }
19270
19271 #[tokio::test]
19272 async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
19273 for tools in ["", "tools: []"] {
19274 let yaml = format!(
19275 r#"
19276name: RuntimeScopeNoGrantAgent
19277system_prompt: "No ordinary tools are granted."
19278{tools}
19279"#
19280 );
19281 let agent = AgentBuilder::from_yaml(&yaml)
19282 .unwrap()
19283 .llm(Arc::new(mock_with_response("done")))
19284 .auto_configure_features()
19285 .unwrap()
19286 .build()
19287 .unwrap();
19288
19289 agent
19290 .runtime_control()
19291 .set_tool_scope(vec!["calculator".to_string()]);
19292
19293 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
19294 }
19295 }
19296
19297 #[tokio::test]
19298 async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
19299 let yaml = r#"
19300name: RuntimeScopeWideningAgent
19301system_prompt: "Runtime scope cannot add authority."
19302tools: [calculator]
19303"#;
19304 let agent = AgentBuilder::from_yaml(yaml)
19305 .unwrap()
19306 .llm(Arc::new(mock_with_response("done")))
19307 .auto_configure_features()
19308 .unwrap()
19309 .build()
19310 .unwrap();
19311
19312 agent
19313 .runtime_control()
19314 .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
19315
19316 assert_eq!(
19317 agent.get_available_tool_ids().await.unwrap(),
19318 vec!["calculator".to_string()]
19319 );
19320 }
19321
19322 #[tokio::test]
19323 async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
19324 let yaml = r#"
19325name: RuntimeScopeIntersectionAgent
19326system_prompt: "Use only declared tools."
19327tools: [calculator, datetime]
19328"#;
19329 let agent = AgentBuilder::from_yaml(yaml)
19330 .unwrap()
19331 .llm(Arc::new(mock_with_response("done")))
19332 .auto_configure_features()
19333 .unwrap()
19334 .build()
19335 .unwrap();
19336 let mut aliases = ai_agents_tools::ToolAliases::default();
19337 aliases
19338 .names
19339 .insert("en".to_string(), "calculate_alias".to_string());
19340 agent.tools.set_tool_aliases("calculator", aliases);
19341 let control = agent.runtime_control();
19342
19343 control.set_tool_scope(vec![
19344 "datetime".to_string(),
19345 "calculate_alias".to_string(),
19346 "calculator".to_string(),
19347 "unknown".to_string(),
19348 "datetime".to_string(),
19349 ]);
19350 assert_eq!(
19351 agent.get_available_tool_ids().await.unwrap(),
19352 vec!["calculator".to_string(), "datetime".to_string()]
19353 );
19354
19355 control.set_tool_scope(vec!["datetime".to_string()]);
19356 assert_eq!(
19357 agent.get_available_tool_ids().await.unwrap(),
19358 vec!["datetime".to_string()]
19359 );
19360
19361 control.clear_tool_scope_override();
19362 assert_eq!(
19363 agent.get_available_tool_ids().await.unwrap(),
19364 vec!["calculator".to_string(), "datetime".to_string()]
19365 );
19366 }
19367
19368 #[tokio::test]
19369 async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
19370 let agent = AgentBuilder::new()
19371 .system_prompt("Use registered tools.")
19372 .llm(Arc::new(mock_with_response("done")))
19373 .tool(Arc::new(ContextEchoTool))
19374 .tool(Arc::new(SlowTool))
19375 .build()
19376 .unwrap();
19377
19378 agent.runtime_control().set_tool_scope(vec![
19379 "Context Echo".to_string(),
19380 "context_echo".to_string(),
19381 "unknown".to_string(),
19382 ]);
19383
19384 assert_eq!(
19385 agent.get_available_tool_ids().await.unwrap(),
19386 vec!["context_echo".to_string()]
19387 );
19388 }
19389
19390 #[tokio::test]
19391 async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
19392 let yaml = r#"
19393name: NestedStateScopeAgent
19394system_prompt: "Honor every state scope."
19395tools: [calculator, datetime, echo]
19396states:
19397 initial: root
19398 states:
19399 root:
19400 tools: [calculate_alias, datetime]
19401 initial: middle
19402 states:
19403 middle:
19404 initial: leaf
19405 states:
19406 leaf:
19407 tools: [datetime_alias, echo]
19408"#;
19409 let agent = AgentBuilder::from_yaml(yaml)
19410 .unwrap()
19411 .llm(Arc::new(mock_with_response("done")))
19412 .auto_configure_features()
19413 .unwrap()
19414 .build()
19415 .unwrap();
19416 let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
19417 calculator_aliases
19418 .names
19419 .insert("en".to_string(), "calculate_alias".to_string());
19420 agent
19421 .tools
19422 .set_tool_aliases("calculator", calculator_aliases);
19423 let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
19424 datetime_aliases
19425 .names
19426 .insert("en".to_string(), "datetime_alias".to_string());
19427 agent.tools.set_tool_aliases("datetime", datetime_aliases);
19428 agent.runtime_control().set_tool_scope(vec![
19429 "unknown".to_string(),
19430 "datetime_alias".to_string(),
19431 "calculate_alias".to_string(),
19432 "datetime".to_string(),
19433 ]);
19434
19435 assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
19436 assert_eq!(
19437 agent.get_available_tool_ids().await.unwrap(),
19438 vec!["datetime".to_string()]
19439 );
19440 }
19441
19442 #[tokio::test]
19443 async fn ancestor_empty_state_scope_denies_omitted_descendants() {
19444 let yaml = r#"
19445name: NestedEmptyStateScopeAgent
19446system_prompt: "An empty ancestor scope denies all tools."
19447tools: [calculator]
19448states:
19449 initial: root
19450 states:
19451 root:
19452 tools: []
19453 initial: middle
19454 states:
19455 middle:
19456 initial: leaf
19457 states:
19458 leaf: {}
19459"#;
19460 let agent = AgentBuilder::from_yaml(yaml)
19461 .unwrap()
19462 .llm(Arc::new(mock_with_response("done")))
19463 .auto_configure_features()
19464 .unwrap()
19465 .build()
19466 .unwrap();
19467
19468 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
19469 }
19470
19471 #[tokio::test]
19472 async fn state_change_during_approval_invalidates_the_reviewed_authority() {
19473 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19474 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19475 let entered = Arc::new(tokio::sync::Barrier::new(2));
19476 let release = Arc::new(tokio::sync::Notify::new());
19477 let handler = Arc::new(BlockingApprovalHandler {
19478 entered: Arc::clone(&entered),
19479 release: Arc::clone(&release),
19480 result: ApprovalResult::Approved,
19481 });
19482 let yaml = r#"
19483name: ApprovalStateGenerationAgent
19484system_prompt: "State authority may change during approval."
19485tools: [locked_write]
19486states:
19487 initial: first
19488 states:
19489 first:
19490 tools: [locked_write]
19491 second:
19492 tools: [locked_write]
19493"#;
19494 let agent = Arc::new(
19495 AgentBuilder::from_yaml(yaml)
19496 .unwrap()
19497 .llm(Arc::new(mock_with_response("done")))
19498 .tool(Arc::new(LockedWriteTool {
19499 active: Arc::clone(&active),
19500 max_active: Arc::clone(&max_active),
19501 }))
19502 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
19503 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
19504 .approval_handler(handler)
19505 .build()
19506 .unwrap(),
19507 );
19508 let running = Arc::clone(&agent);
19509 let call = tokio::spawn(async move {
19510 running
19511 .invoke_tool(ToolExecutionRequest::new(
19512 "approval-state-generation",
19513 "locked_write",
19514 serde_json::json!({"path": "./state-generation.txt"}),
19515 ToolCallSource::Manual,
19516 ))
19517 .await
19518 .unwrap()
19519 });
19520
19521 entered.wait().await;
19522 agent.transition_to("second").await.unwrap();
19523 release.notify_one();
19524 let record = call.await.unwrap();
19525
19526 assert!(!record.executed);
19527 assert!(record.output.contains("Approval became stale"));
19528 assert_eq!(max_active.load(Ordering::SeqCst), 0);
19529 }
19530
19531 #[tokio::test]
19532 async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
19533 let holder_gate = PathMutationGate::new();
19534 let waiter_gate = PathMutationGate::new();
19535 let yaml = r#"
19536name: LockedStateGenerationAgent
19537system_prompt: "State authority must remain stable through admission."
19538tools: [state_lock_holder, state_lock_waiter]
19539states:
19540 initial: first
19541 states:
19542 first:
19543 tools: [state_lock_holder, state_lock_waiter]
19544 second:
19545 tools: [state_lock_holder, state_lock_waiter]
19546"#;
19547 let agent = Arc::new(
19548 AgentBuilder::from_yaml(yaml)
19549 .unwrap()
19550 .llm(Arc::new(mock_with_response("done")))
19551 .tool(Arc::new(BlockingPathMutationTool {
19552 id: "state_lock_holder",
19553 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
19554 gate: holder_gate.clone(),
19555 }))
19556 .tool(Arc::new(BlockingPathMutationTool {
19557 id: "state_lock_waiter",
19558 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
19559 gate: waiter_gate.clone(),
19560 }))
19561 .build()
19562 .unwrap(),
19563 );
19564 let holder_call = {
19565 let agent = Arc::clone(&agent);
19566 tokio::spawn(async move {
19567 agent
19568 .invoke_tool(ToolExecutionRequest::new(
19569 "state-lock-holder",
19570 "state_lock_holder",
19571 serde_json::json!({"path": "./shared-state-path.txt"}),
19572 ToolCallSource::Manual,
19573 ))
19574 .await
19575 .unwrap()
19576 })
19577 };
19578 holder_gate.wait_until_entered().await;
19579 let waiter_call = {
19580 let agent = Arc::clone(&agent);
19581 tokio::spawn(async move {
19582 agent
19583 .invoke_tool(ToolExecutionRequest::new(
19584 "state-lock-waiter",
19585 "state_lock_waiter",
19586 serde_json::json!({"path": "./shared-state-path.txt"}),
19587 ToolCallSource::Manual,
19588 ))
19589 .await
19590 .unwrap()
19591 })
19592 };
19593
19594 wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
19595 agent.transition_to("second").await.unwrap();
19596 holder_gate.release();
19597 let holder_record = holder_call.await.unwrap();
19598 let waiter_record = waiter_call.await.unwrap();
19599
19600 assert!(holder_record.success);
19601 assert!(!waiter_record.executed);
19602 assert!(
19603 waiter_record
19604 .output
19605 .contains("state scope changed before admission")
19606 );
19607 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
19608 }
19609
19610 #[tokio::test]
19611 async fn test_state_tools_cannot_widen_top_level_grant() {
19612 let mock = mock_with_response("hello");
19613 let yaml = r#"
19614name: NarrowToolsAgent
19615system_prompt: "You are helpful."
19616tools:
19617 - calculator
19618states:
19619 initial: current
19620 states:
19621 current:
19622 tools: [datetime]
19623"#;
19624 let agent = AgentBuilder::from_yaml(yaml)
19625 .unwrap()
19626 .llm(Arc::new(mock))
19627 .auto_configure_features()
19628 .unwrap()
19629 .build()
19630 .unwrap();
19631
19632 let available = agent.get_available_tool_ids().await.unwrap();
19633 assert!(available.is_empty());
19634 }
19635
19636 #[tokio::test]
19638 async fn test_integration_tool_execution() {
19639 let mock = mock_with_responses(vec![
19641 r#"I'll calculate that for you.
19643[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
19644 "The answer is 4.",
19646 ]);
19647 let mut tools = ai_agents_tools::ToolRegistry::new();
19648 tools
19649 .register(Arc::new(ai_agents_tools::CalculatorTool))
19650 .unwrap();
19651
19652 let agent = AgentBuilder::new()
19653 .system_prompt("You are a calculator assistant.")
19654 .llm(Arc::new(mock))
19655 .tools(tools)
19656 .build()
19657 .unwrap();
19658
19659 let response = agent.chat("What is 2+2?").await.unwrap();
19660 assert!(!response.content.is_empty());
19662 }
19663
19664 #[tokio::test]
19665 async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
19666 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19667 let hooks = Arc::new(ResponseCountingHooks {
19668 responses: Arc::clone(&responses),
19669 });
19670 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
19671 let yaml = r#"
19672name: ToolRejectAgent
19673system_prompt: "You use tools when requested."
19674tools:
19675 - echo
19676hitl:
19677 tools:
19678 echo:
19679 require_approval: true
19680 approval_message: "Approve echo?"
19681"#;
19682 let agent = AgentBuilder::from_yaml(yaml)
19683 .unwrap()
19684 .llm(Arc::new(mock))
19685 .auto_configure_features()
19686 .unwrap()
19687 .hooks(hooks)
19688 .build()
19689 .unwrap();
19690
19691 let response = agent.chat("echo hello").await.unwrap();
19692
19693 assert!(
19694 response.content.contains("Operation cancelled"),
19695 "unexpected response: {}",
19696 response.content
19697 );
19698 assert_eq!(responses.load(Ordering::SeqCst), 1);
19699 let messages = agent.memory.get_messages(None).await.unwrap();
19700 assert_eq!(messages.len(), 3);
19701 assert_eq!(messages[0].content, "echo hello");
19702 assert!(messages[1].content.contains("\"tool\":\"echo\""));
19703 assert!(messages[2].content.contains("rejected by the approver"));
19704 }
19705
19706 #[tokio::test]
19707 async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
19708 use futures::StreamExt;
19709
19710 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
19711 let hooks = Arc::new(ResponseCountingHooks {
19712 responses: Arc::clone(&responses),
19713 });
19714 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
19715 let yaml = r#"
19716name: ToolRejectStreamingAgent
19717system_prompt: "You use tools when requested."
19718tools:
19719 - echo
19720streaming:
19721 enabled: true
19722hitl:
19723 tools:
19724 echo:
19725 require_approval: true
19726 approval_message: "Approve echo?"
19727"#;
19728 let agent = AgentBuilder::from_yaml(yaml)
19729 .unwrap()
19730 .llm(Arc::new(mock))
19731 .auto_configure_features()
19732 .unwrap()
19733 .hooks(hooks)
19734 .build()
19735 .unwrap();
19736
19737 let mut stream = agent.chat_stream("echo hello").await.unwrap();
19738 let mut terminal_error = String::new();
19739 let mut done = false;
19740 while let Some(chunk) = stream.next().await {
19741 match chunk {
19742 StreamChunk::Error { message } => terminal_error = message,
19743 StreamChunk::Done {} => {
19744 done = true;
19745 break;
19746 }
19747 _ => {}
19748 }
19749 }
19750
19751 assert!(done);
19752 assert!(
19753 terminal_error.contains("Operation cancelled"),
19754 "unexpected terminal error: {}",
19755 terminal_error
19756 );
19757 assert_eq!(responses.load(Ordering::SeqCst), 1);
19758 let messages = agent.memory.get_messages(None).await.unwrap();
19759 assert_eq!(messages.len(), 3);
19760 assert_eq!(messages[0].content, "echo hello");
19761 assert!(messages[1].content.contains("\"tool\":\"echo\""));
19762 assert!(messages[2].content.contains("rejected by the approver"));
19763 }
19764
19765 #[tokio::test]
19766 async fn tool_hitl_rejection_preserves_legacy_error_but_finalizes_event_stream() {
19767 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
19768 let yaml = r#"
19769name: ToolRejectEventAgent
19770system_prompt: "You use tools when requested."
19771tools:
19772 - echo
19773streaming:
19774 enabled: true
19775hitl:
19776 tools:
19777 echo:
19778 require_approval: true
19779 approval_message: "Approve echo?"
19780"#;
19781 let agent = AgentBuilder::from_yaml(yaml)
19782 .unwrap()
19783 .llm(Arc::new(mock))
19784 .auto_configure_features()
19785 .unwrap()
19786 .build()
19787 .unwrap();
19788
19789 let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
19790 let mut error_seen = false;
19791 let mut final_response = None;
19792 while let Some(event) = stream.next().await {
19793 match event {
19794 AgentStreamEvent::Chunk(StreamChunk::Error { .. }) => error_seen = true,
19795 AgentStreamEvent::Final(response) => final_response = Some(response),
19796 AgentStreamEvent::Chunk(_) => {}
19797 }
19798 }
19799
19800 assert!(!error_seen);
19801 assert!(
19802 final_response
19803 .is_some_and(|response| { response.content.contains("Operation cancelled") })
19804 );
19805 }
19806
19807 #[tokio::test]
19808 async fn test_pre_response_guard_transition_skips_old_state_llm() {
19809 let mock = mock_with_response("Billing state response");
19810 let call_counter = mock.clone();
19811 let yaml = r#"
19812name: OptimizedStateAgent
19813system_prompt: "You route before answering."
19814runtime:
19815 optimization:
19816 enabled: true
19817 pre_response_deterministic_transitions: true
19818states:
19819 initial: greeting
19820 states:
19821 greeting:
19822 prompt: "Old state prompt that should be skipped."
19823 transitions:
19824 - to: billing
19825 guard:
19826 context:
19827 topic:
19828 eq: billing
19829 timing: pre_response
19830 billing:
19831 prompt: "Answer from the billing state."
19832"#;
19833 let agent = AgentBuilder::from_yaml(yaml)
19834 .unwrap()
19835 .llm(Arc::new(mock))
19836 .build()
19837 .unwrap();
19838 agent
19839 .set_context("topic", serde_json::json!("billing"))
19840 .unwrap();
19841
19842 let response = agent.chat("I need billing help").await.unwrap();
19843
19844 assert_eq!(agent.current_state().as_deref(), Some("billing"));
19845 assert_eq!(response.content, "Billing state response");
19846 assert_eq!(call_counter.call_count(), 1);
19847 assert_eq!(agent.actor_facts().len(), 0);
19848 }
19849
19850 #[tokio::test]
19851 async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
19852 let mock = mock_with_response("Billing state response");
19853 let call_counter = mock.clone();
19854 let yaml = r#"
19855name: OptimizedStateAgent
19856system_prompt: "You route before answering."
19857runtime:
19858 optimization:
19859 enabled: true
19860 pre_response_deterministic_transitions: true
19861context:
19862 request:
19863 type: runtime
19864 default:
19865 topic: general
19866states:
19867 initial: greeting
19868 states:
19869 greeting:
19870 prompt: "Old state prompt that should be skipped."
19871 transitions:
19872 - to: billing
19873 guard:
19874 context:
19875 request.topic:
19876 eq: billing
19877 timing: pre_response
19878 billing:
19879 prompt: "Answer from the billing state."
19880"#;
19881 let agent = AgentBuilder::from_yaml(yaml)
19882 .unwrap()
19883 .llm(Arc::new(mock))
19884 .build()
19885 .unwrap();
19886 agent
19887 .set_context("request.topic", serde_json::json!("billing"))
19888 .unwrap();
19889
19890 let response = agent.chat("I need billing help").await.unwrap();
19891
19892 assert_eq!(agent.current_state().as_deref(), Some("billing"));
19893 assert_eq!(response.content, "Billing state response");
19894 assert_eq!(call_counter.call_count(), 1);
19895 assert_eq!(
19896 agent.get_context().get("request"),
19897 Some(&serde_json::json!({"topic": "billing"}))
19898 );
19899 }
19900
19901 #[tokio::test]
19902 async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
19903 let mock = mock_with_response("billing");
19904 let yaml = r#"
19905name: OptimizedStateAgent
19906system_prompt: "You route before answering."
19907runtime:
19908 optimization:
19909 enabled: true
19910 pre_response_deterministic_transitions: true
19911hitl:
19912 states:
19913 billing:
19914 on_enter: require_approval
19915 approval_message: "Approve billing route?"
19916states:
19917 initial: greeting
19918 states:
19919 greeting:
19920 prompt: "Old state prompt."
19921 extract:
19922 - key: topic
19923 description: "Support topic"
19924 transitions:
19925 - to: billing
19926 guard:
19927 context:
19928 topic:
19929 eq: billing
19930 timing: pre_response
19931 run_extractors: true
19932 billing:
19933 prompt: "Billing state."
19934"#;
19935 let agent = AgentBuilder::from_yaml(yaml)
19936 .unwrap()
19937 .llm(Arc::new(mock))
19938 .build()
19939 .unwrap();
19940
19941 let response = agent
19942 .try_pre_response_transition("billing please")
19943 .await
19944 .unwrap();
19945
19946 assert!(response.is_none());
19947 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
19948 assert!(!agent.get_context().contains_key("topic"));
19949 assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
19950 }
19951
19952 #[tokio::test]
19953 async fn test_pre_response_extractor_commits_context_on_winning_path() {
19954 let mock = mock_with_responses(vec!["billing", "Billing response"]);
19955 let yaml = r#"
19956name: OptimizedStateAgent
19957system_prompt: "You route before answering."
19958runtime:
19959 optimization:
19960 enabled: true
19961 pre_response_deterministic_transitions: true
19962states:
19963 initial: greeting
19964 states:
19965 greeting:
19966 prompt: "Old state prompt."
19967 extract:
19968 - key: topic
19969 description: "Support topic"
19970 transitions:
19971 - to: billing
19972 guard:
19973 context:
19974 topic:
19975 eq: billing
19976 timing: pre_response
19977 run_extractors: true
19978 billing:
19979 prompt: "Billing state."
19980"#;
19981 let agent = AgentBuilder::from_yaml(yaml)
19982 .unwrap()
19983 .llm(Arc::new(mock))
19984 .build()
19985 .unwrap();
19986
19987 let response = agent.chat("billing please").await.unwrap();
19988
19989 assert_eq!(agent.current_state().as_deref(), Some("billing"));
19990 assert_eq!(response.content, "Billing response");
19991 assert_eq!(
19992 agent.get_context().get("topic"),
19993 Some(&serde_json::json!("billing"))
19994 );
19995 }
19996
19997 #[tokio::test]
19998 async fn test_pre_response_extractor_miss_does_not_mutate_context() {
19999 let mock = mock_with_response("__NONE__");
20000 let yaml = r#"
20001name: OptimizedStateAgent
20002system_prompt: "You route before answering."
20003runtime:
20004 optimization:
20005 enabled: true
20006 pre_response_deterministic_transitions: true
20007states:
20008 initial: greeting
20009 states:
20010 greeting:
20011 prompt: "Old state prompt."
20012 extract:
20013 - key: topic
20014 description: "Support topic"
20015 transitions:
20016 - to: billing
20017 guard:
20018 context:
20019 topic:
20020 eq: billing
20021 timing: pre_response
20022 run_extractors: true
20023 billing:
20024 prompt: "Billing state."
20025"#;
20026 let agent = AgentBuilder::from_yaml(yaml)
20027 .unwrap()
20028 .llm(Arc::new(mock))
20029 .build()
20030 .unwrap();
20031
20032 let response = agent.try_pre_response_transition("hello").await.unwrap();
20033
20034 assert!(response.is_none());
20035 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
20036 assert!(!agent.get_context().contains_key("topic"));
20037 }
20038
20039 #[tokio::test]
20040 async fn test_default_guard_transition_stays_post_response() {
20041 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
20042 let call_counter = mock.clone();
20043 let yaml = r#"
20044name: TimingAgent
20045system_prompt: "You route carefully."
20046runtime:
20047 optimization:
20048 enabled: true
20049 pre_response_deterministic_transitions: true
20050states:
20051 initial: greeting
20052 states:
20053 greeting:
20054 prompt: "Old state prompt."
20055 transitions:
20056 - to: billing
20057 guard:
20058 context:
20059 topic:
20060 eq: billing
20061 billing:
20062 prompt: "Billing state."
20063"#;
20064 let agent = AgentBuilder::from_yaml(yaml)
20065 .unwrap()
20066 .llm(Arc::new(mock))
20067 .build()
20068 .unwrap();
20069 agent
20070 .set_context("topic", serde_json::json!("billing"))
20071 .unwrap();
20072
20073 let response = agent.chat("billing please").await.unwrap();
20074
20075 assert_eq!(agent.current_state().as_deref(), Some("billing"));
20076 assert_eq!(response.content, "Billing response");
20077 assert_eq!(call_counter.call_count(), 2);
20078 }
20079
20080 #[tokio::test]
20081 async fn test_explicit_post_response_guard_transition_stays_post_response() {
20082 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
20083 let call_counter = mock.clone();
20084 let yaml = r#"
20085name: TimingAgent
20086system_prompt: "You route carefully."
20087runtime:
20088 optimization:
20089 enabled: true
20090 pre_response_deterministic_transitions: true
20091states:
20092 initial: greeting
20093 states:
20094 greeting:
20095 prompt: "Old state prompt."
20096 transitions:
20097 - to: billing
20098 guard:
20099 context:
20100 topic:
20101 eq: billing
20102 timing: post_response
20103 billing:
20104 prompt: "Billing state."
20105"#;
20106 let agent = AgentBuilder::from_yaml(yaml)
20107 .unwrap()
20108 .llm(Arc::new(mock))
20109 .build()
20110 .unwrap();
20111 agent
20112 .set_context("topic", serde_json::json!("billing"))
20113 .unwrap();
20114
20115 let response = agent.chat("billing please").await.unwrap();
20116
20117 assert_eq!(agent.current_state().as_deref(), Some("billing"));
20118 assert_eq!(response.content, "Billing response");
20119 assert_eq!(call_counter.call_count(), 2);
20120 }
20121
20122 #[tokio::test]
20123 async fn test_pre_response_extractors_are_transition_scoped() {
20124 let mock = mock_with_responses(vec!["billing", "Billing response"]);
20125 let yaml = r#"
20126name: ScopedExtractorAgent
20127system_prompt: "You route carefully."
20128runtime:
20129 optimization:
20130 enabled: true
20131 pre_response_deterministic_transitions: true
20132states:
20133 initial: greeting
20134 states:
20135 greeting:
20136 prompt: "Old state prompt."
20137 extract:
20138 - key: topic
20139 description: "Support topic"
20140 transitions:
20141 - to: wrong
20142 guard:
20143 context:
20144 topic:
20145 eq: billing
20146 timing: pre_response
20147 - to: billing
20148 guard:
20149 context:
20150 topic:
20151 eq: billing
20152 timing: pre_response
20153 run_extractors: true
20154 wrong:
20155 prompt: "Wrong state."
20156 billing:
20157 prompt: "Billing state."
20158"#;
20159 let agent = AgentBuilder::from_yaml(yaml)
20160 .unwrap()
20161 .llm(Arc::new(mock))
20162 .build()
20163 .unwrap();
20164
20165 let response = agent.chat("billing please").await.unwrap();
20166
20167 assert_eq!(agent.current_state().as_deref(), Some("billing"));
20168 assert_eq!(response.content, "Billing response");
20169 }
20170
20171 #[tokio::test]
20172 async fn test_pre_response_resolved_intent_routes_early() {
20173 let mock = mock_with_response("Billing response");
20174 let yaml = r#"
20175name: IntentAgent
20176system_prompt: "You route carefully."
20177runtime:
20178 optimization:
20179 enabled: true
20180 pre_response_deterministic_transitions: true
20181states:
20182 initial: greeting
20183 states:
20184 greeting:
20185 prompt: "Old state prompt."
20186 transitions:
20187 - to: billing
20188 intent: billing
20189 timing: pre_response
20190 billing:
20191 prompt: "Billing state."
20192"#;
20193 let agent = AgentBuilder::from_yaml(yaml)
20194 .unwrap()
20195 .llm(Arc::new(mock))
20196 .build()
20197 .unwrap();
20198 agent
20199 .set_context("resolved_intent", serde_json::json!("billing"))
20200 .unwrap();
20201
20202 let response = agent
20203 .try_pre_response_transition("I need billing help")
20204 .await
20205 .unwrap()
20206 .unwrap();
20207
20208 assert_eq!(agent.current_state().as_deref(), Some("billing"));
20209 assert_eq!(response.content, "Billing response");
20210 }
20211
20212 #[tokio::test]
20213 async fn test_background_overflow_error_surfaces() {
20214 let mut config = RuntimeConfig::default();
20215 config.optimization.enabled = true;
20216 config.optimization.post_turn.max_background_tasks = 1;
20217 config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
20218 let policy = crate::optimization::MaintenanceTaskPolicy {
20219 mode: MaintenanceMode::Background,
20220 await_before_next_turn: AwaitBeforeNextTurn::Always,
20221 };
20222 let agent = AgentBuilder::new()
20223 .system_prompt("You are helpful.")
20224 .llm(Arc::new(mock_with_response("ok")))
20225 .build()
20226 .unwrap()
20227 .with_runtime_config(config);
20228 agent
20229 .background_maintenance
20230 .spawn(None, async { std::future::pending::<Result<()>>().await })
20231 .unwrap();
20232
20233 let result = agent
20234 .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
20235 .await;
20236
20237 assert!(result.is_err());
20238 }
20239
20240 #[tokio::test]
20241 async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
20242 let default_mock = mock_with_response("Plain draft response");
20243 let router_mock = mock_with_response("cot");
20244 let router_counter = router_mock.clone();
20245 let yaml = r#"
20246name: ReasoningReservationAgent
20247system_prompt: "You answer plainly unless reasoning wins."
20248llm:
20249 default: default
20250 router: router
20251observability:
20252 enabled: true
20253 export:
20254 write_raw_events: true
20255reasoning:
20256 mode: auto
20257 judge_llm: router
20258runtime:
20259 optimization:
20260 enabled: true
20261 max_speculative_llm_calls_per_turn: 1
20262 speculative_reasoning_auto: true
20263 max_parallel_runtime_tasks: 2
20264"#;
20265 let agent = AgentBuilder::from_yaml(yaml)
20266 .unwrap()
20267 .llm_alias("default", Arc::new(default_mock))
20268 .llm_alias("router", Arc::new(router_mock))
20269 .build()
20270 .unwrap();
20271
20272 let response = agent.chat("hello").await.unwrap();
20273
20274 assert_eq!(response.content, "Plain draft response");
20275 assert_eq!(router_counter.call_count(), 1);
20276 let events = agent.observability().unwrap().raw_events();
20277 assert!(!events.iter().any(|event| {
20278 event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
20279 }));
20280 }
20281
20282 #[tokio::test]
20283 async fn test_forced_reasoning_skips_plain_speculative_draft() {
20284 let mock = mock_with_response("Reasoned response");
20285 let yaml = r#"
20286name: ForcedReasoningAgent
20287system_prompt: "You reason before answering."
20288observability:
20289 enabled: true
20290 export:
20291 write_raw_events: true
20292reasoning:
20293 mode: cot
20294runtime:
20295 optimization:
20296 enabled: true
20297 max_speculative_llm_calls_per_turn: 2
20298 speculative_state_transitions: true
20299 max_parallel_runtime_tasks: 2
20300states:
20301 initial: triage
20302 states:
20303 triage:
20304 prompt: "Answer from triage."
20305 transitions:
20306 - to: billing
20307 guard:
20308 context:
20309 route:
20310 eq: billing
20311 timing: parallel
20312 billing:
20313 prompt: "Billing state."
20314"#;
20315 let agent = AgentBuilder::from_yaml(yaml)
20316 .unwrap()
20317 .llm(Arc::new(mock))
20318 .build()
20319 .unwrap();
20320
20321 let response = agent.chat("hello").await.unwrap();
20322
20323 assert_eq!(response.content, "Reasoned response");
20324 let events = agent.observability().unwrap().raw_events();
20325 assert!(
20326 !events
20327 .iter()
20328 .any(|event| event.dimensions.contains_key("branch_status"))
20329 );
20330 }
20331
20332 #[tokio::test]
20333 async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
20334 let default_mock = mock_with_response("Skill committed response");
20335 let router_mock = mock_with_response("helper");
20336 let router_counter = router_mock.clone();
20337 let yaml = r#"
20338name: SkillReservationAgent
20339system_prompt: "Use skills when they match."
20340llm:
20341 default: default
20342 router: router
20343observability:
20344 enabled: true
20345 export:
20346 write_raw_events: true
20347runtime:
20348 optimization:
20349 enabled: true
20350 max_speculative_llm_calls_per_turn: 1
20351 speculative_skill_routing: true
20352 max_parallel_runtime_tasks: 2
20353skills:
20354 - id: helper
20355 description: "Answer helper requests"
20356 trigger: "User asks for helper"
20357 steps:
20358 - prompt: "Answer the helper request: {{ user_input }}"
20359"#;
20360 let agent = AgentBuilder::from_yaml(yaml)
20361 .unwrap()
20362 .llm_alias("default", Arc::new(default_mock))
20363 .llm_alias("router", Arc::new(router_mock))
20364 .build()
20365 .unwrap();
20366
20367 let response = agent.chat("please use helper").await.unwrap();
20368
20369 assert_eq!(response.content, "Skill committed response");
20370 assert_eq!(router_counter.call_count(), 1);
20371 let events = agent.observability().unwrap().raw_events();
20372 assert!(
20373 !events
20374 .iter()
20375 .any(|event| event.dimensions.contains_key("branch_status"))
20376 );
20377 }
20378
20379 #[tokio::test]
20380 async fn test_parallel_transition_low_cap_allows_deterministic_route() {
20381 let mock = mock_with_response("unused");
20382 let call_counter = mock.clone();
20383 let yaml = r#"
20384name: ParallelTransitionLowCapAgent
20385system_prompt: "Route before stale responses when safe."
20386runtime:
20387 optimization:
20388 enabled: true
20389 max_speculative_llm_calls_per_turn: 1
20390 speculative_state_transitions: true
20391 max_parallel_runtime_tasks: 2
20392states:
20393 initial: triage
20394 states:
20395 triage:
20396 prompt: "Triage state."
20397 transitions:
20398 - to: billing
20399 guard:
20400 context:
20401 route:
20402 eq: billing
20403 timing: parallel
20404 billing:
20405 prompt: "Billing state."
20406"#;
20407 let agent = AgentBuilder::from_yaml(yaml)
20408 .unwrap()
20409 .llm(Arc::new(mock))
20410 .build()
20411 .unwrap();
20412 agent
20413 .set_context("route", serde_json::json!("billing"))
20414 .unwrap();
20415 agent.update_active_turn_context("billing help", HashMap::new());
20416 assert!(
20417 agent.reserve_active_speculative_llm_call(
20418 RuntimeOptimizationKind::ParallelStateTransition
20419 )
20420 );
20421
20422 let selection = agent
20423 .select_parallel_transition_candidate("billing help")
20424 .await
20425 .unwrap();
20426 agent.end_root_turn();
20427
20428 match selection {
20429 ParallelTransitionSelection::Candidate(candidate) => {
20430 assert_eq!(candidate.target(), "billing");
20431 }
20432 ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
20433 ParallelTransitionSelection::ReservationExhausted => {
20434 panic!("deterministic route consumed LLM budget")
20435 }
20436 }
20437 assert_eq!(call_counter.call_count(), 0);
20438 }
20439
20440 #[tokio::test]
20441 async fn speculative_transition_drops_loser_before_state_actions() {
20442 let lock = Arc::new(tokio::sync::Mutex::new(()));
20443 let first_started = Arc::new(tokio::sync::Notify::new());
20444 let first_dropped = Arc::new(AtomicBool::new(false));
20445 let committed_after_drop = Arc::new(AtomicBool::new(false));
20446 let default = Arc::new(FirstCallLockingProvider {
20447 lock,
20448 first_started: Arc::clone(&first_started),
20449 first_dropped: Arc::clone(&first_dropped),
20450 committed_after_drop: Arc::clone(&committed_after_drop),
20451 calls: AtomicU64::new(0),
20452 });
20453 let router = Arc::new(RoutingAfterProviderStart {
20454 provider_started: first_started,
20455 });
20456 let yaml = r#"
20457name: SpeculativeCancellationAgent
20458system_prompt: "Route before committed work."
20459llm:
20460 default: default
20461 router: router
20462runtime:
20463 optimization:
20464 enabled: true
20465 max_speculative_llm_calls_per_turn: 2
20466 speculative_state_transitions: true
20467 max_parallel_runtime_tasks: 2
20468states:
20469 initial: triage
20470 states:
20471 triage:
20472 prompt: "Triage state."
20473 transitions:
20474 - to: technical
20475 when: "The request needs technical support"
20476 timing: parallel
20477 technical:
20478 prompt: "Technical state."
20479 on_enter:
20480 - prompt: "Prepare technical context."
20481 llm: default
20482 store_as: preparation
20483"#;
20484 let agent = AgentBuilder::from_yaml(yaml)
20485 .unwrap()
20486 .llm_alias("default", default)
20487 .llm_alias("router", router)
20488 .build()
20489 .unwrap();
20490
20491 let response = tokio::time::timeout(
20492 std::time::Duration::from_secs(2),
20493 agent.chat("I cannot log in because of AUTH-17."),
20494 )
20495 .await
20496 .expect("committed work must not wait on the losing provider future")
20497 .unwrap();
20498
20499 assert_eq!(response.content, "Committed technical response.");
20500 assert_eq!(agent.current_state().as_deref(), Some("technical"));
20501 assert!(first_dropped.load(Ordering::SeqCst));
20502 assert!(committed_after_drop.load(Ordering::SeqCst));
20503 }
20504
20505 #[tokio::test]
20506 async fn buffered_transition_drops_stale_stream_before_redispatch() {
20507 use futures::StreamExt;
20508
20509 let lock = Arc::new(tokio::sync::Mutex::new(()));
20510 let stream_started = Arc::new(tokio::sync::Notify::new());
20511 let stream_dropped = Arc::new(AtomicBool::new(false));
20512 let committed_after_drop = Arc::new(AtomicBool::new(false));
20513 let default = Arc::new(BufferedLockingProvider {
20514 lock,
20515 stream_started: Arc::clone(&stream_started),
20516 stream_dropped: Arc::clone(&stream_dropped),
20517 committed_after_drop: Arc::clone(&committed_after_drop),
20518 });
20519 let router = Arc::new(RoutingAfterProviderStart {
20520 provider_started: stream_started,
20521 });
20522 let yaml = r#"
20523name: BufferedCancellationAgent
20524system_prompt: "Hide stale streamed output."
20525llm:
20526 default: default
20527 router: router
20528streaming:
20529 enabled: true
20530 buffer_size: 8
20531runtime:
20532 optimization:
20533 enabled: true
20534 max_speculative_llm_calls_per_turn: 2
20535 speculative_state_transitions: true
20536 streaming_policy: buffer_until_routing_done
20537 max_parallel_runtime_tasks: 2
20538states:
20539 initial: triage
20540 states:
20541 triage:
20542 prompt: "Triage state."
20543 transitions:
20544 - to: technical
20545 when: "The request needs technical support"
20546 timing: parallel
20547 technical:
20548 prompt: "Technical state."
20549"#;
20550 let agent = AgentBuilder::from_yaml(yaml)
20551 .unwrap()
20552 .llm_alias("default", default)
20553 .llm_alias("router", router)
20554 .build()
20555 .unwrap();
20556
20557 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
20558 let mut stream = agent
20559 .chat_stream("AUTH-17 needs technical help.")
20560 .await
20561 .unwrap();
20562 let mut content = String::new();
20563 while let Some(chunk) = stream.next().await {
20564 match chunk {
20565 StreamChunk::Content { text } => content.push_str(&text),
20566 StreamChunk::Done {} => break,
20567 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
20568 _ => {}
20569 }
20570 }
20571 content
20572 })
20573 .await
20574 .expect("redispatch must not wait on the stale streaming future");
20575
20576 assert_eq!(content, "Committed technical response.");
20577 assert_eq!(agent.current_state().as_deref(), Some("technical"));
20578 assert!(stream_dropped.load(Ordering::SeqCst));
20579 assert!(committed_after_drop.load(Ordering::SeqCst));
20580 }
20581
20582 #[tokio::test]
20583 async fn buffered_transition_drops_established_stream_before_redispatch() {
20584 use futures::StreamExt;
20585
20586 let stream_started = Arc::new(tokio::sync::Notify::new());
20587 let stream_dropped = Arc::new(AtomicBool::new(false));
20588 let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
20589 let committed_after_drop = Arc::new(AtomicBool::new(false));
20590 let default = Arc::new(EstablishedStreamProvider {
20591 stream_started: Arc::clone(&stream_started),
20592 stream_dropped: Arc::clone(&stream_dropped),
20593 stream_dropped_notify,
20594 committed_after_drop: Arc::clone(&committed_after_drop),
20595 });
20596 let router = Arc::new(RoutingAfterProviderStart {
20597 provider_started: stream_started,
20598 });
20599 let yaml = r#"
20600name: EstablishedStreamCancellationAgent
20601system_prompt: "Hide stale streamed output."
20602llm:
20603 default: default
20604 router: router
20605streaming:
20606 enabled: true
20607 buffer_size: 8
20608runtime:
20609 optimization:
20610 enabled: true
20611 max_speculative_llm_calls_per_turn: 2
20612 speculative_state_transitions: true
20613 streaming_policy: buffer_until_routing_done
20614 max_parallel_runtime_tasks: 2
20615states:
20616 initial: triage
20617 states:
20618 triage:
20619 prompt: "Triage state."
20620 transitions:
20621 - to: technical
20622 when: "The request needs technical support"
20623 timing: parallel
20624 technical:
20625 prompt: "Technical state."
20626"#;
20627 let agent = AgentBuilder::from_yaml(yaml)
20628 .unwrap()
20629 .llm_alias("default", default)
20630 .llm_alias("router", router)
20631 .build()
20632 .unwrap();
20633
20634 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
20635 let mut stream = agent
20636 .chat_stream("AUTH-17 needs technical help.")
20637 .await
20638 .unwrap();
20639 let mut content = String::new();
20640 while let Some(chunk) = stream.next().await {
20641 match chunk {
20642 StreamChunk::Content { text } => content.push_str(&text),
20643 StreamChunk::Done {} => break,
20644 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
20645 _ => {}
20646 }
20647 }
20648 content
20649 })
20650 .await
20651 .expect("redispatch must wait for the established stale stream to be dropped");
20652
20653 assert_eq!(content, "Committed technical response.");
20654 assert_eq!(agent.current_state().as_deref(), Some("technical"));
20655 assert!(stream_dropped.load(Ordering::SeqCst));
20656 assert!(committed_after_drop.load(Ordering::SeqCst));
20657 }
20658
20659 #[tokio::test]
20660 async fn test_buffered_streaming_transition_reservation_falls_back() {
20661 use futures::StreamExt;
20662
20663 let mock = mock_with_responses(vec![
20664 "Serial streaming response",
20665 "Serial streaming response",
20666 ]);
20667 let router_mock = mock_with_response("1");
20668 let router_counter = router_mock.clone();
20669 let yaml = r#"
20670name: BufferedReservationFallbackAgent
20671system_prompt: "Stream normally if speculative routing cannot be evaluated."
20672llm:
20673 default: default
20674 router: router
20675observability:
20676 enabled: true
20677 export:
20678 write_raw_events: true
20679streaming:
20680 enabled: true
20681 buffer_size: 8
20682runtime:
20683 optimization:
20684 enabled: true
20685 max_speculative_llm_calls_per_turn: 1
20686 speculative_state_transitions: true
20687 streaming_policy: buffer_until_routing_done
20688 max_parallel_runtime_tasks: 2
20689states:
20690 initial: triage
20691 states:
20692 triage:
20693 prompt: "Triage state."
20694 transitions:
20695 - to: billing
20696 guard:
20697 context:
20698 route:
20699 eq: billing
20700 when: "User asks about billing"
20701 timing: parallel
20702 billing:
20703 prompt: "Billing state."
20704"#;
20705 let agent = AgentBuilder::from_yaml(yaml)
20706 .unwrap()
20707 .llm_alias("default", Arc::new(mock))
20708 .llm_alias("router", Arc::new(router_mock))
20709 .build()
20710 .unwrap();
20711
20712 let mut stream = agent.chat_stream("hello").await.unwrap();
20713 let mut content = String::new();
20714 let mut error = None;
20715 while let Some(chunk) = stream.next().await {
20716 match chunk {
20717 StreamChunk::Content { text } => content.push_str(&text),
20718 StreamChunk::Error { message } => error = Some(message),
20719 StreamChunk::Done {} => break,
20720 _ => {}
20721 }
20722 }
20723
20724 assert_eq!(error, None);
20725 assert_eq!(content, "Serial streaming response");
20726 assert_eq!(router_counter.call_count(), 0);
20727 let events = agent.observability().unwrap().raw_events();
20728 assert!(events.iter().any(|event| {
20729 event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
20730 && event.dimensions.get("commit_behavior")
20731 == Some(&"transition_decision".to_string())
20732 }));
20733 }
20734
20735 #[tokio::test]
20736 async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
20737 let mut mock = mock_with_response("Recovered response");
20738 mock.set_error("boom");
20739 let mut handle = mock.clone();
20740 let agent = AgentBuilder::new()
20741 .system_prompt("You are helpful.")
20742 .llm(Arc::new(mock))
20743 .build()
20744 .unwrap();
20745
20746 assert!(agent.chat("first").await.is_err());
20747 handle.clear_error();
20748 let response = agent.chat("second").await.unwrap();
20749
20750 assert_eq!(response.content, "Recovered response");
20751 let messages = agent.memory.get_messages(None).await.unwrap();
20752 let user_count = messages
20753 .iter()
20754 .filter(|message| message.role == ai_agents_core::Role::User)
20755 .count();
20756 assert_eq!(user_count, 2);
20757 }
20758
20759 #[tokio::test]
20760 async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
20761 use futures::StreamExt;
20762
20763 let mut mock = mock_with_response("Recovered response");
20764 mock.set_error("stream boom");
20765 let mut handle = mock.clone();
20766 let agent = AgentBuilder::new()
20767 .system_prompt("You are helpful.")
20768 .llm(Arc::new(mock))
20769 .build()
20770 .unwrap();
20771
20772 let mut stream = agent.chat_stream("first").await.unwrap();
20773 let mut saw_error = false;
20774 while let Some(chunk) = stream.next().await {
20775 if matches!(chunk, StreamChunk::Error { .. }) {
20776 saw_error = true;
20777 }
20778 }
20779 assert!(saw_error);
20780
20781 handle.clear_error();
20782 let response = agent.chat("second").await.unwrap();
20783
20784 assert_eq!(response.content, "Recovered response");
20785 let messages = agent.memory.get_messages(None).await.unwrap();
20786 let user_count = messages
20787 .iter()
20788 .filter(|message| message.role == ai_agents_core::Role::User)
20789 .count();
20790 assert_eq!(user_count, 2);
20791 }
20792
20793 #[tokio::test]
20794 async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
20795 use futures::StreamExt;
20796
20797 let mut mock = mock_with_response("one two three");
20798 mock.set_latency(10);
20799 let yaml = r#"
20800name: BufferedMissAgent
20801system_prompt: "You stream safely."
20802llm:
20803 default: default
20804streaming:
20805 enabled: true
20806 buffer_size: 1
20807runtime:
20808 optimization:
20809 enabled: true
20810 max_speculative_llm_calls_per_turn: 2
20811 speculative_state_transitions: true
20812 streaming_policy: buffer_until_routing_done
20813 max_parallel_runtime_tasks: 2
20814states:
20815 initial: triage
20816 states:
20817 triage:
20818 prompt: "Answer from triage."
20819 transitions:
20820 - to: billing
20821 guard:
20822 context:
20823 route:
20824 eq: billing
20825 timing: parallel
20826 billing:
20827 prompt: "Billing state."
20828"#;
20829 let agent = AgentBuilder::from_yaml(yaml)
20830 .unwrap()
20831 .llm_alias("default", Arc::new(mock))
20832 .build()
20833 .unwrap();
20834
20835 let mut stream = agent.chat_stream("hello").await.unwrap();
20836 let mut content = String::new();
20837 let mut error = None;
20838 while let Some(chunk) = stream.next().await {
20839 match chunk {
20840 StreamChunk::Content { text } => content.push_str(&text),
20841 StreamChunk::Error { message } => error = Some(message),
20842 StreamChunk::Done {} => break,
20843 _ => {}
20844 }
20845 }
20846
20847 assert_eq!(error, None);
20848 assert_eq!(content, "one two three");
20849 }
20850
20851 #[tokio::test]
20852 async fn test_buffered_streaming_main_failure_finalizes_branch() {
20853 use futures::StreamExt;
20854
20855 let mock = mock_with_response("one two");
20856 let mut router_mock = mock_with_response("0");
20857 router_mock.set_latency(50);
20858 let yaml = r#"
20859name: BufferedFailureAgent
20860system_prompt: "You stream safely."
20861llm:
20862 default: default
20863 router: router
20864observability:
20865 enabled: true
20866 export:
20867 write_raw_events: true
20868streaming:
20869 enabled: true
20870 buffer_size: 1
20871runtime:
20872 optimization:
20873 enabled: true
20874 max_speculative_llm_calls_per_turn: 2
20875 speculative_state_transitions: true
20876 streaming_policy: buffer_until_routing_done
20877 max_parallel_runtime_tasks: 2
20878states:
20879 initial: triage
20880 states:
20881 triage:
20882 prompt: "Ask for the category."
20883 transitions:
20884 - to: billing
20885 when: "User asks about billing"
20886 timing: parallel
20887 billing:
20888 prompt: "Billing state."
20889"#;
20890 let agent = AgentBuilder::from_yaml(yaml)
20891 .unwrap()
20892 .llm_alias("default", Arc::new(mock))
20893 .llm_alias("router", Arc::new(router_mock))
20894 .build()
20895 .unwrap();
20896
20897 let mut stream = agent.chat_stream("hello").await.unwrap();
20898 let mut error = String::new();
20899 while let Some(chunk) = stream.next().await {
20900 if let StreamChunk::Error { message } = chunk {
20901 error = message;
20902 }
20903 }
20904
20905 assert!(
20906 error.contains("stream buffer filled"),
20907 "unexpected stream error: {}",
20908 error
20909 );
20910 let events = agent.observability().unwrap().raw_events();
20911 assert!(events.iter().any(|event| {
20912 event.dimensions.get("branch_status") == Some(&"failed".to_string())
20913 && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
20914 && event.dimensions.get("optimization")
20915 == Some(&"buffered_streaming_routing".to_string())
20916 }));
20917 }
20918
20919 #[tokio::test]
20920 async fn test_streaming_preflight_does_not_emit_old_state_content() {
20921 use futures::StreamExt;
20922
20923 let mock = mock_with_response("Billing streamed response");
20924 let yaml = r#"
20925name: StreamingOptimizedAgent
20926system_prompt: "You route before streaming."
20927runtime:
20928 optimization:
20929 enabled: true
20930 pre_response_deterministic_transitions: true
20931streaming:
20932 enabled: true
20933states:
20934 initial: greeting
20935 states:
20936 greeting:
20937 prompt: "OLD_STATE_SENTINEL"
20938 transitions:
20939 - to: billing
20940 guard:
20941 context:
20942 topic:
20943 eq: billing
20944 timing: pre_response
20945 billing:
20946 prompt: "Billing state."
20947"#;
20948 let agent = AgentBuilder::from_yaml(yaml)
20949 .unwrap()
20950 .llm(Arc::new(mock))
20951 .build()
20952 .unwrap();
20953 agent
20954 .set_context("topic", serde_json::json!("billing"))
20955 .unwrap();
20956
20957 let mut stream = agent.chat_stream("billing please").await.unwrap();
20958 let mut content = String::new();
20959 while let Some(chunk) = stream.next().await {
20960 match chunk {
20961 StreamChunk::Content { text } => content.push_str(&text),
20962 StreamChunk::Error { message } => panic!("stream error: {}", message),
20963 StreamChunk::Done {} => break,
20964 _ => {}
20965 }
20966 }
20967
20968 assert_eq!(agent.current_state().as_deref(), Some("billing"));
20969 assert!(content.contains("Billing streamed response"));
20970 assert!(!content.contains("OLD_STATE_SENTINEL"));
20971 }
20972
20973 #[tokio::test]
20975 async fn test_integration_state_machine_basic() {
20976 let yaml = r#"
20977name: StateAgent
20978system_prompt: "You are a support agent."
20979states:
20980 initial: greeting
20981 states:
20982 greeting:
20983 prompt: "Welcome the user warmly."
20984 transitions:
20985 - to: support
20986 when: "User needs help"
20987 auto: true
20988 support:
20989 prompt: "Help solve the user's problem."
20990"#;
20991 let mock = mock_with_responses(vec![
20992 "Welcome! How can I help?", "1", "I'll help you with that.", ]);
20996 let builder = AgentBuilder::from_yaml(yaml).unwrap();
20997 let agent = builder.llm(Arc::new(mock)).build().unwrap();
20998
20999 assert_eq!(agent.current_state(), Some("greeting".to_string()));
21000 let _ = agent.chat("I need help").await.unwrap();
21001 }
21004
21005 #[tokio::test]
21007 async fn test_integration_state_on_enter_set_context() {
21008 let yaml = r#"
21009name: ActionAgent
21010system_prompt: "You are helpful."
21011states:
21012 initial: step1
21013 states:
21014 step1:
21015 prompt: "Step 1"
21016 on_exit:
21017 - set_context:
21018 step1_exited: true
21019 transitions:
21020 - to: step2
21021 when: "always"
21022 auto: true
21023 step2:
21024 prompt: "Step 2"
21025 on_enter:
21026 - set_context:
21027 step2_entered: true
21028"#;
21029 let mock = mock_with_responses(vec![
21031 "Processing step 1.",
21032 "0", ]);
21034 let builder = AgentBuilder::from_yaml(yaml).unwrap();
21035 let agent = builder.llm(Arc::new(mock)).build().unwrap();
21036
21037 assert_eq!(agent.current_state(), Some("step1".to_string()));
21038
21039 agent.transition_to("step2").await.unwrap();
21041
21042 assert_eq!(agent.current_state(), Some("step2".to_string()));
21043
21044 let ctx = agent.get_context();
21046 assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
21047 assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
21048 }
21049
21050 #[tokio::test]
21051 async fn state_action_tool_preserves_source_in_stored_record() {
21052 let yaml = r#"
21053name: StateActionToolAgent
21054system_prompt: "You are helpful."
21055tools:
21056 - context_echo
21057states:
21058 initial: idle
21059 states:
21060 idle:
21061 prompt: "Idle"
21062 active:
21063 prompt: "Active"
21064 on_enter:
21065 - set_context:
21066 action_started: true
21067 - tool: context_echo
21068 args: {}
21069"#;
21070 let agent = AgentBuilder::from_yaml(yaml)
21071 .unwrap()
21072 .llm(Arc::new(mock_with_response("unused")))
21073 .tool(Arc::new(ContextEchoTool))
21074 .build()
21075 .unwrap();
21076
21077 agent.transition_to("active").await.unwrap();
21078
21079 let record: ToolExecutionRecord = serde_json::from_value(
21080 agent
21081 .get_context()
21082 .get("last_tool_record")
21083 .cloned()
21084 .expect("successful state action must store its execution record"),
21085 )
21086 .unwrap();
21087 assert!(record.executed);
21088 assert!(record.success);
21089 assert_eq!(record.canonical_id, "context_echo");
21090 assert!(matches!(
21091 &record.source,
21092 ToolCallSource::StateAction {
21093 state: Some(state),
21094 action_index: 1,
21095 } if state == "active"
21096 ));
21097 }
21098
21099 #[tokio::test]
21100 async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
21101 let yaml = r#"
21102name: OrdinaryLifecycleAgent
21103system_prompt: "You are helpful."
21104states:
21105 initial: intake
21106 regenerate_on_transition: false
21107 states:
21108 intake:
21109 prompt: "Intake"
21110 transitions:
21111 - to: drafting
21112 guard:
21113 context:
21114 route:
21115 eq: drafting
21116 drafting:
21117 prompt: "Drafting"
21118 on_enter:
21119 - set_context:
21120 draft_version: 1
21121 on_reenter:
21122 - set_context:
21123 draft_version: 2
21124 transitions:
21125 - to: review
21126 guard:
21127 context:
21128 route:
21129 eq: review
21130 review:
21131 prompt: "Review"
21132 on_enter:
21133 - set_context:
21134 review_entry: first
21135 transitions:
21136 - to: drafting
21137 guard:
21138 context:
21139 route:
21140 eq: drafting
21141"#;
21142 let agent = AgentBuilder::from_yaml(yaml)
21143 .unwrap()
21144 .llm(Arc::new(mock_with_responses(vec![
21145 "Intake response",
21146 "Draft response",
21147 "Review response",
21148 ])))
21149 .build()
21150 .unwrap();
21151
21152 agent
21153 .set_context("route", serde_json::json!("drafting"))
21154 .unwrap();
21155 agent.chat("Start a draft").await.unwrap();
21156 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21157 assert_eq!(
21158 agent.get_context().get("draft_version"),
21159 Some(&serde_json::json!(1))
21160 );
21161
21162 agent
21163 .set_context("route", serde_json::json!("review"))
21164 .unwrap();
21165 agent.chat("Review this").await.unwrap();
21166 assert_eq!(agent.current_state().as_deref(), Some("review"));
21167 assert_eq!(
21168 agent.get_context().get("review_entry"),
21169 Some(&serde_json::json!("first"))
21170 );
21171
21172 agent
21173 .set_context("route", serde_json::json!("drafting"))
21174 .unwrap();
21175 agent.chat("Revise this").await.unwrap();
21176 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21177 assert_eq!(
21178 agent.get_context().get("draft_version"),
21179 Some(&serde_json::json!(2))
21180 );
21181 }
21182
21183 #[tokio::test]
21184 async fn test_manual_transition_uses_on_enter_then_on_reenter() {
21185 let yaml = r#"
21186name: ManualLifecycleAgent
21187system_prompt: "You are helpful."
21188states:
21189 initial: intake
21190 states:
21191 intake:
21192 prompt: "Intake"
21193 drafting:
21194 prompt: "Drafting"
21195 on_enter:
21196 - set_context:
21197 draft_version: 1
21198 on_reenter:
21199 - set_context:
21200 draft_version: 2
21201 review:
21202 prompt: "Review"
21203"#;
21204 let agent = AgentBuilder::from_yaml(yaml)
21205 .unwrap()
21206 .llm(Arc::new(mock_with_response("unused")))
21207 .build()
21208 .unwrap();
21209
21210 assert!(!agent.get_context().contains_key("draft_version"));
21211 agent.transition_to("drafting").await.unwrap();
21212 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21213 assert_eq!(
21214 agent.get_context().get("draft_version"),
21215 Some(&serde_json::json!(1))
21216 );
21217
21218 agent.transition_to("review").await.unwrap();
21219 agent.transition_to("drafting").await.unwrap();
21220 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21221 assert_eq!(
21222 agent.get_context().get("draft_version"),
21223 Some(&serde_json::json!(2))
21224 );
21225 }
21226
21227 #[tokio::test]
21228 async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
21229 let yaml = r#"
21230name: TimeoutLifecycleAgent
21231system_prompt: "You are helpful."
21232states:
21233 initial: intake
21234 regenerate_on_transition: false
21235 states:
21236 intake:
21237 prompt: "Intake"
21238 max_turns: 1
21239 timeout_to: drafting
21240 drafting:
21241 prompt: "Drafting"
21242 max_turns: 1
21243 timeout_to: review
21244 on_enter:
21245 - set_context:
21246 draft_version: 1
21247 on_reenter:
21248 - set_context:
21249 draft_version: 2
21250 review:
21251 prompt: "Review"
21252 max_turns: 1
21253 timeout_to: drafting
21254 on_enter:
21255 - set_context:
21256 review_entry: first
21257"#;
21258 let agent = AgentBuilder::from_yaml(yaml)
21259 .unwrap()
21260 .llm(Arc::new(mock_with_responses(vec![
21261 "Intake",
21262 "First draft",
21263 "Review",
21264 "Revised draft",
21265 ])))
21266 .build()
21267 .unwrap();
21268
21269 agent.chat("First turn").await.unwrap();
21270 assert_eq!(agent.current_state().as_deref(), Some("intake"));
21271 assert!(!agent.get_context().contains_key("draft_version"));
21272
21273 agent.chat("Second turn").await.unwrap();
21274 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21275 assert_eq!(
21276 agent.get_context().get("draft_version"),
21277 Some(&serde_json::json!(1))
21278 );
21279
21280 agent.chat("Third turn").await.unwrap();
21281 assert_eq!(agent.current_state().as_deref(), Some("review"));
21282 assert_eq!(
21283 agent.get_context().get("review_entry"),
21284 Some(&serde_json::json!("first"))
21285 );
21286
21287 agent.chat("Fourth turn").await.unwrap();
21288 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
21289 assert_eq!(
21290 agent.get_context().get("draft_version"),
21291 Some(&serde_json::json!(2))
21292 );
21293 }
21294
21295 #[tokio::test]
21297 async fn test_integration_process_normalize() {
21298 let yaml = r#"
21299name: ProcessAgent
21300system_prompt: "You are helpful."
21301process:
21302 input:
21303 - type: normalize
21304 config:
21305 trim: true
21306 collapse_whitespace: true
21307"#;
21308 let mock = mock_with_response("Got your message.");
21309 let builder = AgentBuilder::from_yaml(yaml).unwrap();
21310 let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
21311
21312 let _ = agent.chat(" hello world ").await.unwrap();
21313
21314 let history = mock.call_history();
21316 assert!(!history.is_empty());
21317 let last_call = history.last().unwrap();
21319 let user_msg = last_call
21320 .messages
21321 .iter()
21322 .find(|m| m.role == ai_agents_core::Role::User)
21323 .unwrap();
21324 assert_eq!(user_msg.content, "hello world");
21325 }
21326
21327 #[tokio::test]
21331 async fn test_integration_memory_compression() {
21332 let yaml = r#"
21333name: MemoryAgent
21334system_prompt: "You are helpful."
21335memory:
21336 type: compacting
21337 max_messages: 100
21338 compress_threshold: 5
21339 max_recent_messages: 3
21340 summarize_batch_size: 2
21341"#;
21342 let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
21344 let mock = mock_with_responses(responses);
21345 let builder = AgentBuilder::from_yaml(yaml).unwrap();
21346 let agent = builder.llm(Arc::new(mock)).build().unwrap();
21347
21348 for i in 0..6 {
21350 let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
21351 }
21352
21353 let messages = agent.memory.get_messages(None).await.unwrap();
21356 assert!(messages.len() <= 12); }
21360
21361 #[tokio::test]
21363 async fn test_integration_multi_llm_registry() {
21364 let mut mock_default = MockLLMProvider::new("default");
21365 mock_default.set_response("Default LLM response.");
21366 let mut mock_router = MockLLMProvider::new("router");
21367 mock_router.set_response("Router response.");
21368
21369 let agent = AgentBuilder::new()
21370 .system_prompt("You are helpful.")
21371 .llm_alias("default", Arc::new(mock_default))
21372 .llm_alias("router", Arc::new(mock_router))
21373 .build()
21374 .unwrap();
21375
21376 let response = agent.chat("Hello").await.unwrap();
21377 assert_eq!(response.content, "Default LLM response.");
21378 }
21379
21380 #[tokio::test]
21382 async fn test_integration_agent_reset() {
21383 let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
21384 let agent = AgentBuilder::new()
21385 .system_prompt("You are helpful.")
21386 .llm(Arc::new(mock))
21387 .build()
21388 .unwrap();
21389
21390 let _ = agent.chat("Hi").await.unwrap();
21391 let messages = agent.memory.get_messages(None).await.unwrap();
21392 assert_eq!(messages.len(), 2); agent.reset().await.unwrap();
21395 let messages = agent.memory.get_messages(None).await.unwrap();
21396 assert_eq!(messages.len(), 0);
21397 }
21398
21399 #[tokio::test]
21401 async fn test_integration_process_validate_reject() {
21402 use ai_agents_process::{ProcessConfig, ProcessProcessor};
21403
21404 let validate_config = ai_agents_process::ValidateStage {
21405 id: Some("length_check".to_string()),
21406 condition: None,
21407 config: ai_agents_process::ValidateConfig {
21408 rules: vec![ai_agents_process::ValidationRule::MinLength {
21409 min_length: 10,
21410 on_fail: ai_agents_process::ValidationAction {
21411 action: ai_agents_process::ValidationActionType::Reject,
21412 message: None,
21413 },
21414 }],
21415 ..Default::default()
21416 },
21417 };
21418 let process_config = ProcessConfig {
21419 input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
21420 ..Default::default()
21421 };
21422 let processor = ProcessProcessor::new(process_config);
21423
21424 let mock = mock_with_response("Should not reach here.");
21425 let agent = AgentBuilder::new()
21426 .system_prompt("You are helpful.")
21427 .llm(Arc::new(mock))
21428 .process_processor(processor)
21429 .build()
21430 .unwrap();
21431
21432 let response = agent.chat("Hi").await.unwrap();
21433 assert!(
21435 response.content.contains("rejected")
21436 || response.content.contains("Input rejected")
21437 || response.content.contains("too short")
21438 || response.content.contains("Too short")
21439 || response.content.len() < 50, "Expected rejection response, got: {}",
21441 response.content
21442 );
21443 }
21444
21445 #[tokio::test]
21447 async fn test_llm_fallback_on_failure() {
21448 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
21449
21450 let mut primary = MockLLMProvider::new("primary");
21451 primary.set_error("Primary LLM is unavailable");
21452
21453 let mut fallback = MockLLMProvider::new("fallback");
21454 fallback.set_response("Fallback response works!");
21455
21456 let agent = AgentBuilder::new()
21457 .system_prompt("You are helpful.")
21458 .llm_alias("default", Arc::new(primary))
21459 .llm_alias("backup", Arc::new(fallback))
21460 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
21461 llm: LLMRecoveryConfig {
21462 on_failure: LLMFailureAction::FallbackLlm {
21463 fallback_llm: "backup".to_string(),
21464 },
21465 ..Default::default()
21466 },
21467 ..Default::default()
21468 }))
21469 .build()
21470 .unwrap();
21471
21472 let response = agent.chat("Hello").await.unwrap();
21473 assert!(
21474 response.content.contains("Fallback response"),
21475 "Expected fallback response, got: {}",
21476 response.content
21477 );
21478 }
21479
21480 #[tokio::test]
21482 async fn test_llm_fallback_response_static_message() {
21483 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
21484
21485 let mut primary = MockLLMProvider::new("primary");
21486 primary.set_error("Primary LLM is unavailable");
21487
21488 let agent = AgentBuilder::new()
21489 .system_prompt("You are helpful.")
21490 .llm(Arc::new(primary))
21491 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
21492 llm: LLMRecoveryConfig {
21493 on_failure: LLMFailureAction::FallbackResponse {
21494 message: "I am temporarily unavailable. Please try again later."
21495 .to_string(),
21496 },
21497 ..Default::default()
21498 },
21499 ..Default::default()
21500 }))
21501 .build()
21502 .unwrap();
21503
21504 let response = agent.chat("Hello").await.unwrap();
21505 assert!(
21506 response.content.contains("temporarily unavailable"),
21507 "Expected static fallback message, got: {}",
21508 response.content
21509 );
21510 }
21511
21512 #[tokio::test]
21514 async fn test_tool_failure_skip() {
21515 use ai_agents_recovery::{
21516 ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
21517 };
21518
21519 let mock = mock_with_responses(vec![
21521 r#"I'll use the nonexistent tool.
21522[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
21523 "The tool was unavailable, but I can still help you.",
21524 ]);
21525
21526 let agent = AgentBuilder::new()
21527 .system_prompt("You are helpful.")
21528 .llm(Arc::new(mock))
21529 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
21530 tools: ToolRecoveryConfig {
21531 default: ToolRetryConfig {
21532 max_retries: 0,
21533 timeout_ms: None,
21534 on_failure: ToolFailureAction::Skip,
21535 },
21536 ..Default::default()
21537 },
21538 ..Default::default()
21539 }))
21540 .build()
21541 .unwrap();
21542
21543 let response = agent.chat("Use the nonexistent tool").await;
21545 assert!(
21546 response.is_ok(),
21547 "Expected Ok with skip policy, got: {:?}",
21548 response
21549 );
21550 }
21551}