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, ParallelToolsConfig, StreamChunk, StreamingConfig, ToolCall,
185};
186use crate::optimization::{
187 AwaitBeforeNextTurn, BackgroundMaintenanceQueue, BackgroundOverflowPolicy, MainResponseDraft,
188 MaintenanceMode, MaintenanceSequenceKey, RuntimeBranch, RuntimeBranchResult,
189 RuntimeBranchStatus, RuntimeCommitBehavior, RuntimeConfig, RuntimeOptimizationKind,
190 RuntimeTaskPriority, RuntimeTaskPurpose, ScheduledBranchSet, SkillCandidate,
191 StreamingDraftResult, TransitionCandidate, TurnBranchScheduler, TurnOptimizationContext,
192};
193use crate::spec::StorageConfig;
194
195enum ToolCallOutcome {
197 Continue,
199 TransitionFired,
201 Rejected(AgentResponse),
203}
204
205#[derive(Clone)]
206struct MainToolProtocol {
207 choice: Option<ToolChoice>,
208 tool_ids: Vec<String>,
209 definitions: Vec<LLMToolDefinition>,
210}
211
212struct MainProviderResponse {
213 response: LLMResponse,
214 used_native_tools: bool,
215}
216
217struct CommittedTextResponse<'a> {
221 processed_input: &'a str,
222 input_context: &'a HashMap<String, Value>,
223 answer: String,
224 reasoning_mode: ReasoningMode,
225 auto_detected: bool,
226 iterations: u32,
227 thinking_content: Option<String>,
228 all_tool_calls: Vec<ToolCall>,
229}
230
231struct AgentResponseParts {
235 content: String,
236 all_tool_calls: Vec<ToolCall>,
237 reasoning_mode: ReasoningMode,
238 auto_detected: bool,
239 iterations: u32,
240 thinking: Option<String>,
241 reflection_metadata: Option<ReflectionMetadata>,
242}
243
244#[derive(Clone, Copy)]
245struct DisambiguationOwnership {
246 epoch: u64,
247 state_generation: Option<u64>,
248}
249
250enum SkillRouteResult {
252 NoMatch,
254 Response { skill_id: String, content: String },
256 NeedsClarification {
258 response: AgentResponse,
259 ownership: Option<DisambiguationOwnership>,
260 },
261}
262
263enum ParallelTransitionSelection {
265 Candidate(TransitionCandidate),
267 NoMatch,
269 ReservationExhausted,
271}
272
273enum PostLoopResult {
275 NoTransition(String),
277 Transitioned(String),
279 NeedsRedispatch,
282}
283
284struct StateTransitionReservation<'a> {
285 reserved: &'a AtomicBool,
286}
287
288impl Drop for StateTransitionReservation<'_> {
289 fn drop(&mut self) {
290 self.reserved.store(false, Ordering::SeqCst);
291 }
292}
293
294struct RootTurnCleanup<'a> {
295 agent: &'a RuntimeAgent,
296}
297
298impl<'a> RootTurnCleanup<'a> {
299 fn new(agent: &'a RuntimeAgent) -> Self {
300 Self { agent }
301 }
302}
303
304impl Drop for RootTurnCleanup<'_> {
305 fn drop(&mut self) {
306 self.agent.end_root_turn();
307 }
308}
309
310#[derive(Debug)]
312struct RuntimeControlState {
313 snapshot_guard: RwLock<()>,
315 version: AtomicU64,
317 emergency_deny: Arc<AtomicBool>,
319 tool_security_override: RwLock<Option<ToolSecurityEngine>>,
321 tool_scope_override: RwLock<Option<Vec<String>>>,
323}
324
325impl Default for RuntimeControlState {
326 fn default() -> Self {
327 Self {
328 snapshot_guard: RwLock::new(()),
329 version: AtomicU64::new(1),
330 emergency_deny: Arc::new(AtomicBool::new(false)),
331 tool_security_override: RwLock::new(None),
332 tool_scope_override: RwLock::new(None),
333 }
334 }
335}
336
337#[derive(Clone)]
339pub struct RuntimeControlHandle {
340 state: Arc<RuntimeControlState>,
341}
342
343impl RuntimeControlHandle {
344 pub fn version(&self) -> u64 {
346 self.state.version.load(Ordering::SeqCst)
347 }
348
349 fn bump(&self) -> u64 {
350 self.state.version.fetch_add(1, Ordering::SeqCst) + 1
351 }
352
353 pub fn set_tool_security(&self, config: ToolSecurityConfig) -> u64 {
355 self.try_set_tool_security(config)
356 .expect("invalid tool security configuration")
357 }
358
359 pub fn try_set_tool_security(&self, config: ToolSecurityConfig) -> Result<u64> {
361 config.validate()?;
362 let _guard = self.state.snapshot_guard.write();
363 let generation = self.bump();
364 *self.state.tool_security_override.write() = Some(
365 ToolSecurityEngine::new_with_policy_version(config, generation),
366 );
367 Ok(generation)
368 }
369
370 pub fn clear_tool_security_override(&self) -> u64 {
372 let _guard = self.state.snapshot_guard.write();
373 *self.state.tool_security_override.write() = None;
374 self.bump()
375 }
376
377 pub fn set_tool_scope(&self, tool_ids: Vec<String>) -> u64 {
379 let _guard = self.state.snapshot_guard.write();
380 *self.state.tool_scope_override.write() = Some(tool_ids);
381 self.bump()
382 }
383
384 pub fn clear_tool_scope_override(&self) -> u64 {
386 let _guard = self.state.snapshot_guard.write();
387 *self.state.tool_scope_override.write() = None;
388 self.bump()
389 }
390
391 pub fn set_emergency_deny(&self, enabled: bool) -> u64 {
393 let _guard = self.state.snapshot_guard.write();
394 self.state.emergency_deny.store(enabled, Ordering::SeqCst);
395 self.bump()
396 }
397
398 pub fn cancel_all(&self) -> u64 {
400 self.set_emergency_deny(true)
401 }
402}
403
404pub struct RuntimeAgent {
405 info: AgentInfo,
406 llm_registry: Arc<LLMRegistry>,
407 memory: Arc<dyn Memory>,
408 tools: Arc<ToolRegistry>,
409 skills: Vec<SkillDefinition>,
410 skill_router: Option<SkillRouter>,
411 skill_executor: Option<SkillExecutor>,
412 base_system_prompt: String,
413 max_iterations: u32,
414 iteration_count: RwLock<u32>,
415 max_context_tokens: u32,
416 memory_token_budget: Option<MemoryTokenBudget>,
417 recovery_manager: RecoveryManager,
418 tool_security: ToolSecurityEngine,
419 process_processor: Option<ProcessProcessor>,
420 message_filters: RwLock<HashMap<String, Arc<dyn MessageFilter>>>,
421 state_machine: Option<Arc<StateMachine>>,
422 transition_evaluator: Option<Arc<dyn TransitionEvaluator>>,
423 context_manager: Arc<ContextManager>,
424 template_renderer: TemplateRenderer,
425 tool_call_history: RwLock<Vec<ToolCallRecord>>,
426 parallel_tools: ParallelToolsConfig,
427 streaming: StreamingConfig,
428 hooks: Arc<dyn AgentHooks>,
429 hitl_engine: Option<HITLEngine>,
430 approval_handler: Arc<dyn ApprovalHandler>,
431 storage_config: StorageConfig,
432 storage: RwLock<Option<Arc<dyn AgentStorage>>>,
433 storage_init: tokio::sync::Mutex<()>,
434 reasoning_config: ReasoningConfig,
435 reflection_config: ReflectionConfig,
436 disambiguation_manager: Option<DisambiguationManager>,
437 disambiguation_epoch: AtomicU64,
439 disambiguation_admission: tokio::sync::RwLock<()>,
441 state_transition_reserved: AtomicBool,
443 persona_manager: Option<Arc<ai_agents_persona::PersonaManager>>,
445 pending_skill_id: RwLock<Option<String>>,
449 current_plan: RwLock<Option<Plan>>,
450 declared_tool_ids: Option<Vec<String>>,
452 context_initialized: AtomicBool,
454 spawner: Option<Arc<crate::spawner::AgentSpawner>>,
456 spawner_registry: Option<Arc<crate::spawner::AgentRegistry>>,
458 redispatch_depth: RwLock<u32>,
461 active_turn_context: RwLock<Option<TurnOptimizationContext>>,
463 root_user_message_committed: AtomicBool,
465 actor_id: RwLock<Option<String>>,
467 fact_store: RwLock<Option<Arc<ai_agents_facts::FactStore>>>,
469 fact_extractor: RwLock<Option<Arc<dyn ai_agents_facts::FactExtractor>>>,
472 actor_facts_cache: Arc<RwLock<HashMap<String, Vec<ai_agents_core::KeyFact>>>>,
474 messages_since_extraction: Arc<RwLock<usize>>,
476 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
478 facts_config: Option<ai_agents_facts::FactsConfig>,
480 session_metadata: RwLock<ai_agents_core::SessionMetadata>,
482 current_session_id: RwLock<Option<String>>,
484 relationship_manager: Option<Arc<RelationshipManager>>,
486 observability_manager: Option<Arc<ObservabilityManager>>,
488 runtime_config: RuntimeConfig,
490 background_maintenance: Arc<BackgroundMaintenanceQueue>,
492 resource_locks: ToolResourceLocks,
494 runtime_control: Arc<RuntimeControlState>,
496}
497
498impl std::fmt::Debug for RuntimeAgent {
499 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
500 f.debug_struct("RuntimeAgent")
501 .field("info", &self.info)
502 .field("base_system_prompt", &self.base_system_prompt)
503 .field("max_iterations", &self.max_iterations)
504 .field("skills_count", &self.skills.len())
505 .field("max_context_tokens", &self.max_context_tokens)
506 .field("has_state_machine", &self.state_machine.is_some())
507 .field("parallel_tools", &self.parallel_tools)
508 .field("streaming", &self.streaming)
509 .field("has_hooks", &true)
510 .field("has_hitl", &self.hitl_engine.is_some())
511 .field("storage_type", &self.storage_config.storage_type())
512 .field("reasoning_mode", &self.reasoning_config.mode)
513 .field("reflection_enabled", &self.reflection_config.enabled)
514 .field("declared_tool_ids", &self.declared_tool_ids)
515 .field("has_persona", &self.persona_manager.is_some())
516 .field("has_observability", &self.observability_manager.is_some())
517 .finish_non_exhaustive()
518 }
519}
520
521struct ObservabilityClarificationObserver;
522
523impl ClarificationObserver for ObservabilityClarificationObserver {
524 fn observe_question<'a>(
526 &'a self,
527 future: ClarificationQuestionFuture<'a>,
528 ) -> ClarificationQuestionFuture<'a> {
529 Box::pin(async move {
530 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
531 })
532 }
533
534 fn observe_parse<'a>(
536 &'a self,
537 future: ClarificationParseFuture<'a>,
538 ) -> ClarificationParseFuture<'a> {
539 Box::pin(async move {
540 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
541 })
542 }
543
544 fn observe_confirmation_parse<'a>(
546 &'a self,
547 future: ConfirmationParseFuture<'a>,
548 ) -> ConfirmationParseFuture<'a> {
549 Box::pin(async move {
550 with_observation_purpose(ObservationPurpose::DisambiguationClarification, future).await
551 })
552 }
553}
554
555struct ObservabilityProcessStageObserver;
556
557impl ProcessStageObserver for ObservabilityProcessStageObserver {
558 fn observe<'a>(
560 &'a self,
561 hint: ProcessPurposeHint,
562 future: ProcessStageFuture<'a>,
563 ) -> ProcessStageFuture<'a> {
564 Box::pin(async move {
565 with_observation_purpose(observation_purpose_for_process(hint), future).await
566 })
567 }
568}
569
570struct RegistryLLMGetter {
571 registry: Arc<LLMRegistry>,
572}
573
574impl LLMGetter for RegistryLLMGetter {
575 fn get_llm(&self, alias: &str) -> Option<Arc<dyn LLMProvider>> {
576 self.registry.get(alias).ok()
577 }
578}
579
580impl RuntimeAgent {
581 #[allow(clippy::too_many_arguments)]
582 pub fn new(
583 info: AgentInfo,
584 llm_registry: Arc<LLMRegistry>,
585 memory: Arc<dyn Memory>,
586 tools: Arc<ToolRegistry>,
587 skills: Vec<SkillDefinition>,
588 system_prompt: String,
589 max_iterations: u32,
590 ) -> Self {
591 let (skill_router, skill_executor) = if !skills.is_empty() {
592 let router_llm = llm_registry.router().ok();
593 let router = router_llm.map(|llm| SkillRouter::new(llm, skills.clone()));
594 let executor = SkillExecutor::new(llm_registry.clone(), tools.clone());
595 (router, Some(executor))
596 } else {
597 (None, None)
598 };
599
600 let context_manager =
601 ContextManager::new(HashMap::new(), info.name.clone(), info.version.clone());
602
603 Self {
604 info,
605 llm_registry,
606 memory,
607 tools,
608 skills,
609 skill_router,
610 skill_executor,
611 base_system_prompt: system_prompt,
612 max_iterations,
613 iteration_count: RwLock::new(0),
614 max_context_tokens: 128000,
615 memory_token_budget: None,
616 recovery_manager: RecoveryManager::default(),
617 tool_security: ToolSecurityEngine::default(),
618 process_processor: None,
619 message_filters: RwLock::new(HashMap::new()),
620 state_machine: None,
621 transition_evaluator: None,
622 context_manager: Arc::new(context_manager),
623 template_renderer: TemplateRenderer::new(),
624 tool_call_history: RwLock::new(Vec::new()),
625 parallel_tools: ParallelToolsConfig::default(),
626 streaming: StreamingConfig::default(),
627 hooks: Arc::new(NoopHooks),
628 hitl_engine: None,
629 approval_handler: Arc::new(RejectAllHandler::new()),
630 storage_config: StorageConfig::default(),
631 storage: RwLock::new(None),
632 storage_init: tokio::sync::Mutex::new(()),
633 reasoning_config: ReasoningConfig::default(),
634 reflection_config: ReflectionConfig::default(),
635 disambiguation_manager: None,
636 disambiguation_epoch: AtomicU64::new(0),
637 disambiguation_admission: tokio::sync::RwLock::new(()),
638 state_transition_reserved: AtomicBool::new(false),
639 persona_manager: None,
640 pending_skill_id: RwLock::new(None),
641 current_plan: RwLock::new(None),
642 declared_tool_ids: None,
643 context_initialized: AtomicBool::new(false),
644 spawner: None,
645 spawner_registry: None,
646 redispatch_depth: RwLock::new(0),
647 active_turn_context: RwLock::new(None),
648 root_user_message_committed: AtomicBool::new(false),
649 actor_id: RwLock::new(None),
650 fact_store: RwLock::new(None),
651 fact_extractor: RwLock::new(None),
652 actor_facts_cache: Arc::new(RwLock::new(HashMap::new())),
653 messages_since_extraction: Arc::new(RwLock::new(0)),
654 actor_memory_config: None,
655 facts_config: None,
656 session_metadata: RwLock::new(ai_agents_core::SessionMetadata::default()),
657 current_session_id: RwLock::new(None),
658 relationship_manager: None,
659 observability_manager: None,
660 runtime_config: RuntimeConfig::default(),
661 background_maintenance: Arc::new(BackgroundMaintenanceQueue::default()),
662 resource_locks: new_tool_resource_locks(),
663 runtime_control: Arc::new(RuntimeControlState::default()),
664 }
665 }
666
667 pub fn with_declared_tool_ids(mut self, ids: Option<Vec<String>>) -> Self {
668 self.declared_tool_ids = ids;
669 self
670 }
671
672 pub fn with_storage_config(mut self, config: StorageConfig) -> Self {
673 self.storage_config = config;
674 self
675 }
676
677 pub fn with_storage(self, storage: Arc<dyn AgentStorage>) -> Self {
678 *self.storage.write() = Some(storage);
679 self
680 }
681
682 pub(crate) fn with_shared_resource_locks(mut self, locks: ToolResourceLocks) -> Self {
683 self.resource_locks = locks;
684 self
685 }
686
687 pub fn with_reasoning(mut self, config: ReasoningConfig) -> Self {
688 self.reasoning_config = config;
689 self
690 }
691
692 pub fn with_reflection(mut self, config: ReflectionConfig) -> Self {
693 self.reflection_config = config;
694 self
695 }
696
697 pub fn with_relationships(mut self, manager: Arc<RelationshipManager>) -> Self {
699 self.relationship_manager = Some(manager);
700 self
701 }
702
703 pub fn with_observability(mut self, manager: Arc<ObservabilityManager>) -> Self {
705 self.observability_manager = Some(manager);
706 self
707 }
708
709 pub fn with_runtime_config(mut self, config: RuntimeConfig) -> Self {
711 let max_tasks = config.optimization.post_turn.max_background_tasks;
712 self.background_maintenance = Arc::new(BackgroundMaintenanceQueue::new(max_tasks));
713 self.runtime_config = config;
714 self
715 }
716
717 pub fn runtime_config(&self) -> &RuntimeConfig {
719 &self.runtime_config
720 }
721
722 pub async fn flush_background_tasks(&self) -> Result<()> {
724 self.background_maintenance.flush_all().await
725 }
726
727 pub async fn flush_background_tasks_for_actor(&self, actor_id: &str) -> Result<()> {
729 self.background_maintenance.flush_scope(actor_id).await
730 }
731
732 pub async fn flush_background_tasks_for_purpose(
734 &self,
735 purpose: RuntimeTaskPurpose,
736 ) -> Result<()> {
737 self.background_maintenance.flush_purpose(purpose).await
738 }
739
740 pub async fn flush_background_tasks_for_actor_purpose(
742 &self,
743 actor_id: &str,
744 purpose: RuntimeTaskPurpose,
745 ) -> Result<()> {
746 self.background_maintenance
747 .flush_scope_purpose(actor_id, purpose)
748 .await
749 }
750
751 pub async fn shutdown_background_tasks(&self) -> Result<()> {
753 self.flush_background_tasks().await
754 }
755
756 pub fn observability(&self) -> Option<Arc<ObservabilityManager>> {
758 self.observability_manager.clone()
759 }
760
761 async fn export_observability_if_configured(&self) {
763 let Some(manager) = self.observability_manager.as_ref() else {
764 return;
765 };
766 let export = &manager.config().export;
767 if !export.write_report && !export.write_raw_events {
768 return;
769 }
770 if let Err(error) = manager.export().await {
771 warn!(error = %error, "Observability export failed");
772 }
773 }
774
775 pub fn relationship_manager(&self) -> Option<Arc<RelationshipManager>> {
777 self.relationship_manager.clone()
778 }
779
780 fn current_turn_actor_context(&self) -> Option<crate::TurnActorContext> {
781 current_turn_actor_context()
782 }
783
784 fn effective_actor_id(&self) -> Option<String> {
785 self.current_turn_actor_context()
786 .and_then(|ctx| ctx.effective_actor_id().map(|id| id.to_string()))
787 .or_else(|| self.actor_id.read().clone())
788 }
789
790 fn effective_origin_actor_id(&self) -> Option<String> {
791 self.current_turn_actor_context()
792 .and_then(|ctx| ctx.origin_actor_id.clone())
793 .or_else(|| self.actor_id.read().clone())
794 }
795
796 fn record_session_actor_if_needed(&self) {
797 if let Some(actor_id) = self.effective_origin_actor_id() {
798 let mut meta = self.session_metadata.write();
799 meta.actor_id = Some(actor_id.clone());
800 if !meta.actors.iter().any(|a| a == &actor_id) {
801 meta.actors.push(actor_id);
802 }
803 }
804 }
805
806 fn outbound_actor_context(&self) -> crate::TurnActorContext {
807 let mut context = self.current_turn_actor_context().unwrap_or_default();
808 if context.origin_actor_id.is_none() {
809 context.origin_actor_id = self.effective_origin_actor_id();
810 }
811 context.sender_agent_id = Some(self.info.id.clone());
812 context
813 }
814
815 fn observation_session_id(&self) -> Option<String> {
817 let mut current = self.current_session_id.write();
818 if current.is_none() {
819 *current = Some(new_observation_session_id());
820 }
821 current.clone()
822 }
823
824 fn build_observation_context(&self, actor_id: Option<String>) -> Option<SpanContext> {
826 let manager = self.observability_manager.as_ref()?;
827 let context = self.build_context_with_overlays();
828 let language = resolve_language_from_context(manager.config(), &context);
829 let context = current_observation_context()
830 .map(|parent| parent.child_for_agent(self.info.id.clone()).with_new_turn())
831 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
832 Some(
833 context
834 .with_actor(actor_id.or_else(|| self.effective_actor_id()))
835 .with_session(self.observation_session_id())
836 .with_state(self.current_state())
837 .with_language(Some(language)),
838 )
839 }
840
841 fn current_runtime_observation_context(
843 &self,
844 purpose: ObservationPurpose,
845 ) -> 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 mut observation = current_observation_context()
850 .unwrap_or_else(|| SpanContext::new_root(self.info.id.clone()));
851 observation.agent_id = self.info.id.clone();
852 observation.actor_id = self.effective_actor_id();
853 observation.session_id = self.observation_session_id();
854 observation.state = self.current_state();
855 observation.language = Some(language);
856 observation.purpose = purpose;
857 Some(observation)
858 }
859
860 async fn observe_purpose<F, T>(&self, purpose: ObservationPurpose, future: F) -> T
862 where
863 F: Future<Output = T>,
864 {
865 if let Some(context) = self.current_runtime_observation_context(purpose) {
866 with_observation_context(context, future).await
867 } else {
868 future.await
869 }
870 }
871
872 fn chat_with_actor_context_boxed<'a>(
874 &'a self,
875 input: &'a str,
876 actor_context: crate::TurnActorContext,
877 ) -> Pin<Box<dyn Future<Output = Result<AgentResponse>> + Send + 'a>> {
878 Box::pin(async move {
879 let actor_id = actor_context.effective_actor_id().map(str::to_string);
880 let run = async move {
881 scope_actor_context(
882 actor_context,
883 Box::pin(async move { self.run_loop(input).await }),
884 )
885 .await
886 };
887 let result = if let Some(context) = self.build_observation_context(actor_id) {
888 with_observation_context(context, run).await
889 } else {
890 run.await
891 };
892 self.export_observability_if_configured().await;
893 result
894 })
895 }
896
897 pub async fn chat_with_actor_context(
901 &self,
902 input: &str,
903 actor_context: crate::TurnActorContext,
904 ) -> Result<AgentResponse> {
905 self.chat_with_actor_context_boxed(input, actor_context)
906 .await
907 }
908
909 pub async fn chat_as_actor(&self, actor_id: &str, input: &str) -> Result<AgentResponse> {
911 let actor_context = crate::TurnActorContext::new().with_origin_actor(actor_id);
912 self.chat_with_actor_context(input, actor_context).await
913 }
914
915 pub async fn load_actor_relationship(&self) -> Result<()> {
917 self.maybe_load_actor_relationship().await;
918 Ok(())
919 }
920
921 pub async fn update_relationship_dimension(
923 &self,
924 dimension: &str,
925 delta: f64,
926 reason: Option<&str>,
927 ) -> Result<ai_agents_relationships::DimensionChange> {
928 self.update_relationship_dimension_for_perspective(
929 ai_agents_relationships::RelationshipPerspective::AgentToActor,
930 dimension,
931 delta,
932 reason,
933 )
934 .await
935 }
936
937 pub async fn update_relationship_dimension_for_perspective(
941 &self,
942 perspective: ai_agents_relationships::RelationshipPerspective,
943 dimension: &str,
944 delta: f64,
945 reason: Option<&str>,
946 ) -> Result<ai_agents_relationships::DimensionChange> {
947 let manager = self
948 .relationship_manager
949 .as_ref()
950 .ok_or_else(|| AgentError::Config("Relationship memory is not configured".into()))?;
951 let actor_id = self.effective_actor_id().ok_or_else(|| {
952 AgentError::Config("No actor ID set. Use set_actor_id() first".into())
953 })?;
954 let change = manager.update_dimension_for_perspective(
955 &actor_id,
956 perspective,
957 dimension,
958 delta,
959 1.0,
960 reason.unwrap_or("manual relationship update"),
961 )?;
962 self.persist_actor_relationship(&actor_id).await?;
963 info!(
964 actor_id = %actor_id,
965 perspective = %change.perspective,
966 dimension = %change.dimension,
967 delta = change.delta,
968 current = change.current,
969 "relationship updated manually"
970 );
971 self.hooks
972 .on_relationship_change(&actor_id, std::slice::from_ref(&change))
973 .await;
974 Ok(change)
975 }
976
977 pub fn reasoning_config(&self) -> &ReasoningConfig {
978 &self.reasoning_config
979 }
980
981 pub fn reflection_config(&self) -> &ReflectionConfig {
982 &self.reflection_config
983 }
984
985 pub fn with_facts_config(
988 mut self,
989 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
990 facts_config: Option<ai_agents_facts::FactsConfig>,
991 ) -> Self {
992 self.actor_memory_config = actor_memory_config;
993 self.facts_config = facts_config;
994 self
995 }
996
997 pub fn with_facts(
1000 mut self,
1001 store: Arc<ai_agents_facts::FactStore>,
1002 extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>>,
1003 actor_memory_config: Option<ai_agents_facts::ActorMemoryConfig>,
1004 facts_config: Option<ai_agents_facts::FactsConfig>,
1005 ) -> Self {
1006 *self.fact_store.write() = Some(store);
1007 *self.fact_extractor.write() = extractor;
1008 self.actor_memory_config = actor_memory_config;
1009 self.facts_config = facts_config;
1010 self
1011 }
1012
1013 pub fn fact_store(&self) -> Option<Arc<ai_agents_facts::FactStore>> {
1015 self.fact_store.read().clone()
1016 }
1017
1018 pub fn actor_id(&self) -> Option<String> {
1020 self.actor_id.read().clone()
1021 }
1022
1023 pub fn set_actor_id(&self, actor_id: &str) -> ai_agents_core::Result<()> {
1025 *self.actor_id.write() = Some(actor_id.to_string());
1026 {
1027 let mut meta = self.session_metadata.write();
1028 meta.actor_id = Some(actor_id.to_string());
1029 if !meta.actors.iter().any(|a| a == actor_id) {
1030 meta.actors.push(actor_id.to_string());
1031 }
1032 }
1033 Ok(())
1034 }
1035
1036 pub fn clear_actor_id(&self) {
1038 *self.actor_id.write() = None;
1039 self.session_metadata.write().actor_id = None;
1040 }
1041
1042 pub fn set_user_id(&self, user_id: &str) -> ai_agents_core::Result<()> {
1044 self.set_actor_id(user_id)
1045 }
1046
1047 pub async fn load_actor_memory(&self) -> ai_agents_core::Result<()> {
1049 let actor_id = match self.effective_actor_id() {
1050 Some(id) => id,
1051 None => return Ok(()),
1052 };
1053
1054 let store_opt = self.fact_store.read().clone();
1055 if let Some(store) = store_opt {
1056 let facts = store.get_facts(&actor_id).await?;
1057 let count = facts.len();
1058 self.actor_facts_cache
1059 .write()
1060 .insert(actor_id.clone(), facts);
1061 self.hooks.on_actor_memory_loaded(&actor_id, count).await;
1062 tracing::debug!("loaded {} facts for actor {}", count, actor_id);
1063 }
1064
1065 Ok(())
1066 }
1067
1068 async fn maybe_load_actor_memory(&self) {
1070 let Some(actor_id) = self.effective_actor_id() else {
1071 return;
1072 };
1073 if self.actor_facts_cache.read().contains_key(&actor_id) {
1074 return;
1075 }
1076 let _ = self.load_actor_memory().await;
1077 }
1078
1079 async fn pre_turn_session_lifecycle(&self) {
1081 if *self.redispatch_depth.read() > 0 {
1082 return;
1083 }
1084 self.resolve_actor_id_from_context();
1085 self.await_background_before_next_turn().await;
1086 self.record_session_actor_if_needed();
1087 self.maybe_load_actor_memory().await;
1088 self.maybe_load_actor_relationship().await;
1089 *self.messages_since_extraction.write() += 1;
1090 }
1091
1092 async fn post_turn_session_lifecycle(&self) -> Result<()> {
1094 if *self.redispatch_depth.read() > 0 {
1095 return Ok(());
1096 }
1097 *self.messages_since_extraction.write() += 1;
1098 self.run_post_turn_maintenance().await
1099 }
1100
1101 fn begin_root_turn(&self) {
1103 if *self.redispatch_depth.read() == 0 {
1104 let mut guard = self.active_turn_context.write();
1105 if guard.is_none() {
1106 self.root_user_message_committed
1107 .store(false, Ordering::SeqCst);
1108 let max_calls = self
1109 .runtime_config
1110 .optimization
1111 .max_speculative_llm_calls_per_turn;
1112 *guard = Some(TurnOptimizationContext::new(
1113 String::new(),
1114 HashMap::new(),
1115 max_calls,
1116 ));
1117 }
1118 }
1119 }
1120
1121 fn update_active_turn_context(
1122 &self,
1123 processed_input: &str,
1124 input_context: HashMap<String, Value>,
1125 ) {
1126 if *self.redispatch_depth.read() > 0 {
1127 return;
1128 }
1129 let max_calls = self
1130 .runtime_config
1131 .optimization
1132 .max_speculative_llm_calls_per_turn;
1133 let mut guard = self.active_turn_context.write();
1134 match guard.as_mut() {
1135 Some(context) => {
1136 context.processed_input = processed_input.to_string();
1137 context.input_context = input_context;
1138 context.max_speculative_llm_calls = max_calls;
1139 }
1140 None => {
1141 *guard = Some(TurnOptimizationContext::new(
1142 processed_input,
1143 input_context,
1144 max_calls,
1145 ));
1146 }
1147 }
1148 }
1149
1150 async fn commit_root_user_message(&self, processed_input: &str) -> Result<()> {
1152 if *self.redispatch_depth.read() > 0 {
1153 return Ok(());
1154 }
1155 if !self
1156 .root_user_message_committed
1157 .swap(true, Ordering::SeqCst)
1158 {
1159 self.memory
1160 .add_message(ChatMessage::user(processed_input))
1161 .await?;
1162 if let Some(context) = self.active_turn_context.write().as_mut() {
1163 context.mark_user_message_committed();
1164 }
1165 }
1166 Ok(())
1167 }
1168
1169 fn end_root_turn(&self) {
1171 if *self.redispatch_depth.read() == 0 {
1172 self.root_user_message_committed
1173 .store(false, Ordering::SeqCst);
1174 *self.active_turn_context.write() = None;
1175 }
1176 }
1177
1178 fn reserve_active_speculative_llm_call(&self, kind: RuntimeOptimizationKind) -> bool {
1179 self.begin_root_turn();
1180 let mut guard = self.active_turn_context.write();
1181 let Some(context) = guard.as_mut() else {
1182 return false;
1183 };
1184 context.reserve_speculative_llm_call_for(kind)
1185 }
1186
1187 fn branch_context_preview(&self) -> String {
1188 let context = self.build_context_with_overlays();
1189 let mut value = serde_json::to_string_pretty(&context).unwrap_or_else(|_| "{}".to_string());
1190 const MAX_CONTEXT_PREVIEW_CHARS: usize = 2048;
1191 if value.chars().count() > MAX_CONTEXT_PREVIEW_CHARS {
1192 value = value
1193 .chars()
1194 .take(MAX_CONTEXT_PREVIEW_CHARS)
1195 .collect::<String>();
1196 value.push_str("...");
1197 }
1198 value
1199 }
1200
1201 async fn await_background_before_next_turn(&self) {
1203 let optimization = &self.runtime_config.optimization;
1204 if !optimization.enabled {
1205 return;
1206 }
1207 let actor_id = self.effective_actor_id();
1208 let post = &optimization.post_turn;
1209 self.await_background_task(
1210 post.facts.await_before_next_turn,
1211 RuntimeTaskPurpose::PostTurnFacts,
1212 actor_id.as_deref(),
1213 "facts",
1214 )
1215 .await;
1216 self.await_background_task(
1217 post.relationships.await_before_next_turn,
1218 RuntimeTaskPurpose::PostTurnRelationship,
1219 actor_id.as_deref(),
1220 "relationships",
1221 )
1222 .await;
1223 }
1224
1225 async fn await_background_task(
1226 &self,
1227 policy: AwaitBeforeNextTurn,
1228 purpose: RuntimeTaskPurpose,
1229 actor_id: Option<&str>,
1230 label: &str,
1231 ) {
1232 match policy {
1233 AwaitBeforeNextTurn::Never => {}
1234 AwaitBeforeNextTurn::Always => {
1235 if let Err(error) = self.flush_background_tasks_for_purpose(purpose).await {
1236 warn!(label = label, error = %error, "background maintenance flush failed");
1237 }
1238 }
1239 AwaitBeforeNextTurn::SameActor => {
1240 if let Some(actor_id) = actor_id
1241 && let Err(error) = self
1242 .flush_background_tasks_for_actor_purpose(actor_id, purpose)
1243 .await
1244 {
1245 warn!(label = label, actor_id = %actor_id, error = %error, "actor background maintenance flush failed");
1246 }
1247 }
1248 }
1249 }
1250
1251 async fn run_post_turn_maintenance(&self) -> Result<()> {
1253 let optimization = &self.runtime_config.optimization;
1254 if !optimization.enabled {
1255 self.auto_extract_facts().await;
1256 self.auto_update_relationship().await;
1257 return Ok(());
1258 }
1259
1260 let facts_mode = effective_maintenance_mode(
1261 optimization.post_turn.facts.mode,
1262 optimization.parallel_post_turn_memory,
1263 );
1264 let relationships_mode = effective_maintenance_mode(
1265 optimization.post_turn.relationships.mode,
1266 optimization.parallel_post_turn_memory,
1267 );
1268
1269 match (facts_mode, relationships_mode) {
1270 (MaintenanceMode::InlineSerial, MaintenanceMode::InlineSerial) => {
1271 self.auto_extract_facts().await;
1272 self.auto_update_relationship().await;
1273 }
1274 (MaintenanceMode::InlineParallel, MaintenanceMode::InlineParallel) => {
1275 let facts = self.auto_extract_facts();
1276 let relationships = self.auto_update_relationship();
1277 tokio::join!(facts, relationships);
1278 }
1279 (MaintenanceMode::Background, MaintenanceMode::Background) => {
1280 self.schedule_facts_background().await?;
1281 self.schedule_relationship_background().await?;
1282 }
1283 (MaintenanceMode::Background, MaintenanceMode::InlineParallel)
1284 | (MaintenanceMode::Background, MaintenanceMode::InlineSerial) => {
1285 self.schedule_facts_background().await?;
1286 self.auto_update_relationship().await;
1287 }
1288 (MaintenanceMode::InlineParallel, MaintenanceMode::Background)
1289 | (MaintenanceMode::InlineSerial, MaintenanceMode::Background) => {
1290 self.auto_extract_facts().await;
1291 self.schedule_relationship_background().await?;
1292 }
1293 _ => {
1294 self.auto_extract_facts().await;
1295 self.auto_update_relationship().await;
1296 }
1297 }
1298 Ok(())
1299 }
1300
1301 async fn schedule_facts_background(&self) -> Result<()> {
1302 let policy = self.runtime_config.optimization.post_turn.facts.clone();
1303 let should_extract = self
1304 .facts_config
1305 .as_ref()
1306 .map(|c| c.enabled && c.auto_extract)
1307 .unwrap_or(false);
1308 if !should_extract {
1309 return Ok(());
1310 }
1311 let msgs_since = *self.messages_since_extraction.read();
1312 if msgs_since < 2 {
1313 return Ok(());
1314 }
1315 let Some(actor_id) = self.effective_actor_id() else {
1316 self.record_skipped_maintenance(
1317 "facts",
1318 ObservationPurpose::FactsExtraction,
1319 "missing_actor",
1320 Some(&policy),
1321 );
1322 return Ok(());
1323 };
1324 let Some(extractor) = self.fact_extractor.read().clone() else {
1325 return Ok(());
1326 };
1327 let messages = match self.memory.get_messages(None).await {
1328 Ok(messages) => messages,
1329 Err(error) => {
1330 warn!(error = %error, "failed to snapshot messages for fact extraction");
1331 return Ok(());
1332 }
1333 };
1334 let recent: Vec<_> = messages
1335 .iter()
1336 .rev()
1337 .take(msgs_since)
1338 .rev()
1339 .cloned()
1340 .collect();
1341 if recent.is_empty() {
1342 return Ok(());
1343 }
1344 let existing = self
1345 .actor_facts_cache
1346 .read()
1347 .get(&actor_id)
1348 .cloned()
1349 .unwrap_or_default();
1350 let categories = self
1351 .facts_config
1352 .as_ref()
1353 .map(|c| c.custom_categories.clone())
1354 .unwrap_or_default();
1355 let store = self.fact_store.read().clone();
1356 let cache = Arc::clone(&self.actor_facts_cache);
1357 let counter = Arc::clone(&self.messages_since_extraction);
1358 let hooks = Arc::clone(&self.hooks);
1359 let agent_id = self.info.id.clone();
1360 let observation = current_observation_context();
1361 let key = MaintenanceSequenceKey::actor(
1362 agent_id,
1363 actor_id.clone(),
1364 RuntimeTaskPurpose::PostTurnFacts,
1365 );
1366 let actor_for_task = actor_id.clone();
1367 let task = async move {
1368 let run = async move {
1369 let facts = extractor
1370 .extract(&recent, &existing, Some(&actor_for_task), &categories)
1371 .await?;
1372 if !facts.is_empty() {
1373 if let Some(store) = store {
1374 let authoritative = store.add_facts(&actor_for_task, facts.clone()).await?;
1375 cache.write().insert(actor_for_task.clone(), authoritative);
1376 } else {
1377 cache
1378 .write()
1379 .entry(actor_for_task.clone())
1380 .or_default()
1381 .extend(facts.clone());
1382 }
1383 {
1384 let mut count = counter.write();
1385 if *count <= msgs_since {
1386 *count = 0;
1387 } else {
1388 *count -= msgs_since;
1389 }
1390 }
1391 hooks.on_facts_extracted(&actor_for_task, &facts).await;
1392 }
1393 Ok(())
1394 };
1395 if let Some(context) = observation {
1396 with_observation_context(
1397 context.with_purpose(ObservationPurpose::FactsExtraction),
1398 run,
1399 )
1400 .await
1401 } else {
1402 run.await
1403 }
1404 };
1405 self.spawn_or_handle_background(Some(key), task, "facts", &policy)
1406 .await
1407 }
1408
1409 async fn schedule_relationship_background(&self) -> Result<()> {
1410 let policy = self
1411 .runtime_config
1412 .optimization
1413 .post_turn
1414 .relationships
1415 .clone();
1416 let Some(manager) = self.relationship_manager.as_ref().cloned() else {
1417 return Ok(());
1418 };
1419 let Some(actor_id) = self.effective_actor_id() else {
1420 self.record_skipped_maintenance(
1421 "relationships",
1422 ObservationPurpose::RelationshipUpdate,
1423 "missing_actor",
1424 Some(&policy),
1425 );
1426 return Ok(());
1427 };
1428 let recent_messages = manager.config().auto_update.recent_messages;
1429 let messages = match self.memory.get_messages(Some(recent_messages)).await {
1430 Ok(messages) => messages,
1431 Err(error) => {
1432 warn!(actor = %actor_id, error = %error, "failed to snapshot messages for relationship update");
1433 return Ok(());
1434 }
1435 };
1436 let storage = self.storage.read().clone();
1437 let hooks = Arc::clone(&self.hooks);
1438 let agent_id = self.info.id.clone();
1439 let observation = current_observation_context();
1440 let key = MaintenanceSequenceKey::actor(
1441 agent_id.clone(),
1442 actor_id.clone(),
1443 RuntimeTaskPurpose::PostTurnRelationship,
1444 );
1445 let actor_for_task = actor_id.clone();
1446 let task = async move {
1447 let run = async move {
1448 if manager.config().auto_update.enabled {
1449 let update = manager.auto_update(&actor_for_task, &messages).await?;
1450 if !update.changes.is_empty() {
1451 hooks
1452 .on_relationship_change(&actor_for_task, &update.changes)
1453 .await;
1454 }
1455 if let Some(ref event) = update.event {
1456 hooks.on_notable_event(&actor_for_task, event).await;
1457 }
1458 }
1459 if manager.config().persistence.enabled
1460 && let (Some(storage), Some(value)) =
1461 (storage, manager.relationship_as_value(&actor_for_task)?)
1462 {
1463 storage
1464 .save_relationship(&agent_id, &actor_for_task, &value)
1465 .await?;
1466 }
1467 Ok(())
1468 };
1469 if let Some(context) = observation {
1470 with_observation_context(
1471 context.with_purpose(ObservationPurpose::RelationshipUpdate),
1472 run,
1473 )
1474 .await
1475 } else {
1476 run.await
1477 }
1478 };
1479 self.spawn_or_handle_background(Some(key), task, "relationships", &policy)
1480 .await
1481 }
1482
1483 async fn spawn_or_handle_background<F>(
1485 &self,
1486 key: Option<MaintenanceSequenceKey>,
1487 task: F,
1488 label: &'static str,
1489 policy: &crate::optimization::config::MaintenanceTaskPolicy,
1490 ) -> Result<()>
1491 where
1492 F: Future<Output = Result<()>> + Send + 'static,
1493 {
1494 if self.background_maintenance.is_full() {
1495 match self
1496 .runtime_config
1497 .optimization
1498 .post_turn
1499 .on_background_overflow
1500 {
1501 BackgroundOverflowPolicy::RunInline => {
1502 record_background_maintenance_event(
1503 self.observability_manager.as_ref(),
1504 label,
1505 EventStatus::Success,
1506 0,
1507 "inline_overflow",
1508 None,
1509 Some(policy),
1510 );
1511 let start = Instant::now();
1512 match task.await {
1513 Ok(()) => record_background_maintenance_event(
1514 self.observability_manager.as_ref(),
1515 label,
1516 EventStatus::Success,
1517 start.elapsed().as_millis() as u64,
1518 "inline_completed",
1519 None,
1520 Some(policy),
1521 ),
1522 Err(error) => {
1523 warn!(label = label, error = %error, "inline maintenance fallback failed");
1524 record_background_maintenance_event(
1525 self.observability_manager.as_ref(),
1526 label,
1527 EventStatus::Error,
1528 start.elapsed().as_millis() as u64,
1529 "inline_failed",
1530 Some(error.to_string()),
1531 Some(policy),
1532 );
1533 return Err(error);
1534 }
1535 }
1536 }
1537 BackgroundOverflowPolicy::Drop => {
1538 self.record_skipped_maintenance(
1539 label,
1540 ObservationPurpose::Other(label.to_string()),
1541 "queue_full",
1542 Some(policy),
1543 );
1544 }
1545 BackgroundOverflowPolicy::Error => {
1546 record_background_maintenance_event(
1547 self.observability_manager.as_ref(),
1548 label,
1549 EventStatus::Error,
1550 0,
1551 "queue_full",
1552 None,
1553 Some(policy),
1554 );
1555 warn!(label = label, "background maintenance queue full");
1556 return Err(AgentError::Other(format!(
1557 "background maintenance queue is full for {}",
1558 label
1559 )));
1560 }
1561 }
1562 return Ok(());
1563 }
1564
1565 record_background_maintenance_event(
1566 self.observability_manager.as_ref(),
1567 label,
1568 EventStatus::Success,
1569 0,
1570 "scheduled",
1571 None,
1572 Some(policy),
1573 );
1574 let manager = self.observability_manager.clone();
1575 let policy_for_task = policy.clone();
1576 let observed_task = async move {
1577 let start = Instant::now();
1578 let result = task.await;
1579 match &result {
1580 Ok(()) => record_background_maintenance_event(
1581 manager.as_ref(),
1582 label,
1583 EventStatus::Success,
1584 start.elapsed().as_millis() as u64,
1585 "completed",
1586 None,
1587 Some(&policy_for_task),
1588 ),
1589 Err(error) => record_background_maintenance_event(
1590 manager.as_ref(),
1591 label,
1592 EventStatus::Error,
1593 start.elapsed().as_millis() as u64,
1594 "failed",
1595 Some(error.to_string()),
1596 Some(&policy_for_task),
1597 ),
1598 }
1599 result
1600 };
1601
1602 if let Err(error) = self.background_maintenance.spawn(key, observed_task) {
1603 record_background_maintenance_event(
1604 self.observability_manager.as_ref(),
1605 label,
1606 EventStatus::Error,
1607 0,
1608 "spawn_failed",
1609 Some(error.to_string()),
1610 Some(policy),
1611 );
1612 warn!(label = label, error = %error, "background maintenance spawn failed");
1613 return Err(error);
1614 }
1615 Ok(())
1616 }
1617
1618 fn record_skipped_maintenance(
1620 &self,
1621 label: &str,
1622 purpose: ObservationPurpose,
1623 reason: &str,
1624 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
1625 ) {
1626 if let Some(manager) = self.observability_manager.as_ref() {
1627 let mut tags = background_maintenance_tags(label, "skipped", Some(reason), policy);
1628 tags.insert("runtime.skip_reason".to_string(), reason.to_string());
1629 manager.record_lifecycle_event(
1630 EventType::MemoryOperation {
1631 operation: format!("{}_maintenance", label),
1632 },
1633 purpose,
1634 EventStatus::Skipped,
1635 0,
1636 tags,
1637 None,
1638 );
1639 }
1640 }
1641
1642 pub fn actor_facts(&self) -> Vec<ai_agents_core::KeyFact> {
1644 let Some(actor_id) = self.effective_actor_id() else {
1645 return Vec::new();
1646 };
1647 self.actor_facts_cache
1648 .read()
1649 .get(&actor_id)
1650 .cloned()
1651 .unwrap_or_default()
1652 }
1653
1654 pub fn relationship_memory_text(&self) -> Option<String> {
1656 self.format_relationship_for_context().map(|(_, text)| text)
1657 }
1658
1659 pub async fn extract_facts(
1661 &self,
1662 last_n: usize,
1663 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1664 self.extract_facts_with_source(last_n, "manual").await
1665 }
1666
1667 async fn extract_facts_with_source(
1668 &self,
1669 last_n: usize,
1670 source: &'static str,
1671 ) -> ai_agents_core::Result<Vec<ai_agents_core::KeyFact>> {
1672 let extractor = match self.fact_extractor.read().clone() {
1673 Some(e) => e,
1674 None => return Ok(vec![]),
1675 };
1676
1677 let messages = self.memory.get_messages(None).await?;
1678 let recent: Vec<_> = messages.iter().rev().take(last_n).rev().cloned().collect();
1679
1680 if recent.is_empty() {
1681 return Ok(vec![]);
1682 }
1683
1684 let actor_id = self.effective_actor_id();
1685 let existing = actor_id
1686 .as_ref()
1687 .and_then(|aid| self.actor_facts_cache.read().get(aid).cloned())
1688 .unwrap_or_default();
1689
1690 let categories = self
1691 .facts_config
1692 .as_ref()
1693 .map(|c| c.custom_categories.clone())
1694 .unwrap_or_default();
1695
1696 let facts = self
1697 .observe_purpose(
1698 ObservationPurpose::FactsExtraction,
1699 extractor.extract(&recent, &existing, actor_id.as_deref(), &categories),
1700 )
1701 .await?;
1702
1703 if !facts.is_empty() {
1705 let fact_store_opt = self.fact_store.read().clone();
1706 let mut stored_total = 0usize;
1707 let mut cache_updated = false;
1708 if let (Some(store), Some(aid)) = (fact_store_opt, &actor_id) {
1709 let authoritative = store.add_facts(aid, facts.clone()).await?;
1711 stored_total = authoritative.len();
1712 self.actor_facts_cache
1713 .write()
1714 .insert(aid.clone(), authoritative);
1715 cache_updated = true;
1716 } else if let Some(aid) = &actor_id {
1717 let mut cache = self.actor_facts_cache.write();
1718 let entry = cache.entry(aid.clone()).or_default();
1719 entry.extend(facts.clone());
1720 stored_total = entry.len();
1721 cache_updated = true;
1722 }
1723
1724 info!(
1725 actor_id = %actor_id.as_deref().unwrap_or("<none>"),
1726 source = source,
1727 requested_messages = last_n,
1728 message_count = recent.len(),
1729 extracted_count = facts.len(),
1730 cache_updated = cache_updated,
1731 stored_total = stored_total,
1732 "facts extracted"
1733 );
1734
1735 if let Some(ref aid) = actor_id {
1736 self.hooks.on_facts_extracted(aid, &facts).await;
1737 }
1738 }
1739
1740 Ok(facts)
1741 }
1742
1743 fn resolve_actor_id_from_context(&self) {
1746 if self
1747 .current_turn_actor_context()
1748 .and_then(|ctx| ctx.effective_actor_id().map(str::to_string))
1749 .is_some()
1750 {
1751 return;
1752 }
1753
1754 if let Some(ref am_config) = self.actor_memory_config
1755 && am_config.identification.method == ai_agents_facts::IdentificationMethod::FromContext
1756 && let Some(ref path) = am_config.identification.context_path
1757 {
1758 let val = self
1760 .context_manager
1761 .get_path(path)
1762 .or_else(|| self.context_manager.get(path));
1763 if let Some(val) = val
1764 && let Some(id_str) = val.as_str()
1765 {
1766 let current = self.actor_id.read().clone();
1767 if current.as_deref() != Some(id_str) {
1768 *self.actor_id.write() = Some(id_str.to_string());
1769 let mut meta = self.session_metadata.write();
1770 meta.actor_id = Some(id_str.to_string());
1771 if !meta.actors.iter().any(|a| a == id_str) {
1772 meta.actors.push(id_str.to_string());
1773 }
1774 }
1775 }
1776 }
1777 }
1778
1779 fn format_actor_facts_for_context(&self) -> String {
1781 let should_inject = self
1783 .facts_config
1784 .as_ref()
1785 .map(|c| c.inject_in_context)
1786 .unwrap_or(true);
1787 if !should_inject {
1788 return String::new();
1789 }
1790
1791 let Some(actor_id) = self.effective_actor_id() else {
1792 return String::new();
1793 };
1794
1795 let facts = self
1796 .actor_facts_cache
1797 .read()
1798 .get(&actor_id)
1799 .cloned()
1800 .unwrap_or_default();
1801 if facts.is_empty() {
1802 return String::new();
1803 }
1804
1805 let am_config = self.actor_memory_config.as_ref();
1806 let facts_budget = self
1809 .memory_token_budget
1810 .as_ref()
1811 .map(|b| b.allocation.facts as usize)
1812 .filter(|n| *n > 0);
1813 let default_max = am_config.map(|c| c.injection.max_tokens).unwrap_or(800);
1814 let max_tokens = facts_budget.unwrap_or(default_max);
1815
1816 let filtered: Vec<ai_agents_core::KeyFact> = if let Some(cfg) = am_config {
1818 if cfg.injection.mode == ai_agents_facts::InjectionMode::OnDemand {
1819 return String::new();
1820 }
1821 if cfg.injection.mode == ai_agents_facts::InjectionMode::Category
1822 && !cfg.injection.categories.is_empty()
1823 {
1824 facts
1825 .iter()
1826 .filter(|f| {
1827 cfg.injection
1828 .categories
1829 .iter()
1830 .any(|c| f.category.to_string() == *c)
1831 })
1832 .cloned()
1833 .collect()
1834 } else {
1835 facts.clone()
1836 }
1837 } else {
1838 facts.clone()
1839 };
1840
1841 if filtered.is_empty() {
1842 return String::new();
1843 }
1844
1845 if let Some(store) = self.fact_store.read().clone() {
1846 store.format_for_context(&filtered, max_tokens)
1847 } else {
1848 String::new()
1849 }
1850 }
1851
1852 fn build_context_with_staged(&self, staged: &HashMap<String, Value>) -> HashMap<String, Value> {
1853 let context = self.build_context_with_overlays();
1854 let mut root = Value::Object(context.into_iter().collect());
1855 for (path, value) in staged {
1856 if let Ok(updated) = ai_agents_core::set_dot_path(root.clone(), path, value.clone()) {
1857 root = updated;
1858 }
1859 }
1860 match root {
1861 Value::Object(obj) => obj.into_iter().collect(),
1862 _ => HashMap::new(),
1863 }
1864 }
1865
1866 fn build_context_with_overlays(&self) -> HashMap<String, Value> {
1867 let mut context = self.context_manager.get_all();
1868 let mut root = Value::Object(context.clone().into_iter().collect());
1869
1870 if let Some(turn_ctx) = self.current_turn_actor_context() {
1871 if let Some(ref origin_actor_id) = turn_ctx.origin_actor_id
1872 && let Ok(updated) = ai_agents_core::set_dot_path(
1873 root.clone(),
1874 "interaction.origin_actor_id",
1875 serde_json::json!(origin_actor_id),
1876 )
1877 {
1878 root = updated;
1879 }
1880 if let Some(ref sender_agent_id) = turn_ctx.sender_agent_id
1881 && let Ok(updated) = ai_agents_core::set_dot_path(
1882 root.clone(),
1883 "interaction.sender_agent_id",
1884 serde_json::json!(sender_agent_id),
1885 )
1886 {
1887 root = updated;
1888 }
1889 }
1890
1891 if let Some(ref actor_id) = self.effective_actor_id()
1892 && let Ok(updated) = ai_agents_core::set_dot_path(
1893 root.clone(),
1894 "interaction.actor_id",
1895 serde_json::json!(actor_id),
1896 )
1897 {
1898 root = updated;
1899 }
1900
1901 if let Some(manager) = self.relationship_manager.as_ref()
1902 && let Some(actor_id) = self.effective_actor_id()
1903 && let Some(value) = manager.to_context_value(&actor_id)
1904 && let Ok(updated) = ai_agents_core::set_dot_path(
1905 root.clone(),
1906 &manager.config().injection.context_path,
1907 value,
1908 )
1909 {
1910 root = updated;
1911 }
1912
1913 if let Value::Object(obj) = root {
1914 context = obj.into_iter().collect();
1915 }
1916
1917 context
1918 }
1919
1920 fn resolve_actor_name_from_context(&self) -> Option<String> {
1921 for path in ["actor.name", "user.name", "player.name", "customer.name"] {
1922 if let Some(value) = self.context_manager.get_path(path)
1923 && let Some(name) = value.as_str()
1924 {
1925 return Some(name.to_string());
1926 }
1927 }
1928 None
1929 }
1930
1931 async fn maybe_load_actor_relationship(&self) {
1932 let Some(manager) = self.relationship_manager.as_ref() else {
1933 return;
1934 };
1935 let Some(actor_id) = self.effective_actor_id() else {
1936 return;
1937 };
1938
1939 let mut should_fire_loaded = false;
1940 if manager.get(&actor_id).is_none() {
1941 let mut loaded = false;
1942 if manager.config().persistence.enabled {
1943 let storage = self.storage.read().clone();
1944 if let Some(storage) = storage {
1945 match storage.load_relationship(&self.info.id, &actor_id).await {
1946 Ok(Some(value)) => match manager.insert_from_value(value) {
1947 Ok(_) => loaded = true,
1948 Err(e) => {
1949 warn!(actor = %actor_id, error = %e, "failed to restore relationship")
1950 }
1951 },
1952 Ok(None) => {}
1953 Err(e) => {
1954 warn!(actor = %actor_id, error = %e, "failed to load relationship")
1955 }
1956 }
1957 }
1958 }
1959
1960 if !loaded {
1961 manager.get_or_create(&actor_id, self.resolve_actor_name_from_context().as_deref());
1962 }
1963 should_fire_loaded = true;
1964 }
1965
1966 let actor_name = self.resolve_actor_name_from_context();
1967 let relationship = manager.touch_interaction(&actor_id, actor_name.as_deref());
1968 if should_fire_loaded {
1969 self.hooks
1970 .on_relationship_loaded(&actor_id, &relationship)
1971 .await;
1972 }
1973 }
1974
1975 fn format_relationship_for_context(&self) -> Option<(String, String)> {
1976 let manager = self.relationship_manager.as_ref()?;
1977 if !manager.config().injection.enabled {
1978 return None;
1979 }
1980 let actor_id = self.effective_actor_id()?;
1981 let relationship = manager.get(&actor_id)?;
1982 let local_cap = manager.config().injection.max_tokens;
1983 let global_cap = self
1984 .memory_token_budget
1985 .as_ref()
1986 .map(|b| b.allocation.relationships as usize)
1987 .filter(|n| *n > 0);
1988 let max_tokens = global_cap.map(|g| g.min(local_cap)).unwrap_or(local_cap);
1989 let text = ai_agents_relationships::format_relationship(
1990 &relationship,
1991 &manager.config().injection.format,
1992 max_tokens,
1993 );
1994 if text.is_empty() {
1995 None
1996 } else {
1997 Some((manager.config().injection.prompt_variable.clone(), text))
1998 }
1999 }
2000
2001 async fn persist_actor_relationship(&self, actor_id: &str) -> Result<()> {
2002 let Some(manager) = self.relationship_manager.as_ref() else {
2003 return Ok(());
2004 };
2005 if !manager.config().persistence.enabled {
2006 return Ok(());
2007 }
2008 let storage = self.storage.read().clone();
2009 let Some(storage) = storage else {
2010 return Ok(());
2011 };
2012 if let Some(value) = manager.relationship_as_value(actor_id)? {
2013 storage
2014 .save_relationship(&self.info.id, actor_id, &value)
2015 .await?;
2016 }
2017 Ok(())
2018 }
2019
2020 async fn auto_update_relationship(&self) {
2021 let Some(manager) = self.relationship_manager.as_ref() else {
2022 return;
2023 };
2024 let Some(actor_id) = self.effective_actor_id() else {
2025 return;
2026 };
2027 if !manager.config().auto_update.enabled {
2028 let _ = self.persist_actor_relationship(&actor_id).await;
2029 return;
2030 }
2031
2032 let recent_messages = manager.config().auto_update.recent_messages;
2033 let messages = match self.memory.get_messages(Some(recent_messages)).await {
2034 Ok(messages) => messages,
2035 Err(e) => {
2036 warn!(actor = %actor_id, error = %e, "failed to read messages for relationship update");
2037 return;
2038 }
2039 };
2040
2041 match self
2042 .observe_purpose(
2043 ObservationPurpose::RelationshipUpdate,
2044 manager.auto_update(&actor_id, &messages),
2045 )
2046 .await
2047 {
2048 Ok(update) => {
2049 if !update.changes.is_empty() {
2050 self.hooks
2051 .on_relationship_change(&actor_id, &update.changes)
2052 .await;
2053 }
2054 if let Some(ref event) = update.event {
2055 self.hooks.on_notable_event(&actor_id, event).await;
2056 }
2057 let persisted = match self.persist_actor_relationship(&actor_id).await {
2058 Ok(()) => true,
2059 Err(e) => {
2060 warn!(actor = %actor_id, error = %e, "failed to persist relationship");
2061 false
2062 }
2063 };
2064 if !update.changes.is_empty() || update.event.is_some() {
2065 let changed_dimensions: Vec<String> = update
2066 .changes
2067 .iter()
2068 .map(|change| format!("{}:{}", change.perspective, change.dimension))
2069 .collect();
2070 info!(
2071 actor_id = %actor_id,
2072 change_count = update.changes.len(),
2073 changed_dimensions = ?changed_dimensions,
2074 event_present = update.event.is_some(),
2075 persisted = persisted,
2076 "relationship updated"
2077 );
2078 } else {
2079 debug!(actor_id = %actor_id, persisted = persisted, "relationship evaluation ran but found no changes");
2080 }
2081 }
2082 Err(e) => warn!(actor = %actor_id, error = %e, "relationship update failed"),
2083 }
2084 }
2085
2086 async fn auto_extract_facts(&self) {
2088 let should_extract = self
2089 .facts_config
2090 .as_ref()
2091 .map(|c| c.enabled && c.auto_extract)
2092 .unwrap_or(false);
2093
2094 if !should_extract {
2095 debug!("fact extraction skipped because auto extraction is disabled");
2096 return;
2097 }
2098
2099 let msgs_since = *self.messages_since_extraction.read();
2100 if msgs_since < 2 {
2101 debug!(
2102 messages_since_extraction = msgs_since,
2103 "fact extraction skipped until threshold is reached"
2104 );
2105 return;
2106 }
2107
2108 match self.extract_facts_with_source(msgs_since, "auto").await {
2109 Ok(facts) => {
2110 if !facts.is_empty() {
2111 *self.messages_since_extraction.write() = 0;
2112 } else {
2113 debug!("fact extraction ran but found no new facts");
2114 }
2115 }
2116 Err(e) => {
2117 warn!("fact extraction failed: {}", e);
2118 }
2119 }
2120 }
2121
2122 pub fn with_persona(mut self, manager: Arc<ai_agents_persona::PersonaManager>) -> Self {
2123 self.persona_manager = Some(manager);
2124 self
2125 }
2126
2127 pub fn persona_manager(&self) -> Option<&Arc<ai_agents_persona::PersonaManager>> {
2128 self.persona_manager.as_ref()
2129 }
2130
2131 pub fn with_disambiguation(mut self, config: DisambiguationConfig) -> Self {
2132 if config.is_enabled() {
2133 let manager = DisambiguationManager::new(config, Arc::clone(&self.llm_registry))
2134 .with_clarification_observer(Arc::new(ObservabilityClarificationObserver));
2135 self.disambiguation_manager = Some(manager);
2136 }
2137 self
2138 }
2139
2140 pub fn disambiguation_manager(&self) -> Option<&DisambiguationManager> {
2141 self.disambiguation_manager.as_ref()
2142 }
2143
2144 pub fn has_disambiguation(&self) -> bool {
2145 self.disambiguation_manager
2146 .as_ref()
2147 .is_some_and(|m| m.is_enabled())
2148 }
2149
2150 pub async fn init_storage(&self) -> Result<()> {
2151 let _guard = self.storage_init.lock().await;
2155 let mut storage = self.storage.read().clone();
2156 if storage.is_none() && !self.storage_config.is_none() {
2157 let storage_config = self.convert_storage_config();
2158 storage = create_storage(&storage_config).await?;
2159 *self.storage.write() = storage.clone();
2160 }
2161
2162 self.validate_storage_requirements(storage.as_deref())?;
2163 self.complete_facts_init().await;
2164 Ok(())
2165 }
2166
2167 fn validate_storage_requirements(&self, storage: Option<&dyn AgentStorage>) -> Result<()> {
2168 let facts_required = self
2169 .facts_config
2170 .as_ref()
2171 .is_some_and(|config| config.enabled)
2172 || self
2173 .actor_memory_config
2174 .as_ref()
2175 .is_some_and(|config| config.enabled);
2176 let relationships_required = self
2177 .relationship_manager
2178 .as_ref()
2179 .is_some_and(|manager| manager.config().persistence.enabled);
2180
2181 let Some(storage) = storage else {
2182 let mut requirements = Vec::new();
2183 if facts_required {
2184 requirements.push("actor facts or actor memory");
2185 }
2186 if relationships_required {
2187 requirements.push("persistent relationships");
2188 }
2189 if requirements.is_empty() {
2190 return Ok(());
2191 }
2192 return Err(AgentError::Config(format!(
2193 "Storage is required for enabled {} but none is configured or injected",
2194 requirements.join(" and ")
2195 )));
2196 };
2197
2198 if facts_required && !storage.supports(StorageCapability::ActorFacts) {
2202 return Err(AgentError::UnsupportedStorageCapability(
2203 StorageCapability::ActorFacts,
2204 ));
2205 }
2206 if relationships_required && !storage.supports(StorageCapability::ActorRelationships) {
2207 return Err(AgentError::UnsupportedStorageCapability(
2208 StorageCapability::ActorRelationships,
2209 ));
2210 }
2211 Ok(())
2212 }
2213
2214 async fn complete_facts_init(&self) {
2217 if self.fact_store.read().is_some() {
2218 return;
2219 }
2220 let storage = match self.storage.read().clone() {
2221 Some(s) => s,
2222 None => return,
2223 };
2224
2225 let facts_enabled = self
2226 .facts_config
2227 .as_ref()
2228 .map(|f| f.enabled)
2229 .unwrap_or(false);
2230 let actor_memory_enabled = self
2231 .actor_memory_config
2232 .as_ref()
2233 .map(|a| a.enabled)
2234 .unwrap_or(false);
2235
2236 if !facts_enabled && !actor_memory_enabled {
2237 return;
2238 }
2239
2240 let fc = self.facts_config.clone().unwrap_or_default();
2241 let store = Arc::new(ai_agents_facts::FactStore::new(
2242 storage,
2243 self.info.id.clone(),
2244 fc.clone(),
2245 ));
2246
2247 let extractor: Option<Arc<dyn ai_agents_facts::FactExtractor>> = if facts_enabled {
2248 let extractor_llm = fc
2249 .extractor_llm
2250 .as_ref()
2251 .and_then(|alias| self.llm_registry.get(alias).ok())
2252 .or_else(|| self.llm_registry.router().ok())
2253 .or_else(|| self.llm_registry.default().ok());
2254 extractor_llm.map(|llm| {
2255 Arc::new(ai_agents_facts::LLMFactExtractor::new(llm, fc.clone()))
2256 as Arc<dyn ai_agents_facts::FactExtractor>
2257 })
2258 } else {
2259 None
2260 };
2261
2262 *self.fact_store.write() = Some(store);
2263 *self.fact_extractor.write() = extractor;
2264 debug!(
2265 agent = %self.info.id,
2266 facts_enabled,
2267 actor_memory_enabled,
2268 "facts storage initialized"
2269 );
2270 }
2271
2272 fn convert_storage_config(&self) -> StorageStorageConfig {
2273 crate::spec::storage::to_storage_config(&self.storage_config)
2274 }
2275
2276 pub fn storage(&self) -> Option<Arc<dyn AgentStorage>> {
2277 self.storage.read().clone()
2278 }
2279
2280 pub fn storage_config(&self) -> &StorageConfig {
2281 &self.storage_config
2282 }
2283
2284 pub fn spawner(&self) -> Option<&Arc<crate::spawner::AgentSpawner>> {
2286 self.spawner.as_ref()
2287 }
2288
2289 pub fn spawner_registry(&self) -> Option<&Arc<crate::spawner::AgentRegistry>> {
2291 self.spawner_registry.as_ref()
2292 }
2293
2294 pub fn has_spawner(&self) -> bool {
2295 self.spawner_registry.is_some()
2296 }
2297
2298 pub fn with_spawner_handles(
2299 mut self,
2300 spawner: Arc<crate::spawner::AgentSpawner>,
2301 registry: Arc<crate::spawner::AgentRegistry>,
2302 ) -> Self {
2303 self.spawner = Some(spawner);
2304 self.spawner_registry = Some(registry);
2305 self
2306 }
2307
2308 pub fn with_hooks(mut self, hooks: Arc<dyn AgentHooks>) -> Self {
2309 self.hooks = hooks;
2310 self
2311 }
2312
2313 pub fn with_parallel_tools(mut self, config: ParallelToolsConfig) -> Self {
2314 self.parallel_tools = config;
2315 self
2316 }
2317
2318 pub fn with_streaming(mut self, config: StreamingConfig) -> Self {
2319 self.streaming = config;
2320 self
2321 }
2322
2323 pub fn with_hitl(mut self, engine: HITLEngine, handler: Arc<dyn ApprovalHandler>) -> Self {
2324 self.hitl_engine = Some(engine);
2325 self.approval_handler = handler;
2326 self
2327 }
2328
2329 pub fn with_max_context_tokens(mut self, tokens: u32) -> Self {
2330 self.max_context_tokens = tokens;
2331 self
2332 }
2333
2334 pub fn with_memory_token_budget(mut self, budget: MemoryTokenBudget) -> Self {
2335 self.memory_token_budget = Some(budget);
2336 self
2337 }
2338
2339 pub fn with_recovery_manager(mut self, manager: RecoveryManager) -> Self {
2340 self.recovery_manager = manager;
2341 self
2342 }
2343
2344 pub fn with_tool_security(mut self, engine: ToolSecurityEngine) -> Self {
2345 self.tool_security = engine;
2346 self
2347 }
2348
2349 pub fn runtime_control(&self) -> RuntimeControlHandle {
2351 RuntimeControlHandle {
2352 state: Arc::clone(&self.runtime_control),
2353 }
2354 }
2355
2356 pub fn set_question_handler(&self, handler: Option<Arc<dyn QuestionHandler>>) {
2358 self.tools.set_question_handler(handler);
2359 }
2360
2361 pub fn set_diagnostics_provider(&self, provider: Arc<dyn DiagnosticsProvider>) {
2363 self.tools.set_diagnostics_provider(provider);
2364 }
2365
2366 pub fn set_command_runner(&self, runner: Arc<dyn CommandRunner>) {
2368 self.tools.set_command_runner(runner);
2369 }
2370
2371 pub fn set_web_search_provider(&self, provider: Arc<dyn ai_agents_tools::WebSearchProvider>) {
2373 self.tools.set_web_search_provider(provider);
2374 }
2375
2376 pub fn todos(&self) -> Vec<TodoItem> {
2378 self.tools.todos()
2379 }
2380
2381 fn active_tool_security(&self) -> ToolSecurityEngine {
2383 self.runtime_control
2384 .tool_security_override
2385 .read()
2386 .clone()
2387 .unwrap_or_else(|| self.tool_security.clone())
2388 }
2389
2390 fn runtime_safety_snapshot(&self) -> RuntimeSafetySnapshot {
2392 let _guard = self.runtime_control.snapshot_guard.read();
2393 RuntimeSafetySnapshot {
2394 version: self.runtime_control.version.load(Ordering::SeqCst),
2395 emergency_deny: self.runtime_control.emergency_deny.load(Ordering::SeqCst),
2396 tool_security: self
2397 .runtime_control
2398 .tool_security_override
2399 .read()
2400 .clone()
2401 .unwrap_or_else(|| self.tool_security.clone()),
2402 tool_scope_override: self.runtime_control.tool_scope_override.read().clone(),
2403 }
2404 }
2405
2406 fn admit_tool_execution(
2408 &self,
2409 expected_runtime_version: u64,
2410 expected_policy_version: u64,
2411 expected_state_generation: Option<u64>,
2412 canonical_id: &str,
2413 ) -> SecurityCheckResult {
2414 let _guard = self.runtime_control.snapshot_guard.read();
2415 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
2416 return SecurityCheckResult::Block {
2417 reason: "runtime emergency deny is enabled".to_string(),
2418 };
2419 }
2420 let runtime_version = self.runtime_control.version.load(Ordering::SeqCst);
2421 let security_engine = self
2422 .runtime_control
2423 .tool_security_override
2424 .read()
2425 .clone()
2426 .unwrap_or_else(|| self.tool_security.clone());
2427 if runtime_version != expected_runtime_version
2428 || security_engine.policy_version() != expected_policy_version
2429 {
2430 return SecurityCheckResult::Block {
2431 reason: "runtime safety controls changed before admission".to_string(),
2432 };
2433 }
2434 let current_state_generation = self
2435 .state_machine
2436 .as_ref()
2437 .map(|state_machine| state_machine.generation());
2438 if current_state_generation != expected_state_generation {
2439 return SecurityCheckResult::Block {
2440 reason: "state scope changed before admission".to_string(),
2441 };
2442 }
2443 security_engine.admit_tool_execution(canonical_id)
2444 }
2445
2446 pub fn with_process_processor(mut self, processor: ProcessProcessor) -> Self {
2447 let processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
2448 self.process_processor = Some(processor);
2449 self
2450 }
2451
2452 pub fn with_state_machine(
2453 mut self,
2454 state_machine: Arc<StateMachine>,
2455 evaluator: Arc<dyn TransitionEvaluator>,
2456 ) -> Self {
2457 self.state_machine = Some(state_machine);
2458 self.transition_evaluator = Some(evaluator);
2459 self
2460 }
2461
2462 pub fn with_context_manager(mut self, manager: Arc<ContextManager>) -> Self {
2463 self.context_manager = manager;
2464 self
2465 }
2466
2467 pub fn register_message_filter(&self, name: impl Into<String>, filter: Arc<dyn MessageFilter>) {
2468 self.message_filters.write().insert(name.into(), filter);
2469 }
2470
2471 pub fn set_context(&self, key: &str, value: Value) -> Result<()> {
2472 self.context_manager.update(key, value)
2473 }
2474
2475 pub fn update_context(&self, path: &str, value: Value) -> Result<()> {
2476 self.context_manager.update(path, value)
2477 }
2478
2479 pub fn get_context(&self) -> HashMap<String, Value> {
2480 self.build_context_with_overlays()
2481 }
2482
2483 pub fn remove_context(&self, key: &str) -> Option<Value> {
2484 self.context_manager.remove(key)
2485 }
2486
2487 pub async fn refresh_context(&self, key: &str) -> Result<()> {
2488 self.context_manager.refresh(key).await
2489 }
2490
2491 pub fn register_context_provider(&self, name: &str, provider: Arc<dyn ContextProvider>) {
2492 self.context_manager.register_provider(name, provider);
2493 }
2494
2495 pub fn current_state(&self) -> Option<String> {
2496 self.state_machine.as_ref().map(|sm| sm.current())
2497 }
2498
2499 async fn invalidate_pending_confirmation(&self, reason: &'static str) {
2501 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
2502 let Some(disambiguator) = self.disambiguation_manager.as_ref() else {
2503 return;
2504 };
2505 if disambiguator.has_pending_confirmation().await {
2506 disambiguator.clear_pending().await;
2507 *self.pending_skill_id.write() = None;
2508 info!(
2509 confirmation_event = "invalidated",
2510 invalidation_reason = reason,
2511 "Runtime invalidated pending confirmation"
2512 );
2513 }
2514 }
2515
2516 async fn admit_disambiguation_redispatch(
2518 &self,
2519 expected_epoch: u64,
2520 expected_state_generation: Option<u64>,
2521 ) -> Result<tokio::sync::RwLockReadGuard<'_, ()>> {
2522 let admission = self.disambiguation_admission.read().await;
2523 let state_generation = self
2524 .state_machine
2525 .as_ref()
2526 .map(|state_machine| state_machine.generation());
2527 if self.disambiguation_epoch.load(Ordering::SeqCst) != expected_epoch
2528 || state_generation != expected_state_generation
2529 {
2530 return Err(AgentError::Other(
2531 "Disambiguation ownership changed before redispatch admission".to_string(),
2532 ));
2533 }
2534 Ok(admission)
2535 }
2536
2537 fn reserve_state_transition(&self) -> Option<StateTransitionReservation<'_>> {
2539 self.state_transition_reserved
2540 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
2541 .ok()
2542 .map(|_| StateTransitionReservation {
2543 reserved: &self.state_transition_reserved,
2544 })
2545 }
2546
2547 async fn admit_optional_disambiguation_ownership(
2549 &self,
2550 ownership: Option<DisambiguationOwnership>,
2551 ) -> Result<Option<tokio::sync::RwLockReadGuard<'_, ()>>> {
2552 match ownership {
2553 Some(ownership) => self
2554 .admit_disambiguation_redispatch(ownership.epoch, ownership.state_generation)
2555 .await
2556 .map(Some),
2557 None => Ok(None),
2558 }
2559 }
2560
2561 pub async fn transition_to(&self, state: &str) -> Result<()> {
2563 let Some(ref sm) = self.state_machine else {
2564 return Ok(());
2565 };
2566 let claim_admission = self.disambiguation_admission.write().await;
2567 let reservation = self.reserve_state_transition().ok_or_else(|| {
2568 AgentError::Other("Another state transition is already in progress".to_string())
2569 })?;
2570 let from_state = sm.current();
2571 let expected_state_generation = sm.generation();
2572 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
2573 let history_before = sm.history();
2574 drop(claim_admission);
2575
2576 self.execute_state_exit_actions(&from_state).await;
2577
2578 let admission = self.disambiguation_admission.write().await;
2579 if sm.current() != from_state
2580 || sm.generation() != expected_state_generation
2581 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
2582 {
2583 return Err(AgentError::Other(
2584 "State ownership changed during manual transition preparation".to_string(),
2585 ));
2586 }
2587 sm.transition_to(state, "manual transition")?;
2588 self.invalidate_pending_confirmation("state_transition")
2589 .await;
2590 let entered = sm.current();
2591 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
2592 drop(admission);
2593
2594 self.execute_state_enter_actions(&entered, is_reentry).await;
2595 drop(reservation);
2596 info!(to = %entered, "Manual state transition");
2597 Ok(())
2598 }
2599
2600 pub fn state_history(&self) -> Vec<StateTransitionEvent> {
2601 self.state_machine
2602 .as_ref()
2603 .map(|sm| sm.history())
2604 .unwrap_or_default()
2605 }
2606
2607 pub fn session_metadata(&self) -> ai_agents_core::SessionMetadata {
2609 self.session_metadata.read().clone()
2610 }
2611
2612 pub async fn delete_actor_data(&self, actor_id: &str) -> Result<()> {
2615 let allowed = self
2616 .actor_memory_config
2617 .as_ref()
2618 .map(|c| c.privacy.allow_deletion)
2619 .unwrap_or(true);
2620 if !allowed {
2621 return Err(AgentError::Config(
2622 "privacy.allow_deletion is false; actor data deletion is not permitted".into(),
2623 ));
2624 }
2625 let storage = self.storage.read().clone();
2626 if let Some(storage) = storage {
2627 if !storage.supports(StorageCapability::ActorDataDeletion) {
2631 return Err(AgentError::UnsupportedStorageCapability(
2632 StorageCapability::ActorDataDeletion,
2633 ));
2634 }
2635 storage.delete_actor_data(&self.info.id, actor_id).await?;
2636 } else {
2637 let store = { self.fact_store.read().clone() };
2641 if let Some(store) = store {
2642 store.delete_actor_data(actor_id).await?;
2643 }
2644 }
2645 if let Some(manager) = self.relationship_manager.as_ref() {
2646 manager.remove(actor_id);
2647 }
2648 self.actor_facts_cache.write().remove(actor_id);
2649 Ok(())
2650 }
2651
2652 pub fn set_session_metadata(&self, meta: ai_agents_core::SessionMetadata) {
2654 *self.session_metadata.write() = meta;
2655 }
2656
2657 pub async fn cleanup_expired_sessions(&self) -> Result<usize> {
2659 let storage = self.storage.read().clone();
2660 match storage {
2661 Some(s) => {
2662 let count = s.cleanup_expired().await?;
2663 if count > 0 {
2664 self.hooks.on_sessions_expired(count).await;
2665 }
2666 Ok(count)
2667 }
2668 None => Err(AgentError::Config(
2669 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2670 )),
2671 }
2672 }
2673
2674 pub async fn list_sessions_filtered(
2676 &self,
2677 filter: &ai_agents_core::SessionFilter,
2678 ) -> Result<Vec<ai_agents_core::SessionSummary>> {
2679 let storage = self.storage.read().clone();
2680 match storage {
2681 Some(s) => s.list_sessions_filtered(filter).await,
2682 None => Err(AgentError::Config(
2683 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2684 )),
2685 }
2686 }
2687
2688 pub async fn save_state(&self) -> Result<AgentSnapshot> {
2689 let memory_snapshot = self.memory.snapshot().await?;
2690 let state_machine_snapshot = self.state_machine.as_ref().map(|sm| sm.snapshot());
2691 let context_snapshot = self.context_manager.snapshot();
2692
2693 let mut snapshot = AgentSnapshot::new(self.info.id.clone())
2694 .with_memory(memory_snapshot)
2695 .with_context(context_snapshot)
2696 .with_state_machine(
2697 state_machine_snapshot.unwrap_or_else(|| StateMachineSnapshot {
2698 current_state: String::new(),
2699 previous_state: None,
2700 turn_count: 0,
2701 no_transition_count: 0,
2702 history: vec![],
2703 }),
2704 );
2705
2706 if let Some(ref persona) = self.persona_manager {
2707 snapshot.persona = Some(persona.snapshot_as_value()?);
2708 }
2709
2710 if let Some(ref relationships) = self.relationship_manager {
2711 snapshot.relationships = Some(relationships.snapshot_as_value()?);
2712 }
2713
2714 Ok(snapshot)
2715 }
2716
2717 pub async fn save_state_full(&self) -> Result<AgentSnapshot> {
2719 let mut snapshot = self.save_state().await?;
2720 if let Some(ref registry) = self.spawner_registry {
2721 let entries = registry.list_with_specs();
2722 if !entries.is_empty() {
2723 snapshot = snapshot.with_spawned_agents(entries);
2724 }
2725 }
2726 Ok(snapshot)
2727 }
2728
2729 pub async fn restore_state(&self, snapshot: AgentSnapshot) -> Result<()> {
2731 let _admission = self.disambiguation_admission.write().await;
2732 if self.state_transition_reserved.load(Ordering::SeqCst) {
2733 return Err(AgentError::Other(
2734 "Cannot restore state while a state transition is in progress".to_string(),
2735 ));
2736 }
2737 self.invalidate_pending_confirmation("state_restore").await;
2738 *self.pending_skill_id.write() = None;
2739 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
2740 disambiguator.clear_pending().await;
2741 }
2742 self.memory.restore(snapshot.memory).await?;
2743
2744 if let (Some(sm), Some(sm_snapshot)) = (&self.state_machine, snapshot.state_machine)
2745 && !sm_snapshot.current_state.is_empty()
2746 {
2747 sm.restore(sm_snapshot)?;
2748 }
2749
2750 self.context_manager.restore(snapshot.context);
2751
2752 if let (Some(persona_value), Some(persona_manager)) =
2753 (snapshot.persona, &self.persona_manager)
2754 {
2755 persona_manager.restore_from_value(persona_value)?;
2756 }
2757
2758 if let (Some(relationship_value), Some(relationship_manager)) =
2759 (snapshot.relationships, &self.relationship_manager)
2760 {
2761 relationship_manager.restore_from_value(relationship_value)?;
2762 }
2763
2764 info!(agent_id = %snapshot.agent_id, "State restored");
2765 Ok(())
2766 }
2767
2768 pub async fn save_to(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<()> {
2769 let snapshot = self.save_state().await?;
2770 storage.save(session_id, &snapshot).await
2771 }
2772
2773 async fn load_session_restore(
2774 storage: &dyn AgentStorage,
2775 session_id: &str,
2776 ) -> Result<Option<StoredSessionRestore>> {
2777 let Some(snapshot) = storage.load(session_id).await? else {
2778 return Ok(None);
2779 };
2780 let metadata = if storage.supports(StorageCapability::SessionMetadata) {
2784 storage.load_metadata(session_id).await?
2785 } else {
2786 None
2787 };
2788 Ok(Some(StoredSessionRestore { snapshot, metadata }))
2789 }
2790
2791 async fn capture_session_restore_point(&self) -> Result<RuntimeSessionRestorePoint> {
2792 Ok(RuntimeSessionRestorePoint {
2793 snapshot: self.save_state().await?,
2794 metadata: self.session_metadata(),
2795 actor_id: self.actor_id(),
2796 session_id: self.current_session_id.read().clone(),
2797 })
2798 }
2799
2800 async fn apply_session_restore_unchecked(
2801 &self,
2802 session_id: &str,
2803 stored: StoredSessionRestore,
2804 ) -> Result<()> {
2805 self.restore_state(stored.snapshot).await?;
2806 let metadata = stored.metadata.unwrap_or_default();
2807 if let Some(actor_id) = metadata.actor_id.as_deref() {
2808 self.set_actor_id(actor_id)?;
2809 } else {
2810 self.clear_actor_id();
2811 }
2812 self.set_session_metadata(metadata);
2813 *self.current_session_id.write() = Some(session_id.to_string());
2814 Ok(())
2815 }
2816
2817 async fn restore_session_restore_point(
2818 &self,
2819 restore_point: &RuntimeSessionRestorePoint,
2820 ) -> Result<()> {
2821 self.restore_state(restore_point.snapshot.clone()).await?;
2822 if let Some(actor_id) = restore_point.actor_id.as_deref() {
2823 self.set_actor_id(actor_id)?;
2824 } else {
2825 self.clear_actor_id();
2826 }
2827 self.set_session_metadata(restore_point.metadata.clone());
2828 *self.current_session_id.write() = restore_point.session_id.clone();
2829 Ok(())
2830 }
2831
2832 async fn apply_session_restore(
2833 &self,
2834 session_id: &str,
2835 stored: StoredSessionRestore,
2836 ) -> Result<()> {
2837 let before = self.capture_session_restore_point().await?;
2838 if let Err(error) = self
2839 .apply_session_restore_unchecked(session_id, stored)
2840 .await
2841 {
2842 return match self.restore_session_restore_point(&before).await {
2843 Ok(()) => Err(error),
2844 Err(rollback_error) => Err(AgentError::Other(format!(
2845 "Session restore failed: {error}; rollback failed: {rollback_error}"
2846 ))),
2847 };
2848 }
2849 Ok(())
2850 }
2851
2852 async fn rollback_session_restore_set(
2853 parent: Option<(&RuntimeAgent, &RuntimeSessionRestorePoint)>,
2854 children: &[(String, Arc<RuntimeAgent>, RuntimeSessionRestorePoint)],
2855 ) -> Vec<String> {
2856 let mut errors = Vec::new();
2857 if let Some((agent, restore_point)) = parent
2858 && let Err(error) = agent.restore_session_restore_point(restore_point).await
2859 {
2860 errors.push(format!("parent: {error}"));
2861 }
2862 for (id, agent, restore_point) in children {
2863 if let Err(error) = agent.restore_session_restore_point(restore_point).await {
2864 errors.push(format!("child '{id}': {error}"));
2865 }
2866 }
2867 errors
2868 }
2869
2870 fn restore_failure(error: impl std::fmt::Display, rollback_errors: Vec<String>) -> AgentError {
2871 if rollback_errors.is_empty() {
2872 AgentError::Other(format!(
2873 "Session restore failed: {error}; runtime state was rolled back"
2874 ))
2875 } else {
2876 AgentError::Other(format!(
2877 "Session restore failed: {error}; rollback also failed for {}",
2878 rollback_errors.join(", ")
2879 ))
2880 }
2881 }
2882
2883 pub async fn load_from(&self, storage: &dyn AgentStorage, session_id: &str) -> Result<bool> {
2884 let Some(stored) = Self::load_session_restore(storage, session_id).await? else {
2885 return Ok(false);
2886 };
2887 self.apply_session_restore(session_id, stored).await?;
2888 Ok(true)
2889 }
2890
2891 pub async fn save_session(&self, session_id: &str) -> Result<()> {
2892 let storage = self.storage.read().clone();
2893 match storage {
2894 Some(s) => {
2895 let is_new = {
2897 let cur = self.current_session_id.read().clone();
2898 cur.as_deref() != Some(session_id)
2899 };
2900 if is_new {
2901 *self.current_session_id.write() = Some(session_id.to_string());
2902 self.hooks.on_session_created(session_id).await;
2903 }
2904
2905 {
2907 let now = chrono::Utc::now();
2908 let msg_count = self
2909 .memory
2910 .get_messages(None)
2911 .await
2912 .map(|v| v.len())
2913 .unwrap_or(0);
2914 let mut meta = self.session_metadata.write();
2915 meta.last_active = now;
2916 meta.message_count = msg_count;
2917 if meta.actor_id.is_none() {
2918 meta.actor_id = self.actor_id.read().clone();
2919 }
2920 }
2921
2922 let snapshot = self.save_state().await?;
2923 if s.supports(StorageCapability::SessionMetadata) {
2927 let metadata = self.session_metadata.read().clone();
2928 s.save_snapshot_with_metadata(session_id, &snapshot, &metadata)
2929 .await
2930 } else {
2931 s.save(session_id, &snapshot).await
2932 }
2933 }
2934 None => Err(AgentError::Config(
2935 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2936 )),
2937 }
2938 }
2939
2940 pub async fn load_session(&self, session_id: &str) -> Result<bool> {
2941 let storage = self.storage.read().clone();
2942 match storage {
2943 Some(storage) => self.load_from(storage.as_ref(), session_id).await,
2944 None => Err(AgentError::Config(
2945 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2946 )),
2947 }
2948 }
2949
2950 pub async fn restore_session_full(&self, session_id: &str) -> Result<usize> {
2952 self.init_storage().await?;
2953 let storage = self.storage.read().clone().ok_or_else(|| {
2954 AgentError::Config(
2955 "No storage configured. Use with_storage_config() or with_storage() first".into(),
2956 )
2957 })?;
2958 let target_parent = Self::load_session_restore(storage.as_ref(), session_id)
2959 .await?
2960 .ok_or_else(|| AgentError::Persistence(format!("Session not found: {session_id}")))?;
2961 let manifest = target_parent
2962 .snapshot
2963 .spawned_agents
2964 .clone()
2965 .unwrap_or_default();
2966
2967 let registry = self.spawner_registry.as_ref().cloned();
2968 let spawner = if manifest.is_empty() {
2969 self.spawner.as_ref().cloned()
2970 } else {
2971 Some(self.spawner.as_ref().cloned().ok_or_else(|| {
2972 AgentError::Config(
2973 "Saved session contains child agents but this runtime has no spawner".into(),
2974 )
2975 })?)
2976 };
2977 let registry = if manifest.is_empty() {
2978 registry
2979 } else {
2980 Some(registry.ok_or_else(|| {
2981 AgentError::Config(
2982 "Saved session contains child agents but this runtime has no registry".into(),
2983 )
2984 })?)
2985 };
2986
2987 let mut target_ids = HashSet::with_capacity(manifest.len());
2988 let mut prepared = Vec::with_capacity(manifest.len());
2989 for entry in manifest {
2990 if !target_ids.insert(entry.id.clone()) {
2991 return Err(AgentError::InvalidSpec(format!(
2992 "Saved child manifest contains duplicate ID: {}",
2993 entry.id
2994 )));
2995 }
2996 let spec = crate::spec::AgentSpec::from_yaml_strict(&entry.spec_yaml)?;
2997 spawner
2998 .as_ref()
2999 .expect("non-empty manifests require a spawner")
3000 .validate_explicit_child(&entry.id, &spec)?;
3001 prepared.push((entry.id, spec));
3002 }
3003
3004 let current_ids = registry
3005 .as_ref()
3006 .map(|registry| {
3007 registry
3008 .list()
3009 .into_iter()
3010 .map(|info| info.id)
3011 .collect::<HashSet<_>>()
3012 })
3013 .unwrap_or_default();
3014 let removal_count = current_ids.difference(&target_ids).count();
3015 let additions = prepared
3016 .iter()
3017 .filter(|(id, _)| !current_ids.contains(id))
3018 .cloned()
3019 .collect::<Vec<_>>();
3020
3021 let mut existing = Vec::new();
3022 if let Some(registry) = registry.as_ref() {
3023 for (id, _) in prepared.iter().filter(|(id, _)| current_ids.contains(id)) {
3024 let agent = registry.get(id).ok_or_else(|| {
3025 AgentError::Config(format!("Retained child disappeared during restore: {id}"))
3026 })?;
3027 let child_storage = agent.storage().ok_or_else(|| {
3028 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3029 })?;
3030 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3031 .await?
3032 .ok_or_else(|| {
3033 AgentError::Persistence(format!(
3034 "Child '{id}' has no saved session '{session_id}'"
3035 ))
3036 })?;
3037 existing.push((id.clone(), agent, stored));
3038 }
3039 }
3040
3041 let mut staged = Vec::with_capacity(additions.len());
3042 if !additions.is_empty() {
3043 let spawner = spawner
3044 .as_ref()
3045 .expect("restored additions require a spawner");
3046 let reservations = spawner.reserve_restore_capacity(additions.len(), removal_count)?;
3047 for ((id, spec), reservation) in additions.into_iter().zip(reservations) {
3048 let spawned = spawner
3049 .spawn_with_reserved_capacity(id.clone(), spec, reservation)
3050 .await?;
3051 let child_storage = spawned.agent.storage().ok_or_else(|| {
3052 AgentError::Config(format!("Child '{id}' has no storage for session restore"))
3053 })?;
3054 let stored = Self::load_session_restore(child_storage.as_ref(), session_id)
3055 .await?
3056 .ok_or_else(|| {
3057 AgentError::Persistence(format!(
3058 "Child '{id}' has no saved session '{session_id}'"
3059 ))
3060 })?;
3061 staged.push((spawned, stored));
3062 }
3063 } else if let Some(spawner) = spawner.as_ref() {
3064 spawner.reserve_restore_capacity(0, removal_count)?;
3065 }
3066
3067 let parent_before = self.capture_session_restore_point().await?;
3068 let mut existing_before = Vec::with_capacity(existing.len());
3069 for (id, agent, _) in &existing {
3070 existing_before.push((
3071 id.clone(),
3072 Arc::clone(agent),
3073 agent.capture_session_restore_point().await?,
3074 ));
3075 }
3076
3077 for (_, agent, stored) in &existing {
3081 if let Err(error) = agent
3082 .apply_session_restore_unchecked(session_id, stored.clone())
3083 .await
3084 {
3085 drop(staged);
3086 let rollback_errors =
3087 Self::rollback_session_restore_set(None, &existing_before).await;
3088 return Err(Self::restore_failure(error, rollback_errors));
3089 }
3090 }
3091 for (spawned, stored) in &staged {
3092 if let Err(error) = spawned
3093 .agent
3094 .apply_session_restore_unchecked(session_id, stored.clone())
3095 .await
3096 {
3097 drop(staged);
3098 let rollback_errors =
3099 Self::rollback_session_restore_set(None, &existing_before).await;
3100 return Err(Self::restore_failure(error, rollback_errors));
3101 }
3102 }
3103 if let Err(error) = self
3104 .apply_session_restore_unchecked(session_id, target_parent)
3105 .await
3106 {
3107 drop(staged);
3108 let rollback_errors =
3109 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3110 .await;
3111 return Err(Self::restore_failure(error, rollback_errors));
3112 }
3113
3114 if let Some(registry) = registry.as_ref()
3115 && let Err(error) = registry
3116 .reconcile(
3117 &target_ids,
3118 staged.into_iter().map(|(spawned, _)| spawned).collect(),
3119 )
3120 .await
3121 {
3122 let rollback_errors =
3123 Self::rollback_session_restore_set(Some((self, &parent_before)), &existing_before)
3124 .await;
3125 return Err(Self::restore_failure(error, rollback_errors));
3126 }
3127
3128 Ok(target_ids.len())
3129 }
3130
3131 pub async fn delete_session(&self, session_id: &str) -> Result<()> {
3132 let storage = self.storage.read().clone();
3133 match storage {
3134 Some(s) => s.delete(session_id).await,
3135 None => Err(AgentError::Config(
3136 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3137 )),
3138 }
3139 }
3140
3141 pub async fn list_sessions(&self) -> Result<Vec<String>> {
3142 let storage = self.storage.read().clone();
3143 match storage {
3144 Some(s) => s.list_sessions().await,
3145 None => Err(AgentError::Config(
3146 "No storage configured. Use with_storage_config() or with_storage() first".into(),
3147 )),
3148 }
3149 }
3150
3151 fn estimate_tokens(&self, text: &str) -> u32 {
3152 (text.len() as f32 / 4.0).ceil() as u32
3153 }
3154
3155 fn estimate_total_tokens(&self, messages: &[ChatMessage]) -> u32 {
3156 messages
3157 .iter()
3158 .map(|m| self.estimate_tokens(&m.content))
3159 .sum()
3160 }
3161
3162 fn truncate_context(&self, messages: &mut Vec<ChatMessage>, keep_recent: usize) {
3163 if messages.len() <= keep_recent + 1 {
3164 return;
3165 }
3166 let system_msg = messages.remove(0);
3167 let to_remove = messages.len().saturating_sub(keep_recent);
3168 messages.drain(..to_remove);
3169 messages.insert(0, system_msg);
3170 }
3171
3172 fn get_filter(&self, config: &FilterConfig) -> Arc<dyn MessageFilter> {
3173 match config {
3174 FilterConfig::KeepRecent(n) => Arc::new(KeepRecentFilter::new(*n)),
3175 FilterConfig::ByRole { keep_roles } => Arc::new(ByRoleFilter::new(keep_roles.clone())),
3176 FilterConfig::SkipPattern { skip_if_contains } => {
3177 Arc::new(SkipPatternFilter::new(skip_if_contains.clone()))
3178 }
3179 FilterConfig::Custom { name } => {
3180 let filters = self.message_filters.read();
3181 filters
3182 .get(name)
3183 .cloned()
3184 .unwrap_or_else(|| Arc::new(KeepRecentFilter::new(10)))
3185 }
3186 }
3187 }
3188
3189 async fn summarize_context(
3190 &self,
3191 messages: &mut Vec<ChatMessage>,
3192 summarizer_llm: Option<&str>,
3193 max_summary_tokens: u32,
3194 custom_prompt: Option<&str>,
3195 keep_recent: usize,
3196 filter: Option<&FilterConfig>,
3197 ) -> Result<()> {
3198 let system_msg = messages.remove(0);
3199
3200 let to_summarize_count = messages.len().saturating_sub(keep_recent);
3201 if to_summarize_count == 0 {
3202 messages.insert(0, system_msg);
3203 return Ok(());
3204 }
3205
3206 let recent_msgs: Vec<ChatMessage> = messages.drain(to_summarize_count..).collect();
3207 let mut to_summarize = std::mem::take(messages);
3208
3209 if let Some(filter_config) = filter {
3210 let filter = self.get_filter(filter_config);
3211 to_summarize = filter.filter(to_summarize);
3212 }
3213
3214 if to_summarize.is_empty() {
3215 *messages = recent_msgs;
3216 messages.insert(0, system_msg);
3217 return Ok(());
3218 }
3219
3220 let conversation_text = to_summarize
3221 .iter()
3222 .map(|m| format!("{:?}: {}", m.role, m.content))
3223 .collect::<Vec<_>>()
3224 .join("\n");
3225
3226 let default_prompt = format!(
3227 "Summarize the following conversation in under {} tokens, preserving key information:\n\n{}",
3228 max_summary_tokens, conversation_text
3229 );
3230
3231 let summary_prompt = custom_prompt
3232 .map(|p| format!("{}\n\n{}", p, conversation_text))
3233 .unwrap_or(default_prompt);
3234
3235 let summarizer = if let Some(alias) = summarizer_llm {
3236 self.llm_registry
3237 .get(alias)
3238 .map_err(|e| AgentError::Config(e.to_string()))?
3239 } else {
3240 self.llm_registry
3241 .router()
3242 .or_else(|_| self.llm_registry.default())
3243 .map_err(|e| AgentError::Config(e.to_string()))?
3244 };
3245
3246 let summary_msgs = vec![ChatMessage::user(&summary_prompt)];
3247 let response = self
3248 .observe_purpose(
3249 ObservationPurpose::Summarization,
3250 summarizer.complete(&summary_msgs, None),
3251 )
3252 .await?;
3253
3254 let summary_message = ChatMessage::system(format!(
3255 "[Previous conversation summary]\n{}",
3256 response.content
3257 ));
3258
3259 *messages = vec![system_msg, summary_message];
3260 messages.extend(recent_msgs);
3261
3262 debug!(
3263 summarized_count = to_summarize_count,
3264 kept_recent = keep_recent,
3265 "Context summarized"
3266 );
3267
3268 Ok(())
3269 }
3270
3271 fn render_system_prompt(&self) -> Result<String> {
3272 let mut context = self.build_context_with_overlays();
3273
3274 let facts_text = self.format_actor_facts_for_context();
3276 if !facts_text.is_empty() {
3277 context.insert(
3278 "actor_facts".to_string(),
3279 serde_json::Value::String(facts_text),
3280 );
3281 }
3282
3283 if let Some((key, text)) = self.format_relationship_for_context() {
3284 context.insert(key, serde_json::Value::String(text));
3285 }
3286
3287 self.template_renderer
3288 .render(&self.base_system_prompt, &context)
3289 }
3290
3291 fn canonical_unique_tool_ids(&self, ids: &[String]) -> Vec<String> {
3293 let mut seen = HashSet::new();
3294 ids.iter()
3295 .filter_map(|id| self.tools.canonical_id(id))
3296 .filter(|canonical_id| seen.insert(canonical_id.clone()))
3297 .collect()
3298 }
3299
3300 fn get_top_level_tool_ids_for_scope(&self, scope_override: Option<&[String]>) -> Vec<String> {
3302 let Some(declared) = self.declared_tool_ids.as_deref() else {
3303 return Vec::new();
3304 };
3305 let mut effective = self.canonical_unique_tool_ids(declared);
3306 if let Some(scope) = scope_override {
3307 let scope: HashSet<String> =
3308 self.canonical_unique_tool_ids(scope).into_iter().collect();
3309 effective.retain(|canonical_id| scope.contains(canonical_id));
3310 }
3311 effective
3312 }
3313
3314 async fn get_available_tool_ids(&self) -> Result<Vec<String>> {
3316 Ok(self.get_available_tool_ids_snapshot().await?.tool_ids)
3317 }
3318
3319 async fn get_available_tool_ids_snapshot(&self) -> Result<AvailableToolIdsSnapshot> {
3321 let scope_override = self.runtime_control.tool_scope_override.read().clone();
3322 self.get_available_tool_ids_snapshot_for_scope(scope_override.as_deref())
3323 .await
3324 }
3325
3326 async fn get_available_tool_ids_snapshot_for_scope(
3328 &self,
3329 scope_override: Option<&[String]>,
3330 ) -> Result<AvailableToolIdsSnapshot> {
3331 let mut available = self.get_top_level_tool_ids_for_scope(scope_override);
3332 let (state_generation, state_scopes) = self
3333 .state_machine
3334 .as_ref()
3335 .map(|state_machine| {
3336 let (generation, scopes) = state_machine.current_tool_scope_snapshot();
3337 (Some(generation), scopes)
3338 })
3339 .unwrap_or((None, Vec::new()));
3340
3341 if available.is_empty() || state_scopes.is_empty() {
3342 return Ok(AvailableToolIdsSnapshot {
3343 tool_ids: available,
3344 state_generation,
3345 });
3346 }
3347
3348 let eval_ctx = self.build_evaluation_context().await?;
3349 let llm_getter = RegistryLLMGetter {
3350 registry: self.llm_registry.clone(),
3351 };
3352 let evaluator = ConditionEvaluator::new(llm_getter);
3353
3354 for state_scope in state_scopes {
3355 if state_scope.is_empty() {
3356 available.clear();
3357 break;
3358 }
3359
3360 let mut allowed = HashSet::new();
3361 for tool_ref in &state_scope {
3362 let tool_id = tool_ref.id();
3363 let Some(canonical_id) = self.tools.canonical_id(tool_id) else {
3364 continue;
3365 };
3366 let condition_matches = if let Some(condition) = tool_ref.condition() {
3367 match evaluator.evaluate(condition, &eval_ctx).await {
3368 Ok(matches) => matches,
3369 Err(error) => {
3370 warn!(tool = tool_id, error = %error, "Error evaluating tool condition");
3371 false
3372 }
3373 }
3374 } else {
3375 true
3376 };
3377 if condition_matches {
3378 allowed.insert(canonical_id);
3379 } else {
3380 debug!(tool = tool_id, "Tool condition not met, skipping");
3381 }
3382 }
3383 available.retain(|canonical_id| allowed.contains(canonical_id));
3384 if available.is_empty() {
3385 break;
3386 }
3387 }
3388
3389 Ok(AvailableToolIdsSnapshot {
3390 tool_ids: available,
3391 state_generation,
3392 })
3393 }
3394
3395 async fn build_evaluation_context(&self) -> Result<EvaluationContext> {
3396 let context = self.build_context_with_overlays();
3397 let messages = self.memory.get_messages(Some(10)).await?;
3398 let tool_history = self.tool_call_history.read().clone();
3399
3400 let (state_name, turn_count, previous_state) = if let Some(ref sm) = self.state_machine {
3401 (Some(sm.current()), sm.turn_count(), sm.previous())
3402 } else {
3403 (None, 0, None)
3404 };
3405
3406 Ok(EvaluationContext::default()
3407 .with_context(context)
3408 .with_state(state_name, turn_count, previous_state)
3409 .with_called_tools(tool_history)
3410 .with_messages(messages))
3411 }
3412
3413 fn record_tool_call(&self, tool_id: &str, result: Value) {
3414 self.tool_call_history.write().push(ToolCallRecord {
3415 tool_id: tool_id.to_string(),
3416 result,
3417 timestamp: chrono::Utc::now(),
3418 });
3419 }
3420
3421 async fn get_effective_system_prompt_with_persona_hooks(
3422 &self,
3423 fire_persona_hooks: bool,
3424 include_tool_prompt: bool,
3425 ) -> Result<String> {
3426 let rendered_base = self.render_system_prompt()?;
3427
3428 let persona_prefix = if let Some(ref persona) = self.persona_manager {
3429 let context = self.build_context_with_overlays();
3430 if fire_persona_hooks {
3431 let render_result = persona.render_prompt(&context)?;
3432 for content in &render_result.newly_revealed {
3433 self.hooks.on_secret_revealed(content).await;
3434 }
3435 render_result.prompt
3436 } else {
3437 persona.render_prompt_preview(&context)?
3438 }
3439 } else {
3440 String::new()
3441 };
3442
3443 if let Some(ref sm) = self.state_machine
3444 && let Some(state_def) = sm.current_definition()
3445 {
3446 let state_prompt = if let Some(ref prompt) = state_def.prompt {
3447 let context = self.build_context_with_overlays();
3448 self.template_renderer.render_with_state(
3449 prompt,
3450 &context,
3451 &sm.current(),
3452 sm.previous().as_deref(),
3453 sm.turn_count(),
3454 state_def.max_turns,
3455 )?
3456 } else {
3457 String::new()
3458 };
3459
3460 let combined = match state_def.prompt_mode {
3461 PromptMode::Append => {
3462 if state_prompt.is_empty() {
3463 rendered_base
3464 } else {
3465 format!(
3466 "{}\n\n[Current State: {}]\n{}",
3467 rendered_base,
3468 sm.current(),
3469 state_prompt
3470 )
3471 }
3472 }
3473 PromptMode::Replace => {
3474 if state_prompt.is_empty() {
3475 rendered_base
3476 } else {
3477 state_prompt
3478 }
3479 }
3480 PromptMode::Prepend => {
3481 if state_prompt.is_empty() {
3482 rendered_base
3483 } else {
3484 format!("{}\n\n{}", state_prompt, rendered_base)
3485 }
3486 }
3487 };
3488
3489 let with_persona = if persona_prefix.is_empty() {
3491 combined
3492 } else {
3493 format!("{}\n\n{}", persona_prefix, combined)
3494 };
3495
3496 if include_tool_prompt {
3497 let available_tool_ids = self.get_available_tool_ids().await?;
3498 if !available_tool_ids.is_empty() {
3499 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3500 &available_tool_ids,
3501 None,
3502 self.parallel_tools.enabled,
3503 self.runtime_config.tool_schema_prompt_mode,
3504 );
3505 if !tools_prompt.is_empty() {
3506 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3507 }
3508 }
3509 }
3510 return Ok(with_persona);
3511 }
3512
3513 let with_persona = if persona_prefix.is_empty() {
3515 rendered_base
3516 } else {
3517 format!("{}\n\n{}", persona_prefix, rendered_base)
3518 };
3519
3520 if include_tool_prompt {
3521 let available_tool_ids = self.get_available_tool_ids().await?;
3522 let tools_prompt = self.tools.generate_scoped_prompt_with_mode(
3523 &available_tool_ids,
3524 None,
3525 self.parallel_tools.enabled,
3526 self.runtime_config.tool_schema_prompt_mode,
3527 );
3528 if !tools_prompt.is_empty() {
3529 return Ok(format!("{}\n\n{}", with_persona, tools_prompt));
3530 }
3531 }
3532 Ok(with_persona)
3533 }
3534
3535 fn get_state_llm(&self) -> Result<Arc<dyn LLMProvider>> {
3536 if let Some(ref sm) = self.state_machine
3537 && let Some(state_def) = sm.current_definition()
3538 && let Some(ref llm_alias) = state_def.llm
3539 {
3540 return self
3541 .llm_registry
3542 .get(llm_alias)
3543 .map_err(|e| AgentError::Config(e.to_string()));
3544 }
3545 self.llm_registry
3546 .default()
3547 .map_err(|e| AgentError::Config(e.to_string()))
3548 }
3549
3550 fn get_effective_reasoning_config(&self) -> ReasoningConfig {
3551 if let Some(ref sm) = self.state_machine
3552 && let Some(state_def) = sm.current_definition()
3553 && let Some(ref state_reasoning) = state_def.reasoning
3554 {
3555 return state_reasoning.clone();
3556 }
3557 self.reasoning_config.clone()
3558 }
3559
3560 fn get_effective_reflection_config(&self) -> ReflectionConfig {
3561 if let Some(ref sm) = self.state_machine
3562 && let Some(state_def) = sm.current_definition()
3563 && let Some(ref state_reflection) = state_def.reflection
3564 {
3565 return state_reflection.clone();
3566 }
3567 self.reflection_config.clone()
3568 }
3569
3570 fn get_skill_reasoning_config(&self, skill: &SkillDefinition) -> ReasoningConfig {
3571 skill
3572 .reasoning
3573 .clone()
3574 .unwrap_or_else(|| self.get_effective_reasoning_config())
3575 }
3576
3577 fn get_skill_reflection_config(&self, skill: &SkillDefinition) -> ReflectionConfig {
3578 skill
3579 .reflection
3580 .clone()
3581 .unwrap_or_else(|| self.get_effective_reflection_config())
3582 }
3583
3584 async fn build_disambiguation_context(&self) -> Result<DisambiguationContext> {
3585 let recent_messages: Vec<String> = self
3586 .memory
3587 .get_messages(Some(5))
3588 .await?
3589 .iter()
3590 .rev()
3591 .map(|m| format!("{:?}: {}", m.role, m.content))
3592 .collect();
3593
3594 let current_state = self.current_state().map(|s| s.to_string());
3595
3596 let state_prompt: Option<String> = self
3599 .state_machine
3600 .as_ref()
3601 .and_then(|sm| sm.current_definition())
3602 .and_then(|def| def.prompt.clone());
3603
3604 let available_tools: Vec<String> = self
3605 .get_available_tool_ids()
3606 .await
3607 .unwrap_or_else(|_| self.tools.list_ids());
3608
3609 let available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
3610
3611 let mut user_context = self.build_context_with_overlays();
3612 user_context.remove(DISAMBIGUATION_STATE_GENERATION_KEY);
3613 if let Some(state_generation) = self
3614 .state_machine
3615 .as_ref()
3616 .map(|state_machine| state_machine.generation())
3617 {
3618 user_context.insert(
3619 DISAMBIGUATION_STATE_GENERATION_KEY.to_string(),
3620 serde_json::json!(state_generation),
3621 );
3622 }
3623
3624 let available_intents: Vec<String> = if let Some(ref sm) = self.state_machine {
3626 sm.current_definition()
3627 .map(|def| {
3628 def.transitions
3629 .iter()
3630 .filter_map(|t| t.intent.clone())
3631 .collect()
3632 })
3633 .unwrap_or_default()
3634 } else {
3635 Vec::new()
3636 };
3637
3638 Ok(DisambiguationContext::from_agent_state(
3639 recent_messages,
3640 current_state,
3641 state_prompt,
3642 available_tools,
3643 available_skills,
3644 available_intents,
3645 user_context,
3646 ))
3647 }
3648
3649 fn get_available_skills(&self) -> Vec<&SkillDefinition> {
3650 if let Some(ref sm) = self.state_machine
3651 && let Some(state_def) = sm.current_definition()
3652 {
3653 let parent_def = sm.get_parent_definition();
3654 let effective_skills = state_def.get_effective_skills(parent_def.as_ref());
3655 if !effective_skills.is_empty() {
3656 return self
3657 .skills
3658 .iter()
3659 .filter(|s| effective_skills.contains(&&s.id))
3660 .collect();
3661 }
3662 }
3663 self.skills.iter().collect()
3664 }
3665
3666 async fn build_messages(&self) -> Result<Vec<ChatMessage>> {
3667 self.build_messages_internal(true, None, true).await
3668 }
3669
3670 async fn build_messages_for_draft(&self, user_message: &str) -> Result<Vec<ChatMessage>> {
3671 self.build_messages_internal(false, Some(user_message), true)
3672 .await
3673 }
3674
3675 async fn build_messages_internal(
3676 &self,
3677 fire_persona_hooks: bool,
3678 ephemeral_user_message: Option<&str>,
3679 include_tool_prompt: bool,
3680 ) -> Result<Vec<ChatMessage>> {
3681 let system_prompt = self
3682 .get_effective_system_prompt_with_persona_hooks(fire_persona_hooks, include_tool_prompt)
3683 .await?;
3684 let mut messages = vec![ChatMessage::system(&system_prompt)];
3685
3686 let context = self.memory.get_context().await?;
3687 let history = if let Some(ref budget) = self.memory_token_budget {
3688 context.to_llm_messages_with_allocation(&budget.allocation)
3689 } else {
3690 context.to_llm_messages()
3691 };
3692 messages.extend(history);
3693 if let Some(user_message) = ephemeral_user_message {
3694 messages.push(ChatMessage::user(user_message));
3695 }
3696
3697 let total_tokens = self.estimate_total_tokens(&messages);
3698
3699 if total_tokens > self.max_context_tokens {
3700 debug!(
3701 total = total_tokens,
3702 limit = self.max_context_tokens,
3703 "Context overflow"
3704 );
3705
3706 match &self.recovery_manager.config().llm.on_context_overflow {
3707 ContextOverflowAction::Error => {
3708 return Err(AgentError::LLM(format!(
3709 "Context overflow: {} tokens > {} limit",
3710 total_tokens, self.max_context_tokens
3711 )));
3712 }
3713 ContextOverflowAction::Truncate { keep_recent } => {
3714 self.truncate_context(&mut messages, *keep_recent);
3715 }
3716 ContextOverflowAction::Summarize {
3717 summarizer_llm,
3718 max_summary_tokens,
3719 custom_prompt,
3720 keep_recent,
3721 filter,
3722 } => {
3723 self.summarize_context(
3724 &mut messages,
3725 summarizer_llm.as_deref(),
3726 *max_summary_tokens,
3727 custom_prompt.as_deref(),
3728 *keep_recent,
3729 filter.as_ref(),
3730 )
3731 .await?;
3732 }
3733 }
3734 }
3735
3736 Ok(messages)
3737 }
3738
3739 async fn main_tool_protocol(
3740 &self,
3741 llm: &dyn LLMProvider,
3742 ephemeral_new_turn: bool,
3743 ) -> Result<MainToolProtocol> {
3744 let mut choice = llm.configured_tool_choice();
3745 if matches!(choice.as_ref(), Some(ToolChoice::None)) {
3746 return Ok(MainToolProtocol {
3747 choice,
3748 tool_ids: Vec::new(),
3749 definitions: Vec::new(),
3750 });
3751 }
3752
3753 let mut tool_ids = self.get_available_tool_ids().await?;
3754 tool_ids.sort();
3755 tool_ids.dedup();
3756 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3757 let canonical = self.tools.canonical_id(expected).ok_or_else(|| {
3758 AgentError::Config(format!(
3759 "specific tool choice '{expected}' is not registered"
3760 ))
3761 })?;
3762 if canonical != *expected {
3763 return Err(AgentError::Config(format!(
3764 "specific tool choice must use canonical ID '{canonical}', not '{expected}'"
3765 )));
3766 }
3767 if !tool_ids.iter().any(|tool_id| tool_id == expected) {
3768 return Err(AgentError::Config(format!(
3769 "specific tool choice '{expected}' is outside the effective tool grant"
3770 )));
3771 }
3772 }
3773 if matches!(
3774 choice.as_ref(),
3775 Some(ToolChoice::Required | ToolChoice::Specific(_))
3776 ) && tool_ids.is_empty()
3777 {
3778 return Err(AgentError::Config(
3779 "required tool choice has no tool inside the effective grant".to_string(),
3780 ));
3781 }
3782 if !ephemeral_new_turn
3783 && let Some(configured_choice) = choice.as_ref()
3784 && matches!(
3785 configured_choice,
3786 ToolChoice::Required | ToolChoice::Specific(_)
3787 )
3788 && self
3789 .tool_choice_satisfied_in_current_turn(configured_choice, &tool_ids)
3790 .await?
3791 {
3792 choice = Some(ToolChoice::Auto);
3793 }
3794 if let Some(ToolChoice::Specific(expected)) = choice.as_ref() {
3795 tool_ids.retain(|tool_id| tool_id == expected);
3796 }
3797
3798 let definitions = tool_ids
3799 .iter()
3800 .map(|tool_id| {
3801 let tool = self.tools.get(tool_id).ok_or_else(|| {
3802 AgentError::Config(format!(
3803 "effective tool '{tool_id}' disappeared before provider exposure"
3804 ))
3805 })?;
3806 Ok(LLMToolDefinition {
3807 name: tool_id.clone(),
3808 description: tool.description().to_string(),
3809 input_schema: tool.input_schema(),
3810 })
3811 })
3812 .collect::<Result<Vec<_>>>()?;
3813
3814 Ok(MainToolProtocol {
3818 choice,
3819 tool_ids,
3820 definitions,
3821 })
3822 }
3823
3824 async fn tool_choice_satisfied_in_current_turn(
3825 &self,
3826 choice: &ToolChoice,
3827 effective_tool_ids: &[String],
3828 ) -> Result<bool> {
3829 let messages = self.memory.get_messages(None).await?;
3830 let mut saw_tool_result = false;
3831 for message in messages.iter().rev() {
3832 match message.role {
3833 ai_agents_core::Role::Tool | ai_agents_core::Role::Function => {
3834 saw_tool_result = true;
3835 }
3836 ai_agents_core::Role::Assistant if saw_tool_result => {
3837 let Some(calls) = self.parse_tool_calls(&message.content) else {
3838 continue;
3839 };
3840 let calls_are_effective = !calls.is_empty()
3841 && calls.iter().all(|call| {
3842 self.tools
3843 .canonical_id(&call.name)
3844 .is_some_and(|canonical| effective_tool_ids.contains(&canonical))
3845 });
3846 return Ok(calls_are_effective
3847 && match choice {
3848 ToolChoice::Required => true,
3849 ToolChoice::Specific(expected) => calls.iter().all(|call| {
3850 self.tools.canonical_id(&call.name).as_deref()
3851 == Some(expected.as_str())
3852 }),
3853 _ => false,
3854 });
3855 }
3856 ai_agents_core::Role::User => return Ok(false),
3857 _ => {}
3858 }
3859 }
3860 Ok(false)
3861 }
3862
3863 fn provider_can_use_native_tools(
3864 &self,
3865 llm: &dyn LLMProvider,
3866 protocol: &MainToolProtocol,
3867 ) -> bool {
3868 let Some(choice) = protocol.choice.as_ref() else {
3869 return false;
3870 };
3871 if matches!(choice, ToolChoice::None) || protocol.definitions.is_empty() {
3872 return false;
3873 }
3874 llm.supports_tool_choice(choice)
3875 && protocol.definitions.iter().all(|definition| {
3876 !definition.name.is_empty()
3877 && definition.name.len() <= 64
3878 && definition
3879 .name
3880 .bytes()
3881 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
3882 })
3883 }
3884
3885 fn prompt_messages_for_tool_protocol(
3886 &self,
3887 messages: &[ChatMessage],
3888 protocol: &MainToolProtocol,
3889 corrective: bool,
3890 ) -> Vec<ChatMessage> {
3891 let mut messages = messages.to_vec();
3892 let Some(choice) = protocol.choice.as_ref() else {
3893 return messages;
3894 };
3895 if matches!(choice, ToolChoice::None) || protocol.tool_ids.is_empty() {
3896 return messages;
3897 }
3898
3899 let mut tool_prompt = self.tools.generate_scoped_prompt_with_mode(
3900 &protocol.tool_ids,
3901 None,
3902 self.parallel_tools.enabled,
3903 self.runtime_config.tool_schema_prompt_mode,
3904 );
3905 match choice {
3906 ToolChoice::Required => tool_prompt.push_str(
3907 "\n\nYou must call at least one listed tool before giving a final answer.",
3908 ),
3909 ToolChoice::Specific(tool_id) => tool_prompt.push_str(&format!(
3910 "\n\nYou must call the '{tool_id}' tool before giving a final answer."
3911 )),
3912 ToolChoice::Auto => {}
3913 ToolChoice::None => return messages,
3914 _ => return messages,
3915 }
3916 if let Some(system) = messages
3917 .iter_mut()
3918 .find(|message| message.role == ai_agents_core::Role::System)
3919 {
3920 system.content.push_str("\n\n");
3921 system.content.push_str(&tool_prompt);
3922 } else {
3923 messages.insert(0, ChatMessage::system(tool_prompt));
3924 }
3925 if corrective {
3926 let instruction = match choice {
3927 ToolChoice::Required => {
3928 "Your previous response did not call a required tool. Call at least one listed tool now and return only the JSON tool call."
3929 }
3930 ToolChoice::Specific(tool_id) => {
3931 messages.push(ChatMessage::user(format!(
3932 "Your previous response did not call the required '{tool_id}' tool. Call it now and return only the JSON tool call."
3933 )));
3934 return messages;
3935 }
3936 _ => return messages,
3937 };
3938 messages.push(ChatMessage::user(instruction));
3939 }
3940 messages
3941 }
3942
3943 async fn invoke_main_provider(
3944 &self,
3945 llm: Arc<dyn LLMProvider>,
3946 messages: &[ChatMessage],
3947 protocol: &MainToolProtocol,
3948 corrective: bool,
3949 ) -> std::result::Result<MainProviderResponse, LLMError> {
3950 let use_native = self.provider_can_use_native_tools(llm.as_ref(), protocol);
3951 let response = if use_native {
3952 let request = LLMToolRequest {
3953 tools: protocol.definitions.clone(),
3954 choice: protocol
3955 .choice
3956 .clone()
3957 .expect("native tool requests require an explicit choice"),
3958 };
3959 self.observe_purpose(
3960 ObservationPurpose::MainResponse,
3961 llm.complete_with_tools(messages, None, &request),
3962 )
3963 .await?
3964 } else {
3965 let prompt_messages =
3966 self.prompt_messages_for_tool_protocol(messages, protocol, corrective);
3967 self.observe_purpose(
3968 ObservationPurpose::MainResponse,
3969 llm.complete(&prompt_messages, None),
3970 )
3971 .await?
3972 };
3973 Ok(MainProviderResponse {
3974 response,
3975 used_native_tools: use_native,
3976 })
3977 }
3978
3979 async fn complete_main_attempt_with_recovery(
3980 &self,
3981 llm: Arc<dyn LLMProvider>,
3982 messages: &[ChatMessage],
3983 protocol: &MainToolProtocol,
3984 corrective: bool,
3985 ) -> Result<MainProviderResponse> {
3986 let primary_result = if self.recovery_manager.config().default.max_retries > 0 {
3987 self.recovery_manager
3988 .with_retry("llm_call", None, || {
3989 let llm = Arc::clone(&llm);
3990 async move {
3991 self.invoke_main_provider(llm, messages, protocol, corrective)
3992 .await
3993 .map_err(|error| error.classify())
3994 }
3995 })
3996 .await
3997 .map_err(|error| AgentError::LLM(error.to_string()))
3998 } else {
3999 self.invoke_main_provider(Arc::clone(&llm), messages, protocol, corrective)
4000 .await
4001 .map_err(|error| AgentError::LLM(error.to_string()))
4002 };
4003
4004 match primary_result {
4005 Ok(response) => Ok(response),
4006 Err(primary_error) => match &self.recovery_manager.config().llm.on_failure {
4007 LLMFailureAction::FallbackLlm { fallback_llm } => {
4008 let fallback = self.llm_registry.get(fallback_llm).map_err(|error| {
4009 AgentError::Config(format!(
4010 "Fallback LLM '{fallback_llm}' not found: {error}"
4011 ))
4012 })?;
4013 self.invoke_main_provider(fallback, messages, protocol, corrective)
4014 .await
4015 .map_err(|error| AgentError::LLM(error.to_string()))
4016 }
4017 LLMFailureAction::FallbackResponse { message } => {
4018 if matches!(
4019 protocol.choice.as_ref(),
4020 Some(ToolChoice::Required | ToolChoice::Specific(_))
4021 ) {
4022 Err(AgentError::LLM(format!(
4023 "Required tool selection failed and cannot be satisfied by a static fallback response: {primary_error}"
4024 )))
4025 } else {
4026 Ok(MainProviderResponse {
4027 response: LLMResponse::new(message.clone(), FinishReason::Stop),
4028 used_native_tools: false,
4029 })
4030 }
4031 }
4032 LLMFailureAction::Error => Err(primary_error),
4033 },
4034 }
4035 }
4036
4037 fn normalize_main_provider_response(
4038 &self,
4039 mut response: LLMResponse,
4040 protocol: &MainToolProtocol,
4041 ) -> Result<(LLMResponse, bool)> {
4042 let native_calls = response
4043 .tool_calls()
4044 .map_err(|error| AgentError::LLM(error.to_string()))?;
4045 let calls = match native_calls {
4046 Some(calls) => {
4047 let markers = calls
4048 .iter()
4049 .map(|call| {
4050 serde_json::json!({
4051 "_ai_agents_native_tool_call": true,
4052 "id": call.id,
4053 "tool": call.name,
4054 "arguments": call.arguments,
4055 })
4056 })
4057 .collect::<Vec<_>>();
4058 response.content = if markers.len() == 1 {
4059 markers[0].to_string()
4060 } else {
4061 serde_json::Value::Array(markers).to_string()
4062 };
4063 Some(calls)
4064 }
4065 None if !matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) => {
4066 self.parse_tool_calls(response.content.trim())
4067 }
4068 None => None,
4069 };
4070
4071 if protocol.choice.is_some()
4072 && let Some(calls) = calls.as_ref()
4073 && calls.iter().any(|call| {
4074 self.tools
4075 .canonical_id(&call.name)
4076 .is_none_or(|canonical| !protocol.tool_ids.contains(&canonical))
4077 })
4078 {
4079 return Err(AgentError::LLM(
4080 "Provider returned a tool call outside the effective grant".to_string(),
4081 ));
4082 }
4083
4084 let compliant = match protocol.choice.as_ref() {
4085 Some(ToolChoice::Required) => calls.as_ref().is_some_and(|calls| !calls.is_empty()),
4086 Some(ToolChoice::Specific(expected)) => calls.as_ref().is_some_and(|calls| {
4087 !calls.is_empty()
4088 && calls.iter().all(|call| {
4089 self.tools.canonical_id(&call.name).as_deref() == Some(expected.as_str())
4090 })
4091 }),
4092 _ => true,
4093 };
4094 Ok((response, compliant))
4095 }
4096
4097 async fn complete_main_llm_with_recovery(
4098 &self,
4099 llm: Arc<dyn LLMProvider>,
4100 messages: &[ChatMessage],
4101 protocol: &MainToolProtocol,
4102 ) -> Result<LLMResponse> {
4103 let first = self
4104 .complete_main_attempt_with_recovery(Arc::clone(&llm), messages, protocol, false)
4105 .await?;
4106 let (response, compliant) =
4107 self.normalize_main_provider_response(first.response, protocol)?;
4108 if compliant {
4109 return Ok(response);
4110 }
4111 if first.used_native_tools {
4112 return Err(AgentError::LLM(
4113 "Provider returned no compliant native call for required tool choice".to_string(),
4114 ));
4115 }
4116
4117 let corrected = self
4118 .complete_main_attempt_with_recovery(llm, messages, protocol, true)
4119 .await?;
4120 let (response, compliant) =
4121 self.normalize_main_provider_response(corrected.response, protocol)?;
4122 if compliant {
4123 return Ok(response);
4124 }
4125 Err(AgentError::LLM(
4126 "Provider returned no compliant tool call after one corrective retry".to_string(),
4127 ))
4128 }
4129
4130 fn is_native_tool_call_content(content: &str) -> bool {
4131 let Ok(value) = serde_json::from_str::<serde_json::Value>(content) else {
4132 return false;
4133 };
4134 match value {
4135 serde_json::Value::Array(values) => {
4136 !values.is_empty()
4137 && values.iter().all(|value| {
4138 value
4139 .get("_ai_agents_native_tool_call")
4140 .and_then(|marker| marker.as_bool())
4141 == Some(true)
4142 })
4143 }
4144 serde_json::Value::Object(map) => {
4145 map.get("_ai_agents_native_tool_call")
4146 .and_then(|marker| marker.as_bool())
4147 == Some(true)
4148 }
4149 _ => false,
4150 }
4151 }
4152
4153 fn tool_result_message(
4154 tool_call: &ToolCall,
4155 output: &str,
4156 native_tool_call: bool,
4157 ) -> ChatMessage {
4158 if !native_tool_call {
4159 return ChatMessage::function(&tool_call.name, output);
4160 }
4161 let output = serde_json::from_str::<serde_json::Value>(output)
4162 .unwrap_or_else(|_| serde_json::Value::String(output.to_string()));
4163 ChatMessage::function(
4164 &tool_call.name,
4165 serde_json::json!({
4166 "_ai_agents_native_tool_result": true,
4167 "id": tool_call.id,
4168 "tool": tool_call.name,
4169 "output": output,
4170 })
4171 .to_string(),
4172 )
4173 }
4174
4175 fn parse_main_tool_calls(
4176 &self,
4177 content: &str,
4178 protocol: &MainToolProtocol,
4179 ) -> Option<Vec<ToolCall>> {
4180 if matches!(protocol.choice.as_ref(), Some(ToolChoice::None)) {
4181 None
4182 } else {
4183 self.parse_tool_calls(content)
4184 }
4185 }
4186
4187 fn parse_tool_calls(&self, content: &str) -> Option<Vec<ToolCall>> {
4188 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(content) {
4190 if let Some(arr) = parsed.as_array() {
4192 let calls: Vec<ToolCall> = arr
4193 .iter()
4194 .filter_map(|v| self.extract_tool_call_from_value(v))
4195 .collect();
4196 if !calls.is_empty() {
4197 return Some(calls);
4198 }
4199 }
4200 if let Some(tool_call) = self.extract_tool_call_from_value(&parsed) {
4202 return Some(vec![tool_call]);
4203 }
4204 }
4205
4206 if let Some(json_str) = self.extract_json_from_content(content)
4208 && let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&json_str)
4209 {
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 None
4227 }
4228
4229 fn extract_tool_call_from_value(&self, parsed: &serde_json::Value) -> Option<ToolCall> {
4230 if let Some(tool_name) = parsed.get("tool").and_then(|v| v.as_str()) {
4231 let arguments = parsed
4232 .get("arguments")
4233 .cloned()
4234 .unwrap_or(serde_json::json!({}));
4235 return Some(ToolCall {
4236 id: parsed
4237 .get("id")
4238 .and_then(|value| value.as_str())
4239 .filter(|id| !id.is_empty())
4240 .map(str::to_string)
4241 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
4242 name: tool_name.to_string(),
4243 arguments,
4244 });
4245 }
4246 None
4247 }
4248
4249 fn extract_json_from_content(&self, content: &str) -> Option<String> {
4251 if let Some(result) = self.extract_json_array_from_content(content) {
4253 return Some(result);
4254 }
4255 self.extract_json_object_from_content(content)
4256 }
4257
4258 fn extract_json_array_from_content(&self, content: &str) -> Option<String> {
4260 let start = content.find('[')?;
4261 let content_from_start = &content[start..];
4262
4263 let mut depth = 0;
4264 let mut end = 0;
4265 for (i, ch) in content_from_start.char_indices() {
4266 match ch {
4267 '[' => depth += 1,
4268 ']' => {
4269 depth -= 1;
4270 if depth == 0 {
4271 end = i + 1;
4272 break;
4273 }
4274 }
4275 _ => {}
4276 }
4277 }
4278
4279 if end > 0 {
4280 let json_str = &content_from_start[..end];
4281 if json_str.contains("\"tool\"") {
4283 return Some(json_str.to_string());
4284 }
4285 }
4286
4287 None
4288 }
4289
4290 fn extract_json_object_from_content(&self, content: &str) -> Option<String> {
4292 let start = content.find('{')?;
4293 let content_from_start = &content[start..];
4294
4295 let mut depth = 0;
4297 let mut end = 0;
4298 for (i, ch) in content_from_start.char_indices() {
4299 match ch {
4300 '{' => depth += 1,
4301 '}' => {
4302 depth -= 1;
4303 if depth == 0 {
4304 end = i + 1;
4305 break;
4306 }
4307 }
4308 _ => {}
4309 }
4310 }
4311
4312 if end > 0 {
4313 let json_str = &content_from_start[..end];
4314 if json_str.contains("\"tool\"") {
4316 return Some(json_str.to_string());
4317 }
4318 }
4319
4320 None
4321 }
4322
4323 #[allow(clippy::too_many_arguments)]
4327 fn record_from_parts(
4328 &self,
4329 request: &ToolExecutionRequest,
4330 canonical_id: String,
4331 executed_arguments: Value,
4332 started_at: chrono::DateTime<chrono::Utc>,
4333 start: Instant,
4334 executed: bool,
4335 success: bool,
4336 output: String,
4337 metadata: HashMap<String, Value>,
4338 policy: ToolPolicyDecisionRecord,
4339 approval: Option<ToolApprovalRecord>,
4340 timed_out: bool,
4341 output_truncated: bool,
4342 ) -> ToolExecutionRecord {
4343 let versions = ToolDecisionVersions {
4344 policy: self.active_tool_security().policy_version(),
4345 registry: self.tools.version(),
4346 runtime_control: self.runtime_control.version.load(Ordering::SeqCst),
4347 state: self
4348 .state_machine
4349 .as_ref()
4350 .map(|state_machine| state_machine.generation()),
4351 };
4352 self.record_from_parts_at(
4353 request,
4354 canonical_id,
4355 executed_arguments,
4356 started_at,
4357 start,
4358 executed,
4359 success,
4360 output,
4361 metadata,
4362 policy,
4363 approval,
4364 timed_out,
4365 output_truncated,
4366 versions,
4367 )
4368 }
4369
4370 #[allow(clippy::too_many_arguments)]
4372 fn record_from_parts_at(
4373 &self,
4374 request: &ToolExecutionRequest,
4375 canonical_id: String,
4376 executed_arguments: Value,
4377 started_at: chrono::DateTime<chrono::Utc>,
4378 start: Instant,
4379 executed: bool,
4380 success: bool,
4381 output: String,
4382 metadata: HashMap<String, Value>,
4383 policy: ToolPolicyDecisionRecord,
4384 approval: Option<ToolApprovalRecord>,
4385 timed_out: bool,
4386 output_truncated: bool,
4387 versions: ToolDecisionVersions,
4388 ) -> ToolExecutionRecord {
4389 ToolExecutionRecord {
4390 call_id: request.call_id.clone(),
4391 requested_name: request.requested_name.clone(),
4392 canonical_id,
4393 source: request.source.clone(),
4394 arguments: request.arguments.clone(),
4395 executed_arguments,
4396 policy_version: versions.policy,
4397 registry_version: versions.registry,
4398 runtime_config_version: versions.runtime_control,
4399 executed,
4400 success,
4401 output,
4402 metadata,
4403 policy,
4404 approval,
4405 started_at,
4406 duration_ms: start.elapsed().as_millis() as u64,
4407 timed_out,
4408 cancelled: false,
4409 cancellation_reason: None,
4410 output_truncated,
4411 }
4412 }
4413
4414 async fn finish_tool_record(&self, record: &ToolExecutionRecord) {
4416 let result = ToolResult {
4417 success: record.success,
4418 output: record.model_output_string(),
4419 metadata: if record.metadata.is_empty() {
4420 None
4421 } else {
4422 Some(record.metadata.clone())
4423 },
4424 };
4425 self.hooks
4426 .on_tool_complete(&record.canonical_id, &result, record.duration_ms)
4427 .await;
4428 self.hooks.on_tool_execution_record(record).await;
4429 self.record_tool_call(&record.canonical_id, record.model_output_value());
4430 if !record.success {
4431 self.hooks
4432 .on_error(&AgentError::Tool(record.output.clone()))
4433 .await;
4434 }
4435 }
4436
4437 async fn finish_tool_record_after_resource_guards(
4439 &self,
4440 resource_guards: ToolResourceGuards,
4441 record: &ToolExecutionRecord,
4442 ) {
4443 drop(resource_guards);
4444 self.finish_tool_record(record).await;
4445 }
4446
4447 async fn execute_resolved_tool_once(
4449 &self,
4450 tool: Arc<dyn ai_agents_core::Tool>,
4451 args: Value,
4452 ctx: ToolExecutionContext,
4453 timeout_ms: u64,
4454 ) -> Result<(ToolResult, bool, bool, bool)> {
4455 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4456 return Ok((
4457 ToolResult::error("Tool execution cancelled by runtime control"),
4458 false,
4459 true,
4460 false,
4461 ));
4462 }
4463 let invoked = Arc::new(AtomicBool::new(false));
4467 let invoked_by_future = Arc::clone(&invoked);
4468 let actor_context = current_turn_actor_context();
4469 let future = async move {
4470 invoked_by_future.store(true, Ordering::SeqCst);
4471 if let Some(actor_context) = actor_context {
4472 scope_actor_context(actor_context, tool.execute(args, ctx)).await
4473 } else {
4474 tool.execute(args, ctx).await
4475 }
4476 };
4477 tokio::pin!(future);
4478 let timeout = tokio::time::sleep(std::time::Duration::from_millis(timeout_ms));
4479 tokio::pin!(timeout);
4480 let mut cancel_tick = tokio::time::interval(std::time::Duration::from_millis(50));
4481
4482 loop {
4483 tokio::select! {
4484 result = &mut future => return Ok((result, false, false, true)),
4485 _ = &mut timeout => {
4486 return Ok((
4487 ToolResult::error("Tool execution timed out"),
4488 true,
4489 false,
4490 invoked.load(Ordering::SeqCst),
4491 ));
4492 }
4493 _ = cancel_tick.tick() => {
4494 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4495 return Ok((
4496 ToolResult::error("Tool execution cancelled by runtime control"),
4497 false,
4498 true,
4499 invoked.load(Ordering::SeqCst),
4500 ));
4501 }
4502 }
4503 }
4504 }
4505 }
4506
4507 fn truncate_tool_output(output: String, max_chars: Option<usize>) -> (String, bool) {
4509 let Some(max_chars) = max_chars else {
4510 return (output, false);
4511 };
4512 let mut chars = output.chars();
4513 let truncated: String = chars.by_ref().take(max_chars).collect();
4514 if chars.next().is_some() {
4515 (truncated, true)
4516 } else {
4517 (output, false)
4518 }
4519 }
4520
4521 async fn acquire_tool_resource_locks(&self, keys: &[String]) -> Option<ToolResourceGuards> {
4523 let locks = {
4524 let mut table = self.resource_locks.write();
4525 table.retain(|_, lock| lock.strong_count() > 0);
4526 keys.iter()
4527 .map(|key| {
4528 if let Some(lock) = table.get(key).and_then(Weak::upgrade) {
4529 lock
4530 } else {
4531 let lock = Arc::new(tokio::sync::Mutex::new(()));
4532 table.insert(key.clone(), Arc::downgrade(&lock));
4533 lock
4534 }
4535 })
4536 .collect::<Vec<_>>()
4537 };
4538 let mut resource_guards = ToolResourceGuards {
4539 guards: Vec::with_capacity(locks.len()),
4540 locks: Arc::clone(&self.resource_locks),
4541 };
4542 let mut locks = locks.into_iter();
4543 while let Some(lock) = locks.next() {
4544 let mut lock = Box::pin(lock.lock_owned());
4545 loop {
4546 tokio::select! {
4547 guard = &mut lock => {
4548 resource_guards.guards.push(guard);
4549 break;
4550 }
4551 _ = tokio::time::sleep(std::time::Duration::from_millis(10)) => {
4552 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4553 drop(lock);
4554 drop(locks);
4555 drop(resource_guards);
4556 return None;
4557 }
4558 }
4559 }
4560 }
4561 }
4562 Some(resource_guards)
4563 }
4564
4565 async fn run_tool_with_retries(
4567 &self,
4568 canonical_id: &str,
4569 tool: Arc<dyn ai_agents_core::Tool>,
4570 args: Value,
4571 ctx: ToolExecutionContext,
4572 timeout_ms: u64,
4573 max_retries: u32,
4574 ) -> Result<(ToolResult, bool, bool, bool)> {
4575 let max_retries = if ctx.classification.safely_retryable {
4576 max_retries
4577 } else {
4578 0
4579 };
4580 let mut attempts = 0;
4581 let mut invoked = false;
4582 loop {
4583 let (result, timed_out, cancelled, attempt_invoked) = self
4584 .execute_resolved_tool_once(tool.clone(), args.clone(), ctx.clone(), timeout_ms)
4585 .await?;
4586 invoked |= attempt_invoked;
4587 if result.success || timed_out || cancelled || attempts >= max_retries {
4588 return Ok((result, timed_out, cancelled, invoked));
4589 }
4590 attempts += 1;
4591 warn!(tool = %canonical_id, attempt = attempts, error = %result.output, "Retrying failed tool call");
4592 }
4593 }
4594
4595 fn execute_tool_record(
4597 &self,
4598 request: ToolExecutionRequest,
4599 ) -> Pin<Box<dyn Future<Output = Result<ToolExecutionRecord>> + Send + '_>> {
4600 Box::pin(self.execute_tool_record_inner(request))
4601 }
4602
4603 async fn execute_tool_record_inner(
4604 &self,
4605 request: ToolExecutionRequest,
4606 ) -> Result<ToolExecutionRecord> {
4607 let started_at = chrono::Utc::now();
4608 let start = Instant::now();
4609 info!(tool = %request.requested_name, args = %request.arguments, "Executing tool");
4610
4611 if self.runtime_control.emergency_deny.load(Ordering::SeqCst) {
4612 let record = self.record_from_parts(
4613 &request,
4614 request.requested_name.clone(),
4615 request.arguments.clone(),
4616 started_at,
4617 start,
4618 false,
4619 false,
4620 "Tool execution is disabled by runtime control".to_string(),
4621 HashMap::new(),
4622 ToolPolicyDecisionRecord::deny("runtime emergency deny is enabled"),
4623 None,
4624 false,
4625 false,
4626 );
4627 self.finish_tool_record(&record).await;
4628 return Ok(record);
4629 }
4630
4631 let Some(resolved) = self.tools.resolve(&request.requested_name) else {
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 format!("Tool '{}' is unavailable", request.requested_name),
4641 HashMap::new(),
4642 ToolPolicyDecisionRecord::unavailable(format!(
4643 "Tool '{}' is not registered",
4644 request.requested_name
4645 )),
4646 None,
4647 false,
4648 false,
4649 );
4650 self.finish_tool_record(&record).await;
4651 return Ok(record);
4652 };
4653
4654 let canonical_id = resolved.identity.canonical_id.clone();
4655
4656 let initial_scope_snapshot = self.get_available_tool_ids_snapshot().await?;
4657 if !initial_scope_snapshot
4658 .tool_ids
4659 .iter()
4660 .any(|id| id == &canonical_id)
4661 {
4662 let record = self.record_from_parts(
4663 &request,
4664 canonical_id.clone(),
4665 request.arguments.clone(),
4666 started_at,
4667 start,
4668 false,
4669 false,
4670 format!(
4671 "Tool '{}' is not available in the current scope",
4672 canonical_id
4673 ),
4674 HashMap::new(),
4675 ToolPolicyDecisionRecord::deny(format!(
4676 "Tool '{}' is not granted by the current top-level and state tool scope",
4677 canonical_id
4678 )),
4679 None,
4680 false,
4681 false,
4682 );
4683 self.finish_tool_record(&record).await;
4684 return Ok(record);
4685 }
4686
4687 let approval_control_snapshot = self.runtime_safety_snapshot();
4688 let security_engine = approval_control_snapshot.tool_security.clone();
4689 let bindings = resolved.tool.policy_bindings();
4690 let mut executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
4691 &canonical_id,
4692 &request.arguments,
4693 &bindings,
4694 );
4695 self.hooks
4696 .on_tool_start(&canonical_id, &executed_arguments)
4697 .await;
4698
4699 let mut metadata = HashMap::new();
4700 let safety = resolved.tool.safety_metadata();
4701 let classification = resolved.tool.classify_call(&executed_arguments);
4702 let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
4703 metadata.insert(
4704 "classification".to_string(),
4705 serde_json::to_value(&classification).unwrap_or(Value::Null),
4706 );
4707 metadata.insert(
4708 "effective_limits".to_string(),
4709 serde_json::to_value(&limits).unwrap_or(Value::Null),
4710 );
4711 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
4712 if !policy_snapshot.is_null() {
4713 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
4714 }
4715
4716 let mut approval_record = Some(ToolApprovalRecord {
4717 status: ToolApprovalStatus::NotRequired,
4718 reason: None,
4719 modified_arguments: None,
4720 });
4721
4722 let mut security_result = security_engine
4723 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
4724 .await?;
4725 match &security_result {
4726 SecurityCheckResult::Allow => {}
4727 SecurityCheckResult::Warn { message } => {
4728 warn!(tool = %canonical_id, message = %message, "Tool security warning");
4729 }
4730 SecurityCheckResult::Block { reason } => {
4731 let record = self.record_from_parts(
4732 &request,
4733 canonical_id,
4734 executed_arguments,
4735 started_at,
4736 start,
4737 false,
4738 false,
4739 format!("Denied: {}", reason),
4740 metadata,
4741 ToolPolicyDecisionRecord::deny(reason.clone()),
4742 approval_record,
4743 false,
4744 false,
4745 );
4746 self.finish_tool_record(&record).await;
4747 return Ok(record);
4748 }
4749 SecurityCheckResult::Unavailable { reason } => {
4750 let record = self.record_from_parts(
4751 &request,
4752 canonical_id,
4753 executed_arguments,
4754 started_at,
4755 start,
4756 false,
4757 false,
4758 format!("Unavailable: {}", reason),
4759 metadata,
4760 ToolPolicyDecisionRecord::unavailable(reason.clone()),
4761 approval_record,
4762 false,
4763 false,
4764 );
4765 self.finish_tool_record(&record).await;
4766 return Ok(record);
4767 }
4768 SecurityCheckResult::RequireConfirmation { message } => {
4769 if self.hitl_engine.is_none() {
4770 approval_record = Some(ToolApprovalRecord {
4771 status: ToolApprovalStatus::Unavailable,
4772 reason: Some("No HITL engine configured".to_string()),
4773 modified_arguments: None,
4774 });
4775 let record = self.record_from_parts(
4776 &request,
4777 canonical_id,
4778 executed_arguments,
4779 started_at,
4780 start,
4781 false,
4782 false,
4783 format!("Approval unavailable: {}", message),
4784 metadata,
4785 ToolPolicyDecisionRecord::approval(message.clone()),
4786 approval_record,
4787 false,
4788 false,
4789 );
4790 self.finish_tool_record(&record).await;
4791 return Ok(record);
4792 }
4793
4794 let check_result = HITLCheckResult::required(
4795 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4796 HashMap::new(),
4797 message.clone(),
4798 None,
4799 );
4800 match self.request_hitl_approval(check_result).await? {
4801 ApprovalResult::Approved => {
4802 merge_approved_record(&mut approval_record);
4803 }
4804 ApprovalResult::Modified { changes } => {
4805 if let Some(obj) = executed_arguments.as_object_mut() {
4806 for (key, value) in changes {
4807 obj.insert(key, value);
4808 }
4809 }
4810 security_result = security_engine
4811 .validate_tool_execution_with_bindings(
4812 &canonical_id,
4813 &executed_arguments,
4814 &bindings,
4815 )
4816 .await?;
4817 if !matches!(
4818 security_result,
4819 SecurityCheckResult::Allow
4820 | SecurityCheckResult::Warn { .. }
4821 | SecurityCheckResult::RequireConfirmation { .. }
4822 ) {
4823 let reason = security_result
4824 .reason()
4825 .unwrap_or("modified arguments failed policy")
4826 .to_string();
4827 let record = self.record_from_parts(
4828 &request,
4829 canonical_id,
4830 executed_arguments.clone(),
4831 started_at,
4832 start,
4833 false,
4834 false,
4835 reason.clone(),
4836 metadata,
4837 ToolPolicyDecisionRecord::deny(reason),
4838 Some(ToolApprovalRecord {
4839 status: ToolApprovalStatus::Modified,
4840 reason: None,
4841 modified_arguments: Some(executed_arguments),
4842 }),
4843 false,
4844 false,
4845 );
4846 self.finish_tool_record(&record).await;
4847 return Ok(record);
4848 }
4849 approval_record = Some(ToolApprovalRecord {
4850 status: ToolApprovalStatus::Modified,
4851 reason: None,
4852 modified_arguments: Some(executed_arguments.clone()),
4853 });
4854 }
4855 ApprovalResult::Rejected { reason } => {
4856 let reason = reason.unwrap_or_else(|| "rejected".to_string());
4857 approval_record = Some(ToolApprovalRecord {
4858 status: ToolApprovalStatus::Rejected,
4859 reason: Some(reason.clone()),
4860 modified_arguments: None,
4861 });
4862 let record = self.record_from_parts(
4863 &request,
4864 canonical_id,
4865 executed_arguments,
4866 started_at,
4867 start,
4868 false,
4869 false,
4870 format!("Approval rejected: {}", reason),
4871 metadata,
4872 ToolPolicyDecisionRecord::approval(reason),
4873 approval_record,
4874 false,
4875 false,
4876 );
4877 self.finish_tool_record(&record).await;
4878 return Ok(record);
4879 }
4880 ApprovalResult::Timeout => {
4881 approval_record = Some(ToolApprovalRecord {
4882 status: ToolApprovalStatus::Timeout,
4883 reason: Some("approval timeout".to_string()),
4884 modified_arguments: None,
4885 });
4886 let record = self.record_from_parts(
4887 &request,
4888 canonical_id,
4889 executed_arguments,
4890 started_at,
4891 start,
4892 false,
4893 false,
4894 "Approval timed out".to_string(),
4895 metadata,
4896 ToolPolicyDecisionRecord::approval("approval timeout"),
4897 approval_record,
4898 false,
4899 false,
4900 );
4901 self.finish_tool_record(&record).await;
4902 return Ok(record);
4903 }
4904 }
4905 }
4906 }
4907
4908 if canonical_id == "command" && !self.tools.command_runner_available() {
4909 let record = self.record_from_parts(
4910 &request,
4911 canonical_id.clone(),
4912 executed_arguments.clone(),
4913 started_at,
4914 start,
4915 false,
4916 false,
4917 "Command runner is unavailable".to_string(),
4918 metadata,
4919 ToolPolicyDecisionRecord::unavailable("command runner is unavailable"),
4920 Some(ToolApprovalRecord {
4921 status: ToolApprovalStatus::Unavailable,
4922 reason: Some("command runner is unavailable".to_string()),
4923 modified_arguments: None,
4924 }),
4925 false,
4926 false,
4927 );
4928 self.finish_tool_record(&record).await;
4929 return Ok(record);
4930 }
4931
4932 if approval_record
4933 .as_ref()
4934 .is_some_and(|record| matches!(record.status, ToolApprovalStatus::NotRequired))
4935 && let Some(message) =
4936 security_engine.classification_approval_message(&canonical_id, &classification)
4937 {
4938 if self.hitl_engine.is_none() {
4939 approval_record = Some(ToolApprovalRecord {
4940 status: ToolApprovalStatus::Unavailable,
4941 reason: Some("No HITL engine configured".to_string()),
4942 modified_arguments: None,
4943 });
4944 let record = self.record_from_parts(
4945 &request,
4946 canonical_id,
4947 executed_arguments,
4948 started_at,
4949 start,
4950 false,
4951 false,
4952 format!("Approval unavailable: {}", message),
4953 metadata,
4954 ToolPolicyDecisionRecord::approval(message),
4955 approval_record,
4956 false,
4957 false,
4958 );
4959 self.finish_tool_record(&record).await;
4960 return Ok(record);
4961 }
4962 let check_result = HITLCheckResult::required(
4963 ApprovalTrigger::tool(&canonical_id, executed_arguments.clone()),
4964 HashMap::new(),
4965 message.clone(),
4966 None,
4967 );
4968 match self.request_hitl_approval(check_result).await? {
4969 ApprovalResult::Approved => {
4970 merge_approved_record(&mut approval_record);
4971 }
4972 ApprovalResult::Modified { changes } => {
4973 if let Some(obj) = executed_arguments.as_object_mut() {
4974 for (key, value) in changes {
4975 obj.insert(key, value);
4976 }
4977 }
4978 let modified_security = security_engine
4979 .validate_tool_execution_with_bindings(
4980 &canonical_id,
4981 &executed_arguments,
4982 &bindings,
4983 )
4984 .await?;
4985 if !matches!(
4986 modified_security,
4987 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
4988 ) {
4989 let reason = modified_security
4990 .reason()
4991 .unwrap_or("modified arguments failed policy")
4992 .to_string();
4993 let record = self.record_from_parts(
4994 &request,
4995 canonical_id,
4996 executed_arguments.clone(),
4997 started_at,
4998 start,
4999 false,
5000 false,
5001 reason.clone(),
5002 metadata,
5003 ToolPolicyDecisionRecord::deny(reason),
5004 Some(ToolApprovalRecord {
5005 status: ToolApprovalStatus::Modified,
5006 reason: None,
5007 modified_arguments: Some(executed_arguments),
5008 }),
5009 false,
5010 false,
5011 );
5012 self.finish_tool_record(&record).await;
5013 return Ok(record);
5014 }
5015 approval_record = Some(ToolApprovalRecord {
5016 status: ToolApprovalStatus::Modified,
5017 reason: None,
5018 modified_arguments: Some(executed_arguments.clone()),
5019 });
5020 }
5021 ApprovalResult::Rejected { reason } => {
5022 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5023 let record = self.record_from_parts(
5024 &request,
5025 canonical_id,
5026 executed_arguments,
5027 started_at,
5028 start,
5029 false,
5030 false,
5031 format!("Approval rejected: {}", reason),
5032 metadata,
5033 ToolPolicyDecisionRecord::approval(reason.clone()),
5034 Some(ToolApprovalRecord {
5035 status: ToolApprovalStatus::Rejected,
5036 reason: Some(reason),
5037 modified_arguments: None,
5038 }),
5039 false,
5040 false,
5041 );
5042 self.finish_tool_record(&record).await;
5043 return Ok(record);
5044 }
5045 ApprovalResult::Timeout => {
5046 let record = self.record_from_parts(
5047 &request,
5048 canonical_id,
5049 executed_arguments,
5050 started_at,
5051 start,
5052 false,
5053 false,
5054 "Approval timed out".to_string(),
5055 metadata,
5056 ToolPolicyDecisionRecord::approval("approval timeout"),
5057 Some(ToolApprovalRecord {
5058 status: ToolApprovalStatus::Timeout,
5059 reason: Some("approval timeout".to_string()),
5060 modified_arguments: None,
5061 }),
5062 false,
5063 false,
5064 );
5065 self.finish_tool_record(&record).await;
5066 return Ok(record);
5067 }
5068 }
5069 }
5070
5071 if canonical_id == "diagnostics" && !self.tools.diagnostics_available() {
5072 let record = self.record_from_parts(
5073 &request,
5074 canonical_id.clone(),
5075 executed_arguments.clone(),
5076 started_at,
5077 start,
5078 false,
5079 false,
5080 "Diagnostics provider is unavailable".to_string(),
5081 metadata,
5082 ToolPolicyDecisionRecord::unavailable("diagnostics provider is unavailable"),
5083 Some(ToolApprovalRecord {
5084 status: ToolApprovalStatus::Unavailable,
5085 reason: Some("diagnostics provider is unavailable".to_string()),
5086 modified_arguments: None,
5087 }),
5088 false,
5089 false,
5090 );
5091 self.finish_tool_record(&record).await;
5092 return Ok(record);
5093 }
5094
5095 if canonical_id == "web_search" && !self.tools.web_search_available() {
5096 let record = self.record_from_parts(
5097 &request,
5098 canonical_id.clone(),
5099 executed_arguments.clone(),
5100 started_at,
5101 start,
5102 false,
5103 false,
5104 "Web search provider is unavailable".to_string(),
5105 metadata,
5106 ToolPolicyDecisionRecord::unavailable("web search provider is unavailable"),
5107 Some(ToolApprovalRecord {
5108 status: ToolApprovalStatus::Unavailable,
5109 reason: Some("web search provider is unavailable".to_string()),
5110 modified_arguments: None,
5111 }),
5112 false,
5113 false,
5114 );
5115 self.finish_tool_record(&record).await;
5116 return Ok(record);
5117 }
5118
5119 let hitl_lang_ctx = self.build_hitl_language_context();
5120 if let Some(ref hitl_engine) = self.hitl_engine {
5121 let check_result = self
5122 .observe_purpose(
5123 ObservationPurpose::HitlLocalization,
5124 hitl_engine.check_tool_with_localization(
5125 &canonical_id,
5126 &executed_arguments,
5127 &hitl_lang_ctx,
5128 self.approval_handler.as_ref(),
5129 Some(&self.llm_registry),
5130 ),
5131 )
5132 .await?;
5133 if check_result.is_required() {
5134 match self.request_hitl_approval(check_result).await? {
5135 ApprovalResult::Approved => {
5136 merge_approved_record(&mut approval_record);
5137 }
5138 ApprovalResult::Modified { changes } => {
5139 if let Some(obj) = executed_arguments.as_object_mut() {
5140 for (key, value) in changes {
5141 obj.insert(key, value);
5142 }
5143 }
5144 let modified_security = security_engine
5145 .validate_tool_execution_with_bindings(
5146 &canonical_id,
5147 &executed_arguments,
5148 &bindings,
5149 )
5150 .await?;
5151 if !matches!(
5152 modified_security,
5153 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5154 ) {
5155 let reason = modified_security
5156 .reason()
5157 .unwrap_or("modified arguments failed policy")
5158 .to_string();
5159 let record = self.record_from_parts(
5160 &request,
5161 canonical_id,
5162 executed_arguments.clone(),
5163 started_at,
5164 start,
5165 false,
5166 false,
5167 reason.clone(),
5168 metadata,
5169 ToolPolicyDecisionRecord::deny(reason),
5170 Some(ToolApprovalRecord {
5171 status: ToolApprovalStatus::Modified,
5172 reason: None,
5173 modified_arguments: Some(executed_arguments),
5174 }),
5175 false,
5176 false,
5177 );
5178 self.finish_tool_record(&record).await;
5179 return Ok(record);
5180 }
5181 approval_record = Some(ToolApprovalRecord {
5182 status: ToolApprovalStatus::Modified,
5183 reason: None,
5184 modified_arguments: Some(executed_arguments.clone()),
5185 });
5186 }
5187 ApprovalResult::Rejected { reason } => {
5188 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5189 let record = self.record_from_parts(
5190 &request,
5191 canonical_id,
5192 executed_arguments,
5193 started_at,
5194 start,
5195 false,
5196 false,
5197 format!("Approval rejected: {}", reason),
5198 metadata,
5199 ToolPolicyDecisionRecord::approval(reason.clone()),
5200 Some(ToolApprovalRecord {
5201 status: ToolApprovalStatus::Rejected,
5202 reason: Some(reason),
5203 modified_arguments: None,
5204 }),
5205 false,
5206 false,
5207 );
5208 self.finish_tool_record(&record).await;
5209 return Ok(record);
5210 }
5211 ApprovalResult::Timeout => {
5212 let record = self.record_from_parts(
5213 &request,
5214 canonical_id,
5215 executed_arguments,
5216 started_at,
5217 start,
5218 false,
5219 false,
5220 "Approval timed out".to_string(),
5221 metadata,
5222 ToolPolicyDecisionRecord::approval("approval timeout"),
5223 Some(ToolApprovalRecord {
5224 status: ToolApprovalStatus::Timeout,
5225 reason: Some("approval timeout".to_string()),
5226 modified_arguments: None,
5227 }),
5228 false,
5229 false,
5230 );
5231 self.finish_tool_record(&record).await;
5232 return Ok(record);
5233 }
5234 }
5235 }
5236
5237 let condition_check = self
5238 .observe_purpose(
5239 ObservationPurpose::HitlLocalization,
5240 hitl_engine.check_conditions_with_localization(
5241 &executed_arguments,
5242 &hitl_lang_ctx,
5243 self.approval_handler.as_ref(),
5244 Some(&self.llm_registry),
5245 ),
5246 )
5247 .await?;
5248 if condition_check.is_required() {
5249 match self.request_hitl_approval(condition_check).await? {
5250 ApprovalResult::Approved => {
5251 merge_approved_record(&mut approval_record);
5252 }
5253 ApprovalResult::Modified { changes } => {
5254 if let Some(obj) = executed_arguments.as_object_mut() {
5255 for (key, value) in changes {
5256 obj.insert(key, value);
5257 }
5258 }
5259 let modified_security = security_engine
5260 .validate_tool_execution_with_bindings(
5261 &canonical_id,
5262 &executed_arguments,
5263 &bindings,
5264 )
5265 .await?;
5266 if !matches!(
5267 modified_security,
5268 SecurityCheckResult::Allow | SecurityCheckResult::Warn { .. }
5269 ) {
5270 let reason = modified_security
5271 .reason()
5272 .unwrap_or("modified arguments failed policy")
5273 .to_string();
5274 let record = self.record_from_parts(
5275 &request,
5276 canonical_id,
5277 executed_arguments,
5278 started_at,
5279 start,
5280 false,
5281 false,
5282 reason.clone(),
5283 metadata,
5284 ToolPolicyDecisionRecord::deny(reason),
5285 approval_record,
5286 false,
5287 false,
5288 );
5289 self.finish_tool_record(&record).await;
5290 return Ok(record);
5291 }
5292 approval_record = Some(ToolApprovalRecord {
5293 status: ToolApprovalStatus::Modified,
5294 reason: None,
5295 modified_arguments: Some(executed_arguments.clone()),
5296 });
5297 }
5298 ApprovalResult::Rejected { reason } => {
5299 let reason = reason.unwrap_or_else(|| "rejected".to_string());
5300 let record = self.record_from_parts(
5301 &request,
5302 canonical_id,
5303 executed_arguments,
5304 started_at,
5305 start,
5306 false,
5307 false,
5308 format!("Approval rejected: {}", reason),
5309 metadata,
5310 ToolPolicyDecisionRecord::approval(reason.clone()),
5311 Some(ToolApprovalRecord {
5312 status: ToolApprovalStatus::Rejected,
5313 reason: Some(reason),
5314 modified_arguments: None,
5315 }),
5316 false,
5317 false,
5318 );
5319 self.finish_tool_record(&record).await;
5320 return Ok(record);
5321 }
5322 ApprovalResult::Timeout => {
5323 let record = self.record_from_parts(
5324 &request,
5325 canonical_id,
5326 executed_arguments,
5327 started_at,
5328 start,
5329 false,
5330 false,
5331 "Approval timed out".to_string(),
5332 metadata,
5333 ToolPolicyDecisionRecord::approval("approval timeout"),
5334 Some(ToolApprovalRecord {
5335 status: ToolApprovalStatus::Timeout,
5336 reason: Some("approval timeout".to_string()),
5337 modified_arguments: None,
5338 }),
5339 false,
5340 false,
5341 );
5342 self.finish_tool_record(&record).await;
5343 return Ok(record);
5344 }
5345 }
5346 }
5347 }
5348
5349 executed_arguments = security_engine.prepare_tool_arguments_with_bindings(
5354 &canonical_id,
5355 &executed_arguments,
5356 &bindings,
5357 );
5358 if let Some(record) = approval_record.as_mut()
5359 && matches!(record.status, ToolApprovalStatus::Modified)
5360 {
5361 record.modified_arguments = Some(executed_arguments.clone());
5362 }
5363 let binding_security_result = security_engine
5364 .validate_tool_execution_with_bindings(&canonical_id, &executed_arguments, &bindings)
5365 .await?;
5366 let approval_confirmation_required = matches!(
5367 binding_security_result,
5368 SecurityCheckResult::RequireConfirmation { .. }
5369 ) || security_engine
5370 .classification_approval_message(
5371 &canonical_id,
5372 &resolved.tool.classify_call(&executed_arguments),
5373 )
5374 .is_some();
5375 let approval_binding = approval_record.as_ref().and_then(|record| {
5376 matches!(
5377 record.status,
5378 ToolApprovalStatus::Approved | ToolApprovalStatus::Modified
5379 )
5380 .then(|| ToolApprovalBinding {
5381 canonical_id: canonical_id.clone(),
5382 arguments: executed_arguments.clone(),
5383 confirmation_required: approval_confirmation_required,
5384 policy_version: security_engine.policy_version(),
5385 runtime_control_version: approval_control_snapshot.version,
5386 state_generation: initial_scope_snapshot.state_generation,
5387 reviewed_tool: Arc::clone(&resolved.tool),
5388 })
5389 });
5390
5391 let control_snapshot = self.runtime_safety_snapshot();
5396 let resolved = self.tools.resolve(&request.requested_name);
5397 let registry_version = self.tools.version();
5398 let mut versions = ToolDecisionVersions {
5399 policy: control_snapshot.tool_security.policy_version(),
5400 registry: registry_version,
5401 runtime_control: control_snapshot.version,
5402 state: None,
5403 };
5404 metadata.insert(
5405 "runtime_scope_snapshot".to_string(),
5406 serde_json::to_value(&control_snapshot.tool_scope_override).unwrap_or(Value::Null),
5407 );
5408 let resolved = match resolved {
5409 Some(resolved) => resolved,
5410 None => {
5411 let reason = format!(
5412 "Tool '{}' became unavailable after approval",
5413 request.requested_name
5414 );
5415 let record = self.record_from_parts_at(
5416 &request,
5417 request.requested_name.clone(),
5418 executed_arguments,
5419 started_at,
5420 start,
5421 false,
5422 false,
5423 reason.clone(),
5424 metadata,
5425 ToolPolicyDecisionRecord::unavailable(reason),
5426 approval_record,
5427 false,
5428 false,
5429 versions,
5430 );
5431 self.finish_tool_record(&record).await;
5432 return Ok(record);
5433 }
5434 };
5435
5436 let canonical_id = resolved.identity.canonical_id.clone();
5437 let bindings = resolved.tool.policy_bindings();
5438 let final_arguments = control_snapshot
5439 .tool_security
5440 .prepare_tool_arguments_with_bindings(&canonical_id, &executed_arguments, &bindings);
5441 if let Some(record) = approval_record.as_mut()
5442 && matches!(record.status, ToolApprovalStatus::Modified)
5443 {
5444 record.modified_arguments = Some(final_arguments.clone());
5445 }
5446 let classification = resolved.tool.classify_call(&final_arguments);
5447 let safety = resolved.tool.safety_metadata();
5448 let security_engine = control_snapshot.tool_security;
5449 let limits = security_engine.effective_limits(&canonical_id, &safety, &classification);
5450 let policy_snapshot = security_engine.policy_snapshot(&canonical_id);
5451 let resource_lock_keys =
5452 tool_resource_lock_keys(&canonical_id, &final_arguments, &bindings, &classification);
5453 metadata.insert(
5454 "classification".to_string(),
5455 serde_json::to_value(&classification).unwrap_or(Value::Null),
5456 );
5457 metadata.insert(
5458 "effective_limits".to_string(),
5459 serde_json::to_value(&limits).unwrap_or(Value::Null),
5460 );
5461 metadata.insert(
5462 "resource_lock_keys".to_string(),
5463 serde_json::to_value(&resource_lock_keys).unwrap_or(Value::Null),
5464 );
5465 if policy_snapshot.is_null() {
5466 metadata.remove("policy_snapshot");
5467 } else {
5468 metadata.insert("policy_snapshot".to_string(), policy_snapshot.clone());
5469 }
5470
5471 let final_denial = |canonical_id: String,
5472 output: String,
5473 policy: ToolPolicyDecisionRecord,
5474 metadata: HashMap<String, Value>,
5475 decision_versions: ToolDecisionVersions| {
5476 self.record_from_parts_at(
5477 &request,
5478 canonical_id,
5479 final_arguments.clone(),
5480 started_at,
5481 start,
5482 false,
5483 false,
5484 output,
5485 metadata,
5486 policy,
5487 approval_record.clone(),
5488 false,
5489 false,
5490 decision_versions,
5491 )
5492 };
5493
5494 if control_snapshot.emergency_deny {
5495 let reason = "Tool execution is disabled by runtime control".to_string();
5496 let record = final_denial(
5497 canonical_id,
5498 reason.clone(),
5499 ToolPolicyDecisionRecord::deny(reason),
5500 metadata,
5501 versions,
5502 );
5503 self.finish_tool_record(&record).await;
5504 return Ok(record);
5505 }
5506
5507 let available_snapshot = self
5512 .get_available_tool_ids_snapshot_for_scope(
5513 control_snapshot.tool_scope_override.as_deref(),
5514 )
5515 .await?;
5516 versions.state = available_snapshot.state_generation;
5517 metadata.insert(
5518 "available_tool_ids_snapshot".to_string(),
5519 serde_json::to_value(&available_snapshot.tool_ids).unwrap_or(Value::Null),
5520 );
5521 metadata.insert(
5522 "state_generation_snapshot".to_string(),
5523 serde_json::to_value(available_snapshot.state_generation).unwrap_or(Value::Null),
5524 );
5525 if !available_snapshot
5526 .tool_ids
5527 .iter()
5528 .any(|tool_id| tool_id == &canonical_id)
5529 {
5530 let reason = format!(
5531 "Tool '{}' is not available in the final runtime scope",
5532 canonical_id
5533 );
5534 let record = final_denial(
5535 canonical_id,
5536 reason.clone(),
5537 ToolPolicyDecisionRecord::deny(reason),
5538 metadata,
5539 versions,
5540 );
5541 self.finish_tool_record(&record).await;
5542 return Ok(record);
5543 }
5544
5545 let final_security_result = security_engine
5550 .validate_tool_execution_with_bindings(&canonical_id, &final_arguments, &bindings)
5551 .await?;
5552 match &final_security_result {
5553 SecurityCheckResult::Block { reason } => {
5554 let record = final_denial(
5555 canonical_id,
5556 format!("Denied: {}", reason),
5557 ToolPolicyDecisionRecord::deny(reason.clone()),
5558 metadata,
5559 versions,
5560 );
5561 self.finish_tool_record(&record).await;
5562 return Ok(record);
5563 }
5564 SecurityCheckResult::Unavailable { reason } => {
5565 let record = final_denial(
5566 canonical_id,
5567 format!("Unavailable: {}", reason),
5568 ToolPolicyDecisionRecord::unavailable(reason.clone()),
5569 metadata,
5570 versions,
5571 );
5572 self.finish_tool_record(&record).await;
5573 return Ok(record);
5574 }
5575 SecurityCheckResult::Warn { message } => {
5576 warn!(tool = %canonical_id, message = %message, "Tool security warning after approval");
5577 }
5578 SecurityCheckResult::Allow | SecurityCheckResult::RequireConfirmation { .. } => {}
5579 }
5580 let final_confirmation_required = matches!(
5581 final_security_result,
5582 SecurityCheckResult::RequireConfirmation { .. }
5583 ) || security_engine
5584 .classification_approval_message(&canonical_id, &classification)
5585 .is_some();
5586 let stale_approval = approval_binding.as_ref().is_some_and(|binding| {
5587 binding.is_stale(
5588 &canonical_id,
5589 &final_arguments,
5590 final_confirmation_required,
5591 versions,
5592 &resolved.tool,
5593 )
5594 });
5595 if stale_approval {
5596 let reason = "Approval became stale before final admission".to_string();
5597 let record = final_denial(
5598 canonical_id,
5599 reason.clone(),
5600 ToolPolicyDecisionRecord::deny(reason),
5601 metadata,
5602 versions,
5603 );
5604 self.finish_tool_record(&record).await;
5605 return Ok(record);
5606 }
5607 if final_confirmation_required && approval_binding.is_none() {
5608 let reason = "Final policy requires fresh approval".to_string();
5609 let record = final_denial(
5610 canonical_id,
5611 reason.clone(),
5612 ToolPolicyDecisionRecord::approval(reason),
5613 metadata,
5614 versions,
5615 );
5616 self.finish_tool_record(&record).await;
5617 return Ok(record);
5618 }
5619
5620 let unavailable_reason = match canonical_id.as_str() {
5621 "command" if !self.tools.command_runner_available() => {
5622 Some("command runner is unavailable")
5623 }
5624 "diagnostics" if !self.tools.diagnostics_available() => {
5625 Some("diagnostics provider is unavailable")
5626 }
5627 "web_search" if !self.tools.web_search_available() => {
5628 Some("web search provider is unavailable")
5629 }
5630 _ => None,
5631 };
5632 if let Some(reason) = unavailable_reason {
5633 let record = final_denial(
5634 canonical_id,
5635 reason.to_string(),
5636 ToolPolicyDecisionRecord::unavailable(reason),
5637 metadata,
5638 versions,
5639 );
5640 self.finish_tool_record(&record).await;
5641 return Ok(record);
5642 }
5643
5644 let Some(resource_guards) = self.acquire_tool_resource_locks(&resource_lock_keys).await
5649 else {
5650 let reason = "Tool execution cancelled while waiting for resource locks".to_string();
5651 let record = final_denial(
5652 canonical_id,
5653 reason.clone(),
5654 ToolPolicyDecisionRecord::deny(reason),
5655 metadata,
5656 versions,
5657 );
5658 self.finish_tool_record(&record).await;
5659 return Ok(record);
5660 };
5661
5662 let admission = self.admit_tool_execution(
5667 versions.runtime_control,
5668 versions.policy,
5669 versions.state,
5670 &canonical_id,
5671 );
5672 if !matches!(admission, SecurityCheckResult::Allow) {
5673 let latest_control = self.runtime_safety_snapshot();
5674 let reason = admission
5675 .reason()
5676 .unwrap_or("tool admission was denied")
5677 .to_string();
5678 let policy = if admission.is_unavailable() {
5679 ToolPolicyDecisionRecord::unavailable(reason.clone())
5680 } else {
5681 ToolPolicyDecisionRecord::deny(reason.clone())
5682 };
5683 let record = self.record_from_parts_at(
5684 &request,
5685 canonical_id,
5686 final_arguments,
5687 started_at,
5688 start,
5689 false,
5690 false,
5691 reason,
5692 metadata,
5693 policy,
5694 approval_record,
5695 false,
5696 false,
5697 ToolDecisionVersions {
5698 policy: latest_control.tool_security.policy_version(),
5699 registry: versions.registry,
5700 runtime_control: latest_control.version,
5701 state: self
5702 .state_machine
5703 .as_ref()
5704 .map(|state_machine| state_machine.generation()),
5705 },
5706 );
5707 self.finish_tool_record_after_resource_guards(resource_guards, &record)
5708 .await;
5709 return Ok(record);
5710 }
5711 let executed_arguments = final_arguments;
5712
5713 let tool_config = self.recovery_manager.get_tool_config(&canonical_id);
5714 let timeout_ms = limits
5715 .timeout_ms
5716 .unwrap_or_else(|| security_engine.get_tool_timeout(&canonical_id));
5717 let deadline = Some(started_at + chrono::Duration::milliseconds(timeout_ms as i64));
5718 let turn_actor = current_turn_actor_context();
5719 let actor = ToolActorContext {
5720 actor_id: turn_actor
5721 .as_ref()
5722 .and_then(|context| context.effective_actor_id().map(str::to_string))
5723 .or_else(|| self.actor_id()),
5724 origin_actor_id: turn_actor
5725 .as_ref()
5726 .and_then(|context| context.origin_actor_id.clone()),
5727 sender_agent_id: turn_actor
5728 .as_ref()
5729 .and_then(|context| context.sender_agent_id.clone()),
5730 };
5731 let tool_context = ToolExecutionContext {
5732 requested_name: request.requested_name.clone(),
5733 canonical_id: canonical_id.clone(),
5734 display_name: resolved.identity.display_name.clone(),
5735 provider_id: resolved.identity.provider_id.clone(),
5736 registry_version: versions.registry,
5737 policy_version: versions.policy,
5738 runtime_control_version: versions.runtime_control,
5739 call_id: request.call_id.clone(),
5740 source: request.source.clone(),
5741 actor,
5742 cancellation: ToolCancellationToken::new(
5743 Arc::clone(&self.runtime_control.emergency_deny),
5744 Some("runtime control cancellation".to_string()),
5745 ),
5746 started_at,
5747 deadline,
5748 permission: ToolPolicyDecisionRecord::allow(),
5749 approval: approval_record.clone(),
5750 classification: classification.clone(),
5751 safety,
5752 limits: limits.clone(),
5753 policy_snapshot,
5754 custom_config: security_engine.custom_config(&canonical_id),
5755 };
5756 let (mut result, timed_out, cancelled, invoked) = self
5757 .run_tool_with_retries(
5758 &canonical_id,
5759 resolved.tool.clone(),
5760 executed_arguments.clone(),
5761 tool_context,
5762 timeout_ms,
5763 tool_config.max_retries,
5764 )
5765 .await?;
5766
5767 if !result.success {
5768 match &tool_config.on_failure {
5769 ToolFailureAction::Skip => {
5770 result = ToolResult::ok(format!(
5771 "{{\"skipped\": true, \"reason\": \"Tool '{}' was skipped after failure\"}}",
5772 canonical_id
5773 ));
5774 }
5775 ToolFailureAction::Fallback { fallback_tool } => {
5776 drop(resource_guards);
5777 let fallback_request = ToolExecutionRequest::new(
5778 request.call_id.clone(),
5779 fallback_tool.clone(),
5780 executed_arguments,
5781 ToolCallSource::Fallback {
5782 original_tool: canonical_id,
5783 },
5784 );
5785 return Box::pin(self.execute_tool_record(fallback_request)).await;
5786 }
5787 ToolFailureAction::ReportError => {}
5788 }
5789 }
5790
5791 let output_cap = limits.max_output_chars;
5792 let (output, output_truncated) =
5793 Self::truncate_tool_output(result.output.clone(), output_cap);
5794 if let Some(result_metadata) = result.metadata {
5795 metadata.extend(result_metadata);
5796 }
5797 let mut record = self.record_from_parts_at(
5798 &request,
5799 canonical_id,
5800 executed_arguments,
5801 started_at,
5802 start,
5803 invoked,
5804 result.success,
5805 output,
5806 metadata,
5807 ToolPolicyDecisionRecord::allow(),
5808 approval_record,
5809 timed_out,
5810 output_truncated,
5811 versions,
5812 );
5813 record.cancelled = cancelled;
5814 if cancelled {
5815 record.cancellation_reason = Some("runtime control cancellation".to_string());
5816 }
5817 self.finish_tool_record_after_resource_guards(resource_guards, &record)
5818 .await;
5819 Ok(record)
5820 }
5821
5822 #[instrument(skip(self, tool_call), fields(tool = %tool_call.name))]
5823 async fn execute_tool_smart(&self, tool_call: &ToolCall) -> Result<String> {
5824 let record = self
5825 .execute_tool_record(ToolExecutionRequest::new(
5826 tool_call.id.clone(),
5827 tool_call.name.clone(),
5828 tool_call.arguments.clone(),
5829 ToolCallSource::Model,
5830 ))
5831 .await?;
5832 if record.success {
5833 Ok(record.model_output_string())
5834 } else if matches!(record.policy.outcome, PermissionOutcome::RequiresApproval) {
5835 Err(AgentError::HITLRejected(record.model_output_string()))
5836 } else {
5837 Err(AgentError::Tool(record.model_output_string()))
5838 }
5839 }
5840
5841 async fn select_skill_candidate(&self, input: &str) -> Result<Option<SkillCandidate>> {
5847 let Some(ref router) = self.skill_router else {
5848 return Ok(None);
5849 };
5850 let available_skills = self.get_available_skills();
5851 if available_skills.is_empty() {
5852 return Ok(None);
5853 }
5854 let skill_ids: Vec<&str> = available_skills.iter().map(|s| s.id.as_str()).collect();
5855 let Some(skill_id) = self
5856 .observe_purpose(
5857 ObservationPurpose::SkillRouting,
5858 router.select_skill_filtered(input, &skill_ids),
5859 )
5860 .await?
5861 else {
5862 return Ok(None);
5863 };
5864 let skill = router
5865 .get_skill(&skill_id)
5866 .cloned()
5867 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
5868 info!(skill_id = %skill_id, "Skill selected");
5869 Ok(Some(SkillCandidate::new(skill_id, skill)))
5870 }
5871
5872 async fn commit_skill_candidate_route_result(
5877 &self,
5878 candidate: SkillCandidate,
5879 input: &str,
5880 ) -> Result<SkillRouteResult> {
5881 let skill_id = candidate.skill_id;
5882 let skill = candidate.skill;
5883 let expected_state_generation = self
5884 .state_machine
5885 .as_ref()
5886 .map(|state_machine| state_machine.generation());
5887 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
5888 if let Some(ref skill_disambig) = skill.disambiguation
5889 && skill_disambig.enabled.unwrap_or(false)
5890 && let Some(ref disambiguator) = self.disambiguation_manager
5891 {
5892 let context = self.build_disambiguation_context().await?;
5893 let state_override = self
5894 .state_machine
5895 .as_ref()
5896 .and_then(|sm| sm.current_definition())
5897 .and_then(|def| def.disambiguation.clone());
5898
5899 let disambiguation_result = self
5900 .observe_purpose(
5901 ObservationPurpose::DisambiguationDetection,
5902 disambiguator.process_input_with_override(
5903 input,
5904 &context,
5905 state_override.as_ref(),
5906 Some(skill_disambig),
5907 ),
5908 )
5909 .await?;
5910 let current_state_generation = self
5911 .state_machine
5912 .as_ref()
5913 .map(|state_machine| state_machine.generation());
5914 if current_state_generation != expected_state_generation
5915 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
5916 {
5917 disambiguator.clear_pending().await;
5918 *self.pending_skill_id.write() = None;
5919 return Err(AgentError::Other(
5920 "State or reset ownership changed during skill disambiguation".to_string(),
5921 ));
5922 }
5923 match disambiguation_result {
5924 DisambiguationResult::Clear => {
5925 debug!(skill_id = %skill_id, "Skill disambiguation: clear");
5926 }
5927 DisambiguationResult::NeedsClarification {
5928 question,
5929 detection,
5930 } => {
5931 let admission = self
5932 .admit_disambiguation_redispatch(
5933 expected_disambiguation_epoch,
5934 expected_state_generation,
5935 )
5936 .await?;
5937 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
5938 info!(
5939 skill_id = %skill_id,
5940 ambiguity_type = ?detection.ambiguity_type,
5941 confidence = detection.confidence,
5942 "Skill requires clarification before execution"
5943 );
5944 *self.pending_skill_id.write() = Some(skill_id.clone());
5945 let response = AgentResponse::new(&question.question).with_metadata(
5946 "disambiguation",
5947 serde_json::json!({
5948 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
5949 "skill_id": skill_id,
5950 "options": question.options,
5951 "clarifying": question.clarifying,
5952 "detection": {
5953 "type": detection.ambiguity_type,
5954 "confidence": detection.confidence,
5955 "what_is_unclear": detection.what_is_unclear,
5956 }
5957 }),
5958 );
5959 drop(admission);
5960 return Ok(SkillRouteResult::NeedsClarification {
5961 response,
5962 ownership: Some(DisambiguationOwnership {
5963 epoch: expected_disambiguation_epoch,
5964 state_generation: expected_state_generation,
5965 }),
5966 });
5967 }
5968 DisambiguationResult::Clarified { enriched_input, .. } => {
5969 info!(skill_id = %skill_id, enriched = %enriched_input, "Skill disambiguation clarified");
5970 let admission = self
5971 .admit_disambiguation_redispatch(
5972 expected_disambiguation_epoch,
5973 expected_state_generation,
5974 )
5975 .await?;
5976 drop(admission);
5977 let content = self.execute_skill(&skill, &enriched_input).await?;
5978 return Ok(SkillRouteResult::Response { skill_id, content });
5979 }
5980 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
5981 info!(skill_id = %skill_id, "Skill disambiguation best guess");
5982 let admission = self
5983 .admit_disambiguation_redispatch(
5984 expected_disambiguation_epoch,
5985 expected_state_generation,
5986 )
5987 .await?;
5988 drop(admission);
5989 let content = self.execute_skill(&skill, &enriched_input).await?;
5990 return Ok(SkillRouteResult::Response { skill_id, content });
5991 }
5992 DisambiguationResult::GiveUp { reason } => {
5993 warn!(skill_id = %skill_id, reason = %reason, "Skill disambiguation gave up");
5994 let apology = self
5995 .generate_localized_apology(
5996 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
5997 &reason,
5998 )
5999 .await
6000 .unwrap_or_else(|_| {
6001 format!("I'm sorry, I couldn't understand your request: {}", reason)
6002 });
6003 return Ok(SkillRouteResult::NeedsClarification {
6004 response: AgentResponse::new(&apology),
6005 ownership: None,
6006 });
6007 }
6008 DisambiguationResult::Escalate { reason } => {
6009 info!(skill_id = %skill_id, reason = %reason, "Skill disambiguation escalating");
6010 let apology = self
6011 .generate_localized_apology(
6012 "Explain briefly that you're transferring the user to a human agent for help.",
6013 &reason,
6014 )
6015 .await
6016 .unwrap_or_else(|_| {
6017 format!("I need human assistance to help with your request: {}", reason)
6018 });
6019 return Ok(SkillRouteResult::NeedsClarification {
6020 response: AgentResponse::new(&apology),
6021 ownership: None,
6022 });
6023 }
6024 DisambiguationResult::Abandoned { .. } => {
6025 debug!(skill_id = %skill_id, "Skill disambiguation abandoned");
6026 return Ok(SkillRouteResult::NoMatch);
6027 }
6028 }
6029 }
6030 let admission = self
6031 .admit_disambiguation_redispatch(
6032 expected_disambiguation_epoch,
6033 expected_state_generation,
6034 )
6035 .await?;
6036 drop(admission);
6037 let content = self.execute_skill(&skill, input).await?;
6038 Ok(SkillRouteResult::Response { skill_id, content })
6039 }
6040
6041 async fn try_skill_route(&self, input: &str) -> Result<SkillRouteResult> {
6043 if let Some(candidate) = self.select_skill_candidate(input).await? {
6044 self.commit_skill_candidate_route_result(candidate, input)
6045 .await
6046 } else {
6047 Ok(SkillRouteResult::NoMatch)
6048 }
6049 }
6050
6051 async fn execute_skill(&self, skill: &SkillDefinition, input: &str) -> Result<String> {
6053 if let Some(ref executor) = self.skill_executor {
6054 let skill_reasoning = self.get_skill_reasoning_config(skill);
6055 let skill_reflection = self.get_skill_reflection_config(skill);
6056
6057 debug!(
6058 skill_id = %skill.id,
6059 reasoning_mode = ?skill_reasoning.mode,
6060 reflection_enabled = ?skill_reflection.enabled,
6061 "Skill reasoning/reflection config"
6062 );
6063
6064 let response = self
6065 .observe_purpose(
6066 ObservationPurpose::SkillPrompt,
6067 executor.execute_with_invoker(skill, input, serde_json::json!({}), self),
6068 )
6069 .await?;
6070
6071 if skill_reflection.requires_evaluation() && skill_reflection.is_enabled() {
6072 let should_reflect = self
6073 .should_reflect_with_config(input, &response, &skill_reflection)
6074 .await?;
6075 if should_reflect {
6076 let evaluated = self
6077 .evaluate_and_retry_with_config(input, response, &skill_reflection)
6078 .await?;
6079 return Ok(evaluated);
6080 }
6081 }
6082
6083 return Ok(response);
6084 }
6085 Err(AgentError::Skill(
6086 "No skill executor configured".to_string(),
6087 ))
6088 }
6089
6090 async fn execute_skill_by_id(&self, skill_id: &str, input: &str) -> Result<String> {
6093 let skill = self
6094 .skill_router
6095 .as_ref()
6096 .and_then(|r| r.get_skill(skill_id).cloned())
6097 .ok_or_else(|| AgentError::Skill(format!("Skill not found: {}", skill_id)))?;
6098 self.execute_skill(&skill, input).await
6099 }
6100
6101 async fn should_reflect_with_config(
6102 &self,
6103 input: &str,
6104 response: &str,
6105 config: &ReflectionConfig,
6106 ) -> Result<bool> {
6107 if !config.requires_evaluation() {
6108 return Ok(false);
6109 }
6110
6111 if config.is_enabled() {
6112 return Ok(true);
6113 }
6114
6115 let evaluator_llm = config
6116 .evaluator_llm
6117 .as_ref()
6118 .and_then(|alias| self.llm_registry.get(alias).ok())
6119 .or_else(|| self.llm_registry.router().ok())
6120 .or_else(|| self.llm_registry.default().ok());
6121
6122 let Some(llm) = evaluator_llm else {
6123 return Ok(false);
6124 };
6125
6126 let response_preview: String = response.chars().take(500).collect();
6127 let prompt = format!(
6128 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
6129
6130User query: "{}"
6131Response: "{}"
6132
6133Answer YES or NO only."#,
6134 input, response_preview
6135 );
6136
6137 let messages = vec![ChatMessage::user(&prompt)];
6138 let result = self
6139 .observe_purpose(
6140 ObservationPurpose::ReflectionDecision,
6141 llm.complete(&messages, None),
6142 )
6143 .await;
6144
6145 match result {
6146 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
6147 Err(_) => Ok(false),
6148 }
6149 }
6150
6151 async fn evaluate_and_retry_with_config(
6152 &self,
6153 input: &str,
6154 mut response: String,
6155 config: &ReflectionConfig,
6156 ) -> Result<String> {
6157 let llm = self.get_state_llm()?;
6158 let mut attempts = 0u32;
6159 let max_retries = config.max_retries;
6160
6161 loop {
6162 let evaluation = self
6163 .evaluate_response_with_config(input, &response, config)
6164 .await?;
6165
6166 if evaluation.passed || attempts >= max_retries {
6167 info!(
6168 passed = evaluation.passed,
6169 confidence = evaluation.confidence,
6170 attempts = attempts + 1,
6171 "Skill reflection evaluation complete"
6172 );
6173 return Ok(response);
6174 }
6175
6176 debug!(
6177 attempt = attempts + 1,
6178 failed_criteria = evaluation.failed_criteria().count(),
6179 "Skill response did not meet criteria, retrying"
6180 );
6181
6182 let feedback: Vec<String> = evaluation
6183 .failed_criteria()
6184 .map(|c| format!("- {}", c.criterion))
6185 .collect();
6186
6187 let retry_prompt = format!(
6188 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response to: {}",
6189 feedback.join("\n"),
6190 input
6191 );
6192
6193 let messages = vec![ChatMessage::user(&retry_prompt)];
6194 let retry_response = self
6195 .observe_purpose(
6196 ObservationPurpose::ReflectionEvaluation,
6197 llm.complete(&messages, None),
6198 )
6199 .await
6200 .map_err(|e| AgentError::LLM(e.to_string()))?;
6201
6202 response = retry_response.content.trim().to_string();
6203 attempts += 1;
6204 }
6205 }
6206
6207 async fn evaluate_response_with_config(
6208 &self,
6209 input: &str,
6210 response: &str,
6211 config: &ReflectionConfig,
6212 ) -> Result<EvaluationResult> {
6213 let evaluator_llm = config
6214 .evaluator_llm
6215 .as_ref()
6216 .and_then(|alias| self.llm_registry.get(alias).ok())
6217 .or_else(|| self.llm_registry.router().ok())
6218 .or_else(|| self.llm_registry.default().ok())
6219 .ok_or_else(|| AgentError::Config("No LLM available for evaluation".into()))?;
6220
6221 let criteria = &config.criteria;
6222 let criteria_list = criteria
6223 .iter()
6224 .enumerate()
6225 .map(|(i, c)| format!("{}. {}", i + 1, c))
6226 .collect::<Vec<_>>()
6227 .join("\n");
6228
6229 let prompt = format!(
6230 r#"Evaluate this response against the criteria.
6231
6232User query: "{}"
6233
6234Response to evaluate: "{}"
6235
6236Criteria:
6237{}
6238
6239For each criterion, respond with:
6240- criterion number
6241- PASS or FAIL
6242- brief reason
6243
6244Then provide overall confidence (0.0 to 1.0) and whether it passes overall.
6245
6246Format:
62471. PASS/FAIL - reason
62482. PASS/FAIL - reason
6249...
6250CONFIDENCE: 0.X
6251OVERALL: PASS/FAIL"#,
6252 input, response, criteria_list
6253 );
6254
6255 let messages = vec![ChatMessage::user(&prompt)];
6256 let eval_response = self
6257 .observe_purpose(
6258 ObservationPurpose::ReflectionEvaluation,
6259 evaluator_llm.complete(&messages, None),
6260 )
6261 .await
6262 .map_err(|e| AgentError::LLM(format!("Evaluation failed: {}", e)))?;
6263
6264 let content = eval_response.content.to_uppercase();
6265 let llm_pass = content.contains("OVERALL: PASS");
6266
6267 let confidence = content
6268 .lines()
6269 .find(|l| l.contains("CONFIDENCE:"))
6270 .and_then(|l| {
6271 l.split(':')
6272 .nth(1)
6273 .and_then(|v| v.trim().parse::<f32>().ok())
6274 })
6275 .unwrap_or(if llm_pass { 0.8 } else { 0.4 });
6276
6277 let overall_pass = llm_pass && confidence >= config.pass_threshold;
6280
6281 let mut criteria_results = Vec::new();
6282 for (i, criterion) in criteria.iter().enumerate() {
6283 let line_marker = format!("{}.", i + 1);
6284 let passed = eval_response
6285 .content
6286 .lines()
6287 .find(|l| l.contains(&line_marker))
6288 .map(|l| l.to_uppercase().contains("PASS"))
6289 .unwrap_or(overall_pass);
6290
6291 if passed {
6292 criteria_results.push(CriterionResult::pass(criterion));
6293 } else {
6294 criteria_results.push(CriterionResult::fail(criterion, "Did not meet criterion"));
6295 }
6296 }
6297
6298 Ok(EvaluationResult::new(overall_pass, confidence).with_criteria(criteria_results))
6299 }
6300
6301 async fn process_input(&self, input: &str) -> Result<ProcessData> {
6303 if let Some(processor) = self.get_state_process_processor() {
6304 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6305 return self
6306 .observe_purpose(purpose, processor.process_input(input))
6307 .await;
6308 }
6309 if let Some(ref processor) = self.process_processor {
6310 let purpose = observation_purpose_for_process(processor.input_purpose_hint());
6311 self.observe_purpose(purpose, processor.process_input(input))
6312 .await
6313 } else {
6314 Ok(ProcessData::new(input))
6315 }
6316 }
6317
6318 async fn process_output(
6320 &self,
6321 output: &str,
6322 input_context: &std::collections::HashMap<String, serde_json::Value>,
6323 ) -> Result<ProcessData> {
6324 if let Some(processor) = self.get_state_process_processor() {
6325 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6326 return self
6327 .observe_purpose(purpose, processor.process_output(output, input_context))
6328 .await;
6329 }
6330 if let Some(ref processor) = self.process_processor {
6331 let purpose = observation_purpose_for_process(processor.output_purpose_hint());
6332 self.observe_purpose(purpose, processor.process_output(output, input_context))
6333 .await
6334 } else {
6335 Ok(ProcessData::new(output))
6336 }
6337 }
6338
6339 fn get_state_process_processor(&self) -> Option<ProcessProcessor> {
6341 let sm = self.state_machine.as_ref()?;
6342 let def = sm.current_definition()?;
6343 let config = def.process.as_ref()?;
6344 let mut processor = ProcessProcessor::new(config.clone());
6345 if let Some(ref registry) = Some(self.llm_registry.clone()) {
6346 processor = processor.with_llm_registry(registry.clone());
6347 }
6348 processor = processor.with_stage_observer(Arc::new(ObservabilityProcessStageObserver));
6349 Some(processor)
6350 }
6351
6352 async fn check_turn_timeout(&self) -> Result<()> {
6354 let Some(ref sm) = self.state_machine else {
6355 return Ok(());
6356 };
6357 let Some(timeout_state) = sm.check_timeout() else {
6358 return Ok(());
6359 };
6360 let claim_admission = self.disambiguation_admission.write().await;
6361 if sm.check_timeout().as_deref() != Some(timeout_state.as_str()) {
6362 return Ok(());
6363 }
6364 let Some(reservation) = self.reserve_state_transition() else {
6365 return Ok(());
6366 };
6367 let from_state = sm.current();
6368 let expected_state_generation = sm.generation();
6369 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6370 let history_before = sm.history();
6371 drop(claim_admission);
6372
6373 self.execute_state_exit_actions(&from_state).await;
6374
6375 let admission = self.disambiguation_admission.write().await;
6376 if sm.current() != from_state
6377 || sm.generation() != expected_state_generation
6378 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6379 || sm.check_timeout().as_deref() != Some(timeout_state.as_str())
6380 {
6381 return Ok(());
6382 }
6383 sm.transition_to(&timeout_state, "max_turns exceeded")?;
6384 self.invalidate_pending_confirmation("state_timeout").await;
6385 let entered = sm.current();
6386 let is_reentry = Self::state_was_previously_entered(&entered, &from_state, &history_before);
6387 drop(admission);
6388
6389 self.execute_state_enter_actions(&entered, is_reentry).await;
6390 drop(reservation);
6391 info!(to = %entered, "Timeout transition");
6392 Ok(())
6393 }
6394
6395 fn increment_turn(&self) {
6396 if let Some(ref sm) = self.state_machine {
6397 sm.increment_turn();
6398 }
6399 }
6400
6401 fn transitions_available_for_commit(&self) -> Option<(Vec<Transition>, String)> {
6402 let sm = self.state_machine.as_ref()?;
6403 let current = sm.current();
6404 let transitions: Vec<_> = sm
6405 .auto_transitions()
6406 .into_iter()
6407 .filter(|t| match t.cooldown_turns {
6408 Some(cd) if cd > 0 => {
6409 let resolved = sm.config().resolve_full_path(¤t, &t.to);
6410 !sm.is_on_cooldown(&resolved, cd)
6411 }
6412 _ => true,
6413 })
6414 .collect();
6415 Some((transitions, current))
6416 }
6417
6418 fn transition_reason(transition: &Transition) -> String {
6419 if transition.when.is_empty() {
6420 "guard condition met".to_string()
6421 } else {
6422 transition.when.clone()
6423 }
6424 }
6425
6426 fn build_transition_context(
6428 &self,
6429 user_message: &str,
6430 response: &str,
6431 current_state: &str,
6432 staged: Option<&HashMap<String, Value>>,
6433 ) -> TransitionContext {
6434 let context_map = staged
6435 .map(|writes| self.build_context_with_staged(writes))
6436 .unwrap_or_else(|| self.build_context_with_overlays());
6437 TransitionContext::new(user_message, response, current_state).with_context(context_map)
6438 }
6439
6440 async fn select_transition_candidate(
6442 &self,
6443 user_message: &str,
6444 response: &str,
6445 ) -> Result<Option<TransitionCandidate>> {
6446 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6447 return Ok(None);
6448 };
6449 let transitions: Vec<Transition> = transitions
6450 .into_iter()
6451 .filter(|transition| matches!(transition.timing, TransitionTiming::PostResponse))
6452 .collect();
6453 if transitions.is_empty() {
6454 return Ok(None);
6455 }
6456 let Some(evaluator) = self.transition_evaluator.as_ref() else {
6457 return Ok(None);
6458 };
6459 let context = self.build_transition_context(user_message, response, ¤t_state, None);
6460 let selected = self
6461 .observe_purpose(
6462 ObservationPurpose::StateTransitionEvaluation,
6463 evaluator.select_transition(&transitions, &context),
6464 )
6465 .await?;
6466 Ok(selected.map(|index| {
6467 let transition = transitions[index].clone();
6468 TransitionCandidate::new(
6469 current_state,
6470 transition.clone(),
6471 Self::transition_reason(&transition),
6472 )
6473 }))
6474 }
6475
6476 fn select_deterministic_transition_candidate(
6478 &self,
6479 user_message: &str,
6480 current_state: &str,
6481 transitions: &[Transition],
6482 staged: &HashMap<String, Value>,
6483 ) -> Option<TransitionCandidate> {
6484 let context = self.build_transition_context(user_message, "", current_state, Some(staged));
6485
6486 for transition in transitions {
6487 if let Some(guard) = transition.guard.as_ref()
6488 && evaluate_guard(guard, &context)
6489 {
6490 return Some(TransitionCandidate::new(
6491 current_state,
6492 transition.clone(),
6493 Self::transition_reason(transition),
6494 ));
6495 }
6496 }
6497
6498 let resolved_intent = context
6499 .context
6500 .get("resolved_intent")
6501 .and_then(Value::as_str)
6502 .filter(|value| !value.is_empty());
6503 if let Some(resolved_intent) = resolved_intent {
6504 for transition in transitions {
6505 if transition.intent.as_deref() == Some(resolved_intent) {
6506 return Some(TransitionCandidate::new(
6507 current_state,
6508 transition.clone(),
6509 Self::transition_reason(transition),
6510 ));
6511 }
6512 }
6513 }
6514
6515 None
6516 }
6517
6518 async fn commit_transition_candidate(&self, candidate: &TransitionCandidate) -> Result<bool> {
6520 self.commit_transition_target(&candidate.from_state, candidate.target(), &candidate.reason)
6521 .await
6522 }
6523
6524 async fn approve_transition_target(&self, from_state: &str, target: &str) -> Result<bool> {
6526 let approved = self.check_state_hitl(Some(from_state), target).await?;
6527 if !approved {
6528 info!(to = %target, "State transition rejected by HITL");
6529 }
6530 Ok(approved)
6531 }
6532
6533 async fn apply_transition_target(
6535 &self,
6536 from_state: &str,
6537 target: &str,
6538 reason: &str,
6539 staged: Option<&HashMap<String, Value>>,
6540 ) -> Result<bool> {
6541 let Some(ref sm) = self.state_machine else {
6542 return Ok(false);
6543 };
6544 let claim_admission = self.disambiguation_admission.write().await;
6545 if sm.current() != from_state {
6546 return Ok(false);
6547 }
6548 let Some(reservation) = self.reserve_state_transition() else {
6549 return Ok(false);
6550 };
6551 let expected_state_generation = sm.generation();
6552 let expected_disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
6553 let history_before = sm.history();
6554 drop(claim_admission);
6555
6556 self.execute_state_exit_actions(from_state).await;
6557
6558 let admission = self.disambiguation_admission.write().await;
6559 if sm.current() != from_state
6560 || sm.generation() != expected_state_generation
6561 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
6562 {
6563 return Ok(false);
6564 }
6565 sm.transition_to(target, reason)?;
6566 self.invalidate_pending_confirmation("state_transition")
6567 .await;
6568 sm.reset_no_transition();
6569 if let Some(staged) = staged {
6570 self.commit_staged_context_writes(staged);
6571 }
6572 let entered = sm.current();
6573 let is_reentry = Self::state_was_previously_entered(&entered, from_state, &history_before);
6574 drop(admission);
6575
6576 self.execute_state_enter_actions(&entered, is_reentry).await;
6577 drop(reservation);
6578 self.hooks
6579 .on_state_transition(Some(from_state), &entered, reason)
6580 .await;
6581 info!(from = %from_state, to = %entered, "State transition");
6582 Ok(true)
6583 }
6584
6585 async fn commit_transition_target(
6587 &self,
6588 from_state: &str,
6589 target: &str,
6590 reason: &str,
6591 ) -> Result<bool> {
6592 if !self.approve_transition_target(from_state, target).await? {
6593 return Ok(false);
6594 }
6595 self.apply_transition_target(from_state, target, reason, None)
6596 .await
6597 }
6598
6599 async fn apply_pre_response_transition_candidate(
6601 &self,
6602 candidate: &TransitionCandidate,
6603 staged: &HashMap<String, Value>,
6604 processed_input: &str,
6605 ) -> Result<bool> {
6606 self.commit_root_user_message(processed_input).await?;
6607 self.apply_transition_target(
6608 &candidate.from_state,
6609 candidate.target(),
6610 &candidate.reason,
6611 Some(staged),
6612 )
6613 .await
6614 }
6615
6616 async fn commit_pre_response_transition_candidate(
6618 &self,
6619 candidate: &TransitionCandidate,
6620 staged: &HashMap<String, Value>,
6621 processed_input: &str,
6622 ) -> Result<bool> {
6623 if !self
6624 .approve_transition_target(&candidate.from_state, candidate.target())
6625 .await?
6626 {
6627 return Ok(false);
6628 }
6629 self.apply_pre_response_transition_candidate(candidate, staged, processed_input)
6630 .await
6631 }
6632
6633 async fn handle_transition_miss(&self, current_state: &str) -> Result<bool> {
6635 let Some(ref sm) = self.state_machine else {
6636 return Ok(false);
6637 };
6638 sm.increment_no_transition();
6639 let Some(fallback) = sm.check_fallback() else {
6640 return Ok(false);
6641 };
6642 self.commit_transition_target(current_state, &fallback, "fallback after no transitions")
6643 .await
6644 }
6645
6646 async fn evaluate_transitions(&self, user_message: &str, response: &str) -> Result<bool> {
6648 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6649 return Ok(false);
6650 };
6651 if transitions.is_empty() {
6652 return Ok(false);
6653 }
6654 if let Some(candidate) = self
6655 .select_transition_candidate(user_message, response)
6656 .await?
6657 {
6658 return self.commit_transition_candidate(&candidate).await;
6659 }
6660 self.handle_transition_miss(¤t_state).await
6661 }
6662
6663 async fn try_pre_response_transition(
6665 &self,
6666 processed_input: &str,
6667 ) -> Result<Option<AgentResponse>> {
6668 let optimization = &self.runtime_config.optimization;
6669 if !optimization.enabled || !optimization.pre_response_deterministic_transitions {
6670 return Ok(None);
6671 }
6672 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
6673 return Ok(None);
6674 };
6675 let eligible: Vec<Transition> = transitions
6676 .into_iter()
6677 .filter(|transition| !transition.requires_response)
6678 .filter(|transition| matches!(transition.timing, TransitionTiming::PreResponse))
6679 .collect();
6680 if eligible.is_empty() {
6681 return Ok(None);
6682 }
6683
6684 let empty_staged = HashMap::new();
6685 let mut extracted_staged: Option<HashMap<String, Value>> = None;
6686 let mut selected: Option<(TransitionCandidate, HashMap<String, Value>)> = None;
6687
6688 for transition in &eligible {
6689 let use_extractors = optimization.pre_response_extractors || transition.run_extractors;
6690 let staged_for_eval = if use_extractors {
6691 if extracted_staged.is_none() {
6692 extracted_staged =
6693 Some(self.run_context_extractors_staged(processed_input).await);
6694 }
6695 extracted_staged.as_ref().unwrap_or(&empty_staged)
6696 } else {
6697 &empty_staged
6698 };
6699
6700 if let Some(candidate) = self.select_deterministic_transition_candidate(
6701 processed_input,
6702 ¤t_state,
6703 std::slice::from_ref(transition),
6704 staged_for_eval,
6705 ) {
6706 let staged_for_commit = if use_extractors {
6707 staged_for_eval.clone()
6708 } else {
6709 HashMap::new()
6710 };
6711 selected = Some((candidate, staged_for_commit));
6712 break;
6713 }
6714 }
6715
6716 let Some((candidate, staged)) = selected else {
6717 return Ok(None);
6718 };
6719
6720 if !self
6721 .commit_pre_response_transition_candidate(&candidate, &staged, processed_input)
6722 .await?
6723 {
6724 return Ok(None);
6725 }
6726 self.redispatch_current_state(processed_input)
6727 .await
6728 .map(Some)
6729 }
6730
6731 async fn try_speculative_branches(
6736 &self,
6737 processed_input: &str,
6738 input_context: &HashMap<String, Value>,
6739 ) -> Result<Option<AgentResponse>> {
6740 let optimization = &self.runtime_config.optimization;
6741 if !optimization.enabled {
6742 return Ok(None);
6743 }
6744
6745 let effective_reasoning_mode = self.get_effective_reasoning_config().mode.clone();
6746 if !matches!(
6747 effective_reasoning_mode,
6748 ReasoningMode::None | ReasoningMode::Auto
6749 ) {
6750 return Ok(None);
6751 }
6752
6753 let mut transition_enabled =
6754 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
6755 let mut skill_enabled = optimization.speculative_skill_routing
6756 && self.skill_router.is_some()
6757 && self.pending_skill_id.read().is_none();
6758 let mut reasoning_enabled = optimization.speculative_reasoning_auto
6759 && matches!(effective_reasoning_mode, ReasoningMode::Auto);
6760
6761 if matches!(effective_reasoning_mode, ReasoningMode::Auto)
6762 && (!reasoning_enabled || optimization.max_speculative_llm_calls_per_turn < 2)
6763 {
6764 return Ok(None);
6765 }
6766
6767 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6768 return Ok(None);
6769 }
6770
6771 let mut optional_slots = optimization.max_parallel_runtime_tasks.saturating_sub(1);
6772 let mut speculative_call_slots = optimization
6773 .max_speculative_llm_calls_per_turn
6774 .saturating_sub(1);
6775 if reasoning_enabled {
6776 if optional_slots == 0 || speculative_call_slots == 0 {
6777 return Ok(None);
6778 }
6779 optional_slots -= 1;
6780 speculative_call_slots -= 1;
6781 }
6782 if transition_enabled {
6783 if optional_slots == 0 {
6784 transition_enabled = false;
6785 } else {
6786 optional_slots -= 1;
6787 }
6788 }
6789 if skill_enabled && (optional_slots == 0 || speculative_call_slots == 0) {
6790 skill_enabled = false;
6791 }
6792
6793 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6794 return Ok(None);
6795 }
6796
6797 let main_kind = if transition_enabled {
6798 RuntimeOptimizationKind::ParallelStateTransition
6799 } else if skill_enabled {
6800 RuntimeOptimizationKind::SpeculativeSkillRouting
6801 } else {
6802 RuntimeOptimizationKind::SpeculativeReasoningAuto
6803 };
6804 if !self.reserve_active_speculative_llm_call(main_kind) {
6805 return Ok(None);
6806 }
6807
6808 let mut branch_set = ScheduledBranchSet::new(optimization.max_parallel_runtime_tasks)?;
6809 let main_branch = RuntimeBranch::new(
6810 RuntimeTaskPurpose::MainResponse,
6811 main_kind,
6812 RuntimeTaskPriority::Normal,
6813 RuntimeCommitBehavior::FinalResponse,
6814 );
6815 let transition_branch = RuntimeBranch::new(
6816 RuntimeTaskPurpose::StateTransition,
6817 RuntimeOptimizationKind::ParallelStateTransition,
6818 RuntimeTaskPriority::Critical,
6819 RuntimeCommitBehavior::TransitionDecision,
6820 );
6821 let skill_branch = RuntimeBranch::new(
6822 RuntimeTaskPurpose::SkillRouting,
6823 RuntimeOptimizationKind::SpeculativeSkillRouting,
6824 RuntimeTaskPriority::High,
6825 RuntimeCommitBehavior::SkillSelection,
6826 );
6827 let reasoning_branch = RuntimeBranch::new(
6828 RuntimeTaskPurpose::ReasoningJudge,
6829 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6830 RuntimeTaskPriority::Normal,
6831 RuntimeCommitBehavior::ReasoningDecision,
6832 );
6833 let main_id = main_branch.branch_id();
6834 let transition_id = transition_branch.branch_id();
6835 let skill_id = skill_branch.branch_id();
6836 let reasoning_id = reasoning_branch.branch_id();
6837
6838 let main_id_for_future = main_id.clone();
6839 if !branch_set.schedule(
6840 main_branch,
6841 Box::pin(async move {
6842 match crate::optimization::observability::with_branch_observation(
6843 &main_id_for_future,
6844 main_kind,
6845 RuntimeCommitBehavior::FinalResponse,
6846 self.generate_main_response_draft(processed_input, &ReasoningMode::None),
6847 )
6848 .await
6849 {
6850 Ok(draft) => RuntimeBranchResult::MainDraft(draft),
6851 Err(error) => RuntimeBranchResult::Failed(error),
6852 }
6853 }),
6854 ) {
6855 return Ok(None);
6856 }
6857
6858 if transition_enabled {
6859 let transition_id_for_future = transition_id.clone();
6860 if !branch_set.schedule(
6861 transition_branch,
6862 Box::pin(async move {
6863 match crate::optimization::observability::with_branch_observation(
6864 &transition_id_for_future,
6865 RuntimeOptimizationKind::ParallelStateTransition,
6866 RuntimeCommitBehavior::TransitionDecision,
6867 self.select_parallel_transition_candidate(processed_input),
6868 )
6869 .await
6870 {
6871 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
6872 RuntimeBranchResult::Transition(Some(candidate))
6873 }
6874 Ok(ParallelTransitionSelection::NoMatch) => {
6875 RuntimeBranchResult::Transition(None)
6876 }
6877 Ok(ParallelTransitionSelection::ReservationExhausted) => {
6878 RuntimeBranchResult::Cancelled
6879 }
6880 Err(error) => RuntimeBranchResult::Failed(error),
6881 }
6882 }),
6883 ) {
6884 transition_enabled = false;
6885 }
6886 }
6887
6888 if skill_enabled {
6889 let skill_id_for_future = skill_id.clone();
6890 if !branch_set.schedule(
6891 skill_branch,
6892 Box::pin(async move {
6893 if !self.reserve_active_speculative_llm_call(
6894 RuntimeOptimizationKind::SpeculativeSkillRouting,
6895 ) {
6896 return RuntimeBranchResult::Cancelled;
6897 }
6898 match crate::optimization::observability::with_branch_observation(
6899 &skill_id_for_future,
6900 RuntimeOptimizationKind::SpeculativeSkillRouting,
6901 RuntimeCommitBehavior::SkillSelection,
6902 self.select_skill_candidate(processed_input),
6903 )
6904 .await
6905 {
6906 Ok(candidate) => RuntimeBranchResult::Skill(candidate),
6907 Err(error) => RuntimeBranchResult::Failed(error),
6908 }
6909 }),
6910 ) {
6911 skill_enabled = false;
6912 }
6913 }
6914
6915 if reasoning_enabled {
6916 let reasoning_id_for_future = reasoning_id.clone();
6917 if !branch_set.schedule(
6918 reasoning_branch,
6919 Box::pin(async move {
6920 if !self.reserve_active_speculative_llm_call(
6921 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6922 ) {
6923 return RuntimeBranchResult::Cancelled;
6924 }
6925 match crate::optimization::observability::with_branch_observation(
6926 &reasoning_id_for_future,
6927 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6928 RuntimeCommitBehavior::ReasoningDecision,
6929 self.determine_reasoning_mode_strict(processed_input),
6930 )
6931 .await
6932 {
6933 Ok(mode) => RuntimeBranchResult::Reasoning(mode),
6934 Err(error) => RuntimeBranchResult::Failed(error),
6935 }
6936 }),
6937 ) {
6938 reasoning_enabled = false;
6939 }
6940 }
6941
6942 if matches!(effective_reasoning_mode, ReasoningMode::Auto) && !reasoning_enabled {
6943 self.finalize_pending_branches(branch_set.cancel_pending());
6944 return Ok(None);
6945 }
6946
6947 if !transition_enabled && !skill_enabled && !reasoning_enabled {
6948 self.finalize_pending_branches(branch_set.cancel_pending());
6949 return Ok(None);
6950 }
6951
6952 let mut main_pending = true;
6953 let mut skill_pending = skill_enabled;
6954 let mut reasoning_pending = reasoning_enabled;
6955 let mut transition_finalized = !transition_enabled;
6956 let mut skill_finalized = !skill_enabled;
6957 let mut reasoning_finalized = !reasoning_enabled;
6958 let mut main_result: Option<Result<MainResponseDraft>> = None;
6959 let mut transition_candidate: Option<TransitionCandidate> = None;
6960 let mut skill_candidate: Option<SkillCandidate> = None;
6961 let mut reasoning_decision: Option<ReasoningMode> = None;
6962 let mut transition_fallback_required = false;
6963 let mut skill_fallback_required = false;
6964 let mut reasoning_fallback_required = false;
6965
6966 loop {
6967 if let Some(candidate) = transition_candidate.take() {
6968 if self
6969 .approve_transition_target(&candidate.from_state, candidate.target())
6970 .await?
6971 {
6972 self.finalize_pending_branches(branch_set.cancel_pending());
6974 if !main_pending {
6975 self.finalize_branch_loss(
6976 &main_id,
6977 main_kind,
6978 RuntimeCommitBehavior::FinalResponse,
6979 false,
6980 main_result.as_ref().map(|result| result.is_err()),
6981 );
6982 }
6983 if skill_enabled && !skill_pending {
6984 self.finalize_branch_loss(
6985 &skill_id,
6986 RuntimeOptimizationKind::SpeculativeSkillRouting,
6987 RuntimeCommitBehavior::SkillSelection,
6988 false,
6989 Some(false),
6990 );
6991 }
6992 if reasoning_enabled && !reasoning_pending {
6993 self.finalize_branch_loss(
6994 &reasoning_id,
6995 RuntimeOptimizationKind::SpeculativeReasoningAuto,
6996 RuntimeCommitBehavior::ReasoningDecision,
6997 false,
6998 Some(false),
6999 );
7000 }
7001 if !self
7002 .apply_pre_response_transition_candidate(
7003 &candidate,
7004 &HashMap::new(),
7005 processed_input,
7006 )
7007 .await?
7008 {
7009 self.finalize_optional_branch(
7010 &transition_id,
7011 RuntimeOptimizationKind::ParallelStateTransition,
7012 RuntimeCommitBehavior::TransitionDecision,
7013 "discarded",
7014 false,
7015 );
7016 return Ok(None);
7017 }
7018 self.finalize_optional_branch(
7019 &transition_id,
7020 RuntimeOptimizationKind::ParallelStateTransition,
7021 RuntimeCommitBehavior::TransitionDecision,
7022 "committed",
7023 true,
7024 );
7025 return self
7026 .redispatch_current_state(processed_input)
7027 .await
7028 .map(Some);
7029 }
7030 self.finalize_optional_branch(
7031 &transition_id,
7032 RuntimeOptimizationKind::ParallelStateTransition,
7033 RuntimeCommitBehavior::TransitionDecision,
7034 "discarded",
7035 false,
7036 );
7037 transition_finalized = true;
7038 }
7039
7040 if transition_finalized && skill_candidate.is_some() {
7041 let candidate = skill_candidate.take().unwrap();
7042 self.finalize_optional_branch(
7043 &skill_id,
7044 RuntimeOptimizationKind::SpeculativeSkillRouting,
7045 RuntimeCommitBehavior::SkillSelection,
7046 "committed",
7047 true,
7048 );
7049 if !main_pending {
7050 self.finalize_branch_loss(
7051 &main_id,
7052 main_kind,
7053 RuntimeCommitBehavior::FinalResponse,
7054 false,
7055 main_result.as_ref().map(|result| result.is_err()),
7056 );
7057 }
7058 if reasoning_enabled && !reasoning_pending {
7059 self.finalize_branch_loss(
7060 &reasoning_id,
7061 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7062 RuntimeCommitBehavior::ReasoningDecision,
7063 false,
7064 Some(false),
7065 );
7066 }
7067 self.finalize_pending_branches(branch_set.cancel_pending());
7068 self.commit_root_user_message(processed_input).await?;
7069 return match self
7070 .commit_skill_candidate_route_result(candidate, processed_input)
7071 .await?
7072 {
7073 SkillRouteResult::Response { skill_id, content } => self
7074 .handle_skill_response(processed_input, &skill_id, content, input_context)
7075 .await
7076 .map(Some),
7077 SkillRouteResult::NeedsClarification {
7078 response,
7079 ownership,
7080 } => {
7081 let admission = self
7082 .admit_optional_disambiguation_ownership(ownership)
7083 .await?;
7084 if response
7085 .metadata
7086 .as_ref()
7087 .and_then(|m| m.get("disambiguation"))
7088 .and_then(|d| d.get("status"))
7089 .and_then(|s| s.as_str())
7090 == Some("awaiting_clarification")
7091 {
7092 self.memory
7093 .add_message(ChatMessage::assistant(&response.content))
7094 .await?;
7095 }
7096 drop(admission);
7097 self.finish_turn_if_root(&response).await?;
7098 Ok(Some(response))
7099 }
7100 SkillRouteResult::NoMatch => Ok(None),
7101 };
7102 }
7103
7104 if transition_finalized
7105 && skill_finalized
7106 && let Some(reasoning_mode) = reasoning_decision.take()
7107 {
7108 if !matches!(reasoning_mode, ReasoningMode::None) {
7109 self.finalize_optional_branch(
7110 &reasoning_id,
7111 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7112 RuntimeCommitBehavior::ReasoningDecision,
7113 "committed",
7114 true,
7115 );
7116 if !main_pending {
7117 self.finalize_branch_loss(
7118 &main_id,
7119 main_kind,
7120 RuntimeCommitBehavior::FinalResponse,
7121 false,
7122 main_result.as_ref().map(|result| result.is_err()),
7123 );
7124 }
7125 self.finalize_pending_branches(branch_set.cancel_pending());
7126 self.commit_root_user_message(processed_input).await?;
7127 return if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
7128 self.handle_plan_and_execute(processed_input, input_context, true)
7129 .await
7130 .map(Some)
7131 } else {
7132 self.run_committed_response_loop_with_reasoning(
7133 processed_input,
7134 input_context,
7135 reasoning_mode,
7136 true,
7137 )
7138 .await
7139 .map(Some)
7140 };
7141 }
7142 self.finalize_optional_branch(
7143 &reasoning_id,
7144 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7145 RuntimeCommitBehavior::ReasoningDecision,
7146 "committed",
7147 true,
7148 );
7149 reasoning_finalized = true;
7150 }
7151
7152 if transition_finalized && skill_finalized && reasoning_finalized {
7153 if transition_fallback_required
7154 || skill_fallback_required
7155 || reasoning_fallback_required
7156 {
7157 if !main_pending {
7158 self.finalize_branch_loss(
7159 &main_id,
7160 main_kind,
7161 RuntimeCommitBehavior::FinalResponse,
7162 false,
7163 main_result.as_ref().map(|result| result.is_err()),
7164 );
7165 }
7166 self.finalize_pending_branches(branch_set.cancel_pending());
7167 return Ok(None);
7168 }
7169
7170 if let Some(result) = main_result.take() {
7171 let draft = match result {
7172 Ok(draft) => draft,
7173 Err(error) => {
7174 self.finalize_optional_branch(
7175 &main_id,
7176 main_kind,
7177 RuntimeCommitBehavior::FinalResponse,
7178 "failed",
7179 false,
7180 );
7181 self.finalize_pending_branches(branch_set.cancel_pending());
7182 return Err(error);
7183 }
7184 };
7185 self.finalize_optional_branch(
7186 &main_id,
7187 main_kind,
7188 RuntimeCommitBehavior::FinalResponse,
7189 "committed",
7190 true,
7191 );
7192 self.finalize_pending_branches(branch_set.cancel_pending());
7193 return self
7194 .commit_main_response_draft(
7195 processed_input,
7196 input_context,
7197 draft,
7198 ReasoningMode::None,
7199 reasoning_enabled,
7200 )
7201 .await
7202 .map(Some);
7203 }
7204 }
7205
7206 if branch_set.is_empty() {
7207 return Ok(None);
7208 }
7209
7210 let Some(outcome) = branch_set.next_completed().await else {
7211 return Ok(None);
7212 };
7213 let branch_id = outcome.branch.branch_id();
7214 match outcome.result {
7215 RuntimeBranchResult::MainDraft(draft) => {
7216 main_pending = false;
7217 main_result = Some(Ok(draft));
7218 }
7219 RuntimeBranchResult::Transition(candidate) => {
7220 if let Some(candidate) = candidate {
7221 transition_candidate = Some(candidate);
7222 } else {
7223 self.finalize_optional_branch(
7224 &transition_id,
7225 RuntimeOptimizationKind::ParallelStateTransition,
7226 RuntimeCommitBehavior::TransitionDecision,
7227 "discarded",
7228 false,
7229 );
7230 transition_finalized = true;
7231 }
7232 }
7233 RuntimeBranchResult::Skill(candidate) => {
7234 skill_pending = false;
7235 if let Some(candidate) = candidate {
7236 skill_candidate = Some(candidate);
7237 } else {
7238 self.finalize_optional_branch(
7239 &skill_id,
7240 RuntimeOptimizationKind::SpeculativeSkillRouting,
7241 RuntimeCommitBehavior::SkillSelection,
7242 "discarded",
7243 false,
7244 );
7245 skill_finalized = true;
7246 }
7247 }
7248 RuntimeBranchResult::Reasoning(mode) => {
7249 reasoning_pending = false;
7250 reasoning_decision = Some(mode);
7251 }
7252 RuntimeBranchResult::Failed(error) => {
7253 if branch_id == main_id {
7254 main_pending = false;
7255 main_result = Some(Err(error));
7256 } else if branch_id == transition_id {
7257 self.finalize_optional_branch(
7258 &transition_id,
7259 RuntimeOptimizationKind::ParallelStateTransition,
7260 RuntimeCommitBehavior::TransitionDecision,
7261 "failed",
7262 false,
7263 );
7264 transition_finalized = true;
7265 } else if branch_id == skill_id {
7266 skill_pending = false;
7267 self.finalize_optional_branch(
7268 &skill_id,
7269 RuntimeOptimizationKind::SpeculativeSkillRouting,
7270 RuntimeCommitBehavior::SkillSelection,
7271 "failed",
7272 false,
7273 );
7274 skill_finalized = true;
7275 } else if branch_id == reasoning_id {
7276 reasoning_pending = false;
7277 self.finalize_optional_branch(
7278 &reasoning_id,
7279 RuntimeOptimizationKind::SpeculativeReasoningAuto,
7280 RuntimeCommitBehavior::ReasoningDecision,
7281 "failed",
7282 false,
7283 );
7284 reasoning_finalized = true;
7285 }
7286 }
7287 RuntimeBranchResult::Cancelled => {
7288 self.finalize_optional_branch(
7289 &branch_id,
7290 outcome.branch.optimization,
7291 outcome.branch.commit_behavior,
7292 "cancelled",
7293 false,
7294 );
7295 if branch_id == main_id {
7296 main_pending = false;
7297 main_result =
7298 Some(Err(AgentError::Other("main branch cancelled".to_string())));
7299 } else if branch_id == transition_id {
7300 transition_finalized = true;
7301 transition_fallback_required = true;
7302 } else if branch_id == skill_id {
7303 skill_pending = false;
7304 skill_finalized = true;
7305 skill_fallback_required = true;
7306 } else if branch_id == reasoning_id {
7307 reasoning_pending = false;
7308 reasoning_finalized = true;
7309 reasoning_fallback_required = true;
7310 }
7311 }
7312 }
7313 }
7314 }
7315
7316 fn finalize_pending_branches(&self, branches: Vec<RuntimeBranch>) {
7317 for branch in branches {
7318 self.finalize_optional_branch(
7319 &branch.branch_id(),
7320 branch.optimization,
7321 branch.commit_behavior,
7322 "cancelled",
7323 false,
7324 );
7325 }
7326 }
7327
7328 fn finalize_branch_loss(
7333 &self,
7334 branch_id: &str,
7335 optimization: RuntimeOptimizationKind,
7336 commit_behavior: RuntimeCommitBehavior,
7337 pending: bool,
7338 completed_failed: Option<bool>,
7339 ) {
7340 let status = if pending {
7341 "cancelled"
7342 } else if completed_failed.unwrap_or(false) {
7343 "failed"
7344 } else {
7345 "discarded"
7346 };
7347 self.finalize_optional_branch(branch_id, optimization, commit_behavior, status, false);
7348 }
7349
7350 fn finalize_optional_branch(
7355 &self,
7356 branch_id: &str,
7357 optimization: RuntimeOptimizationKind,
7358 commit_behavior: RuntimeCommitBehavior,
7359 status: &str,
7360 winner: bool,
7361 ) {
7362 crate::optimization::observability::finalize_branch(
7363 self.observability_manager.as_ref(),
7364 branch_id,
7365 status,
7366 winner,
7367 optimization,
7368 commit_behavior,
7369 );
7370 }
7371
7372 fn has_parallel_transition_candidates(&self) -> bool {
7377 self.transitions_available_for_commit()
7378 .map(|(transitions, _)| {
7379 transitions
7380 .iter()
7381 .any(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7382 })
7383 .unwrap_or(false)
7384 }
7385
7386 async fn select_parallel_transition_candidate(
7391 &self,
7392 processed_input: &str,
7393 ) -> Result<ParallelTransitionSelection> {
7394 let Some((transitions, current_state)) = self.transitions_available_for_commit() else {
7395 return Ok(ParallelTransitionSelection::NoMatch);
7396 };
7397 let parallel: Vec<Transition> = transitions
7398 .into_iter()
7399 .filter(|transition| matches!(transition.timing, TransitionTiming::Parallel))
7400 .filter(|transition| !transition.requires_response)
7401 .collect();
7402 if parallel.is_empty() {
7403 return Ok(ParallelTransitionSelection::NoMatch);
7404 }
7405 let empty_staged = HashMap::new();
7406 if let Some(candidate) = self.select_deterministic_transition_candidate(
7407 processed_input,
7408 ¤t_state,
7409 ¶llel,
7410 &empty_staged,
7411 ) {
7412 return Ok(ParallelTransitionSelection::Candidate(candidate));
7413 }
7414 let when_transitions: Vec<(usize, &Transition)> = parallel
7415 .iter()
7416 .enumerate()
7417 .filter(|(_, transition)| !transition.when.trim().is_empty())
7418 .collect();
7419 if when_transitions.is_empty() {
7420 return Ok(ParallelTransitionSelection::NoMatch);
7421 }
7422 let llm = self
7423 .llm_registry
7424 .router()
7425 .or_else(|_| self.llm_registry.default())
7426 .map_err(|e| AgentError::Config(e.to_string()))?;
7427 let conditions = when_transitions
7428 .iter()
7429 .enumerate()
7430 .map(|(display_idx, (_, transition))| {
7431 format!("{}. {}", display_idx + 1, transition.when)
7432 })
7433 .collect::<Vec<_>>()
7434 .join("\n");
7435 if !self
7436 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::ParallelStateTransition)
7437 {
7438 return Ok(ParallelTransitionSelection::ReservationExhausted);
7439 }
7440 let context_preview = self.branch_context_preview();
7441 let prompt = format!(
7442 "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-{}).",
7443 current_state,
7444 processed_input,
7445 context_preview,
7446 conditions,
7447 when_transitions.len()
7448 );
7449 let response = self
7450 .observe_purpose(
7451 ObservationPurpose::StateTransitionEvaluation,
7452 llm.complete(&[ChatMessage::user(prompt)], None),
7453 )
7454 .await
7455 .map_err(|e| AgentError::LLM(e.to_string()))?;
7456 let choice = response.content.trim().parse::<usize>().unwrap_or(0);
7457 if choice == 0 || choice > when_transitions.len() {
7458 return Ok(ParallelTransitionSelection::NoMatch);
7459 }
7460 let transition = when_transitions[choice - 1].1.clone();
7461 Ok(ParallelTransitionSelection::Candidate(
7462 TransitionCandidate::new(
7463 current_state,
7464 transition.clone(),
7465 Self::transition_reason(&transition),
7466 ),
7467 ))
7468 }
7469
7470 async fn redispatch_current_state(&self, processed_input: &str) -> Result<AgentResponse> {
7472 const MAX_REDISPATCH_DEPTH: u32 = 3;
7473 let current_depth = *self.redispatch_depth.read();
7474 if current_depth >= MAX_REDISPATCH_DEPTH {
7475 warn!(depth = current_depth, "Re-dispatch depth limit reached");
7476 let response = AgentResponse::new("");
7477 self.finish_turn_if_root(&response).await?;
7478 return Ok(response);
7479 }
7480 *self.redispatch_depth.write() += 1;
7481 if let Some(context) = self.active_turn_context.write().as_mut() {
7482 context.enter_redispatch();
7483 }
7484 let result = Box::pin(self.run_loop_internal(processed_input)).await;
7485 *self.redispatch_depth.write() -= 1;
7486 if let Some(context) = self.active_turn_context.write().as_mut() {
7487 context.exit_redispatch();
7488 }
7489 let response = result?;
7490 self.finish_turn_if_root(&response).await?;
7491 Ok(response)
7492 }
7493
7494 async fn finish_turn_if_root(&self, response: &AgentResponse) -> Result<()> {
7496 if *self.redispatch_depth.read() == 0 {
7497 self.post_turn_session_lifecycle().await?;
7498 if let Some(context) = self.active_turn_context.write().as_mut() {
7499 context.mark_post_turn_lifecycle_completed();
7500 }
7501 self.hooks.on_response(response).await;
7502 self.end_root_turn();
7503 }
7504 Ok(())
7505 }
7506
7507 async fn execute_state_exit_actions(&self, state_path: &str) {
7509 if let Some(ref sm) = self.state_machine
7510 && let Some(def) = sm.get_definition(state_path)
7511 && !def.on_exit.is_empty()
7512 {
7513 debug!(state = %state_path, count = def.on_exit.len(), "Executing on_exit actions");
7514 self.execute_state_actions(&def.on_exit).await;
7515 }
7516 }
7517
7518 fn state_was_previously_entered(
7520 state_path: &str,
7521 from_state: &str,
7522 history_before: &[StateTransitionEvent],
7523 ) -> bool {
7524 state_path == from_state
7525 || history_before
7526 .iter()
7527 .any(|event| event.from == state_path || event.to == state_path)
7528 }
7529
7530 async fn execute_state_enter_actions(&self, state_path: &str, is_reentry: bool) {
7532 if let Some(ref sm) = self.state_machine
7533 && let Some(def) = sm.get_definition(state_path)
7534 {
7535 if is_reentry && !def.on_reenter.is_empty() {
7536 debug!(state = %state_path, count = def.on_reenter.len(), "Executing on_reenter actions");
7537 self.execute_state_actions(&def.on_reenter).await;
7538 } else if !def.on_enter.is_empty() {
7539 debug!(state = %state_path, count = def.on_enter.len(), "Executing on_enter actions");
7540 self.execute_state_actions(&def.on_enter).await;
7541 }
7542 }
7543 }
7544
7545 async fn execute_state_actions(&self, actions: &[StateAction]) {
7547 for (action_index, action) in actions.iter().enumerate() {
7548 match action {
7549 StateAction::Tool { tool, args } => {
7550 let raw_args = args.clone().unwrap_or(Value::Object(Default::default()));
7551 let args_value = self.render_action_args(&raw_args);
7552 let state = self.state_machine.as_ref().map(|sm| sm.current());
7553 let request = ToolExecutionRequest::new(
7554 uuid::Uuid::new_v4().to_string(),
7555 tool.clone(),
7556 args_value,
7557 ToolCallSource::StateAction {
7558 state,
7559 action_index,
7560 },
7561 );
7562 match self.execute_tool_record(request).await {
7563 Ok(record) if record.success => {
7564 debug!(tool = %record.canonical_id, "State action: tool executed");
7565 let _ = self.context_manager.set(
7566 "last_tool_result",
7567 serde_json::Value::String(record.model_output_string()),
7568 );
7569 let _ = self.context_manager.set(
7570 "last_tool_record",
7571 serde_json::to_value(record).unwrap_or(Value::Null),
7572 );
7573 }
7574 Ok(record) => {
7575 warn!(tool = %record.canonical_id, error = %record.output, "State action: tool failed");
7576 }
7577 Err(e) => {
7578 warn!(tool = %tool, error = %e, "State action: tool failed")
7579 }
7580 }
7581 }
7582 StateAction::Skill { skill } => {
7583 if let Some(ref executor) = self.skill_executor {
7584 if let Some(def) = self.skills.iter().find(|s| s.id == *skill) {
7585 match executor
7586 .execute_with_invoker(def, "", serde_json::json!({}), self)
7587 .await
7588 {
7589 Ok(_) => debug!(skill = %skill, "State action: skill executed"),
7590 Err(e) => {
7591 warn!(skill = %skill, error = %e, "State action: skill failed")
7592 }
7593 }
7594 } else {
7595 warn!(skill = %skill, "State action: skill not found");
7596 }
7597 }
7598 }
7599 StateAction::SetContext { set_context } => {
7600 for (key, value) in set_context {
7601 if let Err(e) = self.context_manager.set(key, value.clone()) {
7602 warn!(key = %key, error = %e, "State action: set_context failed");
7603 } else {
7604 debug!(key = %key, "State action: context set");
7605 }
7606 }
7607 }
7608 StateAction::Prompt {
7609 prompt,
7610 llm,
7611 store_as,
7612 } => {
7613 let llm_result = if let Some(alias) = llm {
7614 self.llm_registry.get(alias)
7615 } else {
7616 self.llm_registry.default()
7617 };
7618 match llm_result {
7619 Ok(llm_provider) => {
7620 let context = self.build_context_with_overlays();
7622 let rendered_prompt = self
7623 .template_renderer
7624 .render(prompt, &context)
7625 .unwrap_or_else(|_| prompt.clone());
7626 let recent =
7627 self.memory.get_messages(Some(5)).await.unwrap_or_default();
7628 let mut messages: Vec<ChatMessage> = recent;
7629 messages.push(ChatMessage::user(&rendered_prompt));
7630 match self
7631 .observe_purpose(
7632 ObservationPurpose::StateAction,
7633 llm_provider.complete(&messages, None),
7634 )
7635 .await
7636 {
7637 Ok(response) => {
7638 if let Some(key) = store_as {
7639 let _ = self
7640 .context_manager
7641 .set(key, Value::String(response.content));
7642 debug!(key = %key, "State action: prompt result stored");
7643 }
7644 }
7645 Err(e) => {
7646 warn!(error = %e, "State action: prompt LLM call failed");
7647 }
7648 }
7649 }
7650 Err(e) => {
7651 warn!(error = %e, "State action: LLM not found for prompt");
7652 }
7653 }
7654 }
7655 }
7656 }
7657 }
7658
7659 async fn run_context_extractors_staged(&self, user_message: &str) -> HashMap<String, Value> {
7660 let extractors = match &self.state_machine {
7661 Some(sm) => match sm.current_definition() {
7662 Some(def) if !def.extract.is_empty() => def.extract.clone(),
7663 _ => return HashMap::new(),
7664 },
7665 None => return HashMap::new(),
7666 };
7667
7668 let mut staged = HashMap::new();
7669 for extractor in &extractors {
7670 let prompt = if let Some(ref custom) = extractor.llm_extract {
7671 format!(
7672 "User message:\n\"{}\"\n\nInstruction:\n{}",
7673 user_message, custom
7674 )
7675 } else if let Some(ref desc) = extractor.description {
7676 format!(
7677 "From the following message, extract: {}\n\n\
7678 Message: \"{}\"\n\n\
7679 If the information is present, return ONLY the extracted value.\n\
7680 If NOT present, return exactly: __NONE__",
7681 desc, user_message
7682 )
7683 } else {
7684 continue;
7685 };
7686
7687 let llm = match self
7688 .llm_registry
7689 .get(&extractor.llm)
7690 .or_else(|_| self.llm_registry.get("router"))
7691 .or_else(|_| self.llm_registry.get("default"))
7692 {
7693 Ok(llm) => llm,
7694 Err(e) => {
7695 warn!(key = %extractor.key, error = %e, "Extractor LLM not found");
7696 continue;
7697 }
7698 };
7699
7700 let messages = vec![ChatMessage::user(&prompt)];
7701 match self
7702 .observe_purpose(
7703 ObservationPurpose::ContextExtraction,
7704 llm.complete(&messages, None),
7705 )
7706 .await
7707 {
7708 Ok(response) => {
7709 let value = response.content.trim().to_string();
7710 if value != "__NONE__" && !value.is_empty() {
7711 staged.insert(
7712 extractor.key.clone(),
7713 serde_json::Value::String(value.clone()),
7714 );
7715 debug!(key = %extractor.key, value = %value, "Context extracted");
7716 } else if extractor.required {
7717 warn!(key = %extractor.key, "Required extraction returned no value");
7718 }
7719 }
7720 Err(e) => {
7721 warn!(key = %extractor.key, error = %e, "Context extraction LLM call failed");
7722 }
7723 }
7724 }
7725 staged
7726 }
7727
7728 fn commit_staged_context_writes(&self, staged: &HashMap<String, Value>) {
7729 for (key, value) in staged {
7730 if let Err(error) = self.context_manager.update(key, value.clone()) {
7731 warn!(key = %key, error = %error, "staged context write failed");
7732 }
7733 }
7734 }
7735
7736 async fn run_context_extractors(&self, user_message: &str) {
7738 let staged = self.run_context_extractors_staged(user_message).await;
7739 self.commit_staged_context_writes(&staged);
7740 }
7741
7742 async fn check_memory_compression(&self) -> Result<()> {
7743 if self.memory.needs_compression() {
7744 let result = self.memory.compress(None).await?;
7745 if let CompressResult::Compressed {
7746 messages_summarized,
7747 new_summary_length,
7748 tokens_saved,
7749 } = result
7750 {
7751 let event = MemoryCompressEvent::new(
7752 messages_summarized,
7753 tokens_saved,
7754 new_summary_length as u32,
7755 );
7756 self.hooks.on_memory_compress(&event).await;
7757 debug!(
7758 messages = messages_summarized,
7759 tokens_saved = tokens_saved,
7760 "Memory compressed"
7761 );
7762 }
7763 }
7764
7765 self.handle_memory_overflow().await?;
7767 self.check_memory_budget().await;
7768
7769 Ok(())
7770 }
7771
7772 async fn check_memory_budget(&self) {
7773 let Some(ref budget) = self.memory_token_budget else {
7774 return;
7775 };
7776
7777 let context = match self.memory.get_context().await {
7778 Ok(ctx) => ctx,
7779 Err(_) => return,
7780 };
7781
7782 let used_tokens = context.estimated_tokens();
7784 if budget.is_over_warn_threshold(used_tokens) {
7785 let event = MemoryBudgetEvent::new("memory", used_tokens, budget.total);
7786 self.hooks.on_memory_budget_warning(&event).await;
7787 debug!(
7788 used = used_tokens,
7789 total = budget.total,
7790 percent = event.usage_percent,
7791 "Memory budget warning"
7792 );
7793 }
7794
7795 if let Some(ref summary) = context.summary {
7797 let summary_tokens = ai_agents_memory::estimate_tokens(summary);
7798 let summary_budget = budget.allocation.summary;
7799 if summary_budget > 0 {
7800 let warn_threshold =
7801 (summary_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7802 if summary_tokens >= warn_threshold {
7803 let event = MemoryBudgetEvent::new("summary", summary_tokens, summary_budget);
7804 self.hooks.on_memory_budget_warning(&event).await;
7805 }
7806 }
7807 }
7808
7809 let recent_tokens: u32 = context
7811 .messages
7812 .iter()
7813 .map(ai_agents_memory::estimate_message_tokens)
7814 .sum();
7815 let recent_budget = budget.allocation.recent_messages;
7816 if recent_budget > 0 {
7817 let warn_threshold =
7818 (recent_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7819 if recent_tokens >= warn_threshold {
7820 let event = MemoryBudgetEvent::new("recent_messages", recent_tokens, recent_budget);
7821 self.hooks.on_memory_budget_warning(&event).await;
7822 }
7823 }
7824
7825 let relationship_budget = budget.allocation.relationships;
7826 if relationship_budget > 0 {
7827 let relationship_tokens = self
7828 .relationship_memory_text()
7829 .map(|text| ai_agents_memory::estimate_tokens(&text))
7830 .unwrap_or(0);
7831 let warn_threshold =
7832 (relationship_budget as f64 * budget.warn_at_percent as f64 / 100.0) as u32;
7833 if relationship_tokens >= warn_threshold {
7834 let event = MemoryBudgetEvent::new(
7835 "relationships",
7836 relationship_tokens,
7837 relationship_budget,
7838 );
7839 self.hooks.on_memory_budget_warning(&event).await;
7840 }
7841 }
7842 }
7843
7844 async fn handle_memory_overflow(&self) -> Result<()> {
7845 let Some(ref budget) = self.memory_token_budget else {
7846 return Ok(());
7847 };
7848
7849 let context = self.memory.get_context().await?;
7850 let used_tokens = context.estimated_tokens();
7851
7852 if used_tokens <= budget.total {
7853 return Ok(());
7854 }
7855
7856 match budget.overflow_strategy {
7857 OverflowStrategy::TruncateOldest => {
7858 let tokens_to_free = used_tokens - budget.total;
7859 let messages_to_evict = self.calculate_eviction_count(tokens_to_free);
7860 if messages_to_evict > 0 {
7861 self.evict_messages(messages_to_evict, EvictionReason::TokenBudgetExceeded)
7862 .await?;
7863 }
7864 }
7865 OverflowStrategy::SummarizeMore => {
7866 let max_attempts = context.total_messages.max(1);
7867 for _ in 0..max_attempts {
7868 match self.memory.compress(None).await? {
7869 CompressResult::Compressed {
7870 messages_summarized,
7871 ..
7872 } if messages_summarized > 0 => {
7873 let context = self.memory.get_context().await?;
7874 if context.estimated_tokens() <= budget.total {
7875 return Ok(());
7876 }
7877 }
7878 _ => break,
7879 }
7880 }
7881 let context = self.memory.get_context().await?;
7882 let used_tokens = context.estimated_tokens();
7883 if used_tokens > budget.total {
7884 return Err(AgentError::MemoryBudgetExceeded {
7885 used: used_tokens,
7886 budget: budget.total,
7887 });
7888 }
7889 }
7890 OverflowStrategy::Error => {
7891 return Err(AgentError::MemoryBudgetExceeded {
7892 used: used_tokens,
7893 budget: budget.total,
7894 });
7895 }
7896 }
7897 Ok(())
7898 }
7899
7900 fn calculate_eviction_count(&self, tokens_to_free: u32) -> usize {
7901 ((tokens_to_free as f64 / 50.0).ceil() as usize).max(1)
7903 }
7904
7905 async fn evict_messages(&self, count: usize, reason: EvictionReason) -> Result<()> {
7906 let evicted = self.memory.evict_oldest(count).await?;
7907 if !evicted.is_empty() {
7908 let event = MemoryEvictEvent {
7909 reason,
7910 messages_evicted: evicted.len(),
7911 importance_scores: vec![],
7912 };
7913 self.hooks.on_memory_evict(&event).await;
7914 debug!(count = evicted.len(), "Messages evicted from memory");
7915 }
7916 Ok(())
7917 }
7918
7919 #[instrument(skip(self, input), fields(agent = %self.info.name))]
7920 async fn determine_reasoning_mode(&self, input: &str) -> Result<ReasoningMode> {
7921 match self.determine_reasoning_mode_strict(input).await {
7922 Ok(mode) => Ok(mode),
7923 Err(_) => Ok(ReasoningMode::None),
7924 }
7925 }
7926
7927 async fn determine_reasoning_mode_strict(&self, input: &str) -> Result<ReasoningMode> {
7928 let effective_config = self.get_effective_reasoning_config();
7929
7930 if !matches!(effective_config.mode, ReasoningMode::Auto) {
7931 return Ok(effective_config.mode.clone());
7932 }
7933
7934 let judge_llm = effective_config
7935 .judge_llm
7936 .as_ref()
7937 .and_then(|alias| self.llm_registry.get(alias).ok())
7938 .or_else(|| self.llm_registry.router().ok())
7939 .or_else(|| self.llm_registry.default().ok());
7940
7941 let Some(llm) = judge_llm else {
7942 return Ok(ReasoningMode::None);
7943 };
7944
7945 let prompt = format!(
7946 r#"Analyze this user request and determine the appropriate reasoning mode.
7947
7948User request: "{}"
7949
7950Choose ONE of these modes:
7951- none: Simple queries, greetings, direct answers (fastest)
7952- cot: Complex analysis, multi-step reasoning, math problems
7953- react: Tasks requiring multiple tool calls with observation
7954- plan_and_execute: Complex multi-step tasks requiring coordination
7955
7956Respond with ONLY the mode name (none, cot, react, or plan_and_execute)."#,
7957 input
7958 );
7959
7960 let messages = vec![ChatMessage::user(&prompt)];
7961 let response = self
7962 .observe_purpose(
7963 ObservationPurpose::ReflectionDecision,
7964 llm.complete(&messages, None),
7965 )
7966 .await
7967 .map_err(|e| AgentError::LLM(e.to_string()))?;
7968
7969 let mode_str = response.content.trim().to_lowercase();
7970 Ok(match mode_str.as_str() {
7971 "cot" => ReasoningMode::CoT,
7972 "react" => ReasoningMode::React,
7973 "plan_and_execute" => ReasoningMode::PlanAndExecute,
7974 _ => ReasoningMode::None,
7975 })
7976 }
7977
7978 async fn should_reflect(&self, input: &str, response: &str) -> Result<bool> {
7979 let effective_config = self.get_effective_reflection_config();
7980
7981 if !effective_config.requires_evaluation() {
7982 return Ok(false);
7983 }
7984
7985 if effective_config.is_enabled() {
7986 return Ok(true);
7987 }
7988
7989 let evaluator_llm = effective_config
7990 .evaluator_llm
7991 .as_ref()
7992 .and_then(|alias| self.llm_registry.get(alias).ok())
7993 .or_else(|| self.llm_registry.router().ok())
7994 .or_else(|| self.llm_registry.default().ok());
7995
7996 let Some(llm) = evaluator_llm else {
7997 return Ok(false);
7998 };
7999
8000 let response_preview: String = response.chars().take(500).collect();
8001 let prompt = format!(
8002 r#"Should this response be evaluated for quality? Consider if it's a complex or important response.
8003
8004User query: "{}"
8005Response: "{}"
8006
8007Answer YES or NO only."#,
8008 input, response_preview
8009 );
8010
8011 let messages = vec![ChatMessage::user(&prompt)];
8012 let result = self
8013 .observe_purpose(
8014 ObservationPurpose::ReflectionDecision,
8015 llm.complete(&messages, None),
8016 )
8017 .await;
8018
8019 match result {
8020 Ok(resp) => Ok(resp.content.trim().to_uppercase().contains("YES")),
8021 Err(_) => Ok(false),
8022 }
8023 }
8024
8025 fn build_cot_system_prompt(&self, base_prompt: &str) -> String {
8026 format!(
8027 "{}\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>",
8028 base_prompt
8029 )
8030 }
8031
8032 fn build_react_system_prompt(&self, base_prompt: &str) -> String {
8033 format!(
8034 "{}\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>",
8035 base_prompt
8036 )
8037 }
8038
8039 async fn generate_plan(&self, input: &str) -> Result<Plan> {
8040 let effective = self.get_effective_reasoning_config();
8041 let planning_config = effective.get_planning();
8042
8043 let planner_llm = planning_config
8044 .and_then(|c| c.planner_llm.as_ref())
8045 .and_then(|alias| self.llm_registry.get(alias).ok())
8046 .or_else(|| self.llm_registry.router().ok())
8047 .or_else(|| self.llm_registry.default().ok())
8048 .ok_or_else(|| AgentError::Config("No LLM available for planning".into()))?;
8049
8050 let mut available_tool_ids: Vec<String> = self
8051 .get_available_tool_ids()
8052 .await
8053 .unwrap_or_else(|_| self.tools.list_ids());
8054 let mut available_skills: Vec<String> = self.skills.iter().map(|s| s.id.clone()).collect();
8055
8056 if let Some(config) = planning_config {
8058 if !config.available.tools.is_all() {
8059 available_tool_ids.retain(|t| config.available.tools.allows(t));
8060 }
8061 if !config.available.skills.is_all() {
8062 available_skills.retain(|s| config.available.skills.allows(s));
8063 }
8064 }
8065
8066 let tool_descriptions: Vec<String> = available_tool_ids
8069 .iter()
8070 .filter_map(|id| {
8071 self.tools.get(id).map(|tool| {
8072 let schema = tool.input_schema();
8073 let args_desc = schema
8074 .get("properties")
8075 .and_then(|p| serde_json::to_string(p).ok())
8076 .unwrap_or_else(|| "{}".to_string());
8077 format!(
8078 "- {} ({}): {}\n Arguments: {}",
8079 id,
8080 tool.name(),
8081 tool.description(),
8082 args_desc
8083 )
8084 })
8085 })
8086 .collect();
8087
8088 let tools_section = if tool_descriptions.is_empty() {
8089 "Available tools: none".to_string()
8090 } else {
8091 format!("Available tools:\n{}", tool_descriptions.join("\n"))
8092 };
8093
8094 let skills_section = if available_skills.is_empty() {
8095 "Available skills: none".to_string()
8096 } else {
8097 format!("Available skills: {}", available_skills.join(", "))
8098 };
8099
8100 let prompt = format!(
8101 r#"Create a step-by-step plan to accomplish this goal.
8102
8103Goal: "{}"
8104
8105{}
8106
8107{}
8108
8109Create a plan with clear steps. For each step, specify:
8110- description: What this step accomplishes
8111- action_type: "tool", "skill", "think", or "respond"
8112- action_target: The tool/skill id (if applicable)
8113- args: The arguments object matching the tool's schema (if action_type is "tool")
8114- dependencies: List of step IDs this depends on (empty if none)
8115
8116Respond in JSON format:
8117{{
8118 "steps": [
8119 {{"id": "step1", "description": "...", "action_type": "tool", "action_target": "tool_id", "args": {{"required_field": "value"}}, "dependencies": []}},
8120 {{"id": "step2", "description": "...", "action_type": "think", "action_target": "...", "dependencies": ["step1"]}}
8121 ]
8122}}"#,
8123 input, tools_section, skills_section,
8124 );
8125
8126 let messages = vec![ChatMessage::user(&prompt)];
8127 let response = self
8128 .observe_purpose(
8129 ObservationPurpose::PlanGeneration,
8130 planner_llm.complete(&messages, None),
8131 )
8132 .await
8133 .map_err(|e| AgentError::LLM(format!("Planning failed: {}", e)))?;
8134
8135 let mut plan = Plan::new(input);
8136
8137 if let Some(json_start) = response.content.find('{')
8138 && let Some(json_end) = response.content.rfind('}')
8139 {
8140 let json_str = &response.content[json_start..=json_end];
8141 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(json_str)
8142 && let Some(steps) = parsed.get("steps").and_then(|s| s.as_array())
8143 {
8144 for step_value in steps {
8145 let id = step_value
8146 .get("id")
8147 .and_then(|v| v.as_str())
8148 .unwrap_or("step");
8149 let desc = step_value
8150 .get("description")
8151 .and_then(|v| v.as_str())
8152 .unwrap_or("");
8153 let action_type = step_value
8154 .get("action_type")
8155 .and_then(|v| v.as_str())
8156 .unwrap_or("think");
8157 let action_target = step_value
8158 .get("action_target")
8159 .and_then(|v| v.as_str())
8160 .unwrap_or("");
8161 let args = step_value
8162 .get("args")
8163 .cloned()
8164 .unwrap_or(serde_json::json!({}));
8165 let deps: Vec<String> = step_value
8166 .get("dependencies")
8167 .and_then(|v| v.as_array())
8168 .map(|arr| {
8169 arr.iter()
8170 .filter_map(|v| v.as_str().map(String::from))
8171 .collect()
8172 })
8173 .unwrap_or_default();
8174
8175 let action = match action_type {
8176 "tool" => PlanAction::tool(action_target, args),
8177 "skill" => PlanAction::skill(action_target),
8178 "respond" => PlanAction::respond(action_target),
8179 _ => PlanAction::think(desc),
8180 };
8181
8182 let step = PlanStep::new(desc, action)
8183 .with_id(id)
8184 .with_dependencies(deps);
8185 plan.add_step(step);
8186 }
8187 }
8188 }
8189
8190 if plan.steps.is_empty() {
8191 plan.add_step(PlanStep::new(
8192 "Process the request",
8193 PlanAction::think(input),
8194 ));
8195 plan.add_step(PlanStep::new(
8196 "Provide response",
8197 PlanAction::respond("Answer based on analysis"),
8198 ));
8199 }
8200
8201 Ok(plan)
8202 }
8203
8204 async fn execute_plan(&self, plan: &mut Plan) -> Result<String> {
8205 let llm = self.get_state_llm()?;
8206 let mut results: HashMap<String, serde_json::Value> = HashMap::new();
8207 let effective = self.get_effective_reasoning_config();
8208 let max_steps = effective.get_planning().map(|c| c.max_steps).unwrap_or(10);
8209
8210 plan.status = PlanStatus::InProgress;
8211
8212 for step_idx in 0..plan.steps.len().min(max_steps as usize) {
8213 let step = &plan.steps[step_idx];
8214
8215 let deps_satisfied = step.dependencies.iter().all(|dep| {
8216 plan.steps
8217 .iter()
8218 .find(|s| &s.id == dep)
8219 .map(|s| s.status.is_completed())
8220 .unwrap_or(false)
8221 });
8222
8223 if !deps_satisfied {
8224 continue;
8225 }
8226
8227 plan.steps[step_idx].mark_running();
8228
8229 let result = match &plan.steps[step_idx].action {
8230 PlanAction::Tool { tool, args } => {
8231 let has_dep_results = plan.steps[step_idx]
8237 .dependencies
8238 .iter()
8239 .any(|dep| results.contains_key(dep));
8240
8241 let final_args = if has_dep_results {
8242 let dep_context: String = plan.steps[step_idx]
8243 .dependencies
8244 .iter()
8245 .filter_map(|dep| results.get(dep).map(|r| format!("{}: {}", dep, r)))
8246 .collect::<Vec<_>>()
8247 .join("\n");
8248
8249 let tool_schema = self
8250 .tools
8251 .get(tool)
8252 .map(|t| {
8253 let schema = t.input_schema();
8254 let props = schema
8255 .get("properties")
8256 .and_then(|p| serde_json::to_string(p).ok())
8257 .unwrap_or_else(|| "{}".to_string());
8258 format!(
8259 "{}: {}\nArguments schema: {}",
8260 t.id(),
8261 t.description(),
8262 props
8263 )
8264 })
8265 .unwrap_or_default();
8266
8267 let step_desc = &plan.steps[step_idx].description;
8268 let arg_prompt = format!(
8269 "Generate the JSON arguments for a tool call.\n\n\
8270 Tool: {}\n\n\
8271 Task: {}\n\n\
8272 Previous step results:\n{}\n\n\
8273 Planner's draft arguments: {}\n\n\
8274 Produce ONLY a valid JSON object with the correct argument values.\n\
8275 Use actual values from the previous step results, not template references.",
8276 tool_schema,
8277 step_desc,
8278 dep_context,
8279 serde_json::to_string(args).unwrap_or_default()
8280 );
8281 let messages = vec![ChatMessage::user(&arg_prompt)];
8282 match self
8283 .observe_purpose(
8284 ObservationPurpose::PlanStep,
8285 llm.complete(&messages, None),
8286 )
8287 .await
8288 {
8289 Ok(resp) => {
8290 let content = resp.content.trim();
8291 let json_start = content.find('{');
8293 let json_end = content.rfind('}');
8294 if let (Some(start), Some(end)) = (json_start, json_end) {
8295 serde_json::from_str(&content[start..=end])
8296 .unwrap_or_else(|_| args.clone())
8297 } else {
8298 args.clone()
8299 }
8300 }
8301 Err(_) => args.clone(),
8302 }
8303 } else {
8304 args.clone()
8305 };
8306
8307 let request = ToolExecutionRequest::new(
8308 uuid::Uuid::new_v4().to_string(),
8309 tool.clone(),
8310 final_args,
8311 ToolCallSource::Plan {
8312 step_index: step_idx,
8313 },
8314 );
8315 match self.execute_tool_record(request).await {
8316 Ok(record) if record.success => {
8317 serde_json::json!({ "output": record.model_output_string() })
8318 }
8319 Ok(record) => {
8320 plan.steps[step_idx].mark_failed(record.model_output_string());
8321 continue;
8322 }
8323 Err(e) => {
8324 plan.steps[step_idx].mark_failed(e.to_string());
8325 continue;
8326 }
8327 }
8328 }
8329 PlanAction::Skill { skill } => {
8330 if let Some(skill_def) = self.skills.iter().find(|s| &s.id == skill) {
8331 if let Some(ref executor) = self.skill_executor {
8332 match executor
8333 .execute_with_invoker(skill_def, "", serde_json::json!({}), self)
8334 .await
8335 {
8336 Ok(output) => serde_json::json!({ "output": output }),
8337 Err(e) => {
8338 plan.steps[step_idx].mark_failed(e.to_string());
8339 continue;
8340 }
8341 }
8342 } else {
8343 serde_json::json!({ "output": "Skill executor not available" })
8344 }
8345 } else {
8346 plan.steps[step_idx].mark_failed("Skill not found");
8347 continue;
8348 }
8349 }
8350 PlanAction::Think { prompt } => {
8351 let context: String = results
8352 .iter()
8353 .map(|(k, v)| format!("{}: {}", k, v))
8354 .collect::<Vec<_>>()
8355 .join("\n");
8356
8357 let think_prompt = format!("Context:\n{}\n\nTask: {}", context, prompt);
8358 let messages = vec![ChatMessage::user(&think_prompt)];
8359
8360 match self
8361 .observe_purpose(
8362 ObservationPurpose::PlanStep,
8363 llm.complete(&messages, None),
8364 )
8365 .await
8366 {
8367 Ok(resp) => serde_json::json!({ "output": resp.content }),
8368 Err(e) => {
8369 plan.steps[step_idx].mark_failed(e.to_string());
8370 continue;
8371 }
8372 }
8373 }
8374 PlanAction::Respond { template } => {
8375 let context: String = results
8376 .iter()
8377 .map(|(k, v)| format!("{}: {}", k, v))
8378 .collect::<Vec<_>>()
8379 .join("\n");
8380
8381 let respond_prompt = format!(
8382 "Based on this context:\n{}\n\nGenerate a response following this template/instruction: {}",
8383 context, template
8384 );
8385 let messages = vec![ChatMessage::user(&respond_prompt)];
8386
8387 match self
8388 .observe_purpose(
8389 ObservationPurpose::PlanStep,
8390 llm.complete(&messages, None),
8391 )
8392 .await
8393 {
8394 Ok(resp) => serde_json::json!({ "output": resp.content }),
8395 Err(e) => {
8396 plan.steps[step_idx].mark_failed(e.to_string());
8397 continue;
8398 }
8399 }
8400 }
8401 };
8402
8403 results.insert(plan.steps[step_idx].id.clone(), result.clone());
8404 plan.steps[step_idx].mark_completed(Some(result));
8405 }
8406
8407 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
8409 if has_failures {
8410 let failed_ids: Vec<String> = plan
8411 .steps
8412 .iter()
8413 .filter(|s| s.status.is_failed())
8414 .map(|s| s.id.clone())
8415 .collect();
8416 plan.status = PlanStatus::Failed {
8417 error: format!("Steps failed: {}", failed_ids.join(", ")),
8418 };
8419 } else {
8420 plan.status = PlanStatus::Completed;
8421 }
8422
8423 let all_outputs: Vec<String> = plan
8425 .steps
8426 .iter()
8427 .filter(|s| s.status.is_completed())
8428 .filter_map(|s| {
8429 s.result
8430 .as_ref()
8431 .and_then(|r| r.get("output"))
8432 .and_then(|o| o.as_str())
8433 .map(|o| format!("{}: {}", s.description, o))
8434 })
8435 .collect();
8436
8437 if all_outputs.is_empty() {
8438 return Ok("Plan execution completed but produced no results.".to_string());
8439 }
8440
8441 if all_outputs.len() == 1 {
8442 return Ok(all_outputs.into_iter().next().unwrap());
8443 }
8444
8445 let context = all_outputs.join("\n\n");
8447 let prompt = format!(
8448 "You completed a multi-step plan for: \"{}\"\n\nStep results:\n{}\n\nProvide a coherent final response that synthesizes these results.",
8449 plan.goal, context
8450 );
8451 let messages = vec![ChatMessage::user(&prompt)];
8452 match self
8453 .observe_purpose(ObservationPurpose::PlanStep, llm.complete(&messages, None))
8454 .await
8455 {
8456 Ok(resp) => Ok(resp.content.trim().to_string()),
8457 Err(_) => Ok(context),
8458 }
8459 }
8460
8461 async fn evaluate_response(&self, input: &str, response: &str) -> Result<EvaluationResult> {
8462 let effective_config = self.get_effective_reflection_config();
8463 self.evaluate_response_with_config(input, response, &effective_config)
8464 .await
8465 }
8466
8467 fn extract_thinking(&self, content: &str) -> (Option<String>, String) {
8468 if let Some(start) = content.find("<thinking>")
8469 && let Some(end) = content.find("</thinking>")
8470 {
8471 let thinking = content[start + 10..end].trim().to_string();
8472 let answer = content[end + 11..].trim().to_string();
8473 return (Some(thinking), answer);
8474 }
8475 (None, content.to_string())
8476 }
8477
8478 fn format_response_with_thinking(&self, thinking: Option<&str>, answer: &str) -> String {
8479 match self.get_effective_reasoning_config().output {
8480 ReasoningOutput::Hidden => answer.to_string(),
8481 ReasoningOutput::Visible => {
8482 if let Some(t) = thinking {
8483 format!("Thinking:\n{}\n\nAnswer:\n{}", t, answer)
8484 } else {
8485 answer.to_string()
8486 }
8487 }
8488 ReasoningOutput::Tagged => {
8489 if let Some(t) = thinking {
8490 format!("<thinking>{}</thinking>\n{}", t, answer)
8491 } else {
8492 answer.to_string()
8493 }
8494 }
8495 }
8496 }
8497
8498 async fn run_loop(&self, input: &str) -> Result<AgentResponse> {
8503 self.init_storage().await?;
8507 self.begin_root_turn();
8508 let _root_cleanup = RootTurnCleanup::new(self);
8509 info!(input_len = input.len(), "Starting chat");
8510
8511 self.hooks.on_message_received(input).await;
8512
8513 if !self.context_initialized.swap(true, Ordering::SeqCst) {
8517 self.context_manager.initialize().await?;
8518 debug!("Context manager initialized (defaults, env, builtins)");
8519 }
8520
8521 self.check_turn_timeout().await?;
8522 self.context_manager.refresh_per_turn().await?;
8523
8524 self.clear_disambiguation_context();
8527
8528 if let Some(ref disambiguator) = self.disambiguation_manager {
8530 let disambiguation_context = self.build_disambiguation_context().await?;
8531
8532 let state_override = self
8534 .state_machine
8535 .as_ref()
8536 .and_then(|sm| sm.current_definition())
8537 .and_then(|def| def.disambiguation.clone());
8538
8539 let state_generation = self
8540 .state_machine
8541 .as_ref()
8542 .map(|state_machine| state_machine.generation());
8543 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
8544 let mut disambiguation_result = self
8545 .observe_purpose(
8546 ObservationPurpose::DisambiguationDetection,
8547 disambiguator.process_input_with_override(
8548 input,
8549 &disambiguation_context,
8550 state_override.as_ref(),
8551 None,
8552 ),
8553 )
8554 .await?;
8555 let current_state_generation = self
8556 .state_machine
8557 .as_ref()
8558 .map(|state_machine| state_machine.generation());
8559 if current_state_generation != state_generation
8560 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
8561 {
8562 disambiguator.clear_pending().await;
8563 *self.pending_skill_id.write() = None;
8564 disambiguation_result = DisambiguationResult::Abandoned { new_input: None };
8565 info!(
8566 confirmation_event = "invalidated",
8567 invalidation_reason = "state_generation_changed",
8568 "Disambiguation result invalidated before redispatch"
8569 );
8570 }
8571 match disambiguation_result {
8572 DisambiguationResult::Clear => {
8573 debug!("Input is clear, proceeding normally");
8574 }
8575 DisambiguationResult::NeedsClarification {
8576 question,
8577 detection,
8578 } => {
8579 let admission = self
8580 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8581 .await?;
8582 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8583 info!(
8584 ambiguity_type = ?detection.ambiguity_type,
8585 confidence = detection.confidence,
8586 "Input requires clarification"
8587 );
8588
8589 self.commit_root_user_message(input).await?;
8592 self.memory
8593 .add_message(ChatMessage::assistant(&question.question))
8594 .await?;
8595
8596 let status = if awaiting_confirmation {
8597 "awaiting_confirmation"
8598 } else {
8599 "awaiting_clarification"
8600 };
8601 let response = AgentResponse::new(&question.question).with_metadata(
8602 "disambiguation",
8603 serde_json::json!({
8604 "status": status,
8605 "options": question.options,
8606 "clarifying": question.clarifying,
8607 "detection": {
8608 "type": detection.ambiguity_type,
8609 "confidence": detection.confidence,
8610 "what_is_unclear": detection.what_is_unclear,
8611 }
8612 }),
8613 );
8614 drop(admission);
8615 self.finish_turn_if_root(&response).await?;
8616 return Ok(response);
8617 }
8618 DisambiguationResult::Clarified {
8619 enriched_input,
8620 resolved,
8621 ..
8622 } => {
8623 let admission = match self
8624 .admit_disambiguation_redispatch(disambiguation_epoch, state_generation)
8625 .await
8626 {
8627 Ok(admission) => admission,
8628 Err(error) => {
8629 *self.pending_skill_id.write() = None;
8630 return Err(error);
8631 }
8632 };
8633 info!(
8634 resolved_count = resolved.len(),
8635 enriched = %enriched_input,
8636 "Input clarified, injecting resolved intent into context"
8637 );
8638
8639 for (key, value) in &resolved {
8642 let context_key = format!("disambiguation.{}", key);
8643 let _ = self.context_manager.set(&context_key, value.clone());
8644 }
8645
8646 if let Some(intent) = resolved.get("intent") {
8647 let _ = self.context_manager.set("resolved_intent", intent.clone());
8648 }
8649
8650 let _ = self
8651 .context_manager
8652 .set("disambiguation.resolved", serde_json::Value::Bool(true));
8653
8654 let skill_id = self.pending_skill_id.read().clone();
8658 if let Some(skill_id) = skill_id {
8659 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input");
8660 drop(admission);
8661 return self
8662 .recheck_skill_disambiguation(
8663 &skill_id,
8664 &enriched_input,
8665 disambiguation_epoch,
8666 state_generation,
8667 )
8668 .await;
8669 }
8670
8671 drop(admission);
8672 return self.run_loop_internal(&enriched_input).await;
8673 }
8674 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
8675 info!("Proceeding with best guess interpretation");
8676
8677 let skill_id = self.pending_skill_id.read().clone();
8679 if let Some(skill_id) = skill_id {
8680 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input");
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 return self.run_loop_internal(&enriched_input).await;
8692 }
8693 DisambiguationResult::GiveUp { reason } => {
8694 *self.pending_skill_id.write() = None;
8695 warn!(reason = %reason, "Disambiguation gave up");
8696 let apology = self
8697 .generate_localized_apology(
8698 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
8699 &reason,
8700 )
8701 .await
8702 .unwrap_or_else(|_| {
8703 format!("I'm sorry, I couldn't understand your request: {}", reason)
8704 });
8705 let response = AgentResponse::new(&apology);
8706 self.finish_turn_if_root(&response).await?;
8707 return Ok(response);
8708 }
8709 DisambiguationResult::Escalate { reason } => {
8710 *self.pending_skill_id.write() = None;
8711 info!(reason = %reason, "Escalating to human");
8712 if let Some(ref hitl) = self.hitl_engine {
8713 let trigger =
8714 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
8715 let mut context_map = HashMap::new();
8716 context_map.insert("original_input".to_string(), serde_json::json!(input));
8717 context_map.insert("reason".to_string(), serde_json::json!(&reason));
8718 let check_result = HITLCheckResult::required(
8719 trigger,
8720 context_map,
8721 format!("User request needs human assistance: {}", reason),
8722 Some(hitl.config().default_timeout_seconds),
8723 );
8724 let result = self.request_hitl_approval(check_result).await?;
8725 if matches!(
8726 result,
8727 ApprovalResult::Approved | ApprovalResult::Modified { .. }
8728 ) {
8729 return self.run_loop_internal(input).await;
8730 }
8731 }
8732 let apology = self
8733 .generate_localized_apology(
8734 "Explain briefly that you're transferring the user to a human agent for help.",
8735 &reason,
8736 )
8737 .await
8738 .unwrap_or_else(|_| {
8739 format!("I need human assistance to help with your request: {}", reason)
8740 });
8741 let response = AgentResponse::new(&apology);
8742 self.finish_turn_if_root(&response).await?;
8743 return Ok(response);
8744 }
8745 DisambiguationResult::Abandoned { new_input } => {
8746 *self.pending_skill_id.write() = None;
8747
8748 info!(
8749 has_new_input = new_input.is_some(),
8750 "Clarification abandoned by user"
8751 );
8752
8753 self.commit_root_user_message(input).await?;
8754
8755 match new_input {
8756 Some(fresh_input) => {
8757 return self.run_loop_internal(&fresh_input).await;
8760 }
8761 None => {
8762 let ack = self
8764 .generate_localized_apology(
8765 "The user changed their mind about their previous request. \
8766 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
8767 Do NOT apologize excessively. Be concise.",
8768 "User abandoned clarification",
8769 )
8770 .await
8771 .unwrap_or_else(|_| {
8772 "OK, no problem. What else can I help with?".to_string()
8773 });
8774
8775 self.memory
8776 .add_message(ChatMessage::assistant(&ack))
8777 .await?;
8778
8779 let response = AgentResponse::new(&ack);
8780 self.finish_turn_if_root(&response).await?;
8781 return Ok(response);
8782 }
8783 }
8784 }
8785 }
8786 }
8787
8788 self.run_loop_internal(input).await
8789 }
8790
8791 async fn generate_localized_apology(&self, instruction: &str, reason: &str) -> Result<String> {
8793 let llm = self.llm_registry.router().map_err(|e| {
8794 AgentError::LLM(format!(
8795 "Router LLM not available for localized response: {}",
8796 e
8797 ))
8798 })?;
8799
8800 let recent: Vec<String> = self
8801 .memory
8802 .get_messages(Some(3))
8803 .await?
8804 .iter()
8805 .map(|m| m.content.clone())
8806 .collect();
8807
8808 let context_hint = if recent.is_empty() {
8809 String::new()
8810 } else {
8811 format!(
8812 "\nRecent conversation (detect the user's language from this):\n{}\n",
8813 recent.join("\n")
8814 )
8815 };
8816
8817 let prompt = format!(
8818 "{}\nReason: {}\n{}Respond in the same language as the user. Output ONLY the message, nothing else.",
8819 instruction, reason, context_hint
8820 );
8821
8822 let messages = vec![ChatMessage::user(&prompt)];
8823 let response = self
8824 .observe_purpose(
8825 ObservationPurpose::DisambiguationClarification,
8826 llm.complete(&messages, None),
8827 )
8828 .await
8829 .map_err(|e| AgentError::LLM(format!("Localized response generation failed: {}", e)))?;
8830
8831 Ok(response.content.trim().to_string())
8832 }
8833
8834 fn render_action_args(&self, args: &Value) -> Value {
8838 let context = self.build_context_with_overlays();
8839 match args {
8840 Value::Object(map) => {
8841 let mut rendered = serde_json::Map::new();
8842 for (k, v) in map {
8843 match v {
8844 Value::String(s) if s.contains("{{") => {
8845 match self.template_renderer.render(s, &context) {
8846 Ok(rendered_str) => {
8847 rendered.insert(k.clone(), Value::String(rendered_str));
8848 }
8849 Err(_) => {
8850 rendered.insert(k.clone(), v.clone());
8851 }
8852 }
8853 }
8854 _ => {
8855 rendered.insert(k.clone(), v.clone());
8856 }
8857 }
8858 }
8859 Value::Object(rendered)
8860 }
8861 _ => args.clone(),
8862 }
8863 }
8864
8865 fn clear_disambiguation_context(&self) {
8867 let _ = self
8868 .context_manager
8869 .set("resolved_intent", serde_json::Value::Null);
8870
8871 let all = self.context_manager.get_all();
8872 for key in all.keys() {
8873 if key.starts_with("disambiguation.") {
8874 let _ = self.context_manager.set(key, serde_json::Value::Null);
8875 }
8876 }
8877 }
8878
8879 async fn recheck_skill_disambiguation(
8885 &self,
8886 skill_id: &str,
8887 enriched_input: &str,
8888 expected_disambiguation_epoch: u64,
8889 expected_state_generation: Option<u64>,
8890 ) -> Result<AgentResponse> {
8891 let skill = self
8892 .skill_router
8893 .as_ref()
8894 .and_then(|r| r.get_skill(skill_id).cloned());
8895
8896 if let Some(ref skill) = skill
8898 && let Some(ref skill_disambig) = skill.disambiguation
8899 && skill_disambig.enabled.unwrap_or(false)
8900 && let Some(ref disambiguator) = self.disambiguation_manager
8901 {
8902 let context = self.build_disambiguation_context().await?;
8903 let state_override = self
8904 .state_machine
8905 .as_ref()
8906 .and_then(|sm| sm.current_definition())
8907 .and_then(|def| def.disambiguation.clone());
8908
8909 let disambiguation_result = self
8910 .observe_purpose(
8911 ObservationPurpose::DisambiguationDetection,
8912 disambiguator.process_input_with_override(
8913 enriched_input,
8914 &context,
8915 state_override.as_ref(),
8916 Some(skill_disambig),
8917 ),
8918 )
8919 .await?;
8920 let current_state_generation = self
8921 .state_machine
8922 .as_ref()
8923 .map(|state_machine| state_machine.generation());
8924 if current_state_generation != expected_state_generation
8925 || self.disambiguation_epoch.load(Ordering::SeqCst) != expected_disambiguation_epoch
8926 {
8927 disambiguator.clear_pending().await;
8928 *self.pending_skill_id.write() = None;
8929 return Err(AgentError::Other(
8930 "State or reset ownership changed during skill disambiguation recheck"
8931 .to_string(),
8932 ));
8933 }
8934 match disambiguation_result {
8935 DisambiguationResult::Clear => {
8936 debug!(skill_id = %skill_id, "Skill re-check: all fields present");
8937 }
8938 DisambiguationResult::NeedsClarification {
8939 question,
8940 detection,
8941 } => {
8942 let admission = self
8943 .admit_disambiguation_redispatch(
8944 expected_disambiguation_epoch,
8945 expected_state_generation,
8946 )
8947 .await?;
8948 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
8949 info!(
8950 skill_id = %skill_id,
8951 ambiguity_type = ?detection.ambiguity_type,
8952 what_is_unclear = ?detection.what_is_unclear,
8953 "Skill re-check: still missing fields, asking again"
8954 );
8955 self.memory
8959 .add_message(ChatMessage::user(enriched_input))
8960 .await?;
8961 self.memory
8962 .add_message(ChatMessage::assistant(&question.question))
8963 .await?;
8964
8965 let response = AgentResponse::new(&question.question).with_metadata(
8966 "disambiguation",
8967 serde_json::json!({
8968 "status": if awaiting_confirmation { "awaiting_confirmation" } else { "awaiting_clarification" },
8969 "skill_id": skill_id,
8970 "options": question.options,
8971 "clarifying": question.clarifying,
8972 "detection": {
8973 "type": detection.ambiguity_type,
8974 "confidence": detection.confidence,
8975 "what_is_unclear": detection.what_is_unclear,
8976 }
8977 }),
8978 );
8979 drop(admission);
8980 self.finish_turn_if_root(&response).await?;
8981 return Ok(response);
8982 }
8983 DisambiguationResult::Clarified {
8984 enriched_input: re_enriched,
8985 ..
8986 } => {
8987 debug!(skill_id = %skill_id, "Skill re-check: clarified immediately, executing");
8988 let admission = self
8989 .admit_disambiguation_redispatch(
8990 expected_disambiguation_epoch,
8991 expected_state_generation,
8992 )
8993 .await?;
8994 *self.pending_skill_id.write() = None;
8995 drop(admission);
8996 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
8997 self.memory
8998 .add_message(ChatMessage::user(&re_enriched))
8999 .await?;
9000 return self
9001 .handle_skill_response(
9002 &re_enriched,
9003 skill_id,
9004 skill_response,
9005 &HashMap::new(),
9006 )
9007 .await;
9008 }
9009 DisambiguationResult::ProceedWithBestGuess {
9010 enriched_input: re_enriched,
9011 } => {
9012 debug!(skill_id = %skill_id, "Skill re-check: proceeding with best guess");
9013 let admission = self
9014 .admit_disambiguation_redispatch(
9015 expected_disambiguation_epoch,
9016 expected_state_generation,
9017 )
9018 .await?;
9019 *self.pending_skill_id.write() = None;
9020 drop(admission);
9021 let skill_response = self.execute_skill_by_id(skill_id, &re_enriched).await?;
9022 self.memory
9023 .add_message(ChatMessage::user(&re_enriched))
9024 .await?;
9025 return self
9026 .handle_skill_response(
9027 &re_enriched,
9028 skill_id,
9029 skill_response,
9030 &HashMap::new(),
9031 )
9032 .await;
9033 }
9034 DisambiguationResult::GiveUp { reason } => {
9035 *self.pending_skill_id.write() = None;
9036 let apology = self
9037 .generate_localized_apology(
9038 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
9039 &reason,
9040 )
9041 .await
9042 .unwrap_or_else(|_| {
9043 format!("I'm sorry, I couldn't understand your request: {}", reason)
9044 });
9045 let response = AgentResponse::new(&apology);
9046 self.finish_turn_if_root(&response).await?;
9047 return Ok(response);
9048 }
9049 DisambiguationResult::Escalate { reason } => {
9050 *self.pending_skill_id.write() = None;
9051 let apology = self
9052 .generate_localized_apology(
9053 "Explain briefly that you're transferring the user to a human agent for help.",
9054 &reason,
9055 )
9056 .await
9057 .unwrap_or_else(|_| {
9058 format!("I need human assistance to help with your request: {}", reason)
9059 });
9060 let response = AgentResponse::new(&apology);
9061 self.finish_turn_if_root(&response).await?;
9062 return Ok(response);
9063 }
9064 DisambiguationResult::Abandoned { new_input } => {
9065 *self.pending_skill_id.write() = None;
9068 debug!(skill_id = %skill_id, "Skill re-check: abandoned by user");
9069 if let Some(fresh) = new_input {
9070 return self.run_loop_internal(&fresh).await;
9071 }
9072 let ack = self
9073 .generate_localized_apology(
9074 "The user changed their mind about their previous request. \
9075 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
9076 Do NOT apologize excessively. Be concise.",
9077 "User abandoned clarification",
9078 )
9079 .await
9080 .unwrap_or_else(|_| {
9081 "OK, no problem. What else can I help with?".to_string()
9082 });
9083 self.memory
9084 .add_message(ChatMessage::assistant(&ack))
9085 .await?;
9086 let response = AgentResponse::new(&ack);
9087 self.finish_turn_if_root(&response).await?;
9088 return Ok(response);
9089 }
9090 }
9091 }
9092
9093 let admission = self
9095 .admit_disambiguation_redispatch(
9096 expected_disambiguation_epoch,
9097 expected_state_generation,
9098 )
9099 .await?;
9100 *self.pending_skill_id.write() = None;
9101 drop(admission);
9102 let skill_response = self.execute_skill_by_id(skill_id, enriched_input).await?;
9103 self.memory
9104 .add_message(ChatMessage::user(enriched_input))
9105 .await?;
9106 self.handle_skill_response(enriched_input, skill_id, skill_response, &HashMap::new())
9107 .await
9108 }
9109
9110 async fn handle_skill_response(
9113 &self,
9114 processed_input: &str,
9115 skill_id: &str,
9116 skill_response: String,
9117 input_context: &HashMap<String, Value>,
9118 ) -> Result<AgentResponse> {
9119 let output_data = self.process_output(&skill_response, input_context).await?;
9120 let final_response = output_data.content;
9121
9122 self.memory
9123 .add_message(ChatMessage::assistant(&final_response))
9124 .await?;
9125
9126 self.check_memory_compression().await?;
9127
9128 self.increment_turn();
9129 self.evaluate_transitions(processed_input, &final_response)
9130 .await?;
9131
9132 let response = AgentResponse::new(final_response)
9133 .with_metadata("skill_id", serde_json::json!(skill_id));
9134 self.finish_turn_if_root(&response).await?;
9135 Ok(response)
9136 }
9137
9138 async fn handle_plan_and_execute(
9141 &self,
9142 processed_input: &str,
9143 input_context: &HashMap<String, Value>,
9144 auto_detected: bool,
9145 ) -> Result<AgentResponse> {
9146 let effective = self.get_effective_reasoning_config();
9147 let plan_reflection = effective
9148 .get_planning()
9149 .map(|c| c.reflection.clone())
9150 .unwrap_or_default();
9151
9152 let max_attempts = if plan_reflection.enabled {
9153 1 + plan_reflection.max_replans
9154 } else {
9155 1
9156 };
9157
9158 let mut plan = self.generate_plan(processed_input).await?;
9159 info!(
9160 plan_id = %plan.id,
9161 steps = plan.steps.len(),
9162 "Plan generated"
9163 );
9164
9165 let mut plan_result = String::new();
9166
9167 for attempt in 0..max_attempts {
9168 *self.current_plan.write() = Some(plan.clone());
9169 plan_result = self.execute_plan(&mut plan).await?;
9170
9171 info!(
9172 plan_status = ?plan.status,
9173 completed_steps = plan.completed_steps().count(),
9174 attempt = attempt + 1,
9175 "Plan execution completed"
9176 );
9177
9178 if !plan_reflection.enabled {
9179 break;
9180 }
9181
9182 let has_failures = plan.steps.iter().any(|s| s.status.is_failed());
9183 if !has_failures {
9184 break;
9185 }
9186
9187 if attempt + 1 >= max_attempts {
9188 break;
9189 }
9190
9191 match plan_reflection.on_step_failure {
9192 StepFailureAction::Replan => {
9193 info!(attempt = attempt + 1, "Plan had failures, replanning");
9194 plan = self.generate_plan(processed_input).await?;
9195 }
9196 StepFailureAction::Abort => {
9197 warn!("Plan step failed, aborting");
9198 break;
9199 }
9200 StepFailureAction::Skip | StepFailureAction::Continue => {
9201 break;
9202 }
9203 }
9204 }
9205
9206 *self.current_plan.write() = Some(plan);
9207
9208 let output_data = self.process_output(&plan_result, input_context).await?;
9209 let final_content = output_data.content;
9210
9211 self.memory
9212 .add_message(ChatMessage::assistant(&final_content))
9213 .await?;
9214
9215 self.check_memory_compression().await?;
9216 self.increment_turn();
9217 self.evaluate_transitions(processed_input, &final_content)
9218 .await?;
9219
9220 let reasoning_metadata =
9221 ReasoningMetadata::new(ReasoningMode::PlanAndExecute).with_auto_detected(auto_detected);
9222
9223 let response = AgentResponse::new(&final_content).with_metadata(
9224 "reasoning",
9225 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9226 );
9227
9228 self.finish_turn_if_root(&response).await?;
9229 Ok(response)
9230 }
9231
9232 fn inject_reasoning_prompt(
9234 &self,
9235 messages: &mut [ChatMessage],
9236 reasoning_mode: &ReasoningMode,
9237 is_first_iteration: bool,
9238 ) {
9239 if !is_first_iteration {
9240 return;
9241 }
9242 match reasoning_mode {
9243 ReasoningMode::CoT => {
9244 if let Some(msg) = messages.first_mut()
9245 && matches!(msg.role, ai_agents_core::Role::System)
9246 {
9247 msg.content = self.build_cot_system_prompt(&msg.content);
9248 debug!("Applied Chain-of-Thought system prompt");
9249 }
9250 }
9251 ReasoningMode::React => {
9252 if let Some(msg) = messages.first_mut()
9253 && matches!(msg.role, ai_agents_core::Role::System)
9254 {
9255 msg.content = self.build_react_system_prompt(&msg.content);
9256 debug!("Applied ReAct system prompt");
9257 }
9258 }
9259 _ => {}
9260 }
9261 }
9262
9263 async fn generate_main_response_draft(
9268 &self,
9269 processed_input: &str,
9270 reasoning_mode: &ReasoningMode,
9271 ) -> Result<MainResponseDraft> {
9272 let llm = self.get_state_llm()?;
9273 let protocol = self.main_tool_protocol(llm.as_ref(), true).await?;
9274 let mut messages = self
9275 .build_messages_internal(false, Some(processed_input), protocol.choice.is_none())
9276 .await?;
9277 self.inject_reasoning_prompt(&mut messages, reasoning_mode, true);
9278 let response = self
9279 .complete_main_llm_with_recovery(llm, &messages, &protocol)
9280 .await?;
9281 let content = response.content.trim().to_string();
9282 let (thinking, answer) = self.extract_thinking(&content);
9283 if let Some(calls) = self.parse_main_tool_calls(&content, &protocol) {
9284 return Ok(MainResponseDraft::ToolCalls {
9285 raw_content: content,
9286 calls,
9287 thinking,
9288 });
9289 }
9290 Ok(MainResponseDraft::Text {
9291 raw_content: answer,
9292 thinking,
9293 })
9294 }
9295
9296 async fn commit_main_response_draft(
9301 &self,
9302 processed_input: &str,
9303 input_context: &HashMap<String, Value>,
9304 draft: MainResponseDraft,
9305 reasoning_mode: ReasoningMode,
9306 auto_detected: bool,
9307 ) -> Result<AgentResponse> {
9308 self.commit_root_user_message(processed_input).await?;
9309 match draft {
9310 MainResponseDraft::Text {
9311 raw_content,
9312 thinking,
9313 } => {
9314 self.finish_text_response_from_model(CommittedTextResponse {
9315 processed_input,
9316 input_context,
9317 answer: raw_content,
9318 reasoning_mode,
9319 auto_detected,
9320 iterations: 1,
9321 thinking_content: thinking,
9322 all_tool_calls: Vec::new(),
9323 })
9324 .await
9325 }
9326 MainResponseDraft::ToolCalls {
9327 raw_content,
9328 calls,
9329 thinking: _,
9330 } => {
9331 let mut all_tool_calls = Vec::new();
9332 match self
9333 .handle_tool_calls(processed_input, &raw_content, calls, &mut all_tool_calls)
9334 .await?
9335 {
9336 ToolCallOutcome::Rejected(response) => {
9337 self.finish_turn_if_root(&response).await?;
9338 Ok(response)
9339 }
9340 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => {
9341 self.continue_after_committed_tool_draft(processed_input)
9342 .await
9343 }
9344 }
9345 }
9346 }
9347 }
9348
9349 async fn continue_after_committed_tool_draft(
9354 &self,
9355 processed_input: &str,
9356 ) -> Result<AgentResponse> {
9357 *self.redispatch_depth.write() += 1;
9358 if let Some(context) = self.active_turn_context.write().as_mut() {
9359 context.enter_redispatch();
9360 }
9361 let result = Box::pin(self.run_loop_internal(processed_input)).await;
9362 *self.redispatch_depth.write() -= 1;
9363 if let Some(context) = self.active_turn_context.write().as_mut() {
9364 context.exit_redispatch();
9365 }
9366 let response = result?;
9367 self.finish_turn_if_root(&response).await?;
9368 Ok(response)
9369 }
9370
9371 async fn finish_text_response_from_model(
9376 &self,
9377 response: CommittedTextResponse<'_>,
9378 ) -> Result<AgentResponse> {
9379 let CommittedTextResponse {
9380 processed_input,
9381 input_context,
9382 answer,
9383 reasoning_mode,
9384 auto_detected,
9385 iterations,
9386 thinking_content,
9387 all_tool_calls,
9388 } = response;
9389 let output_data = self.process_output(&answer, input_context).await?;
9390 let mut final_content = if output_data.metadata.rejected {
9391 output_data
9392 .metadata
9393 .rejection_reason
9394 .unwrap_or_else(|| answer.to_string())
9395 } else {
9396 output_data.content
9397 };
9398 let llm = self.get_state_llm()?;
9399 let reflection_metadata;
9400 (final_content, reflection_metadata) = self
9401 .run_reflection(&*llm, processed_input, final_content)
9402 .await?;
9403 final_content =
9404 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
9405 let final_content = {
9406 let result = self
9407 .post_loop_processing(processed_input, final_content)
9408 .await?;
9409 self.apply_post_loop_result(processed_input, result).await?
9410 };
9411 let response = self.build_agent_response(AgentResponseParts {
9412 content: final_content,
9413 all_tool_calls,
9414 reasoning_mode,
9415 auto_detected,
9416 iterations,
9417 thinking: thinking_content,
9418 reflection_metadata,
9419 });
9420 self.finish_turn_if_root(&response).await?;
9421 Ok(response)
9422 }
9423
9424 async fn run_committed_response_loop_with_reasoning(
9429 &self,
9430 processed_input: &str,
9431 input_context: &HashMap<String, Value>,
9432 reasoning_mode: ReasoningMode,
9433 auto_detected: bool,
9434 ) -> Result<AgentResponse> {
9435 self.commit_root_user_message(processed_input).await?;
9436 let llm = self.get_state_llm()?;
9437 let mut iterations = 0u32;
9438 let mut all_tool_calls = Vec::new();
9439 let mut thinking_content = None;
9440 loop {
9441 let effective_max = if reasoning_mode != ReasoningMode::None {
9442 let rc = self.get_effective_reasoning_config();
9443 self.max_iterations.min(rc.max_iterations)
9444 } else {
9445 self.max_iterations
9446 };
9447 if iterations >= effective_max {
9448 return Err(AgentError::Other(format!(
9449 "Max iterations ({}) exceeded",
9450 effective_max
9451 )));
9452 }
9453 iterations += 1;
9454 *self.iteration_count.write() = iterations;
9455 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
9456 let mut messages = self
9457 .build_messages_internal(true, None, protocol.choice.is_none())
9458 .await?;
9459 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
9460 self.hooks.on_llm_start(&messages).await;
9461 let llm_start = Instant::now();
9462 let response = self
9463 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
9464 .await?;
9465 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
9466 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
9467 let content = response.content.trim();
9468 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
9469 match self
9470 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
9471 .await?
9472 {
9473 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
9474 ToolCallOutcome::Rejected(resp) => {
9475 self.finish_turn_if_root(&resp).await?;
9476 return Ok(resp);
9477 }
9478 }
9479 }
9480 let (extracted_thinking, answer) = self.extract_thinking(content);
9481 if extracted_thinking.is_some() {
9482 thinking_content = extracted_thinking;
9483 }
9484 return self
9485 .finish_text_response_from_model(CommittedTextResponse {
9486 processed_input,
9487 input_context,
9488 answer,
9489 reasoning_mode,
9490 auto_detected,
9491 iterations,
9492 thinking_content,
9493 all_tool_calls,
9494 })
9495 .await;
9496 }
9497 }
9498
9499 async fn handle_tool_calls(
9501 &self,
9502 processed_input: &str,
9503 content: &str,
9504 tool_calls: Vec<ToolCall>,
9505 all_tool_calls: &mut Vec<ToolCall>,
9506 ) -> Result<ToolCallOutcome> {
9507 let transition_fired = self.evaluate_transitions(processed_input, content).await?;
9511 if transition_fired {
9512 self.memory
9513 .add_message(ChatMessage::assistant(
9514 "(Transitioned to new state — tool call handled by workflow)",
9515 ))
9516 .await?;
9517 return Ok(ToolCallOutcome::TransitionFired);
9518 }
9519
9520 self.memory
9522 .add_message(ChatMessage::assistant(content))
9523 .await?;
9524 let native_tool_call = Self::is_native_tool_call_content(content);
9525
9526 let results = self.execute_tools_parallel(&tool_calls).await;
9527
9528 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9529 match result {
9530 Ok(output) => {
9531 self.memory
9532 .add_message(Self::tool_result_message(
9533 tool_call,
9534 &output,
9535 native_tool_call,
9536 ))
9537 .await?;
9538 }
9539 Err(e) => {
9540 if matches!(e, AgentError::HITLRejected(_)) {
9542 self.memory
9543 .add_message(ChatMessage::assistant(format!(
9544 "The operation was rejected by the approver: {}",
9545 e
9546 )))
9547 .await?;
9548 return Ok(ToolCallOutcome::Rejected(AgentResponse {
9550 content: format!("Operation cancelled: {}", e),
9551 metadata: None,
9552 tool_calls: Some(all_tool_calls.clone()),
9553 }));
9554 }
9555 self.memory
9556 .add_message(Self::tool_result_message(
9557 tool_call,
9558 &format!("Error: {}", e),
9559 native_tool_call,
9560 ))
9561 .await?;
9562 }
9563 }
9564 all_tool_calls.push(tool_call.clone());
9565 }
9566 Ok(ToolCallOutcome::Continue)
9567 }
9568
9569 async fn run_reflection(
9571 &self,
9572 llm: &dyn LLMProvider,
9573 processed_input: &str,
9574 mut content: String,
9575 ) -> Result<(String, Option<ReflectionMetadata>)> {
9576 let should_reflect = self.should_reflect(processed_input, &content).await?;
9577 if !should_reflect {
9578 return Ok((content, None));
9579 }
9580
9581 info!("Starting response reflection evaluation");
9582 let mut attempts = 0u32;
9583 let max_retries = self.reflection_config.max_retries;
9584 let mut history: Vec<ReflectionAttempt> = Vec::new();
9585
9586 loop {
9587 let evaluation = self.evaluate_response(processed_input, &content).await?;
9588
9589 if evaluation.passed || attempts >= max_retries {
9590 info!(
9591 passed = evaluation.passed,
9592 confidence = evaluation.confidence,
9593 attempts = attempts + 1,
9594 "Reflection evaluation complete"
9595 );
9596 let reflection_metadata = Some(
9597 ReflectionMetadata::new(evaluation)
9598 .with_attempts(attempts + 1)
9599 .with_history(history),
9600 );
9601 return Ok((content, reflection_metadata));
9602 }
9603
9604 debug!(
9605 attempt = attempts + 1,
9606 failed_criteria = evaluation.failed_criteria().count(),
9607 "Response did not meet criteria, retrying"
9608 );
9609
9610 history.push(
9611 ReflectionAttempt::new(&content, evaluation.clone())
9612 .with_feedback("Response did not meet quality criteria"),
9613 );
9614
9615 let feedback: Vec<String> = evaluation
9616 .failed_criteria()
9617 .map(|c| format!("- {}", c.criterion))
9618 .collect();
9619
9620 let retry_prompt = format!(
9621 "Your previous response did not meet these criteria:\n{}\n\nPlease provide an improved response.",
9622 feedback.join("\n")
9623 );
9624
9625 self.memory
9626 .add_message(ChatMessage::user(&retry_prompt))
9627 .await?;
9628
9629 let retry_messages = self.build_messages().await?;
9630 let retry_response = self
9631 .observe_purpose(
9632 ObservationPurpose::ReflectionEvaluation,
9633 llm.complete(&retry_messages, None),
9634 )
9635 .await
9636 .map_err(|e| AgentError::LLM(e.to_string()))?;
9637
9638 content = retry_response.content.trim().to_string();
9639 attempts += 1;
9640 }
9641 }
9642
9643 async fn post_loop_processing(
9646 &self,
9647 processed_input: &str,
9648 content: String,
9649 ) -> Result<PostLoopResult> {
9650 self.increment_turn();
9655
9656 self.run_context_extractors(processed_input).await;
9658
9659 let transitioned = self.evaluate_transitions(processed_input, &content).await?;
9660
9661 if !transitioned {
9662 self.memory
9663 .add_message(ChatMessage::assistant(&content))
9664 .await?;
9665 self.check_memory_compression().await?;
9666 return Ok(PostLoopResult::NoTransition(content));
9667 }
9668
9669 if !self.should_regenerate_after_transition() {
9671 self.memory
9672 .add_message(ChatMessage::assistant(&content))
9673 .await?;
9674 self.check_memory_compression().await?;
9675 return Ok(PostLoopResult::Transitioned(content));
9676 }
9677
9678 if self.needs_redispatch_for_new_state() {
9682 info!("Post-transition NeedsRedispatch: new state requires full dispatch");
9683 return Ok(PostLoopResult::NeedsRedispatch);
9686 }
9687
9688 self.memory
9691 .add_message(ChatMessage::assistant(&content))
9692 .await?;
9693 self.check_memory_compression().await?;
9694
9695 let new_llm = self.get_state_llm()?;
9701 let mut final_content;
9702
9703 for post_iter in 0..self.max_iterations {
9704 let protocol = self.main_tool_protocol(new_llm.as_ref(), false).await?;
9705 let new_messages = self
9706 .build_messages_internal(true, None, protocol.choice.is_none())
9707 .await?;
9708 if post_iter == 0
9709 && let Some(system_msg) = new_messages.first()
9710 && system_msg.role == ai_agents_core::Role::System
9711 {
9712 debug!(
9713 prompt_preview =
9714 &system_msg.content[system_msg.content.len().saturating_sub(200)..],
9715 "Post-transition system prompt (last 200 chars)"
9716 );
9717 }
9718
9719 let new_response = self
9720 .complete_main_llm_with_recovery(Arc::clone(&new_llm), &new_messages, &protocol)
9721 .await?;
9722 final_content = new_response.content.trim().to_string();
9723
9724 if let Some(tool_calls) = self.parse_main_tool_calls(&final_content, &protocol) {
9727 let native_tool_call = Self::is_native_tool_call_content(&final_content);
9728 debug!(
9729 post_iter = post_iter,
9730 tools = tool_calls.len(),
9731 "Post-transition tool call detected, executing"
9732 );
9733
9734 self.memory
9735 .add_message(ChatMessage::assistant(&final_content))
9736 .await?;
9737
9738 let results = self.execute_tools_parallel(&tool_calls).await;
9739 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
9740 match result {
9741 Ok(output) => {
9742 self.memory
9743 .add_message(Self::tool_result_message(
9744 tool_call,
9745 &output,
9746 native_tool_call,
9747 ))
9748 .await?;
9749 }
9750 Err(e) => {
9751 self.memory
9752 .add_message(Self::tool_result_message(
9753 tool_call,
9754 &format!("Error: {}", e),
9755 native_tool_call,
9756 ))
9757 .await?;
9758 }
9759 }
9760 }
9761 continue;
9763 }
9764
9765 self.memory
9767 .add_message(ChatMessage::assistant(&final_content))
9768 .await?;
9769 return Ok(PostLoopResult::Transitioned(final_content));
9770 }
9771
9772 final_content = "Post-transition processing completed.".to_string();
9774 self.memory
9775 .add_message(ChatMessage::assistant(&final_content))
9776 .await?;
9777
9778 Ok(PostLoopResult::Transitioned(final_content))
9779 }
9780
9781 fn should_regenerate_after_transition(&self) -> bool {
9784 if let Some(ref sm) = self.state_machine {
9785 if !sm.config().regenerate_on_transition {
9787 return false;
9788 }
9789 if let Some(def) = sm.current_definition()
9791 && let Some(regen) = def.regenerate_on_enter
9792 {
9793 return regen;
9794 }
9795 }
9796 true
9797 }
9798
9799 fn needs_redispatch_for_new_state(&self) -> bool {
9802 if let Some(ref sm) = self.state_machine
9803 && let Some(def) = sm.current_definition()
9804 {
9805 if def.concurrent.is_some()
9806 || def.group_chat.is_some()
9807 || def.pipeline.is_some()
9808 || def.handoff.is_some()
9809 || def.delegate.is_some()
9810 {
9811 return true;
9812 }
9813 let effective = self.get_effective_reasoning_config();
9815 if !matches!(effective.mode, ReasoningMode::None) {
9816 return true;
9817 }
9818 }
9819 false
9820 }
9821
9822 async fn apply_post_loop_result(
9825 &self,
9826 processed_input: &str,
9827 result: PostLoopResult,
9828 ) -> Result<String> {
9829 match result {
9830 PostLoopResult::NoTransition(content) | PostLoopResult::Transitioned(content) => {
9831 Ok(content)
9832 }
9833 PostLoopResult::NeedsRedispatch => {
9834 const MAX_REDISPATCH_DEPTH: u32 = 3;
9835 let current_depth = *self.redispatch_depth.read();
9836 if current_depth >= MAX_REDISPATCH_DEPTH {
9837 warn!(
9838 depth = current_depth,
9839 "Post-transition re-dispatch depth limit reached, returning empty response"
9840 );
9841 let content = String::new();
9842 self.memory
9843 .add_message(ChatMessage::assistant(&content))
9844 .await?;
9845 return Ok(content);
9846 }
9847 *self.redispatch_depth.write() += 1;
9848 if let Some(context) = self.active_turn_context.write().as_mut() {
9849 context.enter_redispatch();
9850 }
9851 info!(
9852 depth = current_depth + 1,
9853 "Re-dispatching for new state after transition"
9854 );
9855 let resp = Box::pin(self.run_loop_internal(processed_input)).await;
9856 *self.redispatch_depth.write() -= 1;
9857 if let Some(context) = self.active_turn_context.write().as_mut() {
9858 context.exit_redispatch();
9859 }
9860 resp.map(|r| r.content)
9861 }
9862 }
9863 }
9864
9865 fn build_agent_response(&self, parts: AgentResponseParts) -> AgentResponse {
9867 let AgentResponseParts {
9868 content,
9869 all_tool_calls,
9870 reasoning_mode,
9871 auto_detected,
9872 iterations,
9873 thinking,
9874 reflection_metadata,
9875 } = parts;
9876 let reasoning_metadata = ReasoningMetadata::new(reasoning_mode.clone())
9877 .with_thinking(thinking.clone().unwrap_or_default())
9878 .with_iterations(iterations)
9879 .with_auto_detected(auto_detected);
9880
9881 let mut response = AgentResponse::new(&content);
9882 if !all_tool_calls.is_empty() {
9883 response = response.with_tool_calls(all_tool_calls);
9884 }
9885
9886 if let Some(state) = self.current_state() {
9887 response = response.with_metadata("current_state", serde_json::json!(state));
9888 }
9889
9890 response = response.with_metadata(
9891 "reasoning",
9892 serde_json::to_value(&reasoning_metadata).unwrap_or_default(),
9893 );
9894
9895 if let Some(ref refl_meta) = reflection_metadata {
9896 response = response.with_metadata(
9897 "reflection",
9898 serde_json::to_value(refl_meta).unwrap_or_default(),
9899 );
9900 }
9901
9902 response
9903 }
9904
9905 async fn handle_delegated_state(
9907 &self,
9908 input: &str,
9909 delegate_id: &str,
9910 state_def: &ai_agents_state::StateDefinition,
9911 ) -> Result<AgentResponse> {
9912 use std::time::Instant;
9913
9914 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
9915 AgentError::Config(format!(
9916 "State delegates to '{}' but no agent registry is configured. \
9917 Add a spawner section with auto_spawn to your YAML.",
9918 delegate_id
9919 ))
9920 })?;
9921
9922 let state_name = self
9923 .state_machine
9924 .as_ref()
9925 .map(|sm| sm.current())
9926 .unwrap_or_else(|| "unknown".to_string());
9927
9928 self.hooks.on_delegate_start(delegate_id, &state_name).await;
9929 let start = Instant::now();
9930
9931 let delegate = registry.get(delegate_id).ok_or_else(|| {
9932 AgentError::Other(format!(
9933 "State '{}' delegates to '{}' but no agent with that ID exists in the registry.",
9934 state_name, delegate_id
9935 ))
9936 })?;
9937
9938 let context_mode = state_def.delegate_context.clone().unwrap_or_default();
9940 let effective_input = self
9941 .observe_purpose(
9942 ObservationPurpose::OrchestrationRouting,
9943 crate::orchestration::context::prepare_delegate_input(
9944 input,
9945 &context_mode,
9946 &*self.memory,
9947 self.llm_registry.get("router").ok().as_deref(),
9948 ),
9949 )
9950 .await?;
9951
9952 let response = delegate
9953 .chat_with_actor_context(&effective_input, self.outbound_actor_context())
9954 .await?;
9955
9956 let duration_ms = start.elapsed().as_millis() as u64;
9957 self.hooks
9958 .on_delegate_complete(delegate_id, &state_name, duration_ms)
9959 .await;
9960
9961 let ctx_key = format!("delegation.{}.last_response", delegate_id);
9963 let _ = self.context_manager.set(
9964 &ctx_key,
9965 serde_json::Value::String(response.content.clone()),
9966 );
9967
9968 let _ = self.context_manager.set(
9970 "orchestration",
9971 serde_json::json!({
9972 "type": "delegate",
9973 "agent": delegate_id,
9974 "state": state_name,
9975 "response": response.content,
9976 "duration_ms": duration_ms,
9977 }),
9978 );
9979
9980 self.commit_root_user_message(input).await?;
9981
9982 let post_result = self
9985 .post_loop_processing(
9986 input,
9987 format!("[Delegated to {}]: {}", delegate_id, response.content),
9988 )
9989 .await?;
9990 let final_content = self.apply_post_loop_result(input, post_result).await?;
9991
9992 let mut result = AgentResponse::new(final_content);
9993
9994 let metadata = serde_json::json!({
9995 "orchestration": {
9996 "type": "delegate",
9997 "agent": delegate_id,
9998 "state": state_name,
9999 "response": response.content,
10000 "duration_ms": duration_ms,
10001 }
10002 });
10003 result.metadata = Some(
10004 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10005 metadata,
10006 )
10007 .unwrap_or_default(),
10008 );
10009
10010 self.finish_turn_if_root(&result).await?;
10011 Ok(result)
10012 }
10013
10014 async fn handle_concurrent_state(
10016 &self,
10017 input: &str,
10018 config: &ai_agents_state::ConcurrentStateConfig,
10019 ) -> Result<AgentResponse> {
10020 use std::time::Instant;
10021
10022 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10023 AgentError::Config(
10024 "Concurrent state requires an agent registry. Add a spawner section.".into(),
10025 )
10026 })?;
10027
10028 let context_mode = config.context_mode.clone().unwrap_or_default();
10033 let context_input = self
10034 .observe_purpose(
10035 ObservationPurpose::OrchestrationRouting,
10036 crate::orchestration::context::prepare_delegate_input(
10037 input,
10038 &context_mode,
10039 &*self.memory,
10040 self.llm_registry.get("router").ok().as_deref(),
10041 ),
10042 )
10043 .await?;
10044
10045 let effective_input = if let Some(ref tmpl) = config.input {
10046 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10047 .unwrap_or_else(|_| context_input.clone())
10048 } else {
10049 context_input
10050 };
10051
10052 let start = Instant::now();
10053
10054 let llm_name = config
10055 .aggregation
10056 .synthesizer_llm
10057 .as_deref()
10058 .unwrap_or("router");
10059 let llm_provider = self.llm_registry.get(llm_name).ok();
10060
10061 let vote_parallelism = if self.runtime_config.optimization.enabled
10062 && self
10063 .runtime_config
10064 .optimization
10065 .parallel_orchestration_vote_extraction
10066 {
10067 Some(self.runtime_config.optimization.max_parallel_runtime_tasks)
10068 } else {
10069 None
10070 };
10071
10072 let result = self
10073 .observe_purpose(
10074 ObservationPurpose::OrchestrationAggregation,
10075 scope_actor_context(
10076 self.outbound_actor_context(),
10077 crate::orchestration::concurrent(
10078 registry,
10079 &effective_input,
10080 &config.agents,
10081 &config.aggregation,
10082 llm_provider.as_deref(),
10083 config.min_required,
10084 config.timeout_ms,
10085 config.on_partial_failure.clone(),
10086 vote_parallelism,
10087 ),
10088 ),
10089 )
10090 .await?;
10091
10092 let duration_ms = start.elapsed().as_millis() as u64;
10093 let agent_ids: Vec<String> = config.agents.iter().map(|a| a.id().to_string()).collect();
10094 let strategy = format!("{:?}", config.aggregation.strategy);
10095 self.hooks
10096 .on_concurrent_complete(&agent_ids, &strategy, duration_ms)
10097 .await;
10098
10099 let _ = self.context_manager.set(
10101 "concurrent.result",
10102 serde_json::Value::String(result.response.content.clone()),
10103 );
10104
10105 let agents_json: Vec<serde_json::Value> = result
10107 .agent_results
10108 .iter()
10109 .map(|ar| {
10110 serde_json::json!({
10111 "id": ar.agent_id,
10112 "response": ar.response.as_ref().map(|r| r.content.as_str()),
10113 "success": ar.success,
10114 "error": ar.error,
10115 "duration_ms": ar.duration_ms,
10116 })
10117 })
10118 .collect();
10119
10120 let _ = self.context_manager.set(
10122 "orchestration",
10123 serde_json::json!({
10124 "type": "concurrent",
10125 "result": result.response.content,
10126 "strategy": strategy,
10127 "agents": agents_json,
10128 "duration_ms": duration_ms,
10129 }),
10130 );
10131
10132 self.commit_root_user_message(input).await?;
10133
10134 let post_result = self
10135 .post_loop_processing(input, result.response.content.clone())
10136 .await?;
10137 let final_content = self.apply_post_loop_result(input, post_result).await?;
10138
10139 let mut response = AgentResponse::new(final_content);
10140 let metadata = serde_json::json!({
10141 "orchestration": {
10142 "type": "concurrent",
10143 "result": result.response.content,
10144 "strategy": strategy,
10145 "agents": agents_json,
10146 "duration_ms": duration_ms,
10147 }
10148 });
10149 response.metadata = Some(
10150 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10151 metadata,
10152 )
10153 .unwrap_or_default(),
10154 );
10155
10156 self.finish_turn_if_root(&response).await?;
10157 Ok(response)
10158 }
10159
10160 async fn handle_group_chat_state(
10162 &self,
10163 input: &str,
10164 config: &ai_agents_state::GroupChatStateConfig,
10165 ) -> Result<AgentResponse> {
10166 use std::time::Instant;
10167
10168 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10169 AgentError::Config(
10170 "Group chat state requires an agent registry. Add a spawner section.".into(),
10171 )
10172 })?;
10173
10174 let start = Instant::now();
10175
10176 let llm_provider = self.llm_registry.get("router").ok();
10177
10178 let context_mode = config.context_mode.clone().unwrap_or_default();
10180 let context_input = self
10181 .observe_purpose(
10182 ObservationPurpose::OrchestrationRouting,
10183 crate::orchestration::context::prepare_delegate_input(
10184 input,
10185 &context_mode,
10186 &*self.memory,
10187 self.llm_registry.get("router").ok().as_deref(),
10188 ),
10189 )
10190 .await?;
10191
10192 let effective_topic = if let Some(ref tmpl) = config.input {
10194 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10195 .unwrap_or_else(|_| context_input.clone())
10196 } else {
10197 context_input
10198 };
10199
10200 let result = self
10201 .observe_purpose(
10202 ObservationPurpose::OrchestrationConversation,
10203 scope_actor_context(
10204 self.outbound_actor_context(),
10205 crate::orchestration::group_chat(
10206 registry,
10207 &effective_topic,
10208 config,
10209 llm_provider.as_deref(),
10210 Some(&*self.hooks),
10211 ),
10212 ),
10213 )
10214 .await?;
10215
10216 let duration_ms = start.elapsed().as_millis() as u64;
10217
10218 let _ = self.context_manager.set(
10220 "group_chat.conclusion",
10221 serde_json::Value::String(result.response.content.clone()),
10222 );
10223
10224 let transcript_json: Vec<serde_json::Value> = result
10226 .transcript
10227 .iter()
10228 .map(|t| {
10229 serde_json::json!({
10230 "speaker": t.speaker,
10231 "round": t.round,
10232 "content": t.content,
10233 })
10234 })
10235 .collect();
10236
10237 let _ = self.context_manager.set(
10239 "orchestration",
10240 serde_json::json!({
10241 "type": "group_chat",
10242 "conclusion": result.response.content,
10243 "transcript": transcript_json,
10244 "rounds": result.rounds_completed,
10245 "termination": result.termination_reason,
10246 "duration_ms": duration_ms,
10247 }),
10248 );
10249
10250 self.commit_root_user_message(input).await?;
10251
10252 let post_result = self
10253 .post_loop_processing(input, result.response.content.clone())
10254 .await?;
10255 let final_content = self.apply_post_loop_result(input, post_result).await?;
10256
10257 let mut response = AgentResponse::new(final_content);
10258 let metadata = serde_json::json!({
10259 "orchestration": {
10260 "type": "group_chat",
10261 "conclusion": result.response.content,
10262 "transcript": transcript_json,
10263 "rounds": result.rounds_completed,
10264 "termination": result.termination_reason,
10265 "duration_ms": duration_ms,
10266 }
10267 });
10268 response.metadata = Some(
10269 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10270 metadata,
10271 )
10272 .unwrap_or_default(),
10273 );
10274
10275 self.finish_turn_if_root(&response).await?;
10276 Ok(response)
10277 }
10278
10279 async fn handle_pipeline_state(
10281 &self,
10282 input: &str,
10283 config: &ai_agents_state::PipelineStateConfig,
10284 ) -> Result<AgentResponse> {
10285 use std::time::Instant;
10286
10287 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10288 AgentError::Config(
10289 "Pipeline state requires an agent registry. Add a spawner section.".into(),
10290 )
10291 })?;
10292
10293 let start = Instant::now();
10294
10295 let stages: Vec<crate::orchestration::PipelineStage> = config
10296 .stages
10297 .iter()
10298 .map(|entry| {
10299 let mut stage = crate::orchestration::PipelineStage::id(entry.id());
10300 if let Some(tmpl) = entry.input() {
10301 stage = stage.with_input(tmpl);
10302 }
10303 stage
10304 })
10305 .collect();
10306
10307 let context_mode = config.context_mode.clone().unwrap_or_default();
10309 let context_input = self
10310 .observe_purpose(
10311 ObservationPurpose::OrchestrationRouting,
10312 crate::orchestration::context::prepare_delegate_input(
10313 input,
10314 &context_mode,
10315 &*self.memory,
10316 self.llm_registry.get("router").ok().as_deref(),
10317 ),
10318 )
10319 .await?;
10320
10321 let context_values = self.build_context_with_overlays();
10322 let result = self
10323 .observe_purpose(
10324 ObservationPurpose::OrchestrationRouting,
10325 scope_actor_context(
10326 self.outbound_actor_context(),
10327 crate::orchestration::pipeline(
10328 registry,
10329 &context_input,
10330 &stages,
10331 config.timeout_ms,
10332 Some(&*self.hooks),
10333 Some(&context_values),
10334 ),
10335 ),
10336 )
10337 .await?;
10338
10339 let duration_ms = start.elapsed().as_millis() as u64;
10340
10341 let _ = self.context_manager.set(
10343 "pipeline.result",
10344 serde_json::Value::String(result.response.content.clone()),
10345 );
10346
10347 let stages_json: Vec<serde_json::Value> = result
10349 .stage_outputs
10350 .iter()
10351 .map(|s| {
10352 serde_json::json!({
10353 "agent_id": s.agent_id,
10354 "output": s.output,
10355 "duration_ms": s.duration_ms,
10356 "skipped": s.skipped,
10357 })
10358 })
10359 .collect();
10360
10361 let _ = self.context_manager.set(
10363 "orchestration",
10364 serde_json::json!({
10365 "type": "pipeline",
10366 "result": result.response.content,
10367 "stages": stages_json,
10368 "duration_ms": duration_ms,
10369 }),
10370 );
10371
10372 self.commit_root_user_message(input).await?;
10373
10374 let post_result = self
10375 .post_loop_processing(input, result.response.content.clone())
10376 .await?;
10377 let final_content = self.apply_post_loop_result(input, post_result).await?;
10378
10379 let mut response = AgentResponse::new(final_content);
10380 let metadata = serde_json::json!({
10381 "orchestration": {
10382 "type": "pipeline",
10383 "result": result.response.content,
10384 "stages": stages_json,
10385 "duration_ms": duration_ms,
10386 }
10387 });
10388 response.metadata = Some(
10389 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10390 metadata,
10391 )
10392 .unwrap_or_default(),
10393 );
10394
10395 self.finish_turn_if_root(&response).await?;
10396 Ok(response)
10397 }
10398
10399 async fn handle_handoff_state(
10401 &self,
10402 input: &str,
10403 config: &ai_agents_state::HandoffStateConfig,
10404 ) -> Result<AgentResponse> {
10405 use std::time::Instant;
10406
10407 let registry = self.spawner_registry.as_ref().ok_or_else(|| {
10408 AgentError::Config(
10409 "Handoff state requires an agent registry. Add a spawner section.".into(),
10410 )
10411 })?;
10412
10413 let llm = self
10414 .llm_registry
10415 .get("router")
10416 .map_err(|_| AgentError::Config("Handoff state requires a router LLM.".into()))?;
10417
10418 let start = Instant::now();
10419
10420 let context_mode = config.context_mode.clone().unwrap_or_default();
10422 let context_input = self
10423 .observe_purpose(
10424 ObservationPurpose::OrchestrationRouting,
10425 crate::orchestration::context::prepare_delegate_input(
10426 input,
10427 &context_mode,
10428 &*self.memory,
10429 self.llm_registry.get("router").ok().as_deref(),
10430 ),
10431 )
10432 .await?;
10433
10434 let effective_input = if let Some(ref tmpl) = config.input {
10436 render_concurrent_template(tmpl, &context_input, &self.build_context_with_overlays())
10437 .unwrap_or_else(|_| context_input.clone())
10438 } else {
10439 context_input
10440 };
10441
10442 let result = self
10443 .observe_purpose(
10444 ObservationPurpose::OrchestrationRouting,
10445 scope_actor_context(
10446 self.outbound_actor_context(),
10447 crate::orchestration::handoff(
10448 registry,
10449 &effective_input,
10450 &config.initial_agent,
10451 &config.available_agents,
10452 config.max_handoffs,
10453 llm.as_ref(),
10454 Some(&*self.hooks),
10455 ),
10456 ),
10457 )
10458 .await?;
10459
10460 let duration_ms = start.elapsed().as_millis() as u64;
10461
10462 let _ = self.context_manager.set(
10464 "handoff.result",
10465 serde_json::Value::String(result.response.content.clone()),
10466 );
10467
10468 let chain_json: Vec<serde_json::Value> = result
10470 .handoff_chain
10471 .iter()
10472 .map(|h| {
10473 serde_json::json!({
10474 "from": h.from_agent,
10475 "to": h.to_agent,
10476 "reason": h.reason,
10477 })
10478 })
10479 .collect();
10480
10481 let _ = self.context_manager.set(
10483 "orchestration",
10484 serde_json::json!({
10485 "type": "handoff",
10486 "result": result.response.content,
10487 "final_agent": result.final_agent,
10488 "handoff_chain": chain_json,
10489 "duration_ms": duration_ms,
10490 }),
10491 );
10492
10493 self.commit_root_user_message(input).await?;
10494
10495 let post_result = self
10496 .post_loop_processing(input, result.response.content.clone())
10497 .await?;
10498 let final_content = self.apply_post_loop_result(input, post_result).await?;
10499
10500 let mut response = AgentResponse::new(final_content);
10501 let metadata = serde_json::json!({
10502 "orchestration": {
10503 "type": "handoff",
10504 "result": result.response.content,
10505 "final_agent": result.final_agent,
10506 "handoff_chain": chain_json,
10507 "duration_ms": duration_ms,
10508 }
10509 });
10510 response.metadata = Some(
10511 serde_json::from_value::<std::collections::HashMap<String, serde_json::Value>>(
10512 metadata,
10513 )
10514 .unwrap_or_default(),
10515 );
10516
10517 self.finish_turn_if_root(&response).await?;
10518 Ok(response)
10519 }
10520
10521 async fn run_loop_internal(&self, input: &str) -> Result<AgentResponse> {
10523 self.begin_root_turn();
10524 self.pre_turn_session_lifecycle().await;
10526
10527 let input_data = self.process_input(input).await?;
10528 self.update_active_turn_context(&input_data.content, input_data.context.clone());
10529
10530 for (key, value) in &input_data.context {
10533 let _ = self.context_manager.set(key, value.clone());
10534 }
10535
10536 if input_data.metadata.rejected {
10537 let reason = input_data
10538 .metadata
10539 .rejection_reason
10540 .unwrap_or_else(|| "Input rejected".to_string());
10541 warn!(reason = %reason, "Input rejected");
10542 let response = AgentResponse::new(reason);
10543 self.finish_turn_if_root(&response).await?;
10544 return Ok(response);
10545 }
10546
10547 let processed_input = &input_data.content;
10548
10549 if let Some(response) = self.try_pre_response_transition(processed_input).await? {
10550 return Ok(response);
10551 }
10552
10553 if let Some(ref sm) = self.state_machine
10555 && let Some(def) = sm.current_definition()
10556 {
10557 if let Some(ref delegate_id) = def.delegate {
10558 return self
10559 .handle_delegated_state(processed_input, delegate_id, &def)
10560 .await;
10561 }
10562 if let Some(ref concurrent_config) = def.concurrent {
10563 return self
10564 .handle_concurrent_state(processed_input, concurrent_config)
10565 .await;
10566 }
10567 if let Some(ref group_chat_config) = def.group_chat {
10568 return self
10569 .handle_group_chat_state(processed_input, group_chat_config)
10570 .await;
10571 }
10572 if let Some(ref pipeline_config) = def.pipeline {
10573 return self
10574 .handle_pipeline_state(processed_input, pipeline_config)
10575 .await;
10576 }
10577 if let Some(ref handoff_config) = def.handoff {
10578 return self
10579 .handle_handoff_state(processed_input, handoff_config)
10580 .await;
10581 }
10582 }
10583
10584 if let Some(response) =
10589 Box::pin(self.try_speculative_branches(processed_input, &input_data.context)).await?
10590 {
10591 return Ok(response);
10592 }
10593
10594 match self.try_skill_route(processed_input).await? {
10595 SkillRouteResult::Response { skill_id, content } => {
10596 self.commit_root_user_message(processed_input).await?;
10597 return self
10598 .handle_skill_response(processed_input, &skill_id, content, &input_data.context)
10599 .await;
10600 }
10601 SkillRouteResult::NeedsClarification {
10602 response,
10603 ownership,
10604 } => {
10605 let admission = self
10606 .admit_optional_disambiguation_ownership(ownership)
10607 .await?;
10608 self.commit_root_user_message(processed_input).await?;
10609 if let Some(q) = response
10610 .metadata
10611 .as_ref()
10612 .and_then(|m| m.get("disambiguation"))
10613 .and_then(|d| d.get("status"))
10614 .and_then(|s| s.as_str())
10615 && q == "awaiting_clarification"
10616 {
10617 self.memory
10620 .add_message(ChatMessage::assistant(&response.content))
10621 .await?;
10622 }
10623 drop(admission);
10624 self.finish_turn_if_root(&response).await?;
10625 return Ok(response);
10626 }
10627 SkillRouteResult::NoMatch => {} }
10629
10630 let effective_reasoning = self.get_effective_reasoning_config();
10631 let reasoning_mode = self.determine_reasoning_mode(processed_input).await?;
10632 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
10633
10634 info!(
10635 reasoning_mode = ?reasoning_mode,
10636 auto_detected = auto_detected,
10637 reflection_enabled = ?self.reflection_config.enabled,
10638 "Reasoning mode determined"
10639 );
10640
10641 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
10642 self.commit_root_user_message(processed_input).await?;
10643 return self
10644 .handle_plan_and_execute(processed_input, &input_data.context, auto_detected)
10645 .await;
10646 }
10647
10648 self.commit_root_user_message(processed_input).await?;
10649
10650 let mut iterations = 0u32;
10651 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
10652 let mut thinking_content: Option<String> = None;
10653
10654 let llm = self.get_state_llm()?;
10655
10656 loop {
10657 let effective_max = if reasoning_mode != ReasoningMode::None {
10659 let rc = self.get_effective_reasoning_config();
10660 self.max_iterations.min(rc.max_iterations)
10661 } else {
10662 self.max_iterations
10663 };
10664
10665 if iterations >= effective_max {
10666 let err = AgentError::Other(format!("Max iterations ({}) exceeded", effective_max));
10667 self.hooks.on_error(&err).await;
10668 error!(iterations = iterations, "Max iterations exceeded");
10669 return Err(err);
10670 }
10671 iterations += 1;
10672 *self.iteration_count.write() = iterations;
10673
10674 debug!(iteration = iterations, max = effective_max, "LLM call");
10675
10676 let protocol = self.main_tool_protocol(llm.as_ref(), false).await?;
10677 let mut messages = self
10678 .build_messages_internal(true, None, protocol.choice.is_none())
10679 .await?;
10680 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
10681
10682 self.hooks.on_llm_start(&messages).await;
10683 let llm_start = Instant::now();
10684 let response = self
10685 .complete_main_llm_with_recovery(Arc::clone(&llm), &messages, &protocol)
10686 .await?;
10687
10688 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
10689 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
10690
10691 let content = response.content.trim();
10692
10693 if let Some(tool_calls) = self.parse_main_tool_calls(content, &protocol) {
10694 match self
10695 .handle_tool_calls(processed_input, content, tool_calls, &mut all_tool_calls)
10696 .await?
10697 {
10698 ToolCallOutcome::Continue | ToolCallOutcome::TransitionFired => continue,
10699 ToolCallOutcome::Rejected(resp) => {
10700 self.finish_turn_if_root(&resp).await?;
10701 return Ok(resp);
10702 }
10703 }
10704 }
10705
10706 let (extracted_thinking, answer) = self.extract_thinking(content);
10707 if extracted_thinking.is_some() {
10708 thinking_content = extracted_thinking;
10709 }
10710
10711 let output_data = self.process_output(&answer, &input_data.context).await?;
10712
10713 let mut final_content = if output_data.metadata.rejected {
10714 output_data
10715 .metadata
10716 .rejection_reason
10717 .unwrap_or_else(|| answer.to_string())
10718 } else {
10719 output_data.content
10720 };
10721
10722 let reflection_metadata;
10724 (final_content, reflection_metadata) = self
10725 .run_reflection(&*llm, processed_input, final_content)
10726 .await?;
10727
10728 final_content =
10729 self.format_response_with_thinking(thinking_content.as_deref(), &final_content);
10730
10731 let final_content = {
10735 let result = self
10736 .post_loop_processing(processed_input, final_content)
10737 .await?;
10738 self.apply_post_loop_result(processed_input, result).await?
10739 };
10740
10741 let reflected = reflection_metadata.is_some();
10742 let reasoning_mode_debug = format!("{:?}", reasoning_mode);
10743
10744 let response = self.build_agent_response(AgentResponseParts {
10745 content: final_content,
10746 all_tool_calls,
10747 reasoning_mode,
10748 auto_detected,
10749 iterations,
10750 thinking: thinking_content,
10751 reflection_metadata,
10752 });
10753
10754 self.finish_turn_if_root(&response).await?;
10755
10756 let tool_call_count = response.tool_calls.as_ref().map(|tc| tc.len()).unwrap_or(0);
10757 info!(
10758 tool_calls = tool_call_count,
10759 response_len = response.content.len(),
10760 reasoning_mode = %reasoning_mode_debug,
10761 reflected = reflected,
10762 "Chat completed"
10763 );
10764 return Ok(response);
10765 }
10766 }
10767
10768 async fn generate_buffered_streaming_draft(
10769 &self,
10770 processed_input: &str,
10771 routing_resolved: Arc<AtomicBool>,
10772 ) -> Result<StreamingDraftResult> {
10773 let llm = self.get_state_llm()?;
10774 if llm.configured_tool_choice().is_some() {
10775 let draft = self
10776 .generate_main_response_draft(processed_input, &ReasoningMode::None)
10777 .await?;
10778 return Ok(StreamingDraftResult::new(draft, Vec::new()));
10779 }
10780 let messages = self.build_messages_for_draft(processed_input).await?;
10781 let mut stream = self
10782 .observe_purpose(
10783 ObservationPurpose::MainResponse,
10784 llm.complete_stream(&messages, None),
10785 )
10786 .await
10787 .map_err(|e| AgentError::LLM(e.to_string()))?;
10788 let mut buffer = crate::optimization::StreamBranchBuffer::new(self.streaming.buffer_size)?;
10789 let mut chunks = Vec::new();
10790 let mut accumulated = String::new();
10791 while let Some(chunk_result) = stream.next().await {
10792 let chunk = chunk_result.map_err(|e| AgentError::LLM(e.to_string()))?;
10793 accumulated.push_str(&chunk.delta);
10794 let stream_chunk = StreamChunk::content(chunk.delta);
10795 if routing_resolved.load(Ordering::SeqCst) {
10796 chunks.push(stream_chunk);
10797 } else {
10798 buffer.push(stream_chunk)?;
10799 }
10800 }
10801 chunks.splice(0..0, buffer.drain());
10802 let content = accumulated.trim().to_string();
10803 let draft = if let Some(calls) = self.parse_tool_calls(&content) {
10804 MainResponseDraft::ToolCalls {
10805 raw_content: content,
10806 calls,
10807 thinking: None,
10808 }
10809 } else {
10810 MainResponseDraft::Text {
10811 raw_content: content,
10812 thinking: None,
10813 }
10814 };
10815 Ok(StreamingDraftResult::new(draft, chunks))
10816 }
10817
10818 async fn try_buffered_streaming_branches(
10819 &self,
10820 processed_input: &str,
10821 input_context: &HashMap<String, Value>,
10822 ) -> Result<Option<(AgentResponse, Vec<StreamChunk>)>> {
10823 let optimization = &self.runtime_config.optimization;
10824 if !optimization.enabled {
10825 return Ok(None);
10826 }
10827 let transition_enabled =
10828 optimization.speculative_state_transitions && self.has_parallel_transition_candidates();
10829 if !transition_enabled {
10830 return Ok(None);
10831 }
10832 let mut branch_scheduler =
10833 TurnBranchScheduler::new(optimization.max_parallel_runtime_tasks)?;
10834 if !branch_scheduler.reserve_task() {
10835 return Ok(None);
10836 }
10837 if !self
10838 .reserve_active_speculative_llm_call(RuntimeOptimizationKind::BufferedStreamingRouting)
10839 {
10840 branch_scheduler.release_task();
10841 return Ok(None);
10842 }
10843 if !branch_scheduler.reserve_task() {
10844 branch_scheduler.release_task();
10845 return Ok(None);
10846 }
10847 let mut main_branch = RuntimeBranch::new(
10848 RuntimeTaskPurpose::MainResponse,
10849 RuntimeOptimizationKind::BufferedStreamingRouting,
10850 RuntimeTaskPriority::Normal,
10851 RuntimeCommitBehavior::FinalResponse,
10852 );
10853 let mut transition_branch = RuntimeBranch::new(
10854 RuntimeTaskPurpose::StateTransition,
10855 RuntimeOptimizationKind::ParallelStateTransition,
10856 RuntimeTaskPriority::Critical,
10857 RuntimeCommitBehavior::TransitionDecision,
10858 );
10859 let main_id = main_branch.branch_id();
10860 let transition_id = transition_branch.branch_id();
10861 let routing_resolved = Arc::new(AtomicBool::new(false));
10862 let mut main_future =
10863 Box::pin(crate::optimization::observability::with_branch_observation(
10864 &main_id,
10865 RuntimeOptimizationKind::BufferedStreamingRouting,
10866 RuntimeCommitBehavior::FinalResponse,
10867 self.generate_buffered_streaming_draft(
10868 processed_input,
10869 Arc::clone(&routing_resolved),
10870 ),
10871 ));
10872 let mut transition_future =
10873 Box::pin(crate::optimization::observability::with_branch_observation(
10874 &transition_id,
10875 RuntimeOptimizationKind::ParallelStateTransition,
10876 RuntimeCommitBehavior::TransitionDecision,
10877 self.select_parallel_transition_candidate(processed_input),
10878 ));
10879 let mut main_pending = true;
10880 let mut transition_pending = true;
10881 let mut main_result: Option<Result<StreamingDraftResult>> = None;
10882 let mut transition_finalized = false;
10883 let mut transition_candidate: Option<TransitionCandidate> = None;
10884 loop {
10885 if let Some(candidate) = transition_candidate.take() {
10886 if self
10887 .approve_transition_target(&candidate.from_state, candidate.target())
10888 .await?
10889 {
10890 drop(main_future);
10892 drop(transition_future);
10893 self.finalize_branch_loss(
10894 &main_id,
10895 RuntimeOptimizationKind::BufferedStreamingRouting,
10896 RuntimeCommitBehavior::FinalResponse,
10897 main_pending,
10898 main_result.as_ref().map(|result| result.is_err()),
10899 );
10900 if !self
10901 .apply_pre_response_transition_candidate(
10902 &candidate,
10903 &HashMap::new(),
10904 processed_input,
10905 )
10906 .await?
10907 {
10908 self.finalize_optional_branch(
10909 &transition_id,
10910 RuntimeOptimizationKind::ParallelStateTransition,
10911 RuntimeCommitBehavior::TransitionDecision,
10912 "discarded",
10913 false,
10914 );
10915 return Ok(None);
10916 }
10917 self.finalize_optional_branch(
10918 &transition_id,
10919 RuntimeOptimizationKind::ParallelStateTransition,
10920 RuntimeCommitBehavior::TransitionDecision,
10921 "committed",
10922 true,
10923 );
10924 let response = self.redispatch_current_state(processed_input).await?;
10925 return Ok(Some((
10926 response.clone(),
10927 vec![StreamChunk::content(response.content)],
10928 )));
10929 }
10930 self.finalize_optional_branch(
10931 &transition_id,
10932 RuntimeOptimizationKind::ParallelStateTransition,
10933 RuntimeCommitBehavior::TransitionDecision,
10934 "discarded",
10935 false,
10936 );
10937 routing_resolved.store(true, Ordering::SeqCst);
10938 transition_finalized = true;
10939 }
10940 if transition_finalized && let Some(result) = main_result.take() {
10941 let stream_draft = match result {
10942 Ok(stream_draft) => stream_draft,
10943 Err(error) => {
10944 self.finalize_optional_branch(
10945 &main_id,
10946 RuntimeOptimizationKind::BufferedStreamingRouting,
10947 RuntimeCommitBehavior::FinalResponse,
10948 "failed",
10949 false,
10950 );
10951 return Err(error);
10952 }
10953 };
10954 let raw_draft_content = stream_draft.draft.raw_content().to_string();
10955 let buffered_chunks = stream_draft.chunks;
10956 self.finalize_optional_branch(
10957 &main_id,
10958 RuntimeOptimizationKind::BufferedStreamingRouting,
10959 RuntimeCommitBehavior::FinalResponse,
10960 "committed",
10961 true,
10962 );
10963 let response = self
10964 .commit_main_response_draft(
10965 processed_input,
10966 input_context,
10967 stream_draft.draft,
10968 ReasoningMode::None,
10969 false,
10970 )
10971 .await?;
10972 let chunks = if response.content == raw_draft_content {
10973 buffered_chunks
10974 } else {
10975 vec![StreamChunk::content(response.content.clone())]
10976 };
10977 return Ok(Some((response, chunks)));
10978 }
10979 tokio::select! {
10980 result = &mut main_future, if main_pending => {
10981 main_pending = false;
10982 main_branch.transition_to(RuntimeBranchStatus::Completed)?;
10983 main_result = Some(result);
10984 }
10985 result = &mut transition_future, if transition_pending => {
10986 transition_pending = false;
10987 transition_branch.transition_to(RuntimeBranchStatus::Completed)?;
10988 match result {
10989 Ok(ParallelTransitionSelection::Candidate(candidate)) => {
10990 transition_candidate = Some(candidate)
10991 }
10992 Ok(ParallelTransitionSelection::NoMatch) => {
10993 self.finalize_optional_branch(
10994 &transition_id,
10995 RuntimeOptimizationKind::ParallelStateTransition,
10996 RuntimeCommitBehavior::TransitionDecision,
10997 "discarded",
10998 false,
10999 );
11000 routing_resolved.store(true, Ordering::SeqCst);
11001 transition_finalized = true;
11002 }
11003 Ok(ParallelTransitionSelection::ReservationExhausted) => {
11004 self.finalize_optional_branch(
11005 &transition_id,
11006 RuntimeOptimizationKind::ParallelStateTransition,
11007 RuntimeCommitBehavior::TransitionDecision,
11008 "cancelled",
11009 false,
11010 );
11011 routing_resolved.store(true, Ordering::SeqCst);
11012 self.finalize_branch_loss(
11013 &main_id,
11014 RuntimeOptimizationKind::BufferedStreamingRouting,
11015 RuntimeCommitBehavior::FinalResponse,
11016 main_pending,
11017 main_result.as_ref().map(|result| result.is_err()),
11018 );
11019 return Ok(None);
11020 }
11021 Err(_) => {
11022 self.finalize_optional_branch(
11023 &transition_id,
11024 RuntimeOptimizationKind::ParallelStateTransition,
11025 RuntimeCommitBehavior::TransitionDecision,
11026 "failed",
11027 false,
11028 );
11029 routing_resolved.store(true, Ordering::SeqCst);
11030 transition_finalized = true;
11031 }
11032 }
11033 }
11034 }
11035 }
11036 }
11037
11038 fn run_loop_internal_stream<'a>(
11042 &'a self,
11043 input: &'a str,
11044 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11045 let include_tool_events = self.streaming.include_tool_events;
11046 let include_state_events = self.streaming.include_state_events;
11047
11048 Box::pin(async_stream::stream! {
11049 self.begin_root_turn();
11050 self.pre_turn_session_lifecycle().await;
11052
11053 let input_data = match self.process_input(input).await {
11054 Ok(data) => data,
11055 Err(e) => {
11056 yield StreamChunk::error(e.to_string());
11057 return;
11058 }
11059 };
11060 self.update_active_turn_context(&input_data.content, input_data.context.clone());
11061
11062 for (key, value) in &input_data.context {
11064 let _ = self.context_manager.set(key, value.clone());
11065 }
11066
11067 if input_data.metadata.rejected {
11068 let reason = input_data
11069 .metadata
11070 .rejection_reason
11071 .unwrap_or_else(|| "Input rejected".to_string());
11072 warn!(reason = %reason, "Input rejected (stream)");
11073 yield StreamChunk::error(reason);
11074 return;
11075 }
11076
11077 let processed_input = &input_data.content;
11078
11079 if self.runtime_config.optimization.enabled
11080 && matches!(
11081 self.runtime_config.optimization.streaming_policy,
11082 crate::optimization::StreamingOptimizationPolicy::BufferUntilRoutingDone
11083 )
11084 {
11085 match Box::pin(self.try_buffered_streaming_branches(processed_input, &input_data.context)).await {
11090 Ok(Some((_response, chunks))) => {
11091 for chunk in chunks {
11092 yield chunk;
11093 }
11094 yield StreamChunk::Done {};
11095 return;
11096 }
11097 Ok(None) => {}
11098 Err(e) => {
11099 yield StreamChunk::error(e.to_string());
11100 return;
11101 }
11102 }
11103 }
11104
11105 if self.runtime_config.optimization.enabled
11106 && matches!(
11107 self.runtime_config.optimization.streaming_policy,
11108 crate::optimization::StreamingOptimizationPolicy::PreflightOnly
11109 )
11110 {
11111 match self.try_pre_response_transition(processed_input).await {
11112 Ok(Some(response)) => {
11113 yield StreamChunk::content(&response.content);
11114 yield StreamChunk::Done {};
11115 return;
11116 }
11117 Ok(None) => {}
11118 Err(e) => {
11119 yield StreamChunk::error(e.to_string());
11120 return;
11121 }
11122 }
11123 }
11124
11125 if let Some(ref sm) = self.state_machine
11127 && let Some(def) = sm.current_definition()
11128 {
11129 let orchestration_result = if let Some(ref delegate_id) = def.delegate {
11130 Some(self.handle_delegated_state(processed_input, delegate_id, &def).await)
11131 } else if let Some(ref concurrent_config) = def.concurrent {
11132 Some(self.handle_concurrent_state(processed_input, concurrent_config).await)
11133 } else if let Some(ref group_chat_config) = def.group_chat {
11134 Some(self.handle_group_chat_state(processed_input, group_chat_config).await)
11135 } else if let Some(ref pipeline_config) = def.pipeline {
11136 Some(self.handle_pipeline_state(processed_input, pipeline_config).await)
11137 } else if let Some(ref handoff_config) = def.handoff {
11138 Some(self.handle_handoff_state(processed_input, handoff_config).await)
11139 } else {
11140 None
11141 };
11142
11143 if let Some(result) = orchestration_result {
11144 match result {
11145 Ok(response) => {
11146 yield StreamChunk::content(&response.content);
11147 yield StreamChunk::Done {};
11148 }
11149 Err(e) => {
11150 yield StreamChunk::error(e.to_string());
11151 }
11152 }
11153 return;
11154 }
11155 }
11156
11157 match self.try_skill_route(processed_input).await {
11159 Ok(SkillRouteResult::Response { skill_id, content }) => {
11160 if let Err(e) = self.commit_root_user_message(processed_input).await {
11161 yield StreamChunk::error(e.to_string());
11162 return;
11163 }
11164 match self.handle_skill_response(processed_input, &skill_id, content, &input_data.context).await {
11165 Ok(resp) => {
11166 yield StreamChunk::content(&resp.content);
11167 yield StreamChunk::Done {};
11168 return;
11169 }
11170 Err(e) => {
11171 yield StreamChunk::error(e.to_string());
11172 return;
11173 }
11174 }
11175 }
11176 Ok(SkillRouteResult::NeedsClarification {
11177 response,
11178 ownership,
11179 }) => {
11180 let admission = match self
11181 .admit_optional_disambiguation_ownership(ownership)
11182 .await
11183 {
11184 Ok(admission) => admission,
11185 Err(e) => {
11186 yield StreamChunk::error(e.to_string());
11187 return;
11188 }
11189 };
11190 if let Err(e) = self.commit_root_user_message(processed_input).await {
11191 yield StreamChunk::error(e.to_string());
11192 return;
11193 }
11194 let _ = self.memory.add_message(ChatMessage::assistant(&response.content)).await;
11195 drop(admission);
11196 if let Err(e) = self.finish_turn_if_root(&response).await {
11197 yield StreamChunk::error(e.to_string());
11198 return;
11199 }
11200 yield StreamChunk::content(&response.content);
11201 yield StreamChunk::Done {};
11202 return;
11203 }
11204 Ok(SkillRouteResult::NoMatch) => {} Err(e) => {
11206 yield StreamChunk::error(e.to_string());
11207 return;
11208 }
11209 }
11210
11211 let effective_reasoning = self.get_effective_reasoning_config();
11213 let reasoning_mode = match self.determine_reasoning_mode(processed_input).await {
11214 Ok(mode) => mode,
11215 Err(e) => {
11216 yield StreamChunk::error(e.to_string());
11217 return;
11218 }
11219 };
11220 let auto_detected = matches!(effective_reasoning.mode, ReasoningMode::Auto);
11221
11222 info!(
11223 reasoning_mode = ?reasoning_mode,
11224 auto_detected = auto_detected,
11225 "Reasoning mode determined (stream)"
11226 );
11227
11228 if matches!(reasoning_mode, ReasoningMode::PlanAndExecute) {
11230 if let Err(e) = self.commit_root_user_message(processed_input).await {
11231 yield StreamChunk::error(e.to_string());
11232 return;
11233 }
11234 match self.handle_plan_and_execute(processed_input, &input_data.context, auto_detected).await {
11235 Ok(resp) => {
11236 yield StreamChunk::content(&resp.content);
11237 yield StreamChunk::Done {};
11238 return;
11239 }
11240 Err(e) => {
11241 yield StreamChunk::error(e.to_string());
11242 return;
11243 }
11244 }
11245 }
11246
11247 if let Err(e) = self.commit_root_user_message(processed_input).await {
11248 yield StreamChunk::error(e.to_string());
11249 return;
11250 }
11251
11252 let llm = match self.get_state_llm() {
11253 Ok(llm) => llm,
11254 Err(e) => {
11255 yield StreamChunk::error(e.to_string());
11256 return;
11257 }
11258 };
11259
11260 let mut iterations = 0u32;
11261 let mut all_tool_calls: Vec<ToolCall> = Vec::new();
11262 let mut thinking_content: Option<String> = None;
11263
11264 loop {
11265 let effective_max = if reasoning_mode != ReasoningMode::None {
11267 let rc = self.get_effective_reasoning_config();
11268 self.max_iterations.min(rc.max_iterations)
11269 } else {
11270 self.max_iterations
11271 };
11272
11273 if iterations >= effective_max {
11274 let err_msg = format!("Max iterations ({}) exceeded", effective_max);
11275 let err = AgentError::Other(err_msg.clone());
11276 self.hooks.on_error(&err).await;
11277 error!(iterations = iterations, "Max iterations exceeded (stream)");
11278 yield StreamChunk::error(err_msg);
11279 return;
11280 }
11281 iterations += 1;
11282 *self.iteration_count.write() = iterations;
11283
11284 debug!(iteration = iterations, max = effective_max, "LLM call (stream)");
11285
11286 let protocol = match self.main_tool_protocol(llm.as_ref(), false).await {
11287 Ok(protocol) => protocol,
11288 Err(e) => {
11289 yield StreamChunk::error(e.to_string());
11290 return;
11291 }
11292 };
11293 let mut messages = match self
11294 .build_messages_internal(true, None, protocol.choice.is_none())
11295 .await
11296 {
11297 Ok(m) => m,
11298 Err(e) => {
11299 yield StreamChunk::error(e.to_string());
11300 return;
11301 }
11302 };
11303 self.inject_reasoning_prompt(&mut messages, &reasoning_mode, iterations == 1);
11304
11305 self.hooks.on_llm_start(&messages).await;
11306 let llm_start = Instant::now();
11307
11308 let reflection_active = self
11311 .should_reflect(processed_input, "")
11312 .await
11313 .unwrap_or_default();
11314
11315 let buffered_decision = reflection_active || protocol.choice.is_some();
11316 let content = if buffered_decision {
11317 let response = match self
11321 .complete_main_llm_with_recovery(
11322 Arc::clone(&llm),
11323 &messages,
11324 &protocol,
11325 )
11326 .await
11327 {
11328 Ok(r) => r,
11329 Err(e) => {
11330 yield StreamChunk::error(e.to_string());
11331 return;
11332 }
11333 };
11334 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11335 self.hooks.on_llm_complete(&response, llm_duration_ms).await;
11336 response.content.trim().to_string()
11337 } else {
11338 let llm_stream = match self
11340 .observe_purpose(
11341 ObservationPurpose::MainResponse,
11342 llm.complete_stream(&messages, None),
11343 )
11344 .await
11345 {
11346 Ok(s) => s,
11347 Err(e) => {
11348 yield StreamChunk::error(e.to_string());
11349 return;
11350 }
11351 };
11352 let mut accumulated = String::new();
11353 let mut stream_inner = llm_stream;
11354 while let Some(chunk_result) = stream_inner.next().await {
11355 match chunk_result {
11356 Ok(chunk) => {
11357 accumulated.push_str(&chunk.delta);
11358 yield StreamChunk::content(chunk.delta);
11359 }
11360 Err(e) => {
11361 yield StreamChunk::error(e.to_string());
11362 return;
11363 }
11364 }
11365 }
11366 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
11367 let llm_response = ai_agents_core::LLMResponse::new(
11369 accumulated.trim(),
11370 ai_agents_core::FinishReason::Stop,
11371 );
11372 self.hooks.on_llm_complete(&llm_response, llm_duration_ms).await;
11373 accumulated.trim().to_string()
11374 };
11375
11376 if let Some(tool_calls) = self.parse_main_tool_calls(&content, &protocol) {
11378 let native_tool_call = Self::is_native_tool_call_content(&content);
11379 let transition_fired = match self.evaluate_transitions(processed_input, &content).await {
11382 Ok(v) => v,
11383 Err(e) => {
11384 yield StreamChunk::error(e.to_string());
11385 return;
11386 }
11387 };
11388 if transition_fired {
11389 let _ = self.memory.add_message(ChatMessage::assistant(
11390 "(Transitioned to new state — tool call handled by workflow)",
11391 )).await;
11392
11393 if include_state_events
11394 && let Some(state) = self.current_state()
11395 {
11396 yield StreamChunk::state_transition(None, state);
11397 }
11398 continue;
11399 }
11400
11401 let _ = self.memory.add_message(ChatMessage::assistant(&content)).await;
11403
11404 let results = self.execute_tools_parallel(&tool_calls).await;
11406
11407 for ((_id, result), tool_call) in results.into_iter().zip(tool_calls.iter()) {
11408 if include_tool_events {
11409 yield StreamChunk::tool_start(&tool_call.id, &tool_call.name);
11410 }
11411
11412 match result {
11413 Ok(output) => {
11414 if include_tool_events {
11415 yield StreamChunk::tool_result(
11416 &tool_call.id,
11417 &tool_call.name,
11418 &output,
11419 true,
11420 );
11421 }
11422 let _ = self.memory
11423 .add_message(Self::tool_result_message(
11424 tool_call,
11425 &output,
11426 native_tool_call,
11427 ))
11428 .await;
11429 }
11430 Err(e) => {
11431 if matches!(e, AgentError::HITLRejected(_)) {
11432 let _ = self.memory.add_message(ChatMessage::assistant(
11433 format!("The operation was rejected by the approver: {}", e),
11434 )).await;
11435 let response = AgentResponse {
11436 content: format!("Operation cancelled: {}", e),
11437 metadata: None,
11438 tool_calls: Some(all_tool_calls.clone()),
11439 };
11440 if let Err(finalize_error) = self.finish_turn_if_root(&response).await {
11441 yield StreamChunk::error(finalize_error.to_string());
11442 return;
11443 }
11444 yield StreamChunk::error(response.content);
11445 yield StreamChunk::Done {};
11446 return;
11447 }
11448 if include_tool_events {
11449 yield StreamChunk::tool_result(
11450 &tool_call.id,
11451 &tool_call.name,
11452 e.to_string(),
11453 false,
11454 );
11455 }
11456 let _ = self.memory
11457 .add_message(Self::tool_result_message(
11458 tool_call,
11459 &format!("Error: {}", e),
11460 native_tool_call,
11461 ))
11462 .await;
11463 }
11464 }
11465 all_tool_calls.push(tool_call.clone());
11466
11467 if include_tool_events {
11468 yield StreamChunk::tool_end(&tool_call.id);
11469 }
11470 }
11471 continue;
11472 }
11473
11474 let (extracted_thinking, answer) = self.extract_thinking(&content);
11476 if extracted_thinking.is_some() {
11477 thinking_content = extracted_thinking;
11478 }
11479
11480 let output_data = match self.process_output(&answer, &input_data.context).await {
11481 Ok(d) => d,
11482 Err(e) => {
11483 yield StreamChunk::error(e.to_string());
11484 return;
11485 }
11486 };
11487
11488 let final_content = if output_data.metadata.rejected {
11489 output_data
11490 .metadata
11491 .rejection_reason
11492 .unwrap_or_else(|| answer.to_string())
11493 } else {
11494 output_data.content
11495 };
11496
11497 let (final_content, _reflection_metadata) = match self
11499 .run_reflection(&*llm, processed_input, final_content)
11500 .await
11501 {
11502 Ok(r) => r,
11503 Err(e) => {
11504 yield StreamChunk::error(e.to_string());
11505 return;
11506 }
11507 };
11508
11509 let final_content = self.format_response_with_thinking(
11510 thinking_content.as_deref(),
11511 &final_content,
11512 );
11513
11514 if buffered_decision {
11516 yield StreamChunk::content(&final_content);
11517 }
11518
11519 let post_result = match self
11523 .post_loop_processing(processed_input, final_content)
11524 .await
11525 {
11526 Ok(r) => r,
11527 Err(e) => {
11528 yield StreamChunk::error(e.to_string());
11529 return;
11530 }
11531 };
11532
11533 let (final_content, transitioned) = match post_result {
11534 PostLoopResult::NoTransition(content) => (content, false),
11535 PostLoopResult::Transitioned(content) => (content, true),
11536 PostLoopResult::NeedsRedispatch => {
11537 const MAX_REDISPATCH_DEPTH: u32 = 3;
11538 let current_depth = *self.redispatch_depth.read();
11539 let content = if current_depth >= MAX_REDISPATCH_DEPTH {
11540 warn!(
11541 depth = current_depth,
11542 "Post-transition re-dispatch depth limit reached (stream)"
11543 );
11544 let c = String::new();
11545 let _ = self.memory.add_message(ChatMessage::assistant(&c)).await;
11546 c
11547 } else {
11548 *self.redispatch_depth.write() += 1;
11549 if let Some(context) = self.active_turn_context.write().as_mut() {
11550 context.enter_redispatch();
11551 }
11552 info!(
11553 depth = current_depth + 1,
11554 "Re-dispatching for new state after transition (stream)"
11555 );
11556 let result = self.run_loop_internal(processed_input).await;
11557 *self.redispatch_depth.write() -= 1;
11558 if let Some(context) = self.active_turn_context.write().as_mut() {
11559 context.exit_redispatch();
11560 }
11561 match result {
11562 Ok(resp) => resp.content,
11563 Err(e) => {
11564 yield StreamChunk::error(e.to_string());
11565 return;
11566 }
11567 }
11568 };
11569 (content, true)
11570 }
11571 };
11572
11573 if transitioned {
11574 if include_state_events
11575 && let Some(state) = self.current_state()
11576 {
11577 yield StreamChunk::state_transition(None, state);
11578 }
11579 yield StreamChunk::content(&final_content);
11581 }
11582
11583 let final_response = AgentResponse::new(&final_content);
11585 if let Err(e) = self.finish_turn_if_root(&final_response).await {
11586 yield StreamChunk::error(e.to_string());
11587 return;
11588 }
11589
11590 yield StreamChunk::Done {};
11591 return;
11592 }
11593 })
11594 }
11595
11596 fn run_loop_stream<'a>(
11599 &'a self,
11600 input: &'a str,
11601 ) -> Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> {
11602 Box::pin(async_stream::stream! {
11603 self.begin_root_turn();
11604 let _root_cleanup = RootTurnCleanup::new(self);
11605 self.hooks.on_message_received(input).await;
11606
11607 if !self.context_initialized.swap(true, Ordering::SeqCst) {
11609 if let Err(e) = self.context_manager.initialize().await {
11610 yield StreamChunk::error(e.to_string());
11611 return;
11612 }
11613 debug!("Context manager initialized (defaults, env, builtins)");
11614 }
11615
11616 if let Err(e) = self.check_turn_timeout().await {
11617 yield StreamChunk::error(e.to_string());
11618 return;
11619 }
11620 if let Err(e) = self.context_manager.refresh_per_turn().await {
11621 yield StreamChunk::error(e.to_string());
11622 return;
11623 }
11624
11625 self.clear_disambiguation_context();
11627
11628 if let Some(ref disambiguator) = self.disambiguation_manager {
11630 let disambiguation_context = match self.build_disambiguation_context().await {
11631 Ok(ctx) => ctx,
11632 Err(e) => {
11633 yield StreamChunk::error(e.to_string());
11634 return;
11635 }
11636 };
11637
11638 let state_override = self
11639 .state_machine
11640 .as_ref()
11641 .and_then(|sm| sm.current_definition())
11642 .and_then(|def| def.disambiguation.clone());
11643
11644 let state_generation = self
11645 .state_machine
11646 .as_ref()
11647 .map(|state_machine| state_machine.generation());
11648 let disambiguation_epoch = self.disambiguation_epoch.load(Ordering::SeqCst);
11649 let mut result = match self
11650 .observe_purpose(
11651 ObservationPurpose::DisambiguationDetection,
11652 disambiguator.process_input_with_override(
11653 input,
11654 &disambiguation_context,
11655 state_override.as_ref(),
11656 None,
11657 ),
11658 )
11659 .await
11660 {
11661 Ok(r) => r,
11662 Err(e) => {
11663 yield StreamChunk::error(e.to_string());
11664 return;
11665 }
11666 };
11667 let current_state_generation = self
11668 .state_machine
11669 .as_ref()
11670 .map(|state_machine| state_machine.generation());
11671 if current_state_generation != state_generation
11672 || self.disambiguation_epoch.load(Ordering::SeqCst) != disambiguation_epoch
11673 {
11674 disambiguator.clear_pending().await;
11675 *self.pending_skill_id.write() = None;
11676 result = DisambiguationResult::Abandoned { new_input: None };
11677 info!(
11678 confirmation_event = "invalidated",
11679 invalidation_reason = "state_generation_changed",
11680 "Streaming disambiguation result invalidated before redispatch"
11681 );
11682 }
11683 match result {
11684 DisambiguationResult::Clear => {
11685 debug!("Input is clear, proceeding normally (stream)");
11686 }
11687 DisambiguationResult::NeedsClarification {
11688 question,
11689 detection,
11690 } => {
11691 let admission = match self
11692 .admit_disambiguation_redispatch(
11693 disambiguation_epoch,
11694 state_generation,
11695 )
11696 .await
11697 {
11698 Ok(admission) => admission,
11699 Err(error) => {
11700 *self.pending_skill_id.write() = None;
11701 yield StreamChunk::error(error.to_string());
11702 return;
11703 }
11704 };
11705 let awaiting_confirmation = disambiguator.has_pending_confirmation().await;
11706 info!(
11707 ambiguity_type = ?detection.ambiguity_type,
11708 confidence = detection.confidence,
11709 "Input requires clarification (stream)"
11710 );
11711 if let Err(e) = self.commit_root_user_message(input).await {
11714 yield StreamChunk::error(e.to_string());
11715 return;
11716 }
11717 let _ = self
11718 .memory
11719 .add_message(ChatMessage::assistant(&question.question))
11720 .await;
11721 let status = if awaiting_confirmation {
11722 "awaiting_confirmation"
11723 } else {
11724 "awaiting_clarification"
11725 };
11726 let response = AgentResponse::new(&question.question).with_metadata(
11727 "disambiguation",
11728 serde_json::json!({ "status": status }),
11729 );
11730 drop(admission);
11731 if let Err(e) = self.finish_turn_if_root(&response).await {
11732 yield StreamChunk::error(e.to_string());
11733 return;
11734 }
11735 yield StreamChunk::content(&question.question);
11736 yield StreamChunk::Done {};
11737 return;
11738 }
11739 DisambiguationResult::Clarified {
11740 enriched_input,
11741 resolved,
11742 ..
11743 } => {
11744 let admission = match self
11745 .admit_disambiguation_redispatch(
11746 disambiguation_epoch,
11747 state_generation,
11748 )
11749 .await
11750 {
11751 Ok(admission) => admission,
11752 Err(error) => {
11753 *self.pending_skill_id.write() = None;
11754 yield StreamChunk::error(error.to_string());
11755 return;
11756 }
11757 };
11758 info!(
11759 resolved_count = resolved.len(),
11760 enriched = %enriched_input,
11761 "Input clarified (stream)"
11762 );
11763 for (key, value) in &resolved {
11764 let context_key = format!("disambiguation.{}", key);
11765 let _ = self.context_manager.set(&context_key, value.clone());
11766 }
11767 if let Some(intent) = resolved.get("intent") {
11768 let _ = self.context_manager.set("resolved_intent", intent.clone());
11769 }
11770 let _ = self
11771 .context_manager
11772 .set("disambiguation.resolved", serde_json::Value::Bool(true));
11773
11774 let skill_id = self.pending_skill_id.read().clone();
11778 if let Some(skill_id) = skill_id {
11779 info!(skill_id = %skill_id, "Re-checking skill disambiguation on clarified input (stream)");
11780 drop(admission);
11781 match self
11782 .recheck_skill_disambiguation(
11783 &skill_id,
11784 &enriched_input,
11785 disambiguation_epoch,
11786 state_generation,
11787 )
11788 .await
11789 {
11790 Ok(resp) => {
11791 yield StreamChunk::content(&resp.content);
11792 yield StreamChunk::Done {};
11793 return;
11794 }
11795 Err(e) => {
11796 yield StreamChunk::error(e.to_string());
11797 return;
11798 }
11799 }
11800 }
11801
11802 drop(admission);
11804 let mut inner = self.run_loop_internal_stream(&enriched_input);
11805 while let Some(chunk) = inner.next().await {
11806 yield chunk;
11807 }
11808 return;
11809 }
11810 DisambiguationResult::ProceedWithBestGuess { enriched_input } => {
11811 info!("Proceeding with best guess (stream)");
11812
11813 let skill_id = self.pending_skill_id.read().clone();
11815 if let Some(skill_id) = skill_id {
11816 info!(skill_id = %skill_id, "Re-checking skill disambiguation on best-guess input (stream)");
11817 match self
11818 .recheck_skill_disambiguation(
11819 &skill_id,
11820 &enriched_input,
11821 disambiguation_epoch,
11822 state_generation,
11823 )
11824 .await
11825 {
11826 Ok(resp) => {
11827 yield StreamChunk::content(&resp.content);
11828 yield StreamChunk::Done {};
11829 return;
11830 }
11831 Err(e) => {
11832 yield StreamChunk::error(e.to_string());
11833 return;
11834 }
11835 }
11836 }
11837
11838 let mut inner = self.run_loop_internal_stream(&enriched_input);
11839 while let Some(chunk) = inner.next().await {
11840 yield chunk;
11841 }
11842 return;
11843 }
11844 DisambiguationResult::GiveUp { reason } => {
11845 *self.pending_skill_id.write() = None;
11846 warn!(reason = %reason, "Disambiguation gave up (stream)");
11847 let apology = self
11848 .generate_localized_apology(
11849 "Generate a brief, polite apology saying you couldn't understand the request. Be concise.",
11850 &reason,
11851 )
11852 .await
11853 .unwrap_or_else(|_| {
11854 format!("I'm sorry, I couldn't understand your request: {}", reason)
11855 });
11856 let response = AgentResponse::new(&apology);
11857 if let Err(e) = self.finish_turn_if_root(&response).await {
11858 yield StreamChunk::error(e.to_string());
11859 return;
11860 }
11861 yield StreamChunk::content(&apology);
11862 yield StreamChunk::Done {};
11863 return;
11864 }
11865 DisambiguationResult::Escalate { reason } => {
11866 *self.pending_skill_id.write() = None;
11867 info!(reason = %reason, "Escalating to human (stream)");
11868 if let Some(ref hitl) = self.hitl_engine {
11869 let trigger =
11870 ApprovalTrigger::condition("disambiguation_escalation", reason.clone());
11871 let mut context_map = HashMap::new();
11872 context_map.insert("original_input".to_string(), serde_json::json!(input));
11873 context_map.insert("reason".to_string(), serde_json::json!(&reason));
11874 let check_result = HITLCheckResult::required(
11875 trigger,
11876 context_map,
11877 format!("User request needs human assistance: {}", reason),
11878 Some(hitl.config().default_timeout_seconds),
11879 );
11880 match self.request_hitl_approval(check_result).await {
11881 Ok(ApprovalResult::Approved | ApprovalResult::Modified { .. }) => {
11882 let mut inner = self.run_loop_internal_stream(input);
11883 while let Some(chunk) = inner.next().await {
11884 yield chunk;
11885 }
11886 return;
11887 }
11888 Ok(_) => {}
11889 Err(e) => {
11890 yield StreamChunk::error(e.to_string());
11891 return;
11892 }
11893 }
11894 }
11895 let apology = self
11896 .generate_localized_apology(
11897 "Explain briefly that you're transferring the user to a human agent for help.",
11898 &reason,
11899 )
11900 .await
11901 .unwrap_or_else(|_| {
11902 format!("I need human assistance to help with 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 yield StreamChunk::Done {};
11911 return;
11912 }
11913 DisambiguationResult::Abandoned { new_input } => {
11914 *self.pending_skill_id.write() = None;
11915
11916 info!(
11917 has_new_input = new_input.is_some(),
11918 "Clarification abandoned by user (stream)"
11919 );
11920
11921 if let Err(e) = self.commit_root_user_message(input).await {
11922 yield StreamChunk::error(e.to_string());
11923 return;
11924 }
11925
11926 match new_input {
11927 Some(fresh_input) => {
11928 let mut inner = self.run_loop_internal_stream(&fresh_input);
11930 while let Some(chunk) = inner.next().await {
11931 yield chunk;
11932 }
11933 return;
11934 }
11935 None => {
11936 let ack = self
11938 .generate_localized_apology(
11939 "The user changed their mind about their previous request. \
11940 Generate a brief, friendly acknowledgment (e.g. 'OK, no problem. What else can I help with?'). \
11941 Do NOT apologize excessively. Be concise.",
11942 "User abandoned clarification",
11943 )
11944 .await
11945 .unwrap_or_else(|_| {
11946 "OK, no problem. What else can I help with?".to_string()
11947 });
11948
11949 let _ = self
11950 .memory
11951 .add_message(ChatMessage::assistant(&ack))
11952 .await;
11953
11954 let response = AgentResponse::new(&ack);
11955 if let Err(e) = self.finish_turn_if_root(&response).await {
11956 yield StreamChunk::error(e.to_string());
11957 return;
11958 }
11959 yield StreamChunk::content(&ack);
11960 yield StreamChunk::Done {};
11961 return;
11962 }
11963 }
11964 }
11965 }
11966 }
11967
11968 let mut inner = self.run_loop_internal_stream(input);
11970 while let Some(chunk) = inner.next().await {
11971 yield chunk;
11972 }
11973 })
11974 }
11975
11976 pub fn info(&self) -> AgentInfo {
11977 self.info.clone()
11978 }
11979
11980 pub fn skills(&self) -> &[SkillDefinition] {
11981 &self.skills
11982 }
11983
11984 async fn reset_runtime_state(&self) -> Result<()> {
11986 let _admission = self.disambiguation_admission.write().await;
11987 if self.state_transition_reserved.load(Ordering::SeqCst) {
11988 return Err(AgentError::Other(
11989 "Cannot reset while a state transition is in progress".to_string(),
11990 ));
11991 }
11992 self.disambiguation_epoch.fetch_add(1, Ordering::SeqCst);
11993 *self.pending_skill_id.write() = None;
11994 if let Some(disambiguator) = self.disambiguation_manager.as_ref() {
11995 disambiguator.clear_pending().await;
11996 }
11997 self.memory.clear().await?;
11998 *self.iteration_count.write() = 0;
11999 self.tool_call_history.write().clear();
12000 if let Some(ref sm) = self.state_machine {
12001 sm.reset();
12002 }
12003 Ok(())
12004 }
12005
12006 pub async fn reset(&self) -> Result<()> {
12008 self.reset_runtime_state().await
12009 }
12010
12011 pub fn max_context_tokens(&self) -> u32 {
12012 self.max_context_tokens
12013 }
12014
12015 pub fn llm_registry(&self) -> &Arc<LLMRegistry> {
12016 &self.llm_registry
12017 }
12018
12019 pub fn state_machine(&self) -> Option<&Arc<StateMachine>> {
12020 self.state_machine.as_ref()
12021 }
12022
12023 pub fn context_manager(&self) -> &Arc<ContextManager> {
12024 &self.context_manager
12025 }
12026
12027 pub fn tool_call_history(&self) -> Vec<ToolCallRecord> {
12028 self.tool_call_history.read().clone()
12029 }
12030
12031 pub fn memory_token_budget(&self) -> Option<&MemoryTokenBudget> {
12032 self.memory_token_budget.as_ref()
12033 }
12034
12035 pub fn parallel_tools_config(&self) -> &ParallelToolsConfig {
12036 &self.parallel_tools
12037 }
12038
12039 pub fn streaming_config(&self) -> &StreamingConfig {
12040 &self.streaming
12041 }
12042
12043 pub fn hooks(&self) -> &Arc<dyn AgentHooks> {
12044 &self.hooks
12045 }
12046
12047 pub fn hitl_engine(&self) -> Option<&HITLEngine> {
12048 self.hitl_engine.as_ref()
12049 }
12050
12051 pub fn approval_handler(&self) -> &Arc<dyn ApprovalHandler> {
12052 &self.approval_handler
12053 }
12054
12055 fn build_hitl_language_context(&self) -> HashMap<String, Value> {
12057 let mut ctx = HashMap::new();
12058 for key in &["user.language", "input.detected.language", "language"] {
12059 if let Some(val) = self.context_manager.get(key) {
12060 ctx.insert(key.to_string(), val);
12061 }
12062 }
12063 ctx
12064 }
12065
12066 async fn request_hitl_approval(&self, check_result: HITLCheckResult) -> Result<ApprovalResult> {
12068 let Some(request) = check_result.into_request() else {
12069 return Ok(ApprovalResult::Approved);
12070 };
12071
12072 self.hooks.on_approval_requested(&request).await;
12073
12074 let timeout = request.timeout;
12075
12076 let raw_result = if let Some(duration) = timeout {
12077 match tokio::time::timeout(
12078 duration,
12079 self.approval_handler.request_approval(request.clone()),
12080 )
12081 .await
12082 {
12083 Ok(result) => result,
12084 Err(_) => ApprovalResult::timeout(),
12085 }
12086 } else {
12087 self.approval_handler
12088 .request_approval(request.clone())
12089 .await
12090 };
12091
12092 self.hooks
12093 .on_approval_result(&request.id, &raw_result)
12094 .await;
12095
12096 let (outcome, effective_result): (ApprovalResolvedOutcome, Result<ApprovalResult>) =
12097 match &raw_result {
12098 ApprovalResult::Approved => (
12099 ApprovalResolvedOutcome::Approved,
12100 Ok(ApprovalResult::Approved),
12101 ),
12102 ApprovalResult::Rejected { reason } => (
12103 ApprovalResolvedOutcome::Rejected {
12104 reason: reason.clone(),
12105 },
12106 Ok(ApprovalResult::Rejected {
12107 reason: reason.clone(),
12108 }),
12109 ),
12110 ApprovalResult::Modified { changes } => (
12111 ApprovalResolvedOutcome::Modified {
12112 changes: changes.clone(),
12113 },
12114 Ok(ApprovalResult::Modified {
12115 changes: changes.clone(),
12116 }),
12117 ),
12118 ApprovalResult::Timeout => {
12119 if let Some(ref engine) = self.hitl_engine {
12120 match engine.config().on_timeout {
12121 TimeoutAction::Approve => (
12122 ApprovalResolvedOutcome::Approved,
12123 Ok(ApprovalResult::Approved),
12124 ),
12125 TimeoutAction::Reject => {
12126 let reason = Some("Timeout".to_string());
12127 (
12128 ApprovalResolvedOutcome::Rejected {
12129 reason: reason.clone(),
12130 },
12131 Ok(ApprovalResult::Rejected { reason }),
12132 )
12133 }
12134 TimeoutAction::Error => {
12135 let message = "HITL approval timeout".to_string();
12136 (
12137 ApprovalResolvedOutcome::Error {
12138 message: message.clone(),
12139 },
12140 Err(AgentError::Other(message)),
12141 )
12142 }
12143 }
12144 } else {
12145 let reason = Some("Timeout (no engine)".to_string());
12146 (
12147 ApprovalResolvedOutcome::Rejected {
12148 reason: reason.clone(),
12149 },
12150 Ok(ApprovalResult::Rejected { reason }),
12151 )
12152 }
12153 }
12154 };
12155
12156 self.hooks
12157 .on_approval_resolved(&request, &raw_result, &outcome)
12158 .await;
12159
12160 effective_result
12161 }
12162
12163 pub async fn check_state_hitl(&self, from: Option<&str>, to: &str) -> Result<bool> {
12164 if let Some(ref hitl_engine) = self.hitl_engine {
12165 let hitl_lang_ctx = self.build_hitl_language_context();
12166 let check_result = self
12167 .observe_purpose(
12168 ObservationPurpose::HitlLocalization,
12169 hitl_engine.check_state_transition_with_localization(
12170 from,
12171 to,
12172 &hitl_lang_ctx,
12173 self.approval_handler.as_ref(),
12174 Some(&self.llm_registry),
12175 ),
12176 )
12177 .await?;
12178 if check_result.is_required() {
12179 let result = self.request_hitl_approval(check_result).await?;
12180 return Ok(matches!(
12181 result,
12182 ApprovalResult::Approved | ApprovalResult::Modified { .. }
12183 ));
12184 }
12185 }
12186 Ok(true)
12187 }
12188
12189 async fn execute_tools_parallel(
12191 &self,
12192 tool_calls: &[ToolCall],
12193 ) -> Vec<(String, Result<String>)> {
12194 let can_run_parallel = tool_calls.iter().all(|tc| {
12195 self.tools
12196 .resolve(&tc.name)
12197 .map(|resolved| resolved.tool.classify_call(&tc.arguments).concurrency_safe)
12198 .unwrap_or(false)
12199 });
12200
12201 if !self.parallel_tools.enabled || tool_calls.len() <= 1 || !can_run_parallel {
12202 let mut results = Vec::new();
12203 for tc in tool_calls {
12204 let result = self
12205 .observe_purpose(
12206 current_observation_context()
12207 .map(|context| context.purpose)
12208 .unwrap_or_default(),
12209 self.execute_tool_smart(tc),
12210 )
12211 .await;
12212 results.push((tc.id.clone(), result));
12213 }
12214 return results;
12215 }
12216
12217 let chunks: Vec<_> = tool_calls
12218 .chunks(self.parallel_tools.max_parallel)
12219 .collect();
12220
12221 let mut all_results = Vec::new();
12222
12223 for chunk in chunks {
12224 let futures: Vec<_> = chunk
12225 .iter()
12226 .map(|tc| {
12227 let tc = tc.clone();
12228 async move {
12229 let result = self.execute_tool_smart(&tc).await;
12230 (tc.id.clone(), result)
12231 }
12232 })
12233 .collect();
12234
12235 let results = futures::future::join_all(futures).await;
12236 all_results.extend(results);
12237 }
12238
12239 all_results
12240 }
12241
12242 pub async fn chat_stream<'a>(
12244 &'a self,
12245 input: &'a str,
12246 ) -> Result<Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>>> {
12247 self.init_storage().await?;
12251 info!(input_len = input.len(), "Starting streaming chat");
12252 let inner = self.run_loop_stream(input);
12253 if let Some(context) = self.build_observation_context(None) {
12254 let stream: Pin<Box<dyn Stream<Item = StreamChunk> + Send + 'a>> =
12255 Box::pin(async_stream::stream! {
12256 let mut inner = inner;
12257 loop {
12258 let next = with_observation_context(context.clone(), inner.next()).await;
12259 match next {
12260 Some(chunk) => yield chunk,
12261 None => break,
12262 }
12263 }
12264 self.export_observability_if_configured().await;
12265 });
12266 Ok(stream)
12267 } else {
12268 Ok(inner)
12269 }
12270 }
12271}
12272
12273#[async_trait]
12274impl ToolInvoker for RuntimeAgent {
12275 async fn invoke_tool(&self, request: ToolExecutionRequest) -> Result<ToolExecutionRecord> {
12276 self.execute_tool_record(request).await
12277 }
12278}
12279
12280#[async_trait]
12281impl Agent for RuntimeAgent {
12282 async fn chat(&self, input: &str) -> Result<AgentResponse> {
12283 let result = if let Some(context) = self.build_observation_context(None) {
12284 with_observation_context(context, self.run_loop(input)).await
12285 } else {
12286 self.run_loop(input).await
12287 };
12288 self.export_observability_if_configured().await;
12289 result
12290 }
12291
12292 fn info(&self) -> AgentInfo {
12293 self.info.clone()
12294 }
12295
12296 async fn reset(&self) -> Result<()> {
12298 self.reset_runtime_state().await
12299 }
12300}
12301
12302fn background_maintenance_tags(
12312 label: &str,
12313 stage: &str,
12314 reason: Option<&str>,
12315 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12316) -> HashMap<String, String> {
12317 let mut tags = HashMap::new();
12318 tags.insert("runtime.background".to_string(), "true".to_string());
12319 tags.insert("runtime.maintenance".to_string(), label.to_string());
12320 tags.insert("runtime.maintenance_stage".to_string(), stage.to_string());
12321 if let Some(policy) = policy {
12322 tags.insert(
12323 "runtime.await_before_next_turn".to_string(),
12324 await_before_next_turn_label(policy.await_before_next_turn).to_string(),
12325 );
12326 tags.insert(
12327 "runtime.maintenance_mode".to_string(),
12328 maintenance_mode_label(policy.mode).to_string(),
12329 );
12330 }
12331 if let Some(reason) = reason {
12332 tags.insert("runtime.reason".to_string(), reason.to_string());
12333 }
12334 tags
12335}
12336
12337fn await_before_next_turn_label(policy: AwaitBeforeNextTurn) -> &'static str {
12338 match policy {
12339 AwaitBeforeNextTurn::Never => "never",
12340 AwaitBeforeNextTurn::SameActor => "same_actor",
12341 AwaitBeforeNextTurn::Always => "always",
12342 }
12343}
12344
12345fn maintenance_mode_label(mode: MaintenanceMode) -> &'static str {
12346 match mode {
12347 MaintenanceMode::InlineSerial => "inline_serial",
12348 MaintenanceMode::InlineParallel => "inline_parallel",
12349 MaintenanceMode::Background => "background",
12350 }
12351}
12352
12353fn record_background_maintenance_event(
12355 manager: Option<&Arc<ObservabilityManager>>,
12356 label: &str,
12357 status: EventStatus,
12358 duration_ms: u64,
12359 stage: &str,
12360 reason: Option<String>,
12361 policy: Option<&crate::optimization::config::MaintenanceTaskPolicy>,
12362) {
12363 if let Some(manager) = manager {
12364 manager.record_lifecycle_event(
12365 EventType::MemoryOperation {
12366 operation: format!("{}_background_{}", label, stage),
12367 },
12368 ObservationPurpose::Other(format!("{}_maintenance", label)),
12369 status,
12370 duration_ms,
12371 background_maintenance_tags(label, stage, reason.as_deref(), policy),
12372 None,
12373 );
12374 }
12375}
12376
12377fn effective_maintenance_mode(mode: MaintenanceMode, force_parallel: bool) -> MaintenanceMode {
12378 if force_parallel && matches!(mode, MaintenanceMode::InlineSerial) {
12379 MaintenanceMode::InlineParallel
12380 } else {
12381 mode
12382 }
12383}
12384
12385fn observation_purpose_for_process(hint: ProcessPurposeHint) -> ObservationPurpose {
12386 match hint {
12387 ProcessPurposeHint::Detect => ObservationPurpose::ProcessDetect,
12388 ProcessPurposeHint::Extract => ObservationPurpose::ProcessExtract,
12389 ProcessPurposeHint::Validate => ObservationPurpose::ProcessValidate,
12390 ProcessPurposeHint::Transform | ProcessPurposeHint::Other => {
12391 ObservationPurpose::ProcessTransform
12392 }
12393 }
12394}
12395
12396fn new_tool_resource_locks() -> ToolResourceLocks {
12397 Arc::new(RwLock::new(HashMap::new()))
12398}
12399
12400fn tool_resource_lock_keys(
12405 _canonical_id: &str,
12406 args: &Value,
12407 bindings: &ai_agents_core::ToolPolicyBindings,
12408 classification: &ai_agents_core::ToolCallClassification,
12409) -> Vec<String> {
12410 if classification.concurrency_safe {
12411 return Vec::new();
12412 }
12413
12414 let mut keys = Vec::new();
12415 let mut has_path_resource = false;
12416 for binding in &bindings.path_fields {
12417 let value = value_at_argument_path(args, &binding.field)
12418 .cloned()
12419 .or_else(|| {
12420 binding
12421 .default_path
12422 .as_ref()
12423 .map(|path| Value::String(path.clone()))
12424 });
12425 if let Some(value) = value {
12426 collect_resource_strings(&value, |_| {
12427 has_path_resource = true;
12428 });
12429 }
12430 }
12431 for binding in &bindings.domain_fields {
12432 if let Some(value) = value_at_argument_path(args, &binding.field) {
12433 collect_resource_strings(value, |domain| {
12434 let normalized = if binding.is_url {
12435 normalized_url_resource_key(domain)
12436 } else {
12437 domain.trim().trim_end_matches('.').to_ascii_lowercase()
12438 };
12439 keys.push(format!("domain:{}", normalized));
12440 });
12441 }
12442 }
12443 for binding in &bindings.command_fields {
12444 if !matches!(binding.kind, ai_agents_core::CommandBindingKind::Cwd) {
12445 continue;
12446 }
12447 if let Some(value) = value_at_argument_path(args, &binding.field) {
12448 collect_resource_strings(value, |_| {
12449 has_path_resource = true;
12450 });
12451 }
12452 }
12453 if has_path_resource {
12454 keys.push("path-mutation:global".to_string());
12455 }
12456 if keys.is_empty() {
12457 keys.push("side-effect:unbound".to_string());
12458 }
12459 keys.sort();
12460 keys.dedup();
12461 keys
12462}
12463
12464fn value_at_argument_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
12465 let mut current = value;
12466 for segment in field.split('.') {
12467 if segment.is_empty() {
12468 return None;
12469 }
12470 current = current.get(segment)?;
12471 }
12472 Some(current)
12473}
12474
12475fn collect_resource_strings(value: &Value, mut collect: impl FnMut(&str)) {
12476 match value {
12477 Value::String(value) => collect(value),
12478 Value::Array(values) => {
12479 for value in values {
12480 if let Some(value) = value.as_str() {
12481 collect(value);
12482 }
12483 }
12484 }
12485 _ => {}
12486 }
12487}
12488
12489fn normalized_url_resource_key(value: &str) -> String {
12490 let value = value.trim();
12491 let Some((scheme, remainder)) = value.split_once("://") else {
12492 return value.to_ascii_lowercase();
12493 };
12494 let authority_end = remainder.find(['/', '?', '#']).unwrap_or(remainder.len());
12495 let (authority, suffix) = remainder.split_at(authority_end);
12496 format!(
12497 "{}://{}{}",
12498 scheme.to_ascii_lowercase(),
12499 authority.to_ascii_lowercase(),
12500 suffix
12501 )
12502}
12503
12504fn render_concurrent_template(
12505 template: &str,
12506 user_input: &str,
12507 context_values: &std::collections::HashMap<String, serde_json::Value>,
12508) -> Result<String> {
12509 let mut env = minijinja::Environment::new();
12510 env.add_template("concurrent", template)
12511 .map_err(|e| AgentError::Other(format!("Concurrent template parse error: {}", e)))?;
12512
12513 let mut ctx = std::collections::BTreeMap::new();
12514 ctx.insert("user_input".to_string(), minijinja::Value::from(user_input));
12515
12516 let context_obj = minijinja::Value::from_serialize(context_values);
12518 ctx.insert("context".to_string(), context_obj);
12519
12520 let tmpl = env
12521 .get_template("concurrent")
12522 .map_err(|e| AgentError::Other(format!("Concurrent template error: {}", e)))?;
12523
12524 tmpl.render(minijinja::Value::from_serialize(&ctx))
12525 .map_err(|e| AgentError::Other(format!("Concurrent template render error: {}", e)))
12526}
12527
12528#[cfg(test)]
12529mod tests {
12530 use super::*;
12531 use crate::AgentBuilder;
12532 use ai_agents_core::{LLMChunk, LLMConfig, LLMError, LLMFeature, Tool};
12533 use ai_agents_llm::mock::MockLLMProvider;
12534 use ai_agents_skills::{SkillDefinition, SkillStep};
12535 use ai_agents_tools::{
12536 CalculatorTool, CopyPathTool, DeletePathTool, FileWriteTool, MovePathTool,
12537 WebFetchResolver, WebFetchTool, WebFetchTransport, WebFetchTransportRequest,
12538 WebFetchTransportResponse,
12539 };
12540
12541 fn mock_with_response(response: &str) -> MockLLMProvider {
12542 let mut mock = MockLLMProvider::new("test");
12543 mock.set_response(response);
12544 mock
12545 }
12546
12547 fn mock_with_responses(responses: Vec<&str>) -> MockLLMProvider {
12548 let mut mock = MockLLMProvider::new("test");
12549 mock.set_responses(responses.into_iter().map(String::from).collect(), true);
12550 mock
12551 }
12552
12553 fn disambiguation_state_machine(
12555 state_enabled: Option<bool>,
12556 require_confirmation: bool,
12557 ) -> Arc<StateMachine> {
12558 let definition = ai_agents_state::StateDefinition {
12559 prompt: Some("Handle the resolved request.".to_string()),
12560 disambiguation: Some(ai_agents_disambiguation::StateDisambiguationOverride {
12561 enabled: state_enabled,
12562 require_confirmation,
12563 ..Default::default()
12564 }),
12565 ..Default::default()
12566 };
12567 let review = ai_agents_state::StateDefinition {
12568 prompt: Some("Review a fresh request.".to_string()),
12569 ..Default::default()
12570 };
12571 Arc::new(
12572 StateMachine::new(ai_agents_state::StateConfig {
12573 initial: "active".to_string(),
12574 states: std::collections::HashMap::from([
12575 ("active".to_string(), definition),
12576 ("review".to_string(), review),
12577 ]),
12578 global_transitions: Vec::new(),
12579 fallback: None,
12580 max_no_transition: None,
12581 regenerate_on_transition: true,
12582 })
12583 .unwrap(),
12584 )
12585 }
12586
12587 fn state_disambiguation_agent(
12589 responses: Vec<&str>,
12590 manager_enabled: bool,
12591 state_enabled: Option<bool>,
12592 require_confirmation: bool,
12593 ) -> (RuntimeAgent, MockLLMProvider) {
12594 state_disambiguation_agent_with_skills(
12595 responses,
12596 manager_enabled,
12597 state_enabled,
12598 require_confirmation,
12599 Vec::new(),
12600 )
12601 }
12602
12603 fn state_disambiguation_agent_with_skills(
12605 responses: Vec<&str>,
12606 manager_enabled: bool,
12607 state_enabled: Option<bool>,
12608 require_confirmation: bool,
12609 skills: Vec<SkillDefinition>,
12610 ) -> (RuntimeAgent, MockLLMProvider) {
12611 let mut mock = MockLLMProvider::new("state-confirmation");
12612 mock.set_responses(responses.into_iter().map(String::from).collect(), false);
12613 let observed = mock.clone();
12614 let agent = AgentBuilder::new()
12615 .system_prompt("Handle requests.")
12616 .llm(Arc::new(mock.clone()))
12617 .llm_alias("router", Arc::new(mock))
12618 .state_machine(disambiguation_state_machine(
12619 state_enabled,
12620 require_confirmation,
12621 ))
12622 .skills(skills)
12623 .build()
12624 .unwrap()
12625 .with_disambiguation(DisambiguationConfig {
12626 enabled: manager_enabled,
12627 ..Default::default()
12628 });
12629 (agent, observed)
12630 }
12631
12632 fn confirmation_skill() -> SkillDefinition {
12634 SkillDefinition {
12635 id: "send_report".to_string(),
12636 description: "Send a report after clarification".to_string(),
12637 trigger: "When the user asks to send a report".to_string(),
12638 steps: vec![SkillStep::Prompt {
12639 prompt: "Execute confirmed report skill for: {{ input }}".to_string(),
12640 llm: None,
12641 }],
12642 reasoning: None,
12643 reflection: None,
12644 disambiguation: Some(ai_agents_disambiguation::SkillDisambiguationOverride {
12645 enabled: Some(true),
12646 ..Default::default()
12647 }),
12648 }
12649 }
12650
12651 fn confirmation_skill_call_count(observed: &MockLLMProvider) -> usize {
12653 observed
12654 .call_history()
12655 .iter()
12656 .filter(|call| {
12657 call.messages
12658 .iter()
12659 .any(|message| message.content.contains("Execute confirmed report skill"))
12660 })
12661 .count()
12662 }
12663
12664 struct BlockingRuntimeConfirmationObserver {
12665 entered: tokio::sync::Barrier,
12666 release: tokio::sync::Notify,
12667 }
12668
12669 impl BlockingRuntimeConfirmationObserver {
12670 fn new() -> Self {
12671 Self {
12672 entered: tokio::sync::Barrier::new(2),
12673 release: tokio::sync::Notify::new(),
12674 }
12675 }
12676 }
12677
12678 struct ResetOnTransitionHooks {
12679 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
12680 invoked: AtomicBool,
12681 }
12682
12683 #[async_trait]
12684 impl AgentHooks for ResetOnTransitionHooks {
12685 async fn on_state_transition(&self, _from: Option<&str>, _to: &str, _reason: &str) {
12686 if self.invoked.swap(true, Ordering::SeqCst) {
12687 return;
12688 }
12689 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
12690 if let Some(agent) = agent {
12691 agent.reset().await.unwrap();
12692 }
12693 }
12694 }
12695
12696 impl ClarificationObserver for BlockingRuntimeConfirmationObserver {
12697 fn observe_question<'a>(
12698 &'a self,
12699 future: ClarificationQuestionFuture<'a>,
12700 ) -> ClarificationQuestionFuture<'a> {
12701 future
12702 }
12703
12704 fn observe_parse<'a>(
12705 &'a self,
12706 future: ClarificationParseFuture<'a>,
12707 ) -> ClarificationParseFuture<'a> {
12708 future
12709 }
12710
12711 fn observe_confirmation_parse<'a>(
12712 &'a self,
12713 future: ConfirmationParseFuture<'a>,
12714 ) -> ConfirmationParseFuture<'a> {
12715 Box::pin(async move {
12716 self.entered.wait().await;
12717 self.release.notified().await;
12718 future.await
12719 })
12720 }
12721 }
12722
12723 #[tokio::test]
12724 async fn state_confirmation_blocks_redispatch_until_explicit_agreement() {
12725 let (agent, observed) = state_disambiguation_agent(
12726 vec![
12727 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12728 r#"{"question":"What should I send?","options":null}"#,
12729 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12730 r#"{"question":"Should I send the report to Ada?"}"#,
12731 r#"{"status":"confirmed"}"#,
12732 "Request executed.",
12733 ],
12734 true,
12735 None,
12736 true,
12737 );
12738
12739 let clarification = agent.chat("Send it").await.unwrap();
12740 assert_eq!(clarification.content, "What should I send?");
12741 assert_eq!(observed.call_count(), 2);
12742
12743 let confirmation = agent.chat("The report to Ada").await.unwrap();
12744 assert_eq!(confirmation.content, "Should I send the report to Ada?");
12745 assert_eq!(
12746 confirmation
12747 .metadata
12748 .as_ref()
12749 .and_then(|metadata| metadata.get("disambiguation"))
12750 .and_then(|metadata| metadata.get("status"))
12751 .and_then(Value::as_str),
12752 Some("awaiting_confirmation")
12753 );
12754 assert_eq!(observed.call_count(), 4);
12755
12756 let completed = agent.chat("Yes").await.unwrap();
12757 assert_eq!(completed.content, "Request executed.");
12758 assert_eq!(observed.call_count(), 6);
12759 }
12760
12761 #[tokio::test]
12762 async fn streaming_state_confirmation_ends_the_turn_before_redispatch() {
12763 let (agent, observed) = state_disambiguation_agent(
12764 vec![
12765 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12766 r#"{"question":"What should I send?","options":null}"#,
12767 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12768 r#"{"question":"Should I send the report to Ada?"}"#,
12769 r#"{"status":"confirmed"}"#,
12770 "Request executed.",
12771 ],
12772 true,
12773 None,
12774 true,
12775 );
12776
12777 let mut clarification_stream = agent.chat_stream("Send it").await.unwrap();
12778 let mut clarification = String::new();
12779 while let Some(chunk) = clarification_stream.next().await {
12780 match chunk {
12781 StreamChunk::Content { text } => clarification.push_str(&text),
12782 StreamChunk::Done {} => break,
12783 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12784 _ => {}
12785 }
12786 }
12787 assert_eq!(clarification, "What should I send?");
12788 assert_eq!(observed.call_count(), 2);
12789
12790 let mut confirmation_stream = agent.chat_stream("The report to Ada").await.unwrap();
12791 let mut confirmation = String::new();
12792 while let Some(chunk) = confirmation_stream.next().await {
12793 match chunk {
12794 StreamChunk::Content { text } => confirmation.push_str(&text),
12795 StreamChunk::Done {} => break,
12796 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12797 _ => {}
12798 }
12799 }
12800 assert_eq!(confirmation, "Should I send the report to Ada?");
12801 assert_eq!(observed.call_count(), 4);
12802
12803 let mut completed_stream = agent.chat_stream("Yes").await.unwrap();
12804 let mut completed = String::new();
12805 while let Some(chunk) = completed_stream.next().await {
12806 match chunk {
12807 StreamChunk::Content { text } => completed.push_str(&text),
12808 StreamChunk::Done {} => break,
12809 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12810 _ => {}
12811 }
12812 }
12813 assert_eq!(completed, "Request executed.");
12814 assert_eq!(observed.call_count(), 6);
12815 }
12816
12817 #[tokio::test]
12819 async fn confirmed_skill_route_executes_exactly_once() {
12820 let (agent, observed) = state_disambiguation_agent_with_skills(
12821 vec![
12822 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12823 "send_report",
12824 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12825 r#"{"question":"What should I send?","options":null}"#,
12826 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12827 r#"{"question":"Should I send the report to Ada?"}"#,
12828 r#"{"status":"confirmed"}"#,
12829 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"resolved","what_is_unclear":[],"detected_language":"en"}"#,
12830 "Report skill executed.",
12831 ],
12832 true,
12833 None,
12834 true,
12835 vec![confirmation_skill()],
12836 );
12837
12838 let clarification = agent.chat("Send it").await.unwrap();
12839 assert_eq!(clarification.content, "What should I send?");
12840 assert_eq!(confirmation_skill_call_count(&observed), 0);
12841
12842 let confirmation = agent.chat("The report to Ada").await.unwrap();
12843 assert_eq!(confirmation.content, "Should I send the report to Ada?");
12844 assert_eq!(
12845 confirmation
12846 .metadata
12847 .as_ref()
12848 .and_then(|metadata| metadata.get("disambiguation"))
12849 .and_then(|metadata| metadata.get("status"))
12850 .and_then(Value::as_str),
12851 Some("awaiting_confirmation")
12852 );
12853 assert_eq!(confirmation_skill_call_count(&observed), 0);
12854
12855 let completed = agent.chat("Yes").await.unwrap();
12856 assert_eq!(completed.content, "Report skill executed.");
12857 assert_eq!(confirmation_skill_call_count(&observed), 1);
12858 assert!(agent.pending_skill_id.read().is_none());
12859 let messages = agent.memory.get_messages(None).await.unwrap();
12860 assert!(!messages.iter().any(|message| message.content == "Yes"));
12861 }
12862
12863 #[tokio::test]
12865 async fn confirmed_skill_recheck_preserves_new_clarification_metadata() {
12866 let (agent, observed) = state_disambiguation_agent_with_skills(
12867 vec![
12868 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12869 "send_report",
12870 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12871 r#"{"question":"What should I send?","options":null}"#,
12872 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12873 r#"{"question":"Should I send the report to Ada?"}"#,
12874 r#"{"status":"confirmed"}"#,
12875 r#"{"is_ambiguous":true,"confidence":0.3,"ambiguity_type":"missing_parameters","reasoning":"timing missing","what_is_unclear":["timing"],"detected_language":"en"}"#,
12876 r#"{"question":"When should I send it?","options":null}"#,
12877 ],
12878 true,
12879 None,
12880 true,
12881 vec![confirmation_skill()],
12882 );
12883
12884 agent.chat("Send it").await.unwrap();
12885 agent.chat("The report to Ada").await.unwrap();
12886 let follow_up = agent.chat("Yes").await.unwrap();
12887
12888 assert_eq!(follow_up.content, "When should I send it?");
12889 let metadata = follow_up
12890 .metadata
12891 .as_ref()
12892 .and_then(|metadata| metadata.get("disambiguation"))
12893 .unwrap();
12894 assert_eq!(
12895 metadata.get("status").and_then(Value::as_str),
12896 Some("awaiting_clarification")
12897 );
12898 assert_eq!(
12899 metadata.get("skill_id").and_then(Value::as_str),
12900 Some("send_report")
12901 );
12902 assert!(metadata.get("detection").is_some());
12903 assert_eq!(confirmation_skill_call_count(&observed), 0);
12904 }
12905
12906 #[tokio::test]
12908 async fn rejected_skill_confirmation_never_executes() {
12909 let (agent, observed) = state_disambiguation_agent_with_skills(
12910 vec![
12911 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12912 "send_report",
12913 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12914 r#"{"question":"What should I send?","options":null}"#,
12915 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12916 r#"{"question":"Should I send the report to Ada?"}"#,
12917 r#"{"status":"rejected"}"#,
12918 "Confirmation rejected.",
12919 ],
12920 true,
12921 None,
12922 true,
12923 vec![confirmation_skill()],
12924 );
12925
12926 agent.chat("Send it").await.unwrap();
12927 agent.chat("The report to Ada").await.unwrap();
12928 let rejected = agent.chat("No").await.unwrap();
12929
12930 assert_eq!(rejected.content, "Confirmation rejected.");
12931 assert_eq!(confirmation_skill_call_count(&observed), 0);
12932 assert!(agent.pending_skill_id.read().is_none());
12933 }
12934
12935 #[tokio::test]
12937 async fn reset_invalidates_pending_skill_confirmation_before_streaming_input() {
12938 let (agent, observed) = state_disambiguation_agent_with_skills(
12939 vec![
12940 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12941 "send_report",
12942 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12943 r#"{"question":"What should I send?","options":null}"#,
12944 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12945 r#"{"question":"Should I send the report to Ada?"}"#,
12946 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
12947 "none",
12948 "Fresh response.",
12949 ],
12950 true,
12951 None,
12952 true,
12953 vec![confirmation_skill()],
12954 );
12955
12956 agent.chat("Send it").await.unwrap();
12957 agent.chat("The report to Ada").await.unwrap();
12958 agent.reset().await.unwrap();
12959 assert!(agent.pending_skill_id.read().is_none());
12960 assert!(
12961 !agent
12962 .disambiguation_manager()
12963 .unwrap()
12964 .has_pending_clarification()
12965 .await
12966 );
12967
12968 let mut stream = agent.chat_stream("Yes").await.unwrap();
12969 let mut content = String::new();
12970 while let Some(chunk) = stream.next().await {
12971 match chunk {
12972 StreamChunk::Content { text } => content.push_str(&text),
12973 StreamChunk::Done {} => break,
12974 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
12975 _ => {}
12976 }
12977 }
12978
12979 assert_eq!(content, "Fresh response.");
12980 assert_eq!(confirmation_skill_call_count(&observed), 0);
12981 }
12982
12983 #[tokio::test]
12985 async fn trait_reset_clears_pending_skill_confirmation() {
12986 let (agent, _) = state_disambiguation_agent_with_skills(
12987 vec![
12988 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
12989 "send_report",
12990 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
12991 r#"{"question":"What should I send?","options":null}"#,
12992 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
12993 r#"{"question":"Should I send the report to Ada?"}"#,
12994 ],
12995 true,
12996 None,
12997 true,
12998 vec![confirmation_skill()],
12999 );
13000
13001 agent.chat("Send it").await.unwrap();
13002 agent.chat("The report to Ada").await.unwrap();
13003 <RuntimeAgent as Agent>::reset(&agent).await.unwrap();
13004
13005 assert!(agent.pending_skill_id.read().is_none());
13006 assert!(
13007 !agent
13008 .disambiguation_manager()
13009 .unwrap()
13010 .has_pending_clarification()
13011 .await
13012 );
13013 }
13014
13015 #[tokio::test]
13017 async fn state_change_invalidates_pending_skill_confirmation() {
13018 let (agent, observed) = state_disambiguation_agent_with_skills(
13019 vec![
13020 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13021 "send_report",
13022 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13023 r#"{"question":"What should I send?","options":null}"#,
13024 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13025 r#"{"question":"Should I send the report to Ada?"}"#,
13026 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"fresh input","what_is_unclear":[],"detected_language":"en"}"#,
13027 "none",
13028 "Fresh response.",
13029 ],
13030 true,
13031 None,
13032 true,
13033 vec![confirmation_skill()],
13034 );
13035
13036 agent.chat("Send it").await.unwrap();
13037 agent.chat("The report to Ada").await.unwrap();
13038 agent.transition_to("review").await.unwrap();
13039 let cancelled = agent.chat("Yes").await.unwrap();
13040
13041 assert_eq!(cancelled.content, "Fresh response.");
13042 assert_eq!(confirmation_skill_call_count(&observed), 0);
13043 assert!(agent.pending_skill_id.read().is_none());
13044 }
13045
13046 #[tokio::test]
13048 async fn in_flight_confirmation_cannot_redispatch_after_reset() {
13049 let (mut agent, observed) = state_disambiguation_agent_with_skills(
13050 vec![
13051 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13052 "send_report",
13053 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13054 r#"{"question":"What should I send?","options":null}"#,
13055 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13056 r#"{"question":"Should I send the report to Ada?"}"#,
13057 r#"{"status":"confirmed"}"#,
13058 "Confirmation cancelled.",
13059 ],
13060 true,
13061 None,
13062 true,
13063 vec![confirmation_skill()],
13064 );
13065 let observer = Arc::new(BlockingRuntimeConfirmationObserver::new());
13066 let manager = agent
13067 .disambiguation_manager
13068 .take()
13069 .unwrap()
13070 .with_clarification_observer(observer.clone());
13071 agent.disambiguation_manager = Some(manager);
13072 let agent = Arc::new(agent);
13073
13074 agent.chat("Send it").await.unwrap();
13075 agent.chat("The report to Ada").await.unwrap();
13076
13077 let confirming_agent = Arc::clone(&agent);
13078 let confirmation = tokio::spawn(async move { confirming_agent.chat("Yes").await });
13079 observer.entered.wait().await;
13080 agent.reset().await.unwrap();
13081 observer.release.notify_one();
13082
13083 let response = confirmation.await.unwrap().unwrap();
13084 assert_eq!(response.content, "Confirmation cancelled.");
13085 assert_eq!(confirmation_skill_call_count(&observed), 0);
13086 assert!(agent.pending_skill_id.read().is_none());
13087 }
13088
13089 #[tokio::test]
13091 async fn queued_reset_prevents_stale_confirmation_question_publication() {
13092 let (agent, observed) = state_disambiguation_agent(
13093 vec![
13094 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13095 r#"{"question":"What should I send?","options":null}"#,
13096 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13097 r#"{"question":"Should I send the report to Ada?"}"#,
13098 ],
13099 true,
13100 None,
13101 true,
13102 );
13103 let agent = Arc::new(agent);
13104 agent.chat("Send it").await.unwrap();
13105
13106 let admission = agent.disambiguation_admission.write().await;
13107 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13108 let resetting_agent = Arc::clone(&agent);
13109 let reset = tokio::spawn(async move {
13110 let _ = started_tx.send(());
13111 resetting_agent.reset().await
13112 });
13113 started_rx.await.unwrap();
13114 tokio::task::yield_now().await;
13115
13116 let responding_agent = Arc::clone(&agent);
13117 let response =
13118 tokio::spawn(async move { responding_agent.chat("The report to Ada").await });
13119 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13120 while observed.call_count() < 4 {
13121 tokio::task::yield_now().await;
13122 }
13123 })
13124 .await
13125 .expect("clarification processing must reach terminal publication");
13126 drop(admission);
13127
13128 reset.await.unwrap().unwrap();
13129 let error = response.await.unwrap().unwrap_err();
13130 assert!(error.to_string().contains("ownership changed"));
13131 assert!(
13132 !agent
13133 .disambiguation_manager()
13134 .unwrap()
13135 .has_pending_clarification()
13136 .await
13137 );
13138 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13139 }
13140
13141 #[tokio::test]
13143 async fn queued_reset_prevents_stale_skill_clarification_publication() {
13144 let (agent, observed) = state_disambiguation_agent_with_skills(
13145 vec![
13146 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13147 "send_report",
13148 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13149 r#"{"question":"What should I send?","options":null}"#,
13150 ],
13151 true,
13152 None,
13153 true,
13154 vec![confirmation_skill()],
13155 );
13156 let agent = Arc::new(agent);
13157 let admission = agent.disambiguation_admission.write().await;
13158 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
13159 let resetting_agent = Arc::clone(&agent);
13160 let reset = tokio::spawn(async move {
13161 let _ = started_tx.send(());
13162 resetting_agent.reset().await
13163 });
13164 started_rx.await.unwrap();
13165 tokio::task::yield_now().await;
13166
13167 let responding_agent = Arc::clone(&agent);
13168 let response = tokio::spawn(async move { responding_agent.chat("Send it").await });
13169 tokio::time::timeout(std::time::Duration::from_secs(2), async {
13170 while observed.call_count() < 4 {
13171 tokio::task::yield_now().await;
13172 }
13173 })
13174 .await
13175 .expect("skill clarification must reach terminal publication");
13176 drop(admission);
13177
13178 reset.await.unwrap().unwrap();
13179 let error = response.await.unwrap().unwrap_err();
13180 assert!(error.to_string().contains("ownership changed"));
13181 assert_eq!(confirmation_skill_call_count(&observed), 0);
13182 assert!(agent.pending_skill_id.read().is_none());
13183 assert!(agent.memory.get_messages(None).await.unwrap().is_empty());
13184 }
13185
13186 #[tokio::test]
13188 async fn transition_hook_can_reset_without_admission_deadlock() {
13189 let hooks = Arc::new(ResetOnTransitionHooks {
13190 agent: parking_lot::Mutex::new(None),
13191 invoked: AtomicBool::new(false),
13192 });
13193 let agent = Arc::new(
13194 AgentBuilder::new()
13195 .system_prompt("Test transition hook reentrancy.")
13196 .llm(Arc::new(mock_with_response("done")))
13197 .state_machine(disambiguation_state_machine(None, false))
13198 .build()
13199 .unwrap()
13200 .with_hooks(hooks.clone()),
13201 );
13202 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
13203
13204 let transitioned = tokio::time::timeout(
13205 std::time::Duration::from_secs(2),
13206 agent.apply_transition_target("active", "review", "test transition", None),
13207 )
13208 .await
13209 .expect("transition hook reset must not deadlock")
13210 .unwrap();
13211
13212 assert!(transitioned);
13213 assert!(hooks.invoked.load(Ordering::SeqCst));
13214 assert_eq!(agent.current_state().as_deref(), Some("active"));
13215 }
13216
13217 #[tokio::test]
13219 async fn concurrent_transition_cannot_duplicate_exit_actions() {
13220 let gate = PathMutationGate::new();
13221 let active = ai_agents_state::StateDefinition {
13222 on_exit: vec![StateAction::Tool {
13223 tool: "transition_exit".to_string(),
13224 args: Some(serde_json::json!({"path": "./transition-exit.txt"})),
13225 }],
13226 ..Default::default()
13227 };
13228 let state_machine = Arc::new(
13229 StateMachine::new(ai_agents_state::StateConfig {
13230 initial: "active".to_string(),
13231 states: HashMap::from([
13232 ("active".to_string(), active),
13233 (
13234 "review".to_string(),
13235 ai_agents_state::StateDefinition::default(),
13236 ),
13237 ]),
13238 global_transitions: Vec::new(),
13239 fallback: None,
13240 max_no_transition: None,
13241 regenerate_on_transition: true,
13242 })
13243 .unwrap(),
13244 );
13245 let agent = Arc::new(
13246 AgentBuilder::new()
13247 .system_prompt("Test transition reservation.")
13248 .llm(Arc::new(mock_with_response("done")))
13249 .tool(Arc::new(BlockingPathMutationTool {
13250 id: "transition_exit",
13251 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13252 gate: gate.clone(),
13253 }))
13254 .state_machine(state_machine)
13255 .build()
13256 .unwrap(),
13257 );
13258
13259 let first_agent = Arc::clone(&agent);
13260 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13261 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13262 .await
13263 .expect("reserved transition must enter its exit action");
13264
13265 let second = tokio::time::timeout(
13266 std::time::Duration::from_secs(2),
13267 agent.transition_to("review"),
13268 )
13269 .await
13270 .expect("competing transition must fail without waiting for the exit action")
13271 .unwrap_err();
13272 assert!(second.to_string().contains("already in progress"));
13273
13274 gate.release();
13275 first.await.unwrap().unwrap();
13276 assert_eq!(agent.current_state().as_deref(), Some("review"));
13277 }
13278
13279 #[tokio::test]
13281 async fn concurrent_transition_cannot_overtake_enter_actions() {
13282 let gate = PathMutationGate::new();
13283 let review = ai_agents_state::StateDefinition {
13284 on_enter: vec![StateAction::Tool {
13285 tool: "transition_enter".to_string(),
13286 args: Some(serde_json::json!({"path": "./transition-enter.txt"})),
13287 }],
13288 ..Default::default()
13289 };
13290 let state_machine = Arc::new(
13291 StateMachine::new(ai_agents_state::StateConfig {
13292 initial: "active".to_string(),
13293 states: HashMap::from([
13294 (
13295 "active".to_string(),
13296 ai_agents_state::StateDefinition::default(),
13297 ),
13298 ("review".to_string(), review),
13299 ]),
13300 global_transitions: Vec::new(),
13301 fallback: None,
13302 max_no_transition: None,
13303 regenerate_on_transition: true,
13304 })
13305 .unwrap(),
13306 );
13307 let agent = Arc::new(
13308 AgentBuilder::new()
13309 .system_prompt("Test transition lifecycle reservation.")
13310 .llm(Arc::new(mock_with_response("done")))
13311 .tool(Arc::new(BlockingPathMutationTool {
13312 id: "transition_enter",
13313 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
13314 gate: gate.clone(),
13315 }))
13316 .state_machine(state_machine)
13317 .build()
13318 .unwrap(),
13319 );
13320
13321 let first_agent = Arc::clone(&agent);
13322 let first = tokio::spawn(async move { first_agent.transition_to("review").await });
13323 tokio::time::timeout(std::time::Duration::from_secs(2), gate.wait_until_entered())
13324 .await
13325 .expect("committed transition must enter its destination action");
13326
13327 let second = agent.transition_to("active").await.unwrap_err();
13328 assert!(second.to_string().contains("already in progress"));
13329 assert!(agent.reset().await.is_err());
13330
13331 gate.release();
13332 first.await.unwrap().unwrap();
13333 assert_eq!(agent.current_state().as_deref(), Some("review"));
13334 }
13335
13336 #[tokio::test]
13338 async fn same_state_restore_invalidates_pending_skill_confirmation() {
13339 let (agent, observed) = state_disambiguation_agent_with_skills(
13340 vec![
13341 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13342 "send_report",
13343 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13344 r#"{"question":"What should I send?","options":null}"#,
13345 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13346 r#"{"question":"Should I send the report to Ada?"}"#,
13347 ],
13348 true,
13349 None,
13350 true,
13351 vec![confirmation_skill()],
13352 );
13353
13354 agent.chat("Send it").await.unwrap();
13355 agent.chat("The report to Ada").await.unwrap();
13356 let snapshot = agent.save_state().await.unwrap();
13357 assert_eq!(agent.current_state().as_deref(), Some("active"));
13358
13359 agent.restore_state(snapshot).await.unwrap();
13360
13361 assert_eq!(agent.current_state().as_deref(), Some("active"));
13362 assert!(agent.pending_skill_id.read().is_none());
13363 assert!(
13364 !agent
13365 .disambiguation_manager()
13366 .unwrap()
13367 .has_pending_clarification()
13368 .await
13369 );
13370 assert_eq!(confirmation_skill_call_count(&observed), 0);
13371 }
13372
13373 #[tokio::test]
13375 async fn direct_state_generation_change_invalidates_confirmation() {
13376 let (agent, observed) = state_disambiguation_agent_with_skills(
13377 vec![
13378 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"top-level clear","what_is_unclear":[],"detected_language":"en"}"#,
13379 "send_report",
13380 r#"{"is_ambiguous":true,"confidence":0.2,"ambiguity_type":"missing_target","reasoning":"target missing","what_is_unclear":["target"],"detected_language":"en"}"#,
13381 r#"{"question":"What should I send?","options":null}"#,
13382 r#"{"status":"answered","selected_option":null,"enriched_input":"Send the report to Ada","resolved":{"intent":"send_report"}}"#,
13383 r#"{"question":"Should I send the report to Ada?"}"#,
13384 "Confirmation cancelled.",
13385 ],
13386 true,
13387 None,
13388 true,
13389 vec![confirmation_skill()],
13390 );
13391
13392 agent.chat("Send it").await.unwrap();
13393 agent.chat("The report to Ada").await.unwrap();
13394 let state_machine = agent.state_machine().unwrap();
13395 state_machine
13396 .transition_to("review", "external test")
13397 .unwrap();
13398 state_machine
13399 .transition_to("active", "external test")
13400 .unwrap();
13401
13402 let response = agent.chat("Yes").await.unwrap();
13403
13404 assert_eq!(response.content, "Confirmation cancelled.");
13405 assert_eq!(confirmation_skill_call_count(&observed), 0);
13406 assert!(agent.pending_skill_id.read().is_none());
13407 }
13408
13409 #[tokio::test]
13410 async fn state_confirmation_does_not_add_a_question_for_clear_input() {
13411 let (agent, observed) = state_disambiguation_agent(
13412 vec![
13413 r#"{"is_ambiguous":false,"confidence":0.99,"ambiguity_type":null,"reasoning":"clear","what_is_unclear":[],"detected_language":"en"}"#,
13414 "Request executed.",
13415 ],
13416 true,
13417 None,
13418 true,
13419 );
13420
13421 let response = agent.chat("Send the report to Ada").await.unwrap();
13422
13423 assert_eq!(response.content, "Request executed.");
13424 assert_eq!(observed.call_count(), 2);
13425 }
13426
13427 #[tokio::test]
13428 async fn state_override_cannot_activate_a_disabled_top_level_manager() {
13429 let (agent, observed) =
13430 state_disambiguation_agent(vec!["Request executed."], false, Some(true), true);
13431
13432 assert!(!agent.has_disambiguation());
13433 let response = agent.chat("Send it").await.unwrap();
13434
13435 assert_eq!(response.content, "Request executed.");
13436 assert_eq!(observed.call_count(), 1);
13437 }
13438
13439 #[tokio::test]
13440 async fn native_required_choice_executes_through_the_shared_tool_path() {
13441 let mut mock = MockLLMProvider::new("native-required");
13442 mock.set_tool_choice(Some(ToolChoice::Required));
13443 mock.add_response(
13444 LLMResponse::new("", FinishReason::ToolCall)
13445 .with_tool_calls(vec![ToolCall {
13446 id: "provider-call-1".to_string(),
13447 name: "calculator".to_string(),
13448 arguments: serde_json::json!({"expression": "2 + 2"}),
13449 }])
13450 .unwrap(),
13451 );
13452 mock.add_response(LLMResponse::new("The answer is 4.", FinishReason::Stop));
13453 let observed = mock.clone();
13454 let agent = AgentBuilder::new()
13455 .system_prompt("Use the calculator when needed.")
13456 .llm(Arc::new(mock))
13457 .tool(Arc::new(CalculatorTool::new()))
13458 .build()
13459 .unwrap();
13460
13461 let response = agent.chat("What is 2 + 2?").await.unwrap();
13462
13463 assert_eq!(response.content, "The answer is 4.");
13464 assert_eq!(
13465 response.tool_calls.as_ref().unwrap()[0].id,
13466 "provider-call-1"
13467 );
13468 let calls = observed.call_history();
13469 assert_eq!(calls.len(), 2);
13470 assert!(matches!(
13471 calls[0].request.as_ref().map(|request| &request.choice),
13472 Some(ToolChoice::Required)
13473 ));
13474 assert!(matches!(
13475 calls[1].request.as_ref().map(|request| &request.choice),
13476 Some(ToolChoice::Auto)
13477 ));
13478 }
13479
13480 #[tokio::test]
13481 async fn prompt_fallback_uses_one_corrective_retry() {
13482 let mut mock = MockLLMProvider::new("prompt-required");
13483 mock.set_tool_choice(Some(ToolChoice::Required));
13484 mock.set_native_tool_support(false);
13485 mock.set_responses(
13486 vec![
13487 "I can calculate that.".to_string(),
13488 r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#.to_string(),
13489 "The answer is 4.".to_string(),
13490 ],
13491 false,
13492 );
13493 let observed = mock.clone();
13494 let agent = AgentBuilder::new()
13495 .system_prompt("Use tools.")
13496 .llm(Arc::new(mock))
13497 .tool(Arc::new(CalculatorTool::new()))
13498 .build()
13499 .unwrap();
13500
13501 let response = agent.chat("What is 2 + 2?").await.unwrap();
13502
13503 assert_eq!(response.content, "The answer is 4.");
13504 assert_eq!(observed.call_count(), 3);
13505 let corrective = &observed.call_history()[1].messages;
13506 assert!(
13507 corrective
13508 .last()
13509 .unwrap()
13510 .content
13511 .contains("previous response")
13512 );
13513 }
13514
13515 #[tokio::test]
13516 async fn prompt_fallback_fails_after_one_noncompliant_retry() {
13517 let mut mock = MockLLMProvider::new("prompt-required-failure");
13518 mock.set_tool_choice(Some(ToolChoice::Required));
13519 mock.set_native_tool_support(false);
13520 mock.set_responses(
13521 vec!["No tool.".to_string(), "Still no tool.".to_string()],
13522 false,
13523 );
13524 let observed = mock.clone();
13525 let agent = AgentBuilder::new()
13526 .system_prompt("Use tools.")
13527 .llm(Arc::new(mock))
13528 .tool(Arc::new(CalculatorTool::new()))
13529 .build()
13530 .unwrap();
13531
13532 let error = agent.chat("What is 2 + 2?").await.unwrap_err();
13533
13534 assert!(error.to_string().contains("one corrective retry"));
13535 assert_eq!(observed.call_count(), 2);
13536 }
13537
13538 #[tokio::test]
13539 async fn specific_choice_cannot_widen_the_effective_grant() {
13540 let mut mock = MockLLMProvider::new("specific-outside-grant");
13541 mock.set_tool_choice(Some(ToolChoice::Specific("random".to_string())));
13542 let observed = mock.clone();
13543 let agent = AgentBuilder::new()
13544 .system_prompt("Use tools.")
13545 .llm(Arc::new(mock))
13546 .tool(Arc::new(CalculatorTool::new()))
13547 .build()
13548 .unwrap();
13549
13550 let error = agent.chat("Generate a value.").await.unwrap_err();
13551
13552 assert!(error.to_string().contains("is not registered"));
13553 assert_eq!(observed.call_count(), 0);
13554 }
13555
13556 #[tokio::test]
13557 async fn none_choice_exposes_no_tool_protocol() {
13558 let mut mock = MockLLMProvider::new("no-tools");
13559 mock.set_tool_choice(Some(ToolChoice::None));
13560 mock.set_response(r#"{"tool":"calculator","arguments":{"expression":"2 + 2"}}"#);
13561 let observed = mock.clone();
13562 let agent = AgentBuilder::new()
13563 .system_prompt("Answer directly.")
13564 .llm(Arc::new(mock))
13565 .tool(Arc::new(CalculatorTool::new()))
13566 .build()
13567 .unwrap();
13568
13569 let response = agent.chat("Hello").await.unwrap();
13570
13571 assert!(response.tool_calls.is_none());
13572 assert_eq!(observed.call_count(), 1);
13573 let call = observed.last_call().unwrap();
13574 assert!(call.request.is_none());
13575 assert!(
13576 call.messages
13577 .iter()
13578 .all(|message| !message.content.contains("Available tools:"))
13579 );
13580 }
13581
13582 struct RuntimeStorage {
13583 capabilities: Box<[StorageCapability]>,
13584 snapshots: RwLock<HashMap<String, AgentSnapshot>>,
13585 metadata: RwLock<HashMap<String, ai_agents_core::SessionMetadata>>,
13586 metadata_save_calls: AtomicU64,
13587 metadata_load_calls: AtomicU64,
13588 fail_metadata_save: AtomicBool,
13589 fail_metadata_load: AtomicBool,
13590 }
13591
13592 impl RuntimeStorage {
13593 fn new(capabilities: impl IntoIterator<Item = StorageCapability>) -> Self {
13594 Self {
13595 capabilities: capabilities.into_iter().collect(),
13596 snapshots: RwLock::new(HashMap::new()),
13597 metadata: RwLock::new(HashMap::new()),
13598 metadata_save_calls: AtomicU64::new(0),
13599 metadata_load_calls: AtomicU64::new(0),
13600 fail_metadata_save: AtomicBool::new(false),
13601 fail_metadata_load: AtomicBool::new(false),
13602 }
13603 }
13604 }
13605
13606 #[async_trait]
13607 impl AgentStorage for RuntimeStorage {
13608 fn supports(&self, capability: StorageCapability) -> bool {
13609 self.capabilities.contains(&capability)
13610 }
13611
13612 async fn save(&self, session_id: &str, snapshot: &AgentSnapshot) -> Result<()> {
13613 self.snapshots
13614 .write()
13615 .insert(session_id.to_string(), snapshot.clone());
13616 Ok(())
13617 }
13618
13619 async fn load(&self, session_id: &str) -> Result<Option<AgentSnapshot>> {
13620 Ok(self.snapshots.read().get(session_id).cloned())
13621 }
13622
13623 async fn delete(&self, session_id: &str) -> Result<()> {
13624 self.snapshots.write().remove(session_id);
13625 Ok(())
13626 }
13627
13628 async fn list_sessions(&self) -> Result<Vec<String>> {
13629 Ok(self.snapshots.read().keys().cloned().collect())
13630 }
13631
13632 async fn save_snapshot_with_metadata(
13633 &self,
13634 session_id: &str,
13635 snapshot: &AgentSnapshot,
13636 metadata: &ai_agents_core::SessionMetadata,
13637 ) -> Result<()> {
13638 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13639 if self.fail_metadata_save.load(Ordering::SeqCst) {
13640 return Err(AgentError::Persistence("metadata save failed".into()));
13641 }
13642 self.snapshots
13643 .write()
13644 .insert(session_id.to_string(), snapshot.clone());
13645 self.metadata
13646 .write()
13647 .insert(session_id.to_string(), metadata.clone());
13648 Ok(())
13649 }
13650
13651 async fn save_metadata(
13652 &self,
13653 session_id: &str,
13654 metadata: &ai_agents_core::SessionMetadata,
13655 ) -> Result<()> {
13656 self.metadata_save_calls.fetch_add(1, Ordering::SeqCst);
13657 if self.fail_metadata_save.load(Ordering::SeqCst) {
13658 return Err(AgentError::Persistence("metadata save failed".into()));
13659 }
13660 self.metadata
13661 .write()
13662 .insert(session_id.to_string(), metadata.clone());
13663 Ok(())
13664 }
13665
13666 async fn load_metadata(
13667 &self,
13668 session_id: &str,
13669 ) -> Result<Option<ai_agents_core::SessionMetadata>> {
13670 self.metadata_load_calls.fetch_add(1, Ordering::SeqCst);
13671 if self.fail_metadata_load.load(Ordering::SeqCst) {
13672 return Err(AgentError::Persistence("metadata load failed".into()));
13673 }
13674 Ok(self.metadata.read().get(session_id).cloned())
13675 }
13676 }
13677
13678 fn runtime_storage_agent() -> RuntimeAgent {
13679 AgentBuilder::new()
13680 .system_prompt("Test runtime storage integration.")
13681 .llm(Arc::new(mock_with_response("done")))
13682 .build()
13683 .unwrap()
13684 }
13685
13686 fn restore_spec(id: &str) -> crate::spec::AgentSpec {
13687 crate::spec::AgentSpec {
13688 name: id.to_string(),
13689 system_prompt: format!("Restore child {id}."),
13690 ..crate::spec::AgentSpec::default()
13691 }
13692 }
13693
13694 fn restore_entry(id: &str) -> ai_agents_core::SpawnedAgentEntry {
13695 ai_agents_core::SpawnedAgentEntry {
13696 id: id.to_string(),
13697 name: id.to_string(),
13698 spec_yaml: serde_yaml::to_string(&restore_spec(id)).unwrap(),
13699 }
13700 }
13701
13702 fn restore_spawner(
13703 storage: Arc<RuntimeStorage>,
13704 max_agents: usize,
13705 ) -> (
13706 Arc<crate::spawner::AgentSpawner>,
13707 Arc<crate::spawner::AgentRegistry>,
13708 ) {
13709 let mut llms = LLMRegistry::new();
13710 llms.register("default", Arc::new(mock_with_response("done")));
13711 (
13712 Arc::new(
13713 crate::spawner::AgentSpawner::new()
13714 .with_shared_llms(llms)
13715 .with_shared_storage(storage)
13716 .with_max_agents(max_agents),
13717 ),
13718 Arc::new(crate::spawner::AgentRegistry::new()),
13719 )
13720 }
13721
13722 async fn save_restore_target(
13723 parent: &RuntimeAgent,
13724 storage: &RuntimeStorage,
13725 session_id: &str,
13726 entries: Vec<ai_agents_core::SpawnedAgentEntry>,
13727 ) {
13728 let mut snapshot = parent.save_state().await.unwrap();
13729 snapshot.spawned_agents = Some(entries);
13730 storage.save(session_id, &snapshot).await.unwrap();
13731 storage
13732 .save_metadata(session_id, &ai_agents_core::SessionMetadata::default())
13733 .await
13734 .unwrap();
13735 }
13736
13737 #[tokio::test]
13738 async fn storage_init_requires_storage_for_actor_facts() {
13739 let facts = ai_agents_facts::FactsConfig {
13740 enabled: true,
13741 ..Default::default()
13742 };
13743 let agent = runtime_storage_agent().with_facts_config(None, Some(facts));
13744
13745 let error = agent.init_storage().await.unwrap_err();
13746 assert!(matches!(
13747 error,
13748 AgentError::Config(message)
13749 if message.contains("actor facts or actor memory")
13750 && message.contains("none is configured or injected")
13751 ));
13752 }
13753
13754 #[tokio::test]
13755 async fn storage_init_validates_actor_facts_capability() {
13756 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13757 let actor_memory = ai_agents_facts::ActorMemoryConfig {
13758 enabled: true,
13759 ..Default::default()
13760 };
13761 let agent = runtime_storage_agent()
13762 .with_storage(storage)
13763 .with_facts_config(Some(actor_memory), None);
13764
13765 assert!(matches!(
13766 agent.init_storage().await,
13767 Err(AgentError::UnsupportedStorageCapability(
13768 StorageCapability::ActorFacts
13769 ))
13770 ));
13771 }
13772
13773 #[tokio::test]
13774 async fn blocking_chat_rejects_unsupported_required_storage() {
13775 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13776 let facts = ai_agents_facts::FactsConfig {
13777 enabled: true,
13778 ..Default::default()
13779 };
13780 let agent = runtime_storage_agent()
13781 .with_storage(storage)
13782 .with_facts_config(None, Some(facts));
13783
13784 assert!(matches!(
13785 agent.chat("hello").await,
13786 Err(AgentError::UnsupportedStorageCapability(
13787 StorageCapability::ActorFacts
13788 ))
13789 ));
13790 }
13791
13792 #[tokio::test]
13793 async fn streaming_chat_rejects_unsupported_required_storage_before_stream_creation() {
13794 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13795 let config = ai_agents_relationships::RelationshipConfig {
13796 enabled: true,
13797 ..Default::default()
13798 };
13799 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13800 let agent = runtime_storage_agent()
13801 .with_storage(storage)
13802 .with_relationships(manager);
13803
13804 assert!(matches!(
13805 agent.chat_stream("hello").await,
13806 Err(AgentError::UnsupportedStorageCapability(
13807 StorageCapability::ActorRelationships
13808 ))
13809 ));
13810 }
13811
13812 #[tokio::test]
13813 async fn storage_init_completes_facts_for_injected_storage() {
13814 let storage = Arc::new(RuntimeStorage::new([
13815 StorageCapability::Snapshot,
13816 StorageCapability::ActorFacts,
13817 ]));
13818 let facts = ai_agents_facts::FactsConfig {
13819 enabled: true,
13820 ..Default::default()
13821 };
13822 let agent = runtime_storage_agent()
13823 .with_storage(storage)
13824 .with_facts_config(None, Some(facts));
13825
13826 agent.init_storage().await.unwrap();
13827 assert!(agent.fact_store().is_some());
13828 }
13829
13830 #[tokio::test]
13831 async fn storage_init_requires_storage_for_persistent_relationships() {
13832 let config = ai_agents_relationships::RelationshipConfig {
13833 enabled: true,
13834 ..Default::default()
13835 };
13836 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13837 let agent = runtime_storage_agent().with_relationships(manager);
13838
13839 let error = agent.init_storage().await.unwrap_err();
13840 assert!(matches!(
13841 error,
13842 AgentError::Config(message)
13843 if message.contains("persistent relationships")
13844 && message.contains("none is configured or injected")
13845 ));
13846 }
13847
13848 #[tokio::test]
13849 async fn storage_init_validates_persistent_relationships_capability() {
13850 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13851 let config = ai_agents_relationships::RelationshipConfig {
13852 enabled: true,
13853 ..Default::default()
13854 };
13855 let manager = Arc::new(RelationshipManager::from_config(config).unwrap());
13856 let agent = runtime_storage_agent()
13857 .with_storage(storage)
13858 .with_relationships(manager);
13859
13860 assert!(matches!(
13861 agent.init_storage().await,
13862 Err(AgentError::UnsupportedStorageCapability(
13863 StorageCapability::ActorRelationships
13864 ))
13865 ));
13866 }
13867
13868 #[tokio::test]
13869 async fn session_restore_updates_identity_and_clears_stale_actor_binding() {
13870 let storage = Arc::new(RuntimeStorage::new([
13871 StorageCapability::Snapshot,
13872 StorageCapability::SessionMetadata,
13873 ]));
13874 let agent = runtime_storage_agent().with_storage(storage.clone());
13875 agent.set_actor_id("old-actor").unwrap();
13876 agent.save_session("old").await.unwrap();
13877 storage
13878 .save("target", &agent.save_state().await.unwrap())
13879 .await
13880 .unwrap();
13881 storage
13882 .save_metadata("target", &ai_agents_core::SessionMetadata::default())
13883 .await
13884 .unwrap();
13885
13886 assert!(agent.load_session("target").await.unwrap());
13887
13888 assert_eq!(agent.current_session_id.read().as_deref(), Some("target"));
13889 assert_eq!(agent.actor_id(), None);
13890 }
13891
13892 #[tokio::test]
13893 async fn complete_restore_reconciles_growth_shrink_and_empty_topologies() {
13894 let storage = Arc::new(RuntimeStorage::new([
13895 StorageCapability::Snapshot,
13896 StorageCapability::SessionMetadata,
13897 ]));
13898 let (spawner, registry) = restore_spawner(storage.clone(), 3);
13899 let parent = runtime_storage_agent()
13900 .with_storage(storage.clone())
13901 .with_spawner_handles(Arc::clone(&spawner), Arc::clone(®istry));
13902
13903 for id in ["a", "b"] {
13904 let spawned = spawner
13905 .spawn_with_id(id.to_string(), restore_spec(id))
13906 .await
13907 .unwrap();
13908 spawned.agent.save_session("grow").await.unwrap();
13909 registry.register(spawned).await.unwrap();
13910 }
13911 let staged_c = crate::spawner::storage::NamespacedStorage::new(storage.clone(), "c");
13912 staged_c
13913 .save("grow", &AgentSnapshot::new("c".into()))
13914 .await
13915 .unwrap();
13916 staged_c
13917 .save_metadata("grow", &ai_agents_core::SessionMetadata::default())
13918 .await
13919 .unwrap();
13920 save_restore_target(
13921 &parent,
13922 storage.as_ref(),
13923 "grow",
13924 vec![restore_entry("a"), restore_entry("b"), restore_entry("c")],
13925 )
13926 .await;
13927
13928 assert_eq!(parent.restore_session_full("grow").await.unwrap(), 3);
13929 assert_eq!(registry.count(), 3);
13930 assert!(registry.contains("c"));
13931 assert_eq!(spawner.spawned_count(), 3);
13932
13933 for id in ["a", "b"] {
13934 registry
13935 .get(id)
13936 .unwrap()
13937 .save_session("shrink")
13938 .await
13939 .unwrap();
13940 }
13941 save_restore_target(
13942 &parent,
13943 storage.as_ref(),
13944 "shrink",
13945 vec![restore_entry("a"), restore_entry("b")],
13946 )
13947 .await;
13948
13949 assert_eq!(parent.restore_session_full("shrink").await.unwrap(), 2);
13950 assert_eq!(registry.count(), 2);
13951 assert!(!registry.contains("c"));
13952 assert_eq!(spawner.spawned_count(), 2);
13953
13954 save_restore_target(&parent, storage.as_ref(), "empty", Vec::new()).await;
13955
13956 assert_eq!(parent.restore_session_full("empty").await.unwrap(), 0);
13957 assert_eq!(registry.count(), 0);
13958 assert_eq!(spawner.spawned_count(), 0);
13959 assert_eq!(parent.current_session_id.read().as_deref(), Some("empty"));
13960 }
13961
13962 #[tokio::test]
13963 async fn storage_session_metadata_is_called_only_when_advertised() {
13964 let storage = Arc::new(RuntimeStorage::new([StorageCapability::Snapshot]));
13965 storage.fail_metadata_save.store(true, Ordering::SeqCst);
13966 storage.fail_metadata_load.store(true, Ordering::SeqCst);
13967 let agent = runtime_storage_agent().with_storage(storage.clone());
13968
13969 agent.save_session("session").await.unwrap();
13970 assert!(agent.load_session("session").await.unwrap());
13971 assert_eq!(storage.metadata_save_calls.load(Ordering::SeqCst), 0);
13972 assert_eq!(storage.metadata_load_calls.load(Ordering::SeqCst), 0);
13973 }
13974
13975 #[cfg(feature = "sqlite")]
13976 #[tokio::test]
13977 async fn sqlite_runtime_save_filter_reopen_and_reload_stay_consistent() {
13978 let directory =
13979 std::env::temp_dir().join(format!("ai-agents-runtime-sqlite-{}", uuid::Uuid::new_v4()));
13980 let path = directory.join("sessions.sqlite");
13981 let path_string = path.to_string_lossy().into_owned();
13982 let storage = Arc::new(
13983 ai_agents_storage::SqliteStorage::new(&path_string)
13984 .await
13985 .unwrap(),
13986 );
13987 let agent = runtime_storage_agent().with_storage(storage.clone());
13988 agent.set_session_metadata(ai_agents_core::SessionMetadata {
13989 tags: vec!["initial".into()],
13990 ..Default::default()
13991 });
13992 agent.chat("persist this turn").await.unwrap();
13993 agent.save_session("session").await.unwrap();
13994
13995 agent.set_session_metadata(ai_agents_core::SessionMetadata {
13996 tags: vec!["updated".into()],
13997 ..Default::default()
13998 });
13999 agent.save_session("session").await.unwrap();
14000 assert!(
14001 agent
14002 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14003 tags: Some(vec!["initial".into()]),
14004 ..Default::default()
14005 })
14006 .await
14007 .unwrap()
14008 .is_empty()
14009 );
14010 assert_eq!(
14011 agent
14012 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14013 tags: Some(vec!["updated".into()]),
14014 ..Default::default()
14015 })
14016 .await
14017 .unwrap()
14018 .len(),
14019 1
14020 );
14021 drop(agent);
14022 storage.close().await;
14023 drop(storage);
14024
14025 let reopened_storage = Arc::new(
14026 ai_agents_storage::SqliteStorage::new(&path_string)
14027 .await
14028 .unwrap(),
14029 );
14030 let restored = runtime_storage_agent().with_storage(reopened_storage.clone());
14031 assert!(restored.load_session("session").await.unwrap());
14032 assert_eq!(restored.session_metadata().tags, vec!["updated"]);
14033 assert_eq!(
14034 restored.current_session_id.read().as_deref(),
14035 Some("session")
14036 );
14037 assert!(restored.save_state().await.unwrap().memory.messages.len() >= 2);
14038 assert_eq!(
14039 restored
14040 .list_sessions_filtered(&ai_agents_core::SessionFilter {
14041 tags: Some(vec!["updated".into()]),
14042 ..Default::default()
14043 })
14044 .await
14045 .unwrap()
14046 .len(),
14047 1
14048 );
14049
14050 drop(restored);
14051 reopened_storage.close().await;
14052 drop(reopened_storage);
14053 crate::remove_sqlite_test_directory(&directory)
14054 .await
14055 .unwrap();
14056 }
14057
14058 #[tokio::test]
14059 async fn storage_session_metadata_backend_failures_propagate() {
14060 let storage = Arc::new(RuntimeStorage::new([
14061 StorageCapability::Snapshot,
14062 StorageCapability::SessionMetadata,
14063 ]));
14064 let agent = runtime_storage_agent().with_storage(storage.clone());
14065
14066 agent.save_session("session").await.unwrap();
14067 storage
14068 .save("target", &agent.save_state().await.unwrap())
14069 .await
14070 .unwrap();
14071 storage.fail_metadata_load.store(true, Ordering::SeqCst);
14072 assert!(matches!(
14073 agent.load_session("target").await,
14074 Err(AgentError::Persistence(message)) if message == "metadata load failed"
14075 ));
14076 assert_eq!(agent.current_session_id.read().as_deref(), Some("session"));
14077
14078 storage.fail_metadata_save.store(true, Ordering::SeqCst);
14079 assert!(matches!(
14080 agent.save_session("session").await,
14081 Err(AgentError::Persistence(message)) if message == "metadata save failed"
14082 ));
14083 }
14084
14085 struct ProviderFutureDropSignal {
14086 dropped: Arc<AtomicBool>,
14087 }
14088
14089 impl Drop for ProviderFutureDropSignal {
14090 fn drop(&mut self) {
14091 self.dropped.store(true, Ordering::SeqCst);
14092 }
14093 }
14094
14095 struct BufferedLockingProvider {
14096 lock: Arc<tokio::sync::Mutex<()>>,
14097 stream_started: Arc<tokio::sync::Notify>,
14098 stream_dropped: Arc<AtomicBool>,
14099 committed_after_drop: Arc<AtomicBool>,
14100 }
14101
14102 #[async_trait]
14103 impl LLMProvider for BufferedLockingProvider {
14104 async fn complete(
14105 &self,
14106 _messages: &[ChatMessage],
14107 _config: Option<&LLMConfig>,
14108 ) -> std::result::Result<LLMResponse, LLMError> {
14109 let _guard = self.lock.lock().await;
14110 self.committed_after_drop
14111 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14112 Ok(LLMResponse::new(
14113 "Committed technical response.",
14114 FinishReason::Stop,
14115 ))
14116 }
14117
14118 async fn complete_stream(
14119 &self,
14120 _messages: &[ChatMessage],
14121 _config: Option<&LLMConfig>,
14122 ) -> std::result::Result<
14123 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14124 LLMError,
14125 > {
14126 let _guard = self.lock.lock().await;
14127 let _drop_signal = ProviderFutureDropSignal {
14128 dropped: Arc::clone(&self.stream_dropped),
14129 };
14130 self.stream_started.notify_one();
14131 std::future::pending().await
14132 }
14133
14134 fn provider_name(&self) -> &str {
14135 "buffered-locking"
14136 }
14137
14138 fn supports(&self, _feature: LLMFeature) -> bool {
14139 false
14140 }
14141 }
14142
14143 struct PendingDropStream {
14144 dropped: Arc<AtomicBool>,
14145 dropped_notify: Arc<tokio::sync::Notify>,
14146 }
14147
14148 impl Stream for PendingDropStream {
14149 type Item = std::result::Result<LLMChunk, LLMError>;
14150
14151 fn poll_next(
14152 self: Pin<&mut Self>,
14153 _cx: &mut std::task::Context<'_>,
14154 ) -> std::task::Poll<Option<Self::Item>> {
14155 std::task::Poll::Pending
14156 }
14157 }
14158
14159 impl Drop for PendingDropStream {
14160 fn drop(&mut self) {
14161 self.dropped.store(true, Ordering::SeqCst);
14162 self.dropped_notify.notify_one();
14163 }
14164 }
14165
14166 struct EstablishedStreamProvider {
14167 stream_started: Arc<tokio::sync::Notify>,
14168 stream_dropped: Arc<AtomicBool>,
14169 stream_dropped_notify: Arc<tokio::sync::Notify>,
14170 committed_after_drop: Arc<AtomicBool>,
14171 }
14172
14173 #[async_trait]
14174 impl LLMProvider for EstablishedStreamProvider {
14175 async fn complete(
14176 &self,
14177 _messages: &[ChatMessage],
14178 _config: Option<&LLMConfig>,
14179 ) -> std::result::Result<LLMResponse, LLMError> {
14180 if !self.stream_dropped.load(Ordering::SeqCst) {
14181 self.stream_dropped_notify.notified().await;
14182 }
14183 self.committed_after_drop
14184 .store(self.stream_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14185 Ok(LLMResponse::new(
14186 "Committed technical response.",
14187 FinishReason::Stop,
14188 ))
14189 }
14190
14191 async fn complete_stream(
14192 &self,
14193 _messages: &[ChatMessage],
14194 _config: Option<&LLMConfig>,
14195 ) -> std::result::Result<
14196 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14197 LLMError,
14198 > {
14199 self.stream_started.notify_one();
14200 Ok(Box::new(PendingDropStream {
14201 dropped: Arc::clone(&self.stream_dropped),
14202 dropped_notify: Arc::clone(&self.stream_dropped_notify),
14203 }))
14204 }
14205
14206 fn provider_name(&self) -> &str {
14207 "established-stream"
14208 }
14209
14210 fn supports(&self, _feature: LLMFeature) -> bool {
14211 false
14212 }
14213 }
14214
14215 struct FirstCallLockingProvider {
14216 lock: Arc<tokio::sync::Mutex<()>>,
14217 first_started: Arc<tokio::sync::Notify>,
14218 first_dropped: Arc<AtomicBool>,
14219 committed_after_drop: Arc<AtomicBool>,
14220 calls: AtomicU64,
14221 }
14222
14223 #[async_trait]
14224 impl LLMProvider for FirstCallLockingProvider {
14225 async fn complete(
14226 &self,
14227 _messages: &[ChatMessage],
14228 _config: Option<&LLMConfig>,
14229 ) -> std::result::Result<LLMResponse, LLMError> {
14230 let _guard = self.lock.lock().await;
14231 let call = self.calls.fetch_add(1, Ordering::SeqCst);
14232 if call == 0 {
14233 let _drop_signal = ProviderFutureDropSignal {
14234 dropped: Arc::clone(&self.first_dropped),
14235 };
14236 self.first_started.notify_one();
14237 return std::future::pending().await;
14238 }
14239 self.committed_after_drop
14240 .store(self.first_dropped.load(Ordering::SeqCst), Ordering::SeqCst);
14241 Ok(LLMResponse::new(
14242 "Committed technical response.",
14243 FinishReason::Stop,
14244 ))
14245 }
14246
14247 async fn complete_stream(
14248 &self,
14249 _messages: &[ChatMessage],
14250 _config: Option<&LLMConfig>,
14251 ) -> std::result::Result<
14252 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14253 LLMError,
14254 > {
14255 Err(LLMError::Other(
14256 "streaming is not used in this test".to_string(),
14257 ))
14258 }
14259
14260 fn provider_name(&self) -> &str {
14261 "first-call-locking"
14262 }
14263
14264 fn supports(&self, _feature: LLMFeature) -> bool {
14265 false
14266 }
14267 }
14268
14269 struct RoutingAfterProviderStart {
14270 provider_started: Arc<tokio::sync::Notify>,
14271 }
14272
14273 #[async_trait]
14274 impl LLMProvider for RoutingAfterProviderStart {
14275 async fn complete(
14276 &self,
14277 _messages: &[ChatMessage],
14278 _config: Option<&LLMConfig>,
14279 ) -> std::result::Result<LLMResponse, LLMError> {
14280 self.provider_started.notified().await;
14281 Ok(LLMResponse::new("1", FinishReason::Stop))
14282 }
14283
14284 async fn complete_stream(
14285 &self,
14286 _messages: &[ChatMessage],
14287 _config: Option<&LLMConfig>,
14288 ) -> std::result::Result<
14289 Box<dyn Stream<Item = std::result::Result<LLMChunk, LLMError>> + Unpin + Send>,
14290 LLMError,
14291 > {
14292 Err(LLMError::Other(
14293 "streaming is not used in this test".to_string(),
14294 ))
14295 }
14296
14297 fn provider_name(&self) -> &str {
14298 "routing-after-start"
14299 }
14300
14301 fn supports(&self, _feature: LLMFeature) -> bool {
14302 false
14303 }
14304 }
14305
14306 struct ResponseCountingHooks {
14308 responses: Arc<std::sync::atomic::AtomicUsize>,
14309 }
14310
14311 struct ContextEchoTool;
14313
14314 #[async_trait]
14315 impl ai_agents_core::Tool for ContextEchoTool {
14316 fn id(&self) -> &str {
14317 "context_echo"
14318 }
14319
14320 fn name(&self) -> &str {
14321 "Context Echo"
14322 }
14323
14324 fn description(&self) -> &str {
14325 "Returns selected execution context fields."
14326 }
14327
14328 fn input_schema(&self) -> Value {
14329 serde_json::json!({"type": "object"})
14330 }
14331
14332 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14333 ai_agents_core::ToolPolicyBindings {
14334 path_fields: vec![ai_agents_core::PathPolicyBinding::read("path")],
14335 result_limit_fields: vec![ai_agents_core::ResultLimitBinding::new(
14336 "max_results",
14337 ai_agents_core::ResultLimitKind::MaxResults,
14338 )],
14339 ..Default::default()
14340 }
14341 }
14342
14343 async fn execute(
14344 &self,
14345 _args: Value,
14346 ctx: ai_agents_core::ToolExecutionContext,
14347 ) -> ToolResult {
14348 ToolResult::ok(
14349 serde_json::json!({
14350 "requested_name": ctx.requested_name,
14351 "canonical_id": ctx.canonical_id,
14352 "display_name": ctx.display_name,
14353 "max_results": ctx.limits.max_results,
14354 "custom_config": ctx.custom_config,
14355 })
14356 .to_string(),
14357 )
14358 }
14359 }
14360
14361 struct SlowTool;
14363
14364 struct FlakyWriteTool {
14366 calls: Arc<std::sync::atomic::AtomicUsize>,
14367 }
14368
14369 struct LockedWriteTool {
14371 active: Arc<std::sync::atomic::AtomicUsize>,
14372 max_active: Arc<std::sync::atomic::AtomicUsize>,
14373 }
14374
14375 struct MultiResourceWriteTool {
14376 active: Arc<std::sync::atomic::AtomicUsize>,
14377 max_active: Arc<std::sync::atomic::AtomicUsize>,
14378 }
14379
14380 #[derive(Clone)]
14381 struct PathMutationGate {
14382 entered: Arc<AtomicBool>,
14383 entered_notify: Arc<tokio::sync::Notify>,
14384 release: Arc<tokio::sync::Notify>,
14385 }
14386
14387 impl PathMutationGate {
14388 fn new() -> Self {
14389 Self {
14390 entered: Arc::new(AtomicBool::new(false)),
14391 entered_notify: Arc::new(tokio::sync::Notify::new()),
14392 release: Arc::new(tokio::sync::Notify::new()),
14393 }
14394 }
14395
14396 async fn wait_until_entered(&self) {
14397 if !self.entered.load(Ordering::SeqCst) {
14398 self.entered_notify.notified().await;
14399 }
14400 }
14401
14402 fn release(&self) {
14403 self.release.notify_one();
14404 }
14405 }
14406
14407 struct BlockingPathMutationTool {
14408 id: &'static str,
14409 path_fields: Vec<ai_agents_core::PathPolicyBinding>,
14410 gate: PathMutationGate,
14411 }
14412
14413 struct NoBindingWriteTool {
14414 active: Arc<std::sync::atomic::AtomicUsize>,
14415 max_active: Arc<std::sync::atomic::AtomicUsize>,
14416 }
14417
14418 struct RecoveryTestTool {
14419 id: String,
14420 succeeds: bool,
14421 calls: Arc<std::sync::atomic::AtomicUsize>,
14422 }
14423
14424 struct BlockingApprovalHandler {
14425 entered: Arc<tokio::sync::Barrier>,
14426 release: Arc<tokio::sync::Notify>,
14427 result: ApprovalResult,
14428 }
14429
14430 struct RuntimeWebFetchTransport {
14431 calls: Arc<std::sync::atomic::AtomicUsize>,
14432 }
14433
14434 struct RuntimeWebFetchResolver;
14435
14436 struct ReentrantToolHooks {
14437 agent: parking_lot::Mutex<Option<Weak<RuntimeAgent>>>,
14438 invoked: AtomicBool,
14439 nested_success: AtomicBool,
14440 }
14441
14442 #[async_trait]
14443 impl ai_agents_core::Tool for SlowTool {
14444 fn id(&self) -> &str {
14445 "slow"
14446 }
14447
14448 fn name(&self) -> &str {
14449 "Slow"
14450 }
14451
14452 fn description(&self) -> &str {
14453 "Waits until cancelled or timed out."
14454 }
14455
14456 fn input_schema(&self) -> Value {
14457 serde_json::json!({"type": "object"})
14458 }
14459
14460 async fn execute(
14461 &self,
14462 _args: Value,
14463 _ctx: ai_agents_core::ToolExecutionContext,
14464 ) -> ToolResult {
14465 tokio::time::sleep(std::time::Duration::from_secs(5)).await;
14466 ToolResult::ok("done")
14467 }
14468 }
14469
14470 #[async_trait]
14471 impl ai_agents_core::Tool for FlakyWriteTool {
14472 fn id(&self) -> &str {
14473 "flaky_write"
14474 }
14475
14476 fn name(&self) -> &str {
14477 "Flaky Write"
14478 }
14479
14480 fn description(&self) -> &str {
14481 "Fails on the first write attempt."
14482 }
14483
14484 fn input_schema(&self) -> Value {
14485 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14486 }
14487
14488 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14489 ai_agents_core::ToolPolicyBindings {
14490 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14491 ..Default::default()
14492 }
14493 }
14494
14495 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14496 ai_agents_core::ToolSafetyMetadata {
14497 read_only: false,
14498 concurrency_safe: false,
14499 operation: ai_agents_core::ToolOperationKind::Write,
14500 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14501 requires_network: false,
14502 destructive: false,
14503 open_world: false,
14504 host_dependent: false,
14505 requires_user_interaction: false,
14506 supports_cancellation: true,
14507 default_requires_approval: false,
14508 should_defer_schema: false,
14509 max_output_chars: Some(1024),
14510 max_result_size_chars: Some(1024),
14511 }
14512 }
14513
14514 fn classify_call(&self, _args: &Value) -> ai_agents_core::ToolCallClassification {
14515 let mut classification =
14516 ai_agents_core::ToolCallClassification::from_metadata(&self.safety_metadata());
14517 classification.safely_retryable = false;
14518 classification
14519 }
14520
14521 async fn execute(
14522 &self,
14523 _args: Value,
14524 _ctx: ai_agents_core::ToolExecutionContext,
14525 ) -> ToolResult {
14526 let call = self.calls.fetch_add(1, Ordering::SeqCst);
14527 if call == 0 {
14528 ToolResult::error("first failure")
14529 } else {
14530 ToolResult::ok("second success")
14531 }
14532 }
14533 }
14534
14535 #[async_trait]
14536 impl ai_agents_core::Tool for LockedWriteTool {
14537 fn id(&self) -> &str {
14538 "locked_write"
14539 }
14540
14541 fn name(&self) -> &str {
14542 "Locked Write"
14543 }
14544
14545 fn description(&self) -> &str {
14546 "Tracks concurrent execution on one resource."
14547 }
14548
14549 fn input_schema(&self) -> Value {
14550 serde_json::json!({"type": "object", "properties": {"path": {"type": "string"}}})
14551 }
14552
14553 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14554 ai_agents_core::ToolPolicyBindings {
14555 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14556 ..Default::default()
14557 }
14558 }
14559
14560 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14561 ai_agents_core::ToolSafetyMetadata {
14562 read_only: false,
14563 concurrency_safe: false,
14564 operation: ai_agents_core::ToolOperationKind::Write,
14565 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14566 requires_network: false,
14567 destructive: false,
14568 open_world: false,
14569 host_dependent: false,
14570 requires_user_interaction: false,
14571 supports_cancellation: true,
14572 default_requires_approval: false,
14573 should_defer_schema: false,
14574 max_output_chars: Some(1024),
14575 max_result_size_chars: Some(1024),
14576 }
14577 }
14578
14579 async fn execute(
14580 &self,
14581 _args: Value,
14582 _ctx: ai_agents_core::ToolExecutionContext,
14583 ) -> ToolResult {
14584 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14585 loop {
14586 let current_max = self.max_active.load(Ordering::SeqCst);
14587 if active <= current_max {
14588 break;
14589 }
14590 if self
14591 .max_active
14592 .compare_exchange(current_max, active, Ordering::SeqCst, Ordering::SeqCst)
14593 .is_ok()
14594 {
14595 break;
14596 }
14597 }
14598 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
14599 self.active.fetch_sub(1, Ordering::SeqCst);
14600 ToolResult::ok("done")
14601 }
14602 }
14603
14604 #[async_trait]
14605 impl ai_agents_core::Tool for MultiResourceWriteTool {
14606 fn id(&self) -> &str {
14607 "multi_resource_write"
14608 }
14609
14610 fn name(&self) -> &str {
14611 "Multi Resource Write"
14612 }
14613
14614 fn description(&self) -> &str {
14615 "Tracks concurrent execution across source and destination resources."
14616 }
14617
14618 fn input_schema(&self) -> Value {
14619 serde_json::json!({"type": "object"})
14620 }
14621
14622 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14623 ai_agents_core::ToolPolicyBindings {
14624 path_fields: vec![
14625 ai_agents_core::PathPolicyBinding::read_write("source_path"),
14626 ai_agents_core::PathPolicyBinding::write("destination_path"),
14627 ],
14628 ..Default::default()
14629 }
14630 }
14631
14632 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14633 LockedWriteTool {
14634 active: Arc::clone(&self.active),
14635 max_active: Arc::clone(&self.max_active),
14636 }
14637 .safety_metadata()
14638 }
14639
14640 async fn execute(
14641 &self,
14642 _args: Value,
14643 _ctx: ai_agents_core::ToolExecutionContext,
14644 ) -> ToolResult {
14645 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14646 self.max_active.fetch_max(active, Ordering::SeqCst);
14647 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14648 self.active.fetch_sub(1, Ordering::SeqCst);
14649 ToolResult::ok("done")
14650 }
14651 }
14652
14653 #[async_trait]
14654 impl ai_agents_core::Tool for BlockingPathMutationTool {
14655 fn id(&self) -> &str {
14656 self.id
14657 }
14658
14659 fn name(&self) -> &str {
14660 self.id
14661 }
14662
14663 fn description(&self) -> &str {
14664 "Blocks a path mutation until the test releases it."
14665 }
14666
14667 fn input_schema(&self) -> Value {
14668 serde_json::json!({"type": "object"})
14669 }
14670
14671 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14672 ai_agents_core::ToolPolicyBindings {
14673 path_fields: self.path_fields.clone(),
14674 ..Default::default()
14675 }
14676 }
14677
14678 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14679 ai_agents_core::ToolSafetyMetadata {
14680 read_only: false,
14681 concurrency_safe: false,
14682 operation: ai_agents_core::ToolOperationKind::Write,
14683 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14684 requires_network: false,
14685 destructive: false,
14686 open_world: false,
14687 host_dependent: false,
14688 requires_user_interaction: false,
14689 supports_cancellation: true,
14690 default_requires_approval: false,
14691 should_defer_schema: false,
14692 max_output_chars: Some(1024),
14693 max_result_size_chars: Some(1024),
14694 }
14695 }
14696
14697 async fn execute(
14698 &self,
14699 _args: Value,
14700 _ctx: ai_agents_core::ToolExecutionContext,
14701 ) -> ToolResult {
14702 self.gate.entered.store(true, Ordering::SeqCst);
14703 self.gate.entered_notify.notify_one();
14704 self.gate.release.notified().await;
14705 ToolResult::ok("done")
14706 }
14707 }
14708
14709 #[async_trait]
14710 impl ai_agents_core::Tool for NoBindingWriteTool {
14711 fn id(&self) -> &str {
14712 "no_binding_write"
14713 }
14714
14715 fn name(&self) -> &str {
14716 "No Binding Write"
14717 }
14718
14719 fn description(&self) -> &str {
14720 "Tracks concurrent execution without resource bindings."
14721 }
14722
14723 fn input_schema(&self) -> Value {
14724 serde_json::json!({"type": "object"})
14725 }
14726
14727 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14728 LockedWriteTool {
14729 active: Arc::clone(&self.active),
14730 max_active: Arc::clone(&self.max_active),
14731 }
14732 .safety_metadata()
14733 }
14734
14735 async fn execute(
14736 &self,
14737 _args: Value,
14738 _ctx: ai_agents_core::ToolExecutionContext,
14739 ) -> ToolResult {
14740 let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
14741 self.max_active.fetch_max(active, Ordering::SeqCst);
14742 tokio::time::sleep(std::time::Duration::from_millis(75)).await;
14743 self.active.fetch_sub(1, Ordering::SeqCst);
14744 ToolResult::ok("done")
14745 }
14746 }
14747
14748 #[async_trait]
14749 impl ai_agents_core::Tool for RecoveryTestTool {
14750 fn id(&self) -> &str {
14751 &self.id
14752 }
14753
14754 fn name(&self) -> &str {
14755 &self.id
14756 }
14757
14758 fn description(&self) -> &str {
14759 "Records recovery execution and returns a configured result."
14760 }
14761
14762 fn input_schema(&self) -> Value {
14763 serde_json::json!({"type": "object"})
14764 }
14765
14766 fn policy_bindings(&self) -> ai_agents_core::ToolPolicyBindings {
14767 ai_agents_core::ToolPolicyBindings {
14768 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
14769 ..Default::default()
14770 }
14771 }
14772
14773 fn safety_metadata(&self) -> ai_agents_core::ToolSafetyMetadata {
14774 ai_agents_core::ToolSafetyMetadata {
14775 read_only: false,
14776 concurrency_safe: false,
14777 operation: ai_agents_core::ToolOperationKind::Write,
14778 side_effect_level: ai_agents_core::ToolSideEffectLevel::LocalWrite,
14779 requires_network: false,
14780 destructive: false,
14781 open_world: false,
14782 host_dependent: false,
14783 requires_user_interaction: false,
14784 supports_cancellation: true,
14785 default_requires_approval: false,
14786 should_defer_schema: false,
14787 max_output_chars: Some(1024),
14788 max_result_size_chars: Some(1024),
14789 }
14790 }
14791
14792 async fn execute(
14793 &self,
14794 _args: Value,
14795 _ctx: ai_agents_core::ToolExecutionContext,
14796 ) -> ToolResult {
14797 self.calls.fetch_add(1, Ordering::SeqCst);
14798 if self.succeeds {
14799 ToolResult::ok(format!("{} succeeded", self.id))
14800 } else {
14801 ToolResult::error(format!("{} failed", self.id))
14802 }
14803 }
14804 }
14805
14806 #[async_trait]
14807 impl WebFetchTransport for RuntimeWebFetchTransport {
14808 async fn send(
14810 &self,
14811 _request: WebFetchTransportRequest,
14812 ) -> std::result::Result<WebFetchTransportResponse, String> {
14813 Err("validated addresses are required".to_string())
14814 }
14815
14816 async fn send_validated(
14818 &self,
14819 _request: WebFetchTransportRequest,
14820 _addresses: &[std::net::SocketAddr],
14821 ) -> std::result::Result<WebFetchTransportResponse, String> {
14822 self.calls.fetch_add(1, Ordering::SeqCst);
14823 Ok(WebFetchTransportResponse {
14824 status: 200,
14825 content_type: Some("text/plain".to_string()),
14826 location: None,
14827 body: b"approved".to_vec(),
14828 })
14829 }
14830 }
14831
14832 #[async_trait]
14833 impl WebFetchResolver for RuntimeWebFetchResolver {
14834 async fn resolve(
14836 &self,
14837 _host: &str,
14838 _port: u16,
14839 ) -> std::result::Result<Vec<std::net::IpAddr>, String> {
14840 Ok(vec![std::net::IpAddr::V4(std::net::Ipv4Addr::new(
14841 93, 184, 216, 34,
14842 ))])
14843 }
14844 }
14845
14846 #[async_trait]
14847 impl ApprovalHandler for BlockingApprovalHandler {
14848 async fn request_approval(
14849 &self,
14850 _request: ai_agents_hitl::ApprovalRequest,
14851 ) -> ApprovalResult {
14852 self.entered.wait().await;
14853 self.release.notified().await;
14854 self.result.clone()
14855 }
14856 }
14857
14858 #[async_trait]
14859 impl AgentHooks for ReentrantToolHooks {
14860 async fn on_tool_complete(&self, tool: &str, _result: &ToolResult, _duration_ms: u64) {
14861 if tool != "reentrant_write" || self.invoked.swap(true, Ordering::SeqCst) {
14862 return;
14863 }
14864 let agent = self.agent.lock().as_ref().and_then(Weak::upgrade);
14865 if let Some(agent) = agent {
14866 let result = agent
14867 .invoke_tool(ToolExecutionRequest::new(
14868 "nested-hook-call",
14869 "reentrant_write",
14870 serde_json::json!({"path": "./hook.txt"}),
14871 ToolCallSource::Manual,
14872 ))
14873 .await;
14874 self.nested_success
14875 .store(result.is_ok_and(|record| record.success), Ordering::SeqCst);
14876 }
14877 }
14878 }
14879
14880 #[async_trait]
14881 impl AgentHooks for ResponseCountingHooks {
14882 async fn on_response(&self, _response: &AgentResponse) {
14883 self.responses.fetch_add(1, Ordering::SeqCst);
14884 }
14885 }
14886
14887 struct ApprovalRecordingHooks {
14888 events: parking_lot::Mutex<Vec<String>>,
14889 }
14890
14891 impl ApprovalRecordingHooks {
14892 fn new() -> Self {
14893 Self {
14894 events: parking_lot::Mutex::new(Vec::new()),
14895 }
14896 }
14897
14898 fn events(&self) -> Vec<String> {
14899 self.events.lock().clone()
14900 }
14901 }
14902
14903 #[async_trait]
14904 impl AgentHooks for ApprovalRecordingHooks {
14905 async fn on_approval_result(&self, request_id: &str, result: &ApprovalResult) {
14906 self.events.lock().push(format!(
14907 "raw:{}:{}",
14908 request_id,
14909 approval_result_name(result)
14910 ));
14911 }
14912
14913 async fn on_approval_resolved(
14914 &self,
14915 request: &ai_agents_hitl::ApprovalRequest,
14916 raw_result: &ApprovalResult,
14917 outcome: &ApprovalResolvedOutcome,
14918 ) {
14919 self.events.lock().push(format!(
14920 "resolved:{}:{}:{}",
14921 request.id,
14922 approval_result_name(raw_result),
14923 approval_outcome_name(outcome)
14924 ));
14925 }
14926 }
14927
14928 fn approval_result_name(result: &ApprovalResult) -> &'static str {
14929 match result {
14930 ApprovalResult::Approved => "approved",
14931 ApprovalResult::Rejected { .. } => "rejected",
14932 ApprovalResult::Modified { .. } => "modified",
14933 ApprovalResult::Timeout => "timeout",
14934 }
14935 }
14936
14937 fn approval_outcome_name(outcome: &ApprovalResolvedOutcome) -> &'static str {
14938 match outcome {
14939 ApprovalResolvedOutcome::Approved => "approved",
14940 ApprovalResolvedOutcome::Rejected { .. } => "rejected",
14941 ApprovalResolvedOutcome::Modified { .. } => "modified",
14942 ApprovalResolvedOutcome::Error { .. } => "error",
14943 }
14944 }
14945
14946 fn assert_correlated_approval_events(
14947 events: &[String],
14948 raw_status: &str,
14949 outcome_status: &str,
14950 ) {
14951 assert_eq!(events.len(), 2);
14952 let raw: Vec<_> = events[0].split(':').collect();
14953 let resolved: Vec<_> = events[1].split(':').collect();
14954 assert_eq!(raw[0], "raw");
14955 assert_eq!(resolved[0], "resolved");
14956 assert_eq!(raw[1], resolved[1]);
14957 assert_eq!(raw[2], raw_status);
14958 assert_eq!(resolved[2], raw_status);
14959 assert_eq!(resolved[3], outcome_status);
14960 }
14961
14962 fn approval_security_config(policy_enabled: bool) -> ToolSecurityConfig {
14963 let mut security = ToolSecurityConfig {
14964 enabled: true,
14965 fail_closed: true,
14966 ..Default::default()
14967 };
14968 let policy = ai_agents_tools::ToolPolicyConfig {
14969 enabled: policy_enabled,
14970 write_paths: vec![".".to_string()],
14971 require_confirmation: true,
14972 ..Default::default()
14973 };
14974 security.tools.insert("locked_write".to_string(), policy);
14975 security
14976 }
14977
14978 struct MutationTestWorkspace {
14979 root: std::path::PathBuf,
14980 }
14981
14982 impl MutationTestWorkspace {
14983 fn new() -> Self {
14984 let root = std::env::temp_dir().join(format!(
14985 "ai-agents-runtime-mutation-{}",
14986 uuid::Uuid::new_v4()
14987 ));
14988 std::fs::create_dir_all(&root).unwrap();
14989 Self { root }
14990 }
14991 }
14992
14993 impl Drop for MutationTestWorkspace {
14994 fn drop(&mut self) {
14995 let _ = std::fs::remove_dir_all(&self.root);
14996 }
14997 }
14998
14999 async fn wait_for_resource_lock_strong_count(locks: &ToolResourceLocks, minimum: usize) {
15000 tokio::time::timeout(std::time::Duration::from_secs(2), async {
15001 loop {
15002 let strong_count = locks
15003 .read()
15004 .get("path-mutation:global")
15005 .map_or(0, |lock| lock.strong_count());
15006 if strong_count >= minimum {
15007 break;
15008 }
15009 tokio::task::yield_now().await;
15010 }
15011 })
15012 .await
15013 .expect("path mutation call did not reach the shared lock");
15014 }
15015
15016 async fn assert_path_mutation_pair_serialized(
15017 first_id: &'static str,
15018 first_fields: Vec<ai_agents_core::PathPolicyBinding>,
15019 first_args: Value,
15020 second_id: &'static str,
15021 second_fields: Vec<ai_agents_core::PathPolicyBinding>,
15022 second_args: Value,
15023 ) {
15024 let locks = new_tool_resource_locks();
15025 let first_gate = PathMutationGate::new();
15026 let second_gate = PathMutationGate::new();
15027 second_gate.release();
15028 let agent = Arc::new(
15029 AgentBuilder::new()
15030 .system_prompt("Test global path mutation locking.")
15031 .llm(Arc::new(mock_with_response("done")))
15032 .tool(Arc::new(BlockingPathMutationTool {
15033 id: first_id,
15034 path_fields: first_fields,
15035 gate: first_gate.clone(),
15036 }))
15037 .tool(Arc::new(BlockingPathMutationTool {
15038 id: second_id,
15039 path_fields: second_fields,
15040 gate: second_gate.clone(),
15041 }))
15042 .build()
15043 .unwrap()
15044 .with_shared_resource_locks(Arc::clone(&locks)),
15045 );
15046
15047 let first = {
15048 let agent = Arc::clone(&agent);
15049 tokio::spawn(async move {
15050 agent
15051 .invoke_tool(ToolExecutionRequest::new(
15052 format!("{}-first", first_id),
15053 first_id,
15054 first_args,
15055 ToolCallSource::Manual,
15056 ))
15057 .await
15058 .unwrap()
15059 })
15060 };
15061 first_gate.wait_until_entered().await;
15062
15063 let second = {
15064 let agent = Arc::clone(&agent);
15065 tokio::spawn(async move {
15066 agent
15067 .invoke_tool(ToolExecutionRequest::new(
15068 format!("{}-second", second_id),
15069 second_id,
15070 second_args,
15071 ToolCallSource::Manual,
15072 ))
15073 .await
15074 .unwrap()
15075 })
15076 };
15077 wait_for_resource_lock_strong_count(&locks, 2).await;
15078 assert!(!second_gate.entered.load(Ordering::SeqCst));
15079 assert!(!second.is_finished());
15080
15081 first_gate.release();
15082 let (first, second) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15083 tokio::join!(first, second)
15084 })
15085 .await
15086 .expect("serialized path mutation calls did not finish");
15087 assert!(first.unwrap().success);
15088 assert!(second.unwrap().success);
15089 assert!(second_gate.entered.load(Ordering::SeqCst));
15090 assert!(locks.read().is_empty());
15091 }
15092
15093 #[derive(Clone, Copy)]
15094 enum MutationDenial {
15095 Policy,
15096 Approval,
15097 }
15098
15099 fn mutation_denial_security_config(
15100 tool_id: &str,
15101 workspace: &std::path::Path,
15102 denial: MutationDenial,
15103 ) -> ToolSecurityConfig {
15104 let workspace = workspace.to_string_lossy().into_owned();
15105 let mut policy = ai_agents_tools::ToolPolicyConfig {
15106 read_paths: vec![workspace.clone()],
15107 write_paths: vec![workspace.clone()],
15108 ..Default::default()
15109 };
15110 match denial {
15111 MutationDenial::Policy => policy.blocked_paths = vec![workspace],
15112 MutationDenial::Approval => policy.require_confirmation = true,
15113 }
15114
15115 let mut security = ToolSecurityConfig {
15116 enabled: true,
15117 fail_closed: true,
15118 ..Default::default()
15119 };
15120 security.tools.insert(tool_id.to_string(), policy);
15121 security
15122 }
15123
15124 async fn assert_path_mutation_denied(tool: Arc<dyn Tool>, denial: MutationDenial) {
15125 let workspace = MutationTestWorkspace::new();
15126 let tool_id = tool.id().to_string();
15127 let preserved = workspace.root.join(format!("{}-preserved.txt", tool_id));
15128 let destination = workspace.root.join(format!("{}-destination.txt", tool_id));
15129 std::fs::write(&preserved, "preserved").unwrap();
15130 let arguments = match tool_id.as_str() {
15131 "copy_path" | "move_path" => serde_json::json!({
15132 "source_path": preserved.to_string_lossy(),
15133 "destination_path": destination.to_string_lossy(),
15134 "dry_run": false
15135 }),
15136 "delete_path" => serde_json::json!({
15137 "path": preserved.to_string_lossy(),
15138 "recursive": false,
15139 "dry_run": false
15140 }),
15141 _ => panic!("unsupported mutation tool: {}", tool_id),
15142 };
15143 let security = mutation_denial_security_config(&tool_id, &workspace.root, denial);
15144 let builder = AgentBuilder::new()
15145 .system_prompt("Test mutation denial.")
15146 .llm(Arc::new(mock_with_response("done")))
15147 .tool(tool)
15148 .tool_security(ToolSecurityEngine::new(security));
15149 let builder = match denial {
15150 MutationDenial::Policy => builder,
15151 MutationDenial::Approval => builder
15152 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15153 .approval_handler(Arc::new(RejectAllHandler::new())),
15154 };
15155 let agent = builder.build().unwrap();
15156
15157 let record = agent
15158 .invoke_tool(ToolExecutionRequest::new(
15159 format!("{}-denied", tool_id),
15160 tool_id.clone(),
15161 arguments,
15162 ToolCallSource::Manual,
15163 ))
15164 .await
15165 .unwrap();
15166
15167 assert!(!record.executed, "{} must not be invoked", tool_id);
15168 assert!(!record.success);
15169 match denial {
15170 MutationDenial::Policy => {
15171 assert_eq!(record.policy.outcome, PermissionOutcome::Deny);
15172 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15173 &approval.status,
15174 ToolApprovalStatus::NotRequired
15175 )));
15176 }
15177 MutationDenial::Approval => {
15178 assert_eq!(record.policy.outcome, PermissionOutcome::RequiresApproval);
15179 assert!(record.approval.as_ref().is_some_and(|approval| matches!(
15180 &approval.status,
15181 ToolApprovalStatus::Rejected
15182 )));
15183 }
15184 }
15185 assert_eq!(std::fs::read_to_string(&preserved).unwrap(), "preserved");
15186 assert!(!destination.exists());
15187 }
15188
15189 fn recovery_manager_with_fallbacks(
15190 fallbacks: impl IntoIterator<Item = (String, String)>,
15191 ) -> RecoveryManager {
15192 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15193
15194 let per_tool = fallbacks
15195 .into_iter()
15196 .map(|(tool, fallback_tool)| {
15197 (
15198 tool,
15199 ToolRetryConfig {
15200 max_retries: 0,
15201 timeout_ms: Some(1_000),
15202 on_failure: ToolFailureAction::Fallback { fallback_tool },
15203 },
15204 )
15205 })
15206 .collect();
15207 RecoveryManager::new(ErrorRecoveryConfig {
15208 tools: ToolRecoveryConfig {
15209 per_tool,
15210 ..Default::default()
15211 },
15212 ..Default::default()
15213 })
15214 }
15215
15216 fn approval_check() -> HITLCheckResult {
15217 HITLCheckResult::required(
15218 ApprovalTrigger::tool("test", serde_json::json!({})),
15219 HashMap::new(),
15220 "Approve?",
15221 None,
15222 )
15223 }
15224
15225 fn agent_with_approval_result(
15226 raw_result: ApprovalResult,
15227 timeout_action: TimeoutAction,
15228 hooks: Arc<ApprovalRecordingHooks>,
15229 ) -> RuntimeAgent {
15230 use ai_agents_hitl::{CallbackHandler, HITLConfig};
15231
15232 let config = HITLConfig {
15233 on_timeout: timeout_action,
15234 ..Default::default()
15235 };
15236 let handler = CallbackHandler::new(move |_| raw_result.clone());
15237 AgentBuilder::new()
15238 .system_prompt("Test HITL hooks.")
15239 .llm(Arc::new(mock_with_response("done")))
15240 .build()
15241 .unwrap()
15242 .with_hooks(hooks)
15243 .with_hitl(HITLEngine::new(config), Arc::new(handler))
15244 }
15245
15246 #[tokio::test]
15247 async fn approval_hooks_expose_direct_effective_decisions_after_raw_results() {
15248 let cases = vec![
15249 (ApprovalResult::Approved, "approved"),
15250 (
15251 ApprovalResult::Rejected {
15252 reason: Some("denied".to_string()),
15253 },
15254 "rejected",
15255 ),
15256 (
15257 ApprovalResult::Modified {
15258 changes: HashMap::from([("value".to_string(), serde_json::json!(2))]),
15259 },
15260 "modified",
15261 ),
15262 ];
15263
15264 for (raw_result, expected) in cases {
15265 let hooks = Arc::new(ApprovalRecordingHooks::new());
15266 let agent =
15267 agent_with_approval_result(raw_result, TimeoutAction::Reject, hooks.clone());
15268
15269 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15270
15271 assert_eq!(approval_result_name(&result), expected);
15272 assert_correlated_approval_events(&hooks.events(), expected, expected);
15273 }
15274 }
15275
15276 #[tokio::test]
15277 async fn approval_hooks_expose_timeout_policy_decisions() {
15278 for (timeout_action, expected) in [
15279 (TimeoutAction::Approve, "approved"),
15280 (TimeoutAction::Reject, "rejected"),
15281 ] {
15282 let hooks = Arc::new(ApprovalRecordingHooks::new());
15283 let agent =
15284 agent_with_approval_result(ApprovalResult::Timeout, timeout_action, hooks.clone());
15285
15286 let result = agent.request_hitl_approval(approval_check()).await.unwrap();
15287
15288 assert_eq!(approval_result_name(&result), expected);
15289 assert_correlated_approval_events(&hooks.events(), "timeout", expected);
15290 }
15291 }
15292
15293 #[tokio::test]
15294 async fn timeout_error_fires_correlated_resolved_error_before_returning() {
15295 let hooks = Arc::new(ApprovalRecordingHooks::new());
15296 let agent = agent_with_approval_result(
15297 ApprovalResult::Timeout,
15298 TimeoutAction::Error,
15299 hooks.clone(),
15300 );
15301
15302 let error = agent
15303 .request_hitl_approval(approval_check())
15304 .await
15305 .unwrap_err();
15306
15307 assert!(error.to_string().contains("HITL approval timeout"));
15308 assert_correlated_approval_events(&hooks.events(), "timeout", "error");
15309 }
15310
15311 #[tokio::test]
15313 async fn test_integration_yaml_to_chat_basic() {
15314 let mock = mock_with_response("Hello! How can I help you?");
15315 let agent = AgentBuilder::new()
15316 .system_prompt("You are a test assistant.")
15317 .llm(Arc::new(mock))
15318 .build()
15319 .unwrap();
15320
15321 let response = agent.chat("Hi").await.unwrap();
15322 assert!(!response.content.is_empty());
15323 assert_eq!(response.content, "Hello! How can I help you?");
15324 }
15325
15326 #[tokio::test]
15328 async fn test_integration_multi_turn_conversation() {
15329 let mock = mock_with_responses(vec![
15330 "Hello! I'm your assistant.",
15331 "The weather is sunny today.",
15332 "Goodbye!",
15333 ]);
15334 let agent = AgentBuilder::new()
15335 .system_prompt("You are helpful.")
15336 .llm(Arc::new(mock))
15337 .build()
15338 .unwrap();
15339
15340 let r1 = agent.chat("Hi").await.unwrap();
15341 assert_eq!(r1.content, "Hello! I'm your assistant.");
15342
15343 let r2 = agent.chat("What's the weather?").await.unwrap();
15344 assert_eq!(r2.content, "The weather is sunny today.");
15345
15346 let r3 = agent.chat("Bye").await.unwrap();
15347 assert_eq!(r3.content, "Goodbye!");
15348
15349 let messages = agent.memory.get_messages(None).await.unwrap();
15351 assert_eq!(messages.len(), 6);
15353 }
15354
15355 #[test]
15356 fn later_approval_preserves_modified_evidence() {
15357 let arguments = serde_json::json!({"dry_run": true});
15358 let mut record = Some(ToolApprovalRecord {
15359 status: ToolApprovalStatus::Modified,
15360 reason: None,
15361 modified_arguments: Some(arguments.clone()),
15362 });
15363
15364 merge_approved_record(&mut record);
15365
15366 let record = record.unwrap();
15367 assert!(matches!(record.status, ToolApprovalStatus::Modified));
15368 assert_eq!(record.modified_arguments, Some(arguments));
15369 }
15370
15371 #[test]
15372 fn approval_binding_rejects_replaced_tool_implementation() {
15373 let reviewed_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15374 let same_tool = Arc::clone(&reviewed_tool);
15375 let replacement_tool: Arc<dyn ai_agents_core::Tool> = Arc::new(ContextEchoTool);
15376 let arguments = serde_json::json!({"path": "."});
15377 let versions = ToolDecisionVersions {
15378 policy: 2,
15379 registry: 3,
15380 runtime_control: 4,
15381 state: Some(5),
15382 };
15383 let binding = ToolApprovalBinding {
15384 canonical_id: "context_echo".to_string(),
15385 arguments: arguments.clone(),
15386 confirmation_required: true,
15387 policy_version: versions.policy,
15388 runtime_control_version: versions.runtime_control,
15389 state_generation: versions.state,
15390 reviewed_tool,
15391 };
15392
15393 assert!(!binding.is_stale("context_echo", &arguments, true, versions, &same_tool,));
15394 assert!(binding.is_stale(
15395 "context_echo",
15396 &arguments,
15397 true,
15398 versions,
15399 &replacement_tool,
15400 ));
15401 }
15402
15403 #[tokio::test]
15404 async fn approved_mutation_to_dry_run_remains_executable() {
15405 use ai_agents_hitl::CallbackHandler;
15406
15407 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
15408 changes: HashMap::from([("dry_run".to_string(), serde_json::json!(true))]),
15409 });
15410 let agent = AgentBuilder::new()
15411 .system_prompt("Test safer approval modifications.")
15412 .llm(Arc::new(mock_with_response("done")))
15413 .tool(Arc::new(ai_agents_tools::FileWriteTool::new()))
15414 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
15415 .approval_handler(Arc::new(handler))
15416 .build()
15417 .unwrap();
15418
15419 let record = agent
15420 .invoke_tool(ToolExecutionRequest::new(
15421 "approved-dry-run",
15422 "file_write",
15423 serde_json::json!({
15424 "path": "./approval-dry-run.txt",
15425 "content": "not written"
15426 }),
15427 ToolCallSource::Manual,
15428 ))
15429 .await
15430 .unwrap();
15431
15432 assert!(record.executed);
15433 assert!(record.success);
15434 assert_eq!(record.executed_arguments["dry_run"], true);
15435 assert!(matches!(
15436 record.approval.as_ref().map(|approval| &approval.status),
15437 Some(ToolApprovalStatus::Modified)
15438 ));
15439 let output: Value = serde_json::from_str(&record.output).unwrap();
15440 assert_eq!(output["mutation_performed"], false);
15441 }
15442
15443 #[tokio::test]
15445 async fn shared_executor_approval_reaches_web_fetch_transport() {
15446 use ai_agents_hitl::{CallbackHandler, HITLConfig};
15447 use ai_agents_tools::{DomainPolicyConfig, ToolPolicyConfig};
15448
15449 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15450 let tool = WebFetchTool::with_transport_and_resolver(
15451 Arc::new(RuntimeWebFetchTransport {
15452 calls: Arc::clone(&calls),
15453 }),
15454 Arc::new(RuntimeWebFetchResolver),
15455 );
15456 let mut security = ToolSecurityConfig {
15457 enabled: true,
15458 fail_closed: true,
15459 ..Default::default()
15460 };
15461 security.tools.insert(
15462 "web_fetch".to_string(),
15463 ToolPolicyConfig {
15464 domains: DomainPolicyConfig {
15465 requires_approval: vec!["approval.test".to_string()],
15466 ..Default::default()
15467 },
15468 allowed_schemes: vec!["https".to_string()],
15469 allowed_ports: vec![443],
15470 ..Default::default()
15471 },
15472 );
15473 let handler = CallbackHandler::new(|_| ApprovalResult::Approved);
15474 let agent = AgentBuilder::new()
15475 .system_prompt("Test approved web fetch execution.")
15476 .llm(Arc::new(mock_with_response("done")))
15477 .tool(Arc::new(tool))
15478 .tool_security(ToolSecurityEngine::new(security))
15479 .build()
15480 .unwrap()
15481 .with_hitl(HITLEngine::new(HITLConfig::default()), Arc::new(handler));
15482
15483 let record = agent
15484 .invoke_tool(ToolExecutionRequest::new(
15485 "approved-web-fetch",
15486 "web_fetch",
15487 serde_json::json!({
15488 "url": "https://approval.test/page",
15489 "cache_ttl_seconds": 0
15490 }),
15491 ToolCallSource::Manual,
15492 ))
15493 .await
15494 .unwrap();
15495
15496 assert!(record.success);
15497 assert!(
15498 record
15499 .approval
15500 .as_ref()
15501 .is_some_and(|approval| matches!(approval.status, ToolApprovalStatus::Approved))
15502 );
15503 assert_eq!(calls.load(Ordering::SeqCst), 1);
15504 }
15505
15506 #[tokio::test]
15507 async fn context_preserves_requested_and_canonical_identity() {
15508 let mock = mock_with_response("hello");
15509 let mut tools = ai_agents_tools::ToolRegistry::new();
15510 tools.register(Arc::new(ContextEchoTool)).unwrap();
15511
15512 let mut security = ToolSecurityConfig {
15513 enabled: true,
15514 fail_closed: true,
15515 ..Default::default()
15516 };
15517 let mut policy = ai_agents_tools::ToolPolicyConfig {
15518 read_paths: vec![".".to_string()],
15519 max_results: Some(7),
15520 ..Default::default()
15521 };
15522 policy
15523 .config
15524 .insert("backend".to_string(), serde_json::json!("memory"));
15525 security.tools.insert("context_echo".to_string(), policy);
15526
15527 let agent = AgentBuilder::new()
15528 .system_prompt("You are helpful.")
15529 .llm(Arc::new(mock))
15530 .tools(tools)
15531 .tool_security(ToolSecurityEngine::new(security))
15532 .build()
15533 .unwrap();
15534
15535 let record = agent
15536 .invoke_tool(ToolExecutionRequest::new(
15537 "ctx-call",
15538 "Context Echo",
15539 serde_json::json!({"path": ".", "max_results": 99}),
15540 ToolCallSource::Manual,
15541 ))
15542 .await
15543 .unwrap();
15544
15545 assert!(record.success);
15546 assert!(matches!(&record.source, ToolCallSource::Manual));
15547 assert_eq!(record.requested_name, "Context Echo");
15548 assert_eq!(record.canonical_id, "context_echo");
15549 assert_eq!(record.policy.outcome, PermissionOutcome::Allow);
15550 assert_eq!(record.executed_arguments["max_results"], 7);
15551 let output: Value = serde_json::from_str(&record.output).unwrap();
15552 assert_eq!(output["requested_name"], "Context Echo");
15553 assert_eq!(output["canonical_id"], "context_echo");
15554 assert_eq!(output["max_results"], 7);
15555 assert_eq!(output["custom_config"]["backend"], "memory");
15556 assert!(record.metadata.contains_key("effective_limits"));
15557 assert!(record.metadata.contains_key("policy_snapshot"));
15558 }
15559
15560 #[tokio::test]
15561 async fn test_runtime_control_cancels_active_tool_call() {
15562 let mock = mock_with_response("hello");
15563 let agent = Arc::new(
15564 AgentBuilder::new()
15565 .system_prompt("You are helpful.")
15566 .llm(Arc::new(mock))
15567 .tool(Arc::new(SlowTool))
15568 .build()
15569 .unwrap(),
15570 );
15571 let control = agent.runtime_control();
15572 let running_agent = Arc::clone(&agent);
15573 let handle = tokio::spawn(async move {
15574 running_agent
15575 .invoke_tool(ToolExecutionRequest::new(
15576 "slow-call",
15577 "slow",
15578 serde_json::json!({}),
15579 ToolCallSource::Manual,
15580 ))
15581 .await
15582 .unwrap()
15583 });
15584
15585 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
15586 control.cancel_all();
15587 let record = handle.await.unwrap();
15588
15589 assert!(record.executed);
15590 assert!(record.cancelled);
15591 assert!(!record.success);
15592 assert!(record.cancellation_reason.is_some());
15593 }
15594
15595 #[tokio::test]
15596 async fn non_idempotent_tool_calls_are_not_retried() {
15597 use ai_agents_recovery::{ErrorRecoveryConfig, ToolRecoveryConfig, ToolRetryConfig};
15598
15599 let mock = mock_with_response("hello");
15600 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15601 let agent = AgentBuilder::new()
15602 .system_prompt("You are helpful.")
15603 .llm(Arc::new(mock))
15604 .tool(Arc::new(FlakyWriteTool {
15605 calls: Arc::clone(&calls),
15606 }))
15607 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
15608 tools: ToolRecoveryConfig {
15609 default: ToolRetryConfig {
15610 max_retries: 2,
15611 ..Default::default()
15612 },
15613 ..Default::default()
15614 },
15615 ..Default::default()
15616 }))
15617 .build()
15618 .unwrap();
15619
15620 let record = agent
15621 .invoke_tool(ToolExecutionRequest::new(
15622 "flaky-call",
15623 "flaky_write",
15624 serde_json::json!({"path": "./tmp.txt"}),
15625 ToolCallSource::Manual,
15626 ))
15627 .await
15628 .unwrap();
15629
15630 assert!(!record.success);
15631 assert_eq!(calls.load(Ordering::SeqCst), 1);
15632 }
15633
15634 #[tokio::test]
15635 async fn side_effecting_tools_are_serialized_per_resource() {
15636 let mock = mock_with_response("hello");
15637 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15638 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15639 let agent = Arc::new(
15640 AgentBuilder::new()
15641 .system_prompt("You are helpful.")
15642 .llm(Arc::new(mock))
15643 .tool(Arc::new(LockedWriteTool {
15644 active: Arc::clone(&active),
15645 max_active: Arc::clone(&max_active),
15646 }))
15647 .build()
15648 .unwrap(),
15649 );
15650
15651 let left = {
15652 let agent = Arc::clone(&agent);
15653 tokio::spawn(async move {
15654 agent
15655 .invoke_tool(ToolExecutionRequest::new(
15656 "lock-1",
15657 "locked_write",
15658 serde_json::json!({"path": "./same.txt"}),
15659 ToolCallSource::Manual,
15660 ))
15661 .await
15662 .unwrap()
15663 })
15664 };
15665 let right = {
15666 let agent = Arc::clone(&agent);
15667 tokio::spawn(async move {
15668 agent
15669 .invoke_tool(ToolExecutionRequest::new(
15670 "lock-2",
15671 "locked_write",
15672 serde_json::json!({"path": "./same.txt"}),
15673 ToolCallSource::Manual,
15674 ))
15675 .await
15676 .unwrap()
15677 })
15678 };
15679
15680 let left = left.await.unwrap();
15681 let right = right.await.unwrap();
15682 assert!(left.success);
15683 assert!(right.success);
15684 assert_eq!(max_active.load(Ordering::SeqCst), 1);
15685 }
15686
15687 #[tokio::test]
15688 async fn path_resources_use_shared_global_lock_and_cleanup() {
15689 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15690 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
15691 let bindings = ai_agents_core::ToolPolicyBindings {
15692 path_fields: vec![
15693 ai_agents_core::PathPolicyBinding::read_write("source_path"),
15694 ai_agents_core::PathPolicyBinding::write("destination_path"),
15695 ],
15696 ..Default::default()
15697 };
15698 let classification = ai_agents_core::ToolCallClassification::from_metadata(
15699 &MultiResourceWriteTool {
15700 active: Arc::clone(&active),
15701 max_active: Arc::clone(&max_active),
15702 }
15703 .safety_metadata(),
15704 );
15705 let left_args = serde_json::json!({
15706 "source_path": "./a/../first.txt",
15707 "destination_path": "./second.txt"
15708 });
15709 let right_args = serde_json::json!({
15710 "source_path": "./second.txt",
15711 "destination_path": "./first.txt"
15712 });
15713 let left_keys = tool_resource_lock_keys(
15714 "multi_resource_write",
15715 &left_args,
15716 &bindings,
15717 &classification,
15718 );
15719 let right_keys = tool_resource_lock_keys(
15720 "multi_resource_write",
15721 &right_args,
15722 &bindings,
15723 &classification,
15724 );
15725 assert_eq!(left_keys, right_keys);
15726 assert_eq!(left_keys, vec!["path-mutation:global".to_string()]);
15727
15728 let locks = new_tool_resource_locks();
15729 let build_agent = || {
15730 AgentBuilder::new()
15731 .system_prompt("Test shared resource locks.")
15732 .llm(Arc::new(mock_with_response("done")))
15733 .tool(Arc::new(MultiResourceWriteTool {
15734 active: Arc::clone(&active),
15735 max_active: Arc::clone(&max_active),
15736 }))
15737 .build()
15738 .unwrap()
15739 .with_shared_resource_locks(Arc::clone(&locks))
15740 };
15741 let left_agent = Arc::new(build_agent());
15742 let right_agent = Arc::new(build_agent());
15743 let left = tokio::spawn(async move {
15744 left_agent
15745 .invoke_tool(ToolExecutionRequest::new(
15746 "multi-left",
15747 "multi_resource_write",
15748 left_args,
15749 ToolCallSource::Manual,
15750 ))
15751 .await
15752 .unwrap()
15753 });
15754 let right = tokio::spawn(async move {
15755 right_agent
15756 .invoke_tool(ToolExecutionRequest::new(
15757 "multi-right",
15758 "multi_resource_write",
15759 right_args,
15760 ToolCallSource::Manual,
15761 ))
15762 .await
15763 .unwrap()
15764 });
15765 let (left, right) = tokio::time::timeout(std::time::Duration::from_secs(2), async {
15766 tokio::join!(left, right)
15767 })
15768 .await
15769 .expect("reversed resource acquisition must not deadlock");
15770
15771 assert!(left.unwrap().success);
15772 assert!(right.unwrap().success);
15773 assert_eq!(max_active.load(Ordering::SeqCst), 1);
15774 assert!(locks.read().is_empty());
15775 }
15776
15777 #[tokio::test]
15778 async fn global_path_lock_serializes_copy_destination_with_file_write() {
15779 assert_path_mutation_pair_serialized(
15780 "copy_path",
15781 CopyPathTool::new().policy_bindings().path_fields,
15782 serde_json::json!({
15783 "source_path": "./source.txt",
15784 "destination_path": "./shared.txt"
15785 }),
15786 "file_write",
15787 FileWriteTool::new().policy_bindings().path_fields,
15788 serde_json::json!({"path": "./shared.txt"}),
15789 )
15790 .await;
15791 }
15792
15793 #[tokio::test]
15794 async fn parent_and_spawned_runtime_share_global_path_lock() {
15795 let workspace = MutationTestWorkspace::new();
15796 let destination = workspace.root.join("spawned.txt");
15797 let parent_gate = PathMutationGate::new();
15798 let parent = Arc::new(
15799 AgentBuilder::from_yaml(
15800 r#"
15801name: LockParent
15802system_prompt: parent
15803llm:
15804 default: default
15805tools:
15806 - parent_path_write
15807spawner:
15808 shared_llms: true
15809"#,
15810 )
15811 .unwrap()
15812 .llm(Arc::new(mock_with_response("done")))
15813 .auto_configure_spawner()
15814 .await
15815 .unwrap()
15816 .tool(Arc::new(BlockingPathMutationTool {
15817 id: "parent_path_write",
15818 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15819 gate: parent_gate.clone(),
15820 }))
15821 .build()
15822 .unwrap(),
15823 );
15824
15825 let mut child_spec = crate::spec::AgentSpec {
15826 name: "LockChild".to_string(),
15827 system_prompt: "child".to_string(),
15828 tools: Some(vec![crate::spec::ToolEntry::Simple(
15829 "file_write".to_string(),
15830 )]),
15831 ..Default::default()
15832 };
15833 child_spec.tool_security.enabled = true;
15834 child_spec.tool_security.fail_closed = true;
15835 let file_write_policy = ai_agents_tools::ToolPolicyConfig {
15836 write_paths: vec![workspace.root.to_string_lossy().into_owned()],
15837 allow_without_confirmation: true,
15838 ..Default::default()
15839 };
15840 child_spec
15841 .tool_security
15842 .tools
15843 .insert("file_write".to_string(), file_write_policy);
15844 let spawned = parent
15845 .spawner()
15846 .unwrap()
15847 .spawn_from_spec(child_spec)
15848 .await
15849 .unwrap();
15850 assert!(Arc::ptr_eq(
15851 &parent.resource_locks,
15852 &spawned.agent.resource_locks
15853 ));
15854 assert!(!Arc::ptr_eq(
15855 &parent.runtime_control,
15856 &spawned.agent.runtime_control
15857 ));
15858
15859 let parent_call = {
15860 let parent = Arc::clone(&parent);
15861 let destination = destination.clone();
15862 tokio::spawn(async move {
15863 parent
15864 .invoke_tool(ToolExecutionRequest::new(
15865 "parent-lock-holder",
15866 "parent_path_write",
15867 serde_json::json!({"path": destination}),
15868 ToolCallSource::Manual,
15869 ))
15870 .await
15871 .unwrap()
15872 })
15873 };
15874 parent_gate.wait_until_entered().await;
15875
15876 let child_call = {
15877 let child = Arc::clone(&spawned.agent);
15878 let destination = destination.clone();
15879 tokio::spawn(async move {
15880 child
15881 .invoke_tool(ToolExecutionRequest::new(
15882 "spawned-file-write",
15883 "file_write",
15884 serde_json::json!({
15885 "path": destination,
15886 "content": "spawned",
15887 "dry_run": false
15888 }),
15889 ToolCallSource::Manual,
15890 ))
15891 .await
15892 .unwrap()
15893 })
15894 };
15895 wait_for_resource_lock_strong_count(&parent.resource_locks, 2).await;
15896 assert!(!child_call.is_finished());
15897
15898 parent_gate.release();
15899 let (parent_record, child_record) =
15900 tokio::time::timeout(std::time::Duration::from_secs(2), async {
15901 tokio::join!(parent_call, child_call)
15902 })
15903 .await
15904 .expect("parent and spawned path mutations did not finish");
15905 assert!(parent_record.unwrap().success);
15906 assert!(child_record.unwrap().success);
15907 assert_eq!(std::fs::read_to_string(destination).unwrap(), "spawned");
15908 assert!(parent.resource_locks.read().is_empty());
15909 }
15910
15911 #[tokio::test]
15912 async fn cancelled_global_path_lock_waiter_does_not_retain_weak_entry() {
15913 let locks = new_tool_resource_locks();
15914 let holder_gate = PathMutationGate::new();
15915 let waiter_gate = PathMutationGate::new();
15916 waiter_gate.release();
15917 let holder = Arc::new(
15918 AgentBuilder::new()
15919 .system_prompt("Hold the global path lock.")
15920 .llm(Arc::new(mock_with_response("done")))
15921 .tool(Arc::new(BlockingPathMutationTool {
15922 id: "holder_write",
15923 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15924 gate: holder_gate.clone(),
15925 }))
15926 .build()
15927 .unwrap()
15928 .with_shared_resource_locks(Arc::clone(&locks)),
15929 );
15930 let waiter = Arc::new(
15931 AgentBuilder::new()
15932 .system_prompt("Wait for the global path lock.")
15933 .llm(Arc::new(mock_with_response("done")))
15934 .tool(Arc::new(BlockingPathMutationTool {
15935 id: "waiter_write",
15936 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
15937 gate: waiter_gate.clone(),
15938 }))
15939 .build()
15940 .unwrap()
15941 .with_shared_resource_locks(Arc::clone(&locks)),
15942 );
15943
15944 let holder_call = {
15945 let holder = Arc::clone(&holder);
15946 tokio::spawn(async move {
15947 holder
15948 .invoke_tool(ToolExecutionRequest::new(
15949 "holder-call",
15950 "holder_write",
15951 serde_json::json!({"path": "./shared.txt"}),
15952 ToolCallSource::Manual,
15953 ))
15954 .await
15955 .unwrap()
15956 })
15957 };
15958 holder_gate.wait_until_entered().await;
15959
15960 let waiter_call = {
15961 let waiter = Arc::clone(&waiter);
15962 tokio::spawn(async move {
15963 waiter
15964 .invoke_tool(ToolExecutionRequest::new(
15965 "waiter-call",
15966 "waiter_write",
15967 serde_json::json!({"path": "./shared.txt"}),
15968 ToolCallSource::Manual,
15969 ))
15970 .await
15971 .unwrap()
15972 })
15973 };
15974 wait_for_resource_lock_strong_count(&locks, 2).await;
15975 waiter.runtime_control().cancel_all();
15976
15977 let waiter_record = tokio::time::timeout(std::time::Duration::from_secs(2), waiter_call)
15978 .await
15979 .expect("cancelled lock waiter did not finish")
15980 .unwrap();
15981 assert!(!waiter_record.executed);
15982 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
15983 assert_eq!(
15984 locks
15985 .read()
15986 .get("path-mutation:global")
15987 .map_or(0, |lock| lock.strong_count()),
15988 1
15989 );
15990
15991 holder_gate.release();
15992 let holder_record = tokio::time::timeout(std::time::Duration::from_secs(2), holder_call)
15993 .await
15994 .expect("lock holder did not finish")
15995 .unwrap();
15996 assert!(holder_record.success);
15997 assert!(locks.read().is_empty());
15998 }
15999
16000 #[tokio::test]
16001 async fn path_mutation_policy_and_approval_denials_do_not_invoke_tools() {
16002 for denial in [MutationDenial::Policy, MutationDenial::Approval] {
16003 let tools: [Arc<dyn Tool>; 3] = [
16004 Arc::new(CopyPathTool::new()),
16005 Arc::new(MovePathTool::new()),
16006 Arc::new(DeletePathTool::new()),
16007 ];
16008 for tool in tools {
16009 assert_path_mutation_denied(tool, denial).await;
16010 }
16011 }
16012 }
16013
16014 #[tokio::test]
16015 async fn approval_argument_changes_are_rechecked_against_final_scope() {
16016 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16017 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16018 let entered = Arc::new(tokio::sync::Barrier::new(2));
16019 let release = Arc::new(tokio::sync::Notify::new());
16020 let handler = Arc::new(BlockingApprovalHandler {
16021 entered: Arc::clone(&entered),
16022 release: Arc::clone(&release),
16023 result: ApprovalResult::Modified {
16024 changes: HashMap::from([(
16025 "path".to_string(),
16026 Value::String("./after-approval.txt".to_string()),
16027 )]),
16028 },
16029 });
16030 let agent = Arc::new(
16031 AgentBuilder::new()
16032 .system_prompt("Test final scope validation.")
16033 .llm(Arc::new(mock_with_response("done")))
16034 .tool(Arc::new(LockedWriteTool {
16035 active: Arc::clone(&active),
16036 max_active: Arc::clone(&max_active),
16037 }))
16038 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16039 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16040 .approval_handler(handler)
16041 .build()
16042 .unwrap(),
16043 );
16044 let control = agent.runtime_control();
16045 let running = Arc::clone(&agent);
16046 let call = tokio::spawn(async move {
16047 running
16048 .invoke_tool(ToolExecutionRequest::new(
16049 "approval-scope",
16050 "locked_write",
16051 serde_json::json!({"path": "./before-approval.txt"}),
16052 ToolCallSource::Manual,
16053 ))
16054 .await
16055 .unwrap()
16056 });
16057 entered.wait().await;
16058 let expected_version = control.set_tool_scope(Vec::new());
16059 release.notify_one();
16060 let record = call.await.unwrap();
16061
16062 assert!(!record.executed);
16063 assert!(!record.success);
16064 assert_eq!(record.runtime_config_version, expected_version);
16065 assert_eq!(record.executed_arguments["path"], "./after-approval.txt");
16066 assert_eq!(max_active.load(Ordering::SeqCst), 0);
16067 assert_eq!(
16068 record.metadata["runtime_scope_snapshot"],
16069 serde_json::json!([])
16070 );
16071 }
16072
16073 #[tokio::test]
16074 async fn approval_is_rechecked_against_final_policy_snapshot() {
16075 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16076 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16077 let entered = Arc::new(tokio::sync::Barrier::new(2));
16078 let release = Arc::new(tokio::sync::Notify::new());
16079 let handler = Arc::new(BlockingApprovalHandler {
16080 entered: Arc::clone(&entered),
16081 release: Arc::clone(&release),
16082 result: ApprovalResult::Approved,
16083 });
16084 let agent = Arc::new(
16085 AgentBuilder::new()
16086 .system_prompt("Test final policy validation.")
16087 .llm(Arc::new(mock_with_response("done")))
16088 .tool(Arc::new(LockedWriteTool {
16089 active: Arc::clone(&active),
16090 max_active: Arc::clone(&max_active),
16091 }))
16092 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16093 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16094 .approval_handler(handler)
16095 .build()
16096 .unwrap(),
16097 );
16098 let control = agent.runtime_control();
16099 let running = Arc::clone(&agent);
16100 let call = tokio::spawn(async move {
16101 running
16102 .invoke_tool(ToolExecutionRequest::new(
16103 "approval-policy",
16104 "locked_write",
16105 serde_json::json!({"path": "./policy.txt"}),
16106 ToolCallSource::Manual,
16107 ))
16108 .await
16109 .unwrap()
16110 });
16111 entered.wait().await;
16112 let expected_version = control.set_tool_security(approval_security_config(false));
16113 release.notify_one();
16114 let record = call.await.unwrap();
16115
16116 assert!(!record.executed);
16117 assert!(!record.success);
16118 assert_eq!(record.runtime_config_version, expected_version);
16119 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16120 assert_eq!(max_active.load(Ordering::SeqCst), 0);
16121 assert!(record.metadata.contains_key("policy_snapshot"));
16122 }
16123
16124 #[test]
16125 fn invalid_live_policy_does_not_replace_snapshot_or_generation() {
16126 let agent = AgentBuilder::new()
16127 .system_prompt("Test runtime policy validation.")
16128 .llm(Arc::new(mock_with_response("done")))
16129 .build()
16130 .unwrap();
16131 let control = agent.runtime_control();
16132 let mut valid = ToolSecurityConfig::default();
16133 valid.tools.insert(
16134 "web_search".to_string(),
16135 ai_agents_tools::ToolPolicyConfig {
16136 max_results: Some(5),
16137 ..Default::default()
16138 },
16139 );
16140 let generation = control.try_set_tool_security(valid).unwrap();
16141
16142 let mut invalid = ToolSecurityConfig::default();
16143 invalid.tools.insert(
16144 "web_search".to_string(),
16145 ai_agents_tools::ToolPolicyConfig {
16146 max_results: Some(0),
16147 ..Default::default()
16148 },
16149 );
16150 let error = control.try_set_tool_security(invalid).unwrap_err();
16151
16152 assert!(
16153 error
16154 .to_string()
16155 .contains("max_results must be greater than 0")
16156 );
16157 assert_eq!(control.version(), generation);
16158 assert_eq!(
16159 control
16160 .state
16161 .tool_security_override
16162 .read()
16163 .as_ref()
16164 .unwrap()
16165 .config()
16166 .tools["web_search"]
16167 .max_results,
16168 Some(5)
16169 );
16170 }
16171
16172 #[tokio::test]
16173 async fn persistent_override_preserves_rate_history_within_generation() {
16174 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16175 let agent = AgentBuilder::new()
16176 .system_prompt("Test persistent policy overrides.")
16177 .llm(Arc::new(mock_with_response("done")))
16178 .tool(Arc::new(RecoveryTestTool {
16179 id: "limited_override".to_string(),
16180 succeeds: true,
16181 calls: Arc::clone(&calls),
16182 }))
16183 .build()
16184 .unwrap();
16185 let mut security = ToolSecurityConfig {
16186 enabled: true,
16187 fail_closed: true,
16188 ..Default::default()
16189 };
16190 let policy = ai_agents_tools::ToolPolicyConfig {
16191 write_paths: vec![".".to_string()],
16192 rate_limit: Some(1),
16193 ..Default::default()
16194 };
16195 security
16196 .tools
16197 .insert("limited_override".to_string(), policy);
16198 let generation = agent.runtime_control().set_tool_security(security);
16199
16200 let first = agent
16201 .invoke_tool(ToolExecutionRequest::new(
16202 "limited-first",
16203 "limited_override",
16204 serde_json::json!({"path": "./limited.txt"}),
16205 ToolCallSource::Manual,
16206 ))
16207 .await
16208 .unwrap();
16209 let second = agent
16210 .invoke_tool(ToolExecutionRequest::new(
16211 "limited-second",
16212 "limited_override",
16213 serde_json::json!({"path": "./limited.txt"}),
16214 ToolCallSource::Manual,
16215 ))
16216 .await
16217 .unwrap();
16218
16219 assert!(first.success);
16220 assert_eq!(first.policy_version, generation);
16221 assert!(!second.executed);
16222 assert!(second.output.contains("Rate limit exceeded"));
16223 assert_eq!(second.policy_version, generation);
16224 assert_eq!(calls.load(Ordering::SeqCst), 1);
16225 }
16226
16227 #[tokio::test]
16228 async fn concurrent_rate_admission_consumes_capacity_atomically() {
16229 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16230 let tool = Arc::new(RecoveryTestTool {
16231 id: "atomic_rate".to_string(),
16232 succeeds: true,
16233 calls: Arc::clone(&calls),
16234 });
16235 let arguments = serde_json::json!({"path": "./atomic-rate.txt"});
16236 let bindings = tool.policy_bindings();
16237 let classification = tool.classify_call(&arguments);
16238 let resource_keys =
16239 tool_resource_lock_keys(tool.id(), &arguments, &bindings, &classification);
16240 let mut security = ToolSecurityConfig {
16241 enabled: true,
16242 fail_closed: true,
16243 ..Default::default()
16244 };
16245 let policy = ai_agents_tools::ToolPolicyConfig {
16246 write_paths: vec![".".to_string()],
16247 rate_limit: Some(1),
16248 ..Default::default()
16249 };
16250 security.tools.insert(tool.id().to_string(), policy);
16251 let agent = Arc::new(
16252 AgentBuilder::new()
16253 .system_prompt("Test atomic rate admission.")
16254 .llm(Arc::new(mock_with_response("done")))
16255 .tool(tool)
16256 .tool_security(ToolSecurityEngine::new(security))
16257 .build()
16258 .unwrap(),
16259 );
16260 let held = agent
16261 .acquire_tool_resource_locks(&resource_keys)
16262 .await
16263 .unwrap();
16264 let left = {
16265 let agent = Arc::clone(&agent);
16266 let arguments = arguments.clone();
16267 tokio::spawn(async move {
16268 agent
16269 .invoke_tool(ToolExecutionRequest::new(
16270 "atomic-rate-left",
16271 "atomic_rate",
16272 arguments,
16273 ToolCallSource::Manual,
16274 ))
16275 .await
16276 .unwrap()
16277 })
16278 };
16279 let right = {
16280 let agent = Arc::clone(&agent);
16281 tokio::spawn(async move {
16282 agent
16283 .invoke_tool(ToolExecutionRequest::new(
16284 "atomic-rate-right",
16285 "atomic_rate",
16286 arguments,
16287 ToolCallSource::Manual,
16288 ))
16289 .await
16290 .unwrap()
16291 })
16292 };
16293 tokio::time::sleep(std::time::Duration::from_millis(25)).await;
16294 drop(held);
16295 let (left, right) = tokio::join!(left, right);
16296 let records = [left.unwrap(), right.unwrap()];
16297
16298 assert_eq!(records.iter().filter(|record| record.success).count(), 1);
16299 assert_eq!(records.iter().filter(|record| record.executed).count(), 1);
16300 assert!(
16301 records.iter().any(|record| {
16302 !record.executed && record.output.contains("Rate limit exceeded")
16303 })
16304 );
16305 assert_eq!(calls.load(Ordering::SeqCst), 1);
16306 }
16307
16308 #[tokio::test]
16309 async fn changed_policy_generation_invalidates_pending_approval() {
16310 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16311 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16312 let entered = Arc::new(tokio::sync::Barrier::new(2));
16313 let release = Arc::new(tokio::sync::Notify::new());
16314 let handler = Arc::new(BlockingApprovalHandler {
16315 entered: Arc::clone(&entered),
16316 release: Arc::clone(&release),
16317 result: ApprovalResult::Approved,
16318 });
16319 let agent = Arc::new(
16320 AgentBuilder::new()
16321 .system_prompt("Test stale approval denial.")
16322 .llm(Arc::new(mock_with_response("done")))
16323 .tool(Arc::new(LockedWriteTool {
16324 active: Arc::clone(&active),
16325 max_active: Arc::clone(&max_active),
16326 }))
16327 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
16328 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16329 .approval_handler(handler)
16330 .build()
16331 .unwrap(),
16332 );
16333 let running = Arc::clone(&agent);
16334 let call = tokio::spawn(async move {
16335 running
16336 .invoke_tool(ToolExecutionRequest::new(
16337 "stale-approval",
16338 "locked_write",
16339 serde_json::json!({"path": "./stale.txt"}),
16340 ToolCallSource::Manual,
16341 ))
16342 .await
16343 .unwrap()
16344 });
16345 entered.wait().await;
16346 let generation = agent
16347 .runtime_control()
16348 .set_tool_security(approval_security_config(true));
16349 release.notify_one();
16350 let record = call.await.unwrap();
16351
16352 assert!(!record.executed);
16353 assert!(record.output.contains("Approval became stale"));
16354 assert_eq!(record.policy_version, generation);
16355 assert_eq!(max_active.load(Ordering::SeqCst), 0);
16356 }
16357
16358 #[tokio::test]
16359 async fn final_policy_reapplies_argument_caps_after_approval_changes() {
16360 use ai_agents_hitl::CallbackHandler;
16361
16362 let mut security = ToolSecurityConfig {
16363 enabled: true,
16364 fail_closed: true,
16365 ..Default::default()
16366 };
16367 let policy = ai_agents_tools::ToolPolicyConfig {
16368 read_paths: vec![".".to_string()],
16369 max_results: Some(5),
16370 require_confirmation: true,
16371 ..Default::default()
16372 };
16373 security.tools.insert("context_echo".to_string(), policy);
16374 let handler = CallbackHandler::new(|_| ApprovalResult::Modified {
16375 changes: HashMap::from([("max_results".to_string(), serde_json::json!(99))]),
16376 });
16377 let agent = AgentBuilder::new()
16378 .system_prompt("Test final argument caps.")
16379 .llm(Arc::new(mock_with_response("done")))
16380 .tool(Arc::new(ContextEchoTool))
16381 .tool_security(ToolSecurityEngine::new(security))
16382 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
16383 .approval_handler(Arc::new(handler))
16384 .build()
16385 .unwrap();
16386
16387 let record = agent
16388 .invoke_tool(ToolExecutionRequest::new(
16389 "final-cap",
16390 "context_echo",
16391 serde_json::json!({"path": ".", "max_results": 1}),
16392 ToolCallSource::Manual,
16393 ))
16394 .await
16395 .unwrap();
16396
16397 assert!(record.success);
16398 assert_eq!(record.executed_arguments["max_results"], 5);
16399 assert_eq!(
16400 record.approval.unwrap().modified_arguments.unwrap()["max_results"],
16401 5
16402 );
16403 }
16404
16405 #[tokio::test]
16406 async fn no_binding_writes_use_canonical_fallback_lock() {
16407 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16408 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16409 let agent = Arc::new(
16410 AgentBuilder::new()
16411 .system_prompt("Test fallback resource locks.")
16412 .llm(Arc::new(mock_with_response("done")))
16413 .tool(Arc::new(NoBindingWriteTool {
16414 active: Arc::clone(&active),
16415 max_active: Arc::clone(&max_active),
16416 }))
16417 .build()
16418 .unwrap(),
16419 );
16420 let left = {
16421 let agent = Arc::clone(&agent);
16422 tokio::spawn(async move {
16423 agent
16424 .invoke_tool(ToolExecutionRequest::new(
16425 "no-binding-left",
16426 "no_binding_write",
16427 serde_json::json!({}),
16428 ToolCallSource::Manual,
16429 ))
16430 .await
16431 .unwrap()
16432 })
16433 };
16434 let right = {
16435 let agent = Arc::clone(&agent);
16436 tokio::spawn(async move {
16437 agent
16438 .invoke_tool(ToolExecutionRequest::new(
16439 "no-binding-right",
16440 "no_binding_write",
16441 serde_json::json!({}),
16442 ToolCallSource::Manual,
16443 ))
16444 .await
16445 .unwrap()
16446 })
16447 };
16448 let (left, right) = tokio::join!(left, right);
16449
16450 assert!(left.unwrap().success);
16451 assert!(right.unwrap().success);
16452 assert_eq!(max_active.load(Ordering::SeqCst), 1);
16453 }
16454
16455 #[tokio::test]
16456 async fn parent_and_child_paths_share_a_resource_lock() {
16457 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16458 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16459 let agent = Arc::new(
16460 AgentBuilder::new()
16461 .system_prompt("Test parent child resource locks.")
16462 .llm(Arc::new(mock_with_response("done")))
16463 .tool(Arc::new(LockedWriteTool {
16464 active: Arc::clone(&active),
16465 max_active: Arc::clone(&max_active),
16466 }))
16467 .build()
16468 .unwrap(),
16469 );
16470 let parent = format!("./lock-parent-{}", uuid::Uuid::new_v4());
16471 let child = format!("{}/child.txt", parent);
16472 let left = {
16473 let agent = Arc::clone(&agent);
16474 tokio::spawn(async move {
16475 agent
16476 .invoke_tool(ToolExecutionRequest::new(
16477 "parent-lock",
16478 "locked_write",
16479 serde_json::json!({"path": parent}),
16480 ToolCallSource::Manual,
16481 ))
16482 .await
16483 .unwrap()
16484 })
16485 };
16486 let right = {
16487 let agent = Arc::clone(&agent);
16488 tokio::spawn(async move {
16489 agent
16490 .invoke_tool(ToolExecutionRequest::new(
16491 "child-lock",
16492 "locked_write",
16493 serde_json::json!({"path": child}),
16494 ToolCallSource::Manual,
16495 ))
16496 .await
16497 .unwrap()
16498 })
16499 };
16500 let (left, right) = tokio::join!(left, right);
16501
16502 assert!(left.unwrap().success);
16503 assert!(right.unwrap().success);
16504 assert_eq!(max_active.load(Ordering::SeqCst), 1);
16505 }
16506
16507 #[tokio::test]
16508 async fn tool_hooks_can_reenter_after_resource_guards_are_dropped() {
16509 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16510 let hooks = Arc::new(ReentrantToolHooks {
16511 agent: parking_lot::Mutex::new(None),
16512 invoked: AtomicBool::new(false),
16513 nested_success: AtomicBool::new(false),
16514 });
16515 let agent = Arc::new(
16516 AgentBuilder::new()
16517 .system_prompt("Test hook reentrancy.")
16518 .llm(Arc::new(mock_with_response("done")))
16519 .tool(Arc::new(RecoveryTestTool {
16520 id: "reentrant_write".to_string(),
16521 succeeds: true,
16522 calls: Arc::clone(&calls),
16523 }))
16524 .hooks(hooks.clone())
16525 .build()
16526 .unwrap(),
16527 );
16528 *hooks.agent.lock() = Some(Arc::downgrade(&agent));
16529 let record = tokio::time::timeout(
16530 std::time::Duration::from_secs(2),
16531 agent.invoke_tool(ToolExecutionRequest::new(
16532 "outer-hook-call",
16533 "reentrant_write",
16534 serde_json::json!({"path": "./hook.txt"}),
16535 ToolCallSource::Manual,
16536 )),
16537 )
16538 .await
16539 .expect("tool completion hook must not retain resource guards")
16540 .unwrap();
16541
16542 assert!(record.success);
16543 assert!(hooks.nested_success.load(Ordering::SeqCst));
16544 assert_eq!(calls.load(Ordering::SeqCst), 2);
16545 }
16546
16547 #[tokio::test]
16548 async fn fallback_releases_primary_resource_locks() {
16549 let primary_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16550 let fallback_calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
16551 let agent = AgentBuilder::new()
16552 .system_prompt("Test fallback execution.")
16553 .llm(Arc::new(mock_with_response("done")))
16554 .tool(Arc::new(RecoveryTestTool {
16555 id: "primary".to_string(),
16556 succeeds: false,
16557 calls: Arc::clone(&primary_calls),
16558 }))
16559 .tool(Arc::new(RecoveryTestTool {
16560 id: "fallback".to_string(),
16561 succeeds: true,
16562 calls: Arc::clone(&fallback_calls),
16563 }))
16564 .recovery_manager(recovery_manager_with_fallbacks([(
16565 "primary".to_string(),
16566 "fallback".to_string(),
16567 )]))
16568 .build()
16569 .unwrap();
16570 let record = tokio::time::timeout(
16571 std::time::Duration::from_secs(2),
16572 agent.invoke_tool(ToolExecutionRequest::new(
16573 "fallback-call",
16574 "primary",
16575 serde_json::json!({"path": "./shared.txt"}),
16576 ToolCallSource::Manual,
16577 )),
16578 )
16579 .await
16580 .expect("fallback must not retain the primary resource guard")
16581 .unwrap();
16582
16583 assert!(record.success);
16584 assert_eq!(record.canonical_id, "fallback");
16585 assert_eq!(record.call_id, "fallback-call");
16586 assert!(matches!(record.source, ToolCallSource::Fallback { .. }));
16587 assert_eq!(primary_calls.load(Ordering::SeqCst), 1);
16588 assert_eq!(fallback_calls.load(Ordering::SeqCst), 1);
16589 }
16590
16591 #[tokio::test]
16592 async fn diagnostics_without_provider_records_unavailable_without_execution() {
16593 let mock = mock_with_response("hello");
16594 let yaml = r#"
16595name: DiagnosticsNoProviderAgent
16596system_prompt: "Review diagnostics."
16597tools: [diagnostics]
16598"#;
16599 let agent = AgentBuilder::from_yaml(yaml)
16600 .unwrap()
16601 .llm(Arc::new(mock))
16602 .auto_configure_features()
16603 .unwrap()
16604 .build()
16605 .unwrap();
16606
16607 let record = agent
16608 .invoke_tool(ToolExecutionRequest::new(
16609 "diagnostics-call",
16610 "diagnostics",
16611 serde_json::json!({}),
16612 ToolCallSource::Manual,
16613 ))
16614 .await
16615 .unwrap();
16616
16617 assert!(!record.executed);
16618 assert!(!record.success);
16619 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16620 }
16621
16622 #[tokio::test]
16623 async fn web_search_without_provider_records_unavailable_without_execution() {
16624 let mock = mock_with_response("hello");
16625 let yaml = r#"
16626name: WebSearchNoProviderAgent
16627system_prompt: "You search the web."
16628tools: [web_search]
16629"#;
16630 let agent = AgentBuilder::from_yaml(yaml)
16631 .unwrap()
16632 .llm(Arc::new(mock))
16633 .auto_configure_features()
16634 .unwrap()
16635 .build()
16636 .unwrap();
16637
16638 let record = agent
16639 .invoke_tool(ToolExecutionRequest::new(
16640 "web-search-call",
16641 "web_search",
16642 serde_json::json!({"query": "rust async"}),
16643 ToolCallSource::Manual,
16644 ))
16645 .await
16646 .unwrap();
16647
16648 assert!(!record.executed);
16649 assert!(!record.success);
16650 assert_eq!(record.policy.outcome, PermissionOutcome::Unavailable);
16651 }
16652
16653 #[tokio::test]
16654 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_omitted() {
16655 let mock = mock_with_response("hello");
16656 let yaml = r#"
16657name: SpawnerNoGrantAgent
16658system_prompt: "You manage agents."
16659spawner:
16660 max_agents: 2
16661"#;
16662 let agent = AgentBuilder::from_yaml(yaml)
16663 .unwrap()
16664 .llm(Arc::new(mock))
16665 .auto_configure_features()
16666 .unwrap()
16667 .auto_configure_spawner()
16668 .await
16669 .unwrap()
16670 .build()
16671 .unwrap();
16672
16673 let available = agent.get_available_tool_ids().await.unwrap();
16674 assert!(available.is_empty());
16675 }
16676
16677 #[tokio::test]
16678 async fn test_spawner_section_does_not_grant_core_tools_when_top_level_tools_empty() {
16679 let mock = mock_with_response("hello");
16680 let yaml = r#"
16681name: EmptySpawnerNoGrantAgent
16682system_prompt: "You manage agents."
16683tools: []
16684spawner:
16685 max_agents: 2
16686"#;
16687 let agent = AgentBuilder::from_yaml(yaml)
16688 .unwrap()
16689 .llm(Arc::new(mock))
16690 .auto_configure_features()
16691 .unwrap()
16692 .auto_configure_spawner()
16693 .await
16694 .unwrap()
16695 .build()
16696 .unwrap();
16697
16698 let available = agent.get_available_tool_ids().await.unwrap();
16699 assert!(available.is_empty());
16700 }
16701
16702 #[tokio::test]
16703 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_empty() {
16704 let mock = mock_with_response("hello");
16705 let yaml = r#"
16706name: ManagementGrantAgent
16707system_prompt: "You manage agents."
16708tools: []
16709spawner:
16710 management_tools: true
16711"#;
16712 let agent = AgentBuilder::from_yaml(yaml)
16713 .unwrap()
16714 .llm(Arc::new(mock))
16715 .auto_configure_features()
16716 .unwrap()
16717 .auto_configure_spawner()
16718 .await
16719 .unwrap()
16720 .build()
16721 .unwrap();
16722
16723 let available = agent.get_available_tool_ids().await.unwrap();
16724 assert_eq!(available.len(), 4);
16725 assert!(available.contains(&"spawn_agent".to_string()));
16726 assert!(available.contains(&"send_agent_message".to_string()));
16727 assert!(available.contains(&"list_agents".to_string()));
16728 assert!(available.contains(&"remove_agent".to_string()));
16729 }
16730
16731 #[tokio::test]
16732 async fn test_management_tools_flag_grants_core_tools_when_top_level_tools_omitted() {
16733 let mock = mock_with_response("hello");
16734 let yaml = r#"
16735name: ManagementOmittedToolsGrantAgent
16736system_prompt: "You manage agents."
16737spawner:
16738 management_tools: true
16739"#;
16740 let agent = AgentBuilder::from_yaml(yaml)
16741 .unwrap()
16742 .llm(Arc::new(mock))
16743 .auto_configure_features()
16744 .unwrap()
16745 .auto_configure_spawner()
16746 .await
16747 .unwrap()
16748 .build()
16749 .unwrap();
16750
16751 let available = agent.get_available_tool_ids().await.unwrap();
16752 assert_eq!(available.len(), 4);
16753 assert!(available.contains(&"spawn_agent".to_string()));
16754 assert!(available.contains(&"send_agent_message".to_string()));
16755 assert!(available.contains(&"list_agents".to_string()));
16756 assert!(available.contains(&"remove_agent".to_string()));
16757 }
16758
16759 #[tokio::test]
16760 async fn test_management_tools_selected_grants_only_selected_tools() {
16761 let mock = mock_with_response("hello");
16762 let yaml = r#"
16763name: ManagementSelectedGrantAgent
16764system_prompt: "You manage agents."
16765tools: []
16766spawner:
16767 management_tools:
16768 - spawn_agent
16769 - send_agent_message
16770 - list_agents
16771"#;
16772 let agent = AgentBuilder::from_yaml(yaml)
16773 .unwrap()
16774 .llm(Arc::new(mock))
16775 .auto_configure_features()
16776 .unwrap()
16777 .auto_configure_spawner()
16778 .await
16779 .unwrap()
16780 .build()
16781 .unwrap();
16782
16783 let available = agent.get_available_tool_ids().await.unwrap();
16784 assert_eq!(available.len(), 3);
16785 assert!(available.contains(&"spawn_agent".to_string()));
16786 assert!(available.contains(&"send_agent_message".to_string()));
16787 assert!(available.contains(&"list_agents".to_string()));
16788 assert!(!available.contains(&"remove_agent".to_string()));
16789 }
16790
16791 #[tokio::test]
16792 async fn test_orchestration_tools_flag_grants_tools_when_top_level_tools_empty() {
16793 let mock = mock_with_response("hello");
16794 let yaml = r#"
16795name: OrchestrationGrantAgent
16796system_prompt: "You coordinate agents."
16797llms:
16798 default:
16799 provider: openai
16800 model: gpt-4
16801 router:
16802 provider: openai
16803 model: gpt-4
16804llm:
16805 default: default
16806 router: router
16807tools: []
16808spawner:
16809 orchestration_tools: true
16810"#;
16811 let agent = AgentBuilder::from_yaml(yaml)
16812 .unwrap()
16813 .llm(Arc::new(mock))
16814 .auto_configure_features()
16815 .unwrap()
16816 .auto_configure_spawner()
16817 .await
16818 .unwrap()
16819 .build()
16820 .unwrap();
16821
16822 let available = agent.get_available_tool_ids().await.unwrap();
16823 assert_eq!(available.len(), 5);
16824 assert!(available.contains(&"route_to_agent".to_string()));
16825 assert!(available.contains(&"pipeline_process".to_string()));
16826 assert!(available.contains(&"concurrent_ask".to_string()));
16827 assert!(available.contains(&"group_discussion".to_string()));
16828 assert!(available.contains(&"handoff_conversation".to_string()));
16829 }
16830
16831 #[tokio::test]
16832 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_empty() {
16833 let mock = mock_with_response("hello");
16834 let yaml = r#"
16835name: PersonaGrantAgent
16836system_prompt: "You can evolve persona."
16837llm:
16838 provider: openai
16839 model: gpt-4
16840tools: []
16841persona:
16842 identity:
16843 name: "Guide"
16844 role: "Helper"
16845 evolution:
16846 enabled: true
16847 allow_llm_evolve: true
16848 mutable_fields:
16849 - traits.personality
16850"#;
16851 let agent = AgentBuilder::from_yaml(yaml)
16852 .unwrap()
16853 .llm(Arc::new(mock))
16854 .build()
16855 .unwrap();
16856
16857 let available = agent.get_available_tool_ids().await.unwrap();
16858 assert_eq!(available, vec!["persona_evolve".to_string()]);
16859 }
16860
16861 #[tokio::test]
16862 async fn test_persona_evolve_flag_grants_tool_when_top_level_tools_omitted() {
16863 let mock = mock_with_response("hello");
16864 let yaml = r#"
16865name: PersonaOmittedToolsGrantAgent
16866system_prompt: "You can evolve persona."
16867llm:
16868 provider: openai
16869 model: gpt-4
16870persona:
16871 identity:
16872 name: "Guide"
16873 role: "Helper"
16874 evolution:
16875 enabled: true
16876 allow_llm_evolve: true
16877 mutable_fields:
16878 - traits.personality
16879"#;
16880 let agent = AgentBuilder::from_yaml(yaml)
16881 .unwrap()
16882 .llm(Arc::new(mock))
16883 .build()
16884 .unwrap();
16885
16886 let available = agent.get_available_tool_ids().await.unwrap();
16887 assert_eq!(available, vec!["persona_evolve".to_string()]);
16888 }
16889
16890 #[tokio::test]
16891 async fn test_omitted_yaml_tools_exposes_no_tools() {
16892 let mock = mock_with_response("hello");
16893 let yaml = r#"
16894name: NoToolsAgent
16895system_prompt: "You are helpful."
16896"#;
16897 let agent = AgentBuilder::from_yaml(yaml)
16898 .unwrap()
16899 .llm(Arc::new(mock))
16900 .auto_configure_features()
16901 .unwrap()
16902 .build()
16903 .unwrap();
16904
16905 let available = agent.get_available_tool_ids().await.unwrap();
16906 assert!(available.is_empty());
16907 }
16908
16909 #[tokio::test]
16910 async fn runtime_scope_cannot_widen_omitted_or_empty_yaml_grants() {
16911 for tools in ["", "tools: []"] {
16912 let yaml = format!(
16913 r#"
16914name: RuntimeScopeNoGrantAgent
16915system_prompt: "No ordinary tools are granted."
16916{tools}
16917"#
16918 );
16919 let agent = AgentBuilder::from_yaml(&yaml)
16920 .unwrap()
16921 .llm(Arc::new(mock_with_response("done")))
16922 .auto_configure_features()
16923 .unwrap()
16924 .build()
16925 .unwrap();
16926
16927 agent
16928 .runtime_control()
16929 .set_tool_scope(vec!["calculator".to_string()]);
16930
16931 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
16932 }
16933 }
16934
16935 #[tokio::test]
16936 async fn runtime_scope_widening_attempt_keeps_only_declared_tools() {
16937 let yaml = r#"
16938name: RuntimeScopeWideningAgent
16939system_prompt: "Runtime scope cannot add authority."
16940tools: [calculator]
16941"#;
16942 let agent = AgentBuilder::from_yaml(yaml)
16943 .unwrap()
16944 .llm(Arc::new(mock_with_response("done")))
16945 .auto_configure_features()
16946 .unwrap()
16947 .build()
16948 .unwrap();
16949
16950 agent
16951 .runtime_control()
16952 .set_tool_scope(vec!["calculator".to_string(), "datetime".to_string()]);
16953
16954 assert_eq!(
16955 agent.get_available_tool_ids().await.unwrap(),
16956 vec!["calculator".to_string()]
16957 );
16958 }
16959
16960 #[tokio::test]
16961 async fn runtime_scope_is_canonical_unique_ordered_and_clear_restores_declared_grant() {
16962 let yaml = r#"
16963name: RuntimeScopeIntersectionAgent
16964system_prompt: "Use only declared tools."
16965tools: [calculator, datetime]
16966"#;
16967 let agent = AgentBuilder::from_yaml(yaml)
16968 .unwrap()
16969 .llm(Arc::new(mock_with_response("done")))
16970 .auto_configure_features()
16971 .unwrap()
16972 .build()
16973 .unwrap();
16974 let mut aliases = ai_agents_tools::ToolAliases::default();
16975 aliases
16976 .names
16977 .insert("en".to_string(), "calculate_alias".to_string());
16978 agent.tools.set_tool_aliases("calculator", aliases);
16979 let control = agent.runtime_control();
16980
16981 control.set_tool_scope(vec![
16982 "datetime".to_string(),
16983 "calculate_alias".to_string(),
16984 "calculator".to_string(),
16985 "unknown".to_string(),
16986 "datetime".to_string(),
16987 ]);
16988 assert_eq!(
16989 agent.get_available_tool_ids().await.unwrap(),
16990 vec!["calculator".to_string(), "datetime".to_string()]
16991 );
16992
16993 control.set_tool_scope(vec!["datetime".to_string()]);
16994 assert_eq!(
16995 agent.get_available_tool_ids().await.unwrap(),
16996 vec!["datetime".to_string()]
16997 );
16998
16999 control.clear_tool_scope_override();
17000 assert_eq!(
17001 agent.get_available_tool_ids().await.unwrap(),
17002 vec!["calculator".to_string(), "datetime".to_string()]
17003 );
17004 }
17005
17006 #[tokio::test]
17007 async fn runtime_scope_preserves_programmatic_registration_as_declared_grant() {
17008 let agent = AgentBuilder::new()
17009 .system_prompt("Use registered tools.")
17010 .llm(Arc::new(mock_with_response("done")))
17011 .tool(Arc::new(ContextEchoTool))
17012 .tool(Arc::new(SlowTool))
17013 .build()
17014 .unwrap();
17015
17016 agent.runtime_control().set_tool_scope(vec![
17017 "Context Echo".to_string(),
17018 "context_echo".to_string(),
17019 "unknown".to_string(),
17020 ]);
17021
17022 assert_eq!(
17023 agent.get_available_tool_ids().await.unwrap(),
17024 vec!["context_echo".to_string()]
17025 );
17026 }
17027
17028 #[tokio::test]
17029 async fn nested_state_scopes_intersect_every_ancestor_with_aliases() {
17030 let yaml = r#"
17031name: NestedStateScopeAgent
17032system_prompt: "Honor every state scope."
17033tools: [calculator, datetime, echo]
17034states:
17035 initial: root
17036 states:
17037 root:
17038 tools: [calculate_alias, datetime]
17039 initial: middle
17040 states:
17041 middle:
17042 initial: leaf
17043 states:
17044 leaf:
17045 tools: [datetime_alias, echo]
17046"#;
17047 let agent = AgentBuilder::from_yaml(yaml)
17048 .unwrap()
17049 .llm(Arc::new(mock_with_response("done")))
17050 .auto_configure_features()
17051 .unwrap()
17052 .build()
17053 .unwrap();
17054 let mut calculator_aliases = ai_agents_tools::ToolAliases::default();
17055 calculator_aliases
17056 .names
17057 .insert("en".to_string(), "calculate_alias".to_string());
17058 agent
17059 .tools
17060 .set_tool_aliases("calculator", calculator_aliases);
17061 let mut datetime_aliases = ai_agents_tools::ToolAliases::default();
17062 datetime_aliases
17063 .names
17064 .insert("en".to_string(), "datetime_alias".to_string());
17065 agent.tools.set_tool_aliases("datetime", datetime_aliases);
17066 agent.runtime_control().set_tool_scope(vec![
17067 "unknown".to_string(),
17068 "datetime_alias".to_string(),
17069 "calculate_alias".to_string(),
17070 "datetime".to_string(),
17071 ]);
17072
17073 assert_eq!(agent.current_state().as_deref(), Some("root.middle.leaf"));
17074 assert_eq!(
17075 agent.get_available_tool_ids().await.unwrap(),
17076 vec!["datetime".to_string()]
17077 );
17078 }
17079
17080 #[tokio::test]
17081 async fn ancestor_empty_state_scope_denies_omitted_descendants() {
17082 let yaml = r#"
17083name: NestedEmptyStateScopeAgent
17084system_prompt: "An empty ancestor scope denies all tools."
17085tools: [calculator]
17086states:
17087 initial: root
17088 states:
17089 root:
17090 tools: []
17091 initial: middle
17092 states:
17093 middle:
17094 initial: leaf
17095 states:
17096 leaf: {}
17097"#;
17098 let agent = AgentBuilder::from_yaml(yaml)
17099 .unwrap()
17100 .llm(Arc::new(mock_with_response("done")))
17101 .auto_configure_features()
17102 .unwrap()
17103 .build()
17104 .unwrap();
17105
17106 assert!(agent.get_available_tool_ids().await.unwrap().is_empty());
17107 }
17108
17109 #[tokio::test]
17110 async fn state_change_during_approval_invalidates_the_reviewed_authority() {
17111 let active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17112 let max_active = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17113 let entered = Arc::new(tokio::sync::Barrier::new(2));
17114 let release = Arc::new(tokio::sync::Notify::new());
17115 let handler = Arc::new(BlockingApprovalHandler {
17116 entered: Arc::clone(&entered),
17117 release: Arc::clone(&release),
17118 result: ApprovalResult::Approved,
17119 });
17120 let yaml = r#"
17121name: ApprovalStateGenerationAgent
17122system_prompt: "State authority may change during approval."
17123tools: [locked_write]
17124states:
17125 initial: first
17126 states:
17127 first:
17128 tools: [locked_write]
17129 second:
17130 tools: [locked_write]
17131"#;
17132 let agent = Arc::new(
17133 AgentBuilder::from_yaml(yaml)
17134 .unwrap()
17135 .llm(Arc::new(mock_with_response("done")))
17136 .tool(Arc::new(LockedWriteTool {
17137 active: Arc::clone(&active),
17138 max_active: Arc::clone(&max_active),
17139 }))
17140 .tool_security(ToolSecurityEngine::new(approval_security_config(true)))
17141 .hitl_engine(HITLEngine::new(ai_agents_hitl::HITLConfig::default()))
17142 .approval_handler(handler)
17143 .build()
17144 .unwrap(),
17145 );
17146 let running = Arc::clone(&agent);
17147 let call = tokio::spawn(async move {
17148 running
17149 .invoke_tool(ToolExecutionRequest::new(
17150 "approval-state-generation",
17151 "locked_write",
17152 serde_json::json!({"path": "./state-generation.txt"}),
17153 ToolCallSource::Manual,
17154 ))
17155 .await
17156 .unwrap()
17157 });
17158
17159 entered.wait().await;
17160 agent.transition_to("second").await.unwrap();
17161 release.notify_one();
17162 let record = call.await.unwrap();
17163
17164 assert!(!record.executed);
17165 assert!(record.output.contains("Approval became stale"));
17166 assert_eq!(max_active.load(Ordering::SeqCst), 0);
17167 }
17168
17169 #[tokio::test]
17170 async fn state_change_while_waiting_for_resource_lock_fails_final_admission() {
17171 let holder_gate = PathMutationGate::new();
17172 let waiter_gate = PathMutationGate::new();
17173 let yaml = r#"
17174name: LockedStateGenerationAgent
17175system_prompt: "State authority must remain stable through admission."
17176tools: [state_lock_holder, state_lock_waiter]
17177states:
17178 initial: first
17179 states:
17180 first:
17181 tools: [state_lock_holder, state_lock_waiter]
17182 second:
17183 tools: [state_lock_holder, state_lock_waiter]
17184"#;
17185 let agent = Arc::new(
17186 AgentBuilder::from_yaml(yaml)
17187 .unwrap()
17188 .llm(Arc::new(mock_with_response("done")))
17189 .tool(Arc::new(BlockingPathMutationTool {
17190 id: "state_lock_holder",
17191 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17192 gate: holder_gate.clone(),
17193 }))
17194 .tool(Arc::new(BlockingPathMutationTool {
17195 id: "state_lock_waiter",
17196 path_fields: vec![ai_agents_core::PathPolicyBinding::write("path")],
17197 gate: waiter_gate.clone(),
17198 }))
17199 .build()
17200 .unwrap(),
17201 );
17202 let holder_call = {
17203 let agent = Arc::clone(&agent);
17204 tokio::spawn(async move {
17205 agent
17206 .invoke_tool(ToolExecutionRequest::new(
17207 "state-lock-holder",
17208 "state_lock_holder",
17209 serde_json::json!({"path": "./shared-state-path.txt"}),
17210 ToolCallSource::Manual,
17211 ))
17212 .await
17213 .unwrap()
17214 })
17215 };
17216 holder_gate.wait_until_entered().await;
17217 let waiter_call = {
17218 let agent = Arc::clone(&agent);
17219 tokio::spawn(async move {
17220 agent
17221 .invoke_tool(ToolExecutionRequest::new(
17222 "state-lock-waiter",
17223 "state_lock_waiter",
17224 serde_json::json!({"path": "./shared-state-path.txt"}),
17225 ToolCallSource::Manual,
17226 ))
17227 .await
17228 .unwrap()
17229 })
17230 };
17231
17232 wait_for_resource_lock_strong_count(&agent.resource_locks, 2).await;
17233 agent.transition_to("second").await.unwrap();
17234 holder_gate.release();
17235 let holder_record = holder_call.await.unwrap();
17236 let waiter_record = waiter_call.await.unwrap();
17237
17238 assert!(holder_record.success);
17239 assert!(!waiter_record.executed);
17240 assert!(
17241 waiter_record
17242 .output
17243 .contains("state scope changed before admission")
17244 );
17245 assert!(!waiter_gate.entered.load(Ordering::SeqCst));
17246 }
17247
17248 #[tokio::test]
17249 async fn test_state_tools_cannot_widen_top_level_grant() {
17250 let mock = mock_with_response("hello");
17251 let yaml = r#"
17252name: NarrowToolsAgent
17253system_prompt: "You are helpful."
17254tools:
17255 - calculator
17256states:
17257 initial: current
17258 states:
17259 current:
17260 tools: [datetime]
17261"#;
17262 let agent = AgentBuilder::from_yaml(yaml)
17263 .unwrap()
17264 .llm(Arc::new(mock))
17265 .auto_configure_features()
17266 .unwrap()
17267 .build()
17268 .unwrap();
17269
17270 let available = agent.get_available_tool_ids().await.unwrap();
17271 assert!(available.is_empty());
17272 }
17273
17274 #[tokio::test]
17276 async fn test_integration_tool_execution() {
17277 let mock = mock_with_responses(vec![
17279 r#"I'll calculate that for you.
17281[TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}]"#,
17282 "The answer is 4.",
17284 ]);
17285 let mut tools = ai_agents_tools::ToolRegistry::new();
17286 tools
17287 .register(Arc::new(ai_agents_tools::CalculatorTool))
17288 .unwrap();
17289
17290 let agent = AgentBuilder::new()
17291 .system_prompt("You are a calculator assistant.")
17292 .llm(Arc::new(mock))
17293 .tools(tools)
17294 .build()
17295 .unwrap();
17296
17297 let response = agent.chat("What is 2+2?").await.unwrap();
17298 assert!(!response.content.is_empty());
17300 }
17301
17302 #[tokio::test]
17303 async fn test_tool_hitl_rejection_finalizes_blocking_turn() {
17304 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17305 let hooks = Arc::new(ResponseCountingHooks {
17306 responses: Arc::clone(&responses),
17307 });
17308 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17309 let yaml = r#"
17310name: ToolRejectAgent
17311system_prompt: "You use tools when requested."
17312tools:
17313 - echo
17314hitl:
17315 tools:
17316 echo:
17317 require_approval: true
17318 approval_message: "Approve echo?"
17319"#;
17320 let agent = AgentBuilder::from_yaml(yaml)
17321 .unwrap()
17322 .llm(Arc::new(mock))
17323 .auto_configure_features()
17324 .unwrap()
17325 .hooks(hooks)
17326 .build()
17327 .unwrap();
17328
17329 let response = agent.chat("echo hello").await.unwrap();
17330
17331 assert!(
17332 response.content.contains("Operation cancelled"),
17333 "unexpected response: {}",
17334 response.content
17335 );
17336 assert_eq!(responses.load(Ordering::SeqCst), 1);
17337 let messages = agent.memory.get_messages(None).await.unwrap();
17338 assert_eq!(messages.len(), 3);
17339 assert_eq!(messages[0].content, "echo hello");
17340 assert!(messages[1].content.contains("\"tool\":\"echo\""));
17341 assert!(messages[2].content.contains("rejected by the approver"));
17342 }
17343
17344 #[tokio::test]
17345 async fn test_tool_hitl_rejection_finalizes_streaming_turn() {
17346 use futures::StreamExt;
17347
17348 let responses = Arc::new(std::sync::atomic::AtomicUsize::new(0));
17349 let hooks = Arc::new(ResponseCountingHooks {
17350 responses: Arc::clone(&responses),
17351 });
17352 let mock = mock_with_response(r#"{"tool":"echo","arguments":{"message":"hello"}}"#);
17353 let yaml = r#"
17354name: ToolRejectStreamingAgent
17355system_prompt: "You use tools when requested."
17356tools:
17357 - echo
17358streaming:
17359 enabled: true
17360hitl:
17361 tools:
17362 echo:
17363 require_approval: true
17364 approval_message: "Approve echo?"
17365"#;
17366 let agent = AgentBuilder::from_yaml(yaml)
17367 .unwrap()
17368 .llm(Arc::new(mock))
17369 .auto_configure_features()
17370 .unwrap()
17371 .hooks(hooks)
17372 .build()
17373 .unwrap();
17374
17375 let mut stream = agent.chat_stream("echo hello").await.unwrap();
17376 let mut terminal_error = String::new();
17377 let mut done = false;
17378 while let Some(chunk) = stream.next().await {
17379 match chunk {
17380 StreamChunk::Error { message } => terminal_error = message,
17381 StreamChunk::Done {} => {
17382 done = true;
17383 break;
17384 }
17385 _ => {}
17386 }
17387 }
17388
17389 assert!(done);
17390 assert!(
17391 terminal_error.contains("Operation cancelled"),
17392 "unexpected terminal error: {}",
17393 terminal_error
17394 );
17395 assert_eq!(responses.load(Ordering::SeqCst), 1);
17396 let messages = agent.memory.get_messages(None).await.unwrap();
17397 assert_eq!(messages.len(), 3);
17398 assert_eq!(messages[0].content, "echo hello");
17399 assert!(messages[1].content.contains("\"tool\":\"echo\""));
17400 assert!(messages[2].content.contains("rejected by the approver"));
17401 }
17402
17403 #[tokio::test]
17404 async fn test_pre_response_guard_transition_skips_old_state_llm() {
17405 let mock = mock_with_response("Billing state response");
17406 let call_counter = mock.clone();
17407 let yaml = r#"
17408name: OptimizedStateAgent
17409system_prompt: "You route before answering."
17410runtime:
17411 optimization:
17412 enabled: true
17413 pre_response_deterministic_transitions: true
17414states:
17415 initial: greeting
17416 states:
17417 greeting:
17418 prompt: "Old state prompt that should be skipped."
17419 transitions:
17420 - to: billing
17421 guard:
17422 context:
17423 topic:
17424 eq: billing
17425 timing: pre_response
17426 billing:
17427 prompt: "Answer from the billing state."
17428"#;
17429 let agent = AgentBuilder::from_yaml(yaml)
17430 .unwrap()
17431 .llm(Arc::new(mock))
17432 .build()
17433 .unwrap();
17434 agent
17435 .set_context("topic", serde_json::json!("billing"))
17436 .unwrap();
17437
17438 let response = agent.chat("I need billing help").await.unwrap();
17439
17440 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17441 assert_eq!(response.content, "Billing state response");
17442 assert_eq!(call_counter.call_count(), 1);
17443 assert_eq!(agent.actor_facts().len(), 0);
17444 }
17445
17446 #[tokio::test]
17447 async fn test_set_context_supports_dotted_paths_for_pre_response_guards() {
17448 let mock = mock_with_response("Billing state response");
17449 let call_counter = mock.clone();
17450 let yaml = r#"
17451name: OptimizedStateAgent
17452system_prompt: "You route before answering."
17453runtime:
17454 optimization:
17455 enabled: true
17456 pre_response_deterministic_transitions: true
17457context:
17458 request:
17459 type: runtime
17460 default:
17461 topic: general
17462states:
17463 initial: greeting
17464 states:
17465 greeting:
17466 prompt: "Old state prompt that should be skipped."
17467 transitions:
17468 - to: billing
17469 guard:
17470 context:
17471 request.topic:
17472 eq: billing
17473 timing: pre_response
17474 billing:
17475 prompt: "Answer from the billing state."
17476"#;
17477 let agent = AgentBuilder::from_yaml(yaml)
17478 .unwrap()
17479 .llm(Arc::new(mock))
17480 .build()
17481 .unwrap();
17482 agent
17483 .set_context("request.topic", serde_json::json!("billing"))
17484 .unwrap();
17485
17486 let response = agent.chat("I need billing help").await.unwrap();
17487
17488 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17489 assert_eq!(response.content, "Billing state response");
17490 assert_eq!(call_counter.call_count(), 1);
17491 assert_eq!(
17492 agent.get_context().get("request"),
17493 Some(&serde_json::json!({"topic": "billing"}))
17494 );
17495 }
17496
17497 #[tokio::test]
17498 async fn test_pre_response_rejection_does_not_commit_staged_context_or_user() {
17499 let mock = mock_with_response("billing");
17500 let yaml = r#"
17501name: OptimizedStateAgent
17502system_prompt: "You route before answering."
17503runtime:
17504 optimization:
17505 enabled: true
17506 pre_response_deterministic_transitions: true
17507hitl:
17508 states:
17509 billing:
17510 on_enter: require_approval
17511 approval_message: "Approve billing route?"
17512states:
17513 initial: greeting
17514 states:
17515 greeting:
17516 prompt: "Old state prompt."
17517 extract:
17518 - key: topic
17519 description: "Support topic"
17520 transitions:
17521 - to: billing
17522 guard:
17523 context:
17524 topic:
17525 eq: billing
17526 timing: pre_response
17527 run_extractors: true
17528 billing:
17529 prompt: "Billing state."
17530"#;
17531 let agent = AgentBuilder::from_yaml(yaml)
17532 .unwrap()
17533 .llm(Arc::new(mock))
17534 .build()
17535 .unwrap();
17536
17537 let response = agent
17538 .try_pre_response_transition("billing please")
17539 .await
17540 .unwrap();
17541
17542 assert!(response.is_none());
17543 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17544 assert!(!agent.get_context().contains_key("topic"));
17545 assert_eq!(agent.memory.get_messages(None).await.unwrap().len(), 0);
17546 }
17547
17548 #[tokio::test]
17549 async fn test_pre_response_extractor_commits_context_on_winning_path() {
17550 let mock = mock_with_responses(vec!["billing", "Billing response"]);
17551 let yaml = r#"
17552name: OptimizedStateAgent
17553system_prompt: "You route before answering."
17554runtime:
17555 optimization:
17556 enabled: true
17557 pre_response_deterministic_transitions: true
17558states:
17559 initial: greeting
17560 states:
17561 greeting:
17562 prompt: "Old state prompt."
17563 extract:
17564 - key: topic
17565 description: "Support topic"
17566 transitions:
17567 - to: billing
17568 guard:
17569 context:
17570 topic:
17571 eq: billing
17572 timing: pre_response
17573 run_extractors: true
17574 billing:
17575 prompt: "Billing state."
17576"#;
17577 let agent = AgentBuilder::from_yaml(yaml)
17578 .unwrap()
17579 .llm(Arc::new(mock))
17580 .build()
17581 .unwrap();
17582
17583 let response = agent.chat("billing please").await.unwrap();
17584
17585 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17586 assert_eq!(response.content, "Billing response");
17587 assert_eq!(
17588 agent.get_context().get("topic"),
17589 Some(&serde_json::json!("billing"))
17590 );
17591 }
17592
17593 #[tokio::test]
17594 async fn test_pre_response_extractor_miss_does_not_mutate_context() {
17595 let mock = mock_with_response("__NONE__");
17596 let yaml = r#"
17597name: OptimizedStateAgent
17598system_prompt: "You route before answering."
17599runtime:
17600 optimization:
17601 enabled: true
17602 pre_response_deterministic_transitions: true
17603states:
17604 initial: greeting
17605 states:
17606 greeting:
17607 prompt: "Old state prompt."
17608 extract:
17609 - key: topic
17610 description: "Support topic"
17611 transitions:
17612 - to: billing
17613 guard:
17614 context:
17615 topic:
17616 eq: billing
17617 timing: pre_response
17618 run_extractors: true
17619 billing:
17620 prompt: "Billing state."
17621"#;
17622 let agent = AgentBuilder::from_yaml(yaml)
17623 .unwrap()
17624 .llm(Arc::new(mock))
17625 .build()
17626 .unwrap();
17627
17628 let response = agent.try_pre_response_transition("hello").await.unwrap();
17629
17630 assert!(response.is_none());
17631 assert_eq!(agent.current_state().as_deref(), Some("greeting"));
17632 assert!(!agent.get_context().contains_key("topic"));
17633 }
17634
17635 #[tokio::test]
17636 async fn test_default_guard_transition_stays_post_response() {
17637 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
17638 let call_counter = mock.clone();
17639 let yaml = r#"
17640name: TimingAgent
17641system_prompt: "You route carefully."
17642runtime:
17643 optimization:
17644 enabled: true
17645 pre_response_deterministic_transitions: true
17646states:
17647 initial: greeting
17648 states:
17649 greeting:
17650 prompt: "Old state prompt."
17651 transitions:
17652 - to: billing
17653 guard:
17654 context:
17655 topic:
17656 eq: billing
17657 billing:
17658 prompt: "Billing state."
17659"#;
17660 let agent = AgentBuilder::from_yaml(yaml)
17661 .unwrap()
17662 .llm(Arc::new(mock))
17663 .build()
17664 .unwrap();
17665 agent
17666 .set_context("topic", serde_json::json!("billing"))
17667 .unwrap();
17668
17669 let response = agent.chat("billing please").await.unwrap();
17670
17671 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17672 assert_eq!(response.content, "Billing response");
17673 assert_eq!(call_counter.call_count(), 2);
17674 }
17675
17676 #[tokio::test]
17677 async fn test_explicit_post_response_guard_transition_stays_post_response() {
17678 let mock = mock_with_responses(vec!["Greeting response", "Billing response"]);
17679 let call_counter = mock.clone();
17680 let yaml = r#"
17681name: TimingAgent
17682system_prompt: "You route carefully."
17683runtime:
17684 optimization:
17685 enabled: true
17686 pre_response_deterministic_transitions: true
17687states:
17688 initial: greeting
17689 states:
17690 greeting:
17691 prompt: "Old state prompt."
17692 transitions:
17693 - to: billing
17694 guard:
17695 context:
17696 topic:
17697 eq: billing
17698 timing: post_response
17699 billing:
17700 prompt: "Billing state."
17701"#;
17702 let agent = AgentBuilder::from_yaml(yaml)
17703 .unwrap()
17704 .llm(Arc::new(mock))
17705 .build()
17706 .unwrap();
17707 agent
17708 .set_context("topic", serde_json::json!("billing"))
17709 .unwrap();
17710
17711 let response = agent.chat("billing please").await.unwrap();
17712
17713 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17714 assert_eq!(response.content, "Billing response");
17715 assert_eq!(call_counter.call_count(), 2);
17716 }
17717
17718 #[tokio::test]
17719 async fn test_pre_response_extractors_are_transition_scoped() {
17720 let mock = mock_with_responses(vec!["billing", "Billing response"]);
17721 let yaml = r#"
17722name: ScopedExtractorAgent
17723system_prompt: "You route carefully."
17724runtime:
17725 optimization:
17726 enabled: true
17727 pre_response_deterministic_transitions: true
17728states:
17729 initial: greeting
17730 states:
17731 greeting:
17732 prompt: "Old state prompt."
17733 extract:
17734 - key: topic
17735 description: "Support topic"
17736 transitions:
17737 - to: wrong
17738 guard:
17739 context:
17740 topic:
17741 eq: billing
17742 timing: pre_response
17743 - to: billing
17744 guard:
17745 context:
17746 topic:
17747 eq: billing
17748 timing: pre_response
17749 run_extractors: true
17750 wrong:
17751 prompt: "Wrong state."
17752 billing:
17753 prompt: "Billing state."
17754"#;
17755 let agent = AgentBuilder::from_yaml(yaml)
17756 .unwrap()
17757 .llm(Arc::new(mock))
17758 .build()
17759 .unwrap();
17760
17761 let response = agent.chat("billing please").await.unwrap();
17762
17763 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17764 assert_eq!(response.content, "Billing response");
17765 }
17766
17767 #[tokio::test]
17768 async fn test_pre_response_resolved_intent_routes_early() {
17769 let mock = mock_with_response("Billing response");
17770 let yaml = r#"
17771name: IntentAgent
17772system_prompt: "You route carefully."
17773runtime:
17774 optimization:
17775 enabled: true
17776 pre_response_deterministic_transitions: true
17777states:
17778 initial: greeting
17779 states:
17780 greeting:
17781 prompt: "Old state prompt."
17782 transitions:
17783 - to: billing
17784 intent: billing
17785 timing: pre_response
17786 billing:
17787 prompt: "Billing state."
17788"#;
17789 let agent = AgentBuilder::from_yaml(yaml)
17790 .unwrap()
17791 .llm(Arc::new(mock))
17792 .build()
17793 .unwrap();
17794 agent
17795 .set_context("resolved_intent", serde_json::json!("billing"))
17796 .unwrap();
17797
17798 let response = agent
17799 .try_pre_response_transition("I need billing help")
17800 .await
17801 .unwrap()
17802 .unwrap();
17803
17804 assert_eq!(agent.current_state().as_deref(), Some("billing"));
17805 assert_eq!(response.content, "Billing response");
17806 }
17807
17808 #[tokio::test]
17809 async fn test_background_overflow_error_surfaces() {
17810 let mut config = RuntimeConfig::default();
17811 config.optimization.enabled = true;
17812 config.optimization.post_turn.max_background_tasks = 1;
17813 config.optimization.post_turn.on_background_overflow = BackgroundOverflowPolicy::Error;
17814 let policy = crate::optimization::MaintenanceTaskPolicy {
17815 mode: MaintenanceMode::Background,
17816 await_before_next_turn: AwaitBeforeNextTurn::Always,
17817 };
17818 let agent = AgentBuilder::new()
17819 .system_prompt("You are helpful.")
17820 .llm(Arc::new(mock_with_response("ok")))
17821 .build()
17822 .unwrap()
17823 .with_runtime_config(config);
17824 agent
17825 .background_maintenance
17826 .spawn(None, async { std::future::pending::<Result<()>>().await })
17827 .unwrap();
17828
17829 let result = agent
17830 .spawn_or_handle_background(None, async { Ok(()) }, "facts", &policy)
17831 .await;
17832
17833 assert!(result.is_err());
17834 }
17835
17836 #[tokio::test]
17837 async fn test_speculative_reasoning_low_cap_uses_serial_reasoning() {
17838 let default_mock = mock_with_response("Plain draft response");
17839 let router_mock = mock_with_response("cot");
17840 let router_counter = router_mock.clone();
17841 let yaml = r#"
17842name: ReasoningReservationAgent
17843system_prompt: "You answer plainly unless reasoning wins."
17844llm:
17845 default: default
17846 router: router
17847observability:
17848 enabled: true
17849 export:
17850 write_raw_events: true
17851reasoning:
17852 mode: auto
17853 judge_llm: router
17854runtime:
17855 optimization:
17856 enabled: true
17857 max_speculative_llm_calls_per_turn: 1
17858 speculative_reasoning_auto: true
17859 max_parallel_runtime_tasks: 2
17860"#;
17861 let agent = AgentBuilder::from_yaml(yaml)
17862 .unwrap()
17863 .llm_alias("default", Arc::new(default_mock))
17864 .llm_alias("router", Arc::new(router_mock))
17865 .build()
17866 .unwrap();
17867
17868 let response = agent.chat("hello").await.unwrap();
17869
17870 assert_eq!(response.content, "Plain draft response");
17871 assert_eq!(router_counter.call_count(), 1);
17872 let events = agent.observability().unwrap().raw_events();
17873 assert!(!events.iter().any(|event| {
17874 event.dimensions.get("commit_behavior") == Some(&"reasoning_decision".to_string())
17875 }));
17876 }
17877
17878 #[tokio::test]
17879 async fn test_forced_reasoning_skips_plain_speculative_draft() {
17880 let mock = mock_with_response("Reasoned response");
17881 let yaml = r#"
17882name: ForcedReasoningAgent
17883system_prompt: "You reason before answering."
17884observability:
17885 enabled: true
17886 export:
17887 write_raw_events: true
17888reasoning:
17889 mode: cot
17890runtime:
17891 optimization:
17892 enabled: true
17893 max_speculative_llm_calls_per_turn: 2
17894 speculative_state_transitions: true
17895 max_parallel_runtime_tasks: 2
17896states:
17897 initial: triage
17898 states:
17899 triage:
17900 prompt: "Answer from triage."
17901 transitions:
17902 - to: billing
17903 guard:
17904 context:
17905 route:
17906 eq: billing
17907 timing: parallel
17908 billing:
17909 prompt: "Billing state."
17910"#;
17911 let agent = AgentBuilder::from_yaml(yaml)
17912 .unwrap()
17913 .llm(Arc::new(mock))
17914 .build()
17915 .unwrap();
17916
17917 let response = agent.chat("hello").await.unwrap();
17918
17919 assert_eq!(response.content, "Reasoned response");
17920 let events = agent.observability().unwrap().raw_events();
17921 assert!(
17922 !events
17923 .iter()
17924 .any(|event| event.dimensions.contains_key("branch_status"))
17925 );
17926 }
17927
17928 #[tokio::test]
17929 async fn test_speculative_skill_low_cap_uses_serial_skill_route() {
17930 let default_mock = mock_with_response("Skill committed response");
17931 let router_mock = mock_with_response("helper");
17932 let router_counter = router_mock.clone();
17933 let yaml = r#"
17934name: SkillReservationAgent
17935system_prompt: "Use skills when they match."
17936llm:
17937 default: default
17938 router: router
17939observability:
17940 enabled: true
17941 export:
17942 write_raw_events: true
17943runtime:
17944 optimization:
17945 enabled: true
17946 max_speculative_llm_calls_per_turn: 1
17947 speculative_skill_routing: true
17948 max_parallel_runtime_tasks: 2
17949skills:
17950 - id: helper
17951 description: "Answer helper requests"
17952 trigger: "User asks for helper"
17953 steps:
17954 - prompt: "Answer the helper request: {{ user_input }}"
17955"#;
17956 let agent = AgentBuilder::from_yaml(yaml)
17957 .unwrap()
17958 .llm_alias("default", Arc::new(default_mock))
17959 .llm_alias("router", Arc::new(router_mock))
17960 .build()
17961 .unwrap();
17962
17963 let response = agent.chat("please use helper").await.unwrap();
17964
17965 assert_eq!(response.content, "Skill committed response");
17966 assert_eq!(router_counter.call_count(), 1);
17967 let events = agent.observability().unwrap().raw_events();
17968 assert!(
17969 !events
17970 .iter()
17971 .any(|event| event.dimensions.contains_key("branch_status"))
17972 );
17973 }
17974
17975 #[tokio::test]
17976 async fn test_parallel_transition_low_cap_allows_deterministic_route() {
17977 let mock = mock_with_response("unused");
17978 let call_counter = mock.clone();
17979 let yaml = r#"
17980name: ParallelTransitionLowCapAgent
17981system_prompt: "Route before stale responses when safe."
17982runtime:
17983 optimization:
17984 enabled: true
17985 max_speculative_llm_calls_per_turn: 1
17986 speculative_state_transitions: true
17987 max_parallel_runtime_tasks: 2
17988states:
17989 initial: triage
17990 states:
17991 triage:
17992 prompt: "Triage state."
17993 transitions:
17994 - to: billing
17995 guard:
17996 context:
17997 route:
17998 eq: billing
17999 timing: parallel
18000 billing:
18001 prompt: "Billing state."
18002"#;
18003 let agent = AgentBuilder::from_yaml(yaml)
18004 .unwrap()
18005 .llm(Arc::new(mock))
18006 .build()
18007 .unwrap();
18008 agent
18009 .set_context("route", serde_json::json!("billing"))
18010 .unwrap();
18011 agent.update_active_turn_context("billing help", HashMap::new());
18012 assert!(
18013 agent.reserve_active_speculative_llm_call(
18014 RuntimeOptimizationKind::ParallelStateTransition
18015 )
18016 );
18017
18018 let selection = agent
18019 .select_parallel_transition_candidate("billing help")
18020 .await
18021 .unwrap();
18022 agent.end_root_turn();
18023
18024 match selection {
18025 ParallelTransitionSelection::Candidate(candidate) => {
18026 assert_eq!(candidate.target(), "billing");
18027 }
18028 ParallelTransitionSelection::NoMatch => panic!("deterministic route did not match"),
18029 ParallelTransitionSelection::ReservationExhausted => {
18030 panic!("deterministic route consumed LLM budget")
18031 }
18032 }
18033 assert_eq!(call_counter.call_count(), 0);
18034 }
18035
18036 #[tokio::test]
18037 async fn speculative_transition_drops_loser_before_state_actions() {
18038 let lock = Arc::new(tokio::sync::Mutex::new(()));
18039 let first_started = Arc::new(tokio::sync::Notify::new());
18040 let first_dropped = Arc::new(AtomicBool::new(false));
18041 let committed_after_drop = Arc::new(AtomicBool::new(false));
18042 let default = Arc::new(FirstCallLockingProvider {
18043 lock,
18044 first_started: Arc::clone(&first_started),
18045 first_dropped: Arc::clone(&first_dropped),
18046 committed_after_drop: Arc::clone(&committed_after_drop),
18047 calls: AtomicU64::new(0),
18048 });
18049 let router = Arc::new(RoutingAfterProviderStart {
18050 provider_started: first_started,
18051 });
18052 let yaml = r#"
18053name: SpeculativeCancellationAgent
18054system_prompt: "Route before committed work."
18055llm:
18056 default: default
18057 router: router
18058runtime:
18059 optimization:
18060 enabled: true
18061 max_speculative_llm_calls_per_turn: 2
18062 speculative_state_transitions: true
18063 max_parallel_runtime_tasks: 2
18064states:
18065 initial: triage
18066 states:
18067 triage:
18068 prompt: "Triage state."
18069 transitions:
18070 - to: technical
18071 when: "The request needs technical support"
18072 timing: parallel
18073 technical:
18074 prompt: "Technical state."
18075 on_enter:
18076 - prompt: "Prepare technical context."
18077 llm: default
18078 store_as: preparation
18079"#;
18080 let agent = AgentBuilder::from_yaml(yaml)
18081 .unwrap()
18082 .llm_alias("default", default)
18083 .llm_alias("router", router)
18084 .build()
18085 .unwrap();
18086
18087 let response = tokio::time::timeout(
18088 std::time::Duration::from_secs(2),
18089 agent.chat("I cannot log in because of AUTH-17."),
18090 )
18091 .await
18092 .expect("committed work must not wait on the losing provider future")
18093 .unwrap();
18094
18095 assert_eq!(response.content, "Committed technical response.");
18096 assert_eq!(agent.current_state().as_deref(), Some("technical"));
18097 assert!(first_dropped.load(Ordering::SeqCst));
18098 assert!(committed_after_drop.load(Ordering::SeqCst));
18099 }
18100
18101 #[tokio::test]
18102 async fn buffered_transition_drops_stale_stream_before_redispatch() {
18103 use futures::StreamExt;
18104
18105 let lock = Arc::new(tokio::sync::Mutex::new(()));
18106 let stream_started = Arc::new(tokio::sync::Notify::new());
18107 let stream_dropped = Arc::new(AtomicBool::new(false));
18108 let committed_after_drop = Arc::new(AtomicBool::new(false));
18109 let default = Arc::new(BufferedLockingProvider {
18110 lock,
18111 stream_started: Arc::clone(&stream_started),
18112 stream_dropped: Arc::clone(&stream_dropped),
18113 committed_after_drop: Arc::clone(&committed_after_drop),
18114 });
18115 let router = Arc::new(RoutingAfterProviderStart {
18116 provider_started: stream_started,
18117 });
18118 let yaml = r#"
18119name: BufferedCancellationAgent
18120system_prompt: "Hide stale streamed output."
18121llm:
18122 default: default
18123 router: router
18124streaming:
18125 enabled: true
18126 buffer_size: 8
18127runtime:
18128 optimization:
18129 enabled: true
18130 max_speculative_llm_calls_per_turn: 2
18131 speculative_state_transitions: true
18132 streaming_policy: buffer_until_routing_done
18133 max_parallel_runtime_tasks: 2
18134states:
18135 initial: triage
18136 states:
18137 triage:
18138 prompt: "Triage state."
18139 transitions:
18140 - to: technical
18141 when: "The request needs technical support"
18142 timing: parallel
18143 technical:
18144 prompt: "Technical state."
18145"#;
18146 let agent = AgentBuilder::from_yaml(yaml)
18147 .unwrap()
18148 .llm_alias("default", default)
18149 .llm_alias("router", router)
18150 .build()
18151 .unwrap();
18152
18153 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18154 let mut stream = agent
18155 .chat_stream("AUTH-17 needs technical help.")
18156 .await
18157 .unwrap();
18158 let mut content = String::new();
18159 while let Some(chunk) = stream.next().await {
18160 match chunk {
18161 StreamChunk::Content { text } => content.push_str(&text),
18162 StreamChunk::Done {} => break,
18163 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18164 _ => {}
18165 }
18166 }
18167 content
18168 })
18169 .await
18170 .expect("redispatch must not wait on the stale streaming future");
18171
18172 assert_eq!(content, "Committed technical response.");
18173 assert_eq!(agent.current_state().as_deref(), Some("technical"));
18174 assert!(stream_dropped.load(Ordering::SeqCst));
18175 assert!(committed_after_drop.load(Ordering::SeqCst));
18176 }
18177
18178 #[tokio::test]
18179 async fn buffered_transition_drops_established_stream_before_redispatch() {
18180 use futures::StreamExt;
18181
18182 let stream_started = Arc::new(tokio::sync::Notify::new());
18183 let stream_dropped = Arc::new(AtomicBool::new(false));
18184 let stream_dropped_notify = Arc::new(tokio::sync::Notify::new());
18185 let committed_after_drop = Arc::new(AtomicBool::new(false));
18186 let default = Arc::new(EstablishedStreamProvider {
18187 stream_started: Arc::clone(&stream_started),
18188 stream_dropped: Arc::clone(&stream_dropped),
18189 stream_dropped_notify,
18190 committed_after_drop: Arc::clone(&committed_after_drop),
18191 });
18192 let router = Arc::new(RoutingAfterProviderStart {
18193 provider_started: stream_started,
18194 });
18195 let yaml = r#"
18196name: EstablishedStreamCancellationAgent
18197system_prompt: "Hide stale streamed output."
18198llm:
18199 default: default
18200 router: router
18201streaming:
18202 enabled: true
18203 buffer_size: 8
18204runtime:
18205 optimization:
18206 enabled: true
18207 max_speculative_llm_calls_per_turn: 2
18208 speculative_state_transitions: true
18209 streaming_policy: buffer_until_routing_done
18210 max_parallel_runtime_tasks: 2
18211states:
18212 initial: triage
18213 states:
18214 triage:
18215 prompt: "Triage state."
18216 transitions:
18217 - to: technical
18218 when: "The request needs technical support"
18219 timing: parallel
18220 technical:
18221 prompt: "Technical state."
18222"#;
18223 let agent = AgentBuilder::from_yaml(yaml)
18224 .unwrap()
18225 .llm_alias("default", default)
18226 .llm_alias("router", router)
18227 .build()
18228 .unwrap();
18229
18230 let content = tokio::time::timeout(std::time::Duration::from_secs(2), async {
18231 let mut stream = agent
18232 .chat_stream("AUTH-17 needs technical help.")
18233 .await
18234 .unwrap();
18235 let mut content = String::new();
18236 while let Some(chunk) = stream.next().await {
18237 match chunk {
18238 StreamChunk::Content { text } => content.push_str(&text),
18239 StreamChunk::Done {} => break,
18240 StreamChunk::Error { message } => panic!("unexpected stream error: {message}"),
18241 _ => {}
18242 }
18243 }
18244 content
18245 })
18246 .await
18247 .expect("redispatch must wait for the established stale stream to be dropped");
18248
18249 assert_eq!(content, "Committed technical response.");
18250 assert_eq!(agent.current_state().as_deref(), Some("technical"));
18251 assert!(stream_dropped.load(Ordering::SeqCst));
18252 assert!(committed_after_drop.load(Ordering::SeqCst));
18253 }
18254
18255 #[tokio::test]
18256 async fn test_buffered_streaming_transition_reservation_falls_back() {
18257 use futures::StreamExt;
18258
18259 let mock = mock_with_responses(vec![
18260 "Serial streaming response",
18261 "Serial streaming response",
18262 ]);
18263 let router_mock = mock_with_response("1");
18264 let router_counter = router_mock.clone();
18265 let yaml = r#"
18266name: BufferedReservationFallbackAgent
18267system_prompt: "Stream normally if speculative routing cannot be evaluated."
18268llm:
18269 default: default
18270 router: router
18271observability:
18272 enabled: true
18273 export:
18274 write_raw_events: true
18275streaming:
18276 enabled: true
18277 buffer_size: 8
18278runtime:
18279 optimization:
18280 enabled: true
18281 max_speculative_llm_calls_per_turn: 1
18282 speculative_state_transitions: true
18283 streaming_policy: buffer_until_routing_done
18284 max_parallel_runtime_tasks: 2
18285states:
18286 initial: triage
18287 states:
18288 triage:
18289 prompt: "Triage state."
18290 transitions:
18291 - to: billing
18292 guard:
18293 context:
18294 route:
18295 eq: billing
18296 when: "User asks about billing"
18297 timing: parallel
18298 billing:
18299 prompt: "Billing state."
18300"#;
18301 let agent = AgentBuilder::from_yaml(yaml)
18302 .unwrap()
18303 .llm_alias("default", Arc::new(mock))
18304 .llm_alias("router", Arc::new(router_mock))
18305 .build()
18306 .unwrap();
18307
18308 let mut stream = agent.chat_stream("hello").await.unwrap();
18309 let mut content = String::new();
18310 let mut error = None;
18311 while let Some(chunk) = stream.next().await {
18312 match chunk {
18313 StreamChunk::Content { text } => content.push_str(&text),
18314 StreamChunk::Error { message } => error = Some(message),
18315 StreamChunk::Done {} => break,
18316 _ => {}
18317 }
18318 }
18319
18320 assert_eq!(error, None);
18321 assert_eq!(content, "Serial streaming response");
18322 assert_eq!(router_counter.call_count(), 0);
18323 let events = agent.observability().unwrap().raw_events();
18324 assert!(events.iter().any(|event| {
18325 event.dimensions.get("branch_status") == Some(&"cancelled".to_string())
18326 && event.dimensions.get("commit_behavior")
18327 == Some(&"transition_decision".to_string())
18328 }));
18329 }
18330
18331 #[tokio::test]
18332 async fn test_blocking_error_cleanup_resets_root_turn_for_next_chat() {
18333 let mut mock = mock_with_response("Recovered response");
18334 mock.set_error("boom");
18335 let mut handle = mock.clone();
18336 let agent = AgentBuilder::new()
18337 .system_prompt("You are helpful.")
18338 .llm(Arc::new(mock))
18339 .build()
18340 .unwrap();
18341
18342 assert!(agent.chat("first").await.is_err());
18343 handle.clear_error();
18344 let response = agent.chat("second").await.unwrap();
18345
18346 assert_eq!(response.content, "Recovered response");
18347 let messages = agent.memory.get_messages(None).await.unwrap();
18348 let user_count = messages
18349 .iter()
18350 .filter(|message| message.role == ai_agents_core::Role::User)
18351 .count();
18352 assert_eq!(user_count, 2);
18353 }
18354
18355 #[tokio::test]
18356 async fn test_streaming_error_cleanup_resets_root_turn_for_next_chat() {
18357 use futures::StreamExt;
18358
18359 let mut mock = mock_with_response("Recovered response");
18360 mock.set_error("stream boom");
18361 let mut handle = mock.clone();
18362 let agent = AgentBuilder::new()
18363 .system_prompt("You are helpful.")
18364 .llm(Arc::new(mock))
18365 .build()
18366 .unwrap();
18367
18368 let mut stream = agent.chat_stream("first").await.unwrap();
18369 let mut saw_error = false;
18370 while let Some(chunk) = stream.next().await {
18371 if matches!(chunk, StreamChunk::Error { .. }) {
18372 saw_error = true;
18373 }
18374 }
18375 assert!(saw_error);
18376
18377 handle.clear_error();
18378 let response = agent.chat("second").await.unwrap();
18379
18380 assert_eq!(response.content, "Recovered response");
18381 let messages = agent.memory.get_messages(None).await.unwrap();
18382 let user_count = messages
18383 .iter()
18384 .filter(|message| message.role == ai_agents_core::Role::User)
18385 .count();
18386 assert_eq!(user_count, 2);
18387 }
18388
18389 #[tokio::test]
18390 async fn test_buffered_streaming_route_miss_releases_buffer_limit() {
18391 use futures::StreamExt;
18392
18393 let mut mock = mock_with_response("one two three");
18394 mock.set_latency(10);
18395 let yaml = r#"
18396name: BufferedMissAgent
18397system_prompt: "You stream safely."
18398llm:
18399 default: default
18400streaming:
18401 enabled: true
18402 buffer_size: 1
18403runtime:
18404 optimization:
18405 enabled: true
18406 max_speculative_llm_calls_per_turn: 2
18407 speculative_state_transitions: true
18408 streaming_policy: buffer_until_routing_done
18409 max_parallel_runtime_tasks: 2
18410states:
18411 initial: triage
18412 states:
18413 triage:
18414 prompt: "Answer from triage."
18415 transitions:
18416 - to: billing
18417 guard:
18418 context:
18419 route:
18420 eq: billing
18421 timing: parallel
18422 billing:
18423 prompt: "Billing state."
18424"#;
18425 let agent = AgentBuilder::from_yaml(yaml)
18426 .unwrap()
18427 .llm_alias("default", Arc::new(mock))
18428 .build()
18429 .unwrap();
18430
18431 let mut stream = agent.chat_stream("hello").await.unwrap();
18432 let mut content = String::new();
18433 let mut error = None;
18434 while let Some(chunk) = stream.next().await {
18435 match chunk {
18436 StreamChunk::Content { text } => content.push_str(&text),
18437 StreamChunk::Error { message } => error = Some(message),
18438 StreamChunk::Done {} => break,
18439 _ => {}
18440 }
18441 }
18442
18443 assert_eq!(error, None);
18444 assert_eq!(content, "one two three");
18445 }
18446
18447 #[tokio::test]
18448 async fn test_buffered_streaming_main_failure_finalizes_branch() {
18449 use futures::StreamExt;
18450
18451 let mock = mock_with_response("one two");
18452 let mut router_mock = mock_with_response("0");
18453 router_mock.set_latency(50);
18454 let yaml = r#"
18455name: BufferedFailureAgent
18456system_prompt: "You stream safely."
18457llm:
18458 default: default
18459 router: router
18460observability:
18461 enabled: true
18462 export:
18463 write_raw_events: true
18464streaming:
18465 enabled: true
18466 buffer_size: 1
18467runtime:
18468 optimization:
18469 enabled: true
18470 max_speculative_llm_calls_per_turn: 2
18471 speculative_state_transitions: true
18472 streaming_policy: buffer_until_routing_done
18473 max_parallel_runtime_tasks: 2
18474states:
18475 initial: triage
18476 states:
18477 triage:
18478 prompt: "Ask for the category."
18479 transitions:
18480 - to: billing
18481 when: "User asks about billing"
18482 timing: parallel
18483 billing:
18484 prompt: "Billing state."
18485"#;
18486 let agent = AgentBuilder::from_yaml(yaml)
18487 .unwrap()
18488 .llm_alias("default", Arc::new(mock))
18489 .llm_alias("router", Arc::new(router_mock))
18490 .build()
18491 .unwrap();
18492
18493 let mut stream = agent.chat_stream("hello").await.unwrap();
18494 let mut error = String::new();
18495 while let Some(chunk) = stream.next().await {
18496 if let StreamChunk::Error { message } = chunk {
18497 error = message;
18498 }
18499 }
18500
18501 assert!(
18502 error.contains("stream buffer filled"),
18503 "unexpected stream error: {}",
18504 error
18505 );
18506 let events = agent.observability().unwrap().raw_events();
18507 assert!(events.iter().any(|event| {
18508 event.dimensions.get("branch_status") == Some(&"failed".to_string())
18509 && event.dimensions.get("commit_behavior") == Some(&"final_response".to_string())
18510 && event.dimensions.get("optimization")
18511 == Some(&"buffered_streaming_routing".to_string())
18512 }));
18513 }
18514
18515 #[tokio::test]
18516 async fn test_streaming_preflight_does_not_emit_old_state_content() {
18517 use futures::StreamExt;
18518
18519 let mock = mock_with_response("Billing streamed response");
18520 let yaml = r#"
18521name: StreamingOptimizedAgent
18522system_prompt: "You route before streaming."
18523runtime:
18524 optimization:
18525 enabled: true
18526 pre_response_deterministic_transitions: true
18527streaming:
18528 enabled: true
18529states:
18530 initial: greeting
18531 states:
18532 greeting:
18533 prompt: "OLD_STATE_SENTINEL"
18534 transitions:
18535 - to: billing
18536 guard:
18537 context:
18538 topic:
18539 eq: billing
18540 timing: pre_response
18541 billing:
18542 prompt: "Billing state."
18543"#;
18544 let agent = AgentBuilder::from_yaml(yaml)
18545 .unwrap()
18546 .llm(Arc::new(mock))
18547 .build()
18548 .unwrap();
18549 agent
18550 .set_context("topic", serde_json::json!("billing"))
18551 .unwrap();
18552
18553 let mut stream = agent.chat_stream("billing please").await.unwrap();
18554 let mut content = String::new();
18555 while let Some(chunk) = stream.next().await {
18556 match chunk {
18557 StreamChunk::Content { text } => content.push_str(&text),
18558 StreamChunk::Error { message } => panic!("stream error: {}", message),
18559 StreamChunk::Done {} => break,
18560 _ => {}
18561 }
18562 }
18563
18564 assert_eq!(agent.current_state().as_deref(), Some("billing"));
18565 assert!(content.contains("Billing streamed response"));
18566 assert!(!content.contains("OLD_STATE_SENTINEL"));
18567 }
18568
18569 #[tokio::test]
18571 async fn test_integration_state_machine_basic() {
18572 let yaml = r#"
18573name: StateAgent
18574system_prompt: "You are a support agent."
18575states:
18576 initial: greeting
18577 states:
18578 greeting:
18579 prompt: "Welcome the user warmly."
18580 transitions:
18581 - to: support
18582 when: "User needs help"
18583 auto: true
18584 support:
18585 prompt: "Help solve the user's problem."
18586"#;
18587 let mock = mock_with_responses(vec![
18588 "Welcome! How can I help?", "1", "I'll help you with that.", ]);
18592 let builder = AgentBuilder::from_yaml(yaml).unwrap();
18593 let agent = builder.llm(Arc::new(mock)).build().unwrap();
18594
18595 assert_eq!(agent.current_state(), Some("greeting".to_string()));
18596 let _ = agent.chat("I need help").await.unwrap();
18597 }
18600
18601 #[tokio::test]
18603 async fn test_integration_state_on_enter_set_context() {
18604 let yaml = r#"
18605name: ActionAgent
18606system_prompt: "You are helpful."
18607states:
18608 initial: step1
18609 states:
18610 step1:
18611 prompt: "Step 1"
18612 on_exit:
18613 - set_context:
18614 step1_exited: true
18615 transitions:
18616 - to: step2
18617 when: "always"
18618 auto: true
18619 step2:
18620 prompt: "Step 2"
18621 on_enter:
18622 - set_context:
18623 step2_entered: true
18624"#;
18625 let mock = mock_with_responses(vec![
18627 "Processing step 1.",
18628 "0", ]);
18630 let builder = AgentBuilder::from_yaml(yaml).unwrap();
18631 let agent = builder.llm(Arc::new(mock)).build().unwrap();
18632
18633 assert_eq!(agent.current_state(), Some("step1".to_string()));
18634
18635 agent.transition_to("step2").await.unwrap();
18637
18638 assert_eq!(agent.current_state(), Some("step2".to_string()));
18639
18640 let ctx = agent.get_context();
18642 assert_eq!(ctx.get("step1_exited"), Some(&serde_json::json!(true)));
18643 assert_eq!(ctx.get("step2_entered"), Some(&serde_json::json!(true)));
18644 }
18645
18646 #[tokio::test]
18647 async fn state_action_tool_preserves_source_in_stored_record() {
18648 let yaml = r#"
18649name: StateActionToolAgent
18650system_prompt: "You are helpful."
18651tools:
18652 - context_echo
18653states:
18654 initial: idle
18655 states:
18656 idle:
18657 prompt: "Idle"
18658 active:
18659 prompt: "Active"
18660 on_enter:
18661 - set_context:
18662 action_started: true
18663 - tool: context_echo
18664 args: {}
18665"#;
18666 let agent = AgentBuilder::from_yaml(yaml)
18667 .unwrap()
18668 .llm(Arc::new(mock_with_response("unused")))
18669 .tool(Arc::new(ContextEchoTool))
18670 .build()
18671 .unwrap();
18672
18673 agent.transition_to("active").await.unwrap();
18674
18675 let record: ToolExecutionRecord = serde_json::from_value(
18676 agent
18677 .get_context()
18678 .get("last_tool_record")
18679 .cloned()
18680 .expect("successful state action must store its execution record"),
18681 )
18682 .unwrap();
18683 assert!(record.executed);
18684 assert!(record.success);
18685 assert_eq!(record.canonical_id, "context_echo");
18686 assert!(matches!(
18687 &record.source,
18688 ToolCallSource::StateAction {
18689 state: Some(state),
18690 action_index: 1,
18691 } if state == "active"
18692 ));
18693 }
18694
18695 #[tokio::test]
18696 async fn test_ordinary_transition_uses_on_enter_then_on_reenter() {
18697 let yaml = r#"
18698name: OrdinaryLifecycleAgent
18699system_prompt: "You are helpful."
18700states:
18701 initial: intake
18702 regenerate_on_transition: false
18703 states:
18704 intake:
18705 prompt: "Intake"
18706 transitions:
18707 - to: drafting
18708 guard:
18709 context:
18710 route:
18711 eq: drafting
18712 drafting:
18713 prompt: "Drafting"
18714 on_enter:
18715 - set_context:
18716 draft_version: 1
18717 on_reenter:
18718 - set_context:
18719 draft_version: 2
18720 transitions:
18721 - to: review
18722 guard:
18723 context:
18724 route:
18725 eq: review
18726 review:
18727 prompt: "Review"
18728 on_enter:
18729 - set_context:
18730 review_entry: first
18731 transitions:
18732 - to: drafting
18733 guard:
18734 context:
18735 route:
18736 eq: drafting
18737"#;
18738 let agent = AgentBuilder::from_yaml(yaml)
18739 .unwrap()
18740 .llm(Arc::new(mock_with_responses(vec![
18741 "Intake response",
18742 "Draft response",
18743 "Review response",
18744 ])))
18745 .build()
18746 .unwrap();
18747
18748 agent
18749 .set_context("route", serde_json::json!("drafting"))
18750 .unwrap();
18751 agent.chat("Start a draft").await.unwrap();
18752 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18753 assert_eq!(
18754 agent.get_context().get("draft_version"),
18755 Some(&serde_json::json!(1))
18756 );
18757
18758 agent
18759 .set_context("route", serde_json::json!("review"))
18760 .unwrap();
18761 agent.chat("Review this").await.unwrap();
18762 assert_eq!(agent.current_state().as_deref(), Some("review"));
18763 assert_eq!(
18764 agent.get_context().get("review_entry"),
18765 Some(&serde_json::json!("first"))
18766 );
18767
18768 agent
18769 .set_context("route", serde_json::json!("drafting"))
18770 .unwrap();
18771 agent.chat("Revise this").await.unwrap();
18772 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18773 assert_eq!(
18774 agent.get_context().get("draft_version"),
18775 Some(&serde_json::json!(2))
18776 );
18777 }
18778
18779 #[tokio::test]
18780 async fn test_manual_transition_uses_on_enter_then_on_reenter() {
18781 let yaml = r#"
18782name: ManualLifecycleAgent
18783system_prompt: "You are helpful."
18784states:
18785 initial: intake
18786 states:
18787 intake:
18788 prompt: "Intake"
18789 drafting:
18790 prompt: "Drafting"
18791 on_enter:
18792 - set_context:
18793 draft_version: 1
18794 on_reenter:
18795 - set_context:
18796 draft_version: 2
18797 review:
18798 prompt: "Review"
18799"#;
18800 let agent = AgentBuilder::from_yaml(yaml)
18801 .unwrap()
18802 .llm(Arc::new(mock_with_response("unused")))
18803 .build()
18804 .unwrap();
18805
18806 assert!(!agent.get_context().contains_key("draft_version"));
18807 agent.transition_to("drafting").await.unwrap();
18808 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18809 assert_eq!(
18810 agent.get_context().get("draft_version"),
18811 Some(&serde_json::json!(1))
18812 );
18813
18814 agent.transition_to("review").await.unwrap();
18815 agent.transition_to("drafting").await.unwrap();
18816 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18817 assert_eq!(
18818 agent.get_context().get("draft_version"),
18819 Some(&serde_json::json!(2))
18820 );
18821 }
18822
18823 #[tokio::test]
18824 async fn test_timeout_transition_uses_on_enter_then_on_reenter() {
18825 let yaml = r#"
18826name: TimeoutLifecycleAgent
18827system_prompt: "You are helpful."
18828states:
18829 initial: intake
18830 regenerate_on_transition: false
18831 states:
18832 intake:
18833 prompt: "Intake"
18834 max_turns: 1
18835 timeout_to: drafting
18836 drafting:
18837 prompt: "Drafting"
18838 max_turns: 1
18839 timeout_to: review
18840 on_enter:
18841 - set_context:
18842 draft_version: 1
18843 on_reenter:
18844 - set_context:
18845 draft_version: 2
18846 review:
18847 prompt: "Review"
18848 max_turns: 1
18849 timeout_to: drafting
18850 on_enter:
18851 - set_context:
18852 review_entry: first
18853"#;
18854 let agent = AgentBuilder::from_yaml(yaml)
18855 .unwrap()
18856 .llm(Arc::new(mock_with_responses(vec![
18857 "Intake",
18858 "First draft",
18859 "Review",
18860 "Revised draft",
18861 ])))
18862 .build()
18863 .unwrap();
18864
18865 agent.chat("First turn").await.unwrap();
18866 assert_eq!(agent.current_state().as_deref(), Some("intake"));
18867 assert!(!agent.get_context().contains_key("draft_version"));
18868
18869 agent.chat("Second turn").await.unwrap();
18870 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18871 assert_eq!(
18872 agent.get_context().get("draft_version"),
18873 Some(&serde_json::json!(1))
18874 );
18875
18876 agent.chat("Third turn").await.unwrap();
18877 assert_eq!(agent.current_state().as_deref(), Some("review"));
18878 assert_eq!(
18879 agent.get_context().get("review_entry"),
18880 Some(&serde_json::json!("first"))
18881 );
18882
18883 agent.chat("Fourth turn").await.unwrap();
18884 assert_eq!(agent.current_state().as_deref(), Some("drafting"));
18885 assert_eq!(
18886 agent.get_context().get("draft_version"),
18887 Some(&serde_json::json!(2))
18888 );
18889 }
18890
18891 #[tokio::test]
18893 async fn test_integration_process_normalize() {
18894 let yaml = r#"
18895name: ProcessAgent
18896system_prompt: "You are helpful."
18897process:
18898 input:
18899 - type: normalize
18900 config:
18901 trim: true
18902 collapse_whitespace: true
18903"#;
18904 let mock = mock_with_response("Got your message.");
18905 let builder = AgentBuilder::from_yaml(yaml).unwrap();
18906 let agent = builder.llm(Arc::new(mock.clone())).build().unwrap();
18907
18908 let _ = agent.chat(" hello world ").await.unwrap();
18909
18910 let history = mock.call_history();
18912 assert!(!history.is_empty());
18913 let last_call = history.last().unwrap();
18915 let user_msg = last_call
18916 .messages
18917 .iter()
18918 .find(|m| m.role == ai_agents_core::Role::User)
18919 .unwrap();
18920 assert_eq!(user_msg.content, "hello world");
18921 }
18922
18923 #[tokio::test]
18927 async fn test_integration_memory_compression() {
18928 let yaml = r#"
18929name: MemoryAgent
18930system_prompt: "You are helpful."
18931memory:
18932 type: compacting
18933 max_messages: 100
18934 compress_threshold: 5
18935 max_recent_messages: 3
18936 summarize_batch_size: 2
18937"#;
18938 let responses: Vec<&str> = (0..8).map(|_| "Response from assistant.").collect();
18940 let mock = mock_with_responses(responses);
18941 let builder = AgentBuilder::from_yaml(yaml).unwrap();
18942 let agent = builder.llm(Arc::new(mock)).build().unwrap();
18943
18944 for i in 0..6 {
18946 let _ = agent.chat(&format!("Message {}", i)).await.unwrap();
18947 }
18948
18949 let messages = agent.memory.get_messages(None).await.unwrap();
18952 assert!(messages.len() <= 12); }
18956
18957 #[tokio::test]
18959 async fn test_integration_multi_llm_registry() {
18960 let mut mock_default = MockLLMProvider::new("default");
18961 mock_default.set_response("Default LLM response.");
18962 let mut mock_router = MockLLMProvider::new("router");
18963 mock_router.set_response("Router response.");
18964
18965 let agent = AgentBuilder::new()
18966 .system_prompt("You are helpful.")
18967 .llm_alias("default", Arc::new(mock_default))
18968 .llm_alias("router", Arc::new(mock_router))
18969 .build()
18970 .unwrap();
18971
18972 let response = agent.chat("Hello").await.unwrap();
18973 assert_eq!(response.content, "Default LLM response.");
18974 }
18975
18976 #[tokio::test]
18978 async fn test_integration_agent_reset() {
18979 let mock = mock_with_responses(vec!["Hello!", "Hello again!"]);
18980 let agent = AgentBuilder::new()
18981 .system_prompt("You are helpful.")
18982 .llm(Arc::new(mock))
18983 .build()
18984 .unwrap();
18985
18986 let _ = agent.chat("Hi").await.unwrap();
18987 let messages = agent.memory.get_messages(None).await.unwrap();
18988 assert_eq!(messages.len(), 2); agent.reset().await.unwrap();
18991 let messages = agent.memory.get_messages(None).await.unwrap();
18992 assert_eq!(messages.len(), 0);
18993 }
18994
18995 #[tokio::test]
18997 async fn test_integration_process_validate_reject() {
18998 use ai_agents_process::{ProcessConfig, ProcessProcessor};
18999
19000 let validate_config = ai_agents_process::ValidateStage {
19001 id: Some("length_check".to_string()),
19002 condition: None,
19003 config: ai_agents_process::ValidateConfig {
19004 rules: vec![ai_agents_process::ValidationRule::MinLength {
19005 min_length: 10,
19006 on_fail: ai_agents_process::ValidationAction {
19007 action: ai_agents_process::ValidationActionType::Reject,
19008 message: None,
19009 },
19010 }],
19011 ..Default::default()
19012 },
19013 };
19014 let process_config = ProcessConfig {
19015 input: vec![ai_agents_process::ProcessStage::Validate(validate_config)],
19016 ..Default::default()
19017 };
19018 let processor = ProcessProcessor::new(process_config);
19019
19020 let mock = mock_with_response("Should not reach here.");
19021 let agent = AgentBuilder::new()
19022 .system_prompt("You are helpful.")
19023 .llm(Arc::new(mock))
19024 .process_processor(processor)
19025 .build()
19026 .unwrap();
19027
19028 let response = agent.chat("Hi").await.unwrap();
19029 assert!(
19031 response.content.contains("rejected")
19032 || response.content.contains("Input rejected")
19033 || response.content.contains("too short")
19034 || response.content.contains("Too short")
19035 || response.content.len() < 50, "Expected rejection response, got: {}",
19037 response.content
19038 );
19039 }
19040
19041 #[tokio::test]
19043 async fn test_llm_fallback_on_failure() {
19044 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19045
19046 let mut primary = MockLLMProvider::new("primary");
19047 primary.set_error("Primary LLM is unavailable");
19048
19049 let mut fallback = MockLLMProvider::new("fallback");
19050 fallback.set_response("Fallback response works!");
19051
19052 let agent = AgentBuilder::new()
19053 .system_prompt("You are helpful.")
19054 .llm_alias("default", Arc::new(primary))
19055 .llm_alias("backup", Arc::new(fallback))
19056 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19057 llm: LLMRecoveryConfig {
19058 on_failure: LLMFailureAction::FallbackLlm {
19059 fallback_llm: "backup".to_string(),
19060 },
19061 ..Default::default()
19062 },
19063 ..Default::default()
19064 }))
19065 .build()
19066 .unwrap();
19067
19068 let response = agent.chat("Hello").await.unwrap();
19069 assert!(
19070 response.content.contains("Fallback response"),
19071 "Expected fallback response, got: {}",
19072 response.content
19073 );
19074 }
19075
19076 #[tokio::test]
19078 async fn test_llm_fallback_response_static_message() {
19079 use ai_agents_recovery::{ErrorRecoveryConfig, LLMFailureAction, LLMRecoveryConfig};
19080
19081 let mut primary = MockLLMProvider::new("primary");
19082 primary.set_error("Primary LLM is unavailable");
19083
19084 let agent = AgentBuilder::new()
19085 .system_prompt("You are helpful.")
19086 .llm(Arc::new(primary))
19087 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19088 llm: LLMRecoveryConfig {
19089 on_failure: LLMFailureAction::FallbackResponse {
19090 message: "I am temporarily unavailable. Please try again later."
19091 .to_string(),
19092 },
19093 ..Default::default()
19094 },
19095 ..Default::default()
19096 }))
19097 .build()
19098 .unwrap();
19099
19100 let response = agent.chat("Hello").await.unwrap();
19101 assert!(
19102 response.content.contains("temporarily unavailable"),
19103 "Expected static fallback message, got: {}",
19104 response.content
19105 );
19106 }
19107
19108 #[tokio::test]
19110 async fn test_tool_failure_skip() {
19111 use ai_agents_recovery::{
19112 ErrorRecoveryConfig, ToolFailureAction, ToolRecoveryConfig, ToolRetryConfig,
19113 };
19114
19115 let mock = mock_with_responses(vec![
19117 r#"I'll use the nonexistent tool.
19118[TOOL_CALL: {"name": "nonexistent_tool", "arguments": {}}]"#,
19119 "The tool was unavailable, but I can still help you.",
19120 ]);
19121
19122 let agent = AgentBuilder::new()
19123 .system_prompt("You are helpful.")
19124 .llm(Arc::new(mock))
19125 .recovery_manager(RecoveryManager::new(ErrorRecoveryConfig {
19126 tools: ToolRecoveryConfig {
19127 default: ToolRetryConfig {
19128 max_retries: 0,
19129 timeout_ms: None,
19130 on_failure: ToolFailureAction::Skip,
19131 },
19132 ..Default::default()
19133 },
19134 ..Default::default()
19135 }))
19136 .build()
19137 .unwrap();
19138
19139 let response = agent.chat("Use the nonexistent tool").await;
19141 assert!(
19142 response.is_ok(),
19143 "Expected Ok with skip policy, got: {:?}",
19144 response
19145 );
19146 }
19147}