1use async_trait::async_trait;
2use futures::stream::{Stream, StreamExt};
3use parking_lot::RwLock;
4use serde_json::Value;
5use std::collections::{HashMap, HashSet};
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::{Arc, Weak};
10use std::time::Instant;
11use tracing::{debug, error, info, instrument, warn};
12
13const DISAMBIGUATION_STATE_GENERATION_KEY: &str = "_runtime.disambiguation_state_generation";
14
15pub(crate) type ToolResourceLocks = Arc<RwLock<HashMap<String, Weak<tokio::sync::Mutex<()>>>>>;
17
18struct ToolResourceGuards {
22 guards: Vec<tokio::sync::OwnedMutexGuard<()>>,
23 locks: ToolResourceLocks,
24}
25
26#[derive(Clone)]
27struct StoredSessionRestore {
28 snapshot: AgentSnapshot,
29 metadata: Option<ai_agents_core::SessionMetadata>,
30}
31
32struct RuntimeSessionRestorePoint {
33 snapshot: AgentSnapshot,
34 metadata: ai_agents_core::SessionMetadata,
35 actor_id: Option<String>,
36 session_id: Option<String>,
37}
38
39impl Drop for ToolResourceGuards {
40 fn drop(&mut self) {
41 self.guards.clear();
42 self.locks.write().retain(|_, lock| lock.strong_count() > 0);
43 }
44}
45
46#[derive(Clone)]
50struct RuntimeSafetySnapshot {
51 version: u64,
52 emergency_deny: bool,
53 tool_security: ToolSecurityEngine,
54 tool_scope_override: Option<Vec<String>>,
55}
56
57#[derive(Clone, Copy)]
61struct ToolDecisionVersions {
62 policy: u64,
63 registry: u64,
64 runtime_control: u64,
65 state: Option<u64>,
66}
67
68struct AvailableToolIdsSnapshot {
72 tool_ids: Vec<String>,
73 state_generation: Option<u64>,
74}
75
76#[derive(Clone)]
80struct ToolApprovalBinding {
81 canonical_id: String,
82 arguments: Value,
83 confirmation_required: bool,
84 policy_version: u64,
85 runtime_control_version: u64,
86 state_generation: Option<u64>,
87 reviewed_tool: Arc<dyn ai_agents_core::Tool>,
88}
89
90fn merge_approved_record(record: &mut Option<ToolApprovalRecord>) {
94 if record
95 .as_ref()
96 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::Modified))
97 {
98 return;
99 }
100 *record = Some(ToolApprovalRecord {
101 status: ToolApprovalStatus::Approved,
102 reason: None,
103 modified_arguments: None,
104 });
105}
106
107impl ToolApprovalBinding {
108 fn is_stale(
110 &self,
111 canonical_id: &str,
112 arguments: &Value,
113 confirmation_required: bool,
114 versions: ToolDecisionVersions,
115 resolved_tool: &Arc<dyn ai_agents_core::Tool>,
116 ) -> bool {
117 self.canonical_id != canonical_id
118 || self.arguments != *arguments
119 || self.confirmation_required != confirmation_required
120 || self.policy_version != versions.policy
121 || self.runtime_control_version != versions.runtime_control
122 || self.state_generation != versions.state
123 || !Arc::ptr_eq(&self.reviewed_tool, resolved_tool)
124 }
125}
126
127use crate::turn_context::{current_turn_actor_context, scope_actor_context};
128
129use ai_agents_context::{ContextManager, ContextProvider, TemplateRenderer};
130use ai_agents_core::traits::storage::StorageCapability;
131use ai_agents_core::{
132 AgentError, AgentSnapshot, AgentStorage, ChatMessage, FinishReason, LLMError, LLMProvider,
133 LLMResponse, LLMToolDefinition, LLMToolRequest, PermissionOutcome, Result, ToolActorContext,
134 ToolApprovalRecord, ToolApprovalStatus, ToolCallSource, ToolCancellationToken, ToolChoice,
135 ToolExecutionContext, ToolExecutionRecord, ToolExecutionRequest, ToolInvoker,
136 ToolPolicyDecisionRecord, ToolResult,
137};
138use ai_agents_disambiguation::{
139 ClarificationObserver, ClarificationParseFuture, ClarificationQuestionFuture,
140 ConfirmationParseFuture, DisambiguationConfig, DisambiguationContext, DisambiguationManager,
141 DisambiguationResult,
142};
143use ai_agents_hitl::{
144 ApprovalHandler, ApprovalResolvedOutcome, ApprovalResult, ApprovalTrigger, HITLCheckResult,
145 HITLEngine, RejectAllHandler, TimeoutAction,
146};
147use ai_agents_hooks::{AgentHooks, NoopHooks};
148use ai_agents_llm::LLMRegistry;
149use ai_agents_memory::{
150 CompressResult, EvictionReason, Memory, MemoryBudgetEvent, MemoryCompressEvent,
151 MemoryEvictEvent, MemoryTokenBudget, OverflowStrategy,
152};
153use ai_agents_observability::{
154 EventStatus, EventType, ObservabilityManager, ObservationPurpose, SpanContext,
155 current_observation_context, new_session_id as new_observation_session_id,
156 resolve_language_from_context, with_observation_context, with_observation_purpose,
157};
158use ai_agents_process::{
159 ProcessData, ProcessProcessor, ProcessPurposeHint, ProcessStageFuture, ProcessStageObserver,
160};
161use ai_agents_reasoning::{
162 CriterionResult, EvaluationResult, Plan, PlanAction, PlanStatus, PlanStep, ReasoningConfig,
163 ReasoningMetadata, ReasoningMode, ReasoningOutput, ReflectionAttempt, ReflectionConfig,
164 ReflectionMetadata, StepFailureAction,
165};
166use ai_agents_recovery::{
167 ByRoleFilter, ContextOverflowAction, FilterConfig, IntoClassifiedError, KeepRecentFilter,
168 LLMFailureAction, MessageFilter, RecoveryManager, SkipPatternFilter, ToolFailureAction,
169};
170use ai_agents_relationships::RelationshipManager;
171use ai_agents_skills::{SkillDefinition, SkillExecutor, SkillRouter};
172use ai_agents_state::{
173 PromptMode, StateAction, StateMachine, StateMachineSnapshot, StateTransitionEvent, Transition,
174 TransitionContext, TransitionEvaluator, TransitionTiming, evaluate_guard,
175};
176use ai_agents_storage::{StorageConfig as StorageStorageConfig, create_storage};
177use ai_agents_tools::{
178 CommandRunner, ConditionEvaluator, DiagnosticsProvider, EvaluationContext, LLMGetter,
179 QuestionHandler, SecurityCheckResult, TodoItem, ToolCallRecord, ToolRegistry,
180 ToolSecurityConfig, ToolSecurityEngine,
181};
182
183use super::{
184 Agent, AgentInfo, AgentResponse, AgentStreamEvent, ParallelToolsConfig, StreamChunk,
185 StreamingConfig, ToolCall,
186};
187use crate::optimization::{
188 AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
189 MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
190 RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
191 RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
192 StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
193};
194use crate::spec::StorageConfig;
195
196enum ToolCallOutcome {
198 Continue,
200 TransitionFired,
202 Rejected(AgentResponse),
204}
205
206#[derive(Clone)]
207struct MainToolProtocol {
208 choice: Option<ToolChoice>,
209 tool_ids: Vec<String>,
210 definitions: Vec<LLMToolDefinition>,
211}
212
213struct MainProviderResponse {
214 response: LLMResponse,
215 used_native_tools: bool,
216}
217
218struct CommittedTextResponse<'a> {
222 processed_input: &'a str,
223 input_context: &'a HashMap<String, Value>,
224 answer: String,
225 reasoning_mode: ReasoningMode,
226 auto_detected: bool,
227 iterations: u32,
228 thinking_content: Option<String>,
229 all_tool_calls: Vec<ToolCall>,
230}
231
232struct AgentResponseParts {
236 content: String,
237 all_tool_calls: Vec<ToolCall>,
238 reasoning_mode: ReasoningMode,
239 auto_detected: bool,
240 iterations: u32,
241 thinking: Option<String>,
242 reflection_metadata: Option<ReflectionMetadata>,
243}
244
245type RuntimeStreamTerminalSlot = Arc<RwLock<Option<AgentResponse>>>;
249
250fn new_runtime_stream_terminal_slot() -> RuntimeStreamTerminalSlot {
254 Arc::new(RwLock::new(None))
255}
256
257fn record_runtime_stream_final(slot: &RuntimeStreamTerminalSlot, response: AgentResponse) {
261 *slot.write() = Some(response);
262}
263
264#[derive(Clone, Copy)]
265struct DisambiguationOwnership {
266 epoch: u64,
267 state_generation: Option<u64>,
268}
269
270enum SkillRouteResult {
272 NoMatch,
274 Response { skill_id: String, content: String },
276 NeedsClarification {
278 response: AgentResponse,
279 ownership: Option<DisambiguationOwnership>,
280 },
281}
282
283enum ParallelTransitionSelection {
285 Candidate(TransitionCandidate),
287 NoMatch,
289 ReservationExhausted,
291}
292
293enum PostLoopResult {
295 NoTransition(String),
297 Transitioned(String),
299 NeedsRedispatch,
302}
303
304struct StateTransitionReservation<'a> {
305 reserved: &'a AtomicBool,
306}
307
308impl Drop for StateTransitionReservation<'_> {
309 fn drop(&mut self) {
310 self.reserved.store(false, Ordering::SeqCst);
311 }
312}
313
314struct RootTurnCleanup<'a> {
315 agent: &'a RuntimeAgent,
316}
317
318impl<'a> RootTurnCleanup<'a> {
319 fn new(agent: &'a RuntimeAgent) -> Self {
320 Self { agent }
321 }
322}
323
324impl Drop for RootTurnCleanup<'_> {
325 fn drop(&mut self) {
326 self.agent.end_root_turn();
327 }
328}
329
330#[derive(Debug)]
332struct RuntimeControlState {
333 snapshot_guard: RwLock<()>,
335 version: AtomicU64,
337 emergency_deny: Arc<AtomicBool>,
339 tool_security_override: RwLock<Option<ToolSecurityEngine>>,
341 tool_scope_override: RwLock<Option<Vec<String>>>,
343}
344
345impl Default for RuntimeControlState {
346 fn default() -> Self {
347 Self {
348 snapshot_guard: RwLock::new(()),
349 version: AtomicU64::new(1),
350 emergency_deny: Arc::new(AtomicBool::new(false)),
351 tool_security_override: RwLock::new(None),
352 tool_scope_override: RwLock::new(None),
353 }
354 }
355}
356
357#[derive(Clone)]
359pub struct RuntimeControlHandle {
360 state: Arc<RuntimeControlState>,
361}
362
363impl RuntimeControlHandle {
364 pub fn version(&self) -> u64 {
366 self.state.version.load(Ordering::SeqCst)
367 }
368
369 fn bump(&self) -> u64 {
370 self.state.version.fetch_add(1, Ordering::SeqCst) + 1
371 }
372
373 pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
375 self.try_set_tool_security(config)
376 .expect("invalid tool security configuration")
377 }
378
379 pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
381 config.validate()?;
382 let _guard = self.state.snapshot_guard.write();
383 let generation = self.bump();
384 *self.state.tool_security_override.write() = Some(
385 ToolSecurityEngine::new_with_policy_version(config, generation),
386 );
387 Ok(generation)
388 }
389
390 pub fn clear_tool_security_override(&self) -> u64 {
392 let _guard = self.state.snapshot_guard.write();
393 *self.state.tool_security_override.write() = None;
394 self.bump()
395 }
396
397 pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
399 let _guard = self.state.snapshot_guard.write();
400 *self.state.tool_scope_override.write() = Some(tool_ids);
401 self.bump()
402 }
403
404 pub fn clear_tool_scope_override(&self) -> u64 {
406 let _guard = self.state.snapshot_guard.write();
407 *self.state.tool_scope_override.write() = None;
408 self.bump()
409 }
410
411 pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
413 let _guard = self.state.snapshot_guard.write();
414 self.state.emergency_deny.store(enabled, Ordering::SeqCst);
415 self.bump()
416 }
417
418 pub fn cancel_all(&self) -> u64 {
420 self.set_emergency_deny(true)
421 }
422}
423
424pub struct RuntimeAgent {
425 info: AgentInfo,
426 llm_registry: Arc<LLMRegistry>,
427 memory: Arc<dyn Memory>,
428 tools: Arc<ToolRegistry>,
429 skills: Vec<SkillDefinition>,
430 skill_router: Option<SkillRouter>,
431 skill_executor: Option<SkillExecutor>,
432 base_system_prompt: String,
433 max_iterations: u32,
434 iteration_count: RwLock<u32>,
435 max_context_tokens: u32,
436 memory_token_budget: Option<MemoryTokenBudget>,
437 recovery_manager: RecoveryManager,
438 tool_security: ToolSecurityEngine,
439 process_processor: Option<ProcessProcessor>,
440 message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
441 state_machine: Option<Arc<StateMachine>>,
442 transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
443 context_manager: Arc<ContextManager>,
444 template_renderer: TemplateRenderer,
445 tool_call_history: RwLock<Vec<ToolCallRecord>>,
446 parallel_tools: ParallelToolsConfig,
447 streaming: StreamingConfig,
448 hooks: Arc<dyn AgentHooks>,
449 hitl_engine: Option<HITLEngine>,
450 approval_handler: Arc<dyn ApprovalHandler>,
451 storage_config: StorageConfig,
452 storage: RwLock<Option<Arc<dyn AgentStorage>>>,
453 storage_init: tokio::sync::Mutex<()>,
454 reasoning_config: ReasoningConfig,
455 reflection_config: ReflectionConfig,
456 disambiguation_manager: Option<DisambiguationManager>,
457 disambiguation_epoch: AtomicU64,
459 disambiguation_admission: tokio::sync::RwLock<()>,
461 state_transition_reserved: AtomicBool,
463 persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
465 pending_skill_id: RwLock<Option<String>>,
469 current_plan: RwLock<Option<Plan>>,
470 declared_tool_ids: Option<Vec<String>>,
472 context_initialized: AtomicBool,
474 spawner: Option<Arc<crate::spawner::AgentSpawner>>,
476 spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
478 redispatch_depth: RwLock<u32>,
481 active_turn_context: RwLock<Option<TurnOptimizationContext>>,
483 root_user_message_committed: AtomicBool,
485 actor_id: RwLock<Option<String>>,
487 fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
489 fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
492 actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
494 messages_since_extraction: Arc<RwLock<usize>>,
496 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
498 facts_config: Option<ai_agents_facts::FactsConfig>,
500 session_metadata: RwLock<ai_agents_core::SessionMetadata>,
502 current_session_id: RwLock<Option<String>>,
504 relationship_manager: Option<Arc<RelationshipManager>>,
506 observability_manager: Option<Arc<ObservabilityManager>>,
508 runtime_config: RuntimeConfig,
510 background_maintenance: Arc<BackgroundMaintenanceQueue>,
512 resource_locks: ToolResourceLocks,
514 runtime_control: Arc<RuntimeControlState>,
516}
517
518impl std::fmt::Debug for RuntimeAgent {
519 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
520 f.debug_struct("RuntimeAgent")
521 .field("info", &self.info)
522 .field("base_system_prompt", &self.base_system_prompt)
523 .field("max_iterations", &self.max_iterations)
524 .field("skills_count", &self.skills.len())
525 .field("max_context_tokens", &self.max_context_tokens)
526 .field("has_state_machine", &self.state_machine.is_some())
527 .field("parallel_tools", &self.parallel_tools)
528 .field("streaming", &self.streaming)
529 .field("has_hooks", &true)
530 .field("has_hitl", &self.hitl_engine.is_some())
531 .field("storage_type", &self.storage_config.storage_type())
532 .field("reasoning_mode", &self.reasoning_config.mode)
533 .field("reflection_enabled", &self.reflection_config.enabled)
534 .field("declared_tool_ids", &self.declared_tool_ids)
535 .field("has_persona", &self.persona_manager.is_some())
536 .field("has_observability", &self.observability_manager.is_some())
537 .finish_non_exhaustive()
538 }
539}
540
541struct ObservabilityClarificationObserver;
542
543impl ClarificationObserver for ObservabilityClarificationObserver {
544 fn observe_question<'a>(
546 &'a self,
547 future: ClarificationQuestionFuture<'a>,
548 ) -> ClarificationQuestionFuture<'a> {
549 Box::pin(async move {
550 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
551 })
552 }
553
554 fn observe_parse<'a>(
556 &'a self,
557 future: ClarificationParseFuture<'a>,
558 ) -> ClarificationParseFuture<'a> {
559 Box::pin(async move {
560 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
561 })
562 }
563
564 fn observe_confirmation_parse<'a>(
566 &'a self,
567 future: ConfirmationParseFuture<'a>,
568 ) -> ConfirmationParseFuture<'a> {
569 Box::pin(async move {
570 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
571 })
572 }
573}
574
575struct ObservabilityProcessStageObserver;
576
577impl ProcessStageObserver for ObservabilityProcessStageObserver {
578 fn observe<'a>(
580 &'a self,
581 hint: ProcessPurposeHint,
582 future: ProcessStageFuture<'a>,
583 ) -> ProcessStageFuture<'a> {
584 Box::pin(async move {
585 with_observation_purpose(observation_purpose_for_process(hint), future).await
586 })
587 }
588}
589
590struct RegistryLLMGetter {
591 registry: Arc<LLMRegistry>,
592}
593
594impl LLMGetter for RegistryLLMGetter {
595 fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
596 self.registry.get(alias).ok()
597 }
598}
599
600impl RuntimeAgent {
601 #[allow(clippy::too_many_arguments)]
602 pub fn new(
603 info: AgentInfo,
604 llm_registry: Arc<LLMRegistry>,
605 memory: Arc<dyn Memory>,
606 tools: Arc<ToolRegistry>,
607 skills: Vec<SkillDefinition>,
608 system_prompt: String,
609 max_iterations: u32,
610 ) -> Self {
611 let (skill_router, skill_executor) = if !skills.is_empty() {
612 let router_llm = llm_registry.router().ok();
613 let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
614 let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
615 (router, Some(executor))
616 } else {
617 (None, None)
618 };
619
620 let context_manager =
621 ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
622
623 Self {
624 info,
625 llm_registry,
626 memory,
627 tools,
628 skills,
629 skill_router,
630 skill_executor,
631 base_system_prompt: system_prompt,
632 max_iterations,
633 iteration_count: RwLock::new(0),
634 max_context_tokens: 128000,
635 memory_token_budget: None,
636 recovery_manager: RecoveryManager::default(),
637 tool_security: ToolSecurityEngine::default(),
638 process_processor: None,
639 message_filters: RwLock::new(HashMap::new()),
640 state_machine: None,
641 transition_evaluator: None,
642 context_manager: Arc::new(context_manager),
643 template_renderer: TemplateRenderer::new(),
644 tool_call_history: RwLock::new(Vec::new()),
645 parallel_tools: ParallelToolsConfig::default(),
646 streaming: StreamingConfig::default(),
647 hooks: Arc::new(NoopHooks),
648 hitl_engine: None,
649 approval_handler: Arc::new(RejectAllHandler::new()),
650 storage_config: StorageConfig::default(),
651 storage: RwLock::new(None),
652 storage_init: tokio::sync::Mutex::new(()),
653 reasoning_config: ReasoningConfig::default(),
654 reflection_config: ReflectionConfig::default(),
655 disambiguation_manager: None,
656 disambiguation_epoch: AtomicU64::new(0),
657 disambiguation_admission: tokio::sync::RwLock::new(()),
658 state_transition_reserved: AtomicBool::new(false),
659 persona_manager: None,
660 pending_skill_id: RwLock::new(None),
661 current_plan: RwLock::new(None),
662 declared_tool_ids: None,
663 context_initialized: AtomicBool::new(false),
664 spawner: None,
665 spawner_registry: None,
666 redispatch_depth: RwLock::new(0),
667 active_turn_context: RwLock::new(None),
668 root_user_message_committed: AtomicBool::new(false),
669 actor_id: RwLock::new(None),
670 fact_store: RwLock::new(None),
671 fact_extractor: RwLock::new(None),
672 actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
673 messages_since_extraction: Arc::new(RwLock::new(0)),
674 actor_memory_config: None,
675 facts_config: None,
676 session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
677 current_session_id: RwLock::new(None),
678 relationship_manager: None,
679 observability_manager: None,
680 runtime_config: RuntimeConfig::default(),
681 background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
682 resource_locks: new_tool_resource_locks(),
683 runtime_control: Arc::new(RuntimeControlState::default()),
684 }
685 }
686
687 pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
688 self.declared_tool_ids = ids;
689 self
690 }
691
692 pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
693 self.storage_config = config;
694 self
695 }
696
697 pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
698 *self.storage.write() = Some(storage);
699 self
700 }
701
702 pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
703 self.resource_locks = locks;
704 self
705 }
706
707 pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
708 self.reasoning_config = config;
709 self
710 }
711
712 pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
713 self.reflection_config = config;
714 self
715 }
716
717 pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
719 self.relationship_manager = Some(manager);
720 self
721 }
722
723 pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
725 self.observability_manager = Some(manager);
726 self
727 }
728
729 pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
731 let max_tasks = config.optimization.post_turn.max_background_tasks;
732 self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
733 self.runtime_config = config;
734 self
735 }
736
737 pub fn runtime_config(&self) -> &RuntimeConfig {
739 &self.runtime_config
740 }
741
742 pub async fn flush_background_tasks(&self) -> Result<()> {
744 self.background_maintenance.flush_all().await
745 }
746
747 pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
749 self.background_maintenance.flush_scope(actor_id).await
750 }
751
752 pub async fn flush_background_tasks_for_purpose(
754 &self,
755 purpose: RuntimeTaskPurpose,
756 ) -> Result<()> {
757 self.background_maintenance.flush_purpose(purpose).await
758 }
759
760 pub async fn flush_background_tasks_for_actor_purpose(
762 &self,
763 actor_id: &str,
764 purpose: RuntimeTaskPurpose,
765 ) -> Result<()> {
766 self.background_maintenance
767 .flush_scope_purpose(actor_id, purpose)
768 .await
769 }
770
771 pub async fn shutdown_background_tasks(&self) -> Result<()> {
773 self.flush_background_tasks().await
774 }
775
776 pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
778 self.observability_manager.clone()
779 }
780
781 async fn export_observability_if_configured(&self) {
783 let Some(manager) = self.observability_manager.as_ref() else {
784 return;
785 };
786 let export = &manager.config().export;
787 if !export.write_report && !export.write_raw_events {
788 return;
789 }
790 if let Err(error) = manager.export().await {
791 warn!(error = %error, "Observability export failed");
792 }
793 }
794
795 pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
797 self.relationship_manager.clone()
798 }
799
800 fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
801 current_turn_actor_context()
802 }
803
804 fn effective_actor_id(&self) -> Option<String> {
805 self.current_turn_actor_context()
806 .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
807 .or_else(|| self.actor_id.read().clone())
808 }
809
810 fn effective_origin_actor_id(&self) -> Option<String> {
811 self.current_turn_actor_context()
812 .and_then(|ctx| ctx.origin_actor_id.clone())
813 .or_else(|| self.actor_id.read().clone())
814 }
815
816 fn record_session_actor_if_needed(&self) {
817 if let Some(actor_id) = self.effective_origin_actor_id() {
818 let mut meta = self.session_metadata.write();
819 meta.actor_id = Some(actor_id.clone());
820 if !meta.actors.iter().any(|a| a == &actor_id) {
821 meta.actors.push(actor_id);
822 }
823 }
824 }
825
826 fn outbound_actor_context(&self) -> crate::TurnActorContext {
827 let mut context = self.current_turn_actor_context().unwrap_or_default();
828 if context.origin_actor_id.is_none() {
829 context.origin_actor_id = self.effective_origin_actor_id();
830 }
831 context.sender_agent_id = Some(self.info.id.clone());
832 context
833 }
834
835 fn observation_session_id(&self) -> Option<String> {
837 let mut current = self.current_session_id.write();
838 if current.is_none() {
839 *current = Some(new_observation_session_id());
840 }
841 current.clone()
842 }
843
844 fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
846 let manager = self.observability_manager.as_ref()?;
847 let context = self.build_context_with_overlays();
848 let language = resolve_language_from_context(manager.config(), &context);
849 let context = current_observation_context()
850 .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
851 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
852 Some(
853 context
854 .with_actor(actor_id.or_else(|| self.effective_actor_id()))
855 .with_session(self.observation_session_id())
856 .with_state(self.current_state())
857 .with_language(Some(language)),
858 )
859 }
860
861 fn current_runtime_observation_context(
863 &self,
864 purpose: ObservationPurpose,
865 ) -> Option<SpanContext> {
866 let manager = self.observability_manager.as_ref()?;
867 let context = self.build_context_with_overlays();
868 let language = resolve_language_from_context(manager.config(), &context);
869 let mut observation = current_observation_context()
870 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
871 observation.agent_id = self.info.id.clone();
872 observation.actor_id = self.effective_actor_id();
873 observation.session_id = self.observation_session_id();
874 observation.state = self.current_state();
875 observation.language = Some(language);
876 observation.purpose = purpose;
877 Some(observation)
878 }
879
880 async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
882 where
883 F: Future<Output = T>,
884 {
885 if let Some(context) = self.current_runtime_observation_context(purpose) {
886 with_observation_context(context, future).await
887 } else {
888 future.await
889 }
890 }
891
892 fn chat_with_actor_context_boxed<'a>(
894 &'a self,
895 input: &'a str,
896 actor_context: crate::TurnActorContext,
897 ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
898 Box::pin(async move {
899 let actor_id = actor_context.effective_actor_id().map(str::to_string);
900 let run = async move {
901 scope_actor_context(
902 actor_context,
903 Box::pin(async move { self.run_loop(input).await }),
904 )
905 .await
906 };
907 let result = if let Some(context) = self.build_observation_context(actor_id) {
908 with_observation_context(context, run).await
909 } else {
910 run.await
911 };
912 self.export_observability_if_configured().await;
913 result
914 })
915 }
916
917 pub async fn chat_with_actor_context(
921 &self,
922 input: &str,
923 actor_context: crate::TurnActorContext,
924 ) -> Result<AgentResponse> {
925 self.chat_with_actor_context_boxed(input, actor_context)
926 .await
927 }
928
929 pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
931 let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
932 self.chat_with_actor_context(input, actor_context).await
933 }
934
935 pub async fn load_actor_relationship(&self) -> Result<()> {
937 self.maybe_load_actor_relationship().await;
938 Ok(())
939 }
940
941 pub async fn update_relationship_dimension(
943 &self,
944 dimension: &str,
945 delta: f64,
946 reason: Option<&str>,
947 ) -> Result<ai_agents_relationships::DimensionChange> {
948 self.update_relationship_dimension_for_perspective(
949 ai_agents_relationships::RelationshipPerspective::AgentToActor,
950 dimension,
951 delta,
952 reason,
953 )
954 .await
955 }
956
957 pub async fn update_relationship_dimension_for_perspective(
961 &self,
962 perspective: ai_agents_relationships::RelationshipPerspective,
963 dimension: &str,
964 delta: f64,
965 reason: Option<&str>,
966 ) -> Result<ai_agents_relationships::DimensionChange> {
967 let manager = self
968 .relationship_manager
969 .as_ref()
970 .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
971 let actor_id = self.effective_actor_id().ok_or_else(|| {
972 AgentError::Config("No actor ID set. Use set_actor_id() first".into())
973 })?;
974 let change = manager.update_dimension_for_perspective(
975 &actor_id,
976 perspective,
977 dimension,
978 delta,
979 1.0,
980 reason.unwrap_or("manual relationship update"),
981 )?;
982 self.persist_actor_relationship(&actor_id).await?;
983 info!(
984 actor_id = %actor_id,
985 perspective = %change.perspective,
986 dimension = %change.dimension,
987 delta = change.delta,
988 current = change.current,
989 "relationship updated manually"
990 );
991 self.hooks
992 .on_relationship_change(&actor_id, std::slice::from_ref(&change))
993 .await;
994 Ok(change)
995 }
996
997 pub fn reasoning_config(&self) -> &ReasoningConfig {
998 &self.reasoning_config
999 }
1000
1001 pub fn reflection_config(&self) -> &ReflectionConfig {
1002 &self.reflection_config
1003 }
1004
1005 pub fn with_facts_config(
1008 mut self,
1009 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1010 facts_config: Option<ai_agents_facts::FactsConfig>,
1011 ) -> Self {
1012 self.actor_memory_config = actor_memory_config;
1013 self.facts_config = facts_config;
1014 self
1015 }
1016
1017 pub fn with_facts(
1020 mut self,
1021 store: Arc<ai_agents_facts::FactStore>,
1022 extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1023 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1024 facts_config: Option<ai_agents_facts::FactsConfig>,
1025 ) -> Self {
1026 *self.fact_store.write() = Some(store);
1027 *self.fact_extractor.write() = extractor;
1028 self.actor_memory_config = actor_memory_config;
1029 self.facts_config = facts_config;
1030 self
1031 }
1032
1033 pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1035 self.fact_store.read().clone()
1036 }
1037
1038 pub fn actor_id(&self) -> Option<String> {
1040 self.actor_id.read().clone()
1041 }
1042
1043 pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1045 *self.actor_id.write() = Some(actor_id.to_string());
1046 {
1047 let mut meta = self.session_metadata.write();
1048 meta.actor_id = Some(actor_id.to_string());
1049 if !meta.actors.iter().any(|a| a == actor_id) {
1050 meta.actors.push(actor_id.to_string());
1051 }
1052 }
1053 Ok(())
1054 }
1055
1056 pub fn clear_actor_id(&self) {
1058 *self.actor_id.write() = None;
1059 self.session_metadata.write().actor_id = None;
1060 }
1061
1062 pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1064 self.set_actor_id(user_id)
1065 }
1066
1067 pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1069 let actor_id = match self.effective_actor_id() {
1070 Some(id) => id,
1071 None => return Ok(()),
1072 };
1073
1074 let store_opt = self.fact_store.read().clone();
1075 if let Some(store) = store_opt {
1076 let facts = store.get_facts(&actor_id).await?;
1077 let count = facts.len();
1078 self.actor_facts_cache
1079 .write()
1080 .insert(actor_id.clone(), facts);
1081 self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1082 tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1083 }
1084
1085 Ok(())
1086 }
1087
1088 async fn maybe_load_actor_memory(&self) {
1090 let Some(actor_id) = self.effective_actor_id() else {
1091 return;
1092 };
1093 if self.actor_facts_cache.read().contains_key(&actor_id) {
1094 return;
1095 }
1096 let _ = self.load_actor_memory().await;
1097 }
1098
1099 async fn pre_turn_session_lifecycle(&self) {
1101 if *self.redispatch_depth.read() > 0 {
1102 return;
1103 }
1104 self.resolve_actor_id_from_context();
1105 self.await_background_before_next_turn().await;
1106 self.record_session_actor_if_needed();
1107 self.maybe_load_actor_memory().await;
1108 self.maybe_load_actor_relationship().await;
1109 *self.messages_since_extraction.write() += 1;
1110 }
1111
1112 async fn post_turn_session_lifecycle(&self) -> Result<()> {
1114 if *self.redispatch_depth.read() > 0 {
1115 return Ok(());
1116 }
1117 *self.messages_since_extraction.write() += 1;
1118 self.run_post_turn_maintenance().await
1119 }
1120
1121 fn begin_root_turn(&self) {
1123 if *self.redispatch_depth.read() == 0 {
1124 let mut guard = self.active_turn_context.write();
1125 if guard.is_none() {
1126 self.root_user_message_committed
1127 .store(false, Ordering::SeqCst);
1128 let max_calls = self
1129 .runtime_config
1130 .optimization
1131 .max_speculative_llm_calls_per_turn;
1132 *guard = Some(TurnOptimizationContext::new(
1133 String::new(),
1134 HashMap::new(),
1135 max_calls,
1136 ));
1137 }
1138 }
1139 }
1140
1141 fn update_active_turn_context(
1142 &self,
1143 processed_input: &str,
1144 input_context: HashMap<String, Value>,
1145 ) {
1146 if *self.redispatch_depth.read() > 0 {
1147 return;
1148 }
1149 let max_calls = self
1150 .runtime_config
1151 .optimization
1152 .max_speculative_llm_calls_per_turn;
1153 let mut guard = self.active_turn_context.write();
1154 match guard.as_mut() {
1155 Some(context) => {
1156 context.processed_input = processed_input.to_string();
1157 context.input_context = input_context;
1158 context.max_speculative_llm_calls = max_calls;
1159 }
1160 None => {
1161 *guard = Some(TurnOptimizationContext::new(
1162 processed_input,
1163 input_context,
1164 max_calls,
1165 ));
1166 }
1167 }
1168 }
1169
1170 async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1172 if *self.redispatch_depth.read() > 0 {
1173 return Ok(());
1174 }
1175 if !self
1176 .root_user_message_committed
1177 .swap(true, Ordering::SeqCst)
1178 {
1179 self.memory
1180 .add_message(ChatMessage::user(processed_input))
1181 .await?;
1182 if let Some(context) = self.active_turn_context.write().as_mut() {
1183 context.mark_user_message_committed();
1184 }
1185 }
1186 Ok(())
1187 }
1188
1189 fn end_root_turn(&self) {
1191 if *self.redispatch_depth.read() == 0 {
1192 self.root_user_message_committed
1193 .store(false, Ordering::SeqCst);
1194 *self.active_turn_context.write() = None;
1195 }
1196 }
1197
1198 fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1199 self.begin_root_turn();
1200 let mut guard = self.active_turn_context.write();
1201 let Some(context) = guard.as_mut() else {
1202 return false;
1203 };
1204 context.reserve_speculative_llm_call_for(kind)
1205 }
1206
1207 fn branch_context_preview(&self) -> String {
1208 let context = self.build_context_with_overlays();
1209 let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1210 const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1211 if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1212 value = value
1213 .chars()
1214 .take(MAX_CONTEXT_PREVIEW_CHARS)
1215 .collect::<String>();
1216 value.push_str("...");
1217 }
1218 value
1219 }
1220
1221 async fn await_background_before_next_turn(&self) {
1223 let optimization = &self.runtime_config.optimization;
1224 if !optimization.enabled {
1225 return;
1226 }
1227 let actor_id = self.effective_actor_id();
1228 let post = &optimization.post_turn;
1229 self.await_background_task(
1230 post.facts.await_before_next_turn,
1231 RuntimeTaskPurpose::PostTurnFacts,
1232 actor_id.as_deref(),
1233 "facts",
1234 )
1235 .await;
1236 self.await_background_task(
1237 post.relationships.await_before_next_turn,
1238 RuntimeTaskPurpose::PostTurnRelationship,
1239 actor_id.as_deref(),
1240 "relationships",
1241 )
1242 .await;
1243 }
1244
1245 async fn await_background_task(
1246 &self,
1247 policy: AwaitBeforeNextTurn,
1248 purpose: RuntimeTaskPurpose,
1249 actor_id: Option<&str>,
1250 label: &str,
1251 ) {
1252 match policy {
1253 AwaitBeforeNextTurn::Never => {}
1254 AwaitBeforeNextTurn::Always => {
1255 if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1256 warn!(label = label, error = %error, "background maintenance flush failed");
1257 }
1258 }
1259 AwaitBeforeNextTurn::SameActor => {
1260 if let Some(actor_id) = actor_id
1261 && let Err(error) = self
1262 .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1263 .await
1264 {
1265 warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1266 }
1267 }
1268 }
1269 }
1270
1271 async fn run_post_turn_maintenance(&self) -> Result<()> {
1273 let optimization = &self.runtime_config.optimization;
1274 if !optimization.enabled {
1275 self.auto_extract_facts().await;
1276 self.auto_update_relationship().await;
1277 return Ok(());
1278 }
1279
1280 let facts_mode = effective_maintenance_mode(
1281 optimization.post_turn.facts.mode,
1282 optimization.parallel_post_turn_memory,
1283 );
1284 let relationships_mode = effective_maintenance_mode(
1285 optimization.post_turn.relationships.mode,
1286 optimization.parallel_post_turn_memory,
1287 );
1288
1289 match (facts_mode, relationships_mode) {
1290 (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1291 self.auto_extract_facts().await;
1292 self.auto_update_relationship().await;
1293 }
1294 (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1295 let facts = self.auto_extract_facts();
1296 let relationships = self.auto_update_relationship();
1297 tokio::join!(facts, relationships);
1298 }
1299 (MaintenanceMode::Background, MaintenanceMode::Background) => {
1300 self.schedule_facts_background().await?;
1301 self.schedule_relationship_background().await?;
1302 }
1303 (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1304 | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1305 self.schedule_facts_background().await?;
1306 self.auto_update_relationship().await;
1307 }
1308 (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1309 | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1310 self.auto_extract_facts().await;
1311 self.schedule_relationship_background().await?;
1312 }
1313 _ => {
1314 self.auto_extract_facts().await;
1315 self.auto_update_relationship().await;
1316 }
1317 }
1318 Ok(())
1319 }
1320
1321 async fn schedule_facts_background(&self) -> Result<()> {
1322 let policy = self.runtime_config.optimization.post_turn.facts.clone();
1323 let should_extract = self
1324 .facts_config
1325 .as_ref()
1326 .map(|c| c.enabled && c.auto_extract)
1327 .unwrap_or(false);
1328 if !should_extract {
1329 return Ok(());
1330 }
1331 let msgs_since = *self.messages_since_extraction.read();
1332 if msgs_since < 2 {
1333 return Ok(());
1334 }
1335 let Some(actor_id) = self.effective_actor_id() else {
1336 self.record_skipped_maintenance(
1337 "facts",
1338 ObservationPurpose::FactsExtraction,
1339 "missing_actor",
1340 Some(&policy),
1341 );
1342 return Ok(());
1343 };
1344 let Some(extractor) = self.fact_extractor.read().clone() else {
1345 return Ok(());
1346 };
1347 let messages = match self.memory.get_messages(None).await {
1348 Ok(messages) => messages,
1349 Err(error) => {
1350 warn!(error = %error, "failed to snapshot messages for fact extraction");
1351 return Ok(());
1352 }
1353 };
1354 let recent: Vec<_> = messages
1355 .iter()
1356 .rev()
1357 .take(msgs_since)
1358 .rev()
1359 .cloned()
1360 .collect();
1361 if recent.is_empty() {
1362 return Ok(());
1363 }
1364 let existing = self
1365 .actor_facts_cache
1366 .read()
1367 .get(&actor_id)
1368 .cloned()
1369 .unwrap_or_default();
1370 let categories = self
1371 .facts_config
1372 .as_ref()
1373 .map(|c| c.custom_categories.clone())
1374 .unwrap_or_default();
1375 let store = self.fact_store.read().clone();
1376 let cache = Arc::clone(&self.actor_facts_cache);
1377 let counter = Arc::clone(&self.messages_since_extraction);
1378 let hooks = Arc::clone(&self.hooks);
1379 let agent_id = self.info.id.clone();
1380 let observation = current_observation_context();
1381 let key = MaintenanceSequenceKey::actor(
1382 agent_id,
1383 actor_id.clone(),
1384 RuntimeTaskPurpose::PostTurnFacts,
1385 );
1386 let actor_for_task = actor_id.clone();
1387 let task = async move {
1388 let run = async move {
1389 let facts = extractor
1390 .extract(&recent, &existing, Some(&actor_for_task), &categories)
1391 .await?;
1392 if !facts.is_empty() {
1393 if let Some(store) = store {
1394 let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1395 cache.write().insert(actor_for_task.clone(), authoritative);
1396 } else {
1397 cache
1398 .write()
1399 .entry(actor_for_task.clone())
1400 .or_default()
1401 .extend(facts.clone());
1402 }
1403 {
1404 let mut count = counter.write();
1405 if *count <= msgs_since {
1406 *count = 0;
1407 } else {
1408 *count -= msgs_since;
1409 }
1410 }
1411 hooks.on_facts_extracted(&actor_for_task, &facts).await;
1412 }
1413 Ok(())
1414 };
1415 if let Some(context) = observation {
1416 with_observation_context(
1417 context.with_purpose(ObservationPurpose::FactsExtraction),
1418 run,
1419 )
1420 .await
1421 } else {
1422 run.await
1423 }
1424 };
1425 self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1426 .await
1427 }
1428
1429 async fn schedule_relationship_background(&self) -> Result<()> {
1430 let policy = self
1431 .runtime_config
1432 .optimization
1433 .post_turn
1434 .relationships
1435 .clone();
1436 let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1437 return Ok(());
1438 };
1439 let Some(actor_id) = self.effective_actor_id() else {
1440 self.record_skipped_maintenance(
1441 "relationships",
1442 ObservationPurpose::RelationshipUpdate,
1443 "missing_actor",
1444 Some(&policy),
1445 );
1446 return Ok(());
1447 };
1448 let recent_messages = manager.config().auto_update.recent_messages;
1449 let messages = match self.memory.get_messages(Some(recent_messages)).await {
1450 Ok(messages) => messages,
1451 Err(error) => {
1452 warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1453 return Ok(());
1454 }
1455 };
1456 let storage = self.storage.read().clone();
1457 let hooks = Arc::clone(&self.hooks);
1458 let agent_id = self.info.id.clone();
1459 let observation = current_observation_context();
1460 let key = MaintenanceSequenceKey::actor(
1461 agent_id.clone(),
1462 actor_id.clone(),
1463 RuntimeTaskPurpose::PostTurnRelationship,
1464 );
1465 let actor_for_task = actor_id.clone();
1466 let task = async move {
1467 let run = async move {
1468 if manager.config().auto_update.enabled {
1469 let update = manager.auto_update(&actor_for_task, &messages).await?;
1470 if !update.changes.is_empty() {
1471 hooks
1472 .on_relationship_change(&actor_for_task, &update.changes)
1473 .await;
1474 }
1475 if let Some(ref event) = update.event {
1476 hooks.on_notable_event(&actor_for_task, event).await;
1477 }
1478 }
1479 if manager.config().persistence.enabled
1480 && let (Some(storage), Some(value)) =
1481 (storage, manager.relationship_as_value(&actor_for_task)?)
1482 {
1483 storage
1484 .save_relationship(&agent_id, &actor_for_task, &value)
1485 .await?;
1486 }
1487 Ok(())
1488 };
1489 if let Some(context) = observation {
1490 with_observation_context(
1491 context.with_purpose(ObservationPurpose::RelationshipUpdate),
1492 run,
1493 )
1494 .await
1495 } else {
1496 run.await
1497 }
1498 };
1499 self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1500 .await
1501 }
1502
1503 async fn spawn_or_handle_background<F>(
1505 &self,
1506 key: Option<MaintenanceSequenceKey>,
1507 task: F,
1508 label: &'static str,
1509 policy: &crate::optimization::config::MaintenanceTaskPolicy,
1510 ) -> Result<()>
1511 where
1512 F: Future<Output = Result<()>> + Send + 'static,
1513 {
1514 if self.background_maintenance.is_full() {
1515 match self
1516 .runtime_config
1517 .optimization
1518 .post_turn
1519 .on_background_overflow
1520 {
1521 BackgroundOverflowPolicy::RunInline => {
1522 record_background_maintenance_event(
1523 self.observability_manager.as_ref(),
1524 label,
1525 EventStatus::Success,
1526 0,
1527 "inline_overflow",
1528 None,
1529 Some(policy),
1530 );
1531 let start = Instant::now();
1532 match task.await {
1533 Ok(()) => record_background_maintenance_event(
1534 self.observability_manager.as_ref(),
1535 label,
1536 EventStatus::Success,
1537 start.elapsed().as_millis() as u64,
1538 "inline_completed",
1539 None,
1540 Some(policy),
1541 ),
1542 Err(error) => {
1543 warn!(label = label, error = %error, "inline maintenance fallback failed");
1544 record_background_maintenance_event(
1545 self.observability_manager.as_ref(),
1546 label,
1547 EventStatus::Error,
1548 start.elapsed().as_millis() as u64,
1549 "inline_failed",
1550 Some(error.to_string()),
1551 Some(policy),
1552 );
1553 return Err(error);
1554 }
1555 }
1556 }
1557 BackgroundOverflowPolicy::Drop => {
1558 self.record_skipped_maintenance(
1559 label,
1560 ObservationPurpose::Other(label.to_string()),
1561 "queue_full",
1562 Some(policy),
1563 );
1564 }
1565 BackgroundOverflowPolicy::Error => {
1566 record_background_maintenance_event(
1567 self.observability_manager.as_ref(),
1568 label,
1569 EventStatus::Error,
1570 0,
1571 "queue_full",
1572 None,
1573 Some(policy),
1574 );
1575 warn!(label = label, "background maintenance queue full");
1576 return Err(AgentError::Other(format!(
1577 "background maintenance queue is full for {}",
1578 label
1579 )));
1580 }
1581 }
1582 return Ok(());
1583 }
1584
1585 record_background_maintenance_event(
1586 self.observability_manager.as_ref(),
1587 label,
1588 EventStatus::Success,
1589 0,
1590 "scheduled",
1591 None,
1592 Some(policy),
1593 );
1594 let manager = self.observability_manager.clone();
1595 let policy_for_task = policy.clone();
1596 let observed_task = async move {
1597 let start = Instant::now();
1598 let result = task.await;
1599 match &result {
1600 Ok(()) => record_background_maintenance_event(
1601 manager.as_ref(),
1602 label,
1603 EventStatus::Success,
1604 start.elapsed().as_millis() as u64,
1605 "completed",
1606 None,
1607 Some(&policy_for_task),
1608 ),
1609 Err(error) => record_background_maintenance_event(
1610 manager.as_ref(),
1611 label,
1612 EventStatus::Error,
1613 start.elapsed().as_millis() as u64,
1614 "failed",
1615 Some(error.to_string()),
1616 Some(&policy_for_task),
1617 ),
1618 }
1619 result
1620 };
1621
1622 if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1623 record_background_maintenance_event(
1624 self.observability_manager.as_ref(),
1625 label,
1626 EventStatus::Error,
1627 0,
1628 "spawn_failed",
1629 Some(error.to_string()),
1630 Some(policy),
1631 );
1632 warn!(label = label, error = %error, "background maintenance spawn failed");
1633 return Err(error);
1634 }
1635 Ok(())
1636 }
1637
1638 fn record_skipped_maintenance(
1640 &self,
1641 label: &str,
1642 purpose: ObservationPurpose,
1643 reason: &str,
1644 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1645 ) {
1646 if let Some(manager) = self.observability_manager.as_ref() {
1647 let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1648 tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1649 manager.record_lifecycle_event(
1650 EventType::MemoryOperation {
1651 operation: format!("{}_maintenance", label),
1652 },
1653 purpose,
1654 EventStatus::Skipped,
1655 0,
1656 tags,
1657 None,
1658 );
1659 }
1660 }
1661
1662 pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1664 let Some(actor_id) = self.effective_actor_id() else {
1665 return Vec::new();
1666 };
1667 self.actor_facts_cache
1668 .read()
1669 .get(&actor_id)
1670 .cloned()
1671 .unwrap_or_default()
1672 }
1673
1674 pub fn relationship_memory_text(&self) -> Option<String> {
1676 self.format_relationship_for_context().map(|(_, text)| text)
1677 }
1678
1679 pub async fn extract_facts(
1681 &self,
1682 last_n: usize,
1683 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1684 self.extract_facts_with_source(last_n, "manual").await
1685 }
1686
1687 async fn extract_facts_with_source(
1688 &self,
1689 last_n: usize,
1690 source: &'static str,
1691 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1692 let extractor = match self.fact_extractor.read().clone() {
1693 Some(e) => e,
1694 None => return Ok(vec![]),
1695 };
1696
1697 let messages = self.memory.get_messages(None).await?;
1698 let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1699
1700 if recent.is_empty() {
1701 return Ok(vec![]);
1702 }
1703
1704 let actor_id = self.effective_actor_id();
1705 let existing = actor_id
1706 .as_ref()
1707 .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1708 .unwrap_or_default();
1709
1710 let categories = self
1711 .facts_config
1712 .as_ref()
1713 .map(|c| c.custom_categories.clone())
1714 .unwrap_or_default();
1715
1716 let facts = self
1717 .observe_purpose(
1718 ObservationPurpose::FactsExtraction,
1719 extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1720 )
1721 .await?;
1722
1723 if !facts.is_empty() {
1725 let fact_store_opt = self.fact_store.read().clone();
1726 let mut stored_total = 0usize;
1727 let mut cache_updated = false;
1728 if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1729 let authoritative = store.add_facts(aid, facts.clone()).await?;
1731 stored_total = authoritative.len();
1732 self.actor_facts_cache
1733 .write()
1734 .insert(aid.clone(), authoritative);
1735 cache_updated = true;
1736 } else if let Some(aid) = &actor_id {
1737 let mut cache = self.actor_facts_cache.write();
1738 let entry = cache.entry(aid.clone()).or_default();
1739 entry.extend(facts.clone());
1740 stored_total = entry.len();
1741 cache_updated = true;
1742 }
1743
1744 info!(
1745 actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1746 source = source,
1747 requested_messages = last_n,
1748 message_count = recent.len(),
1749 extracted_count = facts.len(),
1750 cache_updated = cache_updated,
1751 stored_total = stored_total,
1752 "facts extracted"
1753 );
1754
1755 if let Some(ref aid) = actor_id {
1756 self.hooks.on_facts_extracted(aid, &facts).await;
1757 }
1758 }
1759
1760 Ok(facts)
1761 }
1762
1763 fn resolve_actor_id_from_context(&self) {
1766 if self
1767 .current_turn_actor_context()
1768 .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1769 .is_some()
1770 {
1771 return;
1772 }
1773
1774 if let Some(ref am_config) = self.actor_memory_config
1775 && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1776 && let Some(ref path) = am_config.identification.context_path
1777 {
1778 let val = self
1780 .context_manager
1781 .get_path(path)
1782 .or_else(|| self.context_manager.get(path));
1783 if let Some(val) = val
1784 && let Some(id_str) = val.as_str()
1785 {
1786 let current = self.actor_id.read().clone();
1787 if current.as_deref() != Some(id_str) {
1788 *self.actor_id.write() = Some(id_str.to_string());
1789 let mut meta = self.session_metadata.write();
1790 meta.actor_id = Some(id_str.to_string());
1791 if !meta.actors.iter().any(|a| a == id_str) {
1792 meta.actors.push(id_str.to_string());
1793 }
1794 }
1795 }
1796 }
1797 }
1798
1799 fn format_actor_facts_for_context(&self) -> String {
1801 let should_inject = self
1803 .facts_config
1804 .as_ref()
1805 .map(|c| c.inject_in_context)
1806 .unwrap_or(true);
1807 if !should_inject {
1808 return String::new();
1809 }
1810
1811 let Some(actor_id) = self.effective_actor_id() else {
1812 return String::new();
1813 };
1814
1815 let facts = self
1816 .actor_facts_cache
1817 .read()
1818 .get(&actor_id)
1819 .cloned()
1820 .unwrap_or_default();
1821 if facts.is_empty() {
1822 return String::new();
1823 }
1824
1825 let am_config = self.actor_memory_config.as_ref();
1826 let facts_budget = self
1829 .memory_token_budget
1830 .as_ref()
1831 .map(|b| b.allocation.facts as usize)
1832 .filter(|n| *n > 0);
1833 let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1834 let max_tokens = facts_budget.unwrap_or(default_max);
1835
1836 let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1838 if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1839 return String::new();
1840 }
1841 if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1842 && !cfg.injection.categories.is_empty()
1843 {
1844 facts
1845 .iter()
1846 .filter(|f| {
1847 cfg.injection
1848 .categories
1849 .iter()
1850 .any(|c| f.category.to_string() == *c)
1851 })
1852 .cloned()
1853 .collect()
1854 } else {
1855 facts.clone()
1856 }
1857 } else {
1858 facts.clone()
1859 };
1860
1861 if filtered.is_empty() {
1862 return String::new();
1863 }
1864
1865 if let Some(store) = self.fact_store.read().clone() {
1866 store.format_for_context(&filtered, max_tokens)
1867 } else {
1868 String::new()
1869 }
1870 }
1871
1872 fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
1873 let context = self.build_context_with_overlays();
1874 let mut root = Value::Object(context.into_iter().collect());
1875 for (path, value) in staged {
1876 if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
1877 root = updated;
1878 }
1879 }
1880 match root {
1881 Value::Object(obj) => obj.into_iter().collect(),
1882 _ => HashMap::new(),
1883 }
1884 }
1885
1886 fn build_context_with_overlays(&self) -> HashMap<String, Value> {
1887 let mut context = self.context_manager.get_all();
1888 let mut root = Value::Object(context.clone().into_iter().collect());
1889
1890 if let Some(turn_ctx) = self.current_turn_actor_context() {
1891 if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
1892 && let Ok(updated) = ai_agents_core::set_dot_path(
1893 root.clone(),
1894 "interaction.origin_actor_id",
1895 serde_json::json!(origin_actor_id),
1896 )
1897 {
1898 root = updated;
1899 }
1900 if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
1901 && let Ok(updated) = ai_agents_core::set_dot_path(
1902 root.clone(),
1903 "interaction.sender_agent_id",
1904 serde_json::json!(sender_agent_id),
1905 )
1906 {
1907 root = updated;
1908 }
1909 }
1910
1911 if let Some(ref actor_id) = self.effective_actor_id()
1912 && let Ok(updated) = ai_agents_core::set_dot_path(
1913 root.clone(),
1914 "interaction.actor_id",
1915 serde_json::json!(actor_id),
1916 )
1917 {
1918 root = updated;
1919 }
1920
1921 if let Some(manager) = self.relationship_manager.as_ref()
1922 && let Some(actor_id) = self.effective_actor_id()
1923 && let Some(value) = manager.to_context_value(&actor_id)
1924 && let Ok(updated) = ai_agents_core::set_dot_path(
1925 root.clone(),
1926 &manager.config().injection.context_path,
1927 value,
1928 )
1929 {
1930 root = updated;
1931 }
1932
1933 if let Value::Object(obj) = root {
1934 context = obj.into_iter().collect();
1935 }
1936
1937 context
1938 }
1939
1940 fn resolve_actor_name_from_context(&self) -> Option<String> {
1941 for path in ["actor.name", "user.name", "player.name", "customer.name"] {
1942 if let Some(value) = self.context_manager.get_path(path)
1943 && let Some(name) = value.as_str()
1944 {
1945 return Some(name.to_string());
1946 }
1947 }
1948 None
1949 }
1950
1951 async fn maybe_load_actor_relationship(&self) {
1952 let Some(manager) = self.relationship_manager.as_ref() else {
1953 return;
1954 };
1955 let Some(actor_id) = self.effective_actor_id() else {
1956 return;
1957 };
1958
1959 let mut should_fire_loaded = false;
1960 if manager.get(&actor_id).is_none() {
1961 let mut loaded = false;
1962 if manager.config().persistence.enabled {
1963 let storage = self.storage.read().clone();
1964 if let Some(storage) = storage {
1965 match storage.load_relationship(&self.info.id, &actor_id).await {
1966 Ok(Some(value)) => match manager.insert_from_value(value) {
1967 Ok(_) => loaded = true,
1968 Err(e) => {
1969 warn!(actor = %actor_id, error = %e, "failed to restore relationship")
1970 }
1971 },
1972 Ok(None) => {}
1973 Err(e) => {
1974 warn!(actor = %actor_id, error = %e, "failed to load relationship")
1975 }
1976 }
1977 }
1978 }
1979
1980 if !loaded {
1981 manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
1982 }
1983 should_fire_loaded = true;
1984 }
1985
1986 let actor_name = self.resolve_actor_name_from_context();
1987 let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
1988 if should_fire_loaded {
1989 self.hooks
1990 .on_relationship_loaded(&actor_id, &relationship)
1991 .await;
1992 }
1993 }
1994
1995 fn format_relationship_for_context(&self) -> Option<(String, String)> {
1996 let manager = self.relationship_manager.as_ref()?;
1997 if !manager.config().injection.enabled {
1998 return None;
1999 }
2000 let actor_id = self.effective_actor_id()?;
2001 let relationship = manager.get(&actor_id)?;
2002 let local_cap = manager.config().injection.max_tokens;
2003 let global_cap = self
2004 .memory_token_budget
2005 .as_ref()
2006 .map(|b| b.allocation.relationships as usize)
2007 .filter(|n| *n > 0);
2008 let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
2009 let text = ai_agents_relationships::format_relationship(
2010 &relationship,
2011 &manager.config().injection.format,
2012 max_tokens,
2013 );
2014 if text.is_empty() {
2015 None
2016 } else {
2017 Some((manager.config().injection.prompt_variable.clone(), text))
2018 }
2019 }
2020
2021 async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2022 let Some(manager) = self.relationship_manager.as_ref() else {
2023 return Ok(());
2024 };
2025 if !manager.config().persistence.enabled {
2026 return Ok(());
2027 }
2028 let storage = self.storage.read().clone();
2029 let Some(storage) = storage else {
2030 return Ok(());
2031 };
2032 if let Some(value) = manager.relationship_as_value(actor_id)? {
2033 storage
2034 .save_relationship(&self.info.id, actor_id, &value)
2035 .await?;
2036 }
2037 Ok(())
2038 }
2039
2040 async fn auto_update_relationship(&self) {
2041 let Some(manager) = self.relationship_manager.as_ref() else {
2042 return;
2043 };
2044 let Some(actor_id) = self.effective_actor_id() else {
2045 return;
2046 };
2047 if !manager.config().auto_update.enabled {
2048 let _ = self.persist_actor_relationship(&actor_id).await;
2049 return;
2050 }
2051
2052 let recent_messages = manager.config().auto_update.recent_messages;
2053 let messages = match self.memory.get_messages(Some(recent_messages)).await {
2054 Ok(messages) => messages,
2055 Err(e) => {
2056 warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2057 return;
2058 }
2059 };
2060
2061 match self
2062 .observe_purpose(
2063 ObservationPurpose::RelationshipUpdate,
2064 manager.auto_update(&actor_id, &messages),
2065 )
2066 .await
2067 {
2068 Ok(update) => {
2069 if !update.changes.is_empty() {
2070 self.hooks
2071 .on_relationship_change(&actor_id, &update.changes)
2072 .await;
2073 }
2074 if let Some(ref event) = update.event {
2075 self.hooks.on_notable_event(&actor_id, event).await;
2076 }
2077 let persisted = match self.persist_actor_relationship(&actor_id).await {
2078 Ok(()) => true,
2079 Err(e) => {
2080 warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2081 false
2082 }
2083 };
2084 if !update.changes.is_empty() || update.event.is_some() {
2085 let changed_dimensions: Vec<String> = update
2086 .changes
2087 .iter()
2088 .map(|change| format!("{}:{}", change.perspective, change.dimension))
2089 .collect();
2090 info!(
2091 actor_id = %actor_id,
2092 change_count = update.changes.len(),
2093 changed_dimensions = ?changed_dimensions,
2094 event_present = update.event.is_some(),
2095 persisted = persisted,
2096 "relationship updated"
2097 );
2098 } else {
2099 debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2100 }
2101 }
2102 Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2103 }
2104 }
2105
2106 async fn auto_extract_facts(&self) {
2108 let should_extract = self
2109 .facts_config
2110 .as_ref()
2111 .map(|c| c.enabled && c.auto_extract)
2112 .unwrap_or(false);
2113
2114 if !should_extract {
2115 debug!("fact extraction skipped because auto extraction is disabled");
2116 return;
2117 }
2118
2119 let msgs_since = *self.messages_since_extraction.read();
2120 if msgs_since < 2 {
2121 debug!(
2122 messages_since_extraction = msgs_since,
2123 "fact extraction skipped until threshold is reached"
2124 );
2125 return;
2126 }
2127
2128 match self.extract_facts_with_source(msgs_since, "auto").await {
2129 Ok(facts) => {
2130 if !facts.is_empty() {
2131 *self.messages_since_extraction.write() = 0;
2132 } else {
2133 debug!("fact extraction ran but found no new facts");
2134 }
2135 }
2136 Err(e) => {
2137 warn!("fact extraction failed: {}", e);
2138 }
2139 }
2140 }
2141
2142 pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2143 self.persona_manager = Some(manager);
2144 self
2145 }
2146
2147 pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2148 self.persona_manager.as_ref()
2149 }
2150
2151 pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2152 if config.is_enabled() {
2153 let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2154 .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2155 self.disambiguation_manager = Some(manager);
2156 }
2157 self
2158 }
2159
2160 pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2161 self.disambiguation_manager.as_ref()
2162 }
2163
2164 pub fn has_disambiguation(&self) -> bool {
2165 self.disambiguation_manager
2166 .as_ref()
2167 .is_some_and(|m| m.is_enabled())
2168 }
2169
2170 pub async fn init_storage(&self) -> Result<()> {
2171 let _guard = self.storage_init.lock().await;
2175 let mut storage = self.storage.read().clone();
2176 if storage.is_none() && !self.storage_config.is_none() {
2177 let storage_config = self.convert_storage_config();
2178 storage = create_storage(&storage_config).await?;
2179 *self.storage.write() = storage.clone();
2180 }
2181
2182 self.validate_storage_requirements(storage.as_deref())?;
2183 self.complete_facts_init().await;
2184 Ok(())
2185 }
2186
2187 fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2188 let facts_required = self
2189 .facts_config
2190 .as_ref()
2191 .is_some_and(|config| config.enabled)
2192 || self
2193 .actor_memory_config
2194 .as_ref()
2195 .is_some_and(|config| config.enabled);
2196 let relationships_required = self
2197 .relationship_manager
2198 .as_ref()
2199 .is_some_and(|manager| manager.config().persistence.enabled);
2200
2201 let Some(storage) = storage else {
2202 let mut requirements = Vec::new();
2203 if facts_required {
2204 requirements.push("actor facts or actor memory");
2205 }
2206 if relationships_required {
2207 requirements.push("persistent relationships");
2208 }
2209 if requirements.is_empty() {
2210 return Ok(());
2211 }
2212 return Err(AgentError::Config(format!(
2213 "Storage is required for enabled {} but none is configured or injected",
2214 requirements.join(" and ")
2215 )));
2216 };
2217
2218 if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2222 return Err(AgentError::UnsupportedStorageCapability(
2223 StorageCapability::ActorFacts,
2224 ));
2225 }
2226 if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2227 return Err(AgentError::UnsupportedStorageCapability(
2228 StorageCapability::ActorRelationships,
2229 ));
2230 }
2231 Ok(())
2232 }
2233
2234 async fn complete_facts_init(&self) {
2237 if self.fact_store.read().is_some() {
2238 return;
2239 }
2240 let storage = match self.storage.read().clone() {
2241 Some(s) => s,
2242 None => return,
2243 };
2244
2245 let facts_enabled = self
2246 .facts_config
2247 .as_ref()
2248 .map(|f| f.enabled)
2249 .unwrap_or(false);
2250 let actor_memory_enabled = self
2251 .actor_memory_config
2252 .as_ref()
2253 .map(|a| a.enabled)
2254 .unwrap_or(false);
2255
2256 if !facts_enabled && !actor_memory_enabled {
2257 return;
2258 }
2259
2260 let fc = self.facts_config.clone().unwrap_or_default();
2261 let store = Arc::new(ai_agents_facts::FactStore::new(
2262 storage,
2263 self.info.id.clone(),
2264 fc.clone(),
2265 ));
2266
2267 let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2268 let extractor_llm = fc
2269 .extractor_llm
2270 .as_ref()
2271 .and_then(|alias| self.llm_registry.get(alias).ok())
2272 .or_else(|| self.llm_registry.router().ok())
2273 .or_else(|| self.llm_registry.default().ok());
2274 extractor_llm.map(|llm| {
2275 Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2276 as Arc<dyn ai_agents_facts::FactExtractor>
2277 })
2278 } else {
2279 None
2280 };
2281
2282 *self.fact_store.write() = Some(store);
2283 *self.fact_extractor.write() = extractor;
2284 debug!(
2285 agent = %self.info.id,
2286 facts_enabled,
2287 actor_memory_enabled,
2288 "facts storage initialized"
2289 );
2290 }
2291
2292 fn convert_storage_config(&self) -> StorageStorageConfig {
2293 crate::spec::storage::to_storage_config(&self.storage_config)
2294 }
2295
2296 pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2297 self.storage.read().clone()
2298 }
2299
2300 pub fn storage_config(&self) -> &StorageConfig {
2301 &self.storage_config
2302 }
2303
2304 pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2306 self.spawner.as_ref()
2307 }
2308
2309 pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2311 self.spawner_registry.as_ref()
2312 }
2313
2314 pub fn has_spawner(&self) -> bool {
2315 self.spawner_registry.is_some()
2316 }
2317
2318 pub fn with_spawner_handles(
2319 mut self,
2320 spawner: Arc<crate::spawner::AgentSpawner>,
2321 registry: Arc<crate::spawner::AgentRegistry>,
2322 ) -> Self {
2323 self.spawner = Some(spawner);
2324 self.spawner_registry = Some(registry);
2325 self
2326 }
2327
2328 pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2329 self.hooks = hooks;
2330 self
2331 }
2332
2333 pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2334 self.parallel_tools = config;
2335 self
2336 }
2337
2338 pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2339 self.streaming = config;
2340 self
2341 }
2342
2343 pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2344 self.hitl_engine = Some(engine);
2345 self.approval_handler = handler;
2346 self
2347 }
2348
2349 pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2350 self.max_context_tokens = tokens;
2351 self
2352 }
2353
2354 pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2355 self.memory_token_budget = Some(budget);
2356 self
2357 }
2358
2359 pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2360 self.recovery_manager = manager;
2361 self
2362 }
2363
2364 pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2365 self.tool_security = engine;
2366 self
2367 }
2368
2369 pub fn runtime_control(&self) -> RuntimeControlHandle {
2371 RuntimeControlHandle {
2372 state: Arc::clone(&self.runtime_control),
2373 }
2374 }
2375
2376 pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2378 self.tools.set_question_handler(handler);
2379 }
2380
2381 pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2383 self.tools.set_diagnostics_provider(provider);
2384 }
2385
2386 pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2388 self.tools.set_command_runner(runner);
2389 }
2390
2391 pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2393 self.tools.set_web_search_provider(provider);
2394 }
2395
2396 pub fn todos(&self) -> Vec<TodoItem> {
2398 self.tools.todos()
2399 }
2400
2401 fn active_tool_security(&self) -> ToolSecurityEngine {
2403 self.runtime_control
2404 .tool_security_override
2405 .read()
2406 .clone()
2407 .unwrap_or_else(|| self.tool_security.clone())
2408 }
2409
2410 fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2412 let _guard = self.runtime_control.snapshot_guard.read();
2413 RuntimeSafetySnapshot {
2414 version: self.runtime_control.version.load(Ordering::SeqCst),
2415 emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2416 tool_security: self
2417 .runtime_control
2418 .tool_security_override
2419 .read()
2420 .clone()
2421 .unwrap_or_else(|| self.tool_security.clone()),
2422 tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2423 }
2424 }
2425
2426 fn admit_tool_execution(
2428 &self,
2429 expected_runtime_version: u64,
2430 expected_policy_version: u64,
2431 expected_state_generation: Option<u64>,
2432 canonical_id: &str,
2433 ) -> SecurityCheckResult {
2434 let _guard = self.runtime_control.snapshot_guard.read();
2435 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2436 return SecurityCheckResult::Block {
2437 reason: "runtime emergency deny is enabled".to_string(),
2438 };
2439 }
2440 let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2441 let security_engine = self
2442 .runtime_control
2443 .tool_security_override
2444 .read()
2445 .clone()
2446 .unwrap_or_else(|| self.tool_security.clone());
2447 if runtime_version != expected_runtime_version
2448 || security_engine.policy_version() != expected_policy_version
2449 {
2450 return SecurityCheckResult::Block {
2451 reason: "runtime safety controls changed before admission".to_string(),
2452 };
2453 }
2454 let current_state_generation = self
2455 .state_machine
2456 .as_ref()
2457 .map(|state_machine| state_machine.generation());
2458 if current_state_generation != expected_state_generation {
2459 return SecurityCheckResult::Block {
2460 reason: "state scope changed before admission".to_string(),
2461 };
2462 }
2463 security_engine.admit_tool_execution(canonical_id)
2464 }
2465
2466 pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2467 let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2468 self.process_processor = Some(processor);
2469 self
2470 }
2471
2472 pub fn with_state_machine(
2473 mut self,
2474 state_machine: Arc<StateMachine>,
2475 evaluator: Arc<dyn TransitionEvaluator>,
2476 ) -> Self {
2477 self.state_machine = Some(state_machine);
2478 self.transition_evaluator = Some(evaluator);
2479 self
2480 }
2481
2482 pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2483 self.context_manager = manager;
2484 self
2485 }
2486
2487 pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2488 self.message_filters.write().insert(name.into(), filter);
2489 }
2490
2491 pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2492 self.context_manager.update(key, value)
2493 }
2494
2495 pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2496 self.context_manager.update(path, value)
2497 }
2498
2499 pub fn get_context(&self) -> HashMap<String, Value> {
2500 self.build_context_with_overlays()
2501 }
2502
2503 pub fn remove_context(&self, key: &str) -> Option<Value> {
2504 self.context_manager.remove(key)
2505 }
2506
2507 pub async fn refresh_context(&self, key: &str) -> Result<()> {
2508 self.context_manager.refresh(key).await
2509 }
2510
2511 pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2512 self.context_manager.register_provider(name, provider);
2513 }
2514
2515 pub fn current_state(&self) -> Option<String> {
2516 self.state_machine.as_ref().map(|sm| sm.current())
2517 }
2518
2519 async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2521 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2522 let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2523 return;
2524 };
2525 if disambiguator.has_pending_confirmation().await {
2526 disambiguator.clear_pending().await;
2527 *self.pending_skill_id.write() = None;
2528 info!(
2529 confirmation_event = "invalidated",
2530 invalidation_reason = reason,
2531 "Runtime invalidated pending confirmation"
2532 );
2533 }
2534 }
2535
2536 async fn admit_disambiguation_redispatch(
2538 &self,
2539 expected_epoch: u64,
2540 expected_state_generation: Option<u64>,
2541 ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2542 let admission = self.disambiguation_admission.read().await;
2543 let state_generation = self
2544 .state_machine
2545 .as_ref()
2546 .map(|state_machine| state_machine.generation());
2547 if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2548 || state_generation != expected_state_generation
2549 {
2550 return Err(AgentError::Other(
2551 "Disambiguation ownership changed before redispatch admission".to_string(),
2552 ));
2553 }
2554 Ok(admission)
2555 }
2556
2557 fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2559 self.state_transition_reserved
2560 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2561 .ok()
2562 .map(|_| StateTransitionReservation {
2563 reserved: &self.state_transition_reserved,
2564 })
2565 }
2566
2567 async fn admit_optional_disambiguation_ownership(
2569 &self,
2570 ownership: Option<DisambiguationOwnership>,
2571 ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2572 match ownership {
2573 Some(ownership) => self
2574 .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2575 .await
2576 .map(Some),
2577 None => Ok(None),
2578 }
2579 }
2580
2581 pub async fn transition_to(&self, state: &str) -> Result<()> {
2583 let Some(ref sm) = self.state_machine else {
2584 return Ok(());
2585 };
2586 let claim_admission = self.disambiguation_admission.write().await;
2587 let reservation = self.reserve_state_transition().ok_or_else(|| {
2588 AgentError::Other("Another state transition is already in progress".to_string())
2589 })?;
2590 let from_state = sm.current();
2591 let expected_state_generation = sm.generation();
2592 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2593 let history_before = sm.history();
2594 drop(claim_admission);
2595
2596 self.execute_state_exit_actions(&from_state).await;
2597
2598 let admission = self.disambiguation_admission.write().await;
2599 if sm.current() != from_state
2600 || sm.generation() != expected_state_generation
2601 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2602 {
2603 return Err(AgentError::Other(
2604 "State ownership changed during manual transition preparation".to_string(),
2605 ));
2606 }
2607 sm.transition_to(state, "manual transition")?;
2608 self.invalidate_pending_confirmation("state_transition")
2609 .await;
2610 let entered = sm.current();
2611 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2612 drop(admission);
2613
2614 self.execute_state_enter_actions(&entered, is_reentry).await;
2615 drop(reservation);
2616 info!(to = %entered, "Manual state transition");
2617 Ok(())
2618 }
2619
2620 pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2621 self.state_machine
2622 .as_ref()
2623 .map(|sm| sm.history())
2624 .unwrap_or_default()
2625 }
2626
2627 pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2629 self.session_metadata.read().clone()
2630 }
2631
2632 pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2635 let allowed = self
2636 .actor_memory_config
2637 .as_ref()
2638 .map(|c| c.privacy.allow_deletion)
2639 .unwrap_or(true);
2640 if !allowed {
2641 return Err(AgentError::Config(
2642 "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2643 ));
2644 }
2645 let storage = self.storage.read().clone();
2646 if let Some(storage) = storage {
2647 if !storage.supports(StorageCapability::ActorDataDeletion) {
2651 return Err(AgentError::UnsupportedStorageCapability(
2652 StorageCapability::ActorDataDeletion,
2653 ));
2654 }
2655 storage.delete_actor_data(&self.info.id, actor_id).await?;
2656 } else {
2657 let store = { self.fact_store.read().clone() };
2661 if let Some(store) = store {
2662 store.delete_actor_data(actor_id).await?;
2663 }
2664 }
2665 if let Some(manager) = self.relationship_manager.as_ref() {
2666 manager.remove(actor_id);
2667 }
2668 self.actor_facts_cache.write().remove(actor_id);
2669 Ok(())
2670 }
2671
2672 pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2674 *self.session_metadata.write() = meta;
2675 }
2676
2677 pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2679 let storage = self.storage.read().clone();
2680 match storage {
2681 Some(s) => {
2682 let count = s.cleanup_expired().await?;
2683 if count > 0 {
2684 self.hooks.on_sessions_expired(count).await;
2685 }
2686 Ok(count)
2687 }
2688 None => Err(AgentError::Config(
2689 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2690 )),
2691 }
2692 }
2693
2694 pub async fn list_sessions_filtered(
2696 &self,
2697 filter: &ai_agents_core::SessionFilter,
2698 ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2699 let storage = self.storage.read().clone();
2700 match storage {
2701 Some(s) => s.list_sessions_filtered(filter).await,
2702 None => Err(AgentError::Config(
2703 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2704 )),
2705 }
2706 }
2707
2708 pub async fn save_state(&self) -> Result<AgentSnapshot> {
2709 let memory_snapshot = self.memory.snapshot().await?;
2710 let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2711 let context_snapshot = self.context_manager.snapshot();
2712
2713 let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2714 .with_memory(memory_snapshot)
2715 .with_context(context_snapshot)
2716 .with_state_machine(
2717 state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2718 current_state: String::new(),
2719 previous_state: None,
2720 turn_count: 0,
2721 no_transition_count: 0,
2722 history: vec![],
2723 }),
2724 );
2725
2726 if let Some(ref persona) = self.persona_manager {
2727 snapshot.persona = Some(persona.snapshot_as_value()?);
2728 }
2729
2730 if let Some(ref relationships) = self.relationship_manager {
2731 snapshot.relationships = Some(relationships.snapshot_as_value()?);
2732 }
2733
2734 Ok(snapshot)
2735 }
2736
2737 pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2739 let mut snapshot = self.save_state().await?;
2740 if let Some(ref registry) = self.spawner_registry {
2741 let entries = registry.list_with_specs();
2742 if !entries.is_empty() {
2743 snapshot = snapshot.with_spawned_agents(entries);
2744 }
2745 }
2746 Ok(snapshot)
2747 }
2748
2749 pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2751 let _admission = self.disambiguation_admission.write().await;
2752 if self.state_transition_reserved.load(Ordering::SeqCst) {
2753 return Err(AgentError::Other(
2754 "Cannot restore state while a state transition is in progress".to_string(),
2755 ));
2756 }
2757 self.invalidate_pending_confirmation("state_restore").await;
2758 *self.pending_skill_id.write() = None;
2759 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2760 disambiguator.clear_pending().await;
2761 }
2762 self.memory.restore(snapshot.memory).await?;
2763
2764 if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2765 && !sm_snapshot.current_state.is_empty()
2766 {
2767 sm.restore(sm_snapshot)?;
2768 }
2769
2770 self.context_manager.restore(snapshot.context);
2771
2772 if let (Some(persona_value), Some(persona_manager)) =
2773 (snapshot.persona, &self.persona_manager)
2774 {
2775 persona_manager.restore_from_value(persona_value)?;
2776 }
2777
2778 if let (Some(relationship_value), Some(relationship_manager)) =
2779 (snapshot.relationships, &self.relationship_manager)
2780 {
2781 relationship_manager.restore_from_value(relationship_value)?;
2782 }
2783
2784 info!(agent_id = %snapshot.agent_id, "State restored");
2785 Ok(())
2786 }
2787
2788 pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2789 let snapshot = self.save_state().await?;
2790 storage.save(session_id, &snapshot).await
2791 }
2792
2793 async fn load_session_restore(
2794 storage: &dyn AgentStorage,
2795 session_id: &str,
2796 ) -> Result<Option<StoredSessionRestore>> {
2797 let Some(snapshot) = storage.load(session_id).await? else {
2798 return Ok(None);
2799 };
2800 let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2804 storage.load_metadata(session_id).await?
2805 } else {
2806 None
2807 };
2808 Ok(Some(StoredSessionRestore { snapshot, metadata }))
2809 }
2810
2811 async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2812 Ok(RuntimeSessionRestorePoint {
2813 snapshot: self.save_state().await?,
2814 metadata: self.session_metadata(),
2815 actor_id: self.actor_id(),
2816 session_id: self.current_session_id.read().clone(),
2817 })
2818 }
2819
2820 async fn apply_session_restore_unchecked(
2821 &self,
2822 session_id: &str,
2823 stored: StoredSessionRestore,
2824 ) -> Result<()> {
2825 self.restore_state(stored.snapshot).await?;
2826 let metadata = stored.metadata.unwrap_or_default();
2827 if let Some(actor_id) = metadata.actor_id.as_deref() {
2828 self.set_actor_id(actor_id)?;
2829 } else {
2830 self.clear_actor_id();
2831 }
2832 self.set_session_metadata(metadata);
2833 *self.current_session_id.write() = Some(session_id.to_string());
2834 Ok(())
2835 }
2836
2837 async fn restore_session_restore_point(
2838 &self,
2839 restore_point: &RuntimeSessionRestorePoint,
2840 ) -> Result<()> {
2841 self.restore_state(restore_point.snapshot.clone()).await?;
2842 if let Some(actor_id) = restore_point.actor_id.as_deref() {
2843 self.set_actor_id(actor_id)?;
2844 } else {
2845 self.clear_actor_id();
2846 }
2847 self.set_session_metadata(restore_point.metadata.clone());
2848 *self.current_session_id.write() = restore_point.session_id.clone();
2849 Ok(())
2850 }
2851
2852 async fn apply_session_restore(
2853 &self,
2854 session_id: &str,
2855 stored: StoredSessionRestore,
2856 ) -> Result<()> {
2857 let before = self.capture_session_restore_point().await?;
2858 if let Err(error) = self
2859 .apply_session_restore_unchecked(session_id, stored)
2860 .await
2861 {
2862 return match self.restore_session_restore_point(&before).await {
2863 Ok(()) => Err(error),
2864 Err(rollback_error) => Err(AgentError::Other(format!(
2865 "Session restore failed: {error}; rollback failed: {rollback_error}"
2866 ))),
2867 };
2868 }
2869 Ok(())
2870 }
2871
2872 async fn rollback_session_restore_set(
2873 parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
2874 children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
2875 ) -> Vec<String> {
2876 let mut errors = Vec::new();
2877 if let Some((agent, restore_point)) = parent
2878 && let Err(error) = agent.restore_session_restore_point(restore_point).await
2879 {
2880 errors.push(format!("parent: {error}"));
2881 }
2882 for (id, agent, restore_point) in children {
2883 if let Err(error) = agent.restore_session_restore_point(restore_point).await {
2884 errors.push(format!("child '{id}': {error}"));
2885 }
2886 }
2887 errors
2888 }
2889
2890 fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
2891 if rollback_errors.is_empty() {
2892 AgentError::Other(format!(
2893 "Session restore failed: {error}; runtime state was rolled back"
2894 ))
2895 } else {
2896 AgentError::Other(format!(
2897 "Session restore failed: {error}; rollback also failed for {}",
2898 rollback_errors.join(", ")
2899 ))
2900 }
2901 }
2902
2903 pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
2904 let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
2905 return Ok(false);
2906 };
2907 self.apply_session_restore(session_id, stored).await?;
2908 Ok(true)
2909 }
2910
2911 pub async fn save_session(&self, session_id: &str) -> Result<()> {
2912 let storage = self.storage.read().clone();
2913 match storage {
2914 Some(s) => {
2915 let is_new = {
2917 let cur = self.current_session_id.read().clone();
2918 cur.as_deref() != Some(session_id)
2919 };
2920 if is_new {
2921 *self.current_session_id.write() = Some(session_id.to_string());
2922 self.hooks.on_session_created(session_id).await;
2923 }
2924
2925 {
2927 let now = chrono::Utc::now();
2928 let msg_count = self
2929 .memory
2930 .get_messages(None)
2931 .await
2932 .map(|v| v.len())
2933 .unwrap_or(0);
2934 let mut meta = self.session_metadata.write();
2935 meta.last_active = now;
2936 meta.message_count = msg_count;
2937 if meta.actor_id.is_none() {
2938 meta.actor_id = self.actor_id.read().clone();
2939 }
2940 }
2941
2942 let snapshot = self.save_state().await?;
2943 if s.supports(StorageCapability::SessionMetadata) {
2947 let metadata = self.session_metadata.read().clone();
2948 s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
2949 .await
2950 } else {
2951 s.save(session_id, &snapshot).await
2952 }
2953 }
2954 None => Err(AgentError::Config(
2955 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2956 )),
2957 }
2958 }
2959
2960 pub async fn load_session(&self, session_id: &str) -> Result<bool> {
2961 let storage = self.storage.read().clone();
2962 match storage {
2963 Some(storage) => self.load_from(storage.as_ref(), session_id).await,
2964 None => Err(AgentError::Config(
2965 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2966 )),
2967 }
2968 }
2969
2970 pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
2972 self.init_storage().await?;
2973 let storage = self.storage.read().clone().ok_or_else(|| {
2974 AgentError::Config(
2975 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2976 )
2977 })?;
2978 let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
2979 .await?
2980 .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
2981 let manifest = target_parent
2982 .snapshot
2983 .spawned_agents
2984 .clone()
2985 .unwrap_or_default();
2986
2987 let registry = self.spawner_registry.as_ref().cloned();
2988 let spawner = if manifest.is_empty() {
2989 self.spawner.as_ref().cloned()
2990 } else {
2991 Some(self.spawner.as_ref().cloned().ok_or_else(|| {
2992 AgentError::Config(
2993 "Saved session contains child agents but this runtime has no spawner".into(),
2994 )
2995 })?)
2996 };
2997 let registry = if manifest.is_empty() {
2998 registry
2999 } else {
3000 Some(registry.ok_or_else(|| {
3001 AgentError::Config(
3002 "Saved session contains child agents but this runtime has no registry".into(),
3003 )
3004 })?)
3005 };
3006
3007 let mut target_ids = HashSet::with_capacity(manifest.len());
3008 let mut prepared = Vec::with_capacity(manifest.len());
3009 for entry in manifest {
3010 if !target_ids.insert(entry.id.clone()) {
3011 return Err(AgentError::InvalidSpec(format!(
3012 "Saved child manifest contains duplicate ID: {}",
3013 entry.id
3014 )));
3015 }
3016 let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
3017 spawner
3018 .as_ref()
3019 .expect("non-empty manifests require a spawner")
3020 .validate_explicit_child(&entry.id, &spec)?;
3021 prepared.push((entry.id, spec));
3022 }
3023
3024 let current_ids = registry
3025 .as_ref()
3026 .map(|registry| {
3027 registry
3028 .list()
3029 .into_iter()
3030 .map(|info| info.id)
3031 .collect::<HashSet<_>>()
3032 })
3033 .unwrap_or_default();
3034 let removal_count = current_ids.difference(&target_ids).count();
3035 let additions = prepared
3036 .iter()
3037 .filter(|(id, _)| !current_ids.contains(id))
3038 .cloned()
3039 .collect::<Vec<_>>();
3040
3041 let mut existing = Vec::new();
3042 if let Some(registry) = registry.as_ref() {
3043 for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3044 let agent = registry.get(id).ok_or_else(|| {
3045 AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3046 })?;
3047 let child_storage = agent.storage().ok_or_else(|| {
3048 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3049 })?;
3050 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3051 .await?
3052 .ok_or_else(|| {
3053 AgentError::Persistence(format!(
3054 "Child '{id}' has no saved session '{session_id}'"
3055 ))
3056 })?;
3057 existing.push((id.clone(), agent, stored));
3058 }
3059 }
3060
3061 let mut staged = Vec::with_capacity(additions.len());
3062 if !additions.is_empty() {
3063 let spawner = spawner
3064 .as_ref()
3065 .expect("restored additions require a spawner");
3066 let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3067 for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3068 let spawned = spawner
3069 .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3070 .await?;
3071 let child_storage = spawned.agent.storage().ok_or_else(|| {
3072 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3073 })?;
3074 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3075 .await?
3076 .ok_or_else(|| {
3077 AgentError::Persistence(format!(
3078 "Child '{id}' has no saved session '{session_id}'"
3079 ))
3080 })?;
3081 staged.push((spawned, stored));
3082 }
3083 } else if let Some(spawner) = spawner.as_ref() {
3084 spawner.reserve_restore_capacity(0, removal_count)?;
3085 }
3086
3087 let parent_before = self.capture_session_restore_point().await?;
3088 let mut existing_before = Vec::with_capacity(existing.len());
3089 for (id, agent, _) in &existing {
3090 existing_before.push((
3091 id.clone(),
3092 Arc::clone(agent),
3093 agent.capture_session_restore_point().await?,
3094 ));
3095 }
3096
3097 for (_, agent, stored) in &existing {
3101 if let Err(error) = agent
3102 .apply_session_restore_unchecked(session_id, stored.clone())
3103 .await
3104 {
3105 drop(staged);
3106 let rollback_errors =
3107 Self::rollback_session_restore_set(None, &existing_before).await;
3108 return Err(Self::restore_failure(error, rollback_errors));
3109 }
3110 }
3111 for (spawned, stored) in &staged {
3112 if let Err(error) = spawned
3113 .agent
3114 .apply_session_restore_unchecked(session_id, stored.clone())
3115 .await
3116 {
3117 drop(staged);
3118 let rollback_errors =
3119 Self::rollback_session_restore_set(None, &existing_before).await;
3120 return Err(Self::restore_failure(error, rollback_errors));
3121 }
3122 }
3123 if let Err(error) = self
3124 .apply_session_restore_unchecked(session_id, target_parent)
3125 .await
3126 {
3127 drop(staged);
3128 let rollback_errors =
3129 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3130 .await;
3131 return Err(Self::restore_failure(error, rollback_errors));
3132 }
3133
3134 if let Some(registry) = registry.as_ref()
3135 && let Err(error) = registry
3136 .reconcile(
3137 &target_ids,
3138 staged.into_iter().map(|(spawned, _)| spawned).collect(),
3139 )
3140 .await
3141 {
3142 let rollback_errors =
3143 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3144 .await;
3145 return Err(Self::restore_failure(error, rollback_errors));
3146 }
3147
3148 Ok(target_ids.len())
3149 }
3150
3151 pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3152 let storage = self.storage.read().clone();
3153 match storage {
3154 Some(s) => s.delete(session_id).await,
3155 None => Err(AgentError::Config(
3156 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3157 )),
3158 }
3159 }
3160
3161 pub async fn list_sessions(&self) -> Result<Vec<String>> {
3162 let storage = self.storage.read().clone();
3163 match storage {
3164 Some(s) => s.list_sessions().await,
3165 None => Err(AgentError::Config(
3166 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3167 )),
3168 }
3169 }
3170
3171 fn estimate_tokens(&self, text: &str) -> u32 {
3172 (text.len() as f32 / 4.0).ceil() as u32
3173 }
3174
3175 fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3176 messages
3177 .iter()
3178 .map(|m| self.estimate_tokens(&m.content))
3179 .sum()
3180 }
3181
3182 fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3183 if messages.len() <= keep_recent + 1 {
3184 return;
3185 }
3186 let system_msg = messages.remove(0);
3187 let to_remove = messages.len().saturating_sub(keep_recent);
3188 messages.drain(..to_remove);
3189 messages.insert(0, system_msg);
3190 }
3191
3192 fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3193 match config {
3194 FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3195 FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3196 FilterConfig::SkipPattern { skip_if_contains } => {
3197 Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3198 }
3199 FilterConfig::Custom { name } => {
3200 let filters = self.message_filters.read();
3201 filters
3202 .get(name)
3203 .cloned()
3204 .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3205 }
3206 }
3207 }
3208
3209 async fn summarize_context(
3210 &self,
3211 messages: &mut Vec<ChatMessage>,
3212 summarizer_llm: Option<&str>,
3213 max_summary_tokens: u32,
3214 custom_prompt: Option<&str>,
3215 keep_recent: usize,
3216 filter: Option<&FilterConfig>,
3217 ) -> Result<()> {
3218 let system_msg = messages.remove(0);
3219
3220 let to_summarize_count = messages.len().saturating_sub(keep_recent);
3221 if to_summarize_count == 0 {
3222 messages.insert(0, system_msg);
3223 return Ok(());
3224 }
3225
3226 let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3227 let mut to_summarize = std::mem::take(messages);
3228
3229 if let Some(filter_config) = filter {
3230 let filter = self.get_filter(filter_config);
3231 to_summarize = filter.filter(to_summarize);
3232 }
3233
3234 if to_summarize.is_empty() {
3235 *messages = recent_msgs;
3236 messages.insert(0, system_msg);
3237 return Ok(());
3238 }
3239
3240 let conversation_text = to_summarize
3241 .iter()
3242 .map(|m| format!("{:?}: {}", m.role, m.content))
3243 .collect::<Vec<_>>()
3244 .join("\n");
3245
3246 let default_prompt = format!(
3247 "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3248 max_summary_tokens, conversation_text
3249 );
3250
3251 let summary_prompt = custom_prompt
3252 .map(|p| format!("{}\n\n{}", p, conversation_text))
3253 .unwrap_or(default_prompt);
3254
3255 let summarizer = if let Some(alias) = summarizer_llm {
3256 self.llm_registry
3257 .get(alias)
3258 .map_err(|e| AgentError::Config(e.to_string()))?
3259 } else {
3260 self.llm_registry
3261 .router()
3262 .or_else(|_| self.llm_registry.default())
3263 .map_err(|e| AgentError::Config(e.to_string()))?
3264 };
3265
3266 let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3267 let response = self
3268 .observe_purpose(
3269 ObservationPurpose::Summarization,
3270 summarizer.complete(&summary_msgs, None),
3271 )
3272 .await?;
3273
3274 let summary_message = ChatMessage::system(format!(
3275 "[Previous conversation summary]\n{}",
3276 response.content
3277 ));
3278
3279 *messages = vec![system_msg, summary_message];
3280 messages.extend(recent_msgs);
3281
3282 debug!(
3283 summarized_count = to_summarize_count,
3284 kept_recent = keep_recent,
3285 "Context summarized"
3286 );
3287
3288 Ok(())
3289 }
3290
3291 fn render_system_prompt(&self) -> Result<String> {
3292 let mut context = self.build_context_with_overlays();
3293
3294 let facts_text = self.format_actor_facts_for_context();
3296 if !facts_text.is_empty() {
3297 context.insert(
3298 "actor_facts".to_string(),
3299 serde_json::Value::String(facts_text),
3300 );
3301 }
3302
3303 if let Some((key, text)) = self.format_relationship_for_context() {
3304 context.insert(key, serde_json::Value::String(text));
3305 }
3306
3307 self.template_renderer
3308 .render(&self.base_system_prompt, &context)
3309 }
3310
3311 fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3313 let mut seen = HashSet::new();
3314 ids.iter()
3315 .filter_map(|id| self.tools.canonical_id(id))
3316 .filter(|canonical_id| seen.insert(canonical_id.clone()))
3317 .collect()
3318 }
3319
3320 fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3322 let Some(declared) = self.declared_tool_ids.as_deref() else {
3323 return Vec::new();
3324 };
3325 let mut effective = self.canonical_unique_tool_ids(declared);
3326 if let Some(scope) = scope_override {
3327 let scope: HashSet<String> =
3328 self.canonical_unique_tool_ids(scope).into_iter().collect();
3329 effective.retain(|canonical_id| scope.contains(canonical_id));
3330 }
3331 effective
3332 }
3333
3334 async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3336 Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3337 }
3338
3339 async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3341 let scope_override = self.runtime_control.tool_scope_override.read().clone();
3342 self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3343 .await
3344 }
3345
3346 async fn get_available_tool_ids_snapshot_for_scope(
3348 &self,
3349 scope_override: Option<&[String]>,
3350 ) -> Result<AvailableToolIdsSnapshot> {
3351 let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3352 let (state_generation, state_scopes) = self
3353 .state_machine
3354 .as_ref()
3355 .map(|state_machine| {
3356 let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3357 (Some(generation), scopes)
3358 })
3359 .unwrap_or((None, Vec::new()));
3360
3361 if available.is_empty() || state_scopes.is_empty() {
3362 return Ok(AvailableToolIdsSnapshot {
3363 tool_ids: available,
3364 state_generation,
3365 });
3366 }
3367
3368 let eval_ctx = self.build_evaluation_context().await?;
3369 let llm_getter = RegistryLLMGetter {
3370 registry: self.llm_registry.clone(),
3371 };
3372 let evaluator = ConditionEvaluator::new(llm_getter);
3373
3374 for state_scope in state_scopes {
3375 if state_scope.is_empty() {
3376 available.clear();
3377 break;
3378 }
3379
3380 let mut allowed = HashSet::new();
3381 for tool_ref in &state_scope {
3382 let tool_id = tool_ref.id();
3383 let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3384 continue;
3385 };
3386 let condition_matches = if let Some(condition) = tool_ref.condition() {
3387 match evaluator.evaluate(condition, &eval_ctx).await {
3388 Ok(matches) => matches,
3389 Err(error) => {
3390 warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3391 false
3392 }
3393 }
3394 } else {
3395 true
3396 };
3397 if condition_matches {
3398 allowed.insert(canonical_id);
3399 } else {
3400 debug!(tool = tool_id, "Tool condition not met, skipping");
3401 }
3402 }
3403 available.retain(|canonical_id| allowed.contains(canonical_id));
3404 if available.is_empty() {
3405 break;
3406 }
3407 }
3408
3409 Ok(AvailableToolIdsSnapshot {
3410 tool_ids: available,
3411 state_generation,
3412 })
3413 }
3414
3415 async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3416 let context = self.build_context_with_overlays();
3417 let messages = self.memory.get_messages(Some(10)).await?;
3418 let tool_history = self.tool_call_history.read().clone();
3419
3420 let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3421 (Some(sm.current()), sm.turn_count(), sm.previous())
3422 } else {
3423 (None, 0, None)
3424 };
3425
3426 Ok(EvaluationContext::default()
3427 .with_context(context)
3428 .with_state(state_name, turn_count, previous_state)
3429 .with_called_tools(tool_history)
3430 .with_messages(messages))
3431 }
3432
3433 fn record_tool_call(&self, tool_id: &str, result: Value) {
3434 self.tool_call_history.write().push(ToolCallRecord {
3435 tool_id: tool_id.to_string(),
3436 result,
3437 timestamp: chrono::Utc::now(),
3438 });
3439 }
3440
3441 async fn get_effective_system_prompt_with_persona_hooks(
3442 &self,
3443 fire_persona_hooks: bool,
3444 include_tool_prompt: bool,
3445 ) -> Result<String> {
3446 let rendered_base = self.render_system_prompt()?;
3447
3448 let persona_prefix = if let Some(ref persona) = self.persona_manager {
3449 let context = self.build_context_with_overlays();
3450 if fire_persona_hooks {
3451 let render_result = persona.render_prompt(&context)?;
3452 for content in &render_result.newly_revealed {
3453 self.hooks.on_secret_revealed(content).await;
3454 }
3455 render_result.prompt
3456 } else {
3457 persona.render_prompt_preview(&context)?
3458 }
3459 } else {
3460 String::new()
3461 };
3462
3463 if let Some(ref sm) = self.state_machine
3464 && let Some(state_def) = sm.current_definition()
3465 {
3466 let state_prompt = if let Some(ref prompt) = state_def.prompt {
3467 let context = self.build_context_with_overlays();
3468 self.template_renderer.render_with_state(
3469 prompt,
3470 &context,
3471 &sm.current(),
3472 sm.previous().as_deref(),
3473 sm.turn_count(),
3474 state_def.max_turns,
3475 )?
3476 } else {
3477 String::new()
3478 };
3479
3480 let combined = match state_def.prompt_mode {
3481 PromptMode::Append => {
3482 if state_prompt.is_empty() {
3483 rendered_base
3484 } else {
3485 format!(
3486 "{}\n\n[Current State: {}]\n{}",
3487 rendered_base,
3488 sm.current(),
3489 state_prompt
3490 )
3491 }
3492 }
3493 PromptMode::Replace => {
3494 if state_prompt.is_empty() {
3495 rendered_base
3496 } else {
3497 state_prompt
3498 }
3499 }
3500 PromptMode::Prepend => {
3501 if state_prompt.is_empty() {
3502 rendered_base
3503 } else {
3504 format!("{}\n\n{}", state_prompt, rendered_base)
3505 }
3506 }
3507 };
3508
3509 let with_persona = if persona_prefix.is_empty() {
3511 combined
3512 } else {
3513 format!("{}\n\n{}", persona_prefix, combined)
3514 };
3515
3516 if include_tool_prompt {
3517 let available_tool_ids = self.get_available_tool_ids().await?;
3518 if !available_tool_ids.is_empty() {
3519 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3520 &available_tool_ids,
3521 None,
3522 self.parallel_tools.enabled,
3523 self.runtime_config.tool_schema_prompt_mode,
3524 );
3525 if !tools_prompt.is_empty() {
3526 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3527 }
3528 }
3529 }
3530 return Ok(with_persona);
3531 }
3532
3533 let with_persona = if persona_prefix.is_empty() {
3535 rendered_base
3536 } else {
3537 format!("{}\n\n{}", persona_prefix, rendered_base)
3538 };
3539
3540 if include_tool_prompt {
3541 let available_tool_ids = self.get_available_tool_ids().await?;
3542 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3543 &available_tool_ids,
3544 None,
3545 self.parallel_tools.enabled,
3546 self.runtime_config.tool_schema_prompt_mode,
3547 );
3548 if !tools_prompt.is_empty() {
3549 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3550 }
3551 }
3552 Ok(with_persona)
3553 }
3554
3555 fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3556 if let Some(ref sm) = self.state_machine
3557 && let Some(state_def) = sm.current_definition()
3558 && let Some(ref llm_alias) = state_def.llm
3559 {
3560 return self
3561 .llm_registry
3562 .get(llm_alias)
3563 .map_err(|e| AgentError::Config(e.to_string()));
3564 }
3565 self.llm_registry
3566 .default()
3567 .map_err(|e| AgentError::Config(e.to_string()))
3568 }
3569
3570 fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3571 if let Some(ref sm) = self.state_machine
3572 && let Some(state_def) = sm.current_definition()
3573 && let Some(ref state_reasoning) = state_def.reasoning
3574 {
3575 return state_reasoning.clone();
3576 }
3577 self.reasoning_config.clone()
3578 }
3579
3580 fn get_effective_reflection_config(&self) -> ReflectionConfig {
3581 if let Some(ref sm) = self.state_machine
3582 && let Some(state_def) = sm.current_definition()
3583 && let Some(ref state_reflection) = state_def.reflection
3584 {
3585 return state_reflection.clone();
3586 }
3587 self.reflection_config.clone()
3588 }
3589
3590 fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3591 skill
3592 .reasoning
3593 .clone()
3594 .unwrap_or_else(|| self.get_effective_reasoning_config())
3595 }
3596
3597 fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3598 skill
3599 .reflection
3600 .clone()
3601 .unwrap_or_else(|| self.get_effective_reflection_config())
3602 }
3603
3604 async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3605 let recent_messages: Vec<String> = self
3606 .memory
3607 .get_messages(Some(5))
3608 .await?
3609 .iter()
3610 .rev()
3611 .map(|m| format!("{:?}: {}", m.role, m.content))
3612 .collect();
3613
3614 let current_state = self.current_state().map(|s| s.to_string());
3615
3616 let state_prompt: Option<String> = self
3619 .state_machine
3620 .as_ref()
3621 .and_then(|sm| sm.current_definition())
3622 .and_then(|def| def.prompt.clone());
3623
3624 let available_tools: Vec<String> = self
3625 .get_available_tool_ids()
3626 .await
3627 .unwrap_or_else(|_| self.tools.list_ids());
3628
3629 let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3630
3631 let mut user_context = self.build_context_with_overlays();
3632 user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3633 if let Some(state_generation) = self
3634 .state_machine
3635 .as_ref()
3636 .map(|state_machine| state_machine.generation())
3637 {
3638 user_context.insert(
3639 DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3640 serde_json::json!(state_generation),
3641 );
3642 }
3643
3644 let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3646 sm.current_definition()
3647 .map(|def| {
3648 def.transitions
3649 .iter()
3650 .filter_map(|t| t.intent.clone())
3651 .collect()
3652 })
3653 .unwrap_or_default()
3654 } else {
3655 Vec::new()
3656 };
3657
3658 Ok(DisambiguationContext::from_agent_state(
3659 recent_messages,
3660 current_state,
3661 state_prompt,
3662 available_tools,
3663 available_skills,
3664 available_intents,
3665 user_context,
3666 ))
3667 }
3668
3669 fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3670 if let Some(ref sm) = self.state_machine
3671 && let Some(state_def) = sm.current_definition()
3672 {
3673 let parent_def = sm.get_parent_definition();
3674 let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3675 if !effective_skills.is_empty() {
3676 return self
3677 .skills
3678 .iter()
3679 .filter(|s| effective_skills.contains(&&s.id))
3680 .collect();
3681 }
3682 }
3683 self.skills.iter().collect()
3684 }
3685
3686 async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3687 self.build_messages_internal(true, None, true).await
3688 }
3689
3690 async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3691 self.build_messages_internal(false, Some(user_message), true)
3692 .await
3693 }
3694
3695 async fn build_messages_internal(
3696 &self,
3697 fire_persona_hooks: bool,
3698 ephemeral_user_message: Option<&str>,
3699 include_tool_prompt: bool,
3700 ) -> Result<Vec<ChatMessage>> {
3701 let system_prompt = self
3702 .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3703 .await?;
3704 let mut messages = vec![ChatMessage::system(&system_prompt)];
3705
3706 let context = self.memory.get_context().await?;
3707 let history = if let Some(ref budget) = self.memory_token_budget {
3708 context.to_llm_messages_with_allocation(&budget.allocation)
3709 } else {
3710 context.to_llm_messages()
3711 };
3712 messages.extend(history);
3713 if let Some(user_message) = ephemeral_user_message {
3714 messages.push(ChatMessage::user(user_message));
3715 }
3716
3717 let total_tokens = self.estimate_total_tokens(&messages);
3718
3719 if total_tokens > self.max_context_tokens {
3720 debug!(
3721 total = total_tokens,
3722 limit = self.max_context_tokens,
3723 "Context overflow"
3724 );
3725
3726 match &self.recovery_manager.config().llm.on_context_overflow {
3727 ContextOverflowAction::Error => {
3728 return Err(AgentError::LLM(format!(
3729 "Context overflow: {} tokens > {} limit",
3730 total_tokens, self.max_context_tokens
3731 )));
3732 }
3733 ContextOverflowAction::Truncate { keep_recent } => {
3734 self.truncate_context(&mut messages, *keep_recent);
3735 }
3736 ContextOverflowAction::Summarize {
3737 summarizer_llm,
3738 max_summary_tokens,
3739 custom_prompt,
3740 keep_recent,
3741 filter,
3742 } => {
3743 self.summarize_context(
3744 &mut messages,
3745 summarizer_llm.as_deref(),
3746 *max_summary_tokens,
3747 custom_prompt.as_deref(),
3748 *keep_recent,
3749 filter.as_ref(),
3750 )
3751 .await?;
3752 }
3753 }
3754 }
3755
3756 Ok(messages)
3757 }
3758
3759 async fn main_tool_protocol(
3760 &self,
3761 llm: &dyn LLMProvider,
3762 ephemeral_new_turn: bool,
3763 ) -> Result<MainToolProtocol> {
3764 let mut choice = llm.configured_tool_choice();
3765 if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3766 return Ok(MainToolProtocol {
3767 choice,
3768 tool_ids: Vec::new(),
3769 definitions: Vec::new(),
3770 });
3771 }
3772
3773 let mut tool_ids = self.get_available_tool_ids().await?;
3774 tool_ids.sort();
3775 tool_ids.dedup();
3776 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3777 let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3778 AgentError::Config(format!(
3779 "specific tool choice '{expected}' is not registered"
3780 ))
3781 })?;
3782 if canonical != *expected {
3783 return Err(AgentError::Config(format!(
3784 "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3785 )));
3786 }
3787 if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3788 return Err(AgentError::Config(format!(
3789 "specific tool choice '{expected}' is outside the effective tool grant"
3790 )));
3791 }
3792 }
3793 if matches!(
3794 choice.as_ref(),
3795 Some(ToolChoice::Required | ToolChoice::Specific(_))
3796 ) && tool_ids.is_empty()
3797 {
3798 return Err(AgentError::Config(
3799 "required tool choice has no tool inside the effective grant".to_string(),
3800 ));
3801 }
3802 if !ephemeral_new_turn
3803 && let Some(configured_choice) = choice.as_ref()
3804 && matches!(
3805 configured_choice,
3806 ToolChoice::Required | ToolChoice::Specific(_)
3807 )
3808 && self
3809 .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3810 .await?
3811 {
3812 choice = Some(ToolChoice::Auto);
3813 }
3814 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3815 tool_ids.retain(|tool_id| tool_id == expected);
3816 }
3817
3818 let definitions = tool_ids
3819 .iter()
3820 .map(|tool_id| {
3821 let tool = self.tools.get(tool_id).ok_or_else(|| {
3822 AgentError::Config(format!(
3823 "effective tool '{tool_id}' disappeared before provider exposure"
3824 ))
3825 })?;
3826 Ok(LLMToolDefinition {
3827 name: tool_id.clone(),
3828 description: tool.description().to_string(),
3829 input_schema: tool.input_schema(),
3830 })
3831 })
3832 .collect::<Result<Vec<_>>>()?;
3833
3834 Ok(MainToolProtocol {
3838 choice,
3839 tool_ids,
3840 definitions,
3841 })
3842 }
3843
3844 async fn tool_choice_satisfied_in_current_turn(
3845 &self,
3846 choice: &ToolChoice,
3847 effective_tool_ids: &[String],
3848 ) -> Result<bool> {
3849 let messages = self.memory.get_messages(None).await?;
3850 let mut saw_tool_result = false;
3851 for message in messages.iter().rev() {
3852 match message.role {
3853 ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
3854 saw_tool_result = true;
3855 }
3856 ai_agents_core::Role::Assistant if saw_tool_result => {
3857 let Some(calls) = self.parse_tool_calls(&message.content) else {
3858 continue;
3859 };
3860 let calls_are_effective = !calls.is_empty()
3861 && calls.iter().all(|call| {
3862 self.tools
3863 .canonical_id(&call.name)
3864 .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
3865 });
3866 return Ok(calls_are_effective
3867 && match choice {
3868 ToolChoice::Required => true,
3869 ToolChoice::Specific(expected) => calls.iter().all(|call| {
3870 self.tools.canonical_id(&call.name).as_deref()
3871 == Some(expected.as_str())
3872 }),
3873 _ => false,
3874 });
3875 }
3876 ai_agents_core::Role::User => return Ok(false),
3877 _ => {}
3878 }
3879 }
3880 Ok(false)
3881 }
3882
3883 fn provider_can_use_native_tools(
3884 &self,
3885 llm: &dyn LLMProvider,
3886 protocol: &MainToolProtocol,
3887 ) -> bool {
3888 let Some(choice) = protocol.choice.as_ref() else {
3889 return false;
3890 };
3891 if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
3892 return false;
3893 }
3894 llm.supports_tool_choice(choice)
3895 && protocol.definitions.iter().all(|definition| {
3896 !definition.name.is_empty()
3897 && definition.name.len() <= 64
3898 && definition
3899 .name
3900 .bytes()
3901 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
3902 })
3903 }
3904
3905 fn prompt_messages_for_tool_protocol(
3906 &self,
3907 messages: &[ChatMessage],
3908 protocol: &MainToolProtocol,
3909 corrective: bool,
3910 ) -> Vec<ChatMessage> {
3911 let mut messages = messages.to_vec();
3912 let Some(choice) = protocol.choice.as_ref() else {
3913 return messages;
3914 };
3915 if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
3916 return messages;
3917 }
3918
3919 let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
3920 &protocol.tool_ids,
3921 None,
3922 self.parallel_tools.enabled,
3923 self.runtime_config.tool_schema_prompt_mode,
3924 );
3925 match choice {
3926 ToolChoice::Required => tool_prompt.push_str(
3927 "\n\nYou must call at least one listed tool before giving a final answer.",
3928 ),
3929 ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
3930 "\n\nYou must call the '{tool_id}' tool before giving a final answer."
3931 )),
3932 ToolChoice::Auto => {}
3933 ToolChoice::None => return messages,
3934 _ => return messages,
3935 }
3936 if let Some(system) = messages
3937 .iter_mut()
3938 .find(|message| message.role == ai_agents_core::Role::System)
3939 {
3940 system.content.push_str("\n\n");
3941 system.content.push_str(&tool_prompt);
3942 } else {
3943 messages.insert(0, ChatMessage::system(tool_prompt));
3944 }
3945 if corrective {
3946 let instruction = match choice {
3947 ToolChoice::Required => {
3948 "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
3949 }
3950 ToolChoice::Specific(tool_id) => {
3951 messages.push(ChatMessage::user(format!(
3952 "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
3953 )));
3954 return messages;
3955 }
3956 _ => return messages,
3957 };
3958 messages.push(ChatMessage::user(instruction));
3959 }
3960 messages
3961 }
3962
3963 async fn invoke_main_provider(
3964 &self,
3965 llm: Arc<dyn LLMProvider>,
3966 messages: &[ChatMessage],
3967 protocol: &MainToolProtocol,
3968 corrective: bool,
3969 ) -> std::result::Result<MainProviderResponse, LLMError> {
3970 let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
3971 let response = if use_native {
3972 let request = LLMToolRequest {
3973 tools: protocol.definitions.clone(),
3974 choice: protocol
3975 .choice
3976 .clone()
3977 .expect("native tool requests require an explicit choice"),
3978 };
3979 self.observe_purpose(
3980 ObservationPurpose::MainResponse,
3981 llm.complete_with_tools(messages, None, &request),
3982 )
3983 .await?
3984 } else {
3985 let prompt_messages =
3986 self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
3987 self.observe_purpose(
3988 ObservationPurpose::MainResponse,
3989 llm.complete(&prompt_messages, None),
3990 )
3991 .await?
3992 };
3993 Ok(MainProviderResponse {
3994 response,
3995 used_native_tools: use_native,
3996 })
3997 }
3998
3999 async fn complete_main_attempt_with_recovery(
4000 &self,
4001 llm: Arc<dyn LLMProvider>,
4002 messages: &[ChatMessage],
4003 protocol: &MainToolProtocol,
4004 corrective: bool,
4005 ) -> Result<MainProviderResponse> {
4006 let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
4007 self.recovery_manager
4008 .with_retry("llm_call", None, || {
4009 let llm = Arc::clone(&llm);
4010 async move {
4011 self.invoke_main_provider(llm, messages, protocol, corrective)
4012 .await
4013 .map_err(|error| error.classify())
4014 }
4015 })
4016 .await
4017 .map_err(|error| AgentError::LLM(error.to_string()))
4018 } else {
4019 self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4020 .await
4021 .map_err(|error| AgentError::LLM(error.to_string()))
4022 };
4023
4024 match primary_result {
4025 Ok(response) => Ok(response),
4026 Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4027 LLMFailureAction::FallbackLlm { fallback_llm } => {
4028 let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4029 AgentError::Config(format!(
4030 "Fallback LLM '{fallback_llm}' not found: {error}"
4031 ))
4032 })?;
4033 self.invoke_main_provider(fallback, messages, protocol, corrective)
4034 .await
4035 .map_err(|error| AgentError::LLM(error.to_string()))
4036 }
4037 LLMFailureAction::FallbackResponse { message } => {
4038 if matches!(
4039 protocol.choice.as_ref(),
4040 Some(ToolChoice::Required | ToolChoice::Specific(_))
4041 ) {
4042 Err(AgentError::LLM(format!(
4043 "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4044 )))
4045 } else {
4046 Ok(MainProviderResponse {
4047 response: LLMResponse::new(message.clone(), FinishReason::Stop),
4048 used_native_tools: false,
4049 })
4050 }
4051 }
4052 LLMFailureAction::Error => Err(primary_error),
4053 },
4054 }
4055 }
4056
4057 fn normalize_main_provider_response(
4058 &self,
4059 mut response: LLMResponse,
4060 protocol: &MainToolProtocol,
4061 ) -> Result<(LLMResponse, bool)> {
4062 let native_calls = response
4063 .tool_calls()
4064 .map_err(|error| AgentError::LLM(error.to_string()))?;
4065 let calls = match native_calls {
4066 Some(calls) => {
4067 let markers = calls
4068 .iter()
4069 .map(|call| {
4070 serde_json::json!({
4071 "_ai_agents_native_tool_call": true,
4072 "id": call.id,
4073 "tool": call.name,
4074 "arguments": call.arguments,
4075 })
4076 })
4077 .collect::<Vec<_>>();
4078 response.content = if markers.len() == 1 {
4079 markers[0].to_string()
4080 } else {
4081 serde_json::Value::Array(markers).to_string()
4082 };
4083 Some(calls)
4084 }
4085 None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4086 self.parse_tool_calls(response.content.trim())
4087 }
4088 None => None,
4089 };
4090
4091 if protocol.choice.is_some()
4092 && let Some(calls) = calls.as_ref()
4093 && calls.iter().any(|call| {
4094 self.tools
4095 .canonical_id(&call.name)
4096 .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4097 })
4098 {
4099 return Err(AgentError::LLM(
4100 "Provider returned a tool call outside the effective grant".to_string(),
4101 ));
4102 }
4103
4104 let compliant = match protocol.choice.as_ref() {
4105 Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4106 Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4107 !calls.is_empty()
4108 && calls.iter().all(|call| {
4109 self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4110 })
4111 }),
4112 _ => true,
4113 };
4114 Ok((response, compliant))
4115 }
4116
4117 async fn complete_main_llm_with_recovery(
4118 &self,
4119 llm: Arc<dyn LLMProvider>,
4120 messages: &[ChatMessage],
4121 protocol: &MainToolProtocol,
4122 ) -> Result<LLMResponse> {
4123 let first = self
4124 .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4125 .await?;
4126 let (response, compliant) =
4127 self.normalize_main_provider_response(first.response, protocol)?;
4128 if compliant {
4129 return Ok(response);
4130 }
4131 if first.used_native_tools {
4132 return Err(AgentError::LLM(
4133 "Provider returned no compliant native call for required tool choice".to_string(),
4134 ));
4135 }
4136
4137 let corrected = self
4138 .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4139 .await?;
4140 let (response, compliant) =
4141 self.normalize_main_provider_response(corrected.response, protocol)?;
4142 if compliant {
4143 return Ok(response);
4144 }
4145 Err(AgentError::LLM(
4146 "Provider returned no compliant tool call after one corrective retry".to_string(),
4147 ))
4148 }
4149
4150 fn is_native_tool_call_content(content: &str) -> bool {
4151 let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4152 return false;
4153 };
4154 match value {
4155 serde_json::Value::Array(values) => {
4156 !values.is_empty()
4157 && values.iter().all(|value| {
4158 value
4159 .get("_ai_agents_native_tool_call")
4160 .and_then(|marker| marker.as_bool())
4161 == Some(true)
4162 })
4163 }
4164 serde_json::Value::Object(map) => {
4165 map.get("_ai_agents_native_tool_call")
4166 .and_then(|marker| marker.as_bool())
4167 == Some(true)
4168 }
4169 _ => false,
4170 }
4171 }
4172
4173 fn tool_result_message(
4174 tool_call: &ToolCall,
4175 output: &str,
4176 native_tool_call: bool,
4177 ) -> ChatMessage {
4178 if !native_tool_call {
4179 return ChatMessage::function(&tool_call.name, output);
4180 }
4181 let output = serde_json::from_str::<serde_json::Value>(output)
4182 .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4183 ChatMessage::function(
4184 &tool_call.name,
4185 serde_json::json!({
4186 "_ai_agents_native_tool_result": true,
4187 "id": tool_call.id,
4188 "tool": tool_call.name,
4189 "output": output,
4190 })
4191 .to_string(),
4192 )
4193 }
4194
4195 fn parse_main_tool_calls(
4196 &self,
4197 content: &str,
4198 protocol: &MainToolProtocol,
4199 ) -> Option<Vec<ToolCall>> {
4200 if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4201 None
4202 } else {
4203 self.parse_tool_calls(content)
4204 }
4205 }
4206
4207 fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4208 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4210 if let Some(arr) = parsed.as_array() {
4212 let calls: Vec<ToolCall> = arr
4213 .iter()
4214 .filter_map(|v| self.extract_tool_call_from_value(v))
4215 .collect();
4216 if !calls.is_empty() {
4217 return Some(calls);
4218 }
4219 }
4220 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4222 return Some(vec![tool_call]);
4223 }
4224 }
4225
4226 if let Some(json_str) = self.extract_json_from_content(content)
4228 && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4229 {
4230 if let Some(arr) = parsed.as_array() {
4232 let calls: Vec<ToolCall> = arr
4233 .iter()
4234 .filter_map(|v| self.extract_tool_call_from_value(v))
4235 .collect();
4236 if !calls.is_empty() {
4237 return Some(calls);
4238 }
4239 }
4240 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4242 return Some(vec![tool_call]);
4243 }
4244 }
4245
4246 None
4247 }
4248
4249 fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4250 if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4251 let arguments = parsed
4252 .get("arguments")
4253 .cloned()
4254 .unwrap_or(serde_json::json!({}));
4255 return Some(ToolCall {
4256 id: parsed
4257 .get("id")
4258 .and_then(|value| value.as_str())
4259 .filter(|id| !id.is_empty())
4260 .map(str::to_string)
4261 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4262 name: tool_name.to_string(),
4263 arguments,
4264 });
4265 }
4266 None
4267 }
4268
4269 fn extract_json_from_content(&self, content: &str) -> Option<String> {
4271 if let Some(result) = self.extract_json_array_from_content(content) {
4273 return Some(result);
4274 }
4275 self.extract_json_object_from_content(content)
4276 }
4277
4278 fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4280 let start = content.find('[')?;
4281 let content_from_start = &content[start..];
4282
4283 let mut depth = 0;
4284 let mut end = 0;
4285 for (i, ch) in content_from_start.char_indices() {
4286 match ch {
4287 '[' => depth += 1,
4288 ']' => {
4289 depth -= 1;
4290 if depth == 0 {
4291 end = i + 1;
4292 break;
4293 }
4294 }
4295 _ => {}
4296 }
4297 }
4298
4299 if end > 0 {
4300 let json_str = &content_from_start[..end];
4301 if json_str.contains("\"tool\"") {
4303 return Some(json_str.to_string());
4304 }
4305 }
4306
4307 None
4308 }
4309
4310 fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4312 let start = content.find('{')?;
4313 let content_from_start = &content[start..];
4314
4315 let mut depth = 0;
4317 let mut end = 0;
4318 for (i, ch) in content_from_start.char_indices() {
4319 match ch {
4320 '{' => depth += 1,
4321 '}' => {
4322 depth -= 1;
4323 if depth == 0 {
4324 end = i + 1;
4325 break;
4326 }
4327 }
4328 _ => {}
4329 }
4330 }
4331
4332 if end > 0 {
4333 let json_str = &content_from_start[..end];
4334 if json_str.contains("\"tool\"") {
4336 return Some(json_str.to_string());
4337 }
4338 }
4339
4340 None
4341 }
4342
4343 #[allow(clippy::too_many_arguments)]
4347 fn record_from_parts(
4348 &self,
4349 request: &ToolExecutionRequest,
4350 canonical_id: String,
4351 executed_arguments: Value,
4352 started_at: chrono::DateTime<chrono::Utc>,
4353 start: Instant,
4354 executed: bool,
4355 success: bool,
4356 output: String,
4357 metadata: HashMap<String, Value>,
4358 policy: ToolPolicyDecisionRecord,
4359 approval: Option<ToolApprovalRecord>,
4360 timed_out: bool,
4361 output_truncated: bool,
4362 ) -> ToolExecutionRecord {
4363 let versions = ToolDecisionVersions {
4364 policy: self.active_tool_security().policy_version(),
4365 registry: self.tools.version(),
4366 runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4367 state: self
4368 .state_machine
4369 .as_ref()
4370 .map(|state_machine| state_machine.generation()),
4371 };
4372 self.record_from_parts_at(
4373 request,
4374 canonical_id,
4375 executed_arguments,
4376 started_at,
4377 start,
4378 executed,
4379 success,
4380 output,
4381 metadata,
4382 policy,
4383 approval,
4384 timed_out,
4385 output_truncated,
4386 versions,
4387 )
4388 }
4389
4390 #[allow(clippy::too_many_arguments)]
4392 fn record_from_parts_at(
4393 &self,
4394 request: &ToolExecutionRequest,
4395 canonical_id: String,
4396 executed_arguments: Value,
4397 started_at: chrono::DateTime<chrono::Utc>,
4398 start: Instant,
4399 executed: bool,
4400 success: bool,
4401 output: String,
4402 metadata: HashMap<String, Value>,
4403 policy: ToolPolicyDecisionRecord,
4404 approval: Option<ToolApprovalRecord>,
4405 timed_out: bool,
4406 output_truncated: bool,
4407 versions: ToolDecisionVersions,
4408 ) -> ToolExecutionRecord {
4409 ToolExecutionRecord {
4410 call_id: request.call_id.clone(),
4411 requested_name: request.requested_name.clone(),
4412 canonical_id,
4413 source: request.source.clone(),
4414 arguments: request.arguments.clone(),
4415 executed_arguments,
4416 policy_version: versions.policy,
4417 registry_version: versions.registry,
4418 runtime_config_version: versions.runtime_control,
4419 executed,
4420 success,
4421 output,
4422 metadata,
4423 policy,
4424 approval,
4425 started_at,
4426 duration_ms: start.elapsed().as_millis() as u64,
4427 timed_out,
4428 cancelled: false,
4429 cancellation_reason: None,
4430 output_truncated,
4431 }
4432 }
4433
4434 async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4436 let result = ToolResult {
4437 success: record.success,
4438 output: record.model_output_string(),
4439 metadata: if record.metadata.is_empty() {
4440 None
4441 } else {
4442 Some(record.metadata.clone())
4443 },
4444 };
4445 self.hooks
4446 .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4447 .await;
4448 self.hooks.on_tool_execution_record(record).await;
4449 self.record_tool_call(&record.canonical_id, record.model_output_value());
4450 if !record.success {
4451 self.hooks
4452 .on_error(&AgentError::Tool(record.output.clone()))
4453 .await;
4454 }
4455 }
4456
4457 async fn finish_tool_record_after_resource_guards(
4459 &self,
4460 resource_guards: ToolResourceGuards,
4461 record: &ToolExecutionRecord,
4462 ) {
4463 drop(resource_guards);
4464 self.finish_tool_record(record).await;
4465 }
4466
4467 async fn execute_resolved_tool_once(
4469 &self,
4470 tool: Arc<dyn ai_agents_core::Tool>,
4471 args: Value,
4472 ctx: ToolExecutionContext,
4473 timeout_ms: u64,
4474 ) -> Result<(ToolResult, bool, bool, bool)> {
4475 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4476 return Ok((
4477 ToolResult::error("Tool execution cancelled by runtime control"),
4478 false,
4479 true,
4480 false,
4481 ));
4482 }
4483 let invoked = Arc::new(AtomicBool::new(false));
4487 let invoked_by_future = Arc::clone(&invoked);
4488 let actor_context = current_turn_actor_context();
4489 let future = async move {
4490 invoked_by_future.store(true, Ordering::SeqCst);
4491 if let Some(actor_context) = actor_context {
4492 scope_actor_context(actor_context, tool.execute(args, ctx)).await
4493 } else {
4494 tool.execute(args, ctx).await
4495 }
4496 };
4497 tokio::pin!(future);
4498 let timeout = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms));
4499 tokio::pin!(timeout);
4500 let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4501
4502 loop {
4503 tokio::select! {
4504 result = &mut future => return Ok((result, false, false, true)),
4505 _ = &mut timeout => {
4506 return Ok((
4507 ToolResult::error("Tool execution timed out"),
4508 true,
4509 false,
4510 invoked.load(Ordering::SeqCst),
4511 ));
4512 }
4513 _ = cancel_tick.tick() => {
4514 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4515 return Ok((
4516 ToolResult::error("Tool execution cancelled by runtime control"),
4517 false,
4518 true,
4519 invoked.load(Ordering::SeqCst),
4520 ));
4521 }
4522 }
4523 }
4524 }
4525 }
4526
4527 fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4529 let Some(max_chars) = max_chars else {
4530 return (output, false);
4531 };
4532 let mut chars = output.chars();
4533 let truncated: String = chars.by_ref().take(max_chars).collect();
4534 if chars.next().is_some() {
4535 (truncated, true)
4536 } else {
4537 (output, false)
4538 }
4539 }
4540
4541 async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4543 let locks = {
4544 let mut table = self.resource_locks.write();
4545 table.retain(|_, lock| lock.strong_count() > 0);
4546 keys.iter()
4547 .map(|key| {
4548 if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4549 lock
4550 } else {
4551 let lock = Arc::new(tokio::sync::Mutex::new(()));
4552 table.insert(key.clone(), Arc::downgrade(&lock));
4553 lock
4554 }
4555 })
4556 .collect::<Vec<_>>()
4557 };
4558 let mut resource_guards = ToolResourceGuards {
4559 guards: Vec::with_capacity(locks.len()),
4560 locks: Arc::clone(&self.resource_locks),
4561 };
4562 let mut locks = locks.into_iter();
4563 while let Some(lock) = locks.next() {
4564 let mut lock = Box::pin(lock.lock_owned());
4565 loop {
4566 tokio::select! {
4567 guard = &mut lock => {
4568 resource_guards.guards.push(guard);
4569 break;
4570 }
4571 _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4572 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4573 drop(lock);
4574 drop(locks);
4575 drop(resource_guards);
4576 return None;
4577 }
4578 }
4579 }
4580 }
4581 }
4582 Some(resource_guards)
4583 }
4584
4585 async fn run_tool_with_retries(
4587 &self,
4588 canonical_id: &str,
4589 tool: Arc<dyn ai_agents_core::Tool>,
4590 args: Value,
4591 ctx: ToolExecutionContext,
4592 timeout_ms: u64,
4593 max_retries: u32,
4594 ) -> Result<(ToolResult, bool, bool, bool)> {
4595 let max_retries = if ctx.classification.safely_retryable {
4596 max_retries
4597 } else {
4598 0
4599 };
4600 let mut attempts = 0;
4601 let mut invoked = false;
4602 loop {
4603 let (result, timed_out, cancelled, attempt_invoked) = self
4604 .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout_ms)
4605 .await?;
4606 invoked |= attempt_invoked;
4607 if result.success || timed_out || cancelled || attempts >= max_retries {
4608 return Ok((result, timed_out, cancelled, invoked));
4609 }
4610 attempts += 1;
4611 warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4612 }
4613 }
4614
4615 fn execute_tool_record(
4617 &self,
4618 request: ToolExecutionRequest,
4619 ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4620 Box::pin(self.execute_tool_record_inner(request))
4621 }
4622
4623 async fn execute_tool_record_inner(
4624 &self,
4625 request: ToolExecutionRequest,
4626 ) -> Result<ToolExecutionRecord> {
4627 let started_at = chrono::Utc::now();
4628 let start = Instant::now();
4629 info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4630
4631 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4632 let record = self.record_from_parts(
4633 &request,
4634 request.requested_name.clone(),
4635 request.arguments.clone(),
4636 started_at,
4637 start,
4638 false,
4639 false,
4640 "Tool execution is disabled by runtime control".to_string(),
4641 HashMap::new(),
4642 ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4643 None,
4644 false,
4645 false,
4646 );
4647 self.finish_tool_record(&record).await;
4648 return Ok(record);
4649 }
4650
4651 let Some(resolved) = self.tools.resolve(&request.requested_name) else {
4652 let record = self.record_from_parts(
4653 &request,
4654 request.requested_name.clone(),
4655 request.arguments.clone(),
4656 started_at,
4657 start,
4658 false,
4659 false,
4660 format!("Tool '{}' is unavailable", request.requested_name),
4661 HashMap::new(),
4662 ToolPolicyDecisionRecord::unavailable(format!(
4663 "Tool '{}' is not registered",
4664 request.requested_name
4665 )),
4666 None,
4667 false,
4668 false,
4669 );
4670 self.finish_tool_record(&record).await;
4671 return Ok(record);
4672 };
4673
4674 let canonical_id = resolved.identity.canonical_id.clone();
4675
4676 let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4677 if !initial_scope_snapshot
4678 .tool_ids
4679 .iter()
4680 .any(|id| id == &canonical_id)
4681 {
4682 let record = self.record_from_parts(
4683 &request,
4684 canonical_id.clone(),
4685 request.arguments.clone(),
4686 started_at,
4687 start,
4688 false,
4689 false,
4690 format!(
4691 "Tool '{}' is not available in the current scope",
4692 canonical_id
4693 ),
4694 HashMap::new(),
4695 ToolPolicyDecisionRecord::deny(format!(
4696 "Tool '{}' is not granted by the current top-level and state tool scope",
4697 canonical_id
4698 )),
4699 None,
4700 false,
4701 false,
4702 );
4703 self.finish_tool_record(&record).await;
4704 return Ok(record);
4705 }
4706
4707 let approval_control_snapshot = self.runtime_safety_snapshot();
4708 let security_engine = approval_control_snapshot.tool_security.clone();
4709 let bindings = resolved.tool.policy_bindings();
4710 let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4711 &canonical_id,
4712 &request.arguments,
4713 &bindings,
4714 );
4715 self.hooks
4716 .on_tool_start(&canonical_id, &executed_arguments)
4717 .await;
4718
4719 let mut metadata = HashMap::new();
4720 let safety = resolved.tool.safety_metadata();
4721 let classification = resolved.tool.classify_call(&executed_arguments);
4722 let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
4723 metadata.insert(
4724 "classification".to_string(),
4725 serde_json::to_value(&classification).unwrap_or(Value::Null),
4726 );
4727 metadata.insert(
4728 "effective_limits".to_string(),
4729 serde_json::to_value(&limits).unwrap_or(Value::Null),
4730 );
4731 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
4732 if !policy_snapshot.is_null() {
4733 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
4734 }
4735
4736 let mut approval_record = Some(ToolApprovalRecord {
4737 status: ToolApprovalStatus::NotRequired,
4738 reason: None,
4739 modified_arguments: None,
4740 });
4741
4742 let mut security_result = security_engine
4743 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
4744 .await?;
4745 match &security_result {
4746 SecurityCheckResult::Allow => {}
4747 SecurityCheckResult::Warn { message } => {
4748 warn!(tool = %canonical_id, message = %message, "Tool security warning");
4749 }
4750 SecurityCheckResult::Block { reason } => {
4751 let record = self.record_from_parts(
4752 &request,
4753 canonical_id,
4754 executed_arguments,
4755 started_at,
4756 start,
4757 false,
4758 false,
4759 format!("Denied: {}", reason),
4760 metadata,
4761 ToolPolicyDecisionRecord::deny(reason.clone()),
4762 approval_record,
4763 false,
4764 false,
4765 );
4766 self.finish_tool_record(&record).await;
4767 return Ok(record);
4768 }
4769 SecurityCheckResult::Unavailable { reason } => {
4770 let record = self.record_from_parts(
4771 &request,
4772 canonical_id,
4773 executed_arguments,
4774 started_at,
4775 start,
4776 false,
4777 false,
4778 format!("Unavailable: {}", reason),
4779 metadata,
4780 ToolPolicyDecisionRecord::unavailable(reason.clone()),
4781 approval_record,
4782 false,
4783 false,
4784 );
4785 self.finish_tool_record(&record).await;
4786 return Ok(record);
4787 }
4788 SecurityCheckResult::RequireConfirmation { message } => {
4789 if self.hitl_engine.is_none() {
4790 approval_record = Some(ToolApprovalRecord {
4791 status: ToolApprovalStatus::Unavailable,
4792 reason: Some("No HITL engine configured".to_string()),
4793 modified_arguments: None,
4794 });
4795 let record = self.record_from_parts(
4796 &request,
4797 canonical_id,
4798 executed_arguments,
4799 started_at,
4800 start,
4801 false,
4802 false,
4803 format!("Approval unavailable: {}", message),
4804 metadata,
4805 ToolPolicyDecisionRecord::approval(message.clone()),
4806 approval_record,
4807 false,
4808 false,
4809 );
4810 self.finish_tool_record(&record).await;
4811 return Ok(record);
4812 }
4813
4814 let check_result = HITLCheckResult::required(
4815 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4816 HashMap::new(),
4817 message.clone(),
4818 None,
4819 );
4820 match self.request_hitl_approval(check_result).await? {
4821 ApprovalResult::Approved => {
4822 merge_approved_record(&mut approval_record);
4823 }
4824 ApprovalResult::Modified { changes } => {
4825 if let Some(obj) = executed_arguments.as_object_mut() {
4826 for (key, value) in changes {
4827 obj.insert(key, value);
4828 }
4829 }
4830 security_result = security_engine
4831 .validate_tool_execution_with_bindings(
4832 &canonical_id,
4833 &executed_arguments,
4834 &bindings,
4835 )
4836 .await?;
4837 if !matches!(
4838 security_result,
4839 SecurityCheckResult::Allow
4840 | SecurityCheckResult::Warn { .. }
4841 | SecurityCheckResult::RequireConfirmation { .. }
4842 ) {
4843 let reason = security_result
4844 .reason()
4845 .unwrap_or("modified arguments failed policy")
4846 .to_string();
4847 let record = self.record_from_parts(
4848 &request,
4849 canonical_id,
4850 executed_arguments.clone(),
4851 started_at,
4852 start,
4853 false,
4854 false,
4855 reason.clone(),
4856 metadata,
4857 ToolPolicyDecisionRecord::deny(reason),
4858 Some(ToolApprovalRecord {
4859 status: ToolApprovalStatus::Modified,
4860 reason: None,
4861 modified_arguments: Some(executed_arguments),
4862 }),
4863 false,
4864 false,
4865 );
4866 self.finish_tool_record(&record).await;
4867 return Ok(record);
4868 }
4869 approval_record = Some(ToolApprovalRecord {
4870 status: ToolApprovalStatus::Modified,
4871 reason: None,
4872 modified_arguments: Some(executed_arguments.clone()),
4873 });
4874 }
4875 ApprovalResult::Rejected { reason } => {
4876 let reason = reason.unwrap_or_else(|| "rejected".to_string());
4877 approval_record = Some(ToolApprovalRecord {
4878 status: ToolApprovalStatus::Rejected,
4879 reason: Some(reason.clone()),
4880 modified_arguments: None,
4881 });
4882 let record = self.record_from_parts(
4883 &request,
4884 canonical_id,
4885 executed_arguments,
4886 started_at,
4887 start,
4888 false,
4889 false,
4890 format!("Approval rejected: {}", reason),
4891 metadata,
4892 ToolPolicyDecisionRecord::approval(reason),
4893 approval_record,
4894 false,
4895 false,
4896 );
4897 self.finish_tool_record(&record).await;
4898 return Ok(record);
4899 }
4900 ApprovalResult::Timeout => {
4901 approval_record = Some(ToolApprovalRecord {
4902 status: ToolApprovalStatus::Timeout,
4903 reason: Some("approval timeout".to_string()),
4904 modified_arguments: None,
4905 });
4906 let record = self.record_from_parts(
4907 &request,
4908 canonical_id,
4909 executed_arguments,
4910 started_at,
4911 start,
4912 false,
4913 false,
4914 "Approval timed out".to_string(),
4915 metadata,
4916 ToolPolicyDecisionRecord::approval("approval timeout"),
4917 approval_record,
4918 false,
4919 false,
4920 );
4921 self.finish_tool_record(&record).await;
4922 return Ok(record);
4923 }
4924 }
4925 }
4926 }
4927
4928 if canonical_id == "command" && !self.tools.command_runner_available() {
4929 let record = self.record_from_parts(
4930 &request,
4931 canonical_id.clone(),
4932 executed_arguments.clone(),
4933 started_at,
4934 start,
4935 false,
4936 false,
4937 "Command runner is unavailable".to_string(),
4938 metadata,
4939 ToolPolicyDecisionRecord::unavailable("command runner is unavailable"),
4940 Some(ToolApprovalRecord {
4941 status: ToolApprovalStatus::Unavailable,
4942 reason: Some("command runner is unavailable".to_string()),
4943 modified_arguments: None,
4944 }),
4945 false,
4946 false,
4947 );
4948 self.finish_tool_record(&record).await;
4949 return Ok(record);
4950 }
4951
4952 if approval_record
4953 .as_ref()
4954 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
4955 && let Some(message) =
4956 security_engine.classification_approval_message(&canonical_id, &classification)
4957 {
4958 if self.hitl_engine.is_none() {
4959 approval_record = Some(ToolApprovalRecord {
4960 status: ToolApprovalStatus::Unavailable,
4961 reason: Some("No HITL engine configured".to_string()),
4962 modified_arguments: None,
4963 });
4964 let record = self.record_from_parts(
4965 &request,
4966 canonical_id,
4967 executed_arguments,
4968 started_at,
4969 start,
4970 false,
4971 false,
4972 format!("Approval unavailable: {}", message),
4973 metadata,
4974 ToolPolicyDecisionRecord::approval(message),
4975 approval_record,
4976 false,
4977 false,
4978 );
4979 self.finish_tool_record(&record).await;
4980 return Ok(record);
4981 }
4982 let check_result = HITLCheckResult::required(
4983 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4984 HashMap::new(),
4985 message.clone(),
4986 None,
4987 );
4988 match self.request_hitl_approval(check_result).await? {
4989 ApprovalResult::Approved => {
4990 merge_approved_record(&mut approval_record);
4991 }
4992 ApprovalResult::Modified { changes } => {
4993 if let Some(obj) = executed_arguments.as_object_mut() {
4994 for (key, value) in changes {
4995 obj.insert(key, value);
4996 }
4997 }
4998 let modified_security = security_engine
4999 .validate_tool_execution_with_bindings(
5000 &canonical_id,
5001 &executed_arguments,
5002 &bindings,
5003 )
5004 .await?;
5005 if !matches!(
5006 modified_security,
5007 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5008 ) {
5009 let reason = modified_security
5010 .reason()
5011 .unwrap_or("modified arguments failed policy")
5012 .to_string();
5013 let record = self.record_from_parts(
5014 &request,
5015 canonical_id,
5016 executed_arguments.clone(),
5017 started_at,
5018 start,
5019 false,
5020 false,
5021 reason.clone(),
5022 metadata,
5023 ToolPolicyDecisionRecord::deny(reason),
5024 Some(ToolApprovalRecord {
5025 status: ToolApprovalStatus::Modified,
5026 reason: None,
5027 modified_arguments: Some(executed_arguments),
5028 }),
5029 false,
5030 false,
5031 );
5032 self.finish_tool_record(&record).await;
5033 return Ok(record);
5034 }
5035 approval_record = Some(ToolApprovalRecord {
5036 status: ToolApprovalStatus::Modified,
5037 reason: None,
5038 modified_arguments: Some(executed_arguments.clone()),
5039 });
5040 }
5041 ApprovalResult::Rejected { reason } => {
5042 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5043 let record = self.record_from_parts(
5044 &request,
5045 canonical_id,
5046 executed_arguments,
5047 started_at,
5048 start,
5049 false,
5050 false,
5051 format!("Approval rejected: {}", reason),
5052 metadata,
5053 ToolPolicyDecisionRecord::approval(reason.clone()),
5054 Some(ToolApprovalRecord {
5055 status: ToolApprovalStatus::Rejected,
5056 reason: Some(reason),
5057 modified_arguments: None,
5058 }),
5059 false,
5060 false,
5061 );
5062 self.finish_tool_record(&record).await;
5063 return Ok(record);
5064 }
5065 ApprovalResult::Timeout => {
5066 let record = self.record_from_parts(
5067 &request,
5068 canonical_id,
5069 executed_arguments,
5070 started_at,
5071 start,
5072 false,
5073 false,
5074 "Approval timed out".to_string(),
5075 metadata,
5076 ToolPolicyDecisionRecord::approval("approval timeout"),
5077 Some(ToolApprovalRecord {
5078 status: ToolApprovalStatus::Timeout,
5079 reason: Some("approval timeout".to_string()),
5080 modified_arguments: None,
5081 }),
5082 false,
5083 false,
5084 );
5085 self.finish_tool_record(&record).await;
5086 return Ok(record);
5087 }
5088 }
5089 }
5090
5091 if canonical_id == "diagnostics" && !self.tools.diagnostics_available() {
5092 let record = self.record_from_parts(
5093 &request,
5094 canonical_id.clone(),
5095 executed_arguments.clone(),
5096 started_at,
5097 start,
5098 false,
5099 false,
5100 "Diagnostics provider is unavailable".to_string(),
5101 metadata,
5102 ToolPolicyDecisionRecord::unavailable("diagnostics provider is unavailable"),
5103 Some(ToolApprovalRecord {
5104 status: ToolApprovalStatus::Unavailable,
5105 reason: Some("diagnostics provider is unavailable".to_string()),
5106 modified_arguments: None,
5107 }),
5108 false,
5109 false,
5110 );
5111 self.finish_tool_record(&record).await;
5112 return Ok(record);
5113 }
5114
5115 if canonical_id == "web_search" && !self.tools.web_search_available() {
5116 let record = self.record_from_parts(
5117 &request,
5118 canonical_id.clone(),
5119 executed_arguments.clone(),
5120 started_at,
5121 start,
5122 false,
5123 false,
5124 "Web search provider is unavailable".to_string(),
5125 metadata,
5126 ToolPolicyDecisionRecord::unavailable("web search provider is unavailable"),
5127 Some(ToolApprovalRecord {
5128 status: ToolApprovalStatus::Unavailable,
5129 reason: Some("web search provider is unavailable".to_string()),
5130 modified_arguments: None,
5131 }),
5132 false,
5133 false,
5134 );
5135 self.finish_tool_record(&record).await;
5136 return Ok(record);
5137 }
5138
5139 let hitl_lang_ctx = self.build_hitl_language_context();
5140 if let Some(ref hitl_engine) = self.hitl_engine {
5141 let check_result = self
5142 .observe_purpose(
5143 ObservationPurpose::HitlLocalization,
5144 hitl_engine.check_tool_with_localization(
5145 &canonical_id,
5146 &executed_arguments,
5147 &hitl_lang_ctx,
5148 self.approval_handler.as_ref(),
5149 Some(&self.llm_registry),
5150 ),
5151 )
5152 .await?;
5153 if check_result.is_required() {
5154 match self.request_hitl_approval(check_result).await? {
5155 ApprovalResult::Approved => {
5156 merge_approved_record(&mut approval_record);
5157 }
5158 ApprovalResult::Modified { changes } => {
5159 if let Some(obj) = executed_arguments.as_object_mut() {
5160 for (key, value) in changes {
5161 obj.insert(key, value);
5162 }
5163 }
5164 let modified_security = security_engine
5165 .validate_tool_execution_with_bindings(
5166 &canonical_id,
5167 &executed_arguments,
5168 &bindings,
5169 )
5170 .await?;
5171 if !matches!(
5172 modified_security,
5173 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5174 ) {
5175 let reason = modified_security
5176 .reason()
5177 .unwrap_or("modified arguments failed policy")
5178 .to_string();
5179 let record = self.record_from_parts(
5180 &request,
5181 canonical_id,
5182 executed_arguments.clone(),
5183 started_at,
5184 start,
5185 false,
5186 false,
5187 reason.clone(),
5188 metadata,
5189 ToolPolicyDecisionRecord::deny(reason),
5190 Some(ToolApprovalRecord {
5191 status: ToolApprovalStatus::Modified,
5192 reason: None,
5193 modified_arguments: Some(executed_arguments),
5194 }),
5195 false,
5196 false,
5197 );
5198 self.finish_tool_record(&record).await;
5199 return Ok(record);
5200 }
5201 approval_record = Some(ToolApprovalRecord {
5202 status: ToolApprovalStatus::Modified,
5203 reason: None,
5204 modified_arguments: Some(executed_arguments.clone()),
5205 });
5206 }
5207 ApprovalResult::Rejected { reason } => {
5208 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5209 let record = self.record_from_parts(
5210 &request,
5211 canonical_id,
5212 executed_arguments,
5213 started_at,
5214 start,
5215 false,
5216 false,
5217 format!("Approval rejected: {}", reason),
5218 metadata,
5219 ToolPolicyDecisionRecord::approval(reason.clone()),
5220 Some(ToolApprovalRecord {
5221 status: ToolApprovalStatus::Rejected,
5222 reason: Some(reason),
5223 modified_arguments: None,
5224 }),
5225 false,
5226 false,
5227 );
5228 self.finish_tool_record(&record).await;
5229 return Ok(record);
5230 }
5231 ApprovalResult::Timeout => {
5232 let record = self.record_from_parts(
5233 &request,
5234 canonical_id,
5235 executed_arguments,
5236 started_at,
5237 start,
5238 false,
5239 false,
5240 "Approval timed out".to_string(),
5241 metadata,
5242 ToolPolicyDecisionRecord::approval("approval timeout"),
5243 Some(ToolApprovalRecord {
5244 status: ToolApprovalStatus::Timeout,
5245 reason: Some("approval timeout".to_string()),
5246 modified_arguments: None,
5247 }),
5248 false,
5249 false,
5250 );
5251 self.finish_tool_record(&record).await;
5252 return Ok(record);
5253 }
5254 }
5255 }
5256
5257 let condition_check = self
5258 .observe_purpose(
5259 ObservationPurpose::HitlLocalization,
5260 hitl_engine.check_conditions_with_localization(
5261 &executed_arguments,
5262 &hitl_lang_ctx,
5263 self.approval_handler.as_ref(),
5264 Some(&self.llm_registry),
5265 ),
5266 )
5267 .await?;
5268 if condition_check.is_required() {
5269 match self.request_hitl_approval(condition_check).await? {
5270 ApprovalResult::Approved => {
5271 merge_approved_record(&mut approval_record);
5272 }
5273 ApprovalResult::Modified { changes } => {
5274 if let Some(obj) = executed_arguments.as_object_mut() {
5275 for (key, value) in changes {
5276 obj.insert(key, value);
5277 }
5278 }
5279 let modified_security = security_engine
5280 .validate_tool_execution_with_bindings(
5281 &canonical_id,
5282 &executed_arguments,
5283 &bindings,
5284 )
5285 .await?;
5286 if !matches!(
5287 modified_security,
5288 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5289 ) {
5290 let reason = modified_security
5291 .reason()
5292 .unwrap_or("modified arguments failed policy")
5293 .to_string();
5294 let record = self.record_from_parts(
5295 &request,
5296 canonical_id,
5297 executed_arguments,
5298 started_at,
5299 start,
5300 false,
5301 false,
5302 reason.clone(),
5303 metadata,
5304 ToolPolicyDecisionRecord::deny(reason),
5305 approval_record,
5306 false,
5307 false,
5308 );
5309 self.finish_tool_record(&record).await;
5310 return Ok(record);
5311 }
5312 approval_record = Some(ToolApprovalRecord {
5313 status: ToolApprovalStatus::Modified,
5314 reason: None,
5315 modified_arguments: Some(executed_arguments.clone()),
5316 });
5317 }
5318 ApprovalResult::Rejected { reason } => {
5319 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5320 let record = self.record_from_parts(
5321 &request,
5322 canonical_id,
5323 executed_arguments,
5324 started_at,
5325 start,
5326 false,
5327 false,
5328 format!("Approval rejected: {}", reason),
5329 metadata,
5330 ToolPolicyDecisionRecord::approval(reason.clone()),
5331 Some(ToolApprovalRecord {
5332 status: ToolApprovalStatus::Rejected,
5333 reason: Some(reason),
5334 modified_arguments: None,
5335 }),
5336 false,
5337 false,
5338 );
5339 self.finish_tool_record(&record).await;
5340 return Ok(record);
5341 }
5342 ApprovalResult::Timeout => {
5343 let record = self.record_from_parts(
5344 &request,
5345 canonical_id,
5346 executed_arguments,
5347 started_at,
5348 start,
5349 false,
5350 false,
5351 "Approval timed out".to_string(),
5352 metadata,
5353 ToolPolicyDecisionRecord::approval("approval timeout"),
5354 Some(ToolApprovalRecord {
5355 status: ToolApprovalStatus::Timeout,
5356 reason: Some("approval timeout".to_string()),
5357 modified_arguments: None,
5358 }),
5359 false,
5360 false,
5361 );
5362 self.finish_tool_record(&record).await;
5363 return Ok(record);
5364 }
5365 }
5366 }
5367 }
5368
5369 executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5374 &canonical_id,
5375 &executed_arguments,
5376 &bindings,
5377 );
5378 if let Some(record) = approval_record.as_mut()
5379 && matches!(record.status, ToolApprovalStatus::Modified)
5380 {
5381 record.modified_arguments = Some(executed_arguments.clone());
5382 }
5383 let binding_security_result = security_engine
5384 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5385 .await?;
5386 let approval_confirmation_required = matches!(
5387 binding_security_result,
5388 SecurityCheckResult::RequireConfirmation { .. }
5389 ) || security_engine
5390 .classification_approval_message(
5391 &canonical_id,
5392 &resolved.tool.classify_call(&executed_arguments),
5393 )
5394 .is_some();
5395 let approval_binding = approval_record.as_ref().and_then(|record| {
5396 matches!(
5397 record.status,
5398 ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5399 )
5400 .then(|| ToolApprovalBinding {
5401 canonical_id: canonical_id.clone(),
5402 arguments: executed_arguments.clone(),
5403 confirmation_required: approval_confirmation_required,
5404 policy_version: security_engine.policy_version(),
5405 runtime_control_version: approval_control_snapshot.version,
5406 state_generation: initial_scope_snapshot.state_generation,
5407 reviewed_tool: Arc::clone(&resolved.tool),
5408 })
5409 });
5410
5411 let control_snapshot = self.runtime_safety_snapshot();
5416 let resolved = self.tools.resolve(&request.requested_name);
5417 let registry_version = self.tools.version();
5418 let mut versions = ToolDecisionVersions {
5419 policy: control_snapshot.tool_security.policy_version(),
5420 registry: registry_version,
5421 runtime_control: control_snapshot.version,
5422 state: None,
5423 };
5424 metadata.insert(
5425 "runtime_scope_snapshot".to_string(),
5426 serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5427 );
5428 let resolved = match resolved {
5429 Some(resolved) => resolved,
5430 None => {
5431 let reason = format!(
5432 "Tool '{}' became unavailable after approval",
5433 request.requested_name
5434 );
5435 let record = self.record_from_parts_at(
5436 &request,
5437 request.requested_name.clone(),
5438 executed_arguments,
5439 started_at,
5440 start,
5441 false,
5442 false,
5443 reason.clone(),
5444 metadata,
5445 ToolPolicyDecisionRecord::unavailable(reason),
5446 approval_record,
5447 false,
5448 false,
5449 versions,
5450 );
5451 self.finish_tool_record(&record).await;
5452 return Ok(record);
5453 }
5454 };
5455
5456 let canonical_id = resolved.identity.canonical_id.clone();
5457 let bindings = resolved.tool.policy_bindings();
5458 let final_arguments = control_snapshot
5459 .tool_security
5460 .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5461 if let Some(record) = approval_record.as_mut()
5462 && matches!(record.status, ToolApprovalStatus::Modified)
5463 {
5464 record.modified_arguments = Some(final_arguments.clone());
5465 }
5466 let classification = resolved.tool.classify_call(&final_arguments);
5467 let safety = resolved.tool.safety_metadata();
5468 let security_engine = control_snapshot.tool_security;
5469 let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
5470 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5471 let resource_lock_keys =
5472 tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5473 metadata.insert(
5474 "classification".to_string(),
5475 serde_json::to_value(&classification).unwrap_or(Value::Null),
5476 );
5477 metadata.insert(
5478 "effective_limits".to_string(),
5479 serde_json::to_value(&limits).unwrap_or(Value::Null),
5480 );
5481 metadata.insert(
5482 "resource_lock_keys".to_string(),
5483 serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5484 );
5485 if policy_snapshot.is_null() {
5486 metadata.remove("policy_snapshot");
5487 } else {
5488 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5489 }
5490
5491 let final_denial = |canonical_id: String,
5492 output: String,
5493 policy: ToolPolicyDecisionRecord,
5494 metadata: HashMap<String, Value>,
5495 decision_versions: ToolDecisionVersions| {
5496 self.record_from_parts_at(
5497 &request,
5498 canonical_id,
5499 final_arguments.clone(),
5500 started_at,
5501 start,
5502 false,
5503 false,
5504 output,
5505 metadata,
5506 policy,
5507 approval_record.clone(),
5508 false,
5509 false,
5510 decision_versions,
5511 )
5512 };
5513
5514 if control_snapshot.emergency_deny {
5515 let reason = "Tool execution is disabled by runtime control".to_string();
5516 let record = final_denial(
5517 canonical_id,
5518 reason.clone(),
5519 ToolPolicyDecisionRecord::deny(reason),
5520 metadata,
5521 versions,
5522 );
5523 self.finish_tool_record(&record).await;
5524 return Ok(record);
5525 }
5526
5527 let available_snapshot = self
5532 .get_available_tool_ids_snapshot_for_scope(
5533 control_snapshot.tool_scope_override.as_deref(),
5534 )
5535 .await?;
5536 versions.state = available_snapshot.state_generation;
5537 metadata.insert(
5538 "available_tool_ids_snapshot".to_string(),
5539 serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5540 );
5541 metadata.insert(
5542 "state_generation_snapshot".to_string(),
5543 serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5544 );
5545 if !available_snapshot
5546 .tool_ids
5547 .iter()
5548 .any(|tool_id| tool_id == &canonical_id)
5549 {
5550 let reason = format!(
5551 "Tool '{}' is not available in the final runtime scope",
5552 canonical_id
5553 );
5554 let record = final_denial(
5555 canonical_id,
5556 reason.clone(),
5557 ToolPolicyDecisionRecord::deny(reason),
5558 metadata,
5559 versions,
5560 );
5561 self.finish_tool_record(&record).await;
5562 return Ok(record);
5563 }
5564
5565 let final_security_result = security_engine
5570 .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5571 .await?;
5572 match &final_security_result {
5573 SecurityCheckResult::Block { reason } => {
5574 let record = final_denial(
5575 canonical_id,
5576 format!("Denied: {}", reason),
5577 ToolPolicyDecisionRecord::deny(reason.clone()),
5578 metadata,
5579 versions,
5580 );
5581 self.finish_tool_record(&record).await;
5582 return Ok(record);
5583 }
5584 SecurityCheckResult::Unavailable { reason } => {
5585 let record = final_denial(
5586 canonical_id,
5587 format!("Unavailable: {}", reason),
5588 ToolPolicyDecisionRecord::unavailable(reason.clone()),
5589 metadata,
5590 versions,
5591 );
5592 self.finish_tool_record(&record).await;
5593 return Ok(record);
5594 }
5595 SecurityCheckResult::Warn { message } => {
5596 warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5597 }
5598 SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5599 }
5600 let final_confirmation_required = matches!(
5601 final_security_result,
5602 SecurityCheckResult::RequireConfirmation { .. }
5603 ) || security_engine
5604 .classification_approval_message(&canonical_id, &classification)
5605 .is_some();
5606 let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5607 binding.is_stale(
5608 &canonical_id,
5609 &final_arguments,
5610 final_confirmation_required,
5611 versions,
5612 &resolved.tool,
5613 )
5614 });
5615 if stale_approval {
5616 let reason = "Approval became stale before final admission".to_string();
5617 let record = final_denial(
5618 canonical_id,
5619 reason.clone(),
5620 ToolPolicyDecisionRecord::deny(reason),
5621 metadata,
5622 versions,
5623 );
5624 self.finish_tool_record(&record).await;
5625 return Ok(record);
5626 }
5627 if final_confirmation_required && approval_binding.is_none() {
5628 let reason = "Final policy requires fresh approval".to_string();
5629 let record = final_denial(
5630 canonical_id,
5631 reason.clone(),
5632 ToolPolicyDecisionRecord::approval(reason),
5633 metadata,
5634 versions,
5635 );
5636 self.finish_tool_record(&record).await;
5637 return Ok(record);
5638 }
5639
5640 let unavailable_reason = match canonical_id.as_str() {
5641 "command" if !self.tools.command_runner_available() => {
5642 Some("command runner is unavailable")
5643 }
5644 "diagnostics" if !self.tools.diagnostics_available() => {
5645 Some("diagnostics provider is unavailable")
5646 }
5647 "web_search" if !self.tools.web_search_available() => {
5648 Some("web search provider is unavailable")
5649 }
5650 _ => None,
5651 };
5652 if let Some(reason) = unavailable_reason {
5653 let record = final_denial(
5654 canonical_id,
5655 reason.to_string(),
5656 ToolPolicyDecisionRecord::unavailable(reason),
5657 metadata,
5658 versions,
5659 );
5660 self.finish_tool_record(&record).await;
5661 return Ok(record);
5662 }
5663
5664 let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5669 else {
5670 let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5671 let record = final_denial(
5672 canonical_id,
5673 reason.clone(),
5674 ToolPolicyDecisionRecord::deny(reason),
5675 metadata,
5676 versions,
5677 );
5678 self.finish_tool_record(&record).await;
5679 return Ok(record);
5680 };
5681
5682 let admission = self.admit_tool_execution(
5687 versions.runtime_control,
5688 versions.policy,
5689 versions.state,
5690 &canonical_id,
5691 );
5692 if !matches!(admission, SecurityCheckResult::Allow) {
5693 let latest_control = self.runtime_safety_snapshot();
5694 let reason = admission
5695 .reason()
5696 .unwrap_or("tool admission was denied")
5697 .to_string();
5698 let policy = if admission.is_unavailable() {
5699 ToolPolicyDecisionRecord::unavailable(reason.clone())
5700 } else {
5701 ToolPolicyDecisionRecord::deny(reason.clone())
5702 };
5703 let record = self.record_from_parts_at(
5704 &request,
5705 canonical_id,
5706 final_arguments,
5707 started_at,
5708 start,
5709 false,
5710 false,
5711 reason,
5712 metadata,
5713 policy,
5714 approval_record,
5715 false,
5716 false,
5717 ToolDecisionVersions {
5718 policy: latest_control.tool_security.policy_version(),
5719 registry: versions.registry,
5720 runtime_control: latest_control.version,
5721 state: self
5722 .state_machine
5723 .as_ref()
5724 .map(|state_machine| state_machine.generation()),
5725 },
5726 );
5727 self.finish_tool_record_after_resource_guards(resource_guards, &record)
5728 .await;
5729 return Ok(record);
5730 }
5731 let executed_arguments = final_arguments;
5732
5733 let tool_config = self.recovery_manager.get_tool_config(&canonical_id);
5734 let timeout_ms = limits
5735 .timeout_ms
5736 .unwrap_or_else(|| security_engine.get_tool_timeout(&canonical_id));
5737 let deadline = Some(started_at + chrono::Duration::milliseconds(timeout_ms as i64));
5738 let turn_actor = current_turn_actor_context();
5739 let actor = ToolActorContext {
5740 actor_id: turn_actor
5741 .as_ref()
5742 .and_then(|context| context.effective_actor_id().map(str::to_string))
5743 .or_else(|| self.actor_id()),
5744 origin_actor_id: turn_actor
5745 .as_ref()
5746 .and_then(|context| context.origin_actor_id.clone()),
5747 sender_agent_id: turn_actor
5748 .as_ref()
5749 .and_then(|context| context.sender_agent_id.clone()),
5750 };
5751 let tool_context = ToolExecutionContext {
5752 requested_name: request.requested_name.clone(),
5753 canonical_id: canonical_id.clone(),
5754 display_name: resolved.identity.display_name.clone(),
5755 provider_id: resolved.identity.provider_id.clone(),
5756 registry_version: versions.registry,
5757 policy_version: versions.policy,
5758 runtime_control_version: versions.runtime_control,
5759 call_id: request.call_id.clone(),
5760 source: request.source.clone(),
5761 actor,
5762 cancellation: ToolCancellationToken::new(
5763 Arc::clone(&self.runtime_control.emergency_deny),
5764 Some("runtime control cancellation".to_string()),
5765 ),
5766 started_at,
5767 deadline,
5768 permission: ToolPolicyDecisionRecord::allow(),
5769 approval: approval_record.clone(),
5770 classification: classification.clone(),
5771 safety,
5772 limits: limits.clone(),
5773 policy_snapshot,
5774 custom_config: security_engine.custom_config(&canonical_id),
5775 };
5776 let (mut result, timed_out, cancelled, invoked) = self
5777 .run_tool_with_retries(
5778 &canonical_id,
5779 resolved.tool.clone(),
5780 executed_arguments.clone(),
5781 tool_context,
5782 timeout_ms,
5783 tool_config.max_retries,
5784 )
5785 .await?;
5786
5787 if !result.success {
5788 match &tool_config.on_failure {
5789 ToolFailureAction::Skip => {
5790 result = ToolResult::ok(format!(
5791 "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
5792 canonical_id
5793 ));
5794 }
5795 ToolFailureAction::Fallback { fallback_tool } => {
5796 drop(resource_guards);
5797 let fallback_request = ToolExecutionRequest::new(
5798 request.call_id.clone(),
5799 fallback_tool.clone(),
5800 executed_arguments,
5801 ToolCallSource::Fallback {
5802 original_tool: canonical_id,
5803 },
5804 );
5805 return Box::pin(self.execute_tool_record(fallback_request)).await;
5806 }
5807 ToolFailureAction::ReportError => {}
5808 }
5809 }
5810
5811 let output_cap = limits.max_output_chars;
5812 let (output, output_truncated) =
5813 Self::truncate_tool_output(result.output.clone(), output_cap);
5814 if let Some(result_metadata) = result.metadata {
5815 metadata.extend(result_metadata);
5816 }
5817 let mut record = self.record_from_parts_at(
5818 &request,
5819 canonical_id,
5820 executed_arguments,
5821 started_at,
5822 start,
5823 invoked,
5824 result.success,
5825 output,
5826 metadata,
5827 ToolPolicyDecisionRecord::allow(),
5828 approval_record,
5829 timed_out,
5830 output_truncated,
5831 versions,
5832 );
5833 record.cancelled = cancelled;
5834 if cancelled {
5835 record.cancellation_reason = Some("runtime control cancellation".to_string());
5836 }
5837 self.finish_tool_record_after_resource_guards(resource_guards, &record)
5838 .await;
5839 Ok(record)
5840 }
5841
5842 #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
5843 async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
5844 let record = self
5845 .execute_tool_record(ToolExecutionRequest::new(
5846 tool_call.id.clone(),
5847 tool_call.name.clone(),
5848 tool_call.arguments.clone(),
5849 ToolCallSource::Model,
5850 ))
5851 .await?;
5852 if record.success {
5853 Ok(record.model_output_string())
5854 } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
5855 Err(AgentError::HITLRejected(record.model_output_string()))
5856 } else {
5857 Err(AgentError::Tool(record.model_output_string()))
5858 }
5859 }
5860
5861 async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
5867 let Some(ref router) = self.skill_router else {
5868 return Ok(None);
5869 };
5870 let available_skills = self.get_available_skills();
5871 if available_skills.is_empty() {
5872 return Ok(None);
5873 }
5874 let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
5875 let Some(skill_id) = self
5876 .observe_purpose(
5877 ObservationPurpose::SkillRouting,
5878 router.select_skill_filtered(input, &skill_ids),
5879 )
5880 .await?
5881 else {
5882 return Ok(None);
5883 };
5884 let skill = router
5885 .get_skill(&skill_id)
5886 .cloned()
5887 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
5888 info!(skill_id = %skill_id, "Skill selected");
5889 Ok(Some(SkillCandidate::new(skill_id, skill)))
5890 }
5891
5892 async fn commit_skill_candidate_route_result(
5897 &self,
5898 candidate: SkillCandidate,
5899 input: &str,
5900 ) -> Result<SkillRouteResult> {
5901 let skill_id = candidate.skill_id;
5902 let skill = candidate.skill;
5903 let expected_state_generation = self
5904 .state_machine
5905 .as_ref()
5906 .map(|state_machine| state_machine.generation());
5907 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
5908 if let Some(ref skill_disambig) = skill.disambiguation
5909 && skill_disambig.enabled.unwrap_or(false)
5910 && let Some(ref disambiguator) = self.disambiguation_manager
5911 {
5912 let context = self.build_disambiguation_context().await?;
5913 let state_override = self
5914 .state_machine
5915 .as_ref()
5916 .and_then(|sm| sm.current_definition())
5917 .and_then(|def| def.disambiguation.clone());
5918
5919 let disambiguation_result = self
5920 .observe_purpose(
5921 ObservationPurpose::DisambiguationDetection,
5922 disambiguator.process_input_with_override(
5923 input,
5924 &context,
5925 state_override.as_ref(),
5926 Some(skill_disambig),
5927 ),
5928 )
5929 .await?;
5930 let current_state_generation = self
5931 .state_machine
5932 .as_ref()
5933 .map(|state_machine| state_machine.generation());
5934 if current_state_generation != expected_state_generation
5935 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
5936 {
5937 disambiguator.clear_pending().await;
5938 *self.pending_skill_id.write() = None;
5939 return Err(AgentError::Other(
5940 "State or reset ownership changed during skill disambiguation".to_string(),
5941 ));
5942 }
5943 match disambiguation_result {
5944 DisambiguationResult::Clear => {
5945 debug!(skill_id = %skill_id, "Skill disambiguation: clear");
5946 }
5947 DisambiguationResult::NeedsClarification {
5948 question,
5949 detection,
5950 } => {
5951 let admission = self
5952 .admit_disambiguation_redispatch(
5953 expected_disambiguation_epoch,
5954 expected_state_generation,
5955 )
5956 .await?;
5957 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
5958 info!(
5959 skill_id = %skill_id,
5960 ambiguity_type = ?detection.ambiguity_type,
5961 confidence = detection.confidence,
5962 "Skill requires clarification before execution"
5963 );
5964 *self.pending_skill_id.write() = Some(skill_id.clone());
5965 let response = AgentResponse::new(&question.question).with_metadata(
5966 "disambiguation",
5967 serde_json::json!({
5968 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
5969 "skill_id": skill_id,
5970 "options": question.options,
5971 "clarifying": question.clarifying,
5972 "detection": {
5973 "type": detection.ambiguity_type,
5974 "confidence": detection.confidence,
5975 "what_is_unclear": detection.what_is_unclear,
5976 }
5977 }),
5978 );
5979 drop(admission);
5980 return Ok(SkillRouteResult::NeedsClarification {
5981 response,
5982 ownership: Some(DisambiguationOwnership {
5983 epoch: expected_disambiguation_epoch,
5984 state_generation: expected_state_generation,
5985 }),
5986 });
5987 }
5988 DisambiguationResult::Clarified { enriched_input, .. } => {
5989 info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
5990 let admission = self
5991 .admit_disambiguation_redispatch(
5992 expected_disambiguation_epoch,
5993 expected_state_generation,
5994 )
5995 .await?;
5996 drop(admission);
5997 let content = self.execute_skill(&skill, &enriched_input).await?;
5998 return Ok(SkillRouteResult::Response { skill_id, content });
5999 }
6000 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
6001 info!(skill_id = %skill_id, "Skill disambiguation best guess");
6002 let admission = self
6003 .admit_disambiguation_redispatch(
6004 expected_disambiguation_epoch,
6005 expected_state_generation,
6006 )
6007 .await?;
6008 drop(admission);
6009 let content = self.execute_skill(&skill, &enriched_input).await?;
6010 return Ok(SkillRouteResult::Response { skill_id, content });
6011 }
6012 DisambiguationResult::GiveUp { reason } => {
6013 warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
6014 let apology = self
6015 .generate_localized_apology(
6016 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
6017 &reason,
6018 )
6019 .await
6020 .unwrap_or_else(|_| {
6021 format!("I'm sorry, I couldn't understand your request: {}", reason)
6022 });
6023 return Ok(SkillRouteResult::NeedsClarification {
6024 response: AgentResponse::new(&apology),
6025 ownership: None,
6026 });
6027 }
6028 DisambiguationResult::Escalate { reason } => {
6029 info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6030 let apology = self
6031 .generate_localized_apology(
6032 "Explain briefly that you're transferring the user to a human agent for help.",
6033 &reason,
6034 )
6035 .await
6036 .unwrap_or_else(|_| {
6037 format!("I need human assistance to help with your request: {}", reason)
6038 });
6039 return Ok(SkillRouteResult::NeedsClarification {
6040 response: AgentResponse::new(&apology),
6041 ownership: None,
6042 });
6043 }
6044 DisambiguationResult::Abandoned { .. } => {
6045 debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6046 return Ok(SkillRouteResult::NoMatch);
6047 }
6048 }
6049 }
6050 let admission = self
6051 .admit_disambiguation_redispatch(
6052 expected_disambiguation_epoch,
6053 expected_state_generation,
6054 )
6055 .await?;
6056 drop(admission);
6057 let content = self.execute_skill(&skill, input).await?;
6058 Ok(SkillRouteResult::Response { skill_id, content })
6059 }
6060
6061 async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6063 if let Some(candidate) = self.select_skill_candidate(input).await? {
6064 self.commit_skill_candidate_route_result(candidate, input)
6065 .await
6066 } else {
6067 Ok(SkillRouteResult::NoMatch)
6068 }
6069 }
6070
6071 async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6073 if let Some(ref executor) = self.skill_executor {
6074 let skill_reasoning = self.get_skill_reasoning_config(skill);
6075 let skill_reflection = self.get_skill_reflection_config(skill);
6076
6077 debug!(
6078 skill_id = %skill.id,
6079 reasoning_mode = ?skill_reasoning.mode,
6080 reflection_enabled = ?skill_reflection.enabled,
6081 "Skill reasoning/reflection config"
6082 );
6083
6084 let response = self
6085 .observe_purpose(
6086 ObservationPurpose::SkillPrompt,
6087 executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6088 )
6089 .await?;
6090
6091 if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6092 let should_reflect = self
6093 .should_reflect_with_config(input, &response, &skill_reflection)
6094 .await?;
6095 if should_reflect {
6096 let evaluated = self
6097 .evaluate_and_retry_with_config(input, response, &skill_reflection)
6098 .await?;
6099 return Ok(evaluated);
6100 }
6101 }
6102
6103 return Ok(response);
6104 }
6105 Err(AgentError::Skill(
6106 "No skill executor configured".to_string(),
6107 ))
6108 }
6109
6110 async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6113 let skill = self
6114 .skill_router
6115 .as_ref()
6116 .and_then(|r| r.get_skill(skill_id).cloned())
6117 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6118 self.execute_skill(&skill, input).await
6119 }
6120
6121 async fn should_reflect_with_config(
6122 &self,
6123 input: &str,
6124 response: &str,
6125 config: &ReflectionConfig,
6126 ) -> Result<bool> {
6127 if !config.requires_evaluation() {
6128 return Ok(false);
6129 }
6130
6131 if config.is_enabled() {
6132 return Ok(true);
6133 }
6134
6135 let evaluator_llm = config
6136 .evaluator_llm
6137 .as_ref()
6138 .and_then(|alias| self.llm_registry.get(alias).ok())
6139 .or_else(|| self.llm_registry.router().ok())
6140 .or_else(|| self.llm_registry.default().ok());
6141
6142 let Some(llm) = evaluator_llm else {
6143 return Ok(false);
6144 };
6145
6146 let response_preview: String = response.chars().take(500).collect();
6147 let prompt = format!(
6148 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6149
6150User query: "{}"
6151Response: "{}"
6152
6153Answer YES or NO only."#,
6154 input, response_preview
6155 );
6156
6157 let messages = vec![ChatMessage::user(&prompt)];
6158 let result = self
6159 .observe_purpose(
6160 ObservationPurpose::ReflectionDecision,
6161 llm.complete(&messages, None),
6162 )
6163 .await;
6164
6165 match result {
6166 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6167 Err(_) => Ok(false),
6168 }
6169 }
6170
6171 async fn evaluate_and_retry_with_config(
6172 &self,
6173 input: &str,
6174 mut response: String,
6175 config: &ReflectionConfig,
6176 ) -> Result<String> {
6177 let llm = self.get_state_llm()?;
6178 let mut attempts = 0u32;
6179 let max_retries = config.max_retries;
6180
6181 loop {
6182 let evaluation = self
6183 .evaluate_response_with_config(input, &response, config)
6184 .await?;
6185
6186 if evaluation.passed || attempts >= max_retries {
6187 info!(
6188 passed = evaluation.passed,
6189 confidence = evaluation.confidence,
6190 attempts = attempts + 1,
6191 "Skill reflection evaluation complete"
6192 );
6193 return Ok(response);
6194 }
6195
6196 debug!(
6197 attempt = attempts + 1,
6198 failed_criteria = evaluation.failed_criteria().count(),
6199 "Skill response did not meet criteria, retrying"
6200 );
6201
6202 let feedback: Vec<String> = evaluation
6203 .failed_criteria()
6204 .map(|c| format!("- {}", c.criterion))
6205 .collect();
6206
6207 let retry_prompt = format!(
6208 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6209 feedback.join("\n"),
6210 input
6211 );
6212
6213 let messages = vec![ChatMessage::user(&retry_prompt)];
6214 let retry_response = self
6215 .observe_purpose(
6216 ObservationPurpose::ReflectionEvaluation,
6217 llm.complete(&messages, None),
6218 )
6219 .await
6220 .map_err(|e| AgentError::LLM(e.to_string()))?;
6221
6222 response = retry_response.content.trim().to_string();
6223 attempts += 1;
6224 }
6225 }
6226
6227 async fn evaluate_response_with_config(
6228 &self,
6229 input: &str,
6230 response: &str,
6231 config: &ReflectionConfig,
6232 ) -> Result<EvaluationResult> {
6233 let evaluator_llm = config
6234 .evaluator_llm
6235 .as_ref()
6236 .and_then(|alias| self.llm_registry.get(alias).ok())
6237 .or_else(|| self.llm_registry.router().ok())
6238 .or_else(|| self.llm_registry.default().ok())
6239 .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6240
6241 let criteria = &config.criteria;
6242 let criteria_list = criteria
6243 .iter()
6244 .enumerate()
6245 .map(|(i, c)| format!("{}. {}", i + 1, c))
6246 .collect::<Vec<_>>()
6247 .join("\n");
6248
6249 let prompt = format!(
6250 r#"Evaluate this response against the criteria.
6251
6252User query: "{}"
6253
6254Response to evaluate: "{}"
6255
6256Criteria:
6257{}
6258
6259For each criterion, respond with:
6260- criterion number
6261- PASS or FAIL
6262- brief reason
6263
6264Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6265
6266Format:
62671. PASS/FAIL - reason
62682. PASS/FAIL - reason
6269...
6270CONFIDENCE: 0.X
6271OVERALL: PASS/FAIL"#,
6272 input, response, criteria_list
6273 );
6274
6275 let messages = vec![ChatMessage::user(&prompt)];
6276 let eval_response = self
6277 .observe_purpose(
6278 ObservationPurpose::ReflectionEvaluation,
6279 evaluator_llm.complete(&messages, None),
6280 )
6281 .await
6282 .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6283
6284 let content = eval_response.content.to_uppercase();
6285 let llm_pass = content.contains("OVERALL: PASS");
6286
6287 let confidence = content
6288 .lines()
6289 .find(|l| l.contains("CONFIDENCE:"))
6290 .and_then(|l| {
6291 l.split(':')
6292 .nth(1)
6293 .and_then(|v| v.trim().parse::<f32>().ok())
6294 })
6295 .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6296
6297 let overall_pass = llm_pass && confidence >= config.pass_threshold;
6300
6301 let mut criteria_results = Vec::new();
6302 for (i, criterion) in criteria.iter().enumerate() {
6303 let line_marker = format!("{}.", i + 1);
6304 let passed = eval_response
6305 .content
6306 .lines()
6307 .find(|l| l.contains(&line_marker))
6308 .map(|l| l.to_uppercase().contains("PASS"))
6309 .unwrap_or(overall_pass);
6310
6311 if passed {
6312 criteria_results.push(CriterionResult::pass(criterion));
6313 } else {
6314 criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6315 }
6316 }
6317
6318 Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6319 }
6320
6321 async fn process_input(&self, input: &str) -> Result<ProcessData> {
6323 if let Some(processor) = self.get_state_process_processor() {
6324 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6325 return self
6326 .observe_purpose(purpose, processor.process_input(input))
6327 .await;
6328 }
6329 if let Some(ref processor) = self.process_processor {
6330 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6331 self.observe_purpose(purpose, processor.process_input(input))
6332 .await
6333 } else {
6334 Ok(ProcessData::new(input))
6335 }
6336 }
6337
6338 async fn process_output(
6340 &self,
6341 output: &str,
6342 input_context: &std::collections::HashMap<String, serde_json::Value>,
6343 ) -> Result<ProcessData> {
6344 if let Some(processor) = self.get_state_process_processor() {
6345 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6346 return self
6347 .observe_purpose(purpose, processor.process_output(output, input_context))
6348 .await;
6349 }
6350 if let Some(ref processor) = self.process_processor {
6351 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6352 self.observe_purpose(purpose, processor.process_output(output, input_context))
6353 .await
6354 } else {
6355 Ok(ProcessData::new(output))
6356 }
6357 }
6358
6359 fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6361 let sm = self.state_machine.as_ref()?;
6362 let def = sm.current_definition()?;
6363 let config = def.process.as_ref()?;
6364 let mut processor = ProcessProcessor::new(config.clone());
6365 if let Some(ref registry) = Some(self.llm_registry.clone()) {
6366 processor = processor.with_llm_registry(registry.clone());
6367 }
6368 processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6369 Some(processor)
6370 }
6371
6372 async fn check_turn_timeout(&self) -> Result<()> {
6374 let Some(ref sm) = self.state_machine else {
6375 return Ok(());
6376 };
6377 let Some(timeout_state) = sm.check_timeout() else {
6378 return Ok(());
6379 };
6380 let claim_admission = self.disambiguation_admission.write().await;
6381 if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6382 return Ok(());
6383 }
6384 let Some(reservation) = self.reserve_state_transition() else {
6385 return Ok(());
6386 };
6387 let from_state = sm.current();
6388 let expected_state_generation = sm.generation();
6389 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6390 let history_before = sm.history();
6391 drop(claim_admission);
6392
6393 self.execute_state_exit_actions(&from_state).await;
6394
6395 let admission = self.disambiguation_admission.write().await;
6396 if sm.current() != from_state
6397 || sm.generation() != expected_state_generation
6398 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6399 || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6400 {
6401 return Ok(());
6402 }
6403 sm.transition_to(&timeout_state, "max_turns exceeded")?;
6404 self.invalidate_pending_confirmation("state_timeout").await;
6405 let entered = sm.current();
6406 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6407 drop(admission);
6408
6409 self.execute_state_enter_actions(&entered, is_reentry).await;
6410 drop(reservation);
6411 info!(to = %entered, "Timeout transition");
6412 Ok(())
6413 }
6414
6415 fn increment_turn(&self) {
6416 if let Some(ref sm) = self.state_machine {
6417 sm.increment_turn();
6418 }
6419 }
6420
6421 fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6422 let sm = self.state_machine.as_ref()?;
6423 let current = sm.current();
6424 let transitions: Vec<_> = sm
6425 .auto_transitions()
6426 .into_iter()
6427 .filter(|t| match t.cooldown_turns {
6428 Some(cd) if cd > 0 => {
6429 let resolved = sm.config().resolve_full_path(¤t, &t.to);
6430 !sm.is_on_cooldown(&resolved, cd)
6431 }
6432 _ => true,
6433 })
6434 .collect();
6435 Some((transitions, current))
6436 }
6437
6438 fn transition_reason(transition: &Transition) -> String {
6439 if transition.when.is_empty() {
6440 "guard condition met".to_string()
6441 } else {
6442 transition.when.clone()
6443 }
6444 }
6445
6446 fn build_transition_context(
6448 &self,
6449 user_message: &str,
6450 response: &str,
6451 current_state: &str,
6452 staged: Option<&HashMap<String, Value>>,
6453 ) -> TransitionContext {
6454 let context_map = staged
6455 .map(|writes| self.build_context_with_staged(writes))
6456 .unwrap_or_else(|| self.build_context_with_overlays());
6457 TransitionContext::new(user_message, response, current_state).with_context(context_map)
6458 }
6459
6460 async fn select_transition_candidate(
6462 &self,
6463 user_message: &str,
6464 response: &str,
6465 ) -> Result<Option<TransitionCandidate>> {
6466 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6467 return Ok(None);
6468 };
6469 let transitions: Vec<Transition> = transitions
6470 .into_iter()
6471 .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6472 .collect();
6473 if transitions.is_empty() {
6474 return Ok(None);
6475 }
6476 let Some(evaluator) = self.transition_evaluator.as_ref() else {
6477 return Ok(None);
6478 };
6479 let context = self.build_transition_context(user_message, response, ¤t_state, None);
6480 let selected = self
6481 .observe_purpose(
6482 ObservationPurpose::StateTransitionEvaluation,
6483 evaluator.select_transition(&transitions, &context),
6484 )
6485 .await?;
6486 Ok(selected.map(|index| {
6487 let transition = transitions[index].clone();
6488 TransitionCandidate::new(
6489 current_state,
6490 transition.clone(),
6491 Self::transition_reason(&transition),
6492 )
6493 }))
6494 }
6495
6496 fn select_deterministic_transition_candidate(
6498 &self,
6499 user_message: &str,
6500 current_state: &str,
6501 transitions: &[Transition],
6502 staged: &HashMap<String, Value>,
6503 ) -> Option<TransitionCandidate> {
6504 let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6505
6506 for transition in transitions {
6507 if let Some(guard) = transition.guard.as_ref()
6508 && evaluate_guard(guard, &context)
6509 {
6510 return Some(TransitionCandidate::new(
6511 current_state,
6512 transition.clone(),
6513 Self::transition_reason(transition),
6514 ));
6515 }
6516 }
6517
6518 let resolved_intent = context
6519 .context
6520 .get("resolved_intent")
6521 .and_then(Value::as_str)
6522 .filter(|value| !value.is_empty());
6523 if let Some(resolved_intent) = resolved_intent {
6524 for transition in transitions {
6525 if transition.intent.as_deref() == Some(resolved_intent) {
6526 return Some(TransitionCandidate::new(
6527 current_state,
6528 transition.clone(),
6529 Self::transition_reason(transition),
6530 ));
6531 }
6532 }
6533 }
6534
6535 None
6536 }
6537
6538 async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6540 self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6541 .await
6542 }
6543
6544 async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6546 let approved = self.check_state_hitl(Some(from_state), target).await?;
6547 if !approved {
6548 info!(to = %target, "State transition rejected by HITL");
6549 }
6550 Ok(approved)
6551 }
6552
6553 async fn apply_transition_target(
6555 &self,
6556 from_state: &str,
6557 target: &str,
6558 reason: &str,
6559 staged: Option<&HashMap<String, Value>>,
6560 ) -> Result<bool> {
6561 let Some(ref sm) = self.state_machine else {
6562 return Ok(false);
6563 };
6564 let claim_admission = self.disambiguation_admission.write().await;
6565 if sm.current() != from_state {
6566 return Ok(false);
6567 }
6568 let Some(reservation) = self.reserve_state_transition() else {
6569 return Ok(false);
6570 };
6571 let expected_state_generation = sm.generation();
6572 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6573 let history_before = sm.history();
6574 drop(claim_admission);
6575
6576 self.execute_state_exit_actions(from_state).await;
6577
6578 let admission = self.disambiguation_admission.write().await;
6579 if sm.current() != from_state
6580 || sm.generation() != expected_state_generation
6581 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6582 {
6583 return Ok(false);
6584 }
6585 sm.transition_to(target, reason)?;
6586 self.invalidate_pending_confirmation("state_transition")
6587 .await;
6588 sm.reset_no_transition();
6589 if let Some(staged) = staged {
6590 self.commit_staged_context_writes(staged);
6591 }
6592 let entered = sm.current();
6593 let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6594 drop(admission);
6595
6596 self.execute_state_enter_actions(&entered, is_reentry).await;
6597 drop(reservation);
6598 self.hooks
6599 .on_state_transition(Some(from_state), &entered, reason)
6600 .await;
6601 info!(from = %from_state, to = %entered, "State transition");
6602 Ok(true)
6603 }
6604
6605 async fn commit_transition_target(
6607 &self,
6608 from_state: &str,
6609 target: &str,
6610 reason: &str,
6611 ) -> Result<bool> {
6612 if !self.approve_transition_target(from_state, target).await? {
6613 return Ok(false);
6614 }
6615 self.apply_transition_target(from_state, target, reason, None)
6616 .await
6617 }
6618
6619 async fn apply_pre_response_transition_candidate(
6621 &self,
6622 candidate: &TransitionCandidate,
6623 staged: &HashMap<String, Value>,
6624 processed_input: &str,
6625 ) -> Result<bool> {
6626 self.commit_root_user_message(processed_input).await?;
6627 self.apply_transition_target(
6628 &candidate.from_state,
6629 candidate.target(),
6630 &candidate.reason,
6631 Some(staged),
6632 )
6633 .await
6634 }
6635
6636 async fn commit_pre_response_transition_candidate(
6638 &self,
6639 candidate: &TransitionCandidate,
6640 staged: &HashMap<String, Value>,
6641 processed_input: &str,
6642 ) -> Result<bool> {
6643 if !self
6644 .approve_transition_target(&candidate.from_state, candidate.target())
6645 .await?
6646 {
6647 return Ok(false);
6648 }
6649 self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6650 .await
6651 }
6652
6653 async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6655 let Some(ref sm) = self.state_machine else {
6656 return Ok(false);
6657 };
6658 sm.increment_no_transition();
6659 let Some(fallback) = sm.check_fallback() else {
6660 return Ok(false);
6661 };
6662 self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6663 .await
6664 }
6665
6666 async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6668 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6669 return Ok(false);
6670 };
6671 if transitions.is_empty() {
6672 return Ok(false);
6673 }
6674 if let Some(candidate) = self
6675 .select_transition_candidate(user_message, response)
6676 .await?
6677 {
6678 return self.commit_transition_candidate(&candidate).await;
6679 }
6680 self.handle_transition_miss(¤t_state).await
6681 }
6682
6683 async fn try_pre_response_transition(
6685 &self,
6686 processed_input: &str,
6687 ) -> Result<Option<AgentResponse>> {
6688 let optimization = &self.runtime_config.optimization;
6689 if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6690 return Ok(None);
6691 }
6692 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6693 return Ok(None);
6694 };
6695 let eligible: Vec<Transition> = transitions
6696 .into_iter()
6697 .filter(|transition| !transition.requires_response)
6698 .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
6699 .collect();
6700 if eligible.is_empty() {
6701 return Ok(None);
6702 }
6703
6704 let empty_staged = HashMap::new();
6705 let mut extracted_staged: Option<HashMap<String, Value>> = None;
6706 let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
6707
6708 for transition in &eligible {
6709 let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
6710 let staged_for_eval = if use_extractors {
6711 if extracted_staged.is_none() {
6712 extracted_staged =
6713 Some(self.run_context_extractors_staged(processed_input).await);
6714 }
6715 extracted_staged.as_ref().unwrap_or(&empty_staged)
6716 } else {
6717 &empty_staged
6718 };
6719
6720 if let Some(candidate) = self.select_deterministic_transition_candidate(
6721 processed_input,
6722 ¤t_state,
6723 std::slice::from_ref(transition),
6724 staged_for_eval,
6725 ) {
6726 let staged_for_commit = if use_extractors {
6727 staged_for_eval.clone()
6728 } else {
6729 HashMap::new()
6730 };
6731 selected = Some((candidate, staged_for_commit));
6732 break;
6733 }
6734 }
6735
6736 let Some((candidate, staged)) = selected else {
6737 return Ok(None);
6738 };
6739
6740 if !self
6741 .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
6742 .await?
6743 {
6744 return Ok(None);
6745 }
6746 self.redispatch_current_state(processed_input)
6747 .await
6748 .map(Some)
6749 }
6750
6751 async fn try_speculative_branches(
6756 &self,
6757 processed_input: &str,
6758 input_context: &HashMap<String, Value>,
6759 ) -> Result<Option<AgentResponse>> {
6760 let optimization = &self.runtime_config.optimization;
6761 if !optimization.enabled {
6762 return Ok(None);
6763 }
6764
6765 let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
6766 if !matches!(
6767 effective_reasoning_mode,
6768 ReasoningMode::None | ReasoningMode::Auto
6769 ) {
6770 return Ok(None);
6771 }
6772
6773 let mut transition_enabled =
6774 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
6775 let mut skill_enabled = optimization.speculative_skill_routing
6776 && self.skill_router.is_some()
6777 && self.pending_skill_id.read().is_none();
6778 let mut reasoning_enabled = optimization.speculative_reasoning_auto
6779 && matches!(effective_reasoning_mode, ReasoningMode::Auto);
6780
6781 if matches!(effective_reasoning_mode, ReasoningMode::Auto)
6782 && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
6783 {
6784 return Ok(None);
6785 }
6786
6787 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6788 return Ok(None);
6789 }
6790
6791 let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
6792 let mut speculative_call_slots = optimization
6793 .max_speculative_llm_calls_per_turn
6794 .saturating_sub(1);
6795 if reasoning_enabled {
6796 if optional_slots == 0 || speculative_call_slots == 0 {
6797 return Ok(None);
6798 }
6799 optional_slots -= 1;
6800 speculative_call_slots -= 1;
6801 }
6802 if transition_enabled {
6803 if optional_slots == 0 {
6804 transition_enabled = false;
6805 } else {
6806 optional_slots -= 1;
6807 }
6808 }
6809 if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
6810 skill_enabled = false;
6811 }
6812
6813 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6814 return Ok(None);
6815 }
6816
6817 let main_kind = if transition_enabled {
6818 RuntimeOptimizationKind::ParallelStateTransition
6819 } else if skill_enabled {
6820 RuntimeOptimizationKind::SpeculativeSkillRouting
6821 } else {
6822 RuntimeOptimizationKind::SpeculativeReasoningAuto
6823 };
6824 if !self.reserve_active_speculative_llm_call(main_kind) {
6825 return Ok(None);
6826 }
6827
6828 let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
6829 let main_branch = RuntimeBranch::new(
6830 RuntimeTaskPurpose::MainResponse,
6831 main_kind,
6832 RuntimeTaskPriority::Normal,
6833 RuntimeCommitBehavior::FinalResponse,
6834 );
6835 let transition_branch = RuntimeBranch::new(
6836 RuntimeTaskPurpose::StateTransition,
6837 RuntimeOptimizationKind::ParallelStateTransition,
6838 RuntimeTaskPriority::Critical,
6839 RuntimeCommitBehavior::TransitionDecision,
6840 );
6841 let skill_branch = RuntimeBranch::new(
6842 RuntimeTaskPurpose::SkillRouting,
6843 RuntimeOptimizationKind::SpeculativeSkillRouting,
6844 RuntimeTaskPriority::High,
6845 RuntimeCommitBehavior::SkillSelection,
6846 );
6847 let reasoning_branch = RuntimeBranch::new(
6848 RuntimeTaskPurpose::ReasoningJudge,
6849 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6850 RuntimeTaskPriority::Normal,
6851 RuntimeCommitBehavior::ReasoningDecision,
6852 );
6853 let main_id = main_branch.branch_id();
6854 let transition_id = transition_branch.branch_id();
6855 let skill_id = skill_branch.branch_id();
6856 let reasoning_id = reasoning_branch.branch_id();
6857
6858 let main_id_for_future = main_id.clone();
6859 if !branch_set.schedule(
6860 main_branch,
6861 Box::pin(async move {
6862 match crate::optimization::observability::with_branch_observation(
6863 &main_id_for_future,
6864 main_kind,
6865 RuntimeCommitBehavior::FinalResponse,
6866 self.generate_main_response_draft(processed_input, &ReasoningMode::None),
6867 )
6868 .await
6869 {
6870 Ok(draft) => RuntimeBranchResult::MainDraft(draft),
6871 Err(error) => RuntimeBranchResult::Failed(error),
6872 }
6873 }),
6874 ) {
6875 return Ok(None);
6876 }
6877
6878 if transition_enabled {
6879 let transition_id_for_future = transition_id.clone();
6880 if !branch_set.schedule(
6881 transition_branch,
6882 Box::pin(async move {
6883 match crate::optimization::observability::with_branch_observation(
6884 &transition_id_for_future,
6885 RuntimeOptimizationKind::ParallelStateTransition,
6886 RuntimeCommitBehavior::TransitionDecision,
6887 self.select_parallel_transition_candidate(processed_input),
6888 )
6889 .await
6890 {
6891 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
6892 RuntimeBranchResult::Transition(Some(candidate))
6893 }
6894 Ok(ParallelTransitionSelection::NoMatch) => {
6895 RuntimeBranchResult::Transition(None)
6896 }
6897 Ok(ParallelTransitionSelection::ReservationExhausted) => {
6898 RuntimeBranchResult::Cancelled
6899 }
6900 Err(error) => RuntimeBranchResult::Failed(error),
6901 }
6902 }),
6903 ) {
6904 transition_enabled = false;
6905 }
6906 }
6907
6908 if skill_enabled {
6909 let skill_id_for_future = skill_id.clone();
6910 if !branch_set.schedule(
6911 skill_branch,
6912 Box::pin(async move {
6913 if !self.reserve_active_speculative_llm_call(
6914 RuntimeOptimizationKind::SpeculativeSkillRouting,
6915 ) {
6916 return RuntimeBranchResult::Cancelled;
6917 }
6918 match crate::optimization::observability::with_branch_observation(
6919 &skill_id_for_future,
6920 RuntimeOptimizationKind::SpeculativeSkillRouting,
6921 RuntimeCommitBehavior::SkillSelection,
6922 self.select_skill_candidate(processed_input),
6923 )
6924 .await
6925 {
6926 Ok(candidate) => RuntimeBranchResult::Skill(candidate),
6927 Err(error) => RuntimeBranchResult::Failed(error),
6928 }
6929 }),
6930 ) {
6931 skill_enabled = false;
6932 }
6933 }
6934
6935 if reasoning_enabled {
6936 let reasoning_id_for_future = reasoning_id.clone();
6937 if !branch_set.schedule(
6938 reasoning_branch,
6939 Box::pin(async move {
6940 if !self.reserve_active_speculative_llm_call(
6941 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6942 ) {
6943 return RuntimeBranchResult::Cancelled;
6944 }
6945 match crate::optimization::observability::with_branch_observation(
6946 &reasoning_id_for_future,
6947 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6948 RuntimeCommitBehavior::ReasoningDecision,
6949 self.determine_reasoning_mode_strict(processed_input),
6950 )
6951 .await
6952 {
6953 Ok(mode) => RuntimeBranchResult::Reasoning(mode),
6954 Err(error) => RuntimeBranchResult::Failed(error),
6955 }
6956 }),
6957 ) {
6958 reasoning_enabled = false;
6959 }
6960 }
6961
6962 if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
6963 self.finalize_pending_branches(branch_set.cancel_pending());
6964 return Ok(None);
6965 }
6966
6967 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6968 self.finalize_pending_branches(branch_set.cancel_pending());
6969 return Ok(None);
6970 }
6971
6972 let mut main_pending = true;
6973 let mut skill_pending = skill_enabled;
6974 let mut reasoning_pending = reasoning_enabled;
6975 let mut transition_finalized = !transition_enabled;
6976 let mut skill_finalized = !skill_enabled;
6977 let mut reasoning_finalized = !reasoning_enabled;
6978 let mut main_result: Option<Result<MainResponseDraft>> = None;
6979 let mut transition_candidate: Option<TransitionCandidate> = None;
6980 let mut skill_candidate: Option<SkillCandidate> = None;
6981 let mut reasoning_decision: Option<ReasoningMode> = None;
6982 let mut transition_fallback_required = false;
6983 let mut skill_fallback_required = false;
6984 let mut reasoning_fallback_required = false;
6985
6986 loop {
6987 if let Some(candidate) = transition_candidate.take() {
6988 if self
6989 .approve_transition_target(&candidate.from_state, candidate.target())
6990 .await?
6991 {
6992 self.finalize_pending_branches(branch_set.cancel_pending());
6994 if !main_pending {
6995 self.finalize_branch_loss(
6996 &main_id,
6997 main_kind,
6998 RuntimeCommitBehavior::FinalResponse,
6999 false,
7000 main_result.as_ref().map(|result| result.is_err()),
7001 );
7002 }
7003 if skill_enabled && !skill_pending {
7004 self.finalize_branch_loss(
7005 &skill_id,
7006 RuntimeOptimizationKind::SpeculativeSkillRouting,
7007 RuntimeCommitBehavior::SkillSelection,
7008 false,
7009 Some(false),
7010 );
7011 }
7012 if reasoning_enabled && !reasoning_pending {
7013 self.finalize_branch_loss(
7014 &reasoning_id,
7015 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7016 RuntimeCommitBehavior::ReasoningDecision,
7017 false,
7018 Some(false),
7019 );
7020 }
7021 if !self
7022 .apply_pre_response_transition_candidate(
7023 &candidate,
7024 &HashMap::new(),
7025 processed_input,
7026 )
7027 .await?
7028 {
7029 self.finalize_optional_branch(
7030 &transition_id,
7031 RuntimeOptimizationKind::ParallelStateTransition,
7032 RuntimeCommitBehavior::TransitionDecision,
7033 "discarded",
7034 false,
7035 );
7036 return Ok(None);
7037 }
7038 self.finalize_optional_branch(
7039 &transition_id,
7040 RuntimeOptimizationKind::ParallelStateTransition,
7041 RuntimeCommitBehavior::TransitionDecision,
7042 "committed",
7043 true,
7044 );
7045 return self
7046 .redispatch_current_state(processed_input)
7047 .await
7048 .map(Some);
7049 }
7050 self.finalize_optional_branch(
7051 &transition_id,
7052 RuntimeOptimizationKind::ParallelStateTransition,
7053 RuntimeCommitBehavior::TransitionDecision,
7054 "discarded",
7055 false,
7056 );
7057 transition_finalized = true;
7058 }
7059
7060 if transition_finalized && skill_candidate.is_some() {
7061 let candidate = skill_candidate.take().unwrap();
7062 self.finalize_optional_branch(
7063 &skill_id,
7064 RuntimeOptimizationKind::SpeculativeSkillRouting,
7065 RuntimeCommitBehavior::SkillSelection,
7066 "committed",
7067 true,
7068 );
7069 if !main_pending {
7070 self.finalize_branch_loss(
7071 &main_id,
7072 main_kind,
7073 RuntimeCommitBehavior::FinalResponse,
7074 false,
7075 main_result.as_ref().map(|result| result.is_err()),
7076 );
7077 }
7078 if reasoning_enabled && !reasoning_pending {
7079 self.finalize_branch_loss(
7080 &reasoning_id,
7081 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7082 RuntimeCommitBehavior::ReasoningDecision,
7083 false,
7084 Some(false),
7085 );
7086 }
7087 self.finalize_pending_branches(branch_set.cancel_pending());
7088 self.commit_root_user_message(processed_input).await?;
7089 return match self
7090 .commit_skill_candidate_route_result(candidate, processed_input)
7091 .await?
7092 {
7093 SkillRouteResult::Response { skill_id, content } => self
7094 .handle_skill_response(processed_input, &skill_id, content, input_context)
7095 .await
7096 .map(Some),
7097 SkillRouteResult::NeedsClarification {
7098 response,
7099 ownership,
7100 } => {
7101 let admission = self
7102 .admit_optional_disambiguation_ownership(ownership)
7103 .await?;
7104 if response
7105 .metadata
7106 .as_ref()
7107 .and_then(|m| m.get("disambiguation"))
7108 .and_then(|d| d.get("status"))
7109 .and_then(|s| s.as_str())
7110 == Some("awaiting_clarification")
7111 {
7112 self.memory
7113 .add_message(ChatMessage::assistant(&response.content))
7114 .await?;
7115 }
7116 drop(admission);
7117 self.finish_turn_if_root(&response).await?;
7118 Ok(Some(response))
7119 }
7120 SkillRouteResult::NoMatch => Ok(None),
7121 };
7122 }
7123
7124 if transition_finalized
7125 && skill_finalized
7126 && let Some(reasoning_mode) = reasoning_decision.take()
7127 {
7128 if !matches!(reasoning_mode, ReasoningMode::None) {
7129 self.finalize_optional_branch(
7130 &reasoning_id,
7131 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7132 RuntimeCommitBehavior::ReasoningDecision,
7133 "committed",
7134 true,
7135 );
7136 if !main_pending {
7137 self.finalize_branch_loss(
7138 &main_id,
7139 main_kind,
7140 RuntimeCommitBehavior::FinalResponse,
7141 false,
7142 main_result.as_ref().map(|result| result.is_err()),
7143 );
7144 }
7145 self.finalize_pending_branches(branch_set.cancel_pending());
7146 self.commit_root_user_message(processed_input).await?;
7147 return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7148 self.handle_plan_and_execute(processed_input, input_context, true)
7149 .await
7150 .map(Some)
7151 } else {
7152 self.run_committed_response_loop_with_reasoning(
7153 processed_input,
7154 input_context,
7155 reasoning_mode,
7156 true,
7157 )
7158 .await
7159 .map(Some)
7160 };
7161 }
7162 self.finalize_optional_branch(
7163 &reasoning_id,
7164 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7165 RuntimeCommitBehavior::ReasoningDecision,
7166 "committed",
7167 true,
7168 );
7169 reasoning_finalized = true;
7170 }
7171
7172 if transition_finalized && skill_finalized && reasoning_finalized {
7173 if transition_fallback_required
7174 || skill_fallback_required
7175 || reasoning_fallback_required
7176 {
7177 if !main_pending {
7178 self.finalize_branch_loss(
7179 &main_id,
7180 main_kind,
7181 RuntimeCommitBehavior::FinalResponse,
7182 false,
7183 main_result.as_ref().map(|result| result.is_err()),
7184 );
7185 }
7186 self.finalize_pending_branches(branch_set.cancel_pending());
7187 return Ok(None);
7188 }
7189
7190 if let Some(result) = main_result.take() {
7191 let draft = match result {
7192 Ok(draft) => draft,
7193 Err(error) => {
7194 self.finalize_optional_branch(
7195 &main_id,
7196 main_kind,
7197 RuntimeCommitBehavior::FinalResponse,
7198 "failed",
7199 false,
7200 );
7201 self.finalize_pending_branches(branch_set.cancel_pending());
7202 return Err(error);
7203 }
7204 };
7205 self.finalize_optional_branch(
7206 &main_id,
7207 main_kind,
7208 RuntimeCommitBehavior::FinalResponse,
7209 "committed",
7210 true,
7211 );
7212 self.finalize_pending_branches(branch_set.cancel_pending());
7213 return self
7214 .commit_main_response_draft(
7215 processed_input,
7216 input_context,
7217 draft,
7218 ReasoningMode::None,
7219 reasoning_enabled,
7220 )
7221 .await
7222 .map(Some);
7223 }
7224 }
7225
7226 if branch_set.is_empty() {
7227 return Ok(None);
7228 }
7229
7230 let Some(outcome) = branch_set.next_completed().await else {
7231 return Ok(None);
7232 };
7233 let branch_id = outcome.branch.branch_id();
7234 match outcome.result {
7235 RuntimeBranchResult::MainDraft(draft) => {
7236 main_pending = false;
7237 main_result = Some(Ok(draft));
7238 }
7239 RuntimeBranchResult::Transition(candidate) => {
7240 if let Some(candidate) = candidate {
7241 transition_candidate = Some(candidate);
7242 } else {
7243 self.finalize_optional_branch(
7244 &transition_id,
7245 RuntimeOptimizationKind::ParallelStateTransition,
7246 RuntimeCommitBehavior::TransitionDecision,
7247 "discarded",
7248 false,
7249 );
7250 transition_finalized = true;
7251 }
7252 }
7253 RuntimeBranchResult::Skill(candidate) => {
7254 skill_pending = false;
7255 if let Some(candidate) = candidate {
7256 skill_candidate = Some(candidate);
7257 } else {
7258 self.finalize_optional_branch(
7259 &skill_id,
7260 RuntimeOptimizationKind::SpeculativeSkillRouting,
7261 RuntimeCommitBehavior::SkillSelection,
7262 "discarded",
7263 false,
7264 );
7265 skill_finalized = true;
7266 }
7267 }
7268 RuntimeBranchResult::Reasoning(mode) => {
7269 reasoning_pending = false;
7270 reasoning_decision = Some(mode);
7271 }
7272 RuntimeBranchResult::Failed(error) => {
7273 if branch_id == main_id {
7274 main_pending = false;
7275 main_result = Some(Err(error));
7276 } else if branch_id == transition_id {
7277 self.finalize_optional_branch(
7278 &transition_id,
7279 RuntimeOptimizationKind::ParallelStateTransition,
7280 RuntimeCommitBehavior::TransitionDecision,
7281 "failed",
7282 false,
7283 );
7284 transition_finalized = true;
7285 } else if branch_id == skill_id {
7286 skill_pending = false;
7287 self.finalize_optional_branch(
7288 &skill_id,
7289 RuntimeOptimizationKind::SpeculativeSkillRouting,
7290 RuntimeCommitBehavior::SkillSelection,
7291 "failed",
7292 false,
7293 );
7294 skill_finalized = true;
7295 } else if branch_id == reasoning_id {
7296 reasoning_pending = false;
7297 self.finalize_optional_branch(
7298 &reasoning_id,
7299 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7300 RuntimeCommitBehavior::ReasoningDecision,
7301 "failed",
7302 false,
7303 );
7304 reasoning_finalized = true;
7305 }
7306 }
7307 RuntimeBranchResult::Cancelled => {
7308 self.finalize_optional_branch(
7309 &branch_id,
7310 outcome.branch.optimization,
7311 outcome.branch.commit_behavior,
7312 "cancelled",
7313 false,
7314 );
7315 if branch_id == main_id {
7316 main_pending = false;
7317 main_result =
7318 Some(Err(AgentError::Other("main branch cancelled".to_string())));
7319 } else if branch_id == transition_id {
7320 transition_finalized = true;
7321 transition_fallback_required = true;
7322 } else if branch_id == skill_id {
7323 skill_pending = false;
7324 skill_finalized = true;
7325 skill_fallback_required = true;
7326 } else if branch_id == reasoning_id {
7327 reasoning_pending = false;
7328 reasoning_finalized = true;
7329 reasoning_fallback_required = true;
7330 }
7331 }
7332 }
7333 }
7334 }
7335
7336 fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7337 for branch in branches {
7338 self.finalize_optional_branch(
7339 &branch.branch_id(),
7340 branch.optimization,
7341 branch.commit_behavior,
7342 "cancelled",
7343 false,
7344 );
7345 }
7346 }
7347
7348 fn finalize_branch_loss(
7353 &self,
7354 branch_id: &str,
7355 optimization: RuntimeOptimizationKind,
7356 commit_behavior: RuntimeCommitBehavior,
7357 pending: bool,
7358 completed_failed: Option<bool>,
7359 ) {
7360 let status = if pending {
7361 "cancelled"
7362 } else if completed_failed.unwrap_or(false) {
7363 "failed"
7364 } else {
7365 "discarded"
7366 };
7367 self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7368 }
7369
7370 fn finalize_optional_branch(
7375 &self,
7376 branch_id: &str,
7377 optimization: RuntimeOptimizationKind,
7378 commit_behavior: RuntimeCommitBehavior,
7379 status: &str,
7380 winner: bool,
7381 ) {
7382 crate::optimization::observability::finalize_branch(
7383 self.observability_manager.as_ref(),
7384 branch_id,
7385 status,
7386 winner,
7387 optimization,
7388 commit_behavior,
7389 );
7390 }
7391
7392 fn has_parallel_transition_candidates(&self) -> bool {
7397 self.transitions_available_for_commit()
7398 .map(|(transitions, _)| {
7399 transitions
7400 .iter()
7401 .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7402 })
7403 .unwrap_or(false)
7404 }
7405
7406 async fn select_parallel_transition_candidate(
7411 &self,
7412 processed_input: &str,
7413 ) -> Result<ParallelTransitionSelection> {
7414 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7415 return Ok(ParallelTransitionSelection::NoMatch);
7416 };
7417 let parallel: Vec<Transition> = transitions
7418 .into_iter()
7419 .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7420 .filter(|transition| !transition.requires_response)
7421 .collect();
7422 if parallel.is_empty() {
7423 return Ok(ParallelTransitionSelection::NoMatch);
7424 }
7425 let empty_staged = HashMap::new();
7426 if let Some(candidate) = self.select_deterministic_transition_candidate(
7427 processed_input,
7428 ¤t_state,
7429 ¶llel,
7430 &empty_staged,
7431 ) {
7432 return Ok(ParallelTransitionSelection::Candidate(candidate));
7433 }
7434 let when_transitions: Vec<(usize, &Transition)> = parallel
7435 .iter()
7436 .enumerate()
7437 .filter(|(_, transition)| !transition.when.trim().is_empty())
7438 .collect();
7439 if when_transitions.is_empty() {
7440 return Ok(ParallelTransitionSelection::NoMatch);
7441 }
7442 let llm = self
7443 .llm_registry
7444 .router()
7445 .or_else(|_| self.llm_registry.default())
7446 .map_err(|e| AgentError::Config(e.to_string()))?;
7447 let conditions = when_transitions
7448 .iter()
7449 .enumerate()
7450 .map(|(display_idx, (_, transition))| {
7451 format!("{}. {}", display_idx + 1, transition.when)
7452 })
7453 .collect::<Vec<_>>()
7454 .join("\n");
7455 if !self
7456 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7457 {
7458 return Ok(ParallelTransitionSelection::ReservationExhausted);
7459 }
7460 let context_preview = self.branch_context_preview();
7461 let prompt = format!(
7462 "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-{}).",
7463 current_state,
7464 processed_input,
7465 context_preview,
7466 conditions,
7467 when_transitions.len()
7468 );
7469 let response = self
7470 .observe_purpose(
7471 ObservationPurpose::StateTransitionEvaluation,
7472 llm.complete(&[ChatMessage::user(prompt)], None),
7473 )
7474 .await
7475 .map_err(|e| AgentError::LLM(e.to_string()))?;
7476 let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7477 if choice == 0 || choice > when_transitions.len() {
7478 return Ok(ParallelTransitionSelection::NoMatch);
7479 }
7480 let transition = when_transitions[choice - 1].1.clone();
7481 Ok(ParallelTransitionSelection::Candidate(
7482 TransitionCandidate::new(
7483 current_state,
7484 transition.clone(),
7485 Self::transition_reason(&transition),
7486 ),
7487 ))
7488 }
7489
7490 async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7492 const MAX_REDISPATCH_DEPTH: u32 = 3;
7493 let current_depth = *self.redispatch_depth.read();
7494 if current_depth >= MAX_REDISPATCH_DEPTH {
7495 warn!(depth = current_depth, "Re-dispatch depth limit reached");
7496 let response = AgentResponse::new("");
7497 self.finish_turn_if_root(&response).await?;
7498 return Ok(response);
7499 }
7500 *self.redispatch_depth.write() += 1;
7501 if let Some(context) = self.active_turn_context.write().as_mut() {
7502 context.enter_redispatch();
7503 }
7504 let result = Box::pin(self.run_loop_internal(processed_input)).await;
7505 *self.redispatch_depth.write() -= 1;
7506 if let Some(context) = self.active_turn_context.write().as_mut() {
7507 context.exit_redispatch();
7508 }
7509 let response = result?;
7510 self.finish_turn_if_root(&response).await?;
7511 Ok(response)
7512 }
7513
7514 async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7516 if *self.redispatch_depth.read() == 0 {
7517 self.post_turn_session_lifecycle().await?;
7518 if let Some(context) = self.active_turn_context.write().as_mut() {
7519 context.mark_post_turn_lifecycle_completed();
7520 }
7521 self.hooks.on_response(response).await;
7522 self.end_root_turn();
7523 }
7524 Ok(())
7525 }
7526
7527 async fn execute_state_exit_actions(&self, state_path: &str) {
7529 if let Some(ref sm) = self.state_machine
7530 && let Some(def) = sm.get_definition(state_path)
7531 && !def.on_exit.is_empty()
7532 {
7533 debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7534 self.execute_state_actions(&def.on_exit).await;
7535 }
7536 }
7537
7538 fn state_was_previously_entered(
7540 state_path: &str,
7541 from_state: &str,
7542 history_before: &[StateTransitionEvent],
7543 ) -> bool {
7544 state_path == from_state
7545 || history_before
7546 .iter()
7547 .any(|event| event.from == state_path || event.to == state_path)
7548 }
7549
7550 async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7552 if let Some(ref sm) = self.state_machine
7553 && let Some(def) = sm.get_definition(state_path)
7554 {
7555 if is_reentry && !def.on_reenter.is_empty() {
7556 debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7557 self.execute_state_actions(&def.on_reenter).await;
7558 } else if !def.on_enter.is_empty() {
7559 debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7560 self.execute_state_actions(&def.on_enter).await;
7561 }
7562 }
7563 }
7564
7565 async fn execute_state_actions(&self, actions: &[StateAction]) {
7567 for (action_index, action) in actions.iter().enumerate() {
7568 match action {
7569 StateAction::Tool { tool, args } => {
7570 let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7571 let args_value = self.render_action_args(&raw_args);
7572 let state = self.state_machine.as_ref().map(|sm| sm.current());
7573 let request = ToolExecutionRequest::new(
7574 uuid::Uuid::new_v4().to_string(),
7575 tool.clone(),
7576 args_value,
7577 ToolCallSource::StateAction {
7578 state,
7579 action_index,
7580 },
7581 );
7582 match self.execute_tool_record(request).await {
7583 Ok(record) if record.success => {
7584 debug!(tool = %record.canonical_id, "State action: tool executed");
7585 let _ = self.context_manager.set(
7586 "last_tool_result",
7587 serde_json::Value::String(record.model_output_string()),
7588 );
7589 let _ = self.context_manager.set(
7590 "last_tool_record",
7591 serde_json::to_value(record).unwrap_or(Value::Null),
7592 );
7593 }
7594 Ok(record) => {
7595 warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7596 }
7597 Err(e) => {
7598 warn!(tool = %tool, error = %e, "State action: tool failed")
7599 }
7600 }
7601 }
7602 StateAction::Skill { skill } => {
7603 if let Some(ref executor) = self.skill_executor {
7604 if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7605 match executor
7606 .execute_with_invoker(def, "", serde_json::json!({}), self)
7607 .await
7608 {
7609 Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7610 Err(e) => {
7611 warn!(skill = %skill, error = %e, "State action: skill failed")
7612 }
7613 }
7614 } else {
7615 warn!(skill = %skill, "State action: skill not found");
7616 }
7617 }
7618 }
7619 StateAction::SetContext { set_context } => {
7620 for (key, value) in set_context {
7621 if let Err(e) = self.context_manager.set(key, value.clone()) {
7622 warn!(key = %key, error = %e, "State action: set_context failed");
7623 } else {
7624 debug!(key = %key, "State action: context set");
7625 }
7626 }
7627 }
7628 StateAction::Prompt {
7629 prompt,
7630 llm,
7631 store_as,
7632 } => {
7633 let llm_result = if let Some(alias) = llm {
7634 self.llm_registry.get(alias)
7635 } else {
7636 self.llm_registry.default()
7637 };
7638 match llm_result {
7639 Ok(llm_provider) => {
7640 let context = self.build_context_with_overlays();
7642 let rendered_prompt = self
7643 .template_renderer
7644 .render(prompt, &context)
7645 .unwrap_or_else(|_| prompt.clone());
7646 let recent =
7647 self.memory.get_messages(Some(5)).await.unwrap_or_default();
7648 let mut messages: Vec<ChatMessage> = recent;
7649 messages.push(ChatMessage::user(&rendered_prompt));
7650 match self
7651 .observe_purpose(
7652 ObservationPurpose::StateAction,
7653 llm_provider.complete(&messages, None),
7654 )
7655 .await
7656 {
7657 Ok(response) => {
7658 if let Some(key) = store_as {
7659 let _ = self
7660 .context_manager
7661 .set(key, Value::String(response.content));
7662 debug!(key = %key, "State action: prompt result stored");
7663 }
7664 }
7665 Err(e) => {
7666 warn!(error = %e, "State action: prompt LLM call failed");
7667 }
7668 }
7669 }
7670 Err(e) => {
7671 warn!(error = %e, "State action: LLM not found for prompt");
7672 }
7673 }
7674 }
7675 }
7676 }
7677 }
7678
7679 async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7680 let extractors = match &self.state_machine {
7681 Some(sm) => match sm.current_definition() {
7682 Some(def) if !def.extract.is_empty() => def.extract.clone(),
7683 _ => return HashMap::new(),
7684 },
7685 None => return HashMap::new(),
7686 };
7687
7688 let mut staged = HashMap::new();
7689 for extractor in &extractors {
7690 let prompt = if let Some(ref custom) = extractor.llm_extract {
7691 format!(
7692 "User message:\n\"{}\"\n\nInstruction:\n{}",
7693 user_message, custom
7694 )
7695 } else if let Some(ref desc) = extractor.description {
7696 format!(
7697 "From the following message, extract: {}\n\n\
7698 Message: \"{}\"\n\n\
7699 If the information is present, return ONLY the extracted value.\n\
7700 If NOT present, return exactly: __NONE__",
7701 desc, user_message
7702 )
7703 } else {
7704 continue;
7705 };
7706
7707 let llm = match self
7708 .llm_registry
7709 .get(&extractor.llm)
7710 .or_else(|_| self.llm_registry.get("router"))
7711 .or_else(|_| self.llm_registry.get("default"))
7712 {
7713 Ok(llm) => llm,
7714 Err(e) => {
7715 warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
7716 continue;
7717 }
7718 };
7719
7720 let messages = vec![ChatMessage::user(&prompt)];
7721 match self
7722 .observe_purpose(
7723 ObservationPurpose::ContextExtraction,
7724 llm.complete(&messages, None),
7725 )
7726 .await
7727 {
7728 Ok(response) => {
7729 let value = response.content.trim().to_string();
7730 if value != "__NONE__" && !value.is_empty() {
7731 staged.insert(
7732 extractor.key.clone(),
7733 serde_json::Value::String(value.clone()),
7734 );
7735 debug!(key = %extractor.key, value = %value, "Context extracted");
7736 } else if extractor.required {
7737 warn!(key = %extractor.key, "Required extraction returned no value");
7738 }
7739 }
7740 Err(e) => {
7741 warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
7742 }
7743 }
7744 }
7745 staged
7746 }
7747
7748 fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
7749 for (key, value) in staged {
7750 if let Err(error) = self.context_manager.update(key, value.clone()) {
7751 warn!(key = %key, error = %error, "staged context write failed");
7752 }
7753 }
7754 }
7755
7756 async fn run_context_extractors(&self, user_message: &str) {
7758 let staged = self.run_context_extractors_staged(user_message).await;
7759 self.commit_staged_context_writes(&staged);
7760 }
7761
7762 async fn check_memory_compression(&self) -> Result<()> {
7763 if self.memory.needs_compression() {
7764 let result = self.memory.compress(None).await?;
7765 if let CompressResult::Compressed {
7766 messages_summarized,
7767 new_summary_length,
7768 tokens_saved,
7769 } = result
7770 {
7771 let event = MemoryCompressEvent::new(
7772 messages_summarized,
7773 tokens_saved,
7774 new_summary_length as u32,
7775 );
7776 self.hooks.on_memory_compress(&event).await;
7777 debug!(
7778 messages = messages_summarized,
7779 tokens_saved = tokens_saved,
7780 "Memory compressed"
7781 );
7782 }
7783 }
7784
7785 self.handle_memory_overflow().await?;
7787 self.check_memory_budget().await;
7788
7789 Ok(())
7790 }
7791
7792 async fn check_memory_budget(&self) {
7793 let Some(ref budget) = self.memory_token_budget else {
7794 return;
7795 };
7796
7797 let context = match self.memory.get_context().await {
7798 Ok(ctx) => ctx,
7799 Err(_) => return,
7800 };
7801
7802 let used_tokens = context.estimated_tokens();
7804 if budget.is_over_warn_threshold(used_tokens) {
7805 let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
7806 self.hooks.on_memory_budget_warning(&event).await;
7807 debug!(
7808 used = used_tokens,
7809 total = budget.total,
7810 percent = event.usage_percent,
7811 "Memory budget warning"
7812 );
7813 }
7814
7815 if let Some(ref summary) = context.summary {
7817 let summary_tokens = ai_agents_memory::estimate_tokens(summary);
7818 let summary_budget = budget.allocation.summary;
7819 if summary_budget > 0 {
7820 let warn_threshold =
7821 (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7822 if summary_tokens >= warn_threshold {
7823 let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
7824 self.hooks.on_memory_budget_warning(&event).await;
7825 }
7826 }
7827 }
7828
7829 let recent_tokens: u32 = context
7831 .messages
7832 .iter()
7833 .map(ai_agents_memory::estimate_message_tokens)
7834 .sum();
7835 let recent_budget = budget.allocation.recent_messages;
7836 if recent_budget > 0 {
7837 let warn_threshold =
7838 (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7839 if recent_tokens >= warn_threshold {
7840 let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
7841 self.hooks.on_memory_budget_warning(&event).await;
7842 }
7843 }
7844
7845 let relationship_budget = budget.allocation.relationships;
7846 if relationship_budget > 0 {
7847 let relationship_tokens = self
7848 .relationship_memory_text()
7849 .map(|text| ai_agents_memory::estimate_tokens(&text))
7850 .unwrap_or(0);
7851 let warn_threshold =
7852 (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7853 if relationship_tokens >= warn_threshold {
7854 let event = MemoryBudgetEvent::new(
7855 "relationships",
7856 relationship_tokens,
7857 relationship_budget,
7858 );
7859 self.hooks.on_memory_budget_warning(&event).await;
7860 }
7861 }
7862 }
7863
7864 async fn handle_memory_overflow(&self) -> Result<()> {
7865 let Some(ref budget) = self.memory_token_budget else {
7866 return Ok(());
7867 };
7868
7869 let context = self.memory.get_context().await?;
7870 let used_tokens = context.estimated_tokens();
7871
7872 if used_tokens <= budget.total {
7873 return Ok(());
7874 }
7875
7876 match budget.overflow_strategy {
7877 OverflowStrategy::TruncateOldest => {
7878 let tokens_to_free = used_tokens - budget.total;
7879 let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
7880 if messages_to_evict > 0 {
7881 self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
7882 .await?;
7883 }
7884 }
7885 OverflowStrategy::SummarizeMore => {
7886 let max_attempts = context.total_messages.max(1);
7887 for _ in 0..max_attempts {
7888 match self.memory.compress(None).await? {
7889 CompressResult::Compressed {
7890 messages_summarized,
7891 ..
7892 } if messages_summarized > 0 => {
7893 let context = self.memory.get_context().await?;
7894 if context.estimated_tokens() <= budget.total {
7895 return Ok(());
7896 }
7897 }
7898 _ => break,
7899 }
7900 }
7901 let context = self.memory.get_context().await?;
7902 let used_tokens = context.estimated_tokens();
7903 if used_tokens > budget.total {
7904 return Err(AgentError::MemoryBudgetExceeded {
7905 used: used_tokens,
7906 budget: budget.total,
7907 });
7908 }
7909 }
7910 OverflowStrategy::Error => {
7911 return Err(AgentError::MemoryBudgetExceeded {
7912 used: used_tokens,
7913 budget: budget.total,
7914 });
7915 }
7916 }
7917 Ok(())
7918 }
7919
7920 fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
7921 ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
7923 }
7924
7925 async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
7926 let evicted = self.memory.evict_oldest(count).await?;
7927 if !evicted.is_empty() {
7928 let event = MemoryEvictEvent {
7929 reason,
7930 messages_evicted: evicted.len(),
7931 importance_scores: vec![],
7932 };
7933 self.hooks.on_memory_evict(&event).await;
7934 debug!(count = evicted.len(), "Messages evicted from memory");
7935 }
7936 Ok(())
7937 }
7938
7939 #[instrument(skip(self, input), fields(agent = %self.info.name))]
7940 async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
7941 match self.determine_reasoning_mode_strict(input).await {
7942 Ok(mode) => Ok(mode),
7943 Err(_) => Ok(ReasoningMode::None),
7944 }
7945 }
7946
7947 async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
7948 let effective_config = self.get_effective_reasoning_config();
7949
7950 if !matches!(effective_config.mode, ReasoningMode::Auto) {
7951 return Ok(effective_config.mode.clone());
7952 }
7953
7954 let judge_llm = effective_config
7955 .judge_llm
7956 .as_ref()
7957 .and_then(|alias| self.llm_registry.get(alias).ok())
7958 .or_else(|| self.llm_registry.router().ok())
7959 .or_else(|| self.llm_registry.default().ok());
7960
7961 let Some(llm) = judge_llm else {
7962 return Ok(ReasoningMode::None);
7963 };
7964
7965 let prompt = format!(
7966 r#"Analyze this user request and determine the appropriate reasoning mode.
7967
7968User request: "{}"
7969
7970Choose ONE of these modes:
7971- none: Simple queries, greetings, direct answers (fastest)
7972- cot: Complex analysis, multi-step reasoning, math problems
7973- react: Tasks requiring multiple tool calls with observation
7974- plan_and_execute: Complex multi-step tasks requiring coordination
7975
7976Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
7977 input
7978 );
7979
7980 let messages = vec![ChatMessage::user(&prompt)];
7981 let response = self
7982 .observe_purpose(
7983 ObservationPurpose::ReflectionDecision,
7984 llm.complete(&messages, None),
7985 )
7986 .await
7987 .map_err(|e| AgentError::LLM(e.to_string()))?;
7988
7989 let mode_str = response.content.trim().to_lowercase();
7990 Ok(match mode_str.as_str() {
7991 "cot" => ReasoningMode::CoT,
7992 "react" => ReasoningMode::React,
7993 "plan_and_execute" => ReasoningMode::PlanAndExecute,
7994 _ => ReasoningMode::None,
7995 })
7996 }
7997
7998 async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
7999 let effective_config = self.get_effective_reflection_config();
8000
8001 if !effective_config.requires_evaluation() {
8002 return Ok(false);
8003 }
8004
8005 if effective_config.is_enabled() {
8006 return Ok(true);
8007 }
8008
8009 let evaluator_llm = effective_config
8010 .evaluator_llm
8011 .as_ref()
8012 .and_then(|alias| self.llm_registry.get(alias).ok())
8013 .or_else(|| self.llm_registry.router().ok())
8014 .or_else(|| self.llm_registry.default().ok());
8015
8016 let Some(llm) = evaluator_llm else {
8017 return Ok(false);
8018 };
8019
8020 let response_preview: String = response.chars().take(500).collect();
8021 let prompt = format!(
8022 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8023
8024User query: "{}"
8025Response: "{}"
8026
8027Answer YES or NO only."#,
8028 input, response_preview
8029 );
8030
8031 let messages = vec![ChatMessage::user(&prompt)];
8032 let result = self
8033 .observe_purpose(
8034 ObservationPurpose::ReflectionDecision,
8035 llm.complete(&messages, None),
8036 )
8037 .await;
8038
8039 match result {
8040 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8041 Err(_) => Ok(false),
8042 }
8043 }
8044
8045 fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8046 format!(
8047 "{}\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>",
8048 base_prompt
8049 )
8050 }
8051
8052 fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8053 format!(
8054 "{}\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>",
8055 base_prompt
8056 )
8057 }
8058
8059 async fn generate_plan(&self, input: &str) -> Result<Plan> {
8060 let effective = self.get_effective_reasoning_config();
8061 let planning_config = effective.get_planning();
8062
8063 let planner_llm = planning_config
8064 .and_then(|c| c.planner_llm.as_ref())
8065 .and_then(|alias| self.llm_registry.get(alias).ok())
8066 .or_else(|| self.llm_registry.router().ok())
8067 .or_else(|| self.llm_registry.default().ok())
8068 .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8069
8070 let mut available_tool_ids: Vec<String> = self
8071 .get_available_tool_ids()
8072 .await
8073 .unwrap_or_else(|_| self.tools.list_ids());
8074 let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8075
8076 if let Some(config) = planning_config {
8078 if !config.available.tools.is_all() {
8079 available_tool_ids.retain(|t| config.available.tools.allows(t));
8080 }
8081 if !config.available.skills.is_all() {
8082 available_skills.retain(|s| config.available.skills.allows(s));
8083 }
8084 }
8085
8086 let tool_descriptions: Vec<String> = available_tool_ids
8089 .iter()
8090 .filter_map(|id| {
8091 self.tools.get(id).map(|tool| {
8092 let schema = tool.input_schema();
8093 let args_desc = schema
8094 .get("properties")
8095 .and_then(|p| serde_json::to_string(p).ok())
8096 .unwrap_or_else(|| "{}".to_string());
8097 format!(
8098 "- {} ({}): {}\n Arguments: {}",
8099 id,
8100 tool.name(),
8101 tool.description(),
8102 args_desc
8103 )
8104 })
8105 })
8106 .collect();
8107
8108 let tools_section = if tool_descriptions.is_empty() {
8109 "Available tools: none".to_string()
8110 } else {
8111 format!("Available tools:\n{}", tool_descriptions.join("\n"))
8112 };
8113
8114 let skills_section = if available_skills.is_empty() {
8115 "Available skills: none".to_string()
8116 } else {
8117 format!("Available skills: {}", available_skills.join(", "))
8118 };
8119
8120 let prompt = format!(
8121 r#"Create a step-by-step plan to accomplish this goal.
8122
8123Goal: "{}"
8124
8125{}
8126
8127{}
8128
8129Create a plan with clear steps. For each step, specify:
8130- description: What this step accomplishes
8131- action_type: "tool", "skill", "think", or "respond"
8132- action_target: The tool/skill id (if applicable)
8133- args: The arguments object matching the tool's schema (if action_type is "tool")
8134- dependencies: List of step IDs this depends on (empty if none)
8135
8136Respond in JSON format:
8137{{
8138 "steps": [
8139 {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8140 {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8141 ]
8142}}"#,
8143 input, tools_section, skills_section,
8144 );
8145
8146 let messages = vec![ChatMessage::user(&prompt)];
8147 let response = self
8148 .observe_purpose(
8149 ObservationPurpose::PlanGeneration,
8150 planner_llm.complete(&messages, None),
8151 )
8152 .await
8153 .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8154
8155 let mut plan = Plan::new(input);
8156
8157 if let Some(json_start) = response.content.find('{')
8158 && let Some(json_end) = response.content.rfind('}')
8159 {
8160 let json_str = &response.content[json_start..=json_end];
8161 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8162 && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8163 {
8164 for step_value in steps {
8165 let id = step_value
8166 .get("id")
8167 .and_then(|v| v.as_str())
8168 .unwrap_or("step");
8169 let desc = step_value
8170 .get("description")
8171 .and_then(|v| v.as_str())
8172 .unwrap_or("");
8173 let action_type = step_value
8174 .get("action_type")
8175 .and_then(|v| v.as_str())
8176 .unwrap_or("think");
8177 let action_target = step_value
8178 .get("action_target")
8179 .and_then(|v| v.as_str())
8180 .unwrap_or("");
8181 let args = step_value
8182 .get("args")
8183 .cloned()
8184 .unwrap_or(serde_json::json!({}));
8185 let deps: Vec<String> = step_value
8186 .get("dependencies")
8187 .and_then(|v| v.as_array())
8188 .map(|arr| {
8189 arr.iter()
8190 .filter_map(|v| v.as_str().map(String::from))
8191 .collect()
8192 })
8193 .unwrap_or_default();
8194
8195 let action = match action_type {
8196 "tool" => PlanAction::tool(action_target, args),
8197 "skill" => PlanAction::skill(action_target),
8198 "respond" => PlanAction::respond(action_target),
8199 _ => PlanAction::think(desc),
8200 };
8201
8202 let step = PlanStep::new(desc, action)
8203 .with_id(id)
8204 .with_dependencies(deps);
8205 plan.add_step(step);
8206 }
8207 }
8208 }
8209
8210 if plan.steps.is_empty() {
8211 plan.add_step(PlanStep::new(
8212 "Process the request",
8213 PlanAction::think(input),
8214 ));
8215 plan.add_step(PlanStep::new(
8216 "Provide response",
8217 PlanAction::respond("Answer based on analysis"),
8218 ));
8219 }
8220
8221 Ok(plan)
8222 }
8223
8224 async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8225 let llm = self.get_state_llm()?;
8226 let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8227 let effective = self.get_effective_reasoning_config();
8228 let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8229
8230 plan.status = PlanStatus::InProgress;
8231
8232 for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8233 let step = &plan.steps[step_idx];
8234
8235 let deps_satisfied = step.dependencies.iter().all(|dep| {
8236 plan.steps
8237 .iter()
8238 .find(|s| &s.id == dep)
8239 .map(|s| s.status.is_completed())
8240 .unwrap_or(false)
8241 });
8242
8243 if !deps_satisfied {
8244 continue;
8245 }
8246
8247 plan.steps[step_idx].mark_running();
8248
8249 let result = match &plan.steps[step_idx].action {
8250 PlanAction::Tool { tool, args } => {
8251 let has_dep_results = plan.steps[step_idx]
8257 .dependencies
8258 .iter()
8259 .any(|dep| results.contains_key(dep));
8260
8261 let final_args = if has_dep_results {
8262 let dep_context: String = plan.steps[step_idx]
8263 .dependencies
8264 .iter()
8265 .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8266 .collect::<Vec<_>>()
8267 .join("\n");
8268
8269 let tool_schema = self
8270 .tools
8271 .get(tool)
8272 .map(|t| {
8273 let schema = t.input_schema();
8274 let props = schema
8275 .get("properties")
8276 .and_then(|p| serde_json::to_string(p).ok())
8277 .unwrap_or_else(|| "{}".to_string());
8278 format!(
8279 "{}: {}\nArguments schema: {}",
8280 t.id(),
8281 t.description(),
8282 props
8283 )
8284 })
8285 .unwrap_or_default();
8286
8287 let step_desc = &plan.steps[step_idx].description;
8288 let arg_prompt = format!(
8289 "Generate the JSON arguments for a tool call.\n\n\
8290 Tool: {}\n\n\
8291 Task: {}\n\n\
8292 Previous step results:\n{}\n\n\
8293 Planner's draft arguments: {}\n\n\
8294 Produce ONLY a valid JSON object with the correct argument values.\n\
8295 Use actual values from the previous step results, not template references.",
8296 tool_schema,
8297 step_desc,
8298 dep_context,
8299 serde_json::to_string(args).unwrap_or_default()
8300 );
8301 let messages = vec![ChatMessage::user(&arg_prompt)];
8302 match self
8303 .observe_purpose(
8304 ObservationPurpose::PlanStep,
8305 llm.complete(&messages, None),
8306 )
8307 .await
8308 {
8309 Ok(resp) => {
8310 let content = resp.content.trim();
8311 let json_start = content.find('{');
8313 let json_end = content.rfind('}');
8314 if let (Some(start), Some(end)) = (json_start, json_end) {
8315 serde_json::from_str(&content[start..=end])
8316 .unwrap_or_else(|_| args.clone())
8317 } else {
8318 args.clone()
8319 }
8320 }
8321 Err(_) => args.clone(),
8322 }
8323 } else {
8324 args.clone()
8325 };
8326
8327 let request = ToolExecutionRequest::new(
8328 uuid::Uuid::new_v4().to_string(),
8329 tool.clone(),
8330 final_args,
8331 ToolCallSource::Plan {
8332 step_index: step_idx,
8333 },
8334 );
8335 match self.execute_tool_record(request).await {
8336 Ok(record) if record.success => {
8337 serde_json::json!({ "output": record.model_output_string() })
8338 }
8339 Ok(record) => {
8340 plan.steps[step_idx].mark_failed(record.model_output_string());
8341 continue;
8342 }
8343 Err(e) => {
8344 plan.steps[step_idx].mark_failed(e.to_string());
8345 continue;
8346 }
8347 }
8348 }
8349 PlanAction::Skill { skill } => {
8350 if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8351 if let Some(ref executor) = self.skill_executor {
8352 match executor
8353 .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8354 .await
8355 {
8356 Ok(output) => serde_json::json!({ "output": output }),
8357 Err(e) => {
8358 plan.steps[step_idx].mark_failed(e.to_string());
8359 continue;
8360 }
8361 }
8362 } else {
8363 serde_json::json!({ "output": "Skill executor not available" })
8364 }
8365 } else {
8366 plan.steps[step_idx].mark_failed("Skill not found");
8367 continue;
8368 }
8369 }
8370 PlanAction::Think { prompt } => {
8371 let context: String = results
8372 .iter()
8373 .map(|(k, v)| format!("{}: {}", k, v))
8374 .collect::<Vec<_>>()
8375 .join("\n");
8376
8377 let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8378 let messages = vec![ChatMessage::user(&think_prompt)];
8379
8380 match self
8381 .observe_purpose(
8382 ObservationPurpose::PlanStep,
8383 llm.complete(&messages, None),
8384 )
8385 .await
8386 {
8387 Ok(resp) => serde_json::json!({ "output": resp.content }),
8388 Err(e) => {
8389 plan.steps[step_idx].mark_failed(e.to_string());
8390 continue;
8391 }
8392 }
8393 }
8394 PlanAction::Respond { template } => {
8395 let context: String = results
8396 .iter()
8397 .map(|(k, v)| format!("{}: {}", k, v))
8398 .collect::<Vec<_>>()
8399 .join("\n");
8400
8401 let respond_prompt = format!(
8402 "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8403 context, template
8404 );
8405 let messages = vec![ChatMessage::user(&respond_prompt)];
8406
8407 match self
8408 .observe_purpose(
8409 ObservationPurpose::PlanStep,
8410 llm.complete(&messages, None),
8411 )
8412 .await
8413 {
8414 Ok(resp) => serde_json::json!({ "output": resp.content }),
8415 Err(e) => {
8416 plan.steps[step_idx].mark_failed(e.to_string());
8417 continue;
8418 }
8419 }
8420 }
8421 };
8422
8423 results.insert(plan.steps[step_idx].id.clone(), result.clone());
8424 plan.steps[step_idx].mark_completed(Some(result));
8425 }
8426
8427 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8429 if has_failures {
8430 let failed_ids: Vec<String> = plan
8431 .steps
8432 .iter()
8433 .filter(|s| s.status.is_failed())
8434 .map(|s| s.id.clone())
8435 .collect();
8436 plan.status = PlanStatus::Failed {
8437 error: format!("Steps failed: {}", failed_ids.join(", ")),
8438 };
8439 } else {
8440 plan.status = PlanStatus::Completed;
8441 }
8442
8443 let all_outputs: Vec<String> = plan
8445 .steps
8446 .iter()
8447 .filter(|s| s.status.is_completed())
8448 .filter_map(|s| {
8449 s.result
8450 .as_ref()
8451 .and_then(|r| r.get("output"))
8452 .and_then(|o| o.as_str())
8453 .map(|o| format!("{}: {}", s.description, o))
8454 })
8455 .collect();
8456
8457 if all_outputs.is_empty() {
8458 return Ok("Plan execution completed but produced no results.".to_string());
8459 }
8460
8461 if all_outputs.len() == 1 {
8462 return Ok(all_outputs.into_iter().next().unwrap());
8463 }
8464
8465 let context = all_outputs.join("\n\n");
8467 let prompt = format!(
8468 "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8469 plan.goal, context
8470 );
8471 let messages = vec![ChatMessage::user(&prompt)];
8472 match self
8473 .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8474 .await
8475 {
8476 Ok(resp) => Ok(resp.content.trim().to_string()),
8477 Err(_) => Ok(context),
8478 }
8479 }
8480
8481 async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8482 let effective_config = self.get_effective_reflection_config();
8483 self.evaluate_response_with_config(input, response, &effective_config)
8484 .await
8485 }
8486
8487 fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8488 if let Some(start) = content.find("<thinking>")
8489 && let Some(end) = content.find("</thinking>")
8490 {
8491 let thinking = content[start + 10..end].trim().to_string();
8492 let answer = content[end + 11..].trim().to_string();
8493 return (Some(thinking), answer);
8494 }
8495 (None, content.to_string())
8496 }
8497
8498 fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8499 match self.get_effective_reasoning_config().output {
8500 ReasoningOutput::Hidden => answer.to_string(),
8501 ReasoningOutput::Visible => {
8502 if let Some(t) = thinking {
8503 format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8504 } else {
8505 answer.to_string()
8506 }
8507 }
8508 ReasoningOutput::Tagged => {
8509 if let Some(t) = thinking {
8510 format!("<thinking>{}</thinking>\n{}", t, answer)
8511 } else {
8512 answer.to_string()
8513 }
8514 }
8515 }
8516 }
8517
8518 async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8523 self.init_storage().await?;
8527 self.begin_root_turn();
8528 let _root_cleanup = RootTurnCleanup::new(self);
8529 info!(input_len = input.len(), "Starting chat");
8530
8531 self.hooks.on_message_received(input).await;
8532
8533 if !self.context_initialized.swap(true, Ordering::SeqCst) {
8537 self.context_manager.initialize().await?;
8538 debug!("Context manager initialized (defaults, env, builtins)");
8539 }
8540
8541 self.check_turn_timeout().await?;
8542 self.context_manager.refresh_per_turn().await?;
8543
8544 self.clear_disambiguation_context();
8547
8548 if let Some(ref disambiguator) = self.disambiguation_manager {
8550 let disambiguation_context = self.build_disambiguation_context().await?;
8551
8552 let state_override = self
8554 .state_machine
8555 .as_ref()
8556 .and_then(|sm| sm.current_definition())
8557 .and_then(|def| def.disambiguation.clone());
8558
8559 let state_generation = self
8560 .state_machine
8561 .as_ref()
8562 .map(|state_machine| state_machine.generation());
8563 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8564 let mut disambiguation_result = self
8565 .observe_purpose(
8566 ObservationPurpose::DisambiguationDetection,
8567 disambiguator.process_input_with_override(
8568 input,
8569 &disambiguation_context,
8570 state_override.as_ref(),
8571 None,
8572 ),
8573 )
8574 .await?;
8575 let current_state_generation = self
8576 .state_machine
8577 .as_ref()
8578 .map(|state_machine| state_machine.generation());
8579 if current_state_generation != state_generation
8580 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8581 {
8582 disambiguator.clear_pending().await;
8583 *self.pending_skill_id.write() = None;
8584 disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8585 info!(
8586 confirmation_event = "invalidated",
8587 invalidation_reason = "state_generation_changed",
8588 "Disambiguation result invalidated before redispatch"
8589 );
8590 }
8591 match disambiguation_result {
8592 DisambiguationResult::Clear => {
8593 debug!("Input is clear, proceeding normally");
8594 }
8595 DisambiguationResult::NeedsClarification {
8596 question,
8597 detection,
8598 } => {
8599 let admission = self
8600 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8601 .await?;
8602 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8603 info!(
8604 ambiguity_type = ?detection.ambiguity_type,
8605 confidence = detection.confidence,
8606 "Input requires clarification"
8607 );
8608
8609 self.commit_root_user_message(input).await?;
8612 self.memory
8613 .add_message(ChatMessage::assistant(&question.question))
8614 .await?;
8615
8616 let status = if awaiting_confirmation {
8617 "awaiting_confirmation"
8618 } else {
8619 "awaiting_clarification"
8620 };
8621 let response = AgentResponse::new(&question.question).with_metadata(
8622 "disambiguation",
8623 serde_json::json!({
8624 "status": status,
8625 "options": question.options,
8626 "clarifying": question.clarifying,
8627 "detection": {
8628 "type": detection.ambiguity_type,
8629 "confidence": detection.confidence,
8630 "what_is_unclear": detection.what_is_unclear,
8631 }
8632 }),
8633 );
8634 drop(admission);
8635 self.finish_turn_if_root(&response).await?;
8636 return Ok(response);
8637 }
8638 DisambiguationResult::Clarified {
8639 enriched_input,
8640 resolved,
8641 ..
8642 } => {
8643 let admission = match self
8644 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8645 .await
8646 {
8647 Ok(admission) => admission,
8648 Err(error) => {
8649 *self.pending_skill_id.write() = None;
8650 return Err(error);
8651 }
8652 };
8653 info!(
8654 resolved_count = resolved.len(),
8655 enriched = %enriched_input,
8656 "Input clarified, injecting resolved intent into context"
8657 );
8658
8659 for (key, value) in &resolved {
8662 let context_key = format!("disambiguation.{}", key);
8663 let _ = self.context_manager.set(&context_key, value.clone());
8664 }
8665
8666 if let Some(intent) = resolved.get("intent") {
8667 let _ = self.context_manager.set("resolved_intent", intent.clone());
8668 }
8669
8670 let _ = self
8671 .context_manager
8672 .set("disambiguation.resolved", serde_json::Value::Bool(true));
8673
8674 let skill_id = self.pending_skill_id.read().clone();
8678 if let Some(skill_id) = skill_id {
8679 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8680 drop(admission);
8681 return self
8682 .recheck_skill_disambiguation(
8683 &skill_id,
8684 &enriched_input,
8685 disambiguation_epoch,
8686 state_generation,
8687 )
8688 .await;
8689 }
8690
8691 drop(admission);
8692 return self.run_loop_internal(&enriched_input).await;
8693 }
8694 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
8695 info!("Proceeding with best guess interpretation");
8696
8697 let skill_id = self.pending_skill_id.read().clone();
8699 if let Some(skill_id) = skill_id {
8700 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
8701 return self
8702 .recheck_skill_disambiguation(
8703 &skill_id,
8704 &enriched_input,
8705 disambiguation_epoch,
8706 state_generation,
8707 )
8708 .await;
8709 }
8710
8711 return self.run_loop_internal(&enriched_input).await;
8712 }
8713 DisambiguationResult::GiveUp { reason } => {
8714 *self.pending_skill_id.write() = None;
8715 warn!(reason = %reason, "Disambiguation gave up");
8716 let apology = self
8717 .generate_localized_apology(
8718 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
8719 &reason,
8720 )
8721 .await
8722 .unwrap_or_else(|_| {
8723 format!("I'm sorry, I couldn't understand your request: {}", reason)
8724 });
8725 let response = AgentResponse::new(&apology);
8726 self.finish_turn_if_root(&response).await?;
8727 return Ok(response);
8728 }
8729 DisambiguationResult::Escalate { reason } => {
8730 *self.pending_skill_id.write() = None;
8731 info!(reason = %reason, "Escalating to human");
8732 if let Some(ref hitl) = self.hitl_engine {
8733 let trigger =
8734 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
8735 let mut context_map = HashMap::new();
8736 context_map.insert("original_input".to_string(), serde_json::json!(input));
8737 context_map.insert("reason".to_string(), serde_json::json!(&reason));
8738 let check_result = HITLCheckResult::required(
8739 trigger,
8740 context_map,
8741 format!("User request needs human assistance: {}", reason),
8742 Some(hitl.config().default_timeout_seconds),
8743 );
8744 let result = self.request_hitl_approval(check_result).await?;
8745 if matches!(
8746 result,
8747 ApprovalResult::Approved | ApprovalResult::Modified { .. }
8748 ) {
8749 return self.run_loop_internal(input).await;
8750 }
8751 }
8752 let apology = self
8753 .generate_localized_apology(
8754 "Explain briefly that you're transferring the user to a human agent for help.",
8755 &reason,
8756 )
8757 .await
8758 .unwrap_or_else(|_| {
8759 format!("I need human assistance to help with your request: {}", reason)
8760 });
8761 let response = AgentResponse::new(&apology);
8762 self.finish_turn_if_root(&response).await?;
8763 return Ok(response);
8764 }
8765 DisambiguationResult::Abandoned { new_input } => {
8766 *self.pending_skill_id.write() = None;
8767
8768 info!(
8769 has_new_input = new_input.is_some(),
8770 "Clarification abandoned by user"
8771 );
8772
8773 self.commit_root_user_message(input).await?;
8774
8775 match new_input {
8776 Some(fresh_input) => {
8777 return self.run_loop_internal(&fresh_input).await;
8780 }
8781 None => {
8782 let ack = self
8784 .generate_localized_apology(
8785 "The user changed their mind about their previous request. \
8786 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
8787 Do NOT apologize excessively. Be concise.",
8788 "User abandoned clarification",
8789 )
8790 .await
8791 .unwrap_or_else(|_| {
8792 "OK, no problem. What else can I help with?".to_string()
8793 });
8794
8795 self.memory
8796 .add_message(ChatMessage::assistant(&ack))
8797 .await?;
8798
8799 let response = AgentResponse::new(&ack);
8800 self.finish_turn_if_root(&response).await?;
8801 return Ok(response);
8802 }
8803 }
8804 }
8805 }
8806 }
8807
8808 self.run_loop_internal(input).await
8809 }
8810
8811 async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
8813 let llm = self.llm_registry.router().map_err(|e| {
8814 AgentError::LLM(format!(
8815 "Router LLM not available for localized response: {}",
8816 e
8817 ))
8818 })?;
8819
8820 let recent: Vec<String> = self
8821 .memory
8822 .get_messages(Some(3))
8823 .await?
8824 .iter()
8825 .map(|m| m.content.clone())
8826 .collect();
8827
8828 let context_hint = if recent.is_empty() {
8829 String::new()
8830 } else {
8831 format!(
8832 "\nRecent conversation (detect the user's language from this):\n{}\n",
8833 recent.join("\n")
8834 )
8835 };
8836
8837 let prompt = format!(
8838 "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
8839 instruction, reason, context_hint
8840 );
8841
8842 let messages = vec![ChatMessage::user(&prompt)];
8843 let response = self
8844 .observe_purpose(
8845 ObservationPurpose::DisambiguationClarification,
8846 llm.complete(&messages, None),
8847 )
8848 .await
8849 .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
8850
8851 Ok(response.content.trim().to_string())
8852 }
8853
8854 fn render_action_args(&self, args: &Value) -> Value {
8858 let context = self.build_context_with_overlays();
8859 match args {
8860 Value::Object(map) => {
8861 let mut rendered = serde_json::Map::new();
8862 for (k, v) in map {
8863 match v {
8864 Value::String(s) if s.contains("{{") => {
8865 match self.template_renderer.render(s, &context) {
8866 Ok(rendered_str) => {
8867 rendered.insert(k.clone(), Value::String(rendered_str));
8868 }
8869 Err(_) => {
8870 rendered.insert(k.clone(), v.clone());
8871 }
8872 }
8873 }
8874 _ => {
8875 rendered.insert(k.clone(), v.clone());
8876 }
8877 }
8878 }
8879 Value::Object(rendered)
8880 }
8881 _ => args.clone(),
8882 }
8883 }
8884
8885 fn clear_disambiguation_context(&self) {
8887 let _ = self
8888 .context_manager
8889 .set("resolved_intent", serde_json::Value::Null);
8890
8891 let all = self.context_manager.get_all();
8892 for key in all.keys() {
8893 if key.starts_with("disambiguation.") {
8894 let _ = self.context_manager.set(key, serde_json::Value::Null);
8895 }
8896 }
8897 }
8898
8899 async fn recheck_skill_disambiguation(
8905 &self,
8906 skill_id: &str,
8907 enriched_input: &str,
8908 expected_disambiguation_epoch: u64,
8909 expected_state_generation: Option<u64>,
8910 ) -> Result<AgentResponse> {
8911 let skill = self
8912 .skill_router
8913 .as_ref()
8914 .and_then(|r| r.get_skill(skill_id).cloned());
8915
8916 if let Some(ref skill) = skill
8918 && let Some(ref skill_disambig) = skill.disambiguation
8919 && skill_disambig.enabled.unwrap_or(false)
8920 && let Some(ref disambiguator) = self.disambiguation_manager
8921 {
8922 let context = self.build_disambiguation_context().await?;
8923 let state_override = self
8924 .state_machine
8925 .as_ref()
8926 .and_then(|sm| sm.current_definition())
8927 .and_then(|def| def.disambiguation.clone());
8928
8929 let disambiguation_result = self
8930 .observe_purpose(
8931 ObservationPurpose::DisambiguationDetection,
8932 disambiguator.process_input_with_override(
8933 enriched_input,
8934 &context,
8935 state_override.as_ref(),
8936 Some(skill_disambig),
8937 ),
8938 )
8939 .await?;
8940 let current_state_generation = self
8941 .state_machine
8942 .as_ref()
8943 .map(|state_machine| state_machine.generation());
8944 if current_state_generation != expected_state_generation
8945 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
8946 {
8947 disambiguator.clear_pending().await;
8948 *self.pending_skill_id.write() = None;
8949 return Err(AgentError::Other(
8950 "State or reset ownership changed during skill disambiguation recheck"
8951 .to_string(),
8952 ));
8953 }
8954 match disambiguation_result {
8955 DisambiguationResult::Clear => {
8956 debug!(skill_id = %skill_id, "Skill re-check: all fields present");
8957 }
8958 DisambiguationResult::NeedsClarification {
8959 question,
8960 detection,
8961 } => {
8962 let admission = self
8963 .admit_disambiguation_redispatch(
8964 expected_disambiguation_epoch,
8965 expected_state_generation,
8966 )
8967 .await?;
8968 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8969 info!(
8970 skill_id = %skill_id,
8971 ambiguity_type = ?detection.ambiguity_type,
8972 what_is_unclear = ?detection.what_is_unclear,
8973 "Skill re-check: still missing fields, asking again"
8974 );
8975 self.memory
8979 .add_message(ChatMessage::user(enriched_input))
8980 .await?;
8981 self.memory
8982 .add_message(ChatMessage::assistant(&question.question))
8983 .await?;
8984
8985 let response = AgentResponse::new(&question.question).with_metadata(
8986 "disambiguation",
8987 serde_json::json!({
8988 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
8989 "skill_id": skill_id,
8990 "options": question.options,
8991 "clarifying": question.clarifying,
8992 "detection": {
8993 "type": detection.ambiguity_type,
8994 "confidence": detection.confidence,
8995 "what_is_unclear": detection.what_is_unclear,
8996 }
8997 }),
8998 );
8999 drop(admission);
9000 self.finish_turn_if_root(&response).await?;
9001 return Ok(response);
9002 }
9003 DisambiguationResult::Clarified {
9004 enriched_input: re_enriched,
9005 ..
9006 } => {
9007 debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
9008 let admission = self
9009 .admit_disambiguation_redispatch(
9010 expected_disambiguation_epoch,
9011 expected_state_generation,
9012 )
9013 .await?;
9014 *self.pending_skill_id.write() = None;
9015 drop(admission);
9016 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9017 self.memory
9018 .add_message(ChatMessage::user(&re_enriched))
9019 .await?;
9020 return self
9021 .handle_skill_response(
9022 &re_enriched,
9023 skill_id,
9024 skill_response,
9025 &HashMap::new(),
9026 )
9027 .await;
9028 }
9029 DisambiguationResult::ProceedWithBestGuess {
9030 enriched_input: re_enriched,
9031 } => {
9032 debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9033 let admission = self
9034 .admit_disambiguation_redispatch(
9035 expected_disambiguation_epoch,
9036 expected_state_generation,
9037 )
9038 .await?;
9039 *self.pending_skill_id.write() = None;
9040 drop(admission);
9041 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9042 self.memory
9043 .add_message(ChatMessage::user(&re_enriched))
9044 .await?;
9045 return self
9046 .handle_skill_response(
9047 &re_enriched,
9048 skill_id,
9049 skill_response,
9050 &HashMap::new(),
9051 )
9052 .await;
9053 }
9054 DisambiguationResult::GiveUp { reason } => {
9055 *self.pending_skill_id.write() = None;
9056 let apology = self
9057 .generate_localized_apology(
9058 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9059 &reason,
9060 )
9061 .await
9062 .unwrap_or_else(|_| {
9063 format!("I'm sorry, I couldn't understand your request: {}", reason)
9064 });
9065 let response = AgentResponse::new(&apology);
9066 self.finish_turn_if_root(&response).await?;
9067 return Ok(response);
9068 }
9069 DisambiguationResult::Escalate { reason } => {
9070 *self.pending_skill_id.write() = None;
9071 let apology = self
9072 .generate_localized_apology(
9073 "Explain briefly that you're transferring the user to a human agent for help.",
9074 &reason,
9075 )
9076 .await
9077 .unwrap_or_else(|_| {
9078 format!("I need human assistance to help with your request: {}", reason)
9079 });
9080 let response = AgentResponse::new(&apology);
9081 self.finish_turn_if_root(&response).await?;
9082 return Ok(response);
9083 }
9084 DisambiguationResult::Abandoned { new_input } => {
9085 *self.pending_skill_id.write() = None;
9088 debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9089 if let Some(fresh) = new_input {
9090 return self.run_loop_internal(&fresh).await;
9091 }
9092 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 self.memory
9104 .add_message(ChatMessage::assistant(&ack))
9105 .await?;
9106 let response = AgentResponse::new(&ack);
9107 self.finish_turn_if_root(&response).await?;
9108 return Ok(response);
9109 }
9110 }
9111 }
9112
9113 let admission = self
9115 .admit_disambiguation_redispatch(
9116 expected_disambiguation_epoch,
9117 expected_state_generation,
9118 )
9119 .await?;
9120 *self.pending_skill_id.write() = None;
9121 drop(admission);
9122 let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9123 self.memory
9124 .add_message(ChatMessage::user(enriched_input))
9125 .await?;
9126 self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9127 .await
9128 }
9129
9130 async fn handle_skill_response(
9133 &self,
9134 processed_input: &str,
9135 skill_id: &str,
9136 skill_response: String,
9137 input_context: &HashMap<String, Value>,
9138 ) -> Result<AgentResponse> {
9139 let output_data = self.process_output(&skill_response, input_context).await?;
9140 let final_response = output_data.content;
9141
9142 self.memory
9143 .add_message(ChatMessage::assistant(&final_response))
9144 .await?;
9145
9146 self.check_memory_compression().await?;
9147
9148 self.increment_turn();
9149 self.evaluate_transitions(processed_input, &final_response)
9150 .await?;
9151
9152 let response = AgentResponse::new(final_response)
9153 .with_metadata("skill_id", serde_json::json!(skill_id));
9154 self.finish_turn_if_root(&response).await?;
9155 Ok(response)
9156 }
9157
9158 async fn handle_plan_and_execute(
9161 &self,
9162 processed_input: &str,
9163 input_context: &HashMap<String, Value>,
9164 auto_detected: bool,
9165 ) -> Result<AgentResponse> {
9166 let effective = self.get_effective_reasoning_config();
9167 let plan_reflection = effective
9168 .get_planning()
9169 .map(|c| c.reflection.clone())
9170 .unwrap_or_default();
9171
9172 let max_attempts = if plan_reflection.enabled {
9173 1 + plan_reflection.max_replans
9174 } else {
9175 1
9176 };
9177
9178 let mut plan = self.generate_plan(processed_input).await?;
9179 info!(
9180 plan_id = %plan.id,
9181 steps = plan.steps.len(),
9182 "Plan generated"
9183 );
9184
9185 let mut plan_result = String::new();
9186
9187 for attempt in 0..max_attempts {
9188 *self.current_plan.write() = Some(plan.clone());
9189 plan_result = self.execute_plan(&mut plan).await?;
9190
9191 info!(
9192 plan_status = ?plan.status,
9193 completed_steps = plan.completed_steps().count(),
9194 attempt = attempt + 1,
9195 "Plan execution completed"
9196 );
9197
9198 if !plan_reflection.enabled {
9199 break;
9200 }
9201
9202 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9203 if !has_failures {
9204 break;
9205 }
9206
9207 if attempt + 1 >= max_attempts {
9208 break;
9209 }
9210
9211 match plan_reflection.on_step_failure {
9212 StepFailureAction::Replan => {
9213 info!(attempt = attempt + 1, "Plan had failures, replanning");
9214 plan = self.generate_plan(processed_input).await?;
9215 }
9216 StepFailureAction::Abort => {
9217 warn!("Plan step failed, aborting");
9218 break;
9219 }
9220 StepFailureAction::Skip | StepFailureAction::Continue => {
9221 break;
9222 }
9223 }
9224 }
9225
9226 *self.current_plan.write() = Some(plan);
9227
9228 let output_data = self.process_output(&plan_result, input_context).await?;
9229 let final_content = output_data.content;
9230
9231 self.memory
9232 .add_message(ChatMessage::assistant(&final_content))
9233 .await?;
9234
9235 self.check_memory_compression().await?;
9236 self.increment_turn();
9237 self.evaluate_transitions(processed_input, &final_content)
9238 .await?;
9239
9240 let reasoning_metadata =
9241 ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9242
9243 let response = AgentResponse::new(&final_content).with_metadata(
9244 "reasoning",
9245 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9246 );
9247
9248 self.finish_turn_if_root(&response).await?;
9249 Ok(response)
9250 }
9251
9252 fn inject_reasoning_prompt(
9254 &self,
9255 messages: &mut [ChatMessage],
9256 reasoning_mode: &ReasoningMode,
9257 is_first_iteration: bool,
9258 ) {
9259 if !is_first_iteration {
9260 return;
9261 }
9262 match reasoning_mode {
9263 ReasoningMode::CoT => {
9264 if let Some(msg) = messages.first_mut()
9265 && matches!(msg.role, ai_agents_core::Role::System)
9266 {
9267 msg.content = self.build_cot_system_prompt(&msg.content);
9268 debug!("Applied Chain-of-Thought system prompt");
9269 }
9270 }
9271 ReasoningMode::React => {
9272 if let Some(msg) = messages.first_mut()
9273 && matches!(msg.role, ai_agents_core::Role::System)
9274 {
9275 msg.content = self.build_react_system_prompt(&msg.content);
9276 debug!("Applied ReAct system prompt");
9277 }
9278 }
9279 _ => {}
9280 }
9281 }
9282
9283 async fn generate_main_response_draft(
9288 &self,
9289 processed_input: &str,
9290 reasoning_mode: &ReasoningMode,
9291 ) -> Result<MainResponseDraft> {
9292 let llm = self.get_state_llm()?;
9293 let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9294 let mut messages = self
9295 .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9296 .await?;
9297 self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9298 let response = self
9299 .complete_main_llm_with_recovery(llm, &messages, &protocol)
9300 .await?;
9301 let content = response.content.trim().to_string();
9302 let (thinking, answer) = self.extract_thinking(&content);
9303 if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9304 return Ok(MainResponseDraft::ToolCalls {
9305 raw_content: content,
9306 calls,
9307 thinking,
9308 });
9309 }
9310 Ok(MainResponseDraft::Text {
9311 raw_content: answer,
9312 thinking,
9313 })
9314 }
9315
9316 async fn commit_main_response_draft(
9321 &self,
9322 processed_input: &str,
9323 input_context: &HashMap<String, Value>,
9324 draft: MainResponseDraft,
9325 reasoning_mode: ReasoningMode,
9326 auto_detected: bool,
9327 ) -> Result<AgentResponse> {
9328 self.commit_root_user_message(processed_input).await?;
9329 match draft {
9330 MainResponseDraft::Text {
9331 raw_content,
9332 thinking,
9333 } => {
9334 self.finish_text_response_from_model(CommittedTextResponse {
9335 processed_input,
9336 input_context,
9337 answer: raw_content,
9338 reasoning_mode,
9339 auto_detected,
9340 iterations: 1,
9341 thinking_content: thinking,
9342 all_tool_calls: Vec::new(),
9343 })
9344 .await
9345 }
9346 MainResponseDraft::ToolCalls {
9347 raw_content,
9348 calls,
9349 thinking: _,
9350 } => {
9351 let mut all_tool_calls = Vec::new();
9352 match self
9353 .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9354 .await?
9355 {
9356 ToolCallOutcome::Rejected(response) => {
9357 self.finish_turn_if_root(&response).await?;
9358 Ok(response)
9359 }
9360 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9361 self.continue_after_committed_tool_draft(processed_input)
9362 .await
9363 }
9364 }
9365 }
9366 }
9367 }
9368
9369 async fn continue_after_committed_tool_draft(
9374 &self,
9375 processed_input: &str,
9376 ) -> Result<AgentResponse> {
9377 *self.redispatch_depth.write() += 1;
9378 if let Some(context) = self.active_turn_context.write().as_mut() {
9379 context.enter_redispatch();
9380 }
9381 let result = Box::pin(self.run_loop_internal(processed_input)).await;
9382 *self.redispatch_depth.write() -= 1;
9383 if let Some(context) = self.active_turn_context.write().as_mut() {
9384 context.exit_redispatch();
9385 }
9386 let response = result?;
9387 self.finish_turn_if_root(&response).await?;
9388 Ok(response)
9389 }
9390
9391 async fn finish_text_response_from_model(
9396 &self,
9397 response: CommittedTextResponse<'_>,
9398 ) -> Result<AgentResponse> {
9399 let CommittedTextResponse {
9400 processed_input,
9401 input_context,
9402 answer,
9403 reasoning_mode,
9404 auto_detected,
9405 iterations,
9406 thinking_content,
9407 all_tool_calls,
9408 } = response;
9409 let output_data = self.process_output(&answer, input_context).await?;
9410 let mut final_content = if output_data.metadata.rejected {
9411 output_data
9412 .metadata
9413 .rejection_reason
9414 .unwrap_or_else(|| answer.to_string())
9415 } else {
9416 output_data.content
9417 };
9418 let llm = self.get_state_llm()?;
9419 let reflection_metadata;
9420 (final_content, reflection_metadata) = self
9421 .run_reflection(&*llm, processed_input, final_content)
9422 .await?;
9423 final_content =
9424 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9425 let final_content = {
9426 let result = self
9427 .post_loop_processing(processed_input, final_content)
9428 .await?;
9429 self.apply_post_loop_result(processed_input, result).await?
9430 };
9431 let response = self.build_agent_response(AgentResponseParts {
9432 content: final_content,
9433 all_tool_calls,
9434 reasoning_mode,
9435 auto_detected,
9436 iterations,
9437 thinking: thinking_content,
9438 reflection_metadata,
9439 });
9440 self.finish_turn_if_root(&response).await?;
9441 Ok(response)
9442 }
9443
9444 async fn run_committed_response_loop_with_reasoning(
9449 &self,
9450 processed_input: &str,
9451 input_context: &HashMap<String, Value>,
9452 reasoning_mode: ReasoningMode,
9453 auto_detected: bool,
9454 ) -> Result<AgentResponse> {
9455 self.commit_root_user_message(processed_input).await?;
9456 let llm = self.get_state_llm()?;
9457 let mut iterations = 0u32;
9458 let mut all_tool_calls = Vec::new();
9459 let mut thinking_content = None;
9460 loop {
9461 let effective_max = if reasoning_mode != ReasoningMode::None {
9462 let rc = self.get_effective_reasoning_config();
9463 self.max_iterations.min(rc.max_iterations)
9464 } else {
9465 self.max_iterations
9466 };
9467 if iterations >= effective_max {
9468 return Err(AgentError::Other(format!(
9469 "Max iterations ({}) exceeded",
9470 effective_max
9471 )));
9472 }
9473 iterations += 1;
9474 *self.iteration_count.write() = iterations;
9475 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9476 let mut messages = self
9477 .build_messages_internal(true, None, protocol.choice.is_none())
9478 .await?;
9479 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9480 self.hooks.on_llm_start(&messages).await;
9481 let llm_start = Instant::now();
9482 let response = self
9483 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9484 .await?;
9485 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9486 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9487 let content = response.content.trim();
9488 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9489 match self
9490 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9491 .await?
9492 {
9493 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9494 ToolCallOutcome::Rejected(resp) => {
9495 self.finish_turn_if_root(&resp).await?;
9496 return Ok(resp);
9497 }
9498 }
9499 }
9500 let (extracted_thinking, answer) = self.extract_thinking(content);
9501 if extracted_thinking.is_some() {
9502 thinking_content = extracted_thinking;
9503 }
9504 return self
9505 .finish_text_response_from_model(CommittedTextResponse {
9506 processed_input,
9507 input_context,
9508 answer,
9509 reasoning_mode,
9510 auto_detected,
9511 iterations,
9512 thinking_content,
9513 all_tool_calls,
9514 })
9515 .await;
9516 }
9517 }
9518
9519 async fn handle_tool_calls(
9521 &self,
9522 processed_input: &str,
9523 content: &str,
9524 tool_calls: Vec<ToolCall>,
9525 all_tool_calls: &mut Vec<ToolCall>,
9526 ) -> Result<ToolCallOutcome> {
9527 let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9531 if transition_fired {
9532 self.memory
9533 .add_message(ChatMessage::assistant(
9534 "(Transitioned to new state — tool call handled by workflow)",
9535 ))
9536 .await?;
9537 return Ok(ToolCallOutcome::TransitionFired);
9538 }
9539
9540 self.memory
9542 .add_message(ChatMessage::assistant(content))
9543 .await?;
9544 let native_tool_call = Self::is_native_tool_call_content(content);
9545
9546 let results = self.execute_tools_parallel(&tool_calls).await;
9547
9548 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9549 match result {
9550 Ok(output) => {
9551 self.memory
9552 .add_message(Self::tool_result_message(
9553 tool_call,
9554 &output,
9555 native_tool_call,
9556 ))
9557 .await?;
9558 }
9559 Err(e) => {
9560 if matches!(e, AgentError::HITLRejected(_)) {
9562 self.memory
9563 .add_message(ChatMessage::assistant(format!(
9564 "The operation was rejected by the approver: {}",
9565 e
9566 )))
9567 .await?;
9568 return Ok(ToolCallOutcome::Rejected(AgentResponse {
9570 content: format!("Operation cancelled: {}", e),
9571 metadata: None,
9572 tool_calls: Some(all_tool_calls.clone()),
9573 }));
9574 }
9575 self.memory
9576 .add_message(Self::tool_result_message(
9577 tool_call,
9578 &format!("Error: {}", e),
9579 native_tool_call,
9580 ))
9581 .await?;
9582 }
9583 }
9584 all_tool_calls.push(tool_call.clone());
9585 }
9586 Ok(ToolCallOutcome::Continue)
9587 }
9588
9589 async fn run_reflection(
9591 &self,
9592 llm: &dyn LLMProvider,
9593 processed_input: &str,
9594 mut content: String,
9595 ) -> Result<(String, Option<ReflectionMetadata>)> {
9596 let should_reflect = self.should_reflect(processed_input, &content).await?;
9597 if !should_reflect {
9598 return Ok((content, None));
9599 }
9600
9601 info!("Starting response reflection evaluation");
9602 let mut attempts = 0u32;
9603 let max_retries = self.reflection_config.max_retries;
9604 let mut history: Vec<ReflectionAttempt> = Vec::new();
9605
9606 loop {
9607 let evaluation = self.evaluate_response(processed_input, &content).await?;
9608
9609 if evaluation.passed || attempts >= max_retries {
9610 info!(
9611 passed = evaluation.passed,
9612 confidence = evaluation.confidence,
9613 attempts = attempts + 1,
9614 "Reflection evaluation complete"
9615 );
9616 let reflection_metadata = Some(
9617 ReflectionMetadata::new(evaluation)
9618 .with_attempts(attempts + 1)
9619 .with_history(history),
9620 );
9621 return Ok((content, reflection_metadata));
9622 }
9623
9624 debug!(
9625 attempt = attempts + 1,
9626 failed_criteria = evaluation.failed_criteria().count(),
9627 "Response did not meet criteria, retrying"
9628 );
9629
9630 history.push(
9631 ReflectionAttempt::new(&content, evaluation.clone())
9632 .with_feedback("Response did not meet quality criteria"),
9633 );
9634
9635 let feedback: Vec<String> = evaluation
9636 .failed_criteria()
9637 .map(|c| format!("- {}", c.criterion))
9638 .collect();
9639
9640 let retry_prompt = format!(
9641 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9642 feedback.join("\n")
9643 );
9644
9645 self.memory
9646 .add_message(ChatMessage::user(&retry_prompt))
9647 .await?;
9648
9649 let retry_messages = self.build_messages().await?;
9650 let retry_response = self
9651 .observe_purpose(
9652 ObservationPurpose::ReflectionEvaluation,
9653 llm.complete(&retry_messages, None),
9654 )
9655 .await
9656 .map_err(|e| AgentError::LLM(e.to_string()))?;
9657
9658 content = retry_response.content.trim().to_string();
9659 attempts += 1;
9660 }
9661 }
9662
9663 async fn post_loop_processing(
9666 &self,
9667 processed_input: &str,
9668 content: String,
9669 ) -> Result<PostLoopResult> {
9670 self.increment_turn();
9675
9676 self.run_context_extractors(processed_input).await;
9678
9679 let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9680
9681 if !transitioned {
9682 self.memory
9683 .add_message(ChatMessage::assistant(&content))
9684 .await?;
9685 self.check_memory_compression().await?;
9686 return Ok(PostLoopResult::NoTransition(content));
9687 }
9688
9689 if !self.should_regenerate_after_transition() {
9691 self.memory
9692 .add_message(ChatMessage::assistant(&content))
9693 .await?;
9694 self.check_memory_compression().await?;
9695 return Ok(PostLoopResult::Transitioned(content));
9696 }
9697
9698 if self.needs_redispatch_for_new_state() {
9702 info!("Post-transition NeedsRedispatch: new state requires full dispatch");
9703 return Ok(PostLoopResult::NeedsRedispatch);
9706 }
9707
9708 self.memory
9711 .add_message(ChatMessage::assistant(&content))
9712 .await?;
9713 self.check_memory_compression().await?;
9714
9715 let new_llm = self.get_state_llm()?;
9721 let mut final_content;
9722
9723 for post_iter in 0..self.max_iterations {
9724 let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
9725 let new_messages = self
9726 .build_messages_internal(true, None, protocol.choice.is_none())
9727 .await?;
9728 if post_iter == 0
9729 && let Some(system_msg) = new_messages.first()
9730 && system_msg.role == ai_agents_core::Role::System
9731 {
9732 debug!(
9733 prompt_preview =
9734 &system_msg.content[system_msg.content.len().saturating_sub(200)..],
9735 "Post-transition system prompt (last 200 chars)"
9736 );
9737 }
9738
9739 let new_response = self
9740 .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
9741 .await?;
9742 final_content = new_response.content.trim().to_string();
9743
9744 if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
9747 let native_tool_call = Self::is_native_tool_call_content(&final_content);
9748 debug!(
9749 post_iter = post_iter,
9750 tools = tool_calls.len(),
9751 "Post-transition tool call detected, executing"
9752 );
9753
9754 self.memory
9755 .add_message(ChatMessage::assistant(&final_content))
9756 .await?;
9757
9758 let results = self.execute_tools_parallel(&tool_calls).await;
9759 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9760 match result {
9761 Ok(output) => {
9762 self.memory
9763 .add_message(Self::tool_result_message(
9764 tool_call,
9765 &output,
9766 native_tool_call,
9767 ))
9768 .await?;
9769 }
9770 Err(e) => {
9771 self.memory
9772 .add_message(Self::tool_result_message(
9773 tool_call,
9774 &format!("Error: {}", e),
9775 native_tool_call,
9776 ))
9777 .await?;
9778 }
9779 }
9780 }
9781 continue;
9783 }
9784
9785 self.memory
9787 .add_message(ChatMessage::assistant(&final_content))
9788 .await?;
9789 return Ok(PostLoopResult::Transitioned(final_content));
9790 }
9791
9792 final_content = "Post-transition processing completed.".to_string();
9794 self.memory
9795 .add_message(ChatMessage::assistant(&final_content))
9796 .await?;
9797
9798 Ok(PostLoopResult::Transitioned(final_content))
9799 }
9800
9801 fn should_regenerate_after_transition(&self) -> bool {
9804 if let Some(ref sm) = self.state_machine {
9805 if !sm.config().regenerate_on_transition {
9807 return false;
9808 }
9809 if let Some(def) = sm.current_definition()
9811 && let Some(regen) = def.regenerate_on_enter
9812 {
9813 return regen;
9814 }
9815 }
9816 true
9817 }
9818
9819 fn needs_redispatch_for_new_state(&self) -> bool {
9822 if let Some(ref sm) = self.state_machine
9823 && let Some(def) = sm.current_definition()
9824 {
9825 if def.concurrent.is_some()
9826 || def.group_chat.is_some()
9827 || def.pipeline.is_some()
9828 || def.handoff.is_some()
9829 || def.delegate.is_some()
9830 {
9831 return true;
9832 }
9833 let effective = self.get_effective_reasoning_config();
9835 if !matches!(effective.mode, ReasoningMode::None) {
9836 return true;
9837 }
9838 }
9839 false
9840 }
9841
9842 async fn apply_post_loop_result(
9845 &self,
9846 processed_input: &str,
9847 result: PostLoopResult,
9848 ) -> Result<String> {
9849 match result {
9850 PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
9851 Ok(content)
9852 }
9853 PostLoopResult::NeedsRedispatch => {
9854 const MAX_REDISPATCH_DEPTH: u32 = 3;
9855 let current_depth = *self.redispatch_depth.read();
9856 if current_depth >= MAX_REDISPATCH_DEPTH {
9857 warn!(
9858 depth = current_depth,
9859 "Post-transition re-dispatch depth limit reached, returning empty response"
9860 );
9861 let content = String::new();
9862 self.memory
9863 .add_message(ChatMessage::assistant(&content))
9864 .await?;
9865 return Ok(content);
9866 }
9867 *self.redispatch_depth.write() += 1;
9868 if let Some(context) = self.active_turn_context.write().as_mut() {
9869 context.enter_redispatch();
9870 }
9871 info!(
9872 depth = current_depth + 1,
9873 "Re-dispatching for new state after transition"
9874 );
9875 let resp = Box::pin(self.run_loop_internal(processed_input)).await;
9876 *self.redispatch_depth.write() -= 1;
9877 if let Some(context) = self.active_turn_context.write().as_mut() {
9878 context.exit_redispatch();
9879 }
9880 resp.map(|r| r.content)
9881 }
9882 }
9883 }
9884
9885 fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
9887 let AgentResponseParts {
9888 content,
9889 all_tool_calls,
9890 reasoning_mode,
9891 auto_detected,
9892 iterations,
9893 thinking,
9894 reflection_metadata,
9895 } = parts;
9896 let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
9897 .with_thinking(thinking.clone().unwrap_or_default())
9898 .with_iterations(iterations)
9899 .with_auto_detected(auto_detected);
9900
9901 let mut response = AgentResponse::new(&content);
9902 if !all_tool_calls.is_empty() {
9903 response = response.with_tool_calls(all_tool_calls);
9904 }
9905
9906 if let Some(state) = self.current_state() {
9907 response = response.with_metadata("current_state", serde_json::json!(state));
9908 }
9909
9910 response = response.with_metadata(
9911 "reasoning",
9912 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9913 );
9914
9915 if let Some(ref refl_meta) = reflection_metadata {
9916 response = response.with_metadata(
9917 "reflection",
9918 serde_json::to_value(refl_meta).unwrap_or_default(),
9919 );
9920 }
9921
9922 response
9923 }
9924
9925 async fn handle_delegated_state(
9927 &self,
9928 input: &str,
9929 delegate_id: &str,
9930 state_def: &ai_agents_state::StateDefinition,
9931 ) -> Result<AgentResponse> {
9932 use std::time::Instant;
9933
9934 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
9935 AgentError::Config(format!(
9936 "State delegates to '{}' but no agent registry is configured. \
9937 Add a spawner section with auto_spawn to your YAML.",
9938 delegate_id
9939 ))
9940 })?;
9941
9942 let state_name = self
9943 .state_machine
9944 .as_ref()
9945 .map(|sm| sm.current())
9946 .unwrap_or_else(|| "unknown".to_string());
9947
9948 self.hooks.on_delegate_start(delegate_id, &state_name).await;
9949 let start = Instant::now();
9950
9951 let delegate = registry.get(delegate_id).ok_or_else(|| {
9952 AgentError::Other(format!(
9953 "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
9954 state_name, delegate_id
9955 ))
9956 })?;
9957
9958 let context_mode = state_def.delegate_context.clone().unwrap_or_default();
9960 let effective_input = self
9961 .observe_purpose(
9962 ObservationPurpose::OrchestrationRouting,
9963 crate::orchestration::context::prepare_delegate_input(
9964 input,
9965 &context_mode,
9966 &*self.memory,
9967 self.llm_registry.get("router").ok().as_deref(),
9968 ),
9969 )
9970 .await?;
9971
9972 let response = delegate
9973 .chat_with_actor_context(&effective_input, self.outbound_actor_context())
9974 .await?;
9975
9976 let duration_ms = start.elapsed().as_millis() as u64;
9977 self.hooks
9978 .on_delegate_complete(delegate_id, &state_name, duration_ms)
9979 .await;
9980
9981 let ctx_key = format!("delegation.{}.last_response", delegate_id);
9983 let _ = self.context_manager.set(
9984 &ctx_key,
9985 serde_json::Value::String(response.content.clone()),
9986 );
9987
9988 let _ = self.context_manager.set(
9990 "orchestration",
9991 serde_json::json!({
9992 "type": "delegate",
9993 "agent": delegate_id,
9994 "state": state_name,
9995 "response": response.content,
9996 "duration_ms": duration_ms,
9997 }),
9998 );
9999
10000 self.commit_root_user_message(input).await?;
10001
10002 let post_result = self
10005 .post_loop_processing(
10006 input,
10007 format!("[Delegated to {}]: {}", delegate_id, response.content),
10008 )
10009 .await?;
10010 let final_content = self.apply_post_loop_result(input, post_result).await?;
10011
10012 let mut result = AgentResponse::new(final_content);
10013
10014 let metadata = serde_json::json!({
10015 "orchestration": {
10016 "type": "delegate",
10017 "agent": delegate_id,
10018 "state": state_name,
10019 "response": response.content,
10020 "duration_ms": duration_ms,
10021 }
10022 });
10023 result.metadata = Some(
10024 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10025 metadata,
10026 )
10027 .unwrap_or_default(),
10028 );
10029
10030 self.finish_turn_if_root(&result).await?;
10031 Ok(result)
10032 }
10033
10034 async fn handle_concurrent_state(
10036 &self,
10037 input: &str,
10038 config: &ai_agents_state::ConcurrentStateConfig,
10039 ) -> Result<AgentResponse> {
10040 use std::time::Instant;
10041
10042 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10043 AgentError::Config(
10044 "Concurrent state requires an agent registry. Add a spawner section.".into(),
10045 )
10046 })?;
10047
10048 let context_mode = config.context_mode.clone().unwrap_or_default();
10053 let context_input = self
10054 .observe_purpose(
10055 ObservationPurpose::OrchestrationRouting,
10056 crate::orchestration::context::prepare_delegate_input(
10057 input,
10058 &context_mode,
10059 &*self.memory,
10060 self.llm_registry.get("router").ok().as_deref(),
10061 ),
10062 )
10063 .await?;
10064
10065 let effective_input = if let Some(ref tmpl) = config.input {
10066 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10067 .unwrap_or_else(|_| context_input.clone())
10068 } else {
10069 context_input
10070 };
10071
10072 let start = Instant::now();
10073
10074 let llm_name = config
10075 .aggregation
10076 .synthesizer_llm
10077 .as_deref()
10078 .unwrap_or("router");
10079 let llm_provider = self.llm_registry.get(llm_name).ok();
10080
10081 let vote_parallelism = if self.runtime_config.optimization.enabled
10082 && self
10083 .runtime_config
10084 .optimization
10085 .parallel_orchestration_vote_extraction
10086 {
10087 Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10088 } else {
10089 None
10090 };
10091
10092 let result = self
10093 .observe_purpose(
10094 ObservationPurpose::OrchestrationAggregation,
10095 scope_actor_context(
10096 self.outbound_actor_context(),
10097 crate::orchestration::concurrent(
10098 registry,
10099 &effective_input,
10100 &config.agents,
10101 &config.aggregation,
10102 llm_provider.as_deref(),
10103 config.min_required,
10104 config.timeout_ms,
10105 config.on_partial_failure.clone(),
10106 vote_parallelism,
10107 ),
10108 ),
10109 )
10110 .await?;
10111
10112 let duration_ms = start.elapsed().as_millis() as u64;
10113 let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10114 let strategy = format!("{:?}", config.aggregation.strategy);
10115 self.hooks
10116 .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10117 .await;
10118
10119 let _ = self.context_manager.set(
10121 "concurrent.result",
10122 serde_json::Value::String(result.response.content.clone()),
10123 );
10124
10125 let agents_json: Vec<serde_json::Value> = result
10127 .agent_results
10128 .iter()
10129 .map(|ar| {
10130 serde_json::json!({
10131 "id": ar.agent_id,
10132 "response": ar.response.as_ref().map(|r| r.content.as_str()),
10133 "success": ar.success,
10134 "error": ar.error,
10135 "duration_ms": ar.duration_ms,
10136 })
10137 })
10138 .collect();
10139
10140 let _ = self.context_manager.set(
10142 "orchestration",
10143 serde_json::json!({
10144 "type": "concurrent",
10145 "result": result.response.content,
10146 "strategy": strategy,
10147 "agents": agents_json,
10148 "duration_ms": duration_ms,
10149 }),
10150 );
10151
10152 self.commit_root_user_message(input).await?;
10153
10154 let post_result = self
10155 .post_loop_processing(input, result.response.content.clone())
10156 .await?;
10157 let final_content = self.apply_post_loop_result(input, post_result).await?;
10158
10159 let mut response = AgentResponse::new(final_content);
10160 let metadata = serde_json::json!({
10161 "orchestration": {
10162 "type": "concurrent",
10163 "result": result.response.content,
10164 "strategy": strategy,
10165 "agents": agents_json,
10166 "duration_ms": duration_ms,
10167 }
10168 });
10169 response.metadata = Some(
10170 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10171 metadata,
10172 )
10173 .unwrap_or_default(),
10174 );
10175
10176 self.finish_turn_if_root(&response).await?;
10177 Ok(response)
10178 }
10179
10180 async fn handle_group_chat_state(
10182 &self,
10183 input: &str,
10184 config: &ai_agents_state::GroupChatStateConfig,
10185 ) -> Result<AgentResponse> {
10186 use std::time::Instant;
10187
10188 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10189 AgentError::Config(
10190 "Group chat state requires an agent registry. Add a spawner section.".into(),
10191 )
10192 })?;
10193
10194 let start = Instant::now();
10195
10196 let llm_provider = self.llm_registry.get("router").ok();
10197
10198 let context_mode = config.context_mode.clone().unwrap_or_default();
10200 let context_input = self
10201 .observe_purpose(
10202 ObservationPurpose::OrchestrationRouting,
10203 crate::orchestration::context::prepare_delegate_input(
10204 input,
10205 &context_mode,
10206 &*self.memory,
10207 self.llm_registry.get("router").ok().as_deref(),
10208 ),
10209 )
10210 .await?;
10211
10212 let effective_topic = if let Some(ref tmpl) = config.input {
10214 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10215 .unwrap_or_else(|_| context_input.clone())
10216 } else {
10217 context_input
10218 };
10219
10220 let result = self
10221 .observe_purpose(
10222 ObservationPurpose::OrchestrationConversation,
10223 scope_actor_context(
10224 self.outbound_actor_context(),
10225 crate::orchestration::group_chat(
10226 registry,
10227 &effective_topic,
10228 config,
10229 llm_provider.as_deref(),
10230 Some(&*self.hooks),
10231 ),
10232 ),
10233 )
10234 .await?;
10235
10236 let duration_ms = start.elapsed().as_millis() as u64;
10237
10238 let _ = self.context_manager.set(
10240 "group_chat.conclusion",
10241 serde_json::Value::String(result.response.content.clone()),
10242 );
10243
10244 let transcript_json: Vec<serde_json::Value> = result
10246 .transcript
10247 .iter()
10248 .map(|t| {
10249 serde_json::json!({
10250 "speaker": t.speaker,
10251 "round": t.round,
10252 "content": t.content,
10253 })
10254 })
10255 .collect();
10256
10257 let _ = self.context_manager.set(
10259 "orchestration",
10260 serde_json::json!({
10261 "type": "group_chat",
10262 "conclusion": result.response.content,
10263 "transcript": transcript_json,
10264 "rounds": result.rounds_completed,
10265 "termination": result.termination_reason,
10266 "duration_ms": duration_ms,
10267 }),
10268 );
10269
10270 self.commit_root_user_message(input).await?;
10271
10272 let post_result = self
10273 .post_loop_processing(input, result.response.content.clone())
10274 .await?;
10275 let final_content = self.apply_post_loop_result(input, post_result).await?;
10276
10277 let mut response = AgentResponse::new(final_content);
10278 let metadata = serde_json::json!({
10279 "orchestration": {
10280 "type": "group_chat",
10281 "conclusion": result.response.content,
10282 "transcript": transcript_json,
10283 "rounds": result.rounds_completed,
10284 "termination": result.termination_reason,
10285 "duration_ms": duration_ms,
10286 }
10287 });
10288 response.metadata = Some(
10289 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10290 metadata,
10291 )
10292 .unwrap_or_default(),
10293 );
10294
10295 self.finish_turn_if_root(&response).await?;
10296 Ok(response)
10297 }
10298
10299 async fn handle_pipeline_state(
10301 &self,
10302 input: &str,
10303 config: &ai_agents_state::PipelineStateConfig,
10304 ) -> Result<AgentResponse> {
10305 use std::time::Instant;
10306
10307 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10308 AgentError::Config(
10309 "Pipeline state requires an agent registry. Add a spawner section.".into(),
10310 )
10311 })?;
10312
10313 let start = Instant::now();
10314
10315 let stages: Vec<crate::orchestration::PipelineStage> = config
10316 .stages
10317 .iter()
10318 .map(|entry| {
10319 let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10320 if let Some(tmpl) = entry.input() {
10321 stage = stage.with_input(tmpl);
10322 }
10323 stage
10324 })
10325 .collect();
10326
10327 let context_mode = config.context_mode.clone().unwrap_or_default();
10329 let context_input = self
10330 .observe_purpose(
10331 ObservationPurpose::OrchestrationRouting,
10332 crate::orchestration::context::prepare_delegate_input(
10333 input,
10334 &context_mode,
10335 &*self.memory,
10336 self.llm_registry.get("router").ok().as_deref(),
10337 ),
10338 )
10339 .await?;
10340
10341 let context_values = self.build_context_with_overlays();
10342 let result = self
10343 .observe_purpose(
10344 ObservationPurpose::OrchestrationRouting,
10345 scope_actor_context(
10346 self.outbound_actor_context(),
10347 crate::orchestration::pipeline(
10348 registry,
10349 &context_input,
10350 &stages,
10351 config.timeout_ms,
10352 Some(&*self.hooks),
10353 Some(&context_values),
10354 ),
10355 ),
10356 )
10357 .await?;
10358
10359 let duration_ms = start.elapsed().as_millis() as u64;
10360
10361 let _ = self.context_manager.set(
10363 "pipeline.result",
10364 serde_json::Value::String(result.response.content.clone()),
10365 );
10366
10367 let stages_json: Vec<serde_json::Value> = result
10369 .stage_outputs
10370 .iter()
10371 .map(|s| {
10372 serde_json::json!({
10373 "agent_id": s.agent_id,
10374 "output": s.output,
10375 "duration_ms": s.duration_ms,
10376 "skipped": s.skipped,
10377 })
10378 })
10379 .collect();
10380
10381 let _ = self.context_manager.set(
10383 "orchestration",
10384 serde_json::json!({
10385 "type": "pipeline",
10386 "result": result.response.content,
10387 "stages": stages_json,
10388 "duration_ms": duration_ms,
10389 }),
10390 );
10391
10392 self.commit_root_user_message(input).await?;
10393
10394 let post_result = self
10395 .post_loop_processing(input, result.response.content.clone())
10396 .await?;
10397 let final_content = self.apply_post_loop_result(input, post_result).await?;
10398
10399 let mut response = AgentResponse::new(final_content);
10400 let metadata = serde_json::json!({
10401 "orchestration": {
10402 "type": "pipeline",
10403 "result": result.response.content,
10404 "stages": stages_json,
10405 "duration_ms": duration_ms,
10406 }
10407 });
10408 response.metadata = Some(
10409 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10410 metadata,
10411 )
10412 .unwrap_or_default(),
10413 );
10414
10415 self.finish_turn_if_root(&response).await?;
10416 Ok(response)
10417 }
10418
10419 async fn handle_handoff_state(
10421 &self,
10422 input: &str,
10423 config: &ai_agents_state::HandoffStateConfig,
10424 ) -> Result<AgentResponse> {
10425 use std::time::Instant;
10426
10427 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10428 AgentError::Config(
10429 "Handoff state requires an agent registry. Add a spawner section.".into(),
10430 )
10431 })?;
10432
10433 let llm = self
10434 .llm_registry
10435 .get("router")
10436 .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10437
10438 let start = Instant::now();
10439
10440 let context_mode = config.context_mode.clone().unwrap_or_default();
10442 let context_input = self
10443 .observe_purpose(
10444 ObservationPurpose::OrchestrationRouting,
10445 crate::orchestration::context::prepare_delegate_input(
10446 input,
10447 &context_mode,
10448 &*self.memory,
10449 self.llm_registry.get("router").ok().as_deref(),
10450 ),
10451 )
10452 .await?;
10453
10454 let effective_input = if let Some(ref tmpl) = config.input {
10456 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10457 .unwrap_or_else(|_| context_input.clone())
10458 } else {
10459 context_input
10460 };
10461
10462 let result = self
10463 .observe_purpose(
10464 ObservationPurpose::OrchestrationRouting,
10465 scope_actor_context(
10466 self.outbound_actor_context(),
10467 crate::orchestration::handoff(
10468 registry,
10469 &effective_input,
10470 &config.initial_agent,
10471 &config.available_agents,
10472 config.max_handoffs,
10473 llm.as_ref(),
10474 Some(&*self.hooks),
10475 ),
10476 ),
10477 )
10478 .await?;
10479
10480 let duration_ms = start.elapsed().as_millis() as u64;
10481
10482 let _ = self.context_manager.set(
10484 "handoff.result",
10485 serde_json::Value::String(result.response.content.clone()),
10486 );
10487
10488 let chain_json: Vec<serde_json::Value> = result
10490 .handoff_chain
10491 .iter()
10492 .map(|h| {
10493 serde_json::json!({
10494 "from": h.from_agent,
10495 "to": h.to_agent,
10496 "reason": h.reason,
10497 })
10498 })
10499 .collect();
10500
10501 let _ = self.context_manager.set(
10503 "orchestration",
10504 serde_json::json!({
10505 "type": "handoff",
10506 "result": result.response.content,
10507 "final_agent": result.final_agent,
10508 "handoff_chain": chain_json,
10509 "duration_ms": duration_ms,
10510 }),
10511 );
10512
10513 self.commit_root_user_message(input).await?;
10514
10515 let post_result = self
10516 .post_loop_processing(input, result.response.content.clone())
10517 .await?;
10518 let final_content = self.apply_post_loop_result(input, post_result).await?;
10519
10520 let mut response = AgentResponse::new(final_content);
10521 let metadata = serde_json::json!({
10522 "orchestration": {
10523 "type": "handoff",
10524 "result": result.response.content,
10525 "final_agent": result.final_agent,
10526 "handoff_chain": chain_json,
10527 "duration_ms": duration_ms,
10528 }
10529 });
10530 response.metadata = Some(
10531 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10532 metadata,
10533 )
10534 .unwrap_or_default(),
10535 );
10536
10537 self.finish_turn_if_root(&response).await?;
10538 Ok(response)
10539 }
10540
10541 async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10543 self.begin_root_turn();
10544 self.pre_turn_session_lifecycle().await;
10546
10547 let input_data = self.process_input(input).await?;
10548 self.update_active_turn_context(&input_data.content, input_data.context.clone());
10549
10550 for (key, value) in &input_data.context {
10553 let _ = self.context_manager.set(key, value.clone());
10554 }
10555
10556 if input_data.metadata.rejected {
10557 let reason = input_data
10558 .metadata
10559 .rejection_reason
10560 .unwrap_or_else(|| "Input rejected".to_string());
10561 warn!(reason = %reason, "Input rejected");
10562 let response = AgentResponse::new(reason);
10563 self.finish_turn_if_root(&response).await?;
10564 return Ok(response);
10565 }
10566
10567 let processed_input = &input_data.content;
10568
10569 if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10570 return Ok(response);
10571 }
10572
10573 if let Some(ref sm) = self.state_machine
10575 && let Some(def) = sm.current_definition()
10576 {
10577 if let Some(ref delegate_id) = def.delegate {
10578 return self
10579 .handle_delegated_state(processed_input, delegate_id, &def)
10580 .await;
10581 }
10582 if let Some(ref concurrent_config) = def.concurrent {
10583 return self
10584 .handle_concurrent_state(processed_input, concurrent_config)
10585 .await;
10586 }
10587 if let Some(ref group_chat_config) = def.group_chat {
10588 return self
10589 .handle_group_chat_state(processed_input, group_chat_config)
10590 .await;
10591 }
10592 if let Some(ref pipeline_config) = def.pipeline {
10593 return self
10594 .handle_pipeline_state(processed_input, pipeline_config)
10595 .await;
10596 }
10597 if let Some(ref handoff_config) = def.handoff {
10598 return self
10599 .handle_handoff_state(processed_input, handoff_config)
10600 .await;
10601 }
10602 }
10603
10604 if let Some(response) =
10609 Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10610 {
10611 return Ok(response);
10612 }
10613
10614 match self.try_skill_route(processed_input).await? {
10615 SkillRouteResult::Response { skill_id, content } => {
10616 self.commit_root_user_message(processed_input).await?;
10617 return self
10618 .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10619 .await;
10620 }
10621 SkillRouteResult::NeedsClarification {
10622 response,
10623 ownership,
10624 } => {
10625 let admission = self
10626 .admit_optional_disambiguation_ownership(ownership)
10627 .await?;
10628 self.commit_root_user_message(processed_input).await?;
10629 if let Some(q) = response
10630 .metadata
10631 .as_ref()
10632 .and_then(|m| m.get("disambiguation"))
10633 .and_then(|d| d.get("status"))
10634 .and_then(|s| s.as_str())
10635 && q == "awaiting_clarification"
10636 {
10637 self.memory
10640 .add_message(ChatMessage::assistant(&response.content))
10641 .await?;
10642 }
10643 drop(admission);
10644 self.finish_turn_if_root(&response).await?;
10645 return Ok(response);
10646 }
10647 SkillRouteResult::NoMatch => {} }
10649
10650 let effective_reasoning = self.get_effective_reasoning_config();
10651 let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10652 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10653
10654 info!(
10655 reasoning_mode = ?reasoning_mode,
10656 auto_detected = auto_detected,
10657 reflection_enabled = ?self.reflection_config.enabled,
10658 "Reasoning mode determined"
10659 );
10660
10661 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10662 self.commit_root_user_message(processed_input).await?;
10663 return self
10664 .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10665 .await;
10666 }
10667
10668 self.commit_root_user_message(processed_input).await?;
10669
10670 let mut iterations = 0u32;
10671 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10672 let mut thinking_content: Option<String> = None;
10673
10674 let llm = self.get_state_llm()?;
10675
10676 loop {
10677 let effective_max = if reasoning_mode != ReasoningMode::None {
10679 let rc = self.get_effective_reasoning_config();
10680 self.max_iterations.min(rc.max_iterations)
10681 } else {
10682 self.max_iterations
10683 };
10684
10685 if iterations >= effective_max {
10686 let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10687 self.hooks.on_error(&err).await;
10688 error!(iterations = iterations, "Max iterations exceeded");
10689 return Err(err);
10690 }
10691 iterations += 1;
10692 *self.iteration_count.write() = iterations;
10693
10694 debug!(iteration = iterations, max = effective_max, "LLM call");
10695
10696 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
10697 let mut messages = self
10698 .build_messages_internal(true, None, protocol.choice.is_none())
10699 .await?;
10700 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
10701
10702 self.hooks.on_llm_start(&messages).await;
10703 let llm_start = Instant::now();
10704 let response = self
10705 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
10706 .await?;
10707
10708 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
10709 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
10710
10711 let content = response.content.trim();
10712
10713 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
10714 match self
10715 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
10716 .await?
10717 {
10718 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
10719 ToolCallOutcome::Rejected(resp) => {
10720 self.finish_turn_if_root(&resp).await?;
10721 return Ok(resp);
10722 }
10723 }
10724 }
10725
10726 let (extracted_thinking, answer) = self.extract_thinking(content);
10727 if extracted_thinking.is_some() {
10728 thinking_content = extracted_thinking;
10729 }
10730
10731 let output_data = self.process_output(&answer, &input_data.context).await?;
10732
10733 let mut final_content = if output_data.metadata.rejected {
10734 output_data
10735 .metadata
10736 .rejection_reason
10737 .unwrap_or_else(|| answer.to_string())
10738 } else {
10739 output_data.content
10740 };
10741
10742 let reflection_metadata;
10744 (final_content, reflection_metadata) = self
10745 .run_reflection(&*llm, processed_input, final_content)
10746 .await?;
10747
10748 final_content =
10749 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
10750
10751 let final_content = {
10755 let result = self
10756 .post_loop_processing(processed_input, final_content)
10757 .await?;
10758 self.apply_post_loop_result(processed_input, result).await?
10759 };
10760
10761 let reflected = reflection_metadata.is_some();
10762 let reasoning_mode_debug = format!("{:?}", reasoning_mode);
10763
10764 let response = self.build_agent_response(AgentResponseParts {
10765 content: final_content,
10766 all_tool_calls,
10767 reasoning_mode,
10768 auto_detected,
10769 iterations,
10770 thinking: thinking_content,
10771 reflection_metadata,
10772 });
10773
10774 self.finish_turn_if_root(&response).await?;
10775
10776 let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
10777 info!(
10778 tool_calls = tool_call_count,
10779 response_len = response.content.len(),
10780 reasoning_mode = %reasoning_mode_debug,
10781 reflected = reflected,
10782 "Chat completed"
10783 );
10784 return Ok(response);
10785 }
10786 }
10787
10788 async fn generate_buffered_streaming_draft(
10789 &self,
10790 processed_input: &str,
10791 routing_resolved: Arc<AtomicBool>,
10792 ) -> Result<StreamingDraftResult> {
10793 let llm = self.get_state_llm()?;
10794 if llm.configured_tool_choice().is_some() {
10795 let draft = self
10796 .generate_main_response_draft(processed_input, &ReasoningMode::None)
10797 .await?;
10798 return Ok(StreamingDraftResult::new(draft, Vec::new()));
10799 }
10800 let messages = self.build_messages_for_draft(processed_input).await?;
10801 let mut stream = self
10802 .observe_purpose(
10803 ObservationPurpose::MainResponse,
10804 llm.complete_stream(&messages, None),
10805 )
10806 .await
10807 .map_err(|e| AgentError::LLM(e.to_string()))?;
10808 let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
10809 let mut chunks = Vec::new();
10810 let mut accumulated = String::new();
10811 while let Some(chunk_result) = stream.next().await {
10812 let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
10813 accumulated.push_str(&chunk.delta);
10814 let stream_chunk = StreamChunk::content(chunk.delta);
10815 if routing_resolved.load(Ordering::SeqCst) {
10816 chunks.push(stream_chunk);
10817 } else {
10818 buffer.push(stream_chunk)?;
10819 }
10820 }
10821 chunks.splice(0..0, buffer.drain());
10822 let content = accumulated.trim().to_string();
10823 let draft = if let Some(calls) = self.parse_tool_calls(&content) {
10824 MainResponseDraft::ToolCalls {
10825 raw_content: content,
10826 calls,
10827 thinking: None,
10828 }
10829 } else {
10830 MainResponseDraft::Text {
10831 raw_content: content,
10832 thinking: None,
10833 }
10834 };
10835 Ok(StreamingDraftResult::new(draft, chunks))
10836 }
10837
10838 async fn try_buffered_streaming_branches(
10839 &self,
10840 processed_input: &str,
10841 input_context: &HashMap<String, Value>,
10842 ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
10843 let optimization = &self.runtime_config.optimization;
10844 if !optimization.enabled {
10845 return Ok(None);
10846 }
10847 let transition_enabled =
10848 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
10849 if !transition_enabled {
10850 return Ok(None);
10851 }
10852 let mut branch_scheduler =
10853 TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
10854 if !branch_scheduler.reserve_task() {
10855 return Ok(None);
10856 }
10857 if !self
10858 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
10859 {
10860 branch_scheduler.release_task();
10861 return Ok(None);
10862 }
10863 if !branch_scheduler.reserve_task() {
10864 branch_scheduler.release_task();
10865 return Ok(None);
10866 }
10867 let mut main_branch = RuntimeBranch::new(
10868 RuntimeTaskPurpose::MainResponse,
10869 RuntimeOptimizationKind::BufferedStreamingRouting,
10870 RuntimeTaskPriority::Normal,
10871 RuntimeCommitBehavior::FinalResponse,
10872 );
10873 let mut transition_branch = RuntimeBranch::new(
10874 RuntimeTaskPurpose::StateTransition,
10875 RuntimeOptimizationKind::ParallelStateTransition,
10876 RuntimeTaskPriority::Critical,
10877 RuntimeCommitBehavior::TransitionDecision,
10878 );
10879 let main_id = main_branch.branch_id();
10880 let transition_id = transition_branch.branch_id();
10881 let routing_resolved = Arc::new(AtomicBool::new(false));
10882 let mut main_future =
10883 Box::pin(crate::optimization::observability::with_branch_observation(
10884 &main_id,
10885 RuntimeOptimizationKind::BufferedStreamingRouting,
10886 RuntimeCommitBehavior::FinalResponse,
10887 self.generate_buffered_streaming_draft(
10888 processed_input,
10889 Arc::clone(&routing_resolved),
10890 ),
10891 ));
10892 let mut transition_future =
10893 Box::pin(crate::optimization::observability::with_branch_observation(
10894 &transition_id,
10895 RuntimeOptimizationKind::ParallelStateTransition,
10896 RuntimeCommitBehavior::TransitionDecision,
10897 self.select_parallel_transition_candidate(processed_input),
10898 ));
10899 let mut main_pending = true;
10900 let mut transition_pending = true;
10901 let mut main_result: Option<Result<StreamingDraftResult>> = None;
10902 let mut transition_finalized = false;
10903 let mut transition_candidate: Option<TransitionCandidate> = None;
10904 loop {
10905 if let Some(candidate) = transition_candidate.take() {
10906 if self
10907 .approve_transition_target(&candidate.from_state, candidate.target())
10908 .await?
10909 {
10910 drop(main_future);
10912 drop(transition_future);
10913 self.finalize_branch_loss(
10914 &main_id,
10915 RuntimeOptimizationKind::BufferedStreamingRouting,
10916 RuntimeCommitBehavior::FinalResponse,
10917 main_pending,
10918 main_result.as_ref().map(|result| result.is_err()),
10919 );
10920 if !self
10921 .apply_pre_response_transition_candidate(
10922 &candidate,
10923 &HashMap::new(),
10924 processed_input,
10925 )
10926 .await?
10927 {
10928 self.finalize_optional_branch(
10929 &transition_id,
10930 RuntimeOptimizationKind::ParallelStateTransition,
10931 RuntimeCommitBehavior::TransitionDecision,
10932 "discarded",
10933 false,
10934 );
10935 return Ok(None);
10936 }
10937 self.finalize_optional_branch(
10938 &transition_id,
10939 RuntimeOptimizationKind::ParallelStateTransition,
10940 RuntimeCommitBehavior::TransitionDecision,
10941 "committed",
10942 true,
10943 );
10944 let response = self.redispatch_current_state(processed_input).await?;
10945 return Ok(Some((
10946 response.clone(),
10947 vec![StreamChunk::content(response.content)],
10948 )));
10949 }
10950 self.finalize_optional_branch(
10951 &transition_id,
10952 RuntimeOptimizationKind::ParallelStateTransition,
10953 RuntimeCommitBehavior::TransitionDecision,
10954 "discarded",
10955 false,
10956 );
10957 routing_resolved.store(true, Ordering::SeqCst);
10958 transition_finalized = true;
10959 }
10960 if transition_finalized && let Some(result) = main_result.take() {
10961 let stream_draft = match result {
10962 Ok(stream_draft) => stream_draft,
10963 Err(error) => {
10964 self.finalize_optional_branch(
10965 &main_id,
10966 RuntimeOptimizationKind::BufferedStreamingRouting,
10967 RuntimeCommitBehavior::FinalResponse,
10968 "failed",
10969 false,
10970 );
10971 return Err(error);
10972 }
10973 };
10974 let raw_draft_content = stream_draft.draft.raw_content().to_string();
10975 let buffered_chunks = stream_draft.chunks;
10976 self.finalize_optional_branch(
10977 &main_id,
10978 RuntimeOptimizationKind::BufferedStreamingRouting,
10979 RuntimeCommitBehavior::FinalResponse,
10980 "committed",
10981 true,
10982 );
10983 let response = self
10984 .commit_main_response_draft(
10985 processed_input,
10986 input_context,
10987 stream_draft.draft,
10988 ReasoningMode::None,
10989 false,
10990 )
10991 .await?;
10992 let chunks = if response.content == raw_draft_content {
10993 buffered_chunks
10994 } else {
10995 vec![StreamChunk::content(response.content.clone())]
10996 };
10997 return Ok(Some((response, chunks)));
10998 }
10999 tokio::select! {
11000 result = &mut main_future, if main_pending => {
11001 main_pending = false;
11002 main_branch.transition_to(RuntimeBranchStatus::Completed)?;
11003 main_result = Some(result);
11004 }
11005 result = &mut transition_future, if transition_pending => {
11006 transition_pending = false;
11007 transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
11008 match result {
11009 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
11010 transition_candidate = Some(candidate)
11011 }
11012 Ok(ParallelTransitionSelection::NoMatch) => {
11013 self.finalize_optional_branch(
11014 &transition_id,
11015 RuntimeOptimizationKind::ParallelStateTransition,
11016 RuntimeCommitBehavior::TransitionDecision,
11017 "discarded",
11018 false,
11019 );
11020 routing_resolved.store(true, Ordering::SeqCst);
11021 transition_finalized = true;
11022 }
11023 Ok(ParallelTransitionSelection::ReservationExhausted) => {
11024 self.finalize_optional_branch(
11025 &transition_id,
11026 RuntimeOptimizationKind::ParallelStateTransition,
11027 RuntimeCommitBehavior::TransitionDecision,
11028 "cancelled",
11029 false,
11030 );
11031 routing_resolved.store(true, Ordering::SeqCst);
11032 self.finalize_branch_loss(
11033 &main_id,
11034 RuntimeOptimizationKind::BufferedStreamingRouting,
11035 RuntimeCommitBehavior::FinalResponse,
11036 main_pending,
11037 main_result.as_ref().map(|result| result.is_err()),
11038 );
11039 return Ok(None);
11040 }
11041 Err(_) => {
11042 self.finalize_optional_branch(
11043 &transition_id,
11044 RuntimeOptimizationKind::ParallelStateTransition,
11045 RuntimeCommitBehavior::TransitionDecision,
11046 "failed",
11047 false,
11048 );
11049 routing_resolved.store(true, Ordering::SeqCst);
11050 transition_finalized = true;
11051 }
11052 }
11053 }
11054 }
11055 }
11056 }
11057
11058 fn run_loop_internal_stream<'a>(
11062 &'a self,
11063 input: &'a str,
11064 terminal: RuntimeStreamTerminalSlot,
11065 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11066 let include_tool_events = self.streaming.include_tool_events;
11067 let include_state_events = self.streaming.include_state_events;
11068
11069 Box::pin(async_stream::stream! {
11070 self.begin_root_turn();
11071 self.pre_turn_session_lifecycle().await;
11073
11074 let input_data = match self.process_input(input).await {
11075 Ok(data) => data,
11076 Err(e) => {
11077 yield StreamChunk::error(e.to_string());
11078 return;
11079 }
11080 };
11081 self.update_active_turn_context(&input_data.content, input_data.context.clone());
11082
11083 for (key, value) in &input_data.context {
11085 let _ = self.context_manager.set(key, value.clone());
11086 }
11087
11088 if input_data.metadata.rejected {
11089 let reason = input_data
11090 .metadata
11091 .rejection_reason
11092 .unwrap_or_else(|| "Input rejected".to_string());
11093 warn!(reason = %reason, "Input rejected (stream)");
11094 yield StreamChunk::error(reason);
11095 return;
11096 }
11097
11098 let processed_input = &input_data.content;
11099
11100 if self.runtime_config.optimization.enabled
11101 && matches!(
11102 self.runtime_config.optimization.streaming_policy,
11103 crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11104 )
11105 {
11106 match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11111 Ok(Some((response, chunks))) => {
11112 for chunk in chunks {
11113 yield chunk;
11114 }
11115 record_runtime_stream_final(&terminal, response);
11116 yield StreamChunk::Done {};
11117 return;
11118 }
11119 Ok(None) => {}
11120 Err(e) => {
11121 yield StreamChunk::error(e.to_string());
11122 return;
11123 }
11124 }
11125 }
11126
11127 if self.runtime_config.optimization.enabled
11128 && matches!(
11129 self.runtime_config.optimization.streaming_policy,
11130 crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11131 )
11132 {
11133 match self.try_pre_response_transition(processed_input).await {
11134 Ok(Some(response)) => {
11135 yield StreamChunk::content(&response.content);
11136 record_runtime_stream_final(&terminal, response);
11137 yield StreamChunk::Done {};
11138 return;
11139 }
11140 Ok(None) => {}
11141 Err(e) => {
11142 yield StreamChunk::error(e.to_string());
11143 return;
11144 }
11145 }
11146 }
11147
11148 if let Some(ref sm) = self.state_machine
11150 && let Some(def) = sm.current_definition()
11151 {
11152 let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11153 Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11154 } else if let Some(ref concurrent_config) = def.concurrent {
11155 Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11156 } else if let Some(ref group_chat_config) = def.group_chat {
11157 Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11158 } else if let Some(ref pipeline_config) = def.pipeline {
11159 Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11160 } else if let Some(ref handoff_config) = def.handoff {
11161 Some(self.handle_handoff_state(processed_input, handoff_config).await)
11162 } else {
11163 None
11164 };
11165
11166 if let Some(result) = orchestration_result {
11167 match result {
11168 Ok(response) => {
11169 yield StreamChunk::content(&response.content);
11170 record_runtime_stream_final(&terminal, response);
11171 yield StreamChunk::Done {};
11172 }
11173 Err(e) => {
11174 yield StreamChunk::error(e.to_string());
11175 }
11176 }
11177 return;
11178 }
11179 }
11180
11181 match self.try_skill_route(processed_input).await {
11183 Ok(SkillRouteResult::Response { skill_id, content }) => {
11184 if let Err(e) = self.commit_root_user_message(processed_input).await {
11185 yield StreamChunk::error(e.to_string());
11186 return;
11187 }
11188 match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11189 Ok(resp) => {
11190 yield StreamChunk::content(&resp.content);
11191 record_runtime_stream_final(&terminal, resp);
11192 yield StreamChunk::Done {};
11193 return;
11194 }
11195 Err(e) => {
11196 yield StreamChunk::error(e.to_string());
11197 return;
11198 }
11199 }
11200 }
11201 Ok(SkillRouteResult::NeedsClarification {
11202 response,
11203 ownership,
11204 }) => {
11205 let admission = match self
11206 .admit_optional_disambiguation_ownership(ownership)
11207 .await
11208 {
11209 Ok(admission) => admission,
11210 Err(e) => {
11211 yield StreamChunk::error(e.to_string());
11212 return;
11213 }
11214 };
11215 if let Err(e) = self.commit_root_user_message(processed_input).await {
11216 yield StreamChunk::error(e.to_string());
11217 return;
11218 }
11219 let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11220 drop(admission);
11221 if let Err(e) = self.finish_turn_if_root(&response).await {
11222 yield StreamChunk::error(e.to_string());
11223 return;
11224 }
11225 yield StreamChunk::content(&response.content);
11226 record_runtime_stream_final(&terminal, response);
11227 yield StreamChunk::Done {};
11228 return;
11229 }
11230 Ok(SkillRouteResult::NoMatch) => {} Err(e) => {
11232 yield StreamChunk::error(e.to_string());
11233 return;
11234 }
11235 }
11236
11237 let effective_reasoning = self.get_effective_reasoning_config();
11239 let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11240 Ok(mode) => mode,
11241 Err(e) => {
11242 yield StreamChunk::error(e.to_string());
11243 return;
11244 }
11245 };
11246 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11247
11248 info!(
11249 reasoning_mode = ?reasoning_mode,
11250 auto_detected = auto_detected,
11251 "Reasoning mode determined (stream)"
11252 );
11253
11254 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11256 if let Err(e) = self.commit_root_user_message(processed_input).await {
11257 yield StreamChunk::error(e.to_string());
11258 return;
11259 }
11260 match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11261 Ok(resp) => {
11262 yield StreamChunk::content(&resp.content);
11263 record_runtime_stream_final(&terminal, resp);
11264 yield StreamChunk::Done {};
11265 return;
11266 }
11267 Err(e) => {
11268 yield StreamChunk::error(e.to_string());
11269 return;
11270 }
11271 }
11272 }
11273
11274 if let Err(e) = self.commit_root_user_message(processed_input).await {
11275 yield StreamChunk::error(e.to_string());
11276 return;
11277 }
11278
11279 let llm = match self.get_state_llm() {
11280 Ok(llm) => llm,
11281 Err(e) => {
11282 yield StreamChunk::error(e.to_string());
11283 return;
11284 }
11285 };
11286
11287 let mut iterations = 0u32;
11288 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11289 let mut thinking_content: Option<String> = None;
11290
11291 loop {
11292 let effective_max = if reasoning_mode != ReasoningMode::None {
11294 let rc = self.get_effective_reasoning_config();
11295 self.max_iterations.min(rc.max_iterations)
11296 } else {
11297 self.max_iterations
11298 };
11299
11300 if iterations >= effective_max {
11301 let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11302 let err = AgentError::Other(err_msg.clone());
11303 self.hooks.on_error(&err).await;
11304 error!(iterations = iterations, "Max iterations exceeded (stream)");
11305 yield StreamChunk::error(err_msg);
11306 return;
11307 }
11308 iterations += 1;
11309 *self.iteration_count.write() = iterations;
11310
11311 debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11312
11313 let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11314 Ok(protocol) => protocol,
11315 Err(e) => {
11316 yield StreamChunk::error(e.to_string());
11317 return;
11318 }
11319 };
11320 let mut messages = match self
11321 .build_messages_internal(true, None, protocol.choice.is_none())
11322 .await
11323 {
11324 Ok(m) => m,
11325 Err(e) => {
11326 yield StreamChunk::error(e.to_string());
11327 return;
11328 }
11329 };
11330 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11331
11332 self.hooks.on_llm_start(&messages).await;
11333 let llm_start = Instant::now();
11334
11335 let reflection_active = self
11338 .should_reflect(processed_input, "")
11339 .await
11340 .unwrap_or_default();
11341
11342 let buffered_decision = reflection_active || protocol.choice.is_some();
11343 let content = if buffered_decision {
11344 let response = match self
11348 .complete_main_llm_with_recovery(
11349 Arc::clone(&llm),
11350 &messages,
11351 &protocol,
11352 )
11353 .await
11354 {
11355 Ok(r) => r,
11356 Err(e) => {
11357 yield StreamChunk::error(e.to_string());
11358 return;
11359 }
11360 };
11361 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11362 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11363 response.content.trim().to_string()
11364 } else {
11365 let llm_stream = match self
11367 .observe_purpose(
11368 ObservationPurpose::MainResponse,
11369 llm.complete_stream(&messages, None),
11370 )
11371 .await
11372 {
11373 Ok(s) => s,
11374 Err(e) => {
11375 yield StreamChunk::error(e.to_string());
11376 return;
11377 }
11378 };
11379 let mut accumulated = String::new();
11380 let mut stream_inner = llm_stream;
11381 while let Some(chunk_result) = stream_inner.next().await {
11382 match chunk_result {
11383 Ok(chunk) => {
11384 accumulated.push_str(&chunk.delta);
11385 yield StreamChunk::content(chunk.delta);
11386 }
11387 Err(e) => {
11388 yield StreamChunk::error(e.to_string());
11389 return;
11390 }
11391 }
11392 }
11393 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11394 let llm_response = ai_agents_core::LLMResponse::new(
11396 accumulated.trim(),
11397 ai_agents_core::FinishReason::Stop,
11398 );
11399 self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11400 accumulated.trim().to_string()
11401 };
11402
11403 if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11405 let native_tool_call = Self::is_native_tool_call_content(&content);
11406 let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11409 Ok(v) => v,
11410 Err(e) => {
11411 yield StreamChunk::error(e.to_string());
11412 return;
11413 }
11414 };
11415 if transition_fired {
11416 let _ = self.memory.add_message(ChatMessage::assistant(
11417 "(Transitioned to new state — tool call handled by workflow)",
11418 )).await;
11419
11420 if include_state_events
11421 && let Some(state) = self.current_state()
11422 {
11423 yield StreamChunk::state_transition(None, state);
11424 }
11425 continue;
11426 }
11427
11428 let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11430
11431 let results = self.execute_tools_parallel(&tool_calls).await;
11433
11434 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11435 if include_tool_events {
11436 yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11437 }
11438
11439 match result {
11440 Ok(output) => {
11441 if include_tool_events {
11442 yield StreamChunk::tool_result(
11443 &tool_call.id,
11444 &tool_call.name,
11445 &output,
11446 true,
11447 );
11448 }
11449 let _ = self.memory
11450 .add_message(Self::tool_result_message(
11451 tool_call,
11452 &output,
11453 native_tool_call,
11454 ))
11455 .await;
11456 }
11457 Err(e) => {
11458 if matches!(e, AgentError::HITLRejected(_)) {
11459 let _ = self.memory.add_message(ChatMessage::assistant(
11460 format!("The operation was rejected by the approver: {}", e),
11461 )).await;
11462 let response = AgentResponse {
11463 content: format!("Operation cancelled: {}", e),
11464 metadata: None,
11465 tool_calls: Some(all_tool_calls.clone()),
11466 };
11467 if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11468 yield StreamChunk::error(finalize_error.to_string());
11469 return;
11470 }
11471 let legacy_error = response.content.clone();
11472 record_runtime_stream_final(&terminal, response);
11473 yield StreamChunk::error(legacy_error);
11474 yield StreamChunk::Done {};
11475 return;
11476 }
11477 if include_tool_events {
11478 yield StreamChunk::tool_result(
11479 &tool_call.id,
11480 &tool_call.name,
11481 e.to_string(),
11482 false,
11483 );
11484 }
11485 let _ = self.memory
11486 .add_message(Self::tool_result_message(
11487 tool_call,
11488 &format!("Error: {}", e),
11489 native_tool_call,
11490 ))
11491 .await;
11492 }
11493 }
11494 all_tool_calls.push(tool_call.clone());
11495
11496 if include_tool_events {
11497 yield StreamChunk::tool_end(&tool_call.id);
11498 }
11499 }
11500 continue;
11501 }
11502
11503 let (extracted_thinking, answer) = self.extract_thinking(&content);
11505 if extracted_thinking.is_some() {
11506 thinking_content = extracted_thinking;
11507 }
11508
11509 let output_data = match self.process_output(&answer, &input_data.context).await {
11510 Ok(d) => d,
11511 Err(e) => {
11512 yield StreamChunk::error(e.to_string());
11513 return;
11514 }
11515 };
11516
11517 let final_content = if output_data.metadata.rejected {
11518 output_data
11519 .metadata
11520 .rejection_reason
11521 .unwrap_or_else(|| answer.to_string())
11522 } else {
11523 output_data.content
11524 };
11525
11526 let (final_content, reflection_metadata) = match self
11528 .run_reflection(&*llm, processed_input, final_content)
11529 .await
11530 {
11531 Ok(r) => r,
11532 Err(e) => {
11533 yield StreamChunk::error(e.to_string());
11534 return;
11535 }
11536 };
11537
11538 let final_content = self.format_response_with_thinking(
11539 thinking_content.as_deref(),
11540 &final_content,
11541 );
11542
11543 if buffered_decision {
11545 yield StreamChunk::content(&final_content);
11546 }
11547
11548 let post_result = match self
11552 .post_loop_processing(processed_input, final_content)
11553 .await
11554 {
11555 Ok(r) => r,
11556 Err(e) => {
11557 yield StreamChunk::error(e.to_string());
11558 return;
11559 }
11560 };
11561
11562 let (final_content, transitioned) = match post_result {
11563 PostLoopResult::NoTransition(content) => (content, false),
11564 PostLoopResult::Transitioned(content) => (content, true),
11565 PostLoopResult::NeedsRedispatch => {
11566 const MAX_REDISPATCH_DEPTH: u32 = 3;
11567 let current_depth = *self.redispatch_depth.read();
11568 let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11569 warn!(
11570 depth = current_depth,
11571 "Post-transition re-dispatch depth limit reached (stream)"
11572 );
11573 let c = String::new();
11574 let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11575 c
11576 } else {
11577 *self.redispatch_depth.write() += 1;
11578 if let Some(context) = self.active_turn_context.write().as_mut() {
11579 context.enter_redispatch();
11580 }
11581 info!(
11582 depth = current_depth + 1,
11583 "Re-dispatching for new state after transition (stream)"
11584 );
11585 let result = self.run_loop_internal(processed_input).await;
11586 *self.redispatch_depth.write() -= 1;
11587 if let Some(context) = self.active_turn_context.write().as_mut() {
11588 context.exit_redispatch();
11589 }
11590 match result {
11591 Ok(resp) => resp.content,
11592 Err(e) => {
11593 yield StreamChunk::error(e.to_string());
11594 return;
11595 }
11596 }
11597 };
11598 (content, true)
11599 }
11600 };
11601
11602 if transitioned {
11603 if include_state_events
11604 && let Some(state) = self.current_state()
11605 {
11606 yield StreamChunk::state_transition(None, state);
11607 }
11608 yield StreamChunk::content(&final_content);
11610 }
11611
11612 let final_response = self.build_agent_response(AgentResponseParts {
11614 content: final_content,
11615 all_tool_calls,
11616 reasoning_mode,
11617 auto_detected,
11618 iterations,
11619 thinking: thinking_content,
11620 reflection_metadata,
11621 });
11622 if let Err(e) = self.finish_turn_if_root(&final_response).await {
11623 yield StreamChunk::error(e.to_string());
11624 return;
11625 }
11626
11627 record_runtime_stream_final(&terminal, final_response);
11628 yield StreamChunk::Done {};
11629 return;
11630 }
11631 })
11632 }
11633
11634 fn run_loop_stream<'a>(
11637 &'a self,
11638 input: &'a str,
11639 terminal: RuntimeStreamTerminalSlot,
11640 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11641 Box::pin(async_stream::stream! {
11642 self.begin_root_turn();
11643 let _root_cleanup = RootTurnCleanup::new(self);
11644 self.hooks.on_message_received(input).await;
11645
11646 if !self.context_initialized.swap(true, Ordering::SeqCst) {
11648 if let Err(e) = self.context_manager.initialize().await {
11649 yield StreamChunk::error(e.to_string());
11650 return;
11651 }
11652 debug!("Context manager initialized (defaults, env, builtins)");
11653 }
11654
11655 if let Err(e) = self.check_turn_timeout().await {
11656 yield StreamChunk::error(e.to_string());
11657 return;
11658 }
11659 if let Err(e) = self.context_manager.refresh_per_turn().await {
11660 yield StreamChunk::error(e.to_string());
11661 return;
11662 }
11663
11664 self.clear_disambiguation_context();
11666
11667 if let Some(ref disambiguator) = self.disambiguation_manager {
11669 let disambiguation_context = match self.build_disambiguation_context().await {
11670 Ok(ctx) => ctx,
11671 Err(e) => {
11672 yield StreamChunk::error(e.to_string());
11673 return;
11674 }
11675 };
11676
11677 let state_override = self
11678 .state_machine
11679 .as_ref()
11680 .and_then(|sm| sm.current_definition())
11681 .and_then(|def| def.disambiguation.clone());
11682
11683 let state_generation = self
11684 .state_machine
11685 .as_ref()
11686 .map(|state_machine| state_machine.generation());
11687 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11688 let mut result = match self
11689 .observe_purpose(
11690 ObservationPurpose::DisambiguationDetection,
11691 disambiguator.process_input_with_override(
11692 input,
11693 &disambiguation_context,
11694 state_override.as_ref(),
11695 None,
11696 ),
11697 )
11698 .await
11699 {
11700 Ok(r) => r,
11701 Err(e) => {
11702 yield StreamChunk::error(e.to_string());
11703 return;
11704 }
11705 };
11706 let current_state_generation = self
11707 .state_machine
11708 .as_ref()
11709 .map(|state_machine| state_machine.generation());
11710 if current_state_generation != state_generation
11711 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
11712 {
11713 disambiguator.clear_pending().await;
11714 *self.pending_skill_id.write() = None;
11715 result = DisambiguationResult::Abandoned { new_input: None };
11716 info!(
11717 confirmation_event = "invalidated",
11718 invalidation_reason = "state_generation_changed",
11719 "Streaming disambiguation result invalidated before redispatch"
11720 );
11721 }
11722 match result {
11723 DisambiguationResult::Clear => {
11724 debug!("Input is clear, proceeding normally (stream)");
11725 }
11726 DisambiguationResult::NeedsClarification {
11727 question,
11728 detection,
11729 } => {
11730 let admission = match self
11731 .admit_disambiguation_redispatch(
11732 disambiguation_epoch,
11733 state_generation,
11734 )
11735 .await
11736 {
11737 Ok(admission) => admission,
11738 Err(error) => {
11739 *self.pending_skill_id.write() = None;
11740 yield StreamChunk::error(error.to_string());
11741 return;
11742 }
11743 };
11744 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
11745 info!(
11746 ambiguity_type = ?detection.ambiguity_type,
11747 confidence = detection.confidence,
11748 "Input requires clarification (stream)"
11749 );
11750 if let Err(e) = self.commit_root_user_message(input).await {
11753 yield StreamChunk::error(e.to_string());
11754 return;
11755 }
11756 let _ = self
11757 .memory
11758 .add_message(ChatMessage::assistant(&question.question))
11759 .await;
11760 let status = if awaiting_confirmation {
11761 "awaiting_confirmation"
11762 } else {
11763 "awaiting_clarification"
11764 };
11765 let response = AgentResponse::new(&question.question).with_metadata(
11766 "disambiguation",
11767 serde_json::json!({ "status": status }),
11768 );
11769 drop(admission);
11770 if let Err(e) = self.finish_turn_if_root(&response).await {
11771 yield StreamChunk::error(e.to_string());
11772 return;
11773 }
11774 yield StreamChunk::content(&question.question);
11775 record_runtime_stream_final(&terminal, response);
11776 yield StreamChunk::Done {};
11777 return;
11778 }
11779 DisambiguationResult::Clarified {
11780 enriched_input,
11781 resolved,
11782 ..
11783 } => {
11784 let admission = match self
11785 .admit_disambiguation_redispatch(
11786 disambiguation_epoch,
11787 state_generation,
11788 )
11789 .await
11790 {
11791 Ok(admission) => admission,
11792 Err(error) => {
11793 *self.pending_skill_id.write() = None;
11794 yield StreamChunk::error(error.to_string());
11795 return;
11796 }
11797 };
11798 info!(
11799 resolved_count = resolved.len(),
11800 enriched = %enriched_input,
11801 "Input clarified (stream)"
11802 );
11803 for (key, value) in &resolved {
11804 let context_key = format!("disambiguation.{}", key);
11805 let _ = self.context_manager.set(&context_key, value.clone());
11806 }
11807 if let Some(intent) = resolved.get("intent") {
11808 let _ = self.context_manager.set("resolved_intent", intent.clone());
11809 }
11810 let _ = self
11811 .context_manager
11812 .set("disambiguation.resolved", serde_json::Value::Bool(true));
11813
11814 let skill_id = self.pending_skill_id.read().clone();
11818 if let Some(skill_id) = skill_id {
11819 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
11820 drop(admission);
11821 match self
11822 .recheck_skill_disambiguation(
11823 &skill_id,
11824 &enriched_input,
11825 disambiguation_epoch,
11826 state_generation,
11827 )
11828 .await
11829 {
11830 Ok(resp) => {
11831 yield StreamChunk::content(&resp.content);
11832 record_runtime_stream_final(&terminal, resp);
11833 yield StreamChunk::Done {};
11834 return;
11835 }
11836 Err(e) => {
11837 yield StreamChunk::error(e.to_string());
11838 return;
11839 }
11840 }
11841 }
11842
11843 drop(admission);
11845 let mut inner = self.run_loop_internal_stream(
11846 &enriched_input,
11847 Arc::clone(&terminal),
11848 );
11849 while let Some(chunk) = inner.next().await {
11850 yield chunk;
11851 }
11852 return;
11853 }
11854 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
11855 info!("Proceeding with best guess (stream)");
11856
11857 let skill_id = self.pending_skill_id.read().clone();
11859 if let Some(skill_id) = skill_id {
11860 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
11861 match self
11862 .recheck_skill_disambiguation(
11863 &skill_id,
11864 &enriched_input,
11865 disambiguation_epoch,
11866 state_generation,
11867 )
11868 .await
11869 {
11870 Ok(resp) => {
11871 yield StreamChunk::content(&resp.content);
11872 record_runtime_stream_final(&terminal, resp);
11873 yield StreamChunk::Done {};
11874 return;
11875 }
11876 Err(e) => {
11877 yield StreamChunk::error(e.to_string());
11878 return;
11879 }
11880 }
11881 }
11882
11883 let mut inner = self.run_loop_internal_stream(
11884 &enriched_input,
11885 Arc::clone(&terminal),
11886 );
11887 while let Some(chunk) = inner.next().await {
11888 yield chunk;
11889 }
11890 return;
11891 }
11892 DisambiguationResult::GiveUp { reason } => {
11893 *self.pending_skill_id.write() = None;
11894 warn!(reason = %reason, "Disambiguation gave up (stream)");
11895 let apology = self
11896 .generate_localized_apology(
11897 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
11898 &reason,
11899 )
11900 .await
11901 .unwrap_or_else(|_| {
11902 format!("I'm sorry, I couldn't understand your request: {}", reason)
11903 });
11904 let response = AgentResponse::new(&apology);
11905 if let Err(e) = self.finish_turn_if_root(&response).await {
11906 yield StreamChunk::error(e.to_string());
11907 return;
11908 }
11909 yield StreamChunk::content(&apology);
11910 record_runtime_stream_final(&terminal, response);
11911 yield StreamChunk::Done {};
11912 return;
11913 }
11914 DisambiguationResult::Escalate { reason } => {
11915 *self.pending_skill_id.write() = None;
11916 info!(reason = %reason, "Escalating to human (stream)");
11917 if let Some(ref hitl) = self.hitl_engine {
11918 let trigger =
11919 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
11920 let mut context_map = HashMap::new();
11921 context_map.insert("original_input".to_string(), serde_json::json!(input));
11922 context_map.insert("reason".to_string(), serde_json::json!(&reason));
11923 let check_result = HITLCheckResult::required(
11924 trigger,
11925 context_map,
11926 format!("User request needs human assistance: {}", reason),
11927 Some(hitl.config().default_timeout_seconds),
11928 );
11929 match self.request_hitl_approval(check_result).await {
11930 Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
11931 let mut inner = self.run_loop_internal_stream(
11932 input,
11933 Arc::clone(&terminal),
11934 );
11935 while let Some(chunk) = inner.next().await {
11936 yield chunk;
11937 }
11938 return;
11939 }
11940 Ok(_) => {}
11941 Err(e) => {
11942 yield StreamChunk::error(e.to_string());
11943 return;
11944 }
11945 }
11946 }
11947 let apology = self
11948 .generate_localized_apology(
11949 "Explain briefly that you're transferring the user to a human agent for help.",
11950 &reason,
11951 )
11952 .await
11953 .unwrap_or_else(|_| {
11954 format!("I need human assistance to help with your request: {}", reason)
11955 });
11956 let response = AgentResponse::new(&apology);
11957 if let Err(e) = self.finish_turn_if_root(&response).await {
11958 yield StreamChunk::error(e.to_string());
11959 return;
11960 }
11961 yield StreamChunk::content(&apology);
11962 record_runtime_stream_final(&terminal, response);
11963 yield StreamChunk::Done {};
11964 return;
11965 }
11966 DisambiguationResult::Abandoned { new_input } => {
11967 *self.pending_skill_id.write() = None;
11968
11969 info!(
11970 has_new_input = new_input.is_some(),
11971 "Clarification abandoned by user (stream)"
11972 );
11973
11974 if let Err(e) = self.commit_root_user_message(input).await {
11975 yield StreamChunk::error(e.to_string());
11976 return;
11977 }
11978
11979 match new_input {
11980 Some(fresh_input) => {
11981 let mut inner = self.run_loop_internal_stream(
11983 &fresh_input,
11984 Arc::clone(&terminal),
11985 );
11986 while let Some(chunk) = inner.next().await {
11987 yield chunk;
11988 }
11989 return;
11990 }
11991 None => {
11992 let ack = self
11994 .generate_localized_apology(
11995 "The user changed their mind about their previous request. \
11996 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
11997 Do NOT apologize excessively. Be concise.",
11998 "User abandoned clarification",
11999 )
12000 .await
12001 .unwrap_or_else(|_| {
12002 "OK, no problem. What else can I help with?".to_string()
12003 });
12004
12005 let _ = self
12006 .memory
12007 .add_message(ChatMessage::assistant(&ack))
12008 .await;
12009
12010 let response = AgentResponse::new(&ack);
12011 if let Err(e) = self.finish_turn_if_root(&response).await {
12012 yield StreamChunk::error(e.to_string());
12013 return;
12014 }
12015 yield StreamChunk::content(&ack);
12016 record_runtime_stream_final(&terminal, response);
12017 yield StreamChunk::Done {};
12018 return;
12019 }
12020 }
12021 }
12022 }
12023 }
12024
12025 let mut inner = self.run_loop_internal_stream(input, Arc::clone(&terminal));
12027 while let Some(chunk) = inner.next().await {
12028 yield chunk;
12029 }
12030 })
12031 }
12032
12033 pub fn info(&self) -> AgentInfo {
12034 self.info.clone()
12035 }
12036
12037 pub fn skills(&self) -> &[SkillDefinition] {
12038 &self.skills
12039 }
12040
12041 async fn reset_runtime_state(&self) -> Result<()> {
12043 let _admission = self.disambiguation_admission.write().await;
12044 if self.state_transition_reserved.load(Ordering::SeqCst) {
12045 return Err(AgentError::Other(
12046 "Cannot reset while a state transition is in progress".to_string(),
12047 ));
12048 }
12049 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
12050 *self.pending_skill_id.write() = None;
12051 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
12052 disambiguator.clear_pending().await;
12053 }
12054 self.memory.clear().await?;
12055 *self.iteration_count.write() = 0;
12056 self.tool_call_history.write().clear();
12057 if let Some(ref sm) = self.state_machine {
12058 sm.reset();
12059 }
12060 Ok(())
12061 }
12062
12063 pub async fn reset(&self) -> Result<()> {
12065 self.reset_runtime_state().await
12066 }
12067
12068 pub fn max_context_tokens(&self) -> u32 {
12069 self.max_context_tokens
12070 }
12071
12072 pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12073 &self.llm_registry
12074 }
12075
12076 pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12077 self.state_machine.as_ref()
12078 }
12079
12080 pub fn context_manager(&self) -> &Arc<ContextManager> {
12081 &self.context_manager
12082 }
12083
12084 pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12085 self.tool_call_history.read().clone()
12086 }
12087
12088 pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12089 self.memory_token_budget.as_ref()
12090 }
12091
12092 pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12093 &self.parallel_tools
12094 }
12095
12096 pub fn streaming_config(&self) -> &StreamingConfig {
12097 &self.streaming
12098 }
12099
12100 pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12101 &self.hooks
12102 }
12103
12104 pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12105 self.hitl_engine.as_ref()
12106 }
12107
12108 pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12109 &self.approval_handler
12110 }
12111
12112 fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12114 let mut ctx = HashMap::new();
12115 for key in &["user.language", "input.detected.language", "language"] {
12116 if let Some(val) = self.context_manager.get(key) {
12117 ctx.insert(key.to_string(), val);
12118 }
12119 }
12120 ctx
12121 }
12122
12123 async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12125 let Some(request) = check_result.into_request() else {
12126 return Ok(ApprovalResult::Approved);
12127 };
12128
12129 self.hooks.on_approval_requested(&request).await;
12130
12131 let timeout = request.timeout;
12132
12133 let raw_result = if let Some(duration) = timeout {
12134 match tokio::time::timeout(
12135 duration,
12136 self.approval_handler.request_approval(request.clone()),
12137 )
12138 .await
12139 {
12140 Ok(result) => result,
12141 Err(_) => ApprovalResult::timeout(),
12142 }
12143 } else {
12144 self.approval_handler
12145 .request_approval(request.clone())
12146 .await
12147 };
12148
12149 self.hooks
12150 .on_approval_result(&request.id, &raw_result)
12151 .await;
12152
12153 let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12154 match &raw_result {
12155 ApprovalResult::Approved => (
12156 ApprovalResolvedOutcome::Approved,
12157 Ok(ApprovalResult::Approved),
12158 ),
12159 ApprovalResult::Rejected { reason } => (
12160 ApprovalResolvedOutcome::Rejected {
12161 reason: reason.clone(),
12162 },
12163 Ok(ApprovalResult::Rejected {
12164 reason: reason.clone(),
12165 }),
12166 ),
12167 ApprovalResult::Modified { changes } => (
12168 ApprovalResolvedOutcome::Modified {
12169 changes: changes.clone(),
12170 },
12171 Ok(ApprovalResult::Modified {
12172 changes: changes.clone(),
12173 }),
12174 ),
12175 ApprovalResult::Timeout => {
12176 if let Some(ref engine) = self.hitl_engine {
12177 match engine.config().on_timeout {
12178 TimeoutAction::Approve => (
12179 ApprovalResolvedOutcome::Approved,
12180 Ok(ApprovalResult::Approved),
12181 ),
12182 TimeoutAction::Reject => {
12183 let reason = Some("Timeout".to_string());
12184 (
12185 ApprovalResolvedOutcome::Rejected {
12186 reason: reason.clone(),
12187 },
12188 Ok(ApprovalResult::Rejected { reason }),
12189 )
12190 }
12191 TimeoutAction::Error => {
12192 let message = "HITL approval timeout".to_string();
12193 (
12194 ApprovalResolvedOutcome::Error {
12195 message: message.clone(),
12196 },
12197 Err(AgentError::Other(message)),
12198 )
12199 }
12200 }
12201 } else {
12202 let reason = Some("Timeout (no engine)".to_string());
12203 (
12204 ApprovalResolvedOutcome::Rejected {
12205 reason: reason.clone(),
12206 },
12207 Ok(ApprovalResult::Rejected { reason }),
12208 )
12209 }
12210 }
12211 };
12212
12213 self.hooks
12214 .on_approval_resolved(&request, &raw_result, &outcome)
12215 .await;
12216
12217 effective_result
12218 }
12219
12220 pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12221 if let Some(ref hitl_engine) = self.hitl_engine {
12222 let hitl_lang_ctx = self.build_hitl_language_context();
12223 let check_result = self
12224 .observe_purpose(
12225 ObservationPurpose::HitlLocalization,
12226 hitl_engine.check_state_transition_with_localization(
12227 from,
12228 to,
12229 &hitl_lang_ctx,
12230 self.approval_handler.as_ref(),
12231 Some(&self.llm_registry),
12232 ),
12233 )
12234 .await?;
12235 if check_result.is_required() {
12236 let result = self.request_hitl_approval(check_result).await?;
12237 return Ok(matches!(
12238 result,
12239 ApprovalResult::Approved | ApprovalResult::Modified { .. }
12240 ));
12241 }
12242 }
12243 Ok(true)
12244 }
12245
12246 async fn execute_tools_parallel(
12248 &self,
12249 tool_calls: &[ToolCall],
12250 ) -> Vec<(String, Result<String>)> {
12251 let can_run_parallel = tool_calls.iter().all(|tc| {
12252 self.tools
12253 .resolve(&tc.name)
12254 .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12255 .unwrap_or(false)
12256 });
12257
12258 if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12259 let mut results = Vec::new();
12260 for tc in tool_calls {
12261 let result = self
12262 .observe_purpose(
12263 current_observation_context()
12264 .map(|context| context.purpose)
12265 .unwrap_or_default(),
12266 self.execute_tool_smart(tc),
12267 )
12268 .await;
12269 results.push((tc.id.clone(), result));
12270 }
12271 return results;
12272 }
12273
12274 let chunks: Vec<_> = tool_calls
12275 .chunks(self.parallel_tools.max_parallel)
12276 .collect();
12277
12278 let mut all_results = Vec::new();
12279
12280 for chunk in chunks {
12281 let futures: Vec<_> = chunk
12282 .iter()
12283 .map(|tc| {
12284 let tc = tc.clone();
12285 async move {
12286 let result = self.execute_tool_smart(&tc).await;
12287 (tc.id.clone(), result)
12288 }
12289 })
12290 .collect();
12291
12292 let results = futures::future::join_all(futures).await;
12293 all_results.extend(results);
12294 }
12295
12296 all_results
12297 }
12298
12299 pub async fn chat_stream<'a>(
12301 &'a self,
12302 input: &'a str,
12303 ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12304 self.init_storage().await?;
12308 info!(input_len = input.len(), "Starting streaming chat");
12309 let terminal = new_runtime_stream_terminal_slot();
12310 let inner = self.run_loop_stream(input, terminal);
12311 if let Some(context) = self.build_observation_context(None) {
12312 let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12313 Box::pin(async_stream::stream! {
12314 let mut inner = inner;
12315 loop {
12316 let next = with_observation_context(context.clone(), inner.next()).await;
12317 match next {
12318 Some(chunk) => yield chunk,
12319 None => break,
12320 }
12321 }
12322 self.export_observability_if_configured().await;
12323 });
12324 Ok(stream)
12325 } else {
12326 Ok(inner)
12327 }
12328 }
12329
12330 pub async fn chat_stream_events<'a>(
12332 &'a self,
12333 input: &'a str,
12334 ) -> Result<Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>>> {
12335 self.init_storage().await?;
12339 info!(input_len = input.len(), "Starting streaming chat events");
12340 let terminal = new_runtime_stream_terminal_slot();
12341 let mut inner = self.run_loop_stream(input, Arc::clone(&terminal));
12342 let observation_context = self.build_observation_context(None);
12343 let stream: Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send + 'a>> =
12344 Box::pin(async_stream::stream! {
12345 loop {
12346 let next = if let Some(context) = observation_context.as_ref() {
12347 with_observation_context(context.clone(), inner.next()).await
12348 } else {
12349 inner.next().await
12350 };
12351 match next {
12352 Some(StreamChunk::Done {}) => {
12353 let terminal_event = { terminal.write().take() };
12354 if let Some(response) = terminal_event {
12355 while if let Some(context) = observation_context.as_ref() {
12356 with_observation_context(context.clone(), inner.next())
12357 .await
12358 .is_some()
12359 } else {
12360 inner.next().await.is_some()
12361 } {}
12362 if observation_context.is_some() {
12363 self.export_observability_if_configured().await;
12364 }
12365 yield AgentStreamEvent::Final(response);
12366 return;
12367 }
12368 }
12369 Some(StreamChunk::Error { message }) => {
12370 let finalized = { terminal.read().is_some() };
12371 if finalized {
12372 continue;
12373 }
12374 while if let Some(context) = observation_context.as_ref() {
12375 with_observation_context(context.clone(), inner.next())
12376 .await
12377 .is_some()
12378 } else {
12379 inner.next().await.is_some()
12380 } {}
12381 if observation_context.is_some() {
12382 self.export_observability_if_configured().await;
12383 }
12384 yield AgentStreamEvent::Chunk(StreamChunk::Error { message });
12385 return;
12386 }
12387 Some(chunk) => yield AgentStreamEvent::Chunk(chunk),
12388 None => {
12389 if observation_context.is_some() {
12390 self.export_observability_if_configured().await;
12391 }
12392 break;
12393 }
12394 }
12395 }
12396 });
12397 Ok(stream)
12398 }
12399}
12400
12401#[async_trait]
12402impl ToolInvoker for RuntimeAgent {
12403 async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12404 self.execute_tool_record(request).await
12405 }
12406}
12407
12408#[async_trait]
12409impl Agent for RuntimeAgent {
12410 async fn chat(&self, input: &str) -> Result<AgentResponse> {
12411 let result = if let Some(context) = self.build_observation_context(None) {
12412 with_observation_context(context, self.run_loop(input)).await
12413 } else {
12414 self.run_loop(input).await
12415 };
12416 self.export_observability_if_configured().await;
12417 result
12418 }
12419
12420 fn info(&self) -> AgentInfo {
12421 self.info.clone()
12422 }
12423
12424 async fn reset(&self) -> Result<()> {
12426 self.reset_runtime_state().await
12427 }
12428}
12429
12430fn background_maintenance_tags(
12440 label: &str,
12441 stage: &str,
12442 reason: Option<&str>,
12443 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12444) -> HashMap<String, String> {
12445 let mut tags = HashMap::new();
12446 tags.insert("runtime.background".to_string(), "true".to_string());
12447 tags.insert("runtime.maintenance".to_string(), label.to_string());
12448 tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12449 if let Some(policy) = policy {
12450 tags.insert(
12451 "runtime.await_before_next_turn".to_string(),
12452 await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12453 );
12454 tags.insert(
12455 "runtime.maintenance_mode".to_string(),
12456 maintenance_mode_label(policy.mode).to_string(),
12457 );
12458 }
12459 if let Some(reason) = reason {
12460 tags.insert("runtime.reason".to_string(), reason.to_string());
12461 }
12462 tags
12463}
12464
12465fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12466 match policy {
12467 AwaitBeforeNextTurn::Never => "never",
12468 AwaitBeforeNextTurn::SameActor => "same_actor",
12469 AwaitBeforeNextTurn::Always => "always",
12470 }
12471}
12472
12473fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12474 match mode {
12475 MaintenanceMode::InlineSerial => "inline_serial",
12476 MaintenanceMode::InlineParallel => "inline_parallel",
12477 MaintenanceMode::Background => "background",
12478 }
12479}
12480
12481fn record_background_maintenance_event(
12483 manager: Option<&Arc<ObservabilityManager>>,
12484 label: &str,
12485 status: EventStatus,
12486 duration_ms: u64,
12487 stage: &str,
12488 reason: Option<String>,
12489 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12490) {
12491 if let Some(manager) = manager {
12492 manager.record_lifecycle_event(
12493 EventType::MemoryOperation {
12494 operation: format!("{}_background_{}", label, stage),
12495 },
12496 ObservationPurpose::Other(format!("{}_maintenance", label)),
12497 status,
12498 duration_ms,
12499 background_maintenance_tags(label, stage, reason.as_deref(), policy),
12500 None,
12501 );
12502 }
12503}
12504
12505fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12506 if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12507 MaintenanceMode::InlineParallel
12508 } else {
12509 mode
12510 }
12511}
12512
12513fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12514 match hint {
12515 ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12516 ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12517 ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12518 ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12519 ObservationPurpose::ProcessTransform
12520 }
12521 }
12522}
12523
12524fn new_tool_resource_locks() -> ToolResourceLocks {
12525 Arc::new(RwLock::new(HashMap::new()))
12526}
12527
12528fn tool_resource_lock_keys(
12533 _canonical_id: &str,
12534 args: &Value,
12535 bindings: &ai_agents_core::ToolPolicyBindings,
12536 classification: &ai_agents_core::ToolCallClassification,
12537) -> Vec<String> {
12538 if classification.concurrency_safe {
12539 return Vec::new();
12540 }
12541
12542 let mut keys = Vec::new();
12543 let mut has_path_resource = false;
12544 for binding in &bindings.path_fields {
12545 let value = value_at_argument_path(args, &binding.field)
12546 .cloned()
12547 .or_else(|| {
12548 binding
12549 .default_path
12550 .as_ref()
12551 .map(|path| Value::String(path.clone()))
12552 });
12553 if let Some(value) = value {
12554 collect_resource_strings(&value, |_| {
12555 has_path_resource = true;
12556 });
12557 }
12558 }
12559 for binding in &bindings.domain_fields {
12560 if let Some(value) = value_at_argument_path(args, &binding.field) {
12561 collect_resource_strings(value, |domain| {
12562 let normalized = if binding.is_url {
12563 normalized_url_resource_key(domain)
12564 } else {
12565 domain.trim().trim_end_matches('.').to_ascii_lowercase()
12566 };
12567 keys.push(format!("domain:{}", normalized));
12568 });
12569 }
12570 }
12571 for binding in &bindings.command_fields {
12572 if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12573 continue;
12574 }
12575 if let Some(value) = value_at_argument_path(args, &binding.field) {
12576 collect_resource_strings(value, |_| {
12577 has_path_resource = true;
12578 });
12579 }
12580 }
12581 if has_path_resource {
12582 keys.push("path-mutation:global".to_string());
12583 }
12584 if keys.is_empty() {
12585 keys.push("side-effect:unbound".to_string());
12586 }
12587 keys.sort();
12588 keys.dedup();
12589 keys
12590}
12591
12592fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12593 let mut current = value;
12594 for segment in field.split('.') {
12595 if segment.is_empty() {
12596 return None;
12597 }
12598 current = current.get(segment)?;
12599 }
12600 Some(current)
12601}
12602
12603fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12604 match value {
12605 Value::String(value) => collect(value),
12606 Value::Array(values) => {
12607 for value in values {
12608 if let Some(value) = value.as_str() {
12609 collect(value);
12610 }
12611 }
12612 }
12613 _ => {}
12614 }
12615}
12616
12617fn normalized_url_resource_key(value: &str) -> String {
12618 let value = value.trim();
12619 let Some((scheme, remainder)) = value.split_once("://") else {
12620 return value.to_ascii_lowercase();
12621 };
12622 let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
12623 let (authority, suffix) = remainder.split_at(authority_end);
12624 format!(
12625 "{}://{}{}",
12626 scheme.to_ascii_lowercase(),
12627 authority.to_ascii_lowercase(),
12628 suffix
12629 )
12630}
12631
12632fn render_concurrent_template(
12633 template: &str,
12634 user_input: &str,
12635 context_values: &std::collections::HashMap<String, serde_json::Value>,
12636) -> Result<String> {
12637 let mut env = minijinja::Environment::new();
12638 env.add_template("concurrent", template)
12639 .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
12640
12641 let mut ctx = std::collections::BTreeMap::new();
12642 ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
12643
12644 let context_obj = minijinja::Value::from_serialize(context_values);
12646 ctx.insert("context".to_string(), context_obj);
12647
12648 let tmpl = env
12649 .get_template("concurrent")
12650 .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
12651
12652 tmpl.render(minijinja::Value::from_serialize(&ctx))
12653 .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
12654}
12655
12656#[cfg(test)]
12657mod tests {
12658 use super::*;
12659 use crate::AgentBuilder;
12660 use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
12661 use ai_agents_llm::mock::MockLLMProvider;
12662 use ai_agents_skills::{SkillDefinition, SkillStep};
12663 use ai_agents_tools::{
12664 CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool,
12665 WebFetchResolver, WebFetchTool, WebFetchTransport, WebFetchTransportRequest,
12666 WebFetchTransportResponse,
12667 };
12668
12669 fn mock_with_response(response: &str) -> MockLLMProvider {
12670 let mut mock = MockLLMProvider::new("test");
12671 mock.set_response(response);
12672 mock
12673 }
12674
12675 fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
12676 let mut mock = MockLLMProvider::new("test");
12677 mock.set_responses(responses.into_iter().map(String::from).collect(), true);
12678 mock
12679 }
12680
12681 fn disambiguation_state_machine(
12683 state_enabled: Option<bool>,
12684 require_confirmation: bool,
12685 ) -> Arc<StateMachine> {
12686 let definition = ai_agents_state::StateDefinition {
12687 prompt: Some("Handle the resolved request.".to_string()),
12688 disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
12689 enabled: state_enabled,
12690 require_confirmation,
12691 ..Default::default()
12692 }),
12693 ..Default::default()
12694 };
12695 let review = ai_agents_state::StateDefinition {
12696 prompt: Some("Review a fresh request.".to_string()),
12697 ..Default::default()
12698 };
12699 Arc::new(
12700 StateMachine::new(ai_agents_state::StateConfig {
12701 initial: "active".to_string(),
12702 states: std::collections::HashMap::from([
12703 ("active".to_string(), definition),
12704 ("review".to_string(), review),
12705 ]),
12706 global_transitions: Vec::new(),
12707 fallback: None,
12708 max_no_transition: None,
12709 regenerate_on_transition: true,
12710 })
12711 .unwrap(),
12712 )
12713 }
12714
12715 fn state_disambiguation_agent(
12717 responses: Vec<&str>,
12718 manager_enabled: bool,
12719 state_enabled: Option<bool>,
12720 require_confirmation: bool,
12721 ) -> (RuntimeAgent, MockLLMProvider) {
12722 state_disambiguation_agent_with_skills(
12723 responses,
12724 manager_enabled,
12725 state_enabled,
12726 require_confirmation,
12727 Vec::new(),
12728 )
12729 }
12730
12731 fn state_disambiguation_agent_with_skills(
12733 responses: Vec<&str>,
12734 manager_enabled: bool,
12735 state_enabled: Option<bool>,
12736 require_confirmation: bool,
12737 skills: Vec<SkillDefinition>,
12738 ) -> (RuntimeAgent, MockLLMProvider) {
12739 let mut mock = MockLLMProvider::new("state-confirmation");
12740 mock.set_responses(responses.into_iter().map(String::from).collect(), false);
12741 let observed = mock.clone();
12742 let agent = AgentBuilder::new()
12743 .system_prompt("Handle requests.")
12744 .llm(Arc::new(mock.clone()))
12745 .llm_alias("router", Arc::new(mock))
12746 .state_machine(disambiguation_state_machine(
12747 state_enabled,
12748 require_confirmation,
12749 ))
12750 .skills(skills)
12751 .build()
12752 .unwrap()
12753 .with_disambiguation(DisambiguationConfig {
12754 enabled: manager_enabled,
12755 ..Default::default()
12756 });
12757 (agent, observed)
12758 }
12759
12760 fn confirmation_skill() -> SkillDefinition {
12762 SkillDefinition {
12763 id: "send_report".to_string(),
12764 description: "Send a report after clarification".to_string(),
12765 trigger: "When the user asks to send a report".to_string(),
12766 steps: vec![SkillStep::Prompt {
12767 prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
12768 llm: None,
12769 }],
12770 reasoning: None,
12771 reflection: None,
12772 disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
12773 enabled: Some(true),
12774 ..Default::default()
12775 }),
12776 }
12777 }
12778
12779 fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
12781 observed
12782 .call_history()
12783 .iter()
12784 .filter(|call| {
12785 call.messages
12786 .iter()
12787 .any(|message| message.content.contains("Execute confirmed report skill"))
12788 })
12789 .count()
12790 }
12791
12792 struct BlockingRuntimeConfirmationObserver {
12793 entered: tokio::sync::Barrier,
12794 release: tokio::sync::Notify,
12795 }
12796
12797 impl BlockingRuntimeConfirmationObserver {
12798 fn new() -> Self {
12799 Self {
12800 entered: tokio::sync::Barrier::new(2),
12801 release: tokio::sync::Notify::new(),
12802 }
12803 }
12804 }
12805
12806 struct ResetOnTransitionHooks {
12807 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
12808 invoked: AtomicBool,
12809 }
12810
12811 #[async_trait]
12812 impl AgentHooks for ResetOnTransitionHooks {
12813 async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
12814 if self.invoked.swap(true, Ordering::SeqCst) {
12815 return;
12816 }
12817 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
12818 if let Some(agent) = agent {
12819 agent.reset().await.unwrap();
12820 }
12821 }
12822 }
12823
12824 impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
12825 fn observe_question<'a>(
12826 &'a self,
12827 future: ClarificationQuestionFuture<'a>,
12828 ) -> ClarificationQuestionFuture<'a> {
12829 future
12830 }
12831
12832 fn observe_parse<'a>(
12833 &'a self,
12834 future: ClarificationParseFuture<'a>,
12835 ) -> ClarificationParseFuture<'a> {
12836 future
12837 }
12838
12839 fn observe_confirmation_parse<'a>(
12840 &'a self,
12841 future: ConfirmationParseFuture<'a>,
12842 ) -> ConfirmationParseFuture<'a> {
12843 Box::pin(async move {
12844 self.entered.wait().await;
12845 self.release.notified().await;
12846 future.await
12847 })
12848 }
12849 }
12850
12851 #[tokio::test]
12852 async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
12853 let (agent, observed) = state_disambiguation_agent(
12854 vec![
12855 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12856 r#"{"question":"What should I send?","options":null}"#,
12857 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12858 r#"{"question":"Should I send the report to Ada?"}"#,
12859 r#"{"status":"confirmed"}"#,
12860 "Request executed.",
12861 ],
12862 true,
12863 None,
12864 true,
12865 );
12866
12867 let clarification = agent.chat("Send it").await.unwrap();
12868 assert_eq!(clarification.content, "What should I send?");
12869 assert_eq!(observed.call_count(), 2);
12870
12871 let confirmation = agent.chat("The report to Ada").await.unwrap();
12872 assert_eq!(confirmation.content, "Should I send the report to Ada?");
12873 assert_eq!(
12874 confirmation
12875 .metadata
12876 .as_ref()
12877 .and_then(|metadata| metadata.get("disambiguation"))
12878 .and_then(|metadata| metadata.get("status"))
12879 .and_then(Value::as_str),
12880 Some("awaiting_confirmation")
12881 );
12882 assert_eq!(observed.call_count(), 4);
12883
12884 let completed = agent.chat("Yes").await.unwrap();
12885 assert_eq!(completed.content, "Request executed.");
12886 assert_eq!(observed.call_count(), 6);
12887 }
12888
12889 #[tokio::test]
12890 async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
12891 let (agent, observed) = state_disambiguation_agent(
12892 vec![
12893 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12894 r#"{"question":"What should I send?","options":null}"#,
12895 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12896 r#"{"question":"Should I send the report to Ada?"}"#,
12897 r#"{"status":"confirmed"}"#,
12898 "Request executed.",
12899 ],
12900 true,
12901 None,
12902 true,
12903 );
12904
12905 let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
12906 let mut clarification = String::new();
12907 while let Some(chunk) = clarification_stream.next().await {
12908 match chunk {
12909 StreamChunk::Content { text } => clarification.push_str(&text),
12910 StreamChunk::Done {} => break,
12911 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12912 _ => {}
12913 }
12914 }
12915 assert_eq!(clarification, "What should I send?");
12916 assert_eq!(observed.call_count(), 2);
12917
12918 let mut confirmation_stream = agent.chat_stream_events("The report to Ada").await.unwrap();
12919 let mut confirmation = None;
12920 while let Some(event) = confirmation_stream.next().await {
12921 match event {
12922 AgentStreamEvent::Final(response) => confirmation = Some(response),
12923 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
12924 panic!("unexpected stream error: {message}")
12925 }
12926 AgentStreamEvent::Chunk(_) => {}
12927 }
12928 }
12929 let confirmation = confirmation.expect("confirmation must finalize");
12930 assert_eq!(confirmation.content, "Should I send the report to Ada?");
12931 assert_eq!(
12932 confirmation
12933 .metadata
12934 .as_ref()
12935 .and_then(|metadata| metadata.get("disambiguation"))
12936 .and_then(|metadata| metadata.get("status"))
12937 .and_then(Value::as_str),
12938 Some("awaiting_confirmation")
12939 );
12940 assert_eq!(observed.call_count(), 4);
12941
12942 let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
12943 let mut completed = String::new();
12944 while let Some(chunk) = completed_stream.next().await {
12945 match chunk {
12946 StreamChunk::Content { text } => completed.push_str(&text),
12947 StreamChunk::Done {} => break,
12948 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12949 _ => {}
12950 }
12951 }
12952 assert_eq!(completed, "Request executed.");
12953 assert_eq!(observed.call_count(), 6);
12954 }
12955
12956 #[tokio::test]
12958 async fn confirmed_skill_route_executes_exactly_once() {
12959 let (agent, observed) = state_disambiguation_agent_with_skills(
12960 vec![
12961 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12962 "send_report",
12963 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12964 r#"{"question":"What should I send?","options":null}"#,
12965 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12966 r#"{"question":"Should I send the report to Ada?"}"#,
12967 r#"{"status":"confirmed"}"#,
12968 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
12969 "Report skill executed.",
12970 ],
12971 true,
12972 None,
12973 true,
12974 vec![confirmation_skill()],
12975 );
12976
12977 let clarification = agent.chat("Send it").await.unwrap();
12978 assert_eq!(clarification.content, "What should I send?");
12979 assert_eq!(confirmation_skill_call_count(&observed), 0);
12980
12981 let confirmation = agent.chat("The report to Ada").await.unwrap();
12982 assert_eq!(confirmation.content, "Should I send the report to Ada?");
12983 assert_eq!(
12984 confirmation
12985 .metadata
12986 .as_ref()
12987 .and_then(|metadata| metadata.get("disambiguation"))
12988 .and_then(|metadata| metadata.get("status"))
12989 .and_then(Value::as_str),
12990 Some("awaiting_confirmation")
12991 );
12992 assert_eq!(confirmation_skill_call_count(&observed), 0);
12993
12994 let completed = agent.chat("Yes").await.unwrap();
12995 assert_eq!(completed.content, "Report skill executed.");
12996 assert_eq!(confirmation_skill_call_count(&observed), 1);
12997 assert!(agent.pending_skill_id.read().is_none());
12998 let messages = agent.memory.get_messages(None).await.unwrap();
12999 assert!(!messages.iter().any(|message| message.content == "Yes"));
13000 }
13001
13002 #[tokio::test]
13004 async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
13005 let (agent, observed) = state_disambiguation_agent_with_skills(
13006 vec![
13007 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13008 "send_report",
13009 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13010 r#"{"question":"What should I send?","options":null}"#,
13011 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13012 r#"{"question":"Should I send the report to Ada?"}"#,
13013 r#"{"status":"confirmed"}"#,
13014 r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
13015 r#"{"question":"When should I send it?","options":null}"#,
13016 ],
13017 true,
13018 None,
13019 true,
13020 vec![confirmation_skill()],
13021 );
13022
13023 agent.chat("Send it").await.unwrap();
13024 agent.chat("The report to Ada").await.unwrap();
13025 let follow_up = agent.chat("Yes").await.unwrap();
13026
13027 assert_eq!(follow_up.content, "When should I send it?");
13028 let metadata = follow_up
13029 .metadata
13030 .as_ref()
13031 .and_then(|metadata| metadata.get("disambiguation"))
13032 .unwrap();
13033 assert_eq!(
13034 metadata.get("status").and_then(Value::as_str),
13035 Some("awaiting_clarification")
13036 );
13037 assert_eq!(
13038 metadata.get("skill_id").and_then(Value::as_str),
13039 Some("send_report")
13040 );
13041 assert!(metadata.get("detection").is_some());
13042 assert_eq!(confirmation_skill_call_count(&observed), 0);
13043 }
13044
13045 #[tokio::test]
13047 async fn rejected_skill_confirmation_never_executes() {
13048 let (agent, observed) = state_disambiguation_agent_with_skills(
13049 vec![
13050 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13051 "send_report",
13052 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13053 r#"{"question":"What should I send?","options":null}"#,
13054 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13055 r#"{"question":"Should I send the report to Ada?"}"#,
13056 r#"{"status":"rejected"}"#,
13057 "Confirmation rejected.",
13058 ],
13059 true,
13060 None,
13061 true,
13062 vec![confirmation_skill()],
13063 );
13064
13065 agent.chat("Send it").await.unwrap();
13066 agent.chat("The report to Ada").await.unwrap();
13067 let rejected = agent.chat("No").await.unwrap();
13068
13069 assert_eq!(rejected.content, "Confirmation rejected.");
13070 assert_eq!(confirmation_skill_call_count(&observed), 0);
13071 assert!(agent.pending_skill_id.read().is_none());
13072 }
13073
13074 #[tokio::test]
13076 async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
13077 let (agent, observed) = state_disambiguation_agent_with_skills(
13078 vec![
13079 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13080 "send_report",
13081 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13082 r#"{"question":"What should I send?","options":null}"#,
13083 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13084 r#"{"question":"Should I send the report to Ada?"}"#,
13085 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13086 "none",
13087 "Fresh response.",
13088 ],
13089 true,
13090 None,
13091 true,
13092 vec![confirmation_skill()],
13093 );
13094
13095 agent.chat("Send it").await.unwrap();
13096 agent.chat("The report to Ada").await.unwrap();
13097 agent.reset().await.unwrap();
13098 assert!(agent.pending_skill_id.read().is_none());
13099 assert!(
13100 !agent
13101 .disambiguation_manager()
13102 .unwrap()
13103 .has_pending_clarification()
13104 .await
13105 );
13106
13107 let mut stream = agent.chat_stream("Yes").await.unwrap();
13108 let mut content = String::new();
13109 while let Some(chunk) = stream.next().await {
13110 match chunk {
13111 StreamChunk::Content { text } => content.push_str(&text),
13112 StreamChunk::Done {} => break,
13113 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
13114 _ => {}
13115 }
13116 }
13117
13118 assert_eq!(content, "Fresh response.");
13119 assert_eq!(confirmation_skill_call_count(&observed), 0);
13120 }
13121
13122 #[tokio::test]
13124 async fn trait_reset_clears_pending_skill_confirmation() {
13125 let (agent, _) = state_disambiguation_agent_with_skills(
13126 vec![
13127 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13128 "send_report",
13129 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13130 r#"{"question":"What should I send?","options":null}"#,
13131 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13132 r#"{"question":"Should I send the report to Ada?"}"#,
13133 ],
13134 true,
13135 None,
13136 true,
13137 vec![confirmation_skill()],
13138 );
13139
13140 agent.chat("Send it").await.unwrap();
13141 agent.chat("The report to Ada").await.unwrap();
13142 <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13143
13144 assert!(agent.pending_skill_id.read().is_none());
13145 assert!(
13146 !agent
13147 .disambiguation_manager()
13148 .unwrap()
13149 .has_pending_clarification()
13150 .await
13151 );
13152 }
13153
13154 #[tokio::test]
13156 async fn state_change_invalidates_pending_skill_confirmation() {
13157 let (agent, observed) = state_disambiguation_agent_with_skills(
13158 vec![
13159 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13160 "send_report",
13161 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13162 r#"{"question":"What should I send?","options":null}"#,
13163 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13164 r#"{"question":"Should I send the report to Ada?"}"#,
13165 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13166 "none",
13167 "Fresh response.",
13168 ],
13169 true,
13170 None,
13171 true,
13172 vec![confirmation_skill()],
13173 );
13174
13175 agent.chat("Send it").await.unwrap();
13176 agent.chat("The report to Ada").await.unwrap();
13177 agent.transition_to("review").await.unwrap();
13178 let cancelled = agent.chat("Yes").await.unwrap();
13179
13180 assert_eq!(cancelled.content, "Fresh response.");
13181 assert_eq!(confirmation_skill_call_count(&observed), 0);
13182 assert!(agent.pending_skill_id.read().is_none());
13183 }
13184
13185 #[tokio::test]
13187 async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13188 let (mut agent, observed) = state_disambiguation_agent_with_skills(
13189 vec![
13190 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13191 "send_report",
13192 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13193 r#"{"question":"What should I send?","options":null}"#,
13194 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13195 r#"{"question":"Should I send the report to Ada?"}"#,
13196 r#"{"status":"confirmed"}"#,
13197 "Confirmation cancelled.",
13198 ],
13199 true,
13200 None,
13201 true,
13202 vec![confirmation_skill()],
13203 );
13204 let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13205 let manager = agent
13206 .disambiguation_manager
13207 .take()
13208 .unwrap()
13209 .with_clarification_observer(observer.clone());
13210 agent.disambiguation_manager = Some(manager);
13211 let agent = Arc::new(agent);
13212
13213 agent.chat("Send it").await.unwrap();
13214 agent.chat("The report to Ada").await.unwrap();
13215
13216 let confirming_agent = Arc::clone(&agent);
13217 let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13218 observer.entered.wait().await;
13219 agent.reset().await.unwrap();
13220 observer.release.notify_one();
13221
13222 let response = confirmation.await.unwrap().unwrap();
13223 assert_eq!(response.content, "Confirmation cancelled.");
13224 assert_eq!(confirmation_skill_call_count(&observed), 0);
13225 assert!(agent.pending_skill_id.read().is_none());
13226 }
13227
13228 #[tokio::test]
13230 async fn queued_reset_prevents_stale_confirmation_question_publication() {
13231 let (agent, observed) = state_disambiguation_agent(
13232 vec![
13233 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13234 r#"{"question":"What should I send?","options":null}"#,
13235 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13236 r#"{"question":"Should I send the report to Ada?"}"#,
13237 ],
13238 true,
13239 None,
13240 true,
13241 );
13242 let agent = Arc::new(agent);
13243 agent.chat("Send it").await.unwrap();
13244
13245 let admission = agent.disambiguation_admission.write().await;
13246 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13247 let resetting_agent = Arc::clone(&agent);
13248 let reset = tokio::spawn(async move {
13249 let _ = started_tx.send(());
13250 resetting_agent.reset().await
13251 });
13252 started_rx.await.unwrap();
13253 tokio::task::yield_now().await;
13254
13255 let responding_agent = Arc::clone(&agent);
13256 let response =
13257 tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13258 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13259 while observed.call_count() < 4 {
13260 tokio::task::yield_now().await;
13261 }
13262 })
13263 .await
13264 .expect("clarification processing must reach terminal publication");
13265 drop(admission);
13266
13267 reset.await.unwrap().unwrap();
13268 let error = response.await.unwrap().unwrap_err();
13269 assert!(error.to_string().contains("ownership changed"));
13270 assert!(
13271 !agent
13272 .disambiguation_manager()
13273 .unwrap()
13274 .has_pending_clarification()
13275 .await
13276 );
13277 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13278 }
13279
13280 #[tokio::test]
13282 async fn queued_reset_prevents_stale_skill_clarification_publication() {
13283 let (agent, observed) = state_disambiguation_agent_with_skills(
13284 vec![
13285 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13286 "send_report",
13287 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13288 r#"{"question":"What should I send?","options":null}"#,
13289 ],
13290 true,
13291 None,
13292 true,
13293 vec![confirmation_skill()],
13294 );
13295 let agent = Arc::new(agent);
13296 let admission = agent.disambiguation_admission.write().await;
13297 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13298 let resetting_agent = Arc::clone(&agent);
13299 let reset = tokio::spawn(async move {
13300 let _ = started_tx.send(());
13301 resetting_agent.reset().await
13302 });
13303 started_rx.await.unwrap();
13304 tokio::task::yield_now().await;
13305
13306 let responding_agent = Arc::clone(&agent);
13307 let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13308 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13309 while observed.call_count() < 4 {
13310 tokio::task::yield_now().await;
13311 }
13312 })
13313 .await
13314 .expect("skill clarification must reach terminal publication");
13315 drop(admission);
13316
13317 reset.await.unwrap().unwrap();
13318 let error = response.await.unwrap().unwrap_err();
13319 assert!(error.to_string().contains("ownership changed"));
13320 assert_eq!(confirmation_skill_call_count(&observed), 0);
13321 assert!(agent.pending_skill_id.read().is_none());
13322 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13323 }
13324
13325 #[tokio::test]
13327 async fn transition_hook_can_reset_without_admission_deadlock() {
13328 let hooks = Arc::new(ResetOnTransitionHooks {
13329 agent: parking_lot::Mutex::new(None),
13330 invoked: AtomicBool::new(false),
13331 });
13332 let agent = Arc::new(
13333 AgentBuilder::new()
13334 .system_prompt("Test transition hook reentrancy.")
13335 .llm(Arc::new(mock_with_response("done")))
13336 .state_machine(disambiguation_state_machine(None, false))
13337 .build()
13338 .unwrap()
13339 .with_hooks(hooks.clone()),
13340 );
13341 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13342
13343 let transitioned = tokio::time::timeout(
13344 std::time::Duration::from_secs(2),
13345 agent.apply_transition_target("active", "review", "test transition", None),
13346 )
13347 .await
13348 .expect("transition hook reset must not deadlock")
13349 .unwrap();
13350
13351 assert!(transitioned);
13352 assert!(hooks.invoked.load(Ordering::SeqCst));
13353 assert_eq!(agent.current_state().as_deref(), Some("active"));
13354 }
13355
13356 #[tokio::test]
13358 async fn concurrent_transition_cannot_duplicate_exit_actions() {
13359 let gate = PathMutationGate::new();
13360 let active = ai_agents_state::StateDefinition {
13361 on_exit: vec![StateAction::Tool {
13362 tool: "transition_exit".to_string(),
13363 args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13364 }],
13365 ..Default::default()
13366 };
13367 let state_machine = Arc::new(
13368 StateMachine::new(ai_agents_state::StateConfig {
13369 initial: "active".to_string(),
13370 states: HashMap::from([
13371 ("active".to_string(), active),
13372 (
13373 "review".to_string(),
13374 ai_agents_state::StateDefinition::default(),
13375 ),
13376 ]),
13377 global_transitions: Vec::new(),
13378 fallback: None,
13379 max_no_transition: None,
13380 regenerate_on_transition: true,
13381 })
13382 .unwrap(),
13383 );
13384 let agent = Arc::new(
13385 AgentBuilder::new()
13386 .system_prompt("Test transition reservation.")
13387 .llm(Arc::new(mock_with_response("done")))
13388 .tool(Arc::new(BlockingPathMutationTool {
13389 id: "transition_exit",
13390 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13391 gate: gate.clone(),
13392 }))
13393 .state_machine(state_machine)
13394 .build()
13395 .unwrap(),
13396 );
13397
13398 let first_agent = Arc::clone(&agent);
13399 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13400 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13401 .await
13402 .expect("reserved transition must enter its exit action");
13403
13404 let second = tokio::time::timeout(
13405 std::time::Duration::from_secs(2),
13406 agent.transition_to("review"),
13407 )
13408 .await
13409 .expect("competing transition must fail without waiting for the exit action")
13410 .unwrap_err();
13411 assert!(second.to_string().contains("already in progress"));
13412
13413 gate.release();
13414 first.await.unwrap().unwrap();
13415 assert_eq!(agent.current_state().as_deref(), Some("review"));
13416 }
13417
13418 #[tokio::test]
13420 async fn concurrent_transition_cannot_overtake_enter_actions() {
13421 let gate = PathMutationGate::new();
13422 let review = ai_agents_state::StateDefinition {
13423 on_enter: vec![StateAction::Tool {
13424 tool: "transition_enter".to_string(),
13425 args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
13426 }],
13427 ..Default::default()
13428 };
13429 let state_machine = Arc::new(
13430 StateMachine::new(ai_agents_state::StateConfig {
13431 initial: "active".to_string(),
13432 states: HashMap::from([
13433 (
13434 "active".to_string(),
13435 ai_agents_state::StateDefinition::default(),
13436 ),
13437 ("review".to_string(), review),
13438 ]),
13439 global_transitions: Vec::new(),
13440 fallback: None,
13441 max_no_transition: None,
13442 regenerate_on_transition: true,
13443 })
13444 .unwrap(),
13445 );
13446 let agent = Arc::new(
13447 AgentBuilder::new()
13448 .system_prompt("Test transition lifecycle reservation.")
13449 .llm(Arc::new(mock_with_response("done")))
13450 .tool(Arc::new(BlockingPathMutationTool {
13451 id: "transition_enter",
13452 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13453 gate: gate.clone(),
13454 }))
13455 .state_machine(state_machine)
13456 .build()
13457 .unwrap(),
13458 );
13459
13460 let first_agent = Arc::clone(&agent);
13461 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13462 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13463 .await
13464 .expect("committed transition must enter its destination action");
13465
13466 let second = agent.transition_to("active").await.unwrap_err();
13467 assert!(second.to_string().contains("already in progress"));
13468 assert!(agent.reset().await.is_err());
13469
13470 gate.release();
13471 first.await.unwrap().unwrap();
13472 assert_eq!(agent.current_state().as_deref(), Some("review"));
13473 }
13474
13475 #[tokio::test]
13477 async fn same_state_restore_invalidates_pending_skill_confirmation() {
13478 let (agent, observed) = state_disambiguation_agent_with_skills(
13479 vec![
13480 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13481 "send_report",
13482 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13483 r#"{"question":"What should I send?","options":null}"#,
13484 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13485 r#"{"question":"Should I send the report to Ada?"}"#,
13486 ],
13487 true,
13488 None,
13489 true,
13490 vec![confirmation_skill()],
13491 );
13492
13493 agent.chat("Send it").await.unwrap();
13494 agent.chat("The report to Ada").await.unwrap();
13495 let snapshot = agent.save_state().await.unwrap();
13496 assert_eq!(agent.current_state().as_deref(), Some("active"));
13497
13498 agent.restore_state(snapshot).await.unwrap();
13499
13500 assert_eq!(agent.current_state().as_deref(), Some("active"));
13501 assert!(agent.pending_skill_id.read().is_none());
13502 assert!(
13503 !agent
13504 .disambiguation_manager()
13505 .unwrap()
13506 .has_pending_clarification()
13507 .await
13508 );
13509 assert_eq!(confirmation_skill_call_count(&observed), 0);
13510 }
13511
13512 #[tokio::test]
13514 async fn direct_state_generation_change_invalidates_confirmation() {
13515 let (agent, observed) = state_disambiguation_agent_with_skills(
13516 vec![
13517 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13518 "send_report",
13519 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13520 r#"{"question":"What should I send?","options":null}"#,
13521 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13522 r#"{"question":"Should I send the report to Ada?"}"#,
13523 "Confirmation cancelled.",
13524 ],
13525 true,
13526 None,
13527 true,
13528 vec![confirmation_skill()],
13529 );
13530
13531 agent.chat("Send it").await.unwrap();
13532 agent.chat("The report to Ada").await.unwrap();
13533 let state_machine = agent.state_machine().unwrap();
13534 state_machine
13535 .transition_to("review", "external test")
13536 .unwrap();
13537 state_machine
13538 .transition_to("active", "external test")
13539 .unwrap();
13540
13541 let response = agent.chat("Yes").await.unwrap();
13542
13543 assert_eq!(response.content, "Confirmation cancelled.");
13544 assert_eq!(confirmation_skill_call_count(&observed), 0);
13545 assert!(agent.pending_skill_id.read().is_none());
13546 }
13547
13548 #[tokio::test]
13549 async fn state_confirmation_does_not_add_a_question_for_clear_input() {
13550 let (agent, observed) = state_disambiguation_agent(
13551 vec![
13552 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
13553 "Request executed.",
13554 ],
13555 true,
13556 None,
13557 true,
13558 );
13559
13560 let response = agent.chat("Send the report to Ada").await.unwrap();
13561
13562 assert_eq!(response.content, "Request executed.");
13563 assert_eq!(observed.call_count(), 2);
13564 }
13565
13566 #[tokio::test]
13567 async fn state_override_cannot_activate_a_disabled_top_level_manager() {
13568 let (agent, observed) =
13569 state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
13570
13571 assert!(!agent.has_disambiguation());
13572 let response = agent.chat("Send it").await.unwrap();
13573
13574 assert_eq!(response.content, "Request executed.");
13575 assert_eq!(observed.call_count(), 1);
13576 }
13577
13578 #[tokio::test]
13579 async fn native_required_choice_executes_through_the_shared_tool_path() {
13580 let mut mock = MockLLMProvider::new("native-required");
13581 mock.set_tool_choice(Some(ToolChoice::Required));
13582 mock.add_response(
13583 LLMResponse::new("", FinishReason::ToolCall)
13584 .with_tool_calls(vec![ToolCall {
13585 id: "provider-call-1".to_string(),
13586 name: "calculator".to_string(),
13587 arguments: serde_json::json!({"expression": "2 + 2"}),
13588 }])
13589 .unwrap(),
13590 );
13591 mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
13592 let observed = mock.clone();
13593 let agent = AgentBuilder::new()
13594 .system_prompt("Use the calculator when needed.")
13595 .llm(Arc::new(mock))
13596 .tool(Arc::new(CalculatorTool::new()))
13597 .build()
13598 .unwrap();
13599
13600 let response = agent.chat("What is 2 + 2?").await.unwrap();
13601
13602 assert_eq!(response.content, "The answer is 4.");
13603 assert_eq!(
13604 response.tool_calls.as_ref().unwrap()[0].id,
13605 "provider-call-1"
13606 );
13607 let calls = observed.call_history();
13608 assert_eq!(calls.len(), 2);
13609 assert!(matches!(
13610 calls[0].request.as_ref().map(|request| &request.choice),
13611 Some(ToolChoice::Required)
13612 ));
13613 assert!(matches!(
13614 calls[1].request.as_ref().map(|request| &request.choice),
13615 Some(ToolChoice::Auto)
13616 ));
13617 }
13618
13619 #[tokio::test]
13620 async fn prompt_fallback_uses_one_corrective_retry() {
13621 let mut mock = MockLLMProvider::new("prompt-required");
13622 mock.set_tool_choice(Some(ToolChoice::Required));
13623 mock.set_native_tool_support(false);
13624 mock.set_responses(
13625 vec![
13626 "I can calculate that.".to_string(),
13627 r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
13628 "The answer is 4.".to_string(),
13629 ],
13630 false,
13631 );
13632 let observed = mock.clone();
13633 let agent = AgentBuilder::new()
13634 .system_prompt("Use tools.")
13635 .llm(Arc::new(mock))
13636 .tool(Arc::new(CalculatorTool::new()))
13637 .build()
13638 .unwrap();
13639
13640 let response = agent.chat("What is 2 + 2?").await.unwrap();
13641
13642 assert_eq!(response.content, "The answer is 4.");
13643 assert_eq!(observed.call_count(), 3);
13644 let corrective = &observed.call_history()[1].messages;
13645 assert!(
13646 corrective
13647 .last()
13648 .unwrap()
13649 .content
13650 .contains("previous response")
13651 );
13652 }
13653
13654 #[tokio::test]
13655 async fn prompt_fallback_fails_after_one_noncompliant_retry() {
13656 let mut mock = MockLLMProvider::new("prompt-required-failure");
13657 mock.set_tool_choice(Some(ToolChoice::Required));
13658 mock.set_native_tool_support(false);
13659 mock.set_responses(
13660 vec!["No tool.".to_string(), "Still no tool.".to_string()],
13661 false,
13662 );
13663 let observed = mock.clone();
13664 let agent = AgentBuilder::new()
13665 .system_prompt("Use tools.")
13666 .llm(Arc::new(mock))
13667 .tool(Arc::new(CalculatorTool::new()))
13668 .build()
13669 .unwrap();
13670
13671 let error = agent.chat("What is 2 + 2?").await.unwrap_err();
13672
13673 assert!(error.to_string().contains("one corrective retry"));
13674 assert_eq!(observed.call_count(), 2);
13675 }
13676
13677 #[tokio::test]
13678 async fn specific_choice_cannot_widen_the_effective_grant() {
13679 let mut mock = MockLLMProvider::new("specific-outside-grant");
13680 mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
13681 let observed = mock.clone();
13682 let agent = AgentBuilder::new()
13683 .system_prompt("Use tools.")
13684 .llm(Arc::new(mock))
13685 .tool(Arc::new(CalculatorTool::new()))
13686 .build()
13687 .unwrap();
13688
13689 let error = agent.chat("Generate a value.").await.unwrap_err();
13690
13691 assert!(error.to_string().contains("is not registered"));
13692 assert_eq!(observed.call_count(), 0);
13693 }
13694
13695 #[tokio::test]
13696 async fn none_choice_exposes_no_tool_protocol() {
13697 let mut mock = MockLLMProvider::new("no-tools");
13698 mock.set_tool_choice(Some(ToolChoice::None));
13699 mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
13700 let observed = mock.clone();
13701 let agent = AgentBuilder::new()
13702 .system_prompt("Answer directly.")
13703 .llm(Arc::new(mock))
13704 .tool(Arc::new(CalculatorTool::new()))
13705 .build()
13706 .unwrap();
13707
13708 let response = agent.chat("Hello").await.unwrap();
13709
13710 assert!(response.tool_calls.is_none());
13711 assert_eq!(observed.call_count(), 1);
13712 let call = observed.last_call().unwrap();
13713 assert!(call.request.is_none());
13714 assert!(
13715 call.messages
13716 .iter()
13717 .all(|message| !message.content.contains("Available tools:"))
13718 );
13719 }
13720
13721 struct RuntimeStorage {
13722 capabilities: Box<[StorageCapability]>,
13723 snapshots: RwLock<HashMap<String, AgentSnapshot>>,
13724 metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
13725 metadata_save_calls: AtomicU64,
13726 metadata_load_calls: AtomicU64,
13727 fail_metadata_save: AtomicBool,
13728 fail_metadata_load: AtomicBool,
13729 }
13730
13731 impl RuntimeStorage {
13732 fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
13733 Self {
13734 capabilities: capabilities.into_iter().collect(),
13735 snapshots: RwLock::new(HashMap::new()),
13736 metadata: RwLock::new(HashMap::new()),
13737 metadata_save_calls: AtomicU64::new(0),
13738 metadata_load_calls: AtomicU64::new(0),
13739 fail_metadata_save: AtomicBool::new(false),
13740 fail_metadata_load: AtomicBool::new(false),
13741 }
13742 }
13743 }
13744
13745 #[async_trait]
13746 impl AgentStorage for RuntimeStorage {
13747 fn supports(&self, capability: StorageCapability) -> bool {
13748 self.capabilities.contains(&capability)
13749 }
13750
13751 async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
13752 self.snapshots
13753 .write()
13754 .insert(session_id.to_string(), snapshot.clone());
13755 Ok(())
13756 }
13757
13758 async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
13759 Ok(self.snapshots.read().get(session_id).cloned())
13760 }
13761
13762 async fn delete(&self, session_id: &str) -> Result<()> {
13763 self.snapshots.write().remove(session_id);
13764 Ok(())
13765 }
13766
13767 async fn list_sessions(&self) -> Result<Vec<String>> {
13768 Ok(self.snapshots.read().keys().cloned().collect())
13769 }
13770
13771 async fn save_snapshot_with_metadata(
13772 &self,
13773 session_id: &str,
13774 snapshot: &AgentSnapshot,
13775 metadata: &ai_agents_core::SessionMetadata,
13776 ) -> Result<()> {
13777 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13778 if self.fail_metadata_save.load(Ordering::SeqCst) {
13779 return Err(AgentError::Persistence("metadata save failed".into()));
13780 }
13781 self.snapshots
13782 .write()
13783 .insert(session_id.to_string(), snapshot.clone());
13784 self.metadata
13785 .write()
13786 .insert(session_id.to_string(), metadata.clone());
13787 Ok(())
13788 }
13789
13790 async fn save_metadata(
13791 &self,
13792 session_id: &str,
13793 metadata: &ai_agents_core::SessionMetadata,
13794 ) -> Result<()> {
13795 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13796 if self.fail_metadata_save.load(Ordering::SeqCst) {
13797 return Err(AgentError::Persistence("metadata save failed".into()));
13798 }
13799 self.metadata
13800 .write()
13801 .insert(session_id.to_string(), metadata.clone());
13802 Ok(())
13803 }
13804
13805 async fn load_metadata(
13806 &self,
13807 session_id: &str,
13808 ) -> Result<Option<ai_agents_core::SessionMetadata>> {
13809 self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
13810 if self.fail_metadata_load.load(Ordering::SeqCst) {
13811 return Err(AgentError::Persistence("metadata load failed".into()));
13812 }
13813 Ok(self.metadata.read().get(session_id).cloned())
13814 }
13815 }
13816
13817 fn runtime_storage_agent() -> RuntimeAgent {
13818 AgentBuilder::new()
13819 .system_prompt("Test runtime storage integration.")
13820 .llm(Arc::new(mock_with_response("done")))
13821 .build()
13822 .unwrap()
13823 }
13824
13825 fn restore_spec(id: &str) -> crate::spec::AgentSpec {
13826 crate::spec::AgentSpec {
13827 name: id.to_string(),
13828 system_prompt: format!("Restore child {id}."),
13829 ..crate::spec::AgentSpec::default()
13830 }
13831 }
13832
13833 fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
13834 ai_agents_core::SpawnedAgentEntry {
13835 id: id.to_string(),
13836 name: id.to_string(),
13837 spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
13838 }
13839 }
13840
13841 fn restore_spawner(
13842 storage: Arc<RuntimeStorage>,
13843 max_agents: usize,
13844 ) -> (
13845 Arc<crate::spawner::AgentSpawner>,
13846 Arc<crate::spawner::AgentRegistry>,
13847 ) {
13848 let mut llms = LLMRegistry::new();
13849 llms.register("default", Arc::new(mock_with_response("done")));
13850 (
13851 Arc::new(
13852 crate::spawner::AgentSpawner::new()
13853 .with_shared_llms(llms)
13854 .with_shared_storage(storage)
13855 .with_max_agents(max_agents),
13856 ),
13857 Arc::new(crate::spawner::AgentRegistry::new()),
13858 )
13859 }
13860
13861 async fn save_restore_target(
13862 parent: &RuntimeAgent,
13863 storage: &RuntimeStorage,
13864 session_id: &str,
13865 entries: Vec<ai_agents_core::SpawnedAgentEntry>,
13866 ) {
13867 let mut snapshot = parent.save_state().await.unwrap();
13868 snapshot.spawned_agents = Some(entries);
13869 storage.save(session_id, &snapshot).await.unwrap();
13870 storage
13871 .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
13872 .await
13873 .unwrap();
13874 }
13875
13876 #[tokio::test]
13877 async fn storage_init_requires_storage_for_actor_facts() {
13878 let facts = ai_agents_facts::FactsConfig {
13879 enabled: true,
13880 ..Default::default()
13881 };
13882 let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
13883
13884 let error = agent.init_storage().await.unwrap_err();
13885 assert!(matches!(
13886 error,
13887 AgentError::Config(message)
13888 if message.contains("actor facts or actor memory")
13889 && message.contains("none is configured or injected")
13890 ));
13891 }
13892
13893 #[tokio::test]
13894 async fn storage_init_validates_actor_facts_capability() {
13895 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13896 let actor_memory = ai_agents_facts::ActorMemoryConfig {
13897 enabled: true,
13898 ..Default::default()
13899 };
13900 let agent = runtime_storage_agent()
13901 .with_storage(storage)
13902 .with_facts_config(Some(actor_memory), None);
13903
13904 assert!(matches!(
13905 agent.init_storage().await,
13906 Err(AgentError::UnsupportedStorageCapability(
13907 StorageCapability::ActorFacts
13908 ))
13909 ));
13910 }
13911
13912 #[tokio::test]
13913 async fn blocking_chat_rejects_unsupported_required_storage() {
13914 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13915 let facts = ai_agents_facts::FactsConfig {
13916 enabled: true,
13917 ..Default::default()
13918 };
13919 let agent = runtime_storage_agent()
13920 .with_storage(storage)
13921 .with_facts_config(None, Some(facts));
13922
13923 assert!(matches!(
13924 agent.chat("hello").await,
13925 Err(AgentError::UnsupportedStorageCapability(
13926 StorageCapability::ActorFacts
13927 ))
13928 ));
13929 }
13930
13931 #[tokio::test]
13932 async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
13933 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13934 let config = ai_agents_relationships::RelationshipConfig {
13935 enabled: true,
13936 ..Default::default()
13937 };
13938 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13939 let agent = runtime_storage_agent()
13940 .with_storage(storage)
13941 .with_relationships(manager);
13942
13943 assert!(matches!(
13944 agent.chat_stream("hello").await,
13945 Err(AgentError::UnsupportedStorageCapability(
13946 StorageCapability::ActorRelationships
13947 ))
13948 ));
13949 }
13950
13951 #[tokio::test]
13952 async fn storage_init_completes_facts_for_injected_storage() {
13953 let storage = Arc::new(RuntimeStorage::new([
13954 StorageCapability::Snapshot,
13955 StorageCapability::ActorFacts,
13956 ]));
13957 let facts = ai_agents_facts::FactsConfig {
13958 enabled: true,
13959 ..Default::default()
13960 };
13961 let agent = runtime_storage_agent()
13962 .with_storage(storage)
13963 .with_facts_config(None, Some(facts));
13964
13965 agent.init_storage().await.unwrap();
13966 assert!(agent.fact_store().is_some());
13967 }
13968
13969 #[tokio::test]
13970 async fn storage_init_requires_storage_for_persistent_relationships() {
13971 let config = ai_agents_relationships::RelationshipConfig {
13972 enabled: true,
13973 ..Default::default()
13974 };
13975 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13976 let agent = runtime_storage_agent().with_relationships(manager);
13977
13978 let error = agent.init_storage().await.unwrap_err();
13979 assert!(matches!(
13980 error,
13981 AgentError::Config(message)
13982 if message.contains("persistent relationships")
13983 && message.contains("none is configured or injected")
13984 ));
13985 }
13986
13987 #[tokio::test]
13988 async fn storage_init_validates_persistent_relationships_capability() {
13989 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13990 let config = ai_agents_relationships::RelationshipConfig {
13991 enabled: true,
13992 ..Default::default()
13993 };
13994 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13995 let agent = runtime_storage_agent()
13996 .with_storage(storage)
13997 .with_relationships(manager);
13998
13999 assert!(matches!(
14000 agent.init_storage().await,
14001 Err(AgentError::UnsupportedStorageCapability(
14002 StorageCapability::ActorRelationships
14003 ))
14004 ));
14005 }
14006
14007 #[tokio::test]
14008 async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
14009 let storage = Arc::new(RuntimeStorage::new([
14010 StorageCapability::Snapshot,
14011 StorageCapability::SessionMetadata,
14012 ]));
14013 let agent = runtime_storage_agent().with_storage(storage.clone());
14014 agent.set_actor_id("old-actor").unwrap();
14015 agent.save_session("old").await.unwrap();
14016 storage
14017 .save("target", &agent.save_state().await.unwrap())
14018 .await
14019 .unwrap();
14020 storage
14021 .save_metadata("target", &ai_agents_core::SessionMetadata::default())
14022 .await
14023 .unwrap();
14024
14025 assert!(agent.load_session("target").await.unwrap());
14026
14027 assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
14028 assert_eq!(agent.actor_id(), None);
14029 }
14030
14031 #[tokio::test]
14032 async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
14033 let storage = Arc::new(RuntimeStorage::new([
14034 StorageCapability::Snapshot,
14035 StorageCapability::SessionMetadata,
14036 ]));
14037 let (spawner, registry) = restore_spawner(storage.clone(), 3);
14038 let parent = runtime_storage_agent()
14039 .with_storage(storage.clone())
14040 .with_spawner_handles(Arc::clone(&spawner), Arc::clone(®istry));
14041
14042 for id in ["a", "b"] {
14043 let spawned = spawner
14044 .spawn_with_id(id.to_string(), restore_spec(id))
14045 .await
14046 .unwrap();
14047 spawned.agent.save_session("grow").await.unwrap();
14048 registry.register(spawned).await.unwrap();
14049 }
14050 let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
14051 staged_c
14052 .save("grow", &AgentSnapshot::new("c".into()))
14053 .await
14054 .unwrap();
14055 staged_c
14056 .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
14057 .await
14058 .unwrap();
14059 save_restore_target(
14060 &parent,
14061 storage.as_ref(),
14062 "grow",
14063 vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
14064 )
14065 .await;
14066
14067 assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
14068 assert_eq!(registry.count(), 3);
14069 assert!(registry.contains("c"));
14070 assert_eq!(spawner.spawned_count(), 3);
14071
14072 for id in ["a", "b"] {
14073 registry
14074 .get(id)
14075 .unwrap()
14076 .save_session("shrink")
14077 .await
14078 .unwrap();
14079 }
14080 save_restore_target(
14081 &parent,
14082 storage.as_ref(),
14083 "shrink",
14084 vec![restore_entry("a"), restore_entry("b")],
14085 )
14086 .await;
14087
14088 assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
14089 assert_eq!(registry.count(), 2);
14090 assert!(!registry.contains("c"));
14091 assert_eq!(spawner.spawned_count(), 2);
14092
14093 save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
14094
14095 assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
14096 assert_eq!(registry.count(), 0);
14097 assert_eq!(spawner.spawned_count(), 0);
14098 assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
14099 }
14100
14101 #[tokio::test]
14102 async fn storage_session_metadata_is_called_only_when_advertised() {
14103 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
14104 storage.fail_metadata_save.store(true, Ordering::SeqCst);
14105 storage.fail_metadata_load.store(true, Ordering::SeqCst);
14106 let agent = runtime_storage_agent().with_storage(storage.clone());
14107
14108 agent.save_session("session").await.unwrap();
14109 assert!(agent.load_session("session").await.unwrap());
14110 assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
14111 assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
14112 }
14113
14114 #[cfg(feature = "sqlite")]
14115 #[tokio::test]
14116 async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
14117 let directory =
14118 std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
14119 let path = directory.join("sessions.sqlite");
14120 let path_string = path.to_string_lossy().into_owned();
14121 let storage = Arc::new(
14122 ai_agents_storage::SqliteStorage::new(&path_string)
14123 .await
14124 .unwrap(),
14125 );
14126 let agent = runtime_storage_agent().with_storage(storage.clone());
14127 agent.set_session_metadata(ai_agents_core::SessionMetadata {
14128 tags: vec!["initial".into()],
14129 ..Default::default()
14130 });
14131 agent.chat("persist this turn").await.unwrap();
14132 agent.save_session("session").await.unwrap();
14133
14134 agent.set_session_metadata(ai_agents_core::SessionMetadata {
14135 tags: vec!["updated".into()],
14136 ..Default::default()
14137 });
14138 agent.save_session("session").await.unwrap();
14139 assert!(
14140 agent
14141 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14142 tags: Some(vec!["initial".into()]),
14143 ..Default::default()
14144 })
14145 .await
14146 .unwrap()
14147 .is_empty()
14148 );
14149 assert_eq!(
14150 agent
14151 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14152 tags: Some(vec!["updated".into()]),
14153 ..Default::default()
14154 })
14155 .await
14156 .unwrap()
14157 .len(),
14158 1
14159 );
14160 drop(agent);
14161 storage.close().await;
14162 drop(storage);
14163
14164 let reopened_storage = Arc::new(
14165 ai_agents_storage::SqliteStorage::new(&path_string)
14166 .await
14167 .unwrap(),
14168 );
14169 let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14170 assert!(restored.load_session("session").await.unwrap());
14171 assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14172 assert_eq!(
14173 restored.current_session_id.read().as_deref(),
14174 Some("session")
14175 );
14176 assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14177 assert_eq!(
14178 restored
14179 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14180 tags: Some(vec!["updated".into()]),
14181 ..Default::default()
14182 })
14183 .await
14184 .unwrap()
14185 .len(),
14186 1
14187 );
14188
14189 drop(restored);
14190 reopened_storage.close().await;
14191 drop(reopened_storage);
14192 crate::remove_sqlite_test_directory(&directory)
14193 .await
14194 .unwrap();
14195 }
14196
14197 #[tokio::test]
14198 async fn storage_session_metadata_backend_failures_propagate() {
14199 let storage = Arc::new(RuntimeStorage::new([
14200 StorageCapability::Snapshot,
14201 StorageCapability::SessionMetadata,
14202 ]));
14203 let agent = runtime_storage_agent().with_storage(storage.clone());
14204
14205 agent.save_session("session").await.unwrap();
14206 storage
14207 .save("target", &agent.save_state().await.unwrap())
14208 .await
14209 .unwrap();
14210 storage.fail_metadata_load.store(true, Ordering::SeqCst);
14211 assert!(matches!(
14212 agent.load_session("target").await,
14213 Err(AgentError::Persistence(message)) if message == "metadata load failed"
14214 ));
14215 assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14216
14217 storage.fail_metadata_save.store(true, Ordering::SeqCst);
14218 assert!(matches!(
14219 agent.save_session("session").await,
14220 Err(AgentError::Persistence(message)) if message == "metadata save failed"
14221 ));
14222 }
14223
14224 struct ProviderFutureDropSignal {
14225 dropped: Arc<AtomicBool>,
14226 }
14227
14228 impl Drop for ProviderFutureDropSignal {
14229 fn drop(&mut self) {
14230 self.dropped.store(true, Ordering::SeqCst);
14231 }
14232 }
14233
14234 struct BufferedLockingProvider {
14235 lock: Arc<tokio::sync::Mutex<()>>,
14236 stream_started: Arc<tokio::sync::Notify>,
14237 stream_dropped: Arc<AtomicBool>,
14238 committed_after_drop: Arc<AtomicBool>,
14239 }
14240
14241 #[async_trait]
14242 impl LLMProvider for BufferedLockingProvider {
14243 async fn complete(
14244 &self,
14245 _messages: &[ChatMessage],
14246 _config: Option<&LLMConfig>,
14247 ) -> std::result::Result<LLMResponse, LLMError> {
14248 let _guard = self.lock.lock().await;
14249 self.committed_after_drop
14250 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14251 Ok(LLMResponse::new(
14252 "Committed technical response.",
14253 FinishReason::Stop,
14254 ))
14255 }
14256
14257 async fn complete_stream(
14258 &self,
14259 _messages: &[ChatMessage],
14260 _config: Option<&LLMConfig>,
14261 ) -> std::result::Result<
14262 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14263 LLMError,
14264 > {
14265 let _guard = self.lock.lock().await;
14266 let _drop_signal = ProviderFutureDropSignal {
14267 dropped: Arc::clone(&self.stream_dropped),
14268 };
14269 self.stream_started.notify_one();
14270 std::future::pending().await
14271 }
14272
14273 fn provider_name(&self) -> &str {
14274 "buffered-locking"
14275 }
14276
14277 fn supports(&self, _feature: LLMFeature) -> bool {
14278 false
14279 }
14280 }
14281
14282 struct PendingDropStream {
14283 dropped: Arc<AtomicBool>,
14284 dropped_notify: Arc<tokio::sync::Notify>,
14285 }
14286
14287 impl Stream for PendingDropStream {
14288 type Item = std::result::Result<LLMChunk, LLMError>;
14289
14290 fn poll_next(
14291 self: Pin<&mut Self>,
14292 _cx: &mut std::task::Context<'_>,
14293 ) -> std::task::Poll<Option<Self::Item>> {
14294 std::task::Poll::Pending
14295 }
14296 }
14297
14298 impl Drop for PendingDropStream {
14299 fn drop(&mut self) {
14300 self.dropped.store(true, Ordering::SeqCst);
14301 self.dropped_notify.notify_one();
14302 }
14303 }
14304
14305 struct EstablishedStreamProvider {
14306 stream_started: Arc<tokio::sync::Notify>,
14307 stream_dropped: Arc<AtomicBool>,
14308 stream_dropped_notify: Arc<tokio::sync::Notify>,
14309 committed_after_drop: Arc<AtomicBool>,
14310 }
14311
14312 #[async_trait]
14313 impl LLMProvider for EstablishedStreamProvider {
14314 async fn complete(
14315 &self,
14316 _messages: &[ChatMessage],
14317 _config: Option<&LLMConfig>,
14318 ) -> std::result::Result<LLMResponse, LLMError> {
14319 if !self.stream_dropped.load(Ordering::SeqCst) {
14320 self.stream_dropped_notify.notified().await;
14321 }
14322 self.committed_after_drop
14323 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14324 Ok(LLMResponse::new(
14325 "Committed technical response.",
14326 FinishReason::Stop,
14327 ))
14328 }
14329
14330 async fn complete_stream(
14331 &self,
14332 _messages: &[ChatMessage],
14333 _config: Option<&LLMConfig>,
14334 ) -> std::result::Result<
14335 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14336 LLMError,
14337 > {
14338 self.stream_started.notify_one();
14339 Ok(Box::new(PendingDropStream {
14340 dropped: Arc::clone(&self.stream_dropped),
14341 dropped_notify: Arc::clone(&self.stream_dropped_notify),
14342 }))
14343 }
14344
14345 fn provider_name(&self) -> &str {
14346 "established-stream"
14347 }
14348
14349 fn supports(&self, _feature: LLMFeature) -> bool {
14350 false
14351 }
14352 }
14353
14354 struct FirstCallLockingProvider {
14355 lock: Arc<tokio::sync::Mutex<()>>,
14356 first_started: Arc<tokio::sync::Notify>,
14357 first_dropped: Arc<AtomicBool>,
14358 committed_after_drop: Arc<AtomicBool>,
14359 calls: AtomicU64,
14360 }
14361
14362 #[async_trait]
14363 impl LLMProvider for FirstCallLockingProvider {
14364 async fn complete(
14365 &self,
14366 _messages: &[ChatMessage],
14367 _config: Option<&LLMConfig>,
14368 ) -> std::result::Result<LLMResponse, LLMError> {
14369 let _guard = self.lock.lock().await;
14370 let call = self.calls.fetch_add(1, Ordering::SeqCst);
14371 if call == 0 {
14372 let _drop_signal = ProviderFutureDropSignal {
14373 dropped: Arc::clone(&self.first_dropped),
14374 };
14375 self.first_started.notify_one();
14376 return std::future::pending().await;
14377 }
14378 self.committed_after_drop
14379 .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14380 Ok(LLMResponse::new(
14381 "Committed technical response.",
14382 FinishReason::Stop,
14383 ))
14384 }
14385
14386 async fn complete_stream(
14387 &self,
14388 _messages: &[ChatMessage],
14389 _config: Option<&LLMConfig>,
14390 ) -> std::result::Result<
14391 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14392 LLMError,
14393 > {
14394 Err(LLMError::Other(
14395 "streaming is not used in this test".to_string(),
14396 ))
14397 }
14398
14399 fn provider_name(&self) -> &str {
14400 "first-call-locking"
14401 }
14402
14403 fn supports(&self, _feature: LLMFeature) -> bool {
14404 false
14405 }
14406 }
14407
14408 struct RoutingAfterProviderStart {
14409 provider_started: Arc<tokio::sync::Notify>,
14410 }
14411
14412 #[async_trait]
14413 impl LLMProvider for RoutingAfterProviderStart {
14414 async fn complete(
14415 &self,
14416 _messages: &[ChatMessage],
14417 _config: Option<&LLMConfig>,
14418 ) -> std::result::Result<LLMResponse, LLMError> {
14419 self.provider_started.notified().await;
14420 Ok(LLMResponse::new("1", FinishReason::Stop))
14421 }
14422
14423 async fn complete_stream(
14424 &self,
14425 _messages: &[ChatMessage],
14426 _config: Option<&LLMConfig>,
14427 ) -> std::result::Result<
14428 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14429 LLMError,
14430 > {
14431 Err(LLMError::Other(
14432 "streaming is not used in this test".to_string(),
14433 ))
14434 }
14435
14436 fn provider_name(&self) -> &str {
14437 "routing-after-start"
14438 }
14439
14440 fn supports(&self, _feature: LLMFeature) -> bool {
14441 false
14442 }
14443 }
14444
14445 struct ResponseCountingHooks {
14447 responses: Arc<std::sync::atomic::AtomicUsize>,
14448 }
14449
14450 struct ContextEchoTool;
14452
14453 #[async_trait]
14454 impl ai_agents_core::Tool for ContextEchoTool {
14455 fn id(&self) -> &str {
14456 "context_echo"
14457 }
14458
14459 fn name(&self) -> &str {
14460 "Context Echo"
14461 }
14462
14463 fn description(&self) -> &str {
14464 "Returns selected execution context fields."
14465 }
14466
14467 fn input_schema(&self) -> Value {
14468 serde_json::json!({"type": "object"})
14469 }
14470
14471 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14472 ai_agents_core::ToolPolicyBindings {
14473 path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
14474 result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
14475 "max_results",
14476 ai_agents_core::ResultLimitKind::MaxResults,
14477 )],
14478 ..Default::default()
14479 }
14480 }
14481
14482 async fn execute(
14483 &self,
14484 _args: Value,
14485 ctx: ai_agents_core::ToolExecutionContext,
14486 ) -> ToolResult {
14487 ToolResult::ok(
14488 serde_json::json!({
14489 "requested_name": ctx.requested_name,
14490 "canonical_id": ctx.canonical_id,
14491 "display_name": ctx.display_name,
14492 "max_results": ctx.limits.max_results,
14493 "custom_config": ctx.custom_config,
14494 })
14495 .to_string(),
14496 )
14497 }
14498 }
14499
14500 struct SlowTool;
14502
14503 struct FlakyWriteTool {
14505 calls: Arc<std::sync::atomic::AtomicUsize>,
14506 }
14507
14508 struct LockedWriteTool {
14510 active: Arc<std::sync::atomic::AtomicUsize>,
14511 max_active: Arc<std::sync::atomic::AtomicUsize>,
14512 }
14513
14514 struct MultiResourceWriteTool {
14515 active: Arc<std::sync::atomic::AtomicUsize>,
14516 max_active: Arc<std::sync::atomic::AtomicUsize>,
14517 }
14518
14519 #[derive(Clone)]
14520 struct PathMutationGate {
14521 entered: Arc<AtomicBool>,
14522 entered_notify: Arc<tokio::sync::Notify>,
14523 release: Arc<tokio::sync::Notify>,
14524 }
14525
14526 impl PathMutationGate {
14527 fn new() -> Self {
14528 Self {
14529 entered: Arc::new(AtomicBool::new(false)),
14530 entered_notify: Arc::new(tokio::sync::Notify::new()),
14531 release: Arc::new(tokio::sync::Notify::new()),
14532 }
14533 }
14534
14535 async fn wait_until_entered(&self) {
14536 if !self.entered.load(Ordering::SeqCst) {
14537 self.entered_notify.notified().await;
14538 }
14539 }
14540
14541 fn release(&self) {
14542 self.release.notify_one();
14543 }
14544 }
14545
14546 struct BlockingPathMutationTool {
14547 id: &'static str,
14548 path_fields: Vec<ai_agents_core::PathPolicyBinding>,
14549 gate: PathMutationGate,
14550 }
14551
14552 struct NoBindingWriteTool {
14553 active: Arc<std::sync::atomic::AtomicUsize>,
14554 max_active: Arc<std::sync::atomic::AtomicUsize>,
14555 }
14556
14557 struct RecoveryTestTool {
14558 id: String,
14559 succeeds: bool,
14560 calls: Arc<std::sync::atomic::AtomicUsize>,
14561 }
14562
14563 struct BlockingApprovalHandler {
14564 entered: Arc<tokio::sync::Barrier>,
14565 release: Arc<tokio::sync::Notify>,
14566 result: ApprovalResult,
14567 }
14568
14569 struct RuntimeWebFetchTransport {
14570 calls: Arc<std::sync::atomic::AtomicUsize>,
14571 }
14572
14573 struct RuntimeWebFetchResolver;
14574
14575 struct ReentrantToolHooks {
14576 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
14577 invoked: AtomicBool,
14578 nested_success: AtomicBool,
14579 }
14580
14581 #[async_trait]
14582 impl ai_agents_core::Tool for SlowTool {
14583 fn id(&self) -> &str {
14584 "slow"
14585 }
14586
14587 fn name(&self) -> &str {
14588 "Slow"
14589 }
14590
14591 fn description(&self) -> &str {
14592 "Waits until cancelled or timed out."
14593 }
14594
14595 fn input_schema(&self) -> Value {
14596 serde_json::json!({"type": "object"})
14597 }
14598
14599 async fn execute(
14600 &self,
14601 _args: Value,
14602 _ctx: ai_agents_core::ToolExecutionContext,
14603 ) -> ToolResult {
14604 tokio::time::sleep(std::time::Duration::from_secs(5)).await;
14605 ToolResult::ok("done")
14606 }
14607 }
14608
14609 #[async_trait]
14610 impl ai_agents_core::Tool for FlakyWriteTool {
14611 fn id(&self) -> &str {
14612 "flaky_write"
14613 }
14614
14615 fn name(&self) -> &str {
14616 "Flaky Write"
14617 }
14618
14619 fn description(&self) -> &str {
14620 "Fails on the first write attempt."
14621 }
14622
14623 fn input_schema(&self) -> Value {
14624 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14625 }
14626
14627 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14628 ai_agents_core::ToolPolicyBindings {
14629 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14630 ..Default::default()
14631 }
14632 }
14633
14634 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14635 ai_agents_core::ToolSafetyMetadata {
14636 read_only: false,
14637 concurrency_safe: false,
14638 operation: ai_agents_core::ToolOperationKind::Write,
14639 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14640 requires_network: false,
14641 destructive: false,
14642 open_world: false,
14643 host_dependent: false,
14644 requires_user_interaction: false,
14645 supports_cancellation: true,
14646 default_requires_approval: false,
14647 should_defer_schema: false,
14648 max_output_chars: Some(1024),
14649 max_result_size_chars: Some(1024),
14650 }
14651 }
14652
14653 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
14654 let mut classification =
14655 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
14656 classification.safely_retryable = false;
14657 classification
14658 }
14659
14660 async fn execute(
14661 &self,
14662 _args: Value,
14663 _ctx: ai_agents_core::ToolExecutionContext,
14664 ) -> ToolResult {
14665 let call = self.calls.fetch_add(1, Ordering::SeqCst);
14666 if call == 0 {
14667 ToolResult::error("first failure")
14668 } else {
14669 ToolResult::ok("second success")
14670 }
14671 }
14672 }
14673
14674 #[async_trait]
14675 impl ai_agents_core::Tool for LockedWriteTool {
14676 fn id(&self) -> &str {
14677 "locked_write"
14678 }
14679
14680 fn name(&self) -> &str {
14681 "Locked Write"
14682 }
14683
14684 fn description(&self) -> &str {
14685 "Tracks concurrent execution on one resource."
14686 }
14687
14688 fn input_schema(&self) -> Value {
14689 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14690 }
14691
14692 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14693 ai_agents_core::ToolPolicyBindings {
14694 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14695 ..Default::default()
14696 }
14697 }
14698
14699 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14700 ai_agents_core::ToolSafetyMetadata {
14701 read_only: false,
14702 concurrency_safe: false,
14703 operation: ai_agents_core::ToolOperationKind::Write,
14704 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14705 requires_network: false,
14706 destructive: false,
14707 open_world: false,
14708 host_dependent: false,
14709 requires_user_interaction: false,
14710 supports_cancellation: true,
14711 default_requires_approval: false,
14712 should_defer_schema: false,
14713 max_output_chars: Some(1024),
14714 max_result_size_chars: Some(1024),
14715 }
14716 }
14717
14718 async fn execute(
14719 &self,
14720 _args: Value,
14721 _ctx: ai_agents_core::ToolExecutionContext,
14722 ) -> ToolResult {
14723 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14724 loop {
14725 let current_max = self.max_active.load(Ordering::SeqCst);
14726 if active <= current_max {
14727 break;
14728 }
14729 if self
14730 .max_active
14731 .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
14732 .is_ok()
14733 {
14734 break;
14735 }
14736 }
14737 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
14738 self.active.fetch_sub(1, Ordering::SeqCst);
14739 ToolResult::ok("done")
14740 }
14741 }
14742
14743 #[async_trait]
14744 impl ai_agents_core::Tool for MultiResourceWriteTool {
14745 fn id(&self) -> &str {
14746 "multi_resource_write"
14747 }
14748
14749 fn name(&self) -> &str {
14750 "Multi Resource Write"
14751 }
14752
14753 fn description(&self) -> &str {
14754 "Tracks concurrent execution across source and destination resources."
14755 }
14756
14757 fn input_schema(&self) -> Value {
14758 serde_json::json!({"type": "object"})
14759 }
14760
14761 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14762 ai_agents_core::ToolPolicyBindings {
14763 path_fields: vec![
14764 ai_agents_core::PathPolicyBinding::read_write("source_path"),
14765 ai_agents_core::PathPolicyBinding::write("destination_path"),
14766 ],
14767 ..Default::default()
14768 }
14769 }
14770
14771 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14772 LockedWriteTool {
14773 active: Arc::clone(&self.active),
14774 max_active: Arc::clone(&self.max_active),
14775 }
14776 .safety_metadata()
14777 }
14778
14779 async fn execute(
14780 &self,
14781 _args: Value,
14782 _ctx: ai_agents_core::ToolExecutionContext,
14783 ) -> ToolResult {
14784 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14785 self.max_active.fetch_max(active, Ordering::SeqCst);
14786 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14787 self.active.fetch_sub(1, Ordering::SeqCst);
14788 ToolResult::ok("done")
14789 }
14790 }
14791
14792 #[async_trait]
14793 impl ai_agents_core::Tool for BlockingPathMutationTool {
14794 fn id(&self) -> &str {
14795 self.id
14796 }
14797
14798 fn name(&self) -> &str {
14799 self.id
14800 }
14801
14802 fn description(&self) -> &str {
14803 "Blocks a path mutation until the test releases it."
14804 }
14805
14806 fn input_schema(&self) -> Value {
14807 serde_json::json!({"type": "object"})
14808 }
14809
14810 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14811 ai_agents_core::ToolPolicyBindings {
14812 path_fields: self.path_fields.clone(),
14813 ..Default::default()
14814 }
14815 }
14816
14817 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14818 ai_agents_core::ToolSafetyMetadata {
14819 read_only: false,
14820 concurrency_safe: false,
14821 operation: ai_agents_core::ToolOperationKind::Write,
14822 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14823 requires_network: false,
14824 destructive: false,
14825 open_world: false,
14826 host_dependent: false,
14827 requires_user_interaction: false,
14828 supports_cancellation: true,
14829 default_requires_approval: false,
14830 should_defer_schema: false,
14831 max_output_chars: Some(1024),
14832 max_result_size_chars: Some(1024),
14833 }
14834 }
14835
14836 async fn execute(
14837 &self,
14838 _args: Value,
14839 _ctx: ai_agents_core::ToolExecutionContext,
14840 ) -> ToolResult {
14841 self.gate.entered.store(true, Ordering::SeqCst);
14842 self.gate.entered_notify.notify_one();
14843 self.gate.release.notified().await;
14844 ToolResult::ok("done")
14845 }
14846 }
14847
14848 #[async_trait]
14849 impl ai_agents_core::Tool for NoBindingWriteTool {
14850 fn id(&self) -> &str {
14851 "no_binding_write"
14852 }
14853
14854 fn name(&self) -> &str {
14855 "No Binding Write"
14856 }
14857
14858 fn description(&self) -> &str {
14859 "Tracks concurrent execution without resource bindings."
14860 }
14861
14862 fn input_schema(&self) -> Value {
14863 serde_json::json!({"type": "object"})
14864 }
14865
14866 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14867 LockedWriteTool {
14868 active: Arc::clone(&self.active),
14869 max_active: Arc::clone(&self.max_active),
14870 }
14871 .safety_metadata()
14872 }
14873
14874 async fn execute(
14875 &self,
14876 _args: Value,
14877 _ctx: ai_agents_core::ToolExecutionContext,
14878 ) -> ToolResult {
14879 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14880 self.max_active.fetch_max(active, Ordering::SeqCst);
14881 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14882 self.active.fetch_sub(1, Ordering::SeqCst);
14883 ToolResult::ok("done")
14884 }
14885 }
14886
14887 #[async_trait]
14888 impl ai_agents_core::Tool for RecoveryTestTool {
14889 fn id(&self) -> &str {
14890 &self.id
14891 }
14892
14893 fn name(&self) -> &str {
14894 &self.id
14895 }
14896
14897 fn description(&self) -> &str {
14898 "Records recovery execution and returns a configured result."
14899 }
14900
14901 fn input_schema(&self) -> Value {
14902 serde_json::json!({"type": "object"})
14903 }
14904
14905 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14906 ai_agents_core::ToolPolicyBindings {
14907 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14908 ..Default::default()
14909 }
14910 }
14911
14912 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14913 ai_agents_core::ToolSafetyMetadata {
14914 read_only: false,
14915 concurrency_safe: false,
14916 operation: ai_agents_core::ToolOperationKind::Write,
14917 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14918 requires_network: false,
14919 destructive: false,
14920 open_world: false,
14921 host_dependent: false,
14922 requires_user_interaction: false,
14923 supports_cancellation: true,
14924 default_requires_approval: false,
14925 should_defer_schema: false,
14926 max_output_chars: Some(1024),
14927 max_result_size_chars: Some(1024),
14928 }
14929 }
14930
14931 async fn execute(
14932 &self,
14933 _args: Value,
14934 _ctx: ai_agents_core::ToolExecutionContext,
14935 ) -> ToolResult {
14936 self.calls.fetch_add(1, Ordering::SeqCst);
14937 if self.succeeds {
14938 ToolResult::ok(format!("{} succeeded", self.id))
14939 } else {
14940 ToolResult::error(format!("{} failed", self.id))
14941 }
14942 }
14943 }
14944
14945 #[async_trait]
14946 impl WebFetchTransport for RuntimeWebFetchTransport {
14947 async fn send(
14949 &self,
14950 _request: WebFetchTransportRequest,
14951 ) -> std::result::Result<WebFetchTransportResponse, String> {
14952 Err("validated addresses are required".to_string())
14953 }
14954
14955 async fn send_validated(
14957 &self,
14958 _request: WebFetchTransportRequest,
14959 _addresses: &[std::net::SocketAddr],
14960 ) -> std::result::Result<WebFetchTransportResponse, String> {
14961 self.calls.fetch_add(1, Ordering::SeqCst);
14962 Ok(WebFetchTransportResponse {
14963 status: 200,
14964 content_type: Some("text/plain".to_string()),
14965 location: None,
14966 body: b"approved".to_vec(),
14967 })
14968 }
14969 }
14970
14971 #[async_trait]
14972 impl WebFetchResolver for RuntimeWebFetchResolver {
14973 async fn resolve(
14975 &self,
14976 _host: &str,
14977 _port: u16,
14978 ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
14979 Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
14980 93, 184, 216, 34,
14981 ))])
14982 }
14983 }
14984
14985 #[async_trait]
14986 impl ApprovalHandler for BlockingApprovalHandler {
14987 async fn request_approval(
14988 &self,
14989 _request: ai_agents_hitl::ApprovalRequest,
14990 ) -> ApprovalResult {
14991 self.entered.wait().await;
14992 self.release.notified().await;
14993 self.result.clone()
14994 }
14995 }
14996
14997 #[async_trait]
14998 impl AgentHooks for ReentrantToolHooks {
14999 async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
15000 if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
15001 return;
15002 }
15003 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
15004 if let Some(agent) = agent {
15005 let result = agent
15006 .invoke_tool(ToolExecutionRequest::new(
15007 "nested-hook-call",
15008 "reentrant_write",
15009 serde_json::json!({"path": "./hook.txt"}),
15010 ToolCallSource::Manual,
15011 ))
15012 .await;
15013 self.nested_success
15014 .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
15015 }
15016 }
15017 }
15018
15019 #[async_trait]
15020 impl AgentHooks for ResponseCountingHooks {
15021 async fn on_response(&self, _response: &AgentResponse) {
15022 self.responses.fetch_add(1, Ordering::SeqCst);
15023 }
15024 }
15025
15026 struct ApprovalRecordingHooks {
15027 events: parking_lot::Mutex<Vec<String>>,
15028 }
15029
15030 impl ApprovalRecordingHooks {
15031 fn new() -> Self {
15032 Self {
15033 events: parking_lot::Mutex::new(Vec::new()),
15034 }
15035 }
15036
15037 fn events(&self) -> Vec<String> {
15038 self.events.lock().clone()
15039 }
15040 }
15041
15042 #[async_trait]
15043 impl AgentHooks for ApprovalRecordingHooks {
15044 async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
15045 self.events.lock().push(format!(
15046 "raw:{}:{}",
15047 request_id,
15048 approval_result_name(result)
15049 ));
15050 }
15051
15052 async fn on_approval_resolved(
15053 &self,
15054 request: &ai_agents_hitl::ApprovalRequest,
15055 raw_result: &ApprovalResult,
15056 outcome: &ApprovalResolvedOutcome,
15057 ) {
15058 self.events.lock().push(format!(
15059 "resolved:{}:{}:{}",
15060 request.id,
15061 approval_result_name(raw_result),
15062 approval_outcome_name(outcome)
15063 ));
15064 }
15065 }
15066
15067 fn approval_result_name(result: &ApprovalResult) -> &'static str {
15068 match result {
15069 ApprovalResult::Approved => "approved",
15070 ApprovalResult::Rejected { .. } => "rejected",
15071 ApprovalResult::Modified { .. } => "modified",
15072 ApprovalResult::Timeout => "timeout",
15073 }
15074 }
15075
15076 fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
15077 match outcome {
15078 ApprovalResolvedOutcome::Approved => "approved",
15079 ApprovalResolvedOutcome::Rejected { .. } => "rejected",
15080 ApprovalResolvedOutcome::Modified { .. } => "modified",
15081 ApprovalResolvedOutcome::Error { .. } => "error",
15082 }
15083 }
15084
15085 fn assert_correlated_approval_events(
15086 events: &[String],
15087 raw_status: &str,
15088 outcome_status: &str,
15089 ) {
15090 assert_eq!(events.len(), 2);
15091 let raw: Vec<_> = events[0].split(':').collect();
15092 let resolved: Vec<_> = events[1].split(':').collect();
15093 assert_eq!(raw[0], "raw");
15094 assert_eq!(resolved[0], "resolved");
15095 assert_eq!(raw[1], resolved[1]);
15096 assert_eq!(raw[2], raw_status);
15097 assert_eq!(resolved[2], raw_status);
15098 assert_eq!(resolved[3], outcome_status);
15099 }
15100
15101 fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
15102 let mut security = ToolSecurityConfig {
15103 enabled: true,
15104 fail_closed: true,
15105 ..Default::default()
15106 };
15107 let policy = ai_agents_tools::ToolPolicyConfig {
15108 enabled: policy_enabled,
15109 write_paths: vec![".".to_string()],
15110 require_confirmation: true,
15111 ..Default::default()
15112 };
15113 security.tools.insert("locked_write".to_string(), policy);
15114 security
15115 }
15116
15117 struct MutationTestWorkspace {
15118 root: std::path::PathBuf,
15119 }
15120
15121 impl MutationTestWorkspace {
15122 fn new() -> Self {
15123 let root = std::env::temp_dir().join(format!(
15124 "ai-agents-runtime-mutation-{}",
15125 uuid::Uuid::new_v4()
15126 ));
15127 std::fs::create_dir_all(&root).unwrap();
15128 Self { root }
15129 }
15130 }
15131
15132 impl Drop for MutationTestWorkspace {
15133 fn drop(&mut self) {
15134 let _ = std::fs::remove_dir_all(&self.root);
15135 }
15136 }
15137
15138 async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
15139 tokio::time::timeout(std::time::Duration::from_secs(2), async {
15140 loop {
15141 let strong_count = locks
15142 .read()
15143 .get("path-mutation:global")
15144 .map_or(0, |lock| lock.strong_count());
15145 if strong_count >= minimum {
15146 break;
15147 }
15148 tokio::task::yield_now().await;
15149 }
15150 })
15151 .await
15152 .expect("path mutation call did not reach the shared lock");
15153 }
15154
15155 async fn assert_path_mutation_pair_serialized(
15156 first_id: &'static str,
15157 first_fields: Vec<ai_agents_core::PathPolicyBinding>,
15158 first_args: Value,
15159 second_id: &'static str,
15160 second_fields: Vec<ai_agents_core::PathPolicyBinding>,
15161 second_args: Value,
15162 ) {
15163 let locks = new_tool_resource_locks();
15164 let first_gate = PathMutationGate::new();
15165 let second_gate = PathMutationGate::new();
15166 second_gate.release();
15167 let agent = Arc::new(
15168 AgentBuilder::new()
15169 .system_prompt("Test global path mutation locking.")
15170 .llm(Arc::new(mock_with_response("done")))
15171 .tool(Arc::new(BlockingPathMutationTool {
15172 id: first_id,
15173 path_fields: first_fields,
15174 gate: first_gate.clone(),
15175 }))
15176 .tool(Arc::new(BlockingPathMutationTool {
15177 id: second_id,
15178 path_fields: second_fields,
15179 gate: second_gate.clone(),
15180 }))
15181 .build()
15182 .unwrap()
15183 .with_shared_resource_locks(Arc::clone(&locks)),
15184 );
15185
15186 let first = {
15187 let agent = Arc::clone(&agent);
15188 tokio::spawn(async move {
15189 agent
15190 .invoke_tool(ToolExecutionRequest::new(
15191 format!("{}-first", first_id),
15192 first_id,
15193 first_args,
15194 ToolCallSource::Manual,
15195 ))
15196 .await
15197 .unwrap()
15198 })
15199 };
15200 first_gate.wait_until_entered().await;
15201
15202 let second = {
15203 let agent = Arc::clone(&agent);
15204 tokio::spawn(async move {
15205 agent
15206 .invoke_tool(ToolExecutionRequest::new(
15207 format!("{}-second", second_id),
15208 second_id,
15209 second_args,
15210 ToolCallSource::Manual,
15211 ))
15212 .await
15213 .unwrap()
15214 })
15215 };
15216 wait_for_resource_lock_strong_count(&locks, 2).await;
15217 assert!(!second_gate.entered.load(Ordering::SeqCst));
15218 assert!(!second.is_finished());
15219
15220 first_gate.release();
15221 let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15222 tokio::join!(first, second)
15223 })
15224 .await
15225 .expect("serialized path mutation calls did not finish");
15226 assert!(first.unwrap().success);
15227 assert!(second.unwrap().success);
15228 assert!(second_gate.entered.load(Ordering::SeqCst));
15229 assert!(locks.read().is_empty());
15230 }
15231
15232 #[derive(Clone, Copy)]
15233 enum MutationDenial {
15234 Policy,
15235 Approval,
15236 }
15237
15238 fn mutation_denial_security_config(
15239 tool_id: &str,
15240 workspace: &std::path::Path,
15241 denial: MutationDenial,
15242 ) -> ToolSecurityConfig {
15243 let workspace = workspace.to_string_lossy().into_owned();
15244 let mut policy = ai_agents_tools::ToolPolicyConfig {
15245 read_paths: vec![workspace.clone()],
15246 write_paths: vec![workspace.clone()],
15247 ..Default::default()
15248 };
15249 match denial {
15250 MutationDenial::Policy => policy.blocked_paths = vec![workspace],
15251 MutationDenial::Approval => policy.require_confirmation = true,
15252 }
15253
15254 let mut security = ToolSecurityConfig {
15255 enabled: true,
15256 fail_closed: true,
15257 ..Default::default()
15258 };
15259 security.tools.insert(tool_id.to_string(), policy);
15260 security
15261 }
15262
15263 async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
15264 let workspace = MutationTestWorkspace::new();
15265 let tool_id = tool.id().to_string();
15266 let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
15267 let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
15268 std::fs::write(&preserved, "preserved").unwrap();
15269 let arguments = match tool_id.as_str() {
15270 "copy_path" | "move_path" => serde_json::json!({
15271 "source_path": preserved.to_string_lossy(),
15272 "destination_path": destination.to_string_lossy(),
15273 "dry_run": false
15274 }),
15275 "delete_path" => serde_json::json!({
15276 "path": preserved.to_string_lossy(),
15277 "recursive": false,
15278 "dry_run": false
15279 }),
15280 _ => panic!("unsupported mutation tool: {}", tool_id),
15281 };
15282 let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
15283 let builder = AgentBuilder::new()
15284 .system_prompt("Test mutation denial.")
15285 .llm(Arc::new(mock_with_response("done")))
15286 .tool(tool)
15287 .tool_security(ToolSecurityEngine::new(security));
15288 let builder = match denial {
15289 MutationDenial::Policy => builder,
15290 MutationDenial::Approval => builder
15291 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15292 .approval_handler(Arc::new(RejectAllHandler::new())),
15293 };
15294 let agent = builder.build().unwrap();
15295
15296 let record = agent
15297 .invoke_tool(ToolExecutionRequest::new(
15298 format!("{}-denied", tool_id),
15299 tool_id.clone(),
15300 arguments,
15301 ToolCallSource::Manual,
15302 ))
15303 .await
15304 .unwrap();
15305
15306 assert!(!record.executed, "{} must not be invoked", tool_id);
15307 assert!(!record.success);
15308 match denial {
15309 MutationDenial::Policy => {
15310 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
15311 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15312 &approval.status,
15313 ToolApprovalStatus::NotRequired
15314 )));
15315 }
15316 MutationDenial::Approval => {
15317 assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
15318 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15319 &approval.status,
15320 ToolApprovalStatus::Rejected
15321 )));
15322 }
15323 }
15324 assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
15325 assert!(!destination.exists());
15326 }
15327
15328 fn recovery_manager_with_fallbacks(
15329 fallbacks: impl IntoIterator<Item = (String, String)>,
15330 ) -> RecoveryManager {
15331 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15332
15333 let per_tool = fallbacks
15334 .into_iter()
15335 .map(|(tool, fallback_tool)| {
15336 (
15337 tool,
15338 ToolRetryConfig {
15339 max_retries: 0,
15340 timeout_ms: Some(1_000),
15341 on_failure: ToolFailureAction::Fallback { fallback_tool },
15342 },
15343 )
15344 })
15345 .collect();
15346 RecoveryManager::new(ErrorRecoveryConfig {
15347 tools: ToolRecoveryConfig {
15348 per_tool,
15349 ..Default::default()
15350 },
15351 ..Default::default()
15352 })
15353 }
15354
15355 fn approval_check() -> HITLCheckResult {
15356 HITLCheckResult::required(
15357 ApprovalTrigger::tool("test", serde_json::json!({})),
15358 HashMap::new(),
15359 "Approve?",
15360 None,
15361 )
15362 }
15363
15364 fn agent_with_approval_result(
15365 raw_result: ApprovalResult,
15366 timeout_action: TimeoutAction,
15367 hooks: Arc<ApprovalRecordingHooks>,
15368 ) -> RuntimeAgent {
15369 use ai_agents_hitl::{CallbackHandler, HITLConfig};
15370
15371 let config = HITLConfig {
15372 on_timeout: timeout_action,
15373 ..Default::default()
15374 };
15375 let handler = CallbackHandler::new(move |_| raw_result.clone());
15376 AgentBuilder::new()
15377 .system_prompt("Test HITL hooks.")
15378 .llm(Arc::new(mock_with_response("done")))
15379 .build()
15380 .unwrap()
15381 .with_hooks(hooks)
15382 .with_hitl(HITLEngine::new(config), Arc::new(handler))
15383 }
15384
15385 #[tokio::test]
15386 async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
15387 let cases = vec![
15388 (ApprovalResult::Approved, "approved"),
15389 (
15390 ApprovalResult::Rejected {
15391 reason: Some("denied".to_string()),
15392 },
15393 "rejected",
15394 ),
15395 (
15396 ApprovalResult::Modified {
15397 changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
15398 },
15399 "modified",
15400 ),
15401 ];
15402
15403 for (raw_result, expected) in cases {
15404 let hooks = Arc::new(ApprovalRecordingHooks::new());
15405 let agent =
15406 agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
15407
15408 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15409
15410 assert_eq!(approval_result_name(&result), expected);
15411 assert_correlated_approval_events(&hooks.events(), expected, expected);
15412 }
15413 }
15414
15415 #[tokio::test]
15416 async fn approval_hooks_expose_timeout_policy_decisions() {
15417 for (timeout_action, expected) in [
15418 (TimeoutAction::Approve, "approved"),
15419 (TimeoutAction::Reject, "rejected"),
15420 ] {
15421 let hooks = Arc::new(ApprovalRecordingHooks::new());
15422 let agent =
15423 agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
15424
15425 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15426
15427 assert_eq!(approval_result_name(&result), expected);
15428 assert_correlated_approval_events(&hooks.events(), "timeout", expected);
15429 }
15430 }
15431
15432 #[tokio::test]
15433 async fn timeout_error_fires_correlated_resolved_error_before_returning() {
15434 let hooks = Arc::new(ApprovalRecordingHooks::new());
15435 let agent = agent_with_approval_result(
15436 ApprovalResult::Timeout,
15437 TimeoutAction::Error,
15438 hooks.clone(),
15439 );
15440
15441 let error = agent
15442 .request_hitl_approval(approval_check())
15443 .await
15444 .unwrap_err();
15445
15446 assert!(error.to_string().contains("HITL approval timeout"));
15447 assert_correlated_approval_events(&hooks.events(), "timeout", "error");
15448 }
15449
15450 #[tokio::test]
15452 async fn test_integration_yaml_to_chat_basic() {
15453 let mock = mock_with_response("Hello! How can I help you?");
15454 let agent = AgentBuilder::new()
15455 .system_prompt("You are a test assistant.")
15456 .llm(Arc::new(mock))
15457 .build()
15458 .unwrap();
15459
15460 let response = agent.chat("Hi").await.unwrap();
15461 assert!(!response.content.is_empty());
15462 assert_eq!(response.content, "Hello! How can I help you?");
15463 }
15464
15465 #[tokio::test]
15466 async fn stream_events_emit_one_authoritative_final_without_legacy_done() {
15467 let agent = AgentBuilder::new()
15468 .system_prompt("You are a test assistant.")
15469 .llm(Arc::new(mock_with_response(
15470 "Hello from the final response.",
15471 )))
15472 .build()
15473 .unwrap();
15474
15475 let mut stream = agent.chat_stream_events("Hi").await.unwrap();
15476 let mut final_responses = Vec::new();
15477 let mut legacy_done = 0;
15478 while let Some(event) = stream.next().await {
15479 match event {
15480 AgentStreamEvent::Chunk(StreamChunk::Done {}) => legacy_done += 1,
15481 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
15482 panic!("unexpected stream error: {message}")
15483 }
15484 AgentStreamEvent::Final(response) => final_responses.push(response),
15485 AgentStreamEvent::Chunk(_) => {}
15486 }
15487 }
15488
15489 assert_eq!(legacy_done, 0);
15490 assert_eq!(final_responses.len(), 1);
15491 let response = final_responses.pop().unwrap();
15492 assert_eq!(response.content, "Hello from the final response.");
15493 assert!(
15494 response
15495 .metadata
15496 .as_ref()
15497 .is_some_and(|metadata| { metadata.contains_key("reasoning") })
15498 );
15499 }
15500
15501 #[tokio::test]
15502 async fn stream_final_content_includes_output_processing_after_provisional_chunks() {
15503 let yaml = r#"
15504name: ProcessedStreamAgent
15505system_prompt: "Answer directly."
15506process:
15507 output:
15508 - type: format
15509 config:
15510 template: "{{ response }} [finalized]"
15511streaming:
15512 enabled: true
15513"#;
15514 let agent = AgentBuilder::from_yaml(yaml)
15515 .unwrap()
15516 .llm(Arc::new(mock_with_response("provisional answer")))
15517 .auto_configure_features()
15518 .unwrap()
15519 .build()
15520 .unwrap();
15521
15522 let mut stream = agent.chat_stream_events("Hi").await.unwrap();
15523 let mut provisional = String::new();
15524 let mut final_content = None;
15525 while let Some(event) = stream.next().await {
15526 match event {
15527 AgentStreamEvent::Chunk(StreamChunk::Content { text }) => {
15528 provisional.push_str(&text)
15529 }
15530 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
15531 panic!("unexpected stream error: {message}")
15532 }
15533 AgentStreamEvent::Final(response) => final_content = Some(response.content),
15534 AgentStreamEvent::Chunk(_) => {}
15535 }
15536 }
15537
15538 assert_eq!(provisional, "provisional answer");
15539 assert_eq!(
15540 final_content.as_deref(),
15541 Some("provisional answer [finalized]")
15542 );
15543 }
15544
15545 #[tokio::test]
15546 async fn stream_events_preserve_tool_progress_and_final_tool_calls() {
15547 let agent = AgentBuilder::new()
15548 .system_prompt("Use the echo tool once, then answer.")
15549 .llm(Arc::new(mock_with_responses(vec![
15550 r#"{"tool":"echo","arguments":{"message":"hello"}}"#,
15551 "Echo completed.",
15552 ])))
15553 .tool(Arc::new(ai_agents_tools::EchoTool::new()))
15554 .build()
15555 .unwrap();
15556
15557 let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
15558 let mut starts = 0;
15559 let mut results = 0;
15560 let mut ends = 0;
15561 let mut final_response = None;
15562 while let Some(event) = stream.next().await {
15563 match event {
15564 AgentStreamEvent::Chunk(StreamChunk::ToolCallStart { name, .. }) => {
15565 assert_eq!(name, "echo");
15566 starts += 1;
15567 }
15568 AgentStreamEvent::Chunk(StreamChunk::ToolResult { name, success, .. }) => {
15569 assert_eq!(name, "echo");
15570 assert!(success);
15571 results += 1;
15572 }
15573 AgentStreamEvent::Chunk(StreamChunk::ToolCallEnd { .. }) => ends += 1,
15574 AgentStreamEvent::Chunk(StreamChunk::Error { message }) => {
15575 panic!("unexpected stream error: {message}")
15576 }
15577 AgentStreamEvent::Final(response) => final_response = Some(response),
15578 AgentStreamEvent::Chunk(_) => {}
15579 }
15580 }
15581
15582 assert_eq!((starts, results, ends), (1, 1, 1));
15583 let response = final_response.expect("tool stream must finalize");
15584 assert_eq!(response.content, "Echo completed.");
15585 assert_eq!(
15586 response.tool_calls.as_ref().map(|calls| calls
15587 .iter()
15588 .map(|call| call.name.as_str())
15589 .collect::<Vec<_>>()),
15590 Some(vec!["echo"])
15591 );
15592 }
15593
15594 #[tokio::test]
15595 async fn legacy_stream_still_emits_one_done_chunk() {
15596 let agent = AgentBuilder::new()
15597 .system_prompt("You are a test assistant.")
15598 .llm(Arc::new(mock_with_response(
15599 "Hello from the legacy stream.",
15600 )))
15601 .build()
15602 .unwrap();
15603
15604 let mut stream = agent.chat_stream("Hi").await.unwrap();
15605 let mut done = 0;
15606 while let Some(chunk) = stream.next().await {
15607 match chunk {
15608 StreamChunk::Done {} => done += 1,
15609 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
15610 _ => {}
15611 }
15612 }
15613
15614 assert_eq!(done, 1);
15615 }
15616
15617 #[tokio::test]
15619 async fn test_integration_multi_turn_conversation() {
15620 let mock = mock_with_responses(vec![
15621 "Hello! I'm your assistant.",
15622 "The weather is sunny today.",
15623 "Goodbye!",
15624 ]);
15625 let agent = AgentBuilder::new()
15626 .system_prompt("You are helpful.")
15627 .llm(Arc::new(mock))
15628 .build()
15629 .unwrap();
15630
15631 let r1 = agent.chat("Hi").await.unwrap();
15632 assert_eq!(r1.content, "Hello! I'm your assistant.");
15633
15634 let r2 = agent.chat("What's the weather?").await.unwrap();
15635 assert_eq!(r2.content, "The weather is sunny today.");
15636
15637 let r3 = agent.chat("Bye").await.unwrap();
15638 assert_eq!(r3.content, "Goodbye!");
15639
15640 let messages = agent.memory.get_messages(None).await.unwrap();
15642 assert_eq!(messages.len(), 6);
15644 }
15645
15646 #[test]
15647 fn later_approval_preserves_modified_evidence() {
15648 let arguments = serde_json::json!({"dry_run": true});
15649 let mut record = Some(ToolApprovalRecord {
15650 status: ToolApprovalStatus::Modified,
15651 reason: None,
15652 modified_arguments: Some(arguments.clone()),
15653 });
15654
15655 merge_approved_record(&mut record);
15656
15657 let record = record.unwrap();
15658 assert!(matches!(record.status, ToolApprovalStatus::Modified));
15659 assert_eq!(record.modified_arguments, Some(arguments));
15660 }
15661
15662 #[test]
15663 fn approval_binding_rejects_replaced_tool_implementation() {
15664 let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15665 let same_tool = Arc::clone(&reviewed_tool);
15666 let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15667 let arguments = serde_json::json!({"path": "."});
15668 let versions = ToolDecisionVersions {
15669 policy: 2,
15670 registry: 3,
15671 runtime_control: 4,
15672 state: Some(5),
15673 };
15674 let binding = ToolApprovalBinding {
15675 canonical_id: "context_echo".to_string(),
15676 arguments: arguments.clone(),
15677 confirmation_required: true,
15678 policy_version: versions.policy,
15679 runtime_control_version: versions.runtime_control,
15680 state_generation: versions.state,
15681 reviewed_tool,
15682 };
15683
15684 assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
15685 assert!(binding.is_stale(
15686 "context_echo",
15687 &arguments,
15688 true,
15689 versions,
15690 &replacement_tool,
15691 ));
15692 }
15693
15694 #[tokio::test]
15695 async fn approved_mutation_to_dry_run_remains_executable() {
15696 use ai_agents_hitl::CallbackHandler;
15697
15698 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
15699 changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
15700 });
15701 let agent = AgentBuilder::new()
15702 .system_prompt("Test safer approval modifications.")
15703 .llm(Arc::new(mock_with_response("done")))
15704 .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
15705 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15706 .approval_handler(Arc::new(handler))
15707 .build()
15708 .unwrap();
15709
15710 let record = agent
15711 .invoke_tool(ToolExecutionRequest::new(
15712 "approved-dry-run",
15713 "file_write",
15714 serde_json::json!({
15715 "path": "./approval-dry-run.txt",
15716 "content": "not written"
15717 }),
15718 ToolCallSource::Manual,
15719 ))
15720 .await
15721 .unwrap();
15722
15723 assert!(record.executed);
15724 assert!(record.success);
15725 assert_eq!(record.executed_arguments["dry_run"], true);
15726 assert!(matches!(
15727 record.approval.as_ref().map(|approval| &approval.status),
15728 Some(ToolApprovalStatus::Modified)
15729 ));
15730 let output: Value = serde_json::from_str(&record.output).unwrap();
15731 assert_eq!(output["mutation_performed"], false);
15732 }
15733
15734 #[tokio::test]
15736 async fn shared_executor_approval_reaches_web_fetch_transport() {
15737 use ai_agents_hitl::{CallbackHandler, HITLConfig};
15738 use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
15739
15740 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15741 let tool = WebFetchTool::with_transport_and_resolver(
15742 Arc::new(RuntimeWebFetchTransport {
15743 calls: Arc::clone(&calls),
15744 }),
15745 Arc::new(RuntimeWebFetchResolver),
15746 );
15747 let mut security = ToolSecurityConfig {
15748 enabled: true,
15749 fail_closed: true,
15750 ..Default::default()
15751 };
15752 security.tools.insert(
15753 "web_fetch".to_string(),
15754 ToolPolicyConfig {
15755 domains: DomainPolicyConfig {
15756 requires_approval: vec!["approval.test".to_string()],
15757 ..Default::default()
15758 },
15759 allowed_schemes: vec!["https".to_string()],
15760 allowed_ports: vec![443],
15761 ..Default::default()
15762 },
15763 );
15764 let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
15765 let agent = AgentBuilder::new()
15766 .system_prompt("Test approved web fetch execution.")
15767 .llm(Arc::new(mock_with_response("done")))
15768 .tool(Arc::new(tool))
15769 .tool_security(ToolSecurityEngine::new(security))
15770 .build()
15771 .unwrap()
15772 .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
15773
15774 let record = agent
15775 .invoke_tool(ToolExecutionRequest::new(
15776 "approved-web-fetch",
15777 "web_fetch",
15778 serde_json::json!({
15779 "url": "https://approval.test/page",
15780 "cache_ttl_seconds": 0
15781 }),
15782 ToolCallSource::Manual,
15783 ))
15784 .await
15785 .unwrap();
15786
15787 assert!(record.success);
15788 assert!(
15789 record
15790 .approval
15791 .as_ref()
15792 .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
15793 );
15794 assert_eq!(calls.load(Ordering::SeqCst), 1);
15795 }
15796
15797 #[tokio::test]
15798 async fn context_preserves_requested_and_canonical_identity() {
15799 let mock = mock_with_response("hello");
15800 let mut tools = ai_agents_tools::ToolRegistry::new();
15801 tools.register(Arc::new(ContextEchoTool)).unwrap();
15802
15803 let mut security = ToolSecurityConfig {
15804 enabled: true,
15805 fail_closed: true,
15806 ..Default::default()
15807 };
15808 let mut policy = ai_agents_tools::ToolPolicyConfig {
15809 read_paths: vec![".".to_string()],
15810 max_results: Some(7),
15811 ..Default::default()
15812 };
15813 policy
15814 .config
15815 .insert("backend".to_string(), serde_json::json!("memory"));
15816 security.tools.insert("context_echo".to_string(), policy);
15817
15818 let agent = AgentBuilder::new()
15819 .system_prompt("You are helpful.")
15820 .llm(Arc::new(mock))
15821 .tools(tools)
15822 .tool_security(ToolSecurityEngine::new(security))
15823 .build()
15824 .unwrap();
15825
15826 let record = agent
15827 .invoke_tool(ToolExecutionRequest::new(
15828 "ctx-call",
15829 "Context Echo",
15830 serde_json::json!({"path": ".", "max_results": 99}),
15831 ToolCallSource::Manual,
15832 ))
15833 .await
15834 .unwrap();
15835
15836 assert!(record.success);
15837 assert!(matches!(&record.source, ToolCallSource::Manual));
15838 assert_eq!(record.requested_name, "Context Echo");
15839 assert_eq!(record.canonical_id, "context_echo");
15840 assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
15841 assert_eq!(record.executed_arguments["max_results"], 7);
15842 let output: Value = serde_json::from_str(&record.output).unwrap();
15843 assert_eq!(output["requested_name"], "Context Echo");
15844 assert_eq!(output["canonical_id"], "context_echo");
15845 assert_eq!(output["max_results"], 7);
15846 assert_eq!(output["custom_config"]["backend"], "memory");
15847 assert!(record.metadata.contains_key("effective_limits"));
15848 assert!(record.metadata.contains_key("policy_snapshot"));
15849 }
15850
15851 #[tokio::test]
15852 async fn test_runtime_control_cancels_active_tool_call() {
15853 let mock = mock_with_response("hello");
15854 let agent = Arc::new(
15855 AgentBuilder::new()
15856 .system_prompt("You are helpful.")
15857 .llm(Arc::new(mock))
15858 .tool(Arc::new(SlowTool))
15859 .build()
15860 .unwrap(),
15861 );
15862 let control = agent.runtime_control();
15863 let running_agent = Arc::clone(&agent);
15864 let handle = tokio::spawn(async move {
15865 running_agent
15866 .invoke_tool(ToolExecutionRequest::new(
15867 "slow-call",
15868 "slow",
15869 serde_json::json!({}),
15870 ToolCallSource::Manual,
15871 ))
15872 .await
15873 .unwrap()
15874 });
15875
15876 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15877 control.cancel_all();
15878 let record = handle.await.unwrap();
15879
15880 assert!(record.executed);
15881 assert!(record.cancelled);
15882 assert!(!record.success);
15883 assert!(record.cancellation_reason.is_some());
15884 }
15885
15886 #[tokio::test]
15887 async fn non_idempotent_tool_calls_are_not_retried() {
15888 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15889
15890 let mock = mock_with_response("hello");
15891 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15892 let agent = AgentBuilder::new()
15893 .system_prompt("You are helpful.")
15894 .llm(Arc::new(mock))
15895 .tool(Arc::new(FlakyWriteTool {
15896 calls: Arc::clone(&calls),
15897 }))
15898 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
15899 tools: ToolRecoveryConfig {
15900 default: ToolRetryConfig {
15901 max_retries: 2,
15902 ..Default::default()
15903 },
15904 ..Default::default()
15905 },
15906 ..Default::default()
15907 }))
15908 .build()
15909 .unwrap();
15910
15911 let record = agent
15912 .invoke_tool(ToolExecutionRequest::new(
15913 "flaky-call",
15914 "flaky_write",
15915 serde_json::json!({"path": "./tmp.txt"}),
15916 ToolCallSource::Manual,
15917 ))
15918 .await
15919 .unwrap();
15920
15921 assert!(!record.success);
15922 assert_eq!(calls.load(Ordering::SeqCst), 1);
15923 }
15924
15925 #[tokio::test]
15926 async fn side_effecting_tools_are_serialized_per_resource() {
15927 let mock = mock_with_response("hello");
15928 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15929 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15930 let agent = Arc::new(
15931 AgentBuilder::new()
15932 .system_prompt("You are helpful.")
15933 .llm(Arc::new(mock))
15934 .tool(Arc::new(LockedWriteTool {
15935 active: Arc::clone(&active),
15936 max_active: Arc::clone(&max_active),
15937 }))
15938 .build()
15939 .unwrap(),
15940 );
15941
15942 let left = {
15943 let agent = Arc::clone(&agent);
15944 tokio::spawn(async move {
15945 agent
15946 .invoke_tool(ToolExecutionRequest::new(
15947 "lock-1",
15948 "locked_write",
15949 serde_json::json!({"path": "./same.txt"}),
15950 ToolCallSource::Manual,
15951 ))
15952 .await
15953 .unwrap()
15954 })
15955 };
15956 let right = {
15957 let agent = Arc::clone(&agent);
15958 tokio::spawn(async move {
15959 agent
15960 .invoke_tool(ToolExecutionRequest::new(
15961 "lock-2",
15962 "locked_write",
15963 serde_json::json!({"path": "./same.txt"}),
15964 ToolCallSource::Manual,
15965 ))
15966 .await
15967 .unwrap()
15968 })
15969 };
15970
15971 let left = left.await.unwrap();
15972 let right = right.await.unwrap();
15973 assert!(left.success);
15974 assert!(right.success);
15975 assert_eq!(max_active.load(Ordering::SeqCst), 1);
15976 }
15977
15978 #[tokio::test]
15979 async fn path_resources_use_shared_global_lock_and_cleanup() {
15980 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15981 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15982 let bindings = ai_agents_core::ToolPolicyBindings {
15983 path_fields: vec![
15984 ai_agents_core::PathPolicyBinding::read_write("source_path"),
15985 ai_agents_core::PathPolicyBinding::write("destination_path"),
15986 ],
15987 ..Default::default()
15988 };
15989 let classification = ai_agents_core::ToolCallClassification::from_metadata(
15990 &MultiResourceWriteTool {
15991 active: Arc::clone(&active),
15992 max_active: Arc::clone(&max_active),
15993 }
15994 .safety_metadata(),
15995 );
15996 let left_args = serde_json::json!({
15997 "source_path": "./a/../first.txt",
15998 "destination_path": "./second.txt"
15999 });
16000 let right_args = serde_json::json!({
16001 "source_path": "./second.txt",
16002 "destination_path": "./first.txt"
16003 });
16004 let left_keys = tool_resource_lock_keys(
16005 "multi_resource_write",
16006 &left_args,
16007 &bindings,
16008 &classification,
16009 );
16010 let right_keys = tool_resource_lock_keys(
16011 "multi_resource_write",
16012 &right_args,
16013 &bindings,
16014 &classification,
16015 );
16016 assert_eq!(left_keys, right_keys);
16017 assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
16018
16019 let locks = new_tool_resource_locks();
16020 let build_agent = || {
16021 AgentBuilder::new()
16022 .system_prompt("Test shared resource locks.")
16023 .llm(Arc::new(mock_with_response("done")))
16024 .tool(Arc::new(MultiResourceWriteTool {
16025 active: Arc::clone(&active),
16026 max_active: Arc::clone(&max_active),
16027 }))
16028 .build()
16029 .unwrap()
16030 .with_shared_resource_locks(Arc::clone(&locks))
16031 };
16032 let left_agent = Arc::new(build_agent());
16033 let right_agent = Arc::new(build_agent());
16034 let left = tokio::spawn(async move {
16035 left_agent
16036 .invoke_tool(ToolExecutionRequest::new(
16037 "multi-left",
16038 "multi_resource_write",
16039 left_args,
16040 ToolCallSource::Manual,
16041 ))
16042 .await
16043 .unwrap()
16044 });
16045 let right = tokio::spawn(async move {
16046 right_agent
16047 .invoke_tool(ToolExecutionRequest::new(
16048 "multi-right",
16049 "multi_resource_write",
16050 right_args,
16051 ToolCallSource::Manual,
16052 ))
16053 .await
16054 .unwrap()
16055 });
16056 let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
16057 tokio::join!(left, right)
16058 })
16059 .await
16060 .expect("reversed resource acquisition must not deadlock");
16061
16062 assert!(left.unwrap().success);
16063 assert!(right.unwrap().success);
16064 assert_eq!(max_active.load(Ordering::SeqCst), 1);
16065 assert!(locks.read().is_empty());
16066 }
16067
16068 #[tokio::test]
16069 async fn global_path_lock_serializes_copy_destination_with_file_write() {
16070 assert_path_mutation_pair_serialized(
16071 "copy_path",
16072 CopyPathTool::new().policy_bindings().path_fields,
16073 serde_json::json!({
16074 "source_path": "./source.txt",
16075 "destination_path": "./shared.txt"
16076 }),
16077 "file_write",
16078 FileWriteTool::new().policy_bindings().path_fields,
16079 serde_json::json!({"path": "./shared.txt"}),
16080 )
16081 .await;
16082 }
16083
16084 #[tokio::test]
16085 async fn parent_and_spawned_runtime_share_global_path_lock() {
16086 let workspace = MutationTestWorkspace::new();
16087 let destination = workspace.root.join("spawned.txt");
16088 let parent_gate = PathMutationGate::new();
16089 let parent = Arc::new(
16090 AgentBuilder::from_yaml(
16091 r#"
16092name: LockParent
16093system_prompt: parent
16094llm:
16095 default: default
16096tools:
16097 - parent_path_write
16098spawner:
16099 shared_llms: true
16100"#,
16101 )
16102 .unwrap()
16103 .llm(Arc::new(mock_with_response("done")))
16104 .auto_configure_spawner()
16105 .await
16106 .unwrap()
16107 .tool(Arc::new(BlockingPathMutationTool {
16108 id: "parent_path_write",
16109 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16110 gate: parent_gate.clone(),
16111 }))
16112 .build()
16113 .unwrap(),
16114 );
16115
16116 let mut child_spec = crate::spec::AgentSpec {
16117 name: "LockChild".to_string(),
16118 system_prompt: "child".to_string(),
16119 tools: Some(vec![crate::spec::ToolEntry::Simple(
16120 "file_write".to_string(),
16121 )]),
16122 ..Default::default()
16123 };
16124 child_spec.tool_security.enabled = true;
16125 child_spec.tool_security.fail_closed = true;
16126 let file_write_policy = ai_agents_tools::ToolPolicyConfig {
16127 write_paths: vec![workspace.root.to_string_lossy().into_owned()],
16128 allow_without_confirmation: true,
16129 ..Default::default()
16130 };
16131 child_spec
16132 .tool_security
16133 .tools
16134 .insert("file_write".to_string(), file_write_policy);
16135 let spawned = parent
16136 .spawner()
16137 .unwrap()
16138 .spawn_from_spec(child_spec)
16139 .await
16140 .unwrap();
16141 assert!(Arc::ptr_eq(
16142 &parent.resource_locks,
16143 &spawned.agent.resource_locks
16144 ));
16145 assert!(!Arc::ptr_eq(
16146 &parent.runtime_control,
16147 &spawned.agent.runtime_control
16148 ));
16149
16150 let parent_call = {
16151 let parent = Arc::clone(&parent);
16152 let destination = destination.clone();
16153 tokio::spawn(async move {
16154 parent
16155 .invoke_tool(ToolExecutionRequest::new(
16156 "parent-lock-holder",
16157 "parent_path_write",
16158 serde_json::json!({"path": destination}),
16159 ToolCallSource::Manual,
16160 ))
16161 .await
16162 .unwrap()
16163 })
16164 };
16165 parent_gate.wait_until_entered().await;
16166
16167 let child_call = {
16168 let child = Arc::clone(&spawned.agent);
16169 let destination = destination.clone();
16170 tokio::spawn(async move {
16171 child
16172 .invoke_tool(ToolExecutionRequest::new(
16173 "spawned-file-write",
16174 "file_write",
16175 serde_json::json!({
16176 "path": destination,
16177 "content": "spawned",
16178 "dry_run": false
16179 }),
16180 ToolCallSource::Manual,
16181 ))
16182 .await
16183 .unwrap()
16184 })
16185 };
16186 wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
16187 assert!(!child_call.is_finished());
16188
16189 parent_gate.release();
16190 let (parent_record, child_record) =
16191 tokio::time::timeout(std::time::Duration::from_secs(2), async {
16192 tokio::join!(parent_call, child_call)
16193 })
16194 .await
16195 .expect("parent and spawned path mutations did not finish");
16196 assert!(parent_record.unwrap().success);
16197 assert!(child_record.unwrap().success);
16198 assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
16199 assert!(parent.resource_locks.read().is_empty());
16200 }
16201
16202 #[tokio::test]
16203 async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
16204 let locks = new_tool_resource_locks();
16205 let holder_gate = PathMutationGate::new();
16206 let waiter_gate = PathMutationGate::new();
16207 waiter_gate.release();
16208 let holder = Arc::new(
16209 AgentBuilder::new()
16210 .system_prompt("Hold the global path lock.")
16211 .llm(Arc::new(mock_with_response("done")))
16212 .tool(Arc::new(BlockingPathMutationTool {
16213 id: "holder_write",
16214 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16215 gate: holder_gate.clone(),
16216 }))
16217 .build()
16218 .unwrap()
16219 .with_shared_resource_locks(Arc::clone(&locks)),
16220 );
16221 let waiter = Arc::new(
16222 AgentBuilder::new()
16223 .system_prompt("Wait for the global path lock.")
16224 .llm(Arc::new(mock_with_response("done")))
16225 .tool(Arc::new(BlockingPathMutationTool {
16226 id: "waiter_write",
16227 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
16228 gate: waiter_gate.clone(),
16229 }))
16230 .build()
16231 .unwrap()
16232 .with_shared_resource_locks(Arc::clone(&locks)),
16233 );
16234
16235 let holder_call = {
16236 let holder = Arc::clone(&holder);
16237 tokio::spawn(async move {
16238 holder
16239 .invoke_tool(ToolExecutionRequest::new(
16240 "holder-call",
16241 "holder_write",
16242 serde_json::json!({"path": "./shared.txt"}),
16243 ToolCallSource::Manual,
16244 ))
16245 .await
16246 .unwrap()
16247 })
16248 };
16249 holder_gate.wait_until_entered().await;
16250
16251 let waiter_call = {
16252 let waiter = Arc::clone(&waiter);
16253 tokio::spawn(async move {
16254 waiter
16255 .invoke_tool(ToolExecutionRequest::new(
16256 "waiter-call",
16257 "waiter_write",
16258 serde_json::json!({"path": "./shared.txt"}),
16259 ToolCallSource::Manual,
16260 ))
16261 .await
16262 .unwrap()
16263 })
16264 };
16265 wait_for_resource_lock_strong_count(&locks, 2).await;
16266 waiter.runtime_control().cancel_all();
16267
16268 let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
16269 .await
16270 .expect("cancelled lock waiter did not finish")
16271 .unwrap();
16272 assert!(!waiter_record.executed);
16273 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
16274 assert_eq!(
16275 locks
16276 .read()
16277 .get("path-mutation:global")
16278 .map_or(0, |lock| lock.strong_count()),
16279 1
16280 );
16281
16282 holder_gate.release();
16283 let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
16284 .await
16285 .expect("lock holder did not finish")
16286 .unwrap();
16287 assert!(holder_record.success);
16288 assert!(locks.read().is_empty());
16289 }
16290
16291 #[tokio::test]
16292 async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
16293 for denial in [MutationDenial::Policy, MutationDenial::Approval] {
16294 let tools: [Arc<dyn Tool>; 3] = [
16295 Arc::new(CopyPathTool::new()),
16296 Arc::new(MovePathTool::new()),
16297 Arc::new(DeletePathTool::new()),
16298 ];
16299 for tool in tools {
16300 assert_path_mutation_denied(tool, denial).await;
16301 }
16302 }
16303 }
16304
16305 #[tokio::test]
16306 async fn approval_argument_changes_are_rechecked_against_final_scope() {
16307 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16308 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16309 let entered = Arc::new(tokio::sync::Barrier::new(2));
16310 let release = Arc::new(tokio::sync::Notify::new());
16311 let handler = Arc::new(BlockingApprovalHandler {
16312 entered: Arc::clone(&entered),
16313 release: Arc::clone(&release),
16314 result: ApprovalResult::Modified {
16315 changes: HashMap::from([(
16316 "path".to_string(),
16317 Value::String("./after-approval.txt".to_string()),
16318 )]),
16319 },
16320 });
16321 let agent = Arc::new(
16322 AgentBuilder::new()
16323 .system_prompt("Test final scope validation.")
16324 .llm(Arc::new(mock_with_response("done")))
16325 .tool(Arc::new(LockedWriteTool {
16326 active: Arc::clone(&active),
16327 max_active: Arc::clone(&max_active),
16328 }))
16329 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16330 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16331 .approval_handler(handler)
16332 .build()
16333 .unwrap(),
16334 );
16335 let control = agent.runtime_control();
16336 let running = Arc::clone(&agent);
16337 let call = tokio::spawn(async move {
16338 running
16339 .invoke_tool(ToolExecutionRequest::new(
16340 "approval-scope",
16341 "locked_write",
16342 serde_json::json!({"path": "./before-approval.txt"}),
16343 ToolCallSource::Manual,
16344 ))
16345 .await
16346 .unwrap()
16347 });
16348 entered.wait().await;
16349 let expected_version = control.set_tool_scope(Vec::new());
16350 release.notify_one();
16351 let record = call.await.unwrap();
16352
16353 assert!(!record.executed);
16354 assert!(!record.success);
16355 assert_eq!(record.runtime_config_version, expected_version);
16356 assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
16357 assert_eq!(max_active.load(Ordering::SeqCst), 0);
16358 assert_eq!(
16359 record.metadata["runtime_scope_snapshot"],
16360 serde_json::json!([])
16361 );
16362 }
16363
16364 #[tokio::test]
16365 async fn approval_is_rechecked_against_final_policy_snapshot() {
16366 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16367 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16368 let entered = Arc::new(tokio::sync::Barrier::new(2));
16369 let release = Arc::new(tokio::sync::Notify::new());
16370 let handler = Arc::new(BlockingApprovalHandler {
16371 entered: Arc::clone(&entered),
16372 release: Arc::clone(&release),
16373 result: ApprovalResult::Approved,
16374 });
16375 let agent = Arc::new(
16376 AgentBuilder::new()
16377 .system_prompt("Test final policy validation.")
16378 .llm(Arc::new(mock_with_response("done")))
16379 .tool(Arc::new(LockedWriteTool {
16380 active: Arc::clone(&active),
16381 max_active: Arc::clone(&max_active),
16382 }))
16383 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16384 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16385 .approval_handler(handler)
16386 .build()
16387 .unwrap(),
16388 );
16389 let control = agent.runtime_control();
16390 let running = Arc::clone(&agent);
16391 let call = tokio::spawn(async move {
16392 running
16393 .invoke_tool(ToolExecutionRequest::new(
16394 "approval-policy",
16395 "locked_write",
16396 serde_json::json!({"path": "./policy.txt"}),
16397 ToolCallSource::Manual,
16398 ))
16399 .await
16400 .unwrap()
16401 });
16402 entered.wait().await;
16403 let expected_version = control.set_tool_security(approval_security_config(false));
16404 release.notify_one();
16405 let record = call.await.unwrap();
16406
16407 assert!(!record.executed);
16408 assert!(!record.success);
16409 assert_eq!(record.runtime_config_version, expected_version);
16410 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16411 assert_eq!(max_active.load(Ordering::SeqCst), 0);
16412 assert!(record.metadata.contains_key("policy_snapshot"));
16413 }
16414
16415 #[test]
16416 fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
16417 let agent = AgentBuilder::new()
16418 .system_prompt("Test runtime policy validation.")
16419 .llm(Arc::new(mock_with_response("done")))
16420 .build()
16421 .unwrap();
16422 let control = agent.runtime_control();
16423 let mut valid = ToolSecurityConfig::default();
16424 valid.tools.insert(
16425 "web_search".to_string(),
16426 ai_agents_tools::ToolPolicyConfig {
16427 max_results: Some(5),
16428 ..Default::default()
16429 },
16430 );
16431 let generation = control.try_set_tool_security(valid).unwrap();
16432
16433 let mut invalid = ToolSecurityConfig::default();
16434 invalid.tools.insert(
16435 "web_search".to_string(),
16436 ai_agents_tools::ToolPolicyConfig {
16437 max_results: Some(0),
16438 ..Default::default()
16439 },
16440 );
16441 let error = control.try_set_tool_security(invalid).unwrap_err();
16442
16443 assert!(
16444 error
16445 .to_string()
16446 .contains("max_results must be greater than 0")
16447 );
16448 assert_eq!(control.version(), generation);
16449 assert_eq!(
16450 control
16451 .state
16452 .tool_security_override
16453 .read()
16454 .as_ref()
16455 .unwrap()
16456 .config()
16457 .tools["web_search"]
16458 .max_results,
16459 Some(5)
16460 );
16461 }
16462
16463 #[tokio::test]
16464 async fn persistent_override_preserves_rate_history_within_generation() {
16465 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16466 let agent = AgentBuilder::new()
16467 .system_prompt("Test persistent policy overrides.")
16468 .llm(Arc::new(mock_with_response("done")))
16469 .tool(Arc::new(RecoveryTestTool {
16470 id: "limited_override".to_string(),
16471 succeeds: true,
16472 calls: Arc::clone(&calls),
16473 }))
16474 .build()
16475 .unwrap();
16476 let mut security = ToolSecurityConfig {
16477 enabled: true,
16478 fail_closed: true,
16479 ..Default::default()
16480 };
16481 let policy = ai_agents_tools::ToolPolicyConfig {
16482 write_paths: vec![".".to_string()],
16483 rate_limit: Some(1),
16484 ..Default::default()
16485 };
16486 security
16487 .tools
16488 .insert("limited_override".to_string(), policy);
16489 let generation = agent.runtime_control().set_tool_security(security);
16490
16491 let first = agent
16492 .invoke_tool(ToolExecutionRequest::new(
16493 "limited-first",
16494 "limited_override",
16495 serde_json::json!({"path": "./limited.txt"}),
16496 ToolCallSource::Manual,
16497 ))
16498 .await
16499 .unwrap();
16500 let second = agent
16501 .invoke_tool(ToolExecutionRequest::new(
16502 "limited-second",
16503 "limited_override",
16504 serde_json::json!({"path": "./limited.txt"}),
16505 ToolCallSource::Manual,
16506 ))
16507 .await
16508 .unwrap();
16509
16510 assert!(first.success);
16511 assert_eq!(first.policy_version, generation);
16512 assert!(!second.executed);
16513 assert!(second.output.contains("Rate limit exceeded"));
16514 assert_eq!(second.policy_version, generation);
16515 assert_eq!(calls.load(Ordering::SeqCst), 1);
16516 }
16517
16518 #[tokio::test]
16519 async fn concurrent_rate_admission_consumes_capacity_atomically() {
16520 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16521 let tool = Arc::new(RecoveryTestTool {
16522 id: "atomic_rate".to_string(),
16523 succeeds: true,
16524 calls: Arc::clone(&calls),
16525 });
16526 let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
16527 let bindings = tool.policy_bindings();
16528 let classification = tool.classify_call(&arguments);
16529 let resource_keys =
16530 tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
16531 let mut security = ToolSecurityConfig {
16532 enabled: true,
16533 fail_closed: true,
16534 ..Default::default()
16535 };
16536 let policy = ai_agents_tools::ToolPolicyConfig {
16537 write_paths: vec![".".to_string()],
16538 rate_limit: Some(1),
16539 ..Default::default()
16540 };
16541 security.tools.insert(tool.id().to_string(), policy);
16542 let agent = Arc::new(
16543 AgentBuilder::new()
16544 .system_prompt("Test atomic rate admission.")
16545 .llm(Arc::new(mock_with_response("done")))
16546 .tool(tool)
16547 .tool_security(ToolSecurityEngine::new(security))
16548 .build()
16549 .unwrap(),
16550 );
16551 let held = agent
16552 .acquire_tool_resource_locks(&resource_keys)
16553 .await
16554 .unwrap();
16555 let left = {
16556 let agent = Arc::clone(&agent);
16557 let arguments = arguments.clone();
16558 tokio::spawn(async move {
16559 agent
16560 .invoke_tool(ToolExecutionRequest::new(
16561 "atomic-rate-left",
16562 "atomic_rate",
16563 arguments,
16564 ToolCallSource::Manual,
16565 ))
16566 .await
16567 .unwrap()
16568 })
16569 };
16570 let right = {
16571 let agent = Arc::clone(&agent);
16572 tokio::spawn(async move {
16573 agent
16574 .invoke_tool(ToolExecutionRequest::new(
16575 "atomic-rate-right",
16576 "atomic_rate",
16577 arguments,
16578 ToolCallSource::Manual,
16579 ))
16580 .await
16581 .unwrap()
16582 })
16583 };
16584 tokio::time::sleep(std::time::Duration::from_millis(25)).await;
16585 drop(held);
16586 let (left, right) = tokio::join!(left, right);
16587 let records = [left.unwrap(), right.unwrap()];
16588
16589 assert_eq!(records.iter().filter(|record| record.success).count(), 1);
16590 assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
16591 assert!(
16592 records.iter().any(|record| {
16593 !record.executed && record.output.contains("Rate limit exceeded")
16594 })
16595 );
16596 assert_eq!(calls.load(Ordering::SeqCst), 1);
16597 }
16598
16599 #[tokio::test]
16600 async fn changed_policy_generation_invalidates_pending_approval() {
16601 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16602 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16603 let entered = Arc::new(tokio::sync::Barrier::new(2));
16604 let release = Arc::new(tokio::sync::Notify::new());
16605 let handler = Arc::new(BlockingApprovalHandler {
16606 entered: Arc::clone(&entered),
16607 release: Arc::clone(&release),
16608 result: ApprovalResult::Approved,
16609 });
16610 let agent = Arc::new(
16611 AgentBuilder::new()
16612 .system_prompt("Test stale approval denial.")
16613 .llm(Arc::new(mock_with_response("done")))
16614 .tool(Arc::new(LockedWriteTool {
16615 active: Arc::clone(&active),
16616 max_active: Arc::clone(&max_active),
16617 }))
16618 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16619 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16620 .approval_handler(handler)
16621 .build()
16622 .unwrap(),
16623 );
16624 let running = Arc::clone(&agent);
16625 let call = tokio::spawn(async move {
16626 running
16627 .invoke_tool(ToolExecutionRequest::new(
16628 "stale-approval",
16629 "locked_write",
16630 serde_json::json!({"path": "./stale.txt"}),
16631 ToolCallSource::Manual,
16632 ))
16633 .await
16634 .unwrap()
16635 });
16636 entered.wait().await;
16637 let generation = agent
16638 .runtime_control()
16639 .set_tool_security(approval_security_config(true));
16640 release.notify_one();
16641 let record = call.await.unwrap();
16642
16643 assert!(!record.executed);
16644 assert!(record.output.contains("Approval became stale"));
16645 assert_eq!(record.policy_version, generation);
16646 assert_eq!(max_active.load(Ordering::SeqCst), 0);
16647 }
16648
16649 #[tokio::test]
16650 async fn final_policy_reapplies_argument_caps_after_approval_changes() {
16651 use ai_agents_hitl::CallbackHandler;
16652
16653 let mut security = ToolSecurityConfig {
16654 enabled: true,
16655 fail_closed: true,
16656 ..Default::default()
16657 };
16658 let policy = ai_agents_tools::ToolPolicyConfig {
16659 read_paths: vec![".".to_string()],
16660 max_results: Some(5),
16661 require_confirmation: true,
16662 ..Default::default()
16663 };
16664 security.tools.insert("context_echo".to_string(), policy);
16665 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16666 changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
16667 });
16668 let agent = AgentBuilder::new()
16669 .system_prompt("Test final argument caps.")
16670 .llm(Arc::new(mock_with_response("done")))
16671 .tool(Arc::new(ContextEchoTool))
16672 .tool_security(ToolSecurityEngine::new(security))
16673 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16674 .approval_handler(Arc::new(handler))
16675 .build()
16676 .unwrap();
16677
16678 let record = agent
16679 .invoke_tool(ToolExecutionRequest::new(
16680 "final-cap",
16681 "context_echo",
16682 serde_json::json!({"path": ".", "max_results": 1}),
16683 ToolCallSource::Manual,
16684 ))
16685 .await
16686 .unwrap();
16687
16688 assert!(record.success);
16689 assert_eq!(record.executed_arguments["max_results"], 5);
16690 assert_eq!(
16691 record.approval.unwrap().modified_arguments.unwrap()["max_results"],
16692 5
16693 );
16694 }
16695
16696 #[tokio::test]
16697 async fn no_binding_writes_use_canonical_fallback_lock() {
16698 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16699 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16700 let agent = Arc::new(
16701 AgentBuilder::new()
16702 .system_prompt("Test fallback resource locks.")
16703 .llm(Arc::new(mock_with_response("done")))
16704 .tool(Arc::new(NoBindingWriteTool {
16705 active: Arc::clone(&active),
16706 max_active: Arc::clone(&max_active),
16707 }))
16708 .build()
16709 .unwrap(),
16710 );
16711 let left = {
16712 let agent = Arc::clone(&agent);
16713 tokio::spawn(async move {
16714 agent
16715 .invoke_tool(ToolExecutionRequest::new(
16716 "no-binding-left",
16717 "no_binding_write",
16718 serde_json::json!({}),
16719 ToolCallSource::Manual,
16720 ))
16721 .await
16722 .unwrap()
16723 })
16724 };
16725 let right = {
16726 let agent = Arc::clone(&agent);
16727 tokio::spawn(async move {
16728 agent
16729 .invoke_tool(ToolExecutionRequest::new(
16730 "no-binding-right",
16731 "no_binding_write",
16732 serde_json::json!({}),
16733 ToolCallSource::Manual,
16734 ))
16735 .await
16736 .unwrap()
16737 })
16738 };
16739 let (left, right) = tokio::join!(left, right);
16740
16741 assert!(left.unwrap().success);
16742 assert!(right.unwrap().success);
16743 assert_eq!(max_active.load(Ordering::SeqCst), 1);
16744 }
16745
16746 #[tokio::test]
16747 async fn parent_and_child_paths_share_a_resource_lock() {
16748 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16749 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16750 let agent = Arc::new(
16751 AgentBuilder::new()
16752 .system_prompt("Test parent child resource locks.")
16753 .llm(Arc::new(mock_with_response("done")))
16754 .tool(Arc::new(LockedWriteTool {
16755 active: Arc::clone(&active),
16756 max_active: Arc::clone(&max_active),
16757 }))
16758 .build()
16759 .unwrap(),
16760 );
16761 let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
16762 let child = format!("{}/child.txt", parent);
16763 let left = {
16764 let agent = Arc::clone(&agent);
16765 tokio::spawn(async move {
16766 agent
16767 .invoke_tool(ToolExecutionRequest::new(
16768 "parent-lock",
16769 "locked_write",
16770 serde_json::json!({"path": parent}),
16771 ToolCallSource::Manual,
16772 ))
16773 .await
16774 .unwrap()
16775 })
16776 };
16777 let right = {
16778 let agent = Arc::clone(&agent);
16779 tokio::spawn(async move {
16780 agent
16781 .invoke_tool(ToolExecutionRequest::new(
16782 "child-lock",
16783 "locked_write",
16784 serde_json::json!({"path": child}),
16785 ToolCallSource::Manual,
16786 ))
16787 .await
16788 .unwrap()
16789 })
16790 };
16791 let (left, right) = tokio::join!(left, right);
16792
16793 assert!(left.unwrap().success);
16794 assert!(right.unwrap().success);
16795 assert_eq!(max_active.load(Ordering::SeqCst), 1);
16796 }
16797
16798 #[tokio::test]
16799 async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
16800 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16801 let hooks = Arc::new(ReentrantToolHooks {
16802 agent: parking_lot::Mutex::new(None),
16803 invoked: AtomicBool::new(false),
16804 nested_success: AtomicBool::new(false),
16805 });
16806 let agent = Arc::new(
16807 AgentBuilder::new()
16808 .system_prompt("Test hook reentrancy.")
16809 .llm(Arc::new(mock_with_response("done")))
16810 .tool(Arc::new(RecoveryTestTool {
16811 id: "reentrant_write".to_string(),
16812 succeeds: true,
16813 calls: Arc::clone(&calls),
16814 }))
16815 .hooks(hooks.clone())
16816 .build()
16817 .unwrap(),
16818 );
16819 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
16820 let record = tokio::time::timeout(
16821 std::time::Duration::from_secs(2),
16822 agent.invoke_tool(ToolExecutionRequest::new(
16823 "outer-hook-call",
16824 "reentrant_write",
16825 serde_json::json!({"path": "./hook.txt"}),
16826 ToolCallSource::Manual,
16827 )),
16828 )
16829 .await
16830 .expect("tool completion hook must not retain resource guards")
16831 .unwrap();
16832
16833 assert!(record.success);
16834 assert!(hooks.nested_success.load(Ordering::SeqCst));
16835 assert_eq!(calls.load(Ordering::SeqCst), 2);
16836 }
16837
16838 #[tokio::test]
16839 async fn fallback_releases_primary_resource_locks() {
16840 let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16841 let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16842 let agent = AgentBuilder::new()
16843 .system_prompt("Test fallback execution.")
16844 .llm(Arc::new(mock_with_response("done")))
16845 .tool(Arc::new(RecoveryTestTool {
16846 id: "primary".to_string(),
16847 succeeds: false,
16848 calls: Arc::clone(&primary_calls),
16849 }))
16850 .tool(Arc::new(RecoveryTestTool {
16851 id: "fallback".to_string(),
16852 succeeds: true,
16853 calls: Arc::clone(&fallback_calls),
16854 }))
16855 .recovery_manager(recovery_manager_with_fallbacks([(
16856 "primary".to_string(),
16857 "fallback".to_string(),
16858 )]))
16859 .build()
16860 .unwrap();
16861 let record = tokio::time::timeout(
16862 std::time::Duration::from_secs(2),
16863 agent.invoke_tool(ToolExecutionRequest::new(
16864 "fallback-call",
16865 "primary",
16866 serde_json::json!({"path": "./shared.txt"}),
16867 ToolCallSource::Manual,
16868 )),
16869 )
16870 .await
16871 .expect("fallback must not retain the primary resource guard")
16872 .unwrap();
16873
16874 assert!(record.success);
16875 assert_eq!(record.canonical_id, "fallback");
16876 assert_eq!(record.call_id, "fallback-call");
16877 assert!(matches!(record.source, ToolCallSource::Fallback { .. }));
16878 assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
16879 assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
16880 }
16881
16882 #[tokio::test]
16883 async fn diagnostics_without_provider_records_unavailable_without_execution() {
16884 let mock = mock_with_response("hello");
16885 let yaml = r#"
16886name: DiagnosticsNoProviderAgent
16887system_prompt: "Review diagnostics."
16888tools: [diagnostics]
16889"#;
16890 let agent = AgentBuilder::from_yaml(yaml)
16891 .unwrap()
16892 .llm(Arc::new(mock))
16893 .auto_configure_features()
16894 .unwrap()
16895 .build()
16896 .unwrap();
16897
16898 let record = agent
16899 .invoke_tool(ToolExecutionRequest::new(
16900 "diagnostics-call",
16901 "diagnostics",
16902 serde_json::json!({}),
16903 ToolCallSource::Manual,
16904 ))
16905 .await
16906 .unwrap();
16907
16908 assert!(!record.executed);
16909 assert!(!record.success);
16910 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16911 }
16912
16913 #[tokio::test]
16914 async fn web_search_without_provider_records_unavailable_without_execution() {
16915 let mock = mock_with_response("hello");
16916 let yaml = r#"
16917name: WebSearchNoProviderAgent
16918system_prompt: "You search the web."
16919tools: [web_search]
16920"#;
16921 let agent = AgentBuilder::from_yaml(yaml)
16922 .unwrap()
16923 .llm(Arc::new(mock))
16924 .auto_configure_features()
16925 .unwrap()
16926 .build()
16927 .unwrap();
16928
16929 let record = agent
16930 .invoke_tool(ToolExecutionRequest::new(
16931 "web-search-call",
16932 "web_search",
16933 serde_json::json!({"query": "rust async"}),
16934 ToolCallSource::Manual,
16935 ))
16936 .await
16937 .unwrap();
16938
16939 assert!(!record.executed);
16940 assert!(!record.success);
16941 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16942 }
16943
16944 #[tokio::test]
16945 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
16946 let mock = mock_with_response("hello");
16947 let yaml = r#"
16948name: SpawnerNoGrantAgent
16949system_prompt: "You manage agents."
16950spawner:
16951 max_agents: 2
16952"#;
16953 let agent = AgentBuilder::from_yaml(yaml)
16954 .unwrap()
16955 .llm(Arc::new(mock))
16956 .auto_configure_features()
16957 .unwrap()
16958 .auto_configure_spawner()
16959 .await
16960 .unwrap()
16961 .build()
16962 .unwrap();
16963
16964 let available = agent.get_available_tool_ids().await.unwrap();
16965 assert!(available.is_empty());
16966 }
16967
16968 #[tokio::test]
16969 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
16970 let mock = mock_with_response("hello");
16971 let yaml = r#"
16972name: EmptySpawnerNoGrantAgent
16973system_prompt: "You manage agents."
16974tools: []
16975spawner:
16976 max_agents: 2
16977"#;
16978 let agent = AgentBuilder::from_yaml(yaml)
16979 .unwrap()
16980 .llm(Arc::new(mock))
16981 .auto_configure_features()
16982 .unwrap()
16983 .auto_configure_spawner()
16984 .await
16985 .unwrap()
16986 .build()
16987 .unwrap();
16988
16989 let available = agent.get_available_tool_ids().await.unwrap();
16990 assert!(available.is_empty());
16991 }
16992
16993 #[tokio::test]
16994 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
16995 let mock = mock_with_response("hello");
16996 let yaml = r#"
16997name: ManagementGrantAgent
16998system_prompt: "You manage agents."
16999tools: []
17000spawner:
17001 management_tools: true
17002"#;
17003 let agent = AgentBuilder::from_yaml(yaml)
17004 .unwrap()
17005 .llm(Arc::new(mock))
17006 .auto_configure_features()
17007 .unwrap()
17008 .auto_configure_spawner()
17009 .await
17010 .unwrap()
17011 .build()
17012 .unwrap();
17013
17014 let available = agent.get_available_tool_ids().await.unwrap();
17015 assert_eq!(available.len(), 4);
17016 assert!(available.contains(&"spawn_agent".to_string()));
17017 assert!(available.contains(&"send_agent_message".to_string()));
17018 assert!(available.contains(&"list_agents".to_string()));
17019 assert!(available.contains(&"remove_agent".to_string()));
17020 }
17021
17022 #[tokio::test]
17023 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
17024 let mock = mock_with_response("hello");
17025 let yaml = r#"
17026name: ManagementOmittedToolsGrantAgent
17027system_prompt: "You manage agents."
17028spawner:
17029 management_tools: true
17030"#;
17031 let agent = AgentBuilder::from_yaml(yaml)
17032 .unwrap()
17033 .llm(Arc::new(mock))
17034 .auto_configure_features()
17035 .unwrap()
17036 .auto_configure_spawner()
17037 .await
17038 .unwrap()
17039 .build()
17040 .unwrap();
17041
17042 let available = agent.get_available_tool_ids().await.unwrap();
17043 assert_eq!(available.len(), 4);
17044 assert!(available.contains(&"spawn_agent".to_string()));
17045 assert!(available.contains(&"send_agent_message".to_string()));
17046 assert!(available.contains(&"list_agents".to_string()));
17047 assert!(available.contains(&"remove_agent".to_string()));
17048 }
17049
17050 #[tokio::test]
17051 async fn test_management_tools_selected_grants_only_selected_tools() {
17052 let mock = mock_with_response("hello");
17053 let yaml = r#"
17054name: ManagementSelectedGrantAgent
17055system_prompt: "You manage agents."
17056tools: []
17057spawner:
17058 management_tools:
17059 - spawn_agent
17060 - send_agent_message
17061 - list_agents
17062"#;
17063 let agent = AgentBuilder::from_yaml(yaml)
17064 .unwrap()
17065 .llm(Arc::new(mock))
17066 .auto_configure_features()
17067 .unwrap()
17068 .auto_configure_spawner()
17069 .await
17070 .unwrap()
17071 .build()
17072 .unwrap();
17073
17074 let available = agent.get_available_tool_ids().await.unwrap();
17075 assert_eq!(available.len(), 3);
17076 assert!(available.contains(&"spawn_agent".to_string()));
17077 assert!(available.contains(&"send_agent_message".to_string()));
17078 assert!(available.contains(&"list_agents".to_string()));
17079 assert!(!available.contains(&"remove_agent".to_string()));
17080 }
17081
17082 #[tokio::test]
17083 async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
17084 let mock = mock_with_response("hello");
17085 let yaml = r#"
17086name: OrchestrationGrantAgent
17087system_prompt: "You coordinate agents."
17088llms:
17089 default:
17090 provider: openai
17091 model: gpt-4
17092 router:
17093 provider: openai
17094 model: gpt-4
17095llm:
17096 default: default
17097 router: router
17098tools: []
17099spawner:
17100 orchestration_tools: true
17101"#;
17102 let agent = AgentBuilder::from_yaml(yaml)
17103 .unwrap()
17104 .llm(Arc::new(mock))
17105 .auto_configure_features()
17106 .unwrap()
17107 .auto_configure_spawner()
17108 .await
17109 .unwrap()
17110 .build()
17111 .unwrap();
17112
17113 let available = agent.get_available_tool_ids().await.unwrap();
17114 assert_eq!(available.len(), 5);
17115 assert!(available.contains(&"route_to_agent".to_string()));
17116 assert!(available.contains(&"pipeline_process".to_string()));
17117 assert!(available.contains(&"concurrent_ask".to_string()));
17118 assert!(available.contains(&"group_discussion".to_string()));
17119 assert!(available.contains(&"handoff_conversation".to_string()));
17120 }
17121
17122 #[tokio::test]
17123 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
17124 let mock = mock_with_response("hello");
17125 let yaml = r#"
17126name: PersonaGrantAgent
17127system_prompt: "You can evolve persona."
17128llm:
17129 provider: openai
17130 model: gpt-4
17131tools: []
17132persona:
17133 identity:
17134 name: "Guide"
17135 role: "Helper"
17136 evolution:
17137 enabled: true
17138 allow_llm_evolve: true
17139 mutable_fields:
17140 - traits.personality
17141"#;
17142 let agent = AgentBuilder::from_yaml(yaml)
17143 .unwrap()
17144 .llm(Arc::new(mock))
17145 .build()
17146 .unwrap();
17147
17148 let available = agent.get_available_tool_ids().await.unwrap();
17149 assert_eq!(available, vec!["persona_evolve".to_string()]);
17150 }
17151
17152 #[tokio::test]
17153 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
17154 let mock = mock_with_response("hello");
17155 let yaml = r#"
17156name: PersonaOmittedToolsGrantAgent
17157system_prompt: "You can evolve persona."
17158llm:
17159 provider: openai
17160 model: gpt-4
17161persona:
17162 identity:
17163 name: "Guide"
17164 role: "Helper"
17165 evolution:
17166 enabled: true
17167 allow_llm_evolve: true
17168 mutable_fields:
17169 - traits.personality
17170"#;
17171 let agent = AgentBuilder::from_yaml(yaml)
17172 .unwrap()
17173 .llm(Arc::new(mock))
17174 .build()
17175 .unwrap();
17176
17177 let available = agent.get_available_tool_ids().await.unwrap();
17178 assert_eq!(available, vec!["persona_evolve".to_string()]);
17179 }
17180
17181 #[tokio::test]
17182 async fn test_omitted_yaml_tools_exposes_no_tools() {
17183 let mock = mock_with_response("hello");
17184 let yaml = r#"
17185name: NoToolsAgent
17186system_prompt: "You are helpful."
17187"#;
17188 let agent = AgentBuilder::from_yaml(yaml)
17189 .unwrap()
17190 .llm(Arc::new(mock))
17191 .auto_configure_features()
17192 .unwrap()
17193 .build()
17194 .unwrap();
17195
17196 let available = agent.get_available_tool_ids().await.unwrap();
17197 assert!(available.is_empty());
17198 }
17199
17200 #[tokio::test]
17201 async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
17202 for tools in ["", "tools: []"] {
17203 let yaml = format!(
17204 r#"
17205name: RuntimeScopeNoGrantAgent
17206system_prompt: "No ordinary tools are granted."
17207{tools}
17208"#
17209 );
17210 let agent = AgentBuilder::from_yaml(&yaml)
17211 .unwrap()
17212 .llm(Arc::new(mock_with_response("done")))
17213 .auto_configure_features()
17214 .unwrap()
17215 .build()
17216 .unwrap();
17217
17218 agent
17219 .runtime_control()
17220 .set_tool_scope(vec!["calculator".to_string()]);
17221
17222 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
17223 }
17224 }
17225
17226 #[tokio::test]
17227 async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
17228 let yaml = r#"
17229name: RuntimeScopeWideningAgent
17230system_prompt: "Runtime scope cannot add authority."
17231tools: [calculator]
17232"#;
17233 let agent = AgentBuilder::from_yaml(yaml)
17234 .unwrap()
17235 .llm(Arc::new(mock_with_response("done")))
17236 .auto_configure_features()
17237 .unwrap()
17238 .build()
17239 .unwrap();
17240
17241 agent
17242 .runtime_control()
17243 .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
17244
17245 assert_eq!(
17246 agent.get_available_tool_ids().await.unwrap(),
17247 vec!["calculator".to_string()]
17248 );
17249 }
17250
17251 #[tokio::test]
17252 async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
17253 let yaml = r#"
17254name: RuntimeScopeIntersectionAgent
17255system_prompt: "Use only declared tools."
17256tools: [calculator, datetime]
17257"#;
17258 let agent = AgentBuilder::from_yaml(yaml)
17259 .unwrap()
17260 .llm(Arc::new(mock_with_response("done")))
17261 .auto_configure_features()
17262 .unwrap()
17263 .build()
17264 .unwrap();
17265 let mut aliases = ai_agents_tools::ToolAliases::default();
17266 aliases
17267 .names
17268 .insert("en".to_string(), "calculate_alias".to_string());
17269 agent.tools.set_tool_aliases("calculator", aliases);
17270 let control = agent.runtime_control();
17271
17272 control.set_tool_scope(vec![
17273 "datetime".to_string(),
17274 "calculate_alias".to_string(),
17275 "calculator".to_string(),
17276 "unknown".to_string(),
17277 "datetime".to_string(),
17278 ]);
17279 assert_eq!(
17280 agent.get_available_tool_ids().await.unwrap(),
17281 vec!["calculator".to_string(), "datetime".to_string()]
17282 );
17283
17284 control.set_tool_scope(vec!["datetime".to_string()]);
17285 assert_eq!(
17286 agent.get_available_tool_ids().await.unwrap(),
17287 vec!["datetime".to_string()]
17288 );
17289
17290 control.clear_tool_scope_override();
17291 assert_eq!(
17292 agent.get_available_tool_ids().await.unwrap(),
17293 vec!["calculator".to_string(), "datetime".to_string()]
17294 );
17295 }
17296
17297 #[tokio::test]
17298 async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
17299 let agent = AgentBuilder::new()
17300 .system_prompt("Use registered tools.")
17301 .llm(Arc::new(mock_with_response("done")))
17302 .tool(Arc::new(ContextEchoTool))
17303 .tool(Arc::new(SlowTool))
17304 .build()
17305 .unwrap();
17306
17307 agent.runtime_control().set_tool_scope(vec![
17308 "Context Echo".to_string(),
17309 "context_echo".to_string(),
17310 "unknown".to_string(),
17311 ]);
17312
17313 assert_eq!(
17314 agent.get_available_tool_ids().await.unwrap(),
17315 vec!["context_echo".to_string()]
17316 );
17317 }
17318
17319 #[tokio::test]
17320 async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
17321 let yaml = r#"
17322name: NestedStateScopeAgent
17323system_prompt: "Honor every state scope."
17324tools: [calculator, datetime, echo]
17325states:
17326 initial: root
17327 states:
17328 root:
17329 tools: [calculate_alias, datetime]
17330 initial: middle
17331 states:
17332 middle:
17333 initial: leaf
17334 states:
17335 leaf:
17336 tools: [datetime_alias, echo]
17337"#;
17338 let agent = AgentBuilder::from_yaml(yaml)
17339 .unwrap()
17340 .llm(Arc::new(mock_with_response("done")))
17341 .auto_configure_features()
17342 .unwrap()
17343 .build()
17344 .unwrap();
17345 let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
17346 calculator_aliases
17347 .names
17348 .insert("en".to_string(), "calculate_alias".to_string());
17349 agent
17350 .tools
17351 .set_tool_aliases("calculator", calculator_aliases);
17352 let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
17353 datetime_aliases
17354 .names
17355 .insert("en".to_string(), "datetime_alias".to_string());
17356 agent.tools.set_tool_aliases("datetime", datetime_aliases);
17357 agent.runtime_control().set_tool_scope(vec![
17358 "unknown".to_string(),
17359 "datetime_alias".to_string(),
17360 "calculate_alias".to_string(),
17361 "datetime".to_string(),
17362 ]);
17363
17364 assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
17365 assert_eq!(
17366 agent.get_available_tool_ids().await.unwrap(),
17367 vec!["datetime".to_string()]
17368 );
17369 }
17370
17371 #[tokio::test]
17372 async fn ancestor_empty_state_scope_denies_omitted_descendants() {
17373 let yaml = r#"
17374name: NestedEmptyStateScopeAgent
17375system_prompt: "An empty ancestor scope denies all tools."
17376tools: [calculator]
17377states:
17378 initial: root
17379 states:
17380 root:
17381 tools: []
17382 initial: middle
17383 states:
17384 middle:
17385 initial: leaf
17386 states:
17387 leaf: {}
17388"#;
17389 let agent = AgentBuilder::from_yaml(yaml)
17390 .unwrap()
17391 .llm(Arc::new(mock_with_response("done")))
17392 .auto_configure_features()
17393 .unwrap()
17394 .build()
17395 .unwrap();
17396
17397 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
17398 }
17399
17400 #[tokio::test]
17401 async fn state_change_during_approval_invalidates_the_reviewed_authority() {
17402 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17403 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17404 let entered = Arc::new(tokio::sync::Barrier::new(2));
17405 let release = Arc::new(tokio::sync::Notify::new());
17406 let handler = Arc::new(BlockingApprovalHandler {
17407 entered: Arc::clone(&entered),
17408 release: Arc::clone(&release),
17409 result: ApprovalResult::Approved,
17410 });
17411 let yaml = r#"
17412name: ApprovalStateGenerationAgent
17413system_prompt: "State authority may change during approval."
17414tools: [locked_write]
17415states:
17416 initial: first
17417 states:
17418 first:
17419 tools: [locked_write]
17420 second:
17421 tools: [locked_write]
17422"#;
17423 let agent = Arc::new(
17424 AgentBuilder::from_yaml(yaml)
17425 .unwrap()
17426 .llm(Arc::new(mock_with_response("done")))
17427 .tool(Arc::new(LockedWriteTool {
17428 active: Arc::clone(&active),
17429 max_active: Arc::clone(&max_active),
17430 }))
17431 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17432 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17433 .approval_handler(handler)
17434 .build()
17435 .unwrap(),
17436 );
17437 let running = Arc::clone(&agent);
17438 let call = tokio::spawn(async move {
17439 running
17440 .invoke_tool(ToolExecutionRequest::new(
17441 "approval-state-generation",
17442 "locked_write",
17443 serde_json::json!({"path": "./state-generation.txt"}),
17444 ToolCallSource::Manual,
17445 ))
17446 .await
17447 .unwrap()
17448 });
17449
17450 entered.wait().await;
17451 agent.transition_to("second").await.unwrap();
17452 release.notify_one();
17453 let record = call.await.unwrap();
17454
17455 assert!(!record.executed);
17456 assert!(record.output.contains("Approval became stale"));
17457 assert_eq!(max_active.load(Ordering::SeqCst), 0);
17458 }
17459
17460 #[tokio::test]
17461 async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
17462 let holder_gate = PathMutationGate::new();
17463 let waiter_gate = PathMutationGate::new();
17464 let yaml = r#"
17465name: LockedStateGenerationAgent
17466system_prompt: "State authority must remain stable through admission."
17467tools: [state_lock_holder, state_lock_waiter]
17468states:
17469 initial: first
17470 states:
17471 first:
17472 tools: [state_lock_holder, state_lock_waiter]
17473 second:
17474 tools: [state_lock_holder, state_lock_waiter]
17475"#;
17476 let agent = Arc::new(
17477 AgentBuilder::from_yaml(yaml)
17478 .unwrap()
17479 .llm(Arc::new(mock_with_response("done")))
17480 .tool(Arc::new(BlockingPathMutationTool {
17481 id: "state_lock_holder",
17482 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17483 gate: holder_gate.clone(),
17484 }))
17485 .tool(Arc::new(BlockingPathMutationTool {
17486 id: "state_lock_waiter",
17487 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17488 gate: waiter_gate.clone(),
17489 }))
17490 .build()
17491 .unwrap(),
17492 );
17493 let holder_call = {
17494 let agent = Arc::clone(&agent);
17495 tokio::spawn(async move {
17496 agent
17497 .invoke_tool(ToolExecutionRequest::new(
17498 "state-lock-holder",
17499 "state_lock_holder",
17500 serde_json::json!({"path": "./shared-state-path.txt"}),
17501 ToolCallSource::Manual,
17502 ))
17503 .await
17504 .unwrap()
17505 })
17506 };
17507 holder_gate.wait_until_entered().await;
17508 let waiter_call = {
17509 let agent = Arc::clone(&agent);
17510 tokio::spawn(async move {
17511 agent
17512 .invoke_tool(ToolExecutionRequest::new(
17513 "state-lock-waiter",
17514 "state_lock_waiter",
17515 serde_json::json!({"path": "./shared-state-path.txt"}),
17516 ToolCallSource::Manual,
17517 ))
17518 .await
17519 .unwrap()
17520 })
17521 };
17522
17523 wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
17524 agent.transition_to("second").await.unwrap();
17525 holder_gate.release();
17526 let holder_record = holder_call.await.unwrap();
17527 let waiter_record = waiter_call.await.unwrap();
17528
17529 assert!(holder_record.success);
17530 assert!(!waiter_record.executed);
17531 assert!(
17532 waiter_record
17533 .output
17534 .contains("state scope changed before admission")
17535 );
17536 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
17537 }
17538
17539 #[tokio::test]
17540 async fn test_state_tools_cannot_widen_top_level_grant() {
17541 let mock = mock_with_response("hello");
17542 let yaml = r#"
17543name: NarrowToolsAgent
17544system_prompt: "You are helpful."
17545tools:
17546 - calculator
17547states:
17548 initial: current
17549 states:
17550 current:
17551 tools: [datetime]
17552"#;
17553 let agent = AgentBuilder::from_yaml(yaml)
17554 .unwrap()
17555 .llm(Arc::new(mock))
17556 .auto_configure_features()
17557 .unwrap()
17558 .build()
17559 .unwrap();
17560
17561 let available = agent.get_available_tool_ids().await.unwrap();
17562 assert!(available.is_empty());
17563 }
17564
17565 #[tokio::test]
17567 async fn test_integration_tool_execution() {
17568 let mock = mock_with_responses(vec![
17570 r#"I'll calculate that for you.
17572[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
17573 "The answer is 4.",
17575 ]);
17576 let mut tools = ai_agents_tools::ToolRegistry::new();
17577 tools
17578 .register(Arc::new(ai_agents_tools::CalculatorTool))
17579 .unwrap();
17580
17581 let agent = AgentBuilder::new()
17582 .system_prompt("You are a calculator assistant.")
17583 .llm(Arc::new(mock))
17584 .tools(tools)
17585 .build()
17586 .unwrap();
17587
17588 let response = agent.chat("What is 2+2?").await.unwrap();
17589 assert!(!response.content.is_empty());
17591 }
17592
17593 #[tokio::test]
17594 async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
17595 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17596 let hooks = Arc::new(ResponseCountingHooks {
17597 responses: Arc::clone(&responses),
17598 });
17599 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17600 let yaml = r#"
17601name: ToolRejectAgent
17602system_prompt: "You use tools when requested."
17603tools:
17604 - echo
17605hitl:
17606 tools:
17607 echo:
17608 require_approval: true
17609 approval_message: "Approve echo?"
17610"#;
17611 let agent = AgentBuilder::from_yaml(yaml)
17612 .unwrap()
17613 .llm(Arc::new(mock))
17614 .auto_configure_features()
17615 .unwrap()
17616 .hooks(hooks)
17617 .build()
17618 .unwrap();
17619
17620 let response = agent.chat("echo hello").await.unwrap();
17621
17622 assert!(
17623 response.content.contains("Operation cancelled"),
17624 "unexpected response: {}",
17625 response.content
17626 );
17627 assert_eq!(responses.load(Ordering::SeqCst), 1);
17628 let messages = agent.memory.get_messages(None).await.unwrap();
17629 assert_eq!(messages.len(), 3);
17630 assert_eq!(messages[0].content, "echo hello");
17631 assert!(messages[1].content.contains("\"tool\":\"echo\""));
17632 assert!(messages[2].content.contains("rejected by the approver"));
17633 }
17634
17635 #[tokio::test]
17636 async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
17637 use futures::StreamExt;
17638
17639 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17640 let hooks = Arc::new(ResponseCountingHooks {
17641 responses: Arc::clone(&responses),
17642 });
17643 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17644 let yaml = r#"
17645name: ToolRejectStreamingAgent
17646system_prompt: "You use tools when requested."
17647tools:
17648 - echo
17649streaming:
17650 enabled: true
17651hitl:
17652 tools:
17653 echo:
17654 require_approval: true
17655 approval_message: "Approve echo?"
17656"#;
17657 let agent = AgentBuilder::from_yaml(yaml)
17658 .unwrap()
17659 .llm(Arc::new(mock))
17660 .auto_configure_features()
17661 .unwrap()
17662 .hooks(hooks)
17663 .build()
17664 .unwrap();
17665
17666 let mut stream = agent.chat_stream("echo hello").await.unwrap();
17667 let mut terminal_error = String::new();
17668 let mut done = false;
17669 while let Some(chunk) = stream.next().await {
17670 match chunk {
17671 StreamChunk::Error { message } => terminal_error = message,
17672 StreamChunk::Done {} => {
17673 done = true;
17674 break;
17675 }
17676 _ => {}
17677 }
17678 }
17679
17680 assert!(done);
17681 assert!(
17682 terminal_error.contains("Operation cancelled"),
17683 "unexpected terminal error: {}",
17684 terminal_error
17685 );
17686 assert_eq!(responses.load(Ordering::SeqCst), 1);
17687 let messages = agent.memory.get_messages(None).await.unwrap();
17688 assert_eq!(messages.len(), 3);
17689 assert_eq!(messages[0].content, "echo hello");
17690 assert!(messages[1].content.contains("\"tool\":\"echo\""));
17691 assert!(messages[2].content.contains("rejected by the approver"));
17692 }
17693
17694 #[tokio::test]
17695 async fn tool_hitl_rejection_preserves_legacy_error_but_finalizes_event_stream() {
17696 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17697 let yaml = r#"
17698name: ToolRejectEventAgent
17699system_prompt: "You use tools when requested."
17700tools:
17701 - echo
17702streaming:
17703 enabled: true
17704hitl:
17705 tools:
17706 echo:
17707 require_approval: true
17708 approval_message: "Approve echo?"
17709"#;
17710 let agent = AgentBuilder::from_yaml(yaml)
17711 .unwrap()
17712 .llm(Arc::new(mock))
17713 .auto_configure_features()
17714 .unwrap()
17715 .build()
17716 .unwrap();
17717
17718 let mut stream = agent.chat_stream_events("echo hello").await.unwrap();
17719 let mut error_seen = false;
17720 let mut final_response = None;
17721 while let Some(event) = stream.next().await {
17722 match event {
17723 AgentStreamEvent::Chunk(StreamChunk::Error { .. }) => error_seen = true,
17724 AgentStreamEvent::Final(response) => final_response = Some(response),
17725 AgentStreamEvent::Chunk(_) => {}
17726 }
17727 }
17728
17729 assert!(!error_seen);
17730 assert!(
17731 final_response
17732 .is_some_and(|response| { response.content.contains("Operation cancelled") })
17733 );
17734 }
17735
17736 #[tokio::test]
17737 async fn test_pre_response_guard_transition_skips_old_state_llm() {
17738 let mock = mock_with_response("Billing state response");
17739 let call_counter = mock.clone();
17740 let yaml = r#"
17741name: OptimizedStateAgent
17742system_prompt: "You route before answering."
17743runtime:
17744 optimization:
17745 enabled: true
17746 pre_response_deterministic_transitions: true
17747states:
17748 initial: greeting
17749 states:
17750 greeting:
17751 prompt: "Old state prompt that should be skipped."
17752 transitions:
17753 - to: billing
17754 guard:
17755 context:
17756 topic:
17757 eq: billing
17758 timing: pre_response
17759 billing:
17760 prompt: "Answer from the billing state."
17761"#;
17762 let agent = AgentBuilder::from_yaml(yaml)
17763 .unwrap()
17764 .llm(Arc::new(mock))
17765 .build()
17766 .unwrap();
17767 agent
17768 .set_context("topic", serde_json::json!("billing"))
17769 .unwrap();
17770
17771 let response = agent.chat("I need billing help").await.unwrap();
17772
17773 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17774 assert_eq!(response.content, "Billing state response");
17775 assert_eq!(call_counter.call_count(), 1);
17776 assert_eq!(agent.actor_facts().len(), 0);
17777 }
17778
17779 #[tokio::test]
17780 async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
17781 let mock = mock_with_response("Billing state response");
17782 let call_counter = mock.clone();
17783 let yaml = r#"
17784name: OptimizedStateAgent
17785system_prompt: "You route before answering."
17786runtime:
17787 optimization:
17788 enabled: true
17789 pre_response_deterministic_transitions: true
17790context:
17791 request:
17792 type: runtime
17793 default:
17794 topic: general
17795states:
17796 initial: greeting
17797 states:
17798 greeting:
17799 prompt: "Old state prompt that should be skipped."
17800 transitions:
17801 - to: billing
17802 guard:
17803 context:
17804 request.topic:
17805 eq: billing
17806 timing: pre_response
17807 billing:
17808 prompt: "Answer from the billing state."
17809"#;
17810 let agent = AgentBuilder::from_yaml(yaml)
17811 .unwrap()
17812 .llm(Arc::new(mock))
17813 .build()
17814 .unwrap();
17815 agent
17816 .set_context("request.topic", serde_json::json!("billing"))
17817 .unwrap();
17818
17819 let response = agent.chat("I need billing help").await.unwrap();
17820
17821 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17822 assert_eq!(response.content, "Billing state response");
17823 assert_eq!(call_counter.call_count(), 1);
17824 assert_eq!(
17825 agent.get_context().get("request"),
17826 Some(&serde_json::json!({"topic": "billing"}))
17827 );
17828 }
17829
17830 #[tokio::test]
17831 async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
17832 let mock = mock_with_response("billing");
17833 let yaml = r#"
17834name: OptimizedStateAgent
17835system_prompt: "You route before answering."
17836runtime:
17837 optimization:
17838 enabled: true
17839 pre_response_deterministic_transitions: true
17840hitl:
17841 states:
17842 billing:
17843 on_enter: require_approval
17844 approval_message: "Approve billing route?"
17845states:
17846 initial: greeting
17847 states:
17848 greeting:
17849 prompt: "Old state prompt."
17850 extract:
17851 - key: topic
17852 description: "Support topic"
17853 transitions:
17854 - to: billing
17855 guard:
17856 context:
17857 topic:
17858 eq: billing
17859 timing: pre_response
17860 run_extractors: true
17861 billing:
17862 prompt: "Billing state."
17863"#;
17864 let agent = AgentBuilder::from_yaml(yaml)
17865 .unwrap()
17866 .llm(Arc::new(mock))
17867 .build()
17868 .unwrap();
17869
17870 let response = agent
17871 .try_pre_response_transition("billing please")
17872 .await
17873 .unwrap();
17874
17875 assert!(response.is_none());
17876 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17877 assert!(!agent.get_context().contains_key("topic"));
17878 assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
17879 }
17880
17881 #[tokio::test]
17882 async fn test_pre_response_extractor_commits_context_on_winning_path() {
17883 let mock = mock_with_responses(vec!["billing", "Billing response"]);
17884 let yaml = r#"
17885name: OptimizedStateAgent
17886system_prompt: "You route before answering."
17887runtime:
17888 optimization:
17889 enabled: true
17890 pre_response_deterministic_transitions: true
17891states:
17892 initial: greeting
17893 states:
17894 greeting:
17895 prompt: "Old state prompt."
17896 extract:
17897 - key: topic
17898 description: "Support topic"
17899 transitions:
17900 - to: billing
17901 guard:
17902 context:
17903 topic:
17904 eq: billing
17905 timing: pre_response
17906 run_extractors: true
17907 billing:
17908 prompt: "Billing state."
17909"#;
17910 let agent = AgentBuilder::from_yaml(yaml)
17911 .unwrap()
17912 .llm(Arc::new(mock))
17913 .build()
17914 .unwrap();
17915
17916 let response = agent.chat("billing please").await.unwrap();
17917
17918 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17919 assert_eq!(response.content, "Billing response");
17920 assert_eq!(
17921 agent.get_context().get("topic"),
17922 Some(&serde_json::json!("billing"))
17923 );
17924 }
17925
17926 #[tokio::test]
17927 async fn test_pre_response_extractor_miss_does_not_mutate_context() {
17928 let mock = mock_with_response("__NONE__");
17929 let yaml = r#"
17930name: OptimizedStateAgent
17931system_prompt: "You route before answering."
17932runtime:
17933 optimization:
17934 enabled: true
17935 pre_response_deterministic_transitions: true
17936states:
17937 initial: greeting
17938 states:
17939 greeting:
17940 prompt: "Old state prompt."
17941 extract:
17942 - key: topic
17943 description: "Support topic"
17944 transitions:
17945 - to: billing
17946 guard:
17947 context:
17948 topic:
17949 eq: billing
17950 timing: pre_response
17951 run_extractors: true
17952 billing:
17953 prompt: "Billing state."
17954"#;
17955 let agent = AgentBuilder::from_yaml(yaml)
17956 .unwrap()
17957 .llm(Arc::new(mock))
17958 .build()
17959 .unwrap();
17960
17961 let response = agent.try_pre_response_transition("hello").await.unwrap();
17962
17963 assert!(response.is_none());
17964 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17965 assert!(!agent.get_context().contains_key("topic"));
17966 }
17967
17968 #[tokio::test]
17969 async fn test_default_guard_transition_stays_post_response() {
17970 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
17971 let call_counter = mock.clone();
17972 let yaml = r#"
17973name: TimingAgent
17974system_prompt: "You route carefully."
17975runtime:
17976 optimization:
17977 enabled: true
17978 pre_response_deterministic_transitions: true
17979states:
17980 initial: greeting
17981 states:
17982 greeting:
17983 prompt: "Old state prompt."
17984 transitions:
17985 - to: billing
17986 guard:
17987 context:
17988 topic:
17989 eq: billing
17990 billing:
17991 prompt: "Billing state."
17992"#;
17993 let agent = AgentBuilder::from_yaml(yaml)
17994 .unwrap()
17995 .llm(Arc::new(mock))
17996 .build()
17997 .unwrap();
17998 agent
17999 .set_context("topic", serde_json::json!("billing"))
18000 .unwrap();
18001
18002 let response = agent.chat("billing please").await.unwrap();
18003
18004 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18005 assert_eq!(response.content, "Billing response");
18006 assert_eq!(call_counter.call_count(), 2);
18007 }
18008
18009 #[tokio::test]
18010 async fn test_explicit_post_response_guard_transition_stays_post_response() {
18011 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
18012 let call_counter = mock.clone();
18013 let yaml = r#"
18014name: TimingAgent
18015system_prompt: "You route carefully."
18016runtime:
18017 optimization:
18018 enabled: true
18019 pre_response_deterministic_transitions: true
18020states:
18021 initial: greeting
18022 states:
18023 greeting:
18024 prompt: "Old state prompt."
18025 transitions:
18026 - to: billing
18027 guard:
18028 context:
18029 topic:
18030 eq: billing
18031 timing: post_response
18032 billing:
18033 prompt: "Billing state."
18034"#;
18035 let agent = AgentBuilder::from_yaml(yaml)
18036 .unwrap()
18037 .llm(Arc::new(mock))
18038 .build()
18039 .unwrap();
18040 agent
18041 .set_context("topic", serde_json::json!("billing"))
18042 .unwrap();
18043
18044 let response = agent.chat("billing please").await.unwrap();
18045
18046 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18047 assert_eq!(response.content, "Billing response");
18048 assert_eq!(call_counter.call_count(), 2);
18049 }
18050
18051 #[tokio::test]
18052 async fn test_pre_response_extractors_are_transition_scoped() {
18053 let mock = mock_with_responses(vec!["billing", "Billing response"]);
18054 let yaml = r#"
18055name: ScopedExtractorAgent
18056system_prompt: "You route carefully."
18057runtime:
18058 optimization:
18059 enabled: true
18060 pre_response_deterministic_transitions: true
18061states:
18062 initial: greeting
18063 states:
18064 greeting:
18065 prompt: "Old state prompt."
18066 extract:
18067 - key: topic
18068 description: "Support topic"
18069 transitions:
18070 - to: wrong
18071 guard:
18072 context:
18073 topic:
18074 eq: billing
18075 timing: pre_response
18076 - to: billing
18077 guard:
18078 context:
18079 topic:
18080 eq: billing
18081 timing: pre_response
18082 run_extractors: true
18083 wrong:
18084 prompt: "Wrong state."
18085 billing:
18086 prompt: "Billing state."
18087"#;
18088 let agent = AgentBuilder::from_yaml(yaml)
18089 .unwrap()
18090 .llm(Arc::new(mock))
18091 .build()
18092 .unwrap();
18093
18094 let response = agent.chat("billing please").await.unwrap();
18095
18096 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18097 assert_eq!(response.content, "Billing response");
18098 }
18099
18100 #[tokio::test]
18101 async fn test_pre_response_resolved_intent_routes_early() {
18102 let mock = mock_with_response("Billing response");
18103 let yaml = r#"
18104name: IntentAgent
18105system_prompt: "You route carefully."
18106runtime:
18107 optimization:
18108 enabled: true
18109 pre_response_deterministic_transitions: true
18110states:
18111 initial: greeting
18112 states:
18113 greeting:
18114 prompt: "Old state prompt."
18115 transitions:
18116 - to: billing
18117 intent: billing
18118 timing: pre_response
18119 billing:
18120 prompt: "Billing state."
18121"#;
18122 let agent = AgentBuilder::from_yaml(yaml)
18123 .unwrap()
18124 .llm(Arc::new(mock))
18125 .build()
18126 .unwrap();
18127 agent
18128 .set_context("resolved_intent", serde_json::json!("billing"))
18129 .unwrap();
18130
18131 let response = agent
18132 .try_pre_response_transition("I need billing help")
18133 .await
18134 .unwrap()
18135 .unwrap();
18136
18137 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18138 assert_eq!(response.content, "Billing response");
18139 }
18140
18141 #[tokio::test]
18142 async fn test_background_overflow_error_surfaces() {
18143 let mut config = RuntimeConfig::default();
18144 config.optimization.enabled = true;
18145 config.optimization.post_turn.max_background_tasks = 1;
18146 config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
18147 let policy = crate::optimization::MaintenanceTaskPolicy {
18148 mode: MaintenanceMode::Background,
18149 await_before_next_turn: AwaitBeforeNextTurn::Always,
18150 };
18151 let agent = AgentBuilder::new()
18152 .system_prompt("You are helpful.")
18153 .llm(Arc::new(mock_with_response("ok")))
18154 .build()
18155 .unwrap()
18156 .with_runtime_config(config);
18157 agent
18158 .background_maintenance
18159 .spawn(None, async { std::future::pending::<Result<()>>().await })
18160 .unwrap();
18161
18162 let result = agent
18163 .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
18164 .await;
18165
18166 assert!(result.is_err());
18167 }
18168
18169 #[tokio::test]
18170 async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
18171 let default_mock = mock_with_response("Plain draft response");
18172 let router_mock = mock_with_response("cot");
18173 let router_counter = router_mock.clone();
18174 let yaml = r#"
18175name: ReasoningReservationAgent
18176system_prompt: "You answer plainly unless reasoning wins."
18177llm:
18178 default: default
18179 router: router
18180observability:
18181 enabled: true
18182 export:
18183 write_raw_events: true
18184reasoning:
18185 mode: auto
18186 judge_llm: router
18187runtime:
18188 optimization:
18189 enabled: true
18190 max_speculative_llm_calls_per_turn: 1
18191 speculative_reasoning_auto: true
18192 max_parallel_runtime_tasks: 2
18193"#;
18194 let agent = AgentBuilder::from_yaml(yaml)
18195 .unwrap()
18196 .llm_alias("default", Arc::new(default_mock))
18197 .llm_alias("router", Arc::new(router_mock))
18198 .build()
18199 .unwrap();
18200
18201 let response = agent.chat("hello").await.unwrap();
18202
18203 assert_eq!(response.content, "Plain draft response");
18204 assert_eq!(router_counter.call_count(), 1);
18205 let events = agent.observability().unwrap().raw_events();
18206 assert!(!events.iter().any(|event| {
18207 event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
18208 }));
18209 }
18210
18211 #[tokio::test]
18212 async fn test_forced_reasoning_skips_plain_speculative_draft() {
18213 let mock = mock_with_response("Reasoned response");
18214 let yaml = r#"
18215name: ForcedReasoningAgent
18216system_prompt: "You reason before answering."
18217observability:
18218 enabled: true
18219 export:
18220 write_raw_events: true
18221reasoning:
18222 mode: cot
18223runtime:
18224 optimization:
18225 enabled: true
18226 max_speculative_llm_calls_per_turn: 2
18227 speculative_state_transitions: true
18228 max_parallel_runtime_tasks: 2
18229states:
18230 initial: triage
18231 states:
18232 triage:
18233 prompt: "Answer from triage."
18234 transitions:
18235 - to: billing
18236 guard:
18237 context:
18238 route:
18239 eq: billing
18240 timing: parallel
18241 billing:
18242 prompt: "Billing state."
18243"#;
18244 let agent = AgentBuilder::from_yaml(yaml)
18245 .unwrap()
18246 .llm(Arc::new(mock))
18247 .build()
18248 .unwrap();
18249
18250 let response = agent.chat("hello").await.unwrap();
18251
18252 assert_eq!(response.content, "Reasoned response");
18253 let events = agent.observability().unwrap().raw_events();
18254 assert!(
18255 !events
18256 .iter()
18257 .any(|event| event.dimensions.contains_key("branch_status"))
18258 );
18259 }
18260
18261 #[tokio::test]
18262 async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
18263 let default_mock = mock_with_response("Skill committed response");
18264 let router_mock = mock_with_response("helper");
18265 let router_counter = router_mock.clone();
18266 let yaml = r#"
18267name: SkillReservationAgent
18268system_prompt: "Use skills when they match."
18269llm:
18270 default: default
18271 router: router
18272observability:
18273 enabled: true
18274 export:
18275 write_raw_events: true
18276runtime:
18277 optimization:
18278 enabled: true
18279 max_speculative_llm_calls_per_turn: 1
18280 speculative_skill_routing: true
18281 max_parallel_runtime_tasks: 2
18282skills:
18283 - id: helper
18284 description: "Answer helper requests"
18285 trigger: "User asks for helper"
18286 steps:
18287 - prompt: "Answer the helper request: {{ user_input }}"
18288"#;
18289 let agent = AgentBuilder::from_yaml(yaml)
18290 .unwrap()
18291 .llm_alias("default", Arc::new(default_mock))
18292 .llm_alias("router", Arc::new(router_mock))
18293 .build()
18294 .unwrap();
18295
18296 let response = agent.chat("please use helper").await.unwrap();
18297
18298 assert_eq!(response.content, "Skill committed response");
18299 assert_eq!(router_counter.call_count(), 1);
18300 let events = agent.observability().unwrap().raw_events();
18301 assert!(
18302 !events
18303 .iter()
18304 .any(|event| event.dimensions.contains_key("branch_status"))
18305 );
18306 }
18307
18308 #[tokio::test]
18309 async fn test_parallel_transition_low_cap_allows_deterministic_route() {
18310 let mock = mock_with_response("unused");
18311 let call_counter = mock.clone();
18312 let yaml = r#"
18313name: ParallelTransitionLowCapAgent
18314system_prompt: "Route before stale responses when safe."
18315runtime:
18316 optimization:
18317 enabled: true
18318 max_speculative_llm_calls_per_turn: 1
18319 speculative_state_transitions: true
18320 max_parallel_runtime_tasks: 2
18321states:
18322 initial: triage
18323 states:
18324 triage:
18325 prompt: "Triage state."
18326 transitions:
18327 - to: billing
18328 guard:
18329 context:
18330 route:
18331 eq: billing
18332 timing: parallel
18333 billing:
18334 prompt: "Billing state."
18335"#;
18336 let agent = AgentBuilder::from_yaml(yaml)
18337 .unwrap()
18338 .llm(Arc::new(mock))
18339 .build()
18340 .unwrap();
18341 agent
18342 .set_context("route", serde_json::json!("billing"))
18343 .unwrap();
18344 agent.update_active_turn_context("billing help", HashMap::new());
18345 assert!(
18346 agent.reserve_active_speculative_llm_call(
18347 RuntimeOptimizationKind::ParallelStateTransition
18348 )
18349 );
18350
18351 let selection = agent
18352 .select_parallel_transition_candidate("billing help")
18353 .await
18354 .unwrap();
18355 agent.end_root_turn();
18356
18357 match selection {
18358 ParallelTransitionSelection::Candidate(candidate) => {
18359 assert_eq!(candidate.target(), "billing");
18360 }
18361 ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
18362 ParallelTransitionSelection::ReservationExhausted => {
18363 panic!("deterministic route consumed LLM budget")
18364 }
18365 }
18366 assert_eq!(call_counter.call_count(), 0);
18367 }
18368
18369 #[tokio::test]
18370 async fn speculative_transition_drops_loser_before_state_actions() {
18371 let lock = Arc::new(tokio::sync::Mutex::new(()));
18372 let first_started = Arc::new(tokio::sync::Notify::new());
18373 let first_dropped = Arc::new(AtomicBool::new(false));
18374 let committed_after_drop = Arc::new(AtomicBool::new(false));
18375 let default = Arc::new(FirstCallLockingProvider {
18376 lock,
18377 first_started: Arc::clone(&first_started),
18378 first_dropped: Arc::clone(&first_dropped),
18379 committed_after_drop: Arc::clone(&committed_after_drop),
18380 calls: AtomicU64::new(0),
18381 });
18382 let router = Arc::new(RoutingAfterProviderStart {
18383 provider_started: first_started,
18384 });
18385 let yaml = r#"
18386name: SpeculativeCancellationAgent
18387system_prompt: "Route before committed work."
18388llm:
18389 default: default
18390 router: router
18391runtime:
18392 optimization:
18393 enabled: true
18394 max_speculative_llm_calls_per_turn: 2
18395 speculative_state_transitions: true
18396 max_parallel_runtime_tasks: 2
18397states:
18398 initial: triage
18399 states:
18400 triage:
18401 prompt: "Triage state."
18402 transitions:
18403 - to: technical
18404 when: "The request needs technical support"
18405 timing: parallel
18406 technical:
18407 prompt: "Technical state."
18408 on_enter:
18409 - prompt: "Prepare technical context."
18410 llm: default
18411 store_as: preparation
18412"#;
18413 let agent = AgentBuilder::from_yaml(yaml)
18414 .unwrap()
18415 .llm_alias("default", default)
18416 .llm_alias("router", router)
18417 .build()
18418 .unwrap();
18419
18420 let response = tokio::time::timeout(
18421 std::time::Duration::from_secs(2),
18422 agent.chat("I cannot log in because of AUTH-17."),
18423 )
18424 .await
18425 .expect("committed work must not wait on the losing provider future")
18426 .unwrap();
18427
18428 assert_eq!(response.content, "Committed technical response.");
18429 assert_eq!(agent.current_state().as_deref(), Some("technical"));
18430 assert!(first_dropped.load(Ordering::SeqCst));
18431 assert!(committed_after_drop.load(Ordering::SeqCst));
18432 }
18433
18434 #[tokio::test]
18435 async fn buffered_transition_drops_stale_stream_before_redispatch() {
18436 use futures::StreamExt;
18437
18438 let lock = Arc::new(tokio::sync::Mutex::new(()));
18439 let stream_started = Arc::new(tokio::sync::Notify::new());
18440 let stream_dropped = Arc::new(AtomicBool::new(false));
18441 let committed_after_drop = Arc::new(AtomicBool::new(false));
18442 let default = Arc::new(BufferedLockingProvider {
18443 lock,
18444 stream_started: Arc::clone(&stream_started),
18445 stream_dropped: Arc::clone(&stream_dropped),
18446 committed_after_drop: Arc::clone(&committed_after_drop),
18447 });
18448 let router = Arc::new(RoutingAfterProviderStart {
18449 provider_started: stream_started,
18450 });
18451 let yaml = r#"
18452name: BufferedCancellationAgent
18453system_prompt: "Hide stale streamed output."
18454llm:
18455 default: default
18456 router: router
18457streaming:
18458 enabled: true
18459 buffer_size: 8
18460runtime:
18461 optimization:
18462 enabled: true
18463 max_speculative_llm_calls_per_turn: 2
18464 speculative_state_transitions: true
18465 streaming_policy: buffer_until_routing_done
18466 max_parallel_runtime_tasks: 2
18467states:
18468 initial: triage
18469 states:
18470 triage:
18471 prompt: "Triage state."
18472 transitions:
18473 - to: technical
18474 when: "The request needs technical support"
18475 timing: parallel
18476 technical:
18477 prompt: "Technical state."
18478"#;
18479 let agent = AgentBuilder::from_yaml(yaml)
18480 .unwrap()
18481 .llm_alias("default", default)
18482 .llm_alias("router", router)
18483 .build()
18484 .unwrap();
18485
18486 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18487 let mut stream = agent
18488 .chat_stream("AUTH-17 needs technical help.")
18489 .await
18490 .unwrap();
18491 let mut content = String::new();
18492 while let Some(chunk) = stream.next().await {
18493 match chunk {
18494 StreamChunk::Content { text } => content.push_str(&text),
18495 StreamChunk::Done {} => break,
18496 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18497 _ => {}
18498 }
18499 }
18500 content
18501 })
18502 .await
18503 .expect("redispatch must not wait on the stale streaming future");
18504
18505 assert_eq!(content, "Committed technical response.");
18506 assert_eq!(agent.current_state().as_deref(), Some("technical"));
18507 assert!(stream_dropped.load(Ordering::SeqCst));
18508 assert!(committed_after_drop.load(Ordering::SeqCst));
18509 }
18510
18511 #[tokio::test]
18512 async fn buffered_transition_drops_established_stream_before_redispatch() {
18513 use futures::StreamExt;
18514
18515 let stream_started = Arc::new(tokio::sync::Notify::new());
18516 let stream_dropped = Arc::new(AtomicBool::new(false));
18517 let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
18518 let committed_after_drop = Arc::new(AtomicBool::new(false));
18519 let default = Arc::new(EstablishedStreamProvider {
18520 stream_started: Arc::clone(&stream_started),
18521 stream_dropped: Arc::clone(&stream_dropped),
18522 stream_dropped_notify,
18523 committed_after_drop: Arc::clone(&committed_after_drop),
18524 });
18525 let router = Arc::new(RoutingAfterProviderStart {
18526 provider_started: stream_started,
18527 });
18528 let yaml = r#"
18529name: EstablishedStreamCancellationAgent
18530system_prompt: "Hide stale streamed output."
18531llm:
18532 default: default
18533 router: router
18534streaming:
18535 enabled: true
18536 buffer_size: 8
18537runtime:
18538 optimization:
18539 enabled: true
18540 max_speculative_llm_calls_per_turn: 2
18541 speculative_state_transitions: true
18542 streaming_policy: buffer_until_routing_done
18543 max_parallel_runtime_tasks: 2
18544states:
18545 initial: triage
18546 states:
18547 triage:
18548 prompt: "Triage state."
18549 transitions:
18550 - to: technical
18551 when: "The request needs technical support"
18552 timing: parallel
18553 technical:
18554 prompt: "Technical state."
18555"#;
18556 let agent = AgentBuilder::from_yaml(yaml)
18557 .unwrap()
18558 .llm_alias("default", default)
18559 .llm_alias("router", router)
18560 .build()
18561 .unwrap();
18562
18563 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18564 let mut stream = agent
18565 .chat_stream("AUTH-17 needs technical help.")
18566 .await
18567 .unwrap();
18568 let mut content = String::new();
18569 while let Some(chunk) = stream.next().await {
18570 match chunk {
18571 StreamChunk::Content { text } => content.push_str(&text),
18572 StreamChunk::Done {} => break,
18573 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18574 _ => {}
18575 }
18576 }
18577 content
18578 })
18579 .await
18580 .expect("redispatch must wait for the established stale stream to be dropped");
18581
18582 assert_eq!(content, "Committed technical response.");
18583 assert_eq!(agent.current_state().as_deref(), Some("technical"));
18584 assert!(stream_dropped.load(Ordering::SeqCst));
18585 assert!(committed_after_drop.load(Ordering::SeqCst));
18586 }
18587
18588 #[tokio::test]
18589 async fn test_buffered_streaming_transition_reservation_falls_back() {
18590 use futures::StreamExt;
18591
18592 let mock = mock_with_responses(vec![
18593 "Serial streaming response",
18594 "Serial streaming response",
18595 ]);
18596 let router_mock = mock_with_response("1");
18597 let router_counter = router_mock.clone();
18598 let yaml = r#"
18599name: BufferedReservationFallbackAgent
18600system_prompt: "Stream normally if speculative routing cannot be evaluated."
18601llm:
18602 default: default
18603 router: router
18604observability:
18605 enabled: true
18606 export:
18607 write_raw_events: true
18608streaming:
18609 enabled: true
18610 buffer_size: 8
18611runtime:
18612 optimization:
18613 enabled: true
18614 max_speculative_llm_calls_per_turn: 1
18615 speculative_state_transitions: true
18616 streaming_policy: buffer_until_routing_done
18617 max_parallel_runtime_tasks: 2
18618states:
18619 initial: triage
18620 states:
18621 triage:
18622 prompt: "Triage state."
18623 transitions:
18624 - to: billing
18625 guard:
18626 context:
18627 route:
18628 eq: billing
18629 when: "User asks about billing"
18630 timing: parallel
18631 billing:
18632 prompt: "Billing state."
18633"#;
18634 let agent = AgentBuilder::from_yaml(yaml)
18635 .unwrap()
18636 .llm_alias("default", Arc::new(mock))
18637 .llm_alias("router", Arc::new(router_mock))
18638 .build()
18639 .unwrap();
18640
18641 let mut stream = agent.chat_stream("hello").await.unwrap();
18642 let mut content = String::new();
18643 let mut error = None;
18644 while let Some(chunk) = stream.next().await {
18645 match chunk {
18646 StreamChunk::Content { text } => content.push_str(&text),
18647 StreamChunk::Error { message } => error = Some(message),
18648 StreamChunk::Done {} => break,
18649 _ => {}
18650 }
18651 }
18652
18653 assert_eq!(error, None);
18654 assert_eq!(content, "Serial streaming response");
18655 assert_eq!(router_counter.call_count(), 0);
18656 let events = agent.observability().unwrap().raw_events();
18657 assert!(events.iter().any(|event| {
18658 event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
18659 && event.dimensions.get("commit_behavior")
18660 == Some(&"transition_decision".to_string())
18661 }));
18662 }
18663
18664 #[tokio::test]
18665 async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
18666 let mut mock = mock_with_response("Recovered response");
18667 mock.set_error("boom");
18668 let mut handle = mock.clone();
18669 let agent = AgentBuilder::new()
18670 .system_prompt("You are helpful.")
18671 .llm(Arc::new(mock))
18672 .build()
18673 .unwrap();
18674
18675 assert!(agent.chat("first").await.is_err());
18676 handle.clear_error();
18677 let response = agent.chat("second").await.unwrap();
18678
18679 assert_eq!(response.content, "Recovered response");
18680 let messages = agent.memory.get_messages(None).await.unwrap();
18681 let user_count = messages
18682 .iter()
18683 .filter(|message| message.role == ai_agents_core::Role::User)
18684 .count();
18685 assert_eq!(user_count, 2);
18686 }
18687
18688 #[tokio::test]
18689 async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
18690 use futures::StreamExt;
18691
18692 let mut mock = mock_with_response("Recovered response");
18693 mock.set_error("stream boom");
18694 let mut handle = mock.clone();
18695 let agent = AgentBuilder::new()
18696 .system_prompt("You are helpful.")
18697 .llm(Arc::new(mock))
18698 .build()
18699 .unwrap();
18700
18701 let mut stream = agent.chat_stream("first").await.unwrap();
18702 let mut saw_error = false;
18703 while let Some(chunk) = stream.next().await {
18704 if matches!(chunk, StreamChunk::Error { .. }) {
18705 saw_error = true;
18706 }
18707 }
18708 assert!(saw_error);
18709
18710 handle.clear_error();
18711 let response = agent.chat("second").await.unwrap();
18712
18713 assert_eq!(response.content, "Recovered response");
18714 let messages = agent.memory.get_messages(None).await.unwrap();
18715 let user_count = messages
18716 .iter()
18717 .filter(|message| message.role == ai_agents_core::Role::User)
18718 .count();
18719 assert_eq!(user_count, 2);
18720 }
18721
18722 #[tokio::test]
18723 async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
18724 use futures::StreamExt;
18725
18726 let mut mock = mock_with_response("one two three");
18727 mock.set_latency(10);
18728 let yaml = r#"
18729name: BufferedMissAgent
18730system_prompt: "You stream safely."
18731llm:
18732 default: default
18733streaming:
18734 enabled: true
18735 buffer_size: 1
18736runtime:
18737 optimization:
18738 enabled: true
18739 max_speculative_llm_calls_per_turn: 2
18740 speculative_state_transitions: true
18741 streaming_policy: buffer_until_routing_done
18742 max_parallel_runtime_tasks: 2
18743states:
18744 initial: triage
18745 states:
18746 triage:
18747 prompt: "Answer from triage."
18748 transitions:
18749 - to: billing
18750 guard:
18751 context:
18752 route:
18753 eq: billing
18754 timing: parallel
18755 billing:
18756 prompt: "Billing state."
18757"#;
18758 let agent = AgentBuilder::from_yaml(yaml)
18759 .unwrap()
18760 .llm_alias("default", Arc::new(mock))
18761 .build()
18762 .unwrap();
18763
18764 let mut stream = agent.chat_stream("hello").await.unwrap();
18765 let mut content = String::new();
18766 let mut error = None;
18767 while let Some(chunk) = stream.next().await {
18768 match chunk {
18769 StreamChunk::Content { text } => content.push_str(&text),
18770 StreamChunk::Error { message } => error = Some(message),
18771 StreamChunk::Done {} => break,
18772 _ => {}
18773 }
18774 }
18775
18776 assert_eq!(error, None);
18777 assert_eq!(content, "one two three");
18778 }
18779
18780 #[tokio::test]
18781 async fn test_buffered_streaming_main_failure_finalizes_branch() {
18782 use futures::StreamExt;
18783
18784 let mock = mock_with_response("one two");
18785 let mut router_mock = mock_with_response("0");
18786 router_mock.set_latency(50);
18787 let yaml = r#"
18788name: BufferedFailureAgent
18789system_prompt: "You stream safely."
18790llm:
18791 default: default
18792 router: router
18793observability:
18794 enabled: true
18795 export:
18796 write_raw_events: true
18797streaming:
18798 enabled: true
18799 buffer_size: 1
18800runtime:
18801 optimization:
18802 enabled: true
18803 max_speculative_llm_calls_per_turn: 2
18804 speculative_state_transitions: true
18805 streaming_policy: buffer_until_routing_done
18806 max_parallel_runtime_tasks: 2
18807states:
18808 initial: triage
18809 states:
18810 triage:
18811 prompt: "Ask for the category."
18812 transitions:
18813 - to: billing
18814 when: "User asks about billing"
18815 timing: parallel
18816 billing:
18817 prompt: "Billing state."
18818"#;
18819 let agent = AgentBuilder::from_yaml(yaml)
18820 .unwrap()
18821 .llm_alias("default", Arc::new(mock))
18822 .llm_alias("router", Arc::new(router_mock))
18823 .build()
18824 .unwrap();
18825
18826 let mut stream = agent.chat_stream("hello").await.unwrap();
18827 let mut error = String::new();
18828 while let Some(chunk) = stream.next().await {
18829 if let StreamChunk::Error { message } = chunk {
18830 error = message;
18831 }
18832 }
18833
18834 assert!(
18835 error.contains("stream buffer filled"),
18836 "unexpected stream error: {}",
18837 error
18838 );
18839 let events = agent.observability().unwrap().raw_events();
18840 assert!(events.iter().any(|event| {
18841 event.dimensions.get("branch_status") == Some(&"failed".to_string())
18842 && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
18843 && event.dimensions.get("optimization")
18844 == Some(&"buffered_streaming_routing".to_string())
18845 }));
18846 }
18847
18848 #[tokio::test]
18849 async fn test_streaming_preflight_does_not_emit_old_state_content() {
18850 use futures::StreamExt;
18851
18852 let mock = mock_with_response("Billing streamed response");
18853 let yaml = r#"
18854name: StreamingOptimizedAgent
18855system_prompt: "You route before streaming."
18856runtime:
18857 optimization:
18858 enabled: true
18859 pre_response_deterministic_transitions: true
18860streaming:
18861 enabled: true
18862states:
18863 initial: greeting
18864 states:
18865 greeting:
18866 prompt: "OLD_STATE_SENTINEL"
18867 transitions:
18868 - to: billing
18869 guard:
18870 context:
18871 topic:
18872 eq: billing
18873 timing: pre_response
18874 billing:
18875 prompt: "Billing state."
18876"#;
18877 let agent = AgentBuilder::from_yaml(yaml)
18878 .unwrap()
18879 .llm(Arc::new(mock))
18880 .build()
18881 .unwrap();
18882 agent
18883 .set_context("topic", serde_json::json!("billing"))
18884 .unwrap();
18885
18886 let mut stream = agent.chat_stream("billing please").await.unwrap();
18887 let mut content = String::new();
18888 while let Some(chunk) = stream.next().await {
18889 match chunk {
18890 StreamChunk::Content { text } => content.push_str(&text),
18891 StreamChunk::Error { message } => panic!("stream error: {}", message),
18892 StreamChunk::Done {} => break,
18893 _ => {}
18894 }
18895 }
18896
18897 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18898 assert!(content.contains("Billing streamed response"));
18899 assert!(!content.contains("OLD_STATE_SENTINEL"));
18900 }
18901
18902 #[tokio::test]
18904 async fn test_integration_state_machine_basic() {
18905 let yaml = r#"
18906name: StateAgent
18907system_prompt: "You are a support agent."
18908states:
18909 initial: greeting
18910 states:
18911 greeting:
18912 prompt: "Welcome the user warmly."
18913 transitions:
18914 - to: support
18915 when: "User needs help"
18916 auto: true
18917 support:
18918 prompt: "Help solve the user's problem."
18919"#;
18920 let mock = mock_with_responses(vec![
18921 "Welcome! How can I help?", "1", "I'll help you with that.", ]);
18925 let builder = AgentBuilder::from_yaml(yaml).unwrap();
18926 let agent = builder.llm(Arc::new(mock)).build().unwrap();
18927
18928 assert_eq!(agent.current_state(), Some("greeting".to_string()));
18929 let _ = agent.chat("I need help").await.unwrap();
18930 }
18933
18934 #[tokio::test]
18936 async fn test_integration_state_on_enter_set_context() {
18937 let yaml = r#"
18938name: ActionAgent
18939system_prompt: "You are helpful."
18940states:
18941 initial: step1
18942 states:
18943 step1:
18944 prompt: "Step 1"
18945 on_exit:
18946 - set_context:
18947 step1_exited: true
18948 transitions:
18949 - to: step2
18950 when: "always"
18951 auto: true
18952 step2:
18953 prompt: "Step 2"
18954 on_enter:
18955 - set_context:
18956 step2_entered: true
18957"#;
18958 let mock = mock_with_responses(vec![
18960 "Processing step 1.",
18961 "0", ]);
18963 let builder = AgentBuilder::from_yaml(yaml).unwrap();
18964 let agent = builder.llm(Arc::new(mock)).build().unwrap();
18965
18966 assert_eq!(agent.current_state(), Some("step1".to_string()));
18967
18968 agent.transition_to("step2").await.unwrap();
18970
18971 assert_eq!(agent.current_state(), Some("step2".to_string()));
18972
18973 let ctx = agent.get_context();
18975 assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
18976 assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
18977 }
18978
18979 #[tokio::test]
18980 async fn state_action_tool_preserves_source_in_stored_record() {
18981 let yaml = r#"
18982name: StateActionToolAgent
18983system_prompt: "You are helpful."
18984tools:
18985 - context_echo
18986states:
18987 initial: idle
18988 states:
18989 idle:
18990 prompt: "Idle"
18991 active:
18992 prompt: "Active"
18993 on_enter:
18994 - set_context:
18995 action_started: true
18996 - tool: context_echo
18997 args: {}
18998"#;
18999 let agent = AgentBuilder::from_yaml(yaml)
19000 .unwrap()
19001 .llm(Arc::new(mock_with_response("unused")))
19002 .tool(Arc::new(ContextEchoTool))
19003 .build()
19004 .unwrap();
19005
19006 agent.transition_to("active").await.unwrap();
19007
19008 let record: ToolExecutionRecord = serde_json::from_value(
19009 agent
19010 .get_context()
19011 .get("last_tool_record")
19012 .cloned()
19013 .expect("successful state action must store its execution record"),
19014 )
19015 .unwrap();
19016 assert!(record.executed);
19017 assert!(record.success);
19018 assert_eq!(record.canonical_id, "context_echo");
19019 assert!(matches!(
19020 &record.source,
19021 ToolCallSource::StateAction {
19022 state: Some(state),
19023 action_index: 1,
19024 } if state == "active"
19025 ));
19026 }
19027
19028 #[tokio::test]
19029 async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
19030 let yaml = r#"
19031name: OrdinaryLifecycleAgent
19032system_prompt: "You are helpful."
19033states:
19034 initial: intake
19035 regenerate_on_transition: false
19036 states:
19037 intake:
19038 prompt: "Intake"
19039 transitions:
19040 - to: drafting
19041 guard:
19042 context:
19043 route:
19044 eq: drafting
19045 drafting:
19046 prompt: "Drafting"
19047 on_enter:
19048 - set_context:
19049 draft_version: 1
19050 on_reenter:
19051 - set_context:
19052 draft_version: 2
19053 transitions:
19054 - to: review
19055 guard:
19056 context:
19057 route:
19058 eq: review
19059 review:
19060 prompt: "Review"
19061 on_enter:
19062 - set_context:
19063 review_entry: first
19064 transitions:
19065 - to: drafting
19066 guard:
19067 context:
19068 route:
19069 eq: drafting
19070"#;
19071 let agent = AgentBuilder::from_yaml(yaml)
19072 .unwrap()
19073 .llm(Arc::new(mock_with_responses(vec![
19074 "Intake response",
19075 "Draft response",
19076 "Review response",
19077 ])))
19078 .build()
19079 .unwrap();
19080
19081 agent
19082 .set_context("route", serde_json::json!("drafting"))
19083 .unwrap();
19084 agent.chat("Start a draft").await.unwrap();
19085 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19086 assert_eq!(
19087 agent.get_context().get("draft_version"),
19088 Some(&serde_json::json!(1))
19089 );
19090
19091 agent
19092 .set_context("route", serde_json::json!("review"))
19093 .unwrap();
19094 agent.chat("Review this").await.unwrap();
19095 assert_eq!(agent.current_state().as_deref(), Some("review"));
19096 assert_eq!(
19097 agent.get_context().get("review_entry"),
19098 Some(&serde_json::json!("first"))
19099 );
19100
19101 agent
19102 .set_context("route", serde_json::json!("drafting"))
19103 .unwrap();
19104 agent.chat("Revise this").await.unwrap();
19105 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19106 assert_eq!(
19107 agent.get_context().get("draft_version"),
19108 Some(&serde_json::json!(2))
19109 );
19110 }
19111
19112 #[tokio::test]
19113 async fn test_manual_transition_uses_on_enter_then_on_reenter() {
19114 let yaml = r#"
19115name: ManualLifecycleAgent
19116system_prompt: "You are helpful."
19117states:
19118 initial: intake
19119 states:
19120 intake:
19121 prompt: "Intake"
19122 drafting:
19123 prompt: "Drafting"
19124 on_enter:
19125 - set_context:
19126 draft_version: 1
19127 on_reenter:
19128 - set_context:
19129 draft_version: 2
19130 review:
19131 prompt: "Review"
19132"#;
19133 let agent = AgentBuilder::from_yaml(yaml)
19134 .unwrap()
19135 .llm(Arc::new(mock_with_response("unused")))
19136 .build()
19137 .unwrap();
19138
19139 assert!(!agent.get_context().contains_key("draft_version"));
19140 agent.transition_to("drafting").await.unwrap();
19141 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19142 assert_eq!(
19143 agent.get_context().get("draft_version"),
19144 Some(&serde_json::json!(1))
19145 );
19146
19147 agent.transition_to("review").await.unwrap();
19148 agent.transition_to("drafting").await.unwrap();
19149 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19150 assert_eq!(
19151 agent.get_context().get("draft_version"),
19152 Some(&serde_json::json!(2))
19153 );
19154 }
19155
19156 #[tokio::test]
19157 async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
19158 let yaml = r#"
19159name: TimeoutLifecycleAgent
19160system_prompt: "You are helpful."
19161states:
19162 initial: intake
19163 regenerate_on_transition: false
19164 states:
19165 intake:
19166 prompt: "Intake"
19167 max_turns: 1
19168 timeout_to: drafting
19169 drafting:
19170 prompt: "Drafting"
19171 max_turns: 1
19172 timeout_to: review
19173 on_enter:
19174 - set_context:
19175 draft_version: 1
19176 on_reenter:
19177 - set_context:
19178 draft_version: 2
19179 review:
19180 prompt: "Review"
19181 max_turns: 1
19182 timeout_to: drafting
19183 on_enter:
19184 - set_context:
19185 review_entry: first
19186"#;
19187 let agent = AgentBuilder::from_yaml(yaml)
19188 .unwrap()
19189 .llm(Arc::new(mock_with_responses(vec![
19190 "Intake",
19191 "First draft",
19192 "Review",
19193 "Revised draft",
19194 ])))
19195 .build()
19196 .unwrap();
19197
19198 agent.chat("First turn").await.unwrap();
19199 assert_eq!(agent.current_state().as_deref(), Some("intake"));
19200 assert!(!agent.get_context().contains_key("draft_version"));
19201
19202 agent.chat("Second turn").await.unwrap();
19203 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19204 assert_eq!(
19205 agent.get_context().get("draft_version"),
19206 Some(&serde_json::json!(1))
19207 );
19208
19209 agent.chat("Third turn").await.unwrap();
19210 assert_eq!(agent.current_state().as_deref(), Some("review"));
19211 assert_eq!(
19212 agent.get_context().get("review_entry"),
19213 Some(&serde_json::json!("first"))
19214 );
19215
19216 agent.chat("Fourth turn").await.unwrap();
19217 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
19218 assert_eq!(
19219 agent.get_context().get("draft_version"),
19220 Some(&serde_json::json!(2))
19221 );
19222 }
19223
19224 #[tokio::test]
19226 async fn test_integration_process_normalize() {
19227 let yaml = r#"
19228name: ProcessAgent
19229system_prompt: "You are helpful."
19230process:
19231 input:
19232 - type: normalize
19233 config:
19234 trim: true
19235 collapse_whitespace: true
19236"#;
19237 let mock = mock_with_response("Got your message.");
19238 let builder = AgentBuilder::from_yaml(yaml).unwrap();
19239 let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
19240
19241 let _ = agent.chat(" hello world ").await.unwrap();
19242
19243 let history = mock.call_history();
19245 assert!(!history.is_empty());
19246 let last_call = history.last().unwrap();
19248 let user_msg = last_call
19249 .messages
19250 .iter()
19251 .find(|m| m.role == ai_agents_core::Role::User)
19252 .unwrap();
19253 assert_eq!(user_msg.content, "hello world");
19254 }
19255
19256 #[tokio::test]
19260 async fn test_integration_memory_compression() {
19261 let yaml = r#"
19262name: MemoryAgent
19263system_prompt: "You are helpful."
19264memory:
19265 type: compacting
19266 max_messages: 100
19267 compress_threshold: 5
19268 max_recent_messages: 3
19269 summarize_batch_size: 2
19270"#;
19271 let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
19273 let mock = mock_with_responses(responses);
19274 let builder = AgentBuilder::from_yaml(yaml).unwrap();
19275 let agent = builder.llm(Arc::new(mock)).build().unwrap();
19276
19277 for i in 0..6 {
19279 let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
19280 }
19281
19282 let messages = agent.memory.get_messages(None).await.unwrap();
19285 assert!(messages.len() <= 12); }
19289
19290 #[tokio::test]
19292 async fn test_integration_multi_llm_registry() {
19293 let mut mock_default = MockLLMProvider::new("default");
19294 mock_default.set_response("Default LLM response.");
19295 let mut mock_router = MockLLMProvider::new("router");
19296 mock_router.set_response("Router response.");
19297
19298 let agent = AgentBuilder::new()
19299 .system_prompt("You are helpful.")
19300 .llm_alias("default", Arc::new(mock_default))
19301 .llm_alias("router", Arc::new(mock_router))
19302 .build()
19303 .unwrap();
19304
19305 let response = agent.chat("Hello").await.unwrap();
19306 assert_eq!(response.content, "Default LLM response.");
19307 }
19308
19309 #[tokio::test]
19311 async fn test_integration_agent_reset() {
19312 let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
19313 let agent = AgentBuilder::new()
19314 .system_prompt("You are helpful.")
19315 .llm(Arc::new(mock))
19316 .build()
19317 .unwrap();
19318
19319 let _ = agent.chat("Hi").await.unwrap();
19320 let messages = agent.memory.get_messages(None).await.unwrap();
19321 assert_eq!(messages.len(), 2); agent.reset().await.unwrap();
19324 let messages = agent.memory.get_messages(None).await.unwrap();
19325 assert_eq!(messages.len(), 0);
19326 }
19327
19328 #[tokio::test]
19330 async fn test_integration_process_validate_reject() {
19331 use ai_agents_process::{ProcessConfig, ProcessProcessor};
19332
19333 let validate_config = ai_agents_process::ValidateStage {
19334 id: Some("length_check".to_string()),
19335 condition: None,
19336 config: ai_agents_process::ValidateConfig {
19337 rules: vec![ai_agents_process::ValidationRule::MinLength {
19338 min_length: 10,
19339 on_fail: ai_agents_process::ValidationAction {
19340 action: ai_agents_process::ValidationActionType::Reject,
19341 message: None,
19342 },
19343 }],
19344 ..Default::default()
19345 },
19346 };
19347 let process_config = ProcessConfig {
19348 input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
19349 ..Default::default()
19350 };
19351 let processor = ProcessProcessor::new(process_config);
19352
19353 let mock = mock_with_response("Should not reach here.");
19354 let agent = AgentBuilder::new()
19355 .system_prompt("You are helpful.")
19356 .llm(Arc::new(mock))
19357 .process_processor(processor)
19358 .build()
19359 .unwrap();
19360
19361 let response = agent.chat("Hi").await.unwrap();
19362 assert!(
19364 response.content.contains("rejected")
19365 || response.content.contains("Input rejected")
19366 || response.content.contains("too short")
19367 || response.content.contains("Too short")
19368 || response.content.len() < 50, "Expected rejection response, got: {}",
19370 response.content
19371 );
19372 }
19373
19374 #[tokio::test]
19376 async fn test_llm_fallback_on_failure() {
19377 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19378
19379 let mut primary = MockLLMProvider::new("primary");
19380 primary.set_error("Primary LLM is unavailable");
19381
19382 let mut fallback = MockLLMProvider::new("fallback");
19383 fallback.set_response("Fallback response works!");
19384
19385 let agent = AgentBuilder::new()
19386 .system_prompt("You are helpful.")
19387 .llm_alias("default", Arc::new(primary))
19388 .llm_alias("backup", Arc::new(fallback))
19389 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19390 llm: LLMRecoveryConfig {
19391 on_failure: LLMFailureAction::FallbackLlm {
19392 fallback_llm: "backup".to_string(),
19393 },
19394 ..Default::default()
19395 },
19396 ..Default::default()
19397 }))
19398 .build()
19399 .unwrap();
19400
19401 let response = agent.chat("Hello").await.unwrap();
19402 assert!(
19403 response.content.contains("Fallback response"),
19404 "Expected fallback response, got: {}",
19405 response.content
19406 );
19407 }
19408
19409 #[tokio::test]
19411 async fn test_llm_fallback_response_static_message() {
19412 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19413
19414 let mut primary = MockLLMProvider::new("primary");
19415 primary.set_error("Primary LLM is unavailable");
19416
19417 let agent = AgentBuilder::new()
19418 .system_prompt("You are helpful.")
19419 .llm(Arc::new(primary))
19420 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19421 llm: LLMRecoveryConfig {
19422 on_failure: LLMFailureAction::FallbackResponse {
19423 message: "I am temporarily unavailable. Please try again later."
19424 .to_string(),
19425 },
19426 ..Default::default()
19427 },
19428 ..Default::default()
19429 }))
19430 .build()
19431 .unwrap();
19432
19433 let response = agent.chat("Hello").await.unwrap();
19434 assert!(
19435 response.content.contains("temporarily unavailable"),
19436 "Expected static fallback message, got: {}",
19437 response.content
19438 );
19439 }
19440
19441 #[tokio::test]
19443 async fn test_tool_failure_skip() {
19444 use ai_agents_recovery::{
19445 ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
19446 };
19447
19448 let mock = mock_with_responses(vec![
19450 r#"I'll use the nonexistent tool.
19451[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
19452 "The tool was unavailable, but I can still help you.",
19453 ]);
19454
19455 let agent = AgentBuilder::new()
19456 .system_prompt("You are helpful.")
19457 .llm(Arc::new(mock))
19458 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19459 tools: ToolRecoveryConfig {
19460 default: ToolRetryConfig {
19461 max_retries: 0,
19462 timeout_ms: None,
19463 on_failure: ToolFailureAction::Skip,
19464 },
19465 ..Default::default()
19466 },
19467 ..Default::default()
19468 }))
19469 .build()
19470 .unwrap();
19471
19472 let response = agent.chat("Use the nonexistent tool").await;
19474 assert!(
19475 response.is_ok(),
19476 "Expected Ok with skip policy, got: {:?}",
19477 response
19478 );
19479 }
19480}